brahmap.core.InvNoiseCovLO_Toeplitz01

Bases: InvNoiseCovLinearOperator

Linear operator for the inverse of Toeplitz noise covariance

The input covariance array must be at least of the size n. The input power spectrum array must be of the size 2n-2 or 2n-1.

Parameters:

Name	Type	Description	Default
size	int	description	required
input	Union[ndarray, List]	description	required
input_type	Literal['covariance', 'power_spectrum']	description, by default "power_spectrum"	'power_spectrum'
precond_op	Union[LinearOperator, Literal[None, 'Strang', 'TChan', 'RChan', 'KK2']]	description, by default None	None
precond_maxiter	int	description, by default 50	50
precond_rtol	float	description, by default 1.0e-10	1e-10
precond_atol	float	description, by default 1.0e-10	1e-10
precond_callback	Callable	description, by default None	None
dtype	DTypeFloat	description, by default np.float64	float64

Source code in brahmap/core/noise_ops_toeplitz.py

class InvNoiseCovLO_Toeplitz01(InvNoiseCovLinearOperator):
    """Linear operator for the inverse of Toeplitz noise covariance

    The input covariance array must be at least of the size n. The input power
    spectrum array must be of the size 2n-2 or 2n-1.

    Parameters
    ----------
    size : int
        _description_
    input : Union[np.ndarray, List]
        _description_
    input_type : Literal["covariance", "power_spectrum"], optional
        _description_, by default "power_spectrum"
    precond_op : Union[ LinearOperator, Literal[None, "Strang", "TChan", "RChan", "KK2"] ], optional
        _description_, by default None
    precond_maxiter : int, optional
        _description_, by default 50
    precond_rtol : float, optional
        _description_, by default 1.0e-10
    precond_atol : float, optional
        _description_, by default 1.0e-10
    precond_callback : Callable, optional
        _description_, by default None
    dtype : DTypeFloat, optional
        _description_, by default np.float64
    """

    def __init__(
        self,
        size: int,
        input: Union[np.ndarray, List],
        input_type: Literal["covariance", "power_spectrum"] = "power_spectrum",
        precond_op: Union[
            LinearOperator, Literal[None, "Strang", "TChan", "RChan", "KK2"]
        ] = None,
        precond_maxiter: int = 50,
        precond_rtol: float = 1.0e-10,
        precond_atol: float = 1.0e-10,
        precond_callback: Callable = None,
        dtype: DTypeFloat = np.float64,
    ):
        self.__toeplitz_op = NoiseCovLO_Toeplitz01(
            size=size,
            input=input,
            input_type=input_type,
            dtype=dtype,
        )

        self.__precond_rtol = precond_rtol
        self.__precond_atol = precond_atol
        self.__precond_maxiter = precond_maxiter
        self.__precond_callback = precond_callback

        if precond_op is None:
            self.__precond_op = None
        elif isinstance(precond_op, LinearOperator) or isinstance(
            precond_op, np.ndarray
        ):
            self.__precond_op = precond_op
        elif precond_op in ["Strang", "TChan", "RChan", "KK2"]:
            if input_type == "power_spectrum":
                cov = np.fft.ifft(input).real[:size]
            else:
                cov = input[:size]

            if precond_op == "Strang":
                temp_size = int(np.floor(cov.size / 2))
                if cov.size % 2 == 0:
                    new_cov = np.concatenate(
                        [cov[:temp_size], cov[1 : temp_size + 1][::-1]]
                    )
                else:
                    new_cov = np.concatenate(
                        [cov[: temp_size + 1], cov[1 : temp_size + 1][::-1]]
                    )
            elif precond_op == "TChan":
                new_cov = np.empty_like(cov)
                new_cov[0] = cov[0]
                n = cov.size
                for idx in range(1, n):
                    new_cov[idx] = ((n - idx) * cov[idx] + idx * cov[n - idx]) / n
            elif precond_op == "RChan":
                new_cov = np.roll(np.flip(cov), 1)
                new_cov += cov
                new_cov[0] = cov[0]
            elif precond_op == "KK2":  # Circulant but not symmetric
                new_cov = np.roll(np.flip(cov), 1)
                new_cov[0] = 0
                new_cov = cov - new_cov

            self.__precond_op = InvNoiseCovLO_Circulant(
                size=size,
                input=new_cov,
                input_type="covariance",
                dtype=dtype,
            )
        else:
            MPI_RAISE_EXCEPTION(
                condition=True,
                exception=ValueError,
                message="Invalid preconditioner operator provided!",
            )

        super(InvNoiseCovLO_Toeplitz01, self).__init__(
            nargin=size,
            matvec=self._mult,
            input_type=input_type,
            dtype=dtype,
        )

    @property
    def diag(self) -> np.ndarray:
        factor = 1.0
        return factor * np.ones(self.size, dtype=self.dtype)

    def get_inverse(self):
        return self.__toeplitz_op

    def _mult(self, vec: np.ndarray):
        MPI_RAISE_EXCEPTION(
            condition=(len(vec) != self.shape[0]),
            exception=ValueError,
            message=f"Dimensions of `vec` is not compatible with the dimensions of this `InvNoiseCovLO_Toeplitz` instance.\nShape of `InvNoiseCovLO_Toeplitz` instance: {self.shape}\nShape of `vec`: {vec.shape}",
        )

        if vec.dtype != self.dtype:
            if MPI_UTILS.rank == 0:
                warnings.warn(
                    f"dtype of `vec` will be changed to {self.dtype}",
                    TypeChangeWarning,
                )
            vec = vec.astype(dtype=self.dtype, copy=False)

        prod, _ = cg(
            A=self.__toeplitz_op,
            b=vec,
            rtol=self.__precond_rtol,
            atol=self.__precond_atol,
            maxiter=self.__precond_maxiter,
            M=self.__precond_op,
            callback=self.__precond_callback,
            parallel=False,
        )

        return prod