Source code for openmdao.matrices.csr_matrix

"""Define the CSRmatrix class."""
from scipy.sparse import coo_matrix

from openmdao.matrices.coo_matrix import COOMatrix


class CSRMatrix(COOMatrix):
    """
    Sparse matrix in Compressed Row Storage format.

    Parameters
    ----------
    comm : MPI.Comm or <FakeComm>
        Communicator of the top-level system that owns the <Jacobian>.
    is_internal : bool
        If True, this is the int_mtx of an AssembledJacobian.
    """

    def _build(self, num_rows, num_cols, in_ranges, out_ranges):
        """
        Allocate the matrix.

        Parameters
        ----------
        num_rows : int
            number of rows in the matrix.
        num_cols : int
            number of cols in the matrix.
        in_ranges : dict
            Maps input var name to column range.
        out_ranges : dict
            Maps output var name to row range.
        """
        super()._build(num_rows, num_cols, in_ranges, out_ranges)

        self._coo = self._matrix

    def _pre_update(self):
        """
        Do anything that needs to be done at the end of AssembledJacobian._update.
        """
        self._matrix = self._coo

    def _post_update(self):
        """
        Do anything that needs to be done at the end of AssembledJacobian._update.
        """
        # this will add any repeated entries together
        self._matrix = self._coo.tocsc()