brahmap.core.BlockDiagonalPreconditionerLO

Bases: LinearOperator

A block-diagonal preconditioner operator for iterative map-making solvers.

Computes the standard map-making preconditioner defined as:

\[M_{BD} = (P^T \text{diag}(N)^{-1} P)^{-1}\]

where \(P\) is the pointing matrix and \(N\) is the noise covariance.

Parameters:

Name	Type	Description	Default
processed_samples	ProcessTimeSamples | SharedMemProcessTimeSamples	The pre-processed time samples object containing accumulated map-making weights	required
solver_type	SolverType | None	The map-making solver configuration to use. If None, it falls back to the solver_type of processed_samples, by default None	None
return_copy	bool	If True, the operator application (matvec) is computed locally by each process and it returns the product as a new numpy array. If False and processed_samples is a SharedMemProcessTimeSamples object, the computation is performed only on the node root process using a node-level shared memory buffer, and a reference to this shared memory buffer is returned directly (zero-copy). By default True	True

Attributes:

Name	Type	Description
solver_type	SolverType	The active map-making solver configuration
return_copy	bool	If True, matvec returns a locally computed copy. If False, it returns the raw shared memory buffer reference

Methods:

Name	Description
reset_counters	Resets matrix-vector product counter to zero.
dot	Numpy-like dot() method.
matvec	Matrix-vector multiplication method.
to_array	Returns the dense form of the linear operator as a 2D NumPy array.

Source code in brahmap/core/linearoperators.py

class BlockDiagonalPreconditionerLO(LinearOperator):
    r"""A block-diagonal preconditioner operator for iterative map-making solvers.

    Computes the standard map-making preconditioner defined as:

    $$M_{BD} = (P^T \text{diag}(N)^{-1} P)^{-1}$$

    where $P$ is the pointing matrix and $N$ is the
    noise covariance.

    Parameters
    ----------
    processed_samples : ProcessTimeSamples | SharedMemProcessTimeSamples
        The pre-processed time samples object containing accumulated map-making weights
    solver_type : SolverType | None, optional
        The map-making solver configuration to use. If `None`, it falls
        back to the `solver_type` of `processed_samples`, by default None
    return_copy : bool, optional
        If `True`, the operator application (`matvec`) is computed locally by
        each process and it returns the product as a new numpy array. If
        `False` and `processed_samples` is a
        [`SharedMemProcessTimeSamples`][brahmap.core.SharedMemProcessTimeSamples]
        object, the computation is performed only on the node root process
        using a node-level shared memory buffer, and a reference to this
        shared memory buffer is returned directly (zero-copy). By default `True`

    Attributes
    ----------
    solver_type : SolverType
        The active map-making solver configuration
    return_copy : bool
        If `True`, `matvec` returns a locally computed copy. If `False`, it
        returns the raw shared memory buffer reference
    """

    def __init__(
        self,
        processed_samples: ProcessTimeSamples | SharedMemProcessTimeSamples,
        solver_type: None | SolverType = None,
        return_copy: bool = True,
    ) -> None:
        ### Some of the functionalities of this class are implemented with C++
        ### extensions. A corresponding full Python implementation is provided in
        ### `tests/py_BlkDiagPrecondLO.py` for reference.

        if solver_type is None:
            self.__solver_type = processed_samples.solver_type
        else:
            if int(processed_samples.solver_type) < int(solver_type):
                raise ValueError(
                    "`solver_type` must be lower than or equal to the"
                    "`solver_type` of `processed_samples` object"
                )
            self.__solver_type = solver_type

        self.new_npix = processed_samples.new_npix
        self.size = processed_samples.new_npix * self.solver_type

        self.__return_copy = return_copy

        # If a return copy is requested, there is not need to use shared
        # memory, we can just use standard operations
        self._shmem_mode = not return_copy and isinstance(
            processed_samples, SharedMemProcessTimeSamples
        )
        if self._shmem_mode:
            assert isinstance(processed_samples, SharedMemProcessTimeSamples)
            self.__shared_mem_mgr = processed_samples.shared_mem_manager
            (
                self._node_prod,
                self._win_node_prod,
            ) = self.__shared_mem_mgr.alloc_shared_node(
                self.size,
                processed_samples.dtype_float,
            )

        if self.solver_type == 1:
            self.weighted_counts = processed_samples.weighted_counts  # type: ignore
        else:
            self.weighted_sin_sq = processed_samples.weighted_sin_sq
            self.weighted_cos_sq = processed_samples.weighted_cos_sq
            self.weighted_sincos = processed_samples.weighted_sincos
            self.one_over_determinant = processed_samples.one_over_determinant
            if self.solver_type == 3:
                self.weighted_counts = processed_samples.weighted_counts  # type: ignore
                self.weighted_sin = processed_samples.weighted_sin  # type: ignore
                self.weighted_cos = processed_samples.weighted_cos  # type: ignore

        if self.solver_type == 1:
            super().__init__(
                nargin=self.size,
                nargout=self.size,
                symmetric=True,
                matvec=self._mult_I,
                dtype=processed_samples.dtype_float,
            )
        elif self.solver_type == 2:
            super().__init__(
                nargin=self.size,
                nargout=self.size,
                symmetric=True,
                matvec=self._mult_QU,
                dtype=processed_samples.dtype_float,
            )
        else:
            super().__init__(
                nargin=self.size,
                nargout=self.size,
                symmetric=True,
                matvec=self._mult_IQU,
                dtype=processed_samples.dtype_float,
            )

    def _mult_I(self, vec: npt.NDArray[np.number]) -> npt.NDArray[np.number]:
        r"""Applies the block-diagonal preconditioner for temperature-only
        ($I$) map-making.

        Computes the action of $y = (P^T \text{diag}(N^{-1}) P)^{-1} v$.

        Parameters
        ----------
        vec : npt.NDArray[np.number]
            The input vector `new_npix`

        Returns
        -------
        npt.NDArray[np.number]
            The resulting array of size `new_npix`
        """
        if self._shmem_mode:
            if self.__shared_mem_mgr.node_rank == 0:
                self._node_prod[:] = vec / self.weighted_counts
            self._win_node_prod.Fence(0)
            return self._node_prod
        else:
            prod = vec / self.weighted_counts
            return prod

    def _mult_QU(self, vec: npt.NDArray[np.number]) -> npt.NDArray[np.number]:
        r"""Applies the block-diagonal preconditioner for linear
        polarization ($QU$) map-making.

        Computes the action of $y = (P^T \text{diag}(N^{-1}) P)^{-1} v$.

        Parameters
        ----------
        vec : npt.NDArray[np.number]
            The input vector `2*new_npix`

        Returns
        -------
        npt.NDArray[np.number]
            The resulting array of size `2*new_npix`
        """
        if self._shmem_mode:
            if self.__shared_mem_mgr.node_rank == 0:
                BlkDiagPrecondLO_tools.BDPLO_mult_QU(
                    new_npix=self.new_npix,
                    weighted_sin_sq=self.weighted_sin_sq,
                    weighted_cos_sq=self.weighted_cos_sq,
                    weighted_sincos=self.weighted_sincos,
                    one_over_determinant=self.one_over_determinant,
                    vec=vec,
                    prod=self._node_prod,
                )
            self._win_node_prod.Fence(0)
            return self._node_prod
        else:
            prod = np.zeros(self.size, dtype=self.dtype)
            BlkDiagPrecondLO_tools.BDPLO_mult_QU(
                new_npix=self.new_npix,
                weighted_sin_sq=self.weighted_sin_sq,
                weighted_cos_sq=self.weighted_cos_sq,
                weighted_sincos=self.weighted_sincos,
                one_over_determinant=self.one_over_determinant,
                vec=vec,
                prod=prod,
            )
            return prod

    def _mult_IQU(self, vec: npt.NDArray[np.number]) -> npt.NDArray[np.number]:
        r"""Applies the block-diagonal preconditioner for temperature and
        linear polarization map-making.

        Computes the action of $y = (P^T \text{diag}(N^{-1}) P)^{-1} v$.

        Parameters
        ----------
        vec : npt.NDArray[np.number]
            The input vector `3*new_npix`

        Returns
        -------
        npt.NDArray[np.number]
            The resulting array of size `3*new_npix`
        """
        if self._shmem_mode:
            if self.__shared_mem_mgr.node_rank == 0:
                BlkDiagPrecondLO_tools.BDPLO_mult_IQU(
                    new_npix=self.new_npix,
                    weighted_counts=self.weighted_counts,
                    weighted_sin_sq=self.weighted_sin_sq,
                    weighted_cos_sq=self.weighted_cos_sq,
                    weighted_sincos=self.weighted_sincos,
                    weighted_sin=self.weighted_sin,
                    weighted_cos=self.weighted_cos,
                    one_over_determinant=self.one_over_determinant,
                    vec=vec,
                    prod=self._node_prod,
                )
            self._win_node_prod.Fence(0)
            return self._node_prod
        else:
            prod = np.zeros(self.size, dtype=self.dtype)
            BlkDiagPrecondLO_tools.BDPLO_mult_IQU(
                new_npix=self.new_npix,
                weighted_counts=self.weighted_counts,
                weighted_sin_sq=self.weighted_sin_sq,
                weighted_cos_sq=self.weighted_cos_sq,
                weighted_sincos=self.weighted_sincos,
                weighted_sin=self.weighted_sin,
                weighted_cos=self.weighted_cos,
                one_over_determinant=self.one_over_determinant,
                vec=vec,
                prod=prod,
            )
            return prod

    @property
    def solver_type(self) -> SolverType:
        """The current map-making solver configuration.

        Returns
        -------
        SolverType
            The map-making solver type
        """
        return self.__solver_type

    @property
    def return_copy(self) -> bool:
        """Whether the operator application returns a copy of the shared
        memory buffer.

        Returns
        -------
        bool
            `True` if a copy is returned, `False` otherwise.
        """
        return self.__return_copy

Attributes

dtype: npt.DTypeLike property writable

The data type of the operator.

Returns:

Type	Description
DTypeLike	The NumPy data type of the operator

nargin: int property

Size of the input vector \(x\), i.e. the number of columns of the operator

Returns:

Type	Description
int	The number of input columns

nargout: int property

Size of the output vector \(A(x)\), i.e. the number of rows of the operator

Returns:

Type	Description
int	The number of output rows

symmetric: bool property

Indicates whether the operator is symmetric or not

Returns:

Type	Description
bool	True if symmetric, False otherwise

shape: Tuple[int, int] property

A tuple (nargout, nargin) representing the shape of the operator

Returns:

Type	Description
tuple[int, int]	A tuple (nrows, ncols)

nMatvec: int property

The number of matrix-vector multiplications computed so far

Returns:

Type	Description
int	The number of matrix-vector multiplications performed

T: LinearOperator property

The transpose operator

Returns:

Type	Description
LinearOperator	The transpose of this linear operator

H: LinearOperator property

The adjoint operator

Returns:

Type	Description
LinearOperator	The Hermitian adjoint of this linear operator

solver_type: SolverType property

The current map-making solver configuration.

Returns:

Type	Description
SolverType	The map-making solver type

return_copy: bool property

Whether the operator application returns a copy of the shared memory buffer.

Returns:

Type	Description
bool	True if a copy is returned, False otherwise.

Methods:

reset_counters() -> None

Resets matrix-vector product counter to zero.

Source code in brahmap/base/linop.py

def reset_counters(self) -> None:
    """Resets matrix-vector product counter to zero."""
    self._nMatvec = 0

dot(x) -> npt.NDArray[np.number]

Numpy-like dot() method.

Parameters:

Name	Type	Description	Default
x	Any	The input vector or object to multiply with.	required

Returns:

Type	Description
NDArray[number]	The result of the dot product.

Source code in brahmap/base/linop.py

def dot(self, x) -> npt.NDArray[np.number]:
    """Numpy-like dot() method.

    Parameters
    ----------
    x : Any
        The input vector or object to multiply with.
    Returns
    -------
    npt.NDArray[np.number]
        The result of the dot product.
    """
    return self.__mul__(x)

matvec(x) -> npt.NDArray[np.number]

Matrix-vector multiplication method.

The matvec method encapsulates the matvec routine specified at construct time, to ensure the consistency of the input and output arrays with the operator's shape.

Parameters:

Name	Type	Description	Default
x	NDArray[number]	The input vector \(x\) to be multiplied by the operator	required

Returns:

Type	Description
NDArray[number]	The result of the matrix-vector multiplication \(A(x)\)

Source code in brahmap/base/linop.py

def matvec(self, x) -> npt.NDArray[np.number]:
    """
    Matrix-vector multiplication method.

    The `matvec` method encapsulates the `matvec`
    routine specified at construct time, to ensure the
    consistency of the input and output arrays with the
    operator's shape.

    Parameters
    ----------
    x : npt.NDArray[np.number]
        The input vector $x$ to be multiplied by the operator

    Returns
    -------
    npt.NDArray[np.number]
        The result of the matrix-vector multiplication $A(x)$
    """
    x = np.asanyarray(x, dtype=self.dtype)
    M, N = self.shape

    # check input data consistency
    N = int(N)
    try:
        x = x.reshape(N)
    except ValueError:
        msg = (
            f"The size of the input array is incompatible with the "
            f"dimensions required by the operator\n"
            f"size of the input array: {x.size}\n"
            f"shape of the operator: {self.shape}"
        )
        msg = f"{self.__class__.__name__}: " + msg
        raise ValueError(msg)

    y = self.__matvec(x)

    # check output data consistency
    M = int(M)
    try:
        y = y.reshape(M)
    except ValueError:
        msg = (
            f"The size of the output array is incompatible with the "
            f"dimensions required by the operator\n"
            f"size of the output array: {y.size}\n"
            f"shape of the operator: {self.shape}"
        )
        msg = f"{self.__class__.__name__}: " + msg
        raise ValueError(msg)

    return y

to_array() -> npt.NDArray[np.number]

Returns the dense form of the linear operator as a 2D NumPy array.

Warning

This method first allocates a NumPy array of shape self.shape and data-type self.dtype, and then fills them with numbers. As such, for a large linear operator, it can occupy an enormous amount of memory and crash your system. Don't use it unless you understand the risk!

Returns:

Type	Description
NDArray[number]	The dense 2D array representation of the linear operator

Source code in brahmap/base/linop.py

def to_array(self) -> npt.NDArray[np.number]:
    """Returns the dense form of the linear operator as a 2D NumPy array.

    !!! Warning

        This method first allocates a NumPy array of shape `self.shape`
        and data-type `self.dtype`, and then fills them with numbers. As
        such, for a large linear operator, it can occupy an enormous
        amount of memory and crash your system. Don't use it unless you
        understand the risk!

    Returns
    -------
    npt.NDArray[np.number]
        The dense 2D array representation of the linear operator
    """
    n, m = self.shape
    H = np.empty((n, m), dtype=self.dtype)
    ej = np.zeros(m, dtype=self.dtype)
    for j in range(m):
        ej[j] = 1.0
        H[:, j] = self * ej
        ej[j] = 0.0
    return H