Skip to content

Sparse¤

tatva.sparse.ColoredMatrix dataclass ¤

ColoredMatrix(
    data: Array,
    indptr: Array,
    indices: Array,
    shape: tuple[int, int],
    colors: Array,
)

Class to represent the sparsity pattern of a matrix, including the row pointers, column indices, and colors for graph coloring.

Methods:

  • from_csr –

    Create a SparseMatrix instance from a SciPy CSR matrix and optional colors.

  • to_csr –

    Convert the sparse matrix to SciPy's CSR format.

  • to_bcoo –

    Convert the sparse matrix to JAX's BCOO format.

  • to_bcsr –

    Convert the sparse matrix to JAX's BCSR format.

  • to_dense –

    Convert the sparse matrix to a dense array.

Attributes:

  • data (Array) –

    Data values of the sparse matrix

  • indptr (Array) –

    Row pointers for the original sparsity pattern (CSR format)

  • indices (Array) –

    Column indices for the original sparsity pattern (CSR format)

  • shape (tuple[int, int]) –

    Shape of the sparse matrix

  • colors (Array) –

    Colors assigned to each degree of freedom (DOF) for graph coloring

data class-attribute instance-attribute ¤

data: Array = field(repr=False)

Data values of the sparse matrix

indptr class-attribute instance-attribute ¤

indptr: Array = field(repr=False)

Row pointers for the original sparsity pattern (CSR format)

indices class-attribute instance-attribute ¤

indices: Array = field(repr=False)

Column indices for the original sparsity pattern (CSR format)

shape class-attribute instance-attribute ¤

shape: tuple[int, int] = field(metadata=dict(static=True))

Shape of the sparse matrix

colors class-attribute instance-attribute ¤

colors: Array = field(repr=False)

Colors assigned to each degree of freedom (DOF) for graph coloring

from_csr classmethod ¤

from_csr(
    csr_matrix: csr_matrix, colors: NDArray | None = None
) -> Self

Create a SparseMatrix instance from a SciPy CSR matrix and optional colors.

to_csr ¤

to_csr() -> sp.csr_matrix

Convert the sparse matrix to SciPy's CSR format.

to_bcoo ¤

to_bcoo() -> jsp.BCOO

Convert the sparse matrix to JAX's BCOO format.

to_bcsr ¤

to_bcsr() -> jsp.BCSR

Convert the sparse matrix to JAX's BCSR format.

to_dense ¤

to_dense() -> Array

Convert the sparse matrix to a dense array.

tatva.sparse.jacfwd ¤

jacfwd(
    fn: Callable[Concatenate[Array, P], Array],
    colored_matrix: ColoredMatrix,
    *,
    color_batch_size: int | None = None,
) -> Callable[Concatenate[Array, P], ColoredMatrix]

Returns a function that computes the Jacobian of fn using forward-mode automatic differentiation and graph coloring. The returned function takes the same arguments as fn and returns a sparse Jacobian as a new instance of Sparsity.

Parameters:

  • fn ¤

    (Callable[Concatenate[Array, P], Array]) –

    Function for which to compute the Jacobian. Must take an Array as its first argument and return an Array. Will be differentiated with respect to the first argument.

  • colored_matrix ¤

    (ColoredMatrix) –

    An instance of ColoredMatrix representing the sparsity pattern and coloring of the Jacobian.

  • color_batch_size ¤

    (int | None, default: None ) –

    Optional batch size for processing colors. If None, processes all colors at once. If memory usage is a concern, set to a smaller value to process colors in batches.

Returns:

  • Callable[Concatenate[Array, P], ColoredMatrix] –

    A function that computes the sparse Jacobian of fn at a given input, returning

  • Callable[Concatenate[Array, P], ColoredMatrix] –

    a new instance of ColoredMatrix containing the Jacobian values in the data field.

tatva.sparse.linearized_jacfwd ¤

linearized_jacfwd(
    fn: Callable[Concatenate[Array, P], Array],
    colored_matrix: ColoredMatrix,
    *,
    color_batch_size: int | None = None,
) -> Callable[
    Concatenate[Array, P], tuple[Array, ColoredMatrix]
]

Like sparse.jacfwd but uses jax.linearize to avoid redundant forward passes. In general that means the memory usage scales with size of the computation.

Parameters:

  • fn ¤

    (Callable[Concatenate[Array, P], Array]) –

    Function for which to compute the Jacobian. Must take an Array as its first argument and return an Array. Will be differentiated with respect to the first argument.

  • colored_matrix ¤

    (ColoredMatrix) –

    An instance of ColoredMatrix representing the sparsity pattern and coloring of the Jacobian.

  • color_batch_size ¤

    (int | None, default: None ) –

    Optional batch size for processing colors. If None, processes all colors at once. If memory usage is a concern, set to a smaller value to process colors in batches.

Returns:

  • Callable[Concatenate[Array, P], tuple[Array, ColoredMatrix]] –

    a function that computes both the primal values and the sparse Jacobian in a single

  • Callable[Concatenate[Array, P], tuple[Array, ColoredMatrix]] –

    call, sharing the forward pass.

tatva.sparse.create_sparsity_pattern ¤

create_sparsity_pattern(
    mesh: Mesh, n_dofs_per_node: int
) -> sps.csr_matrix

Create a sparsity pattern using SciPy's COO format for efficient setup on CPU.

Parameters:

  • mesh ¤

    (Mesh) –

    Mesh object

  • n_dofs_per_node ¤

    (int) –

    Number of degrees of freedom per node

tatva.sparse.augment_sparsity_with_lifter ¤

augment_sparsity_with_lifter(
    sparsity: csr_matrix, lifter: Lifter
) -> sps.csr_matrix

Augment the sparsity pattern with constraints from a lifter.

Parameters:

  • sparsity ¤

    (csr_matrix) –

    Sparsity pattern in SciPy CSR format.

  • lifter ¤

    (Lifter) –

    Lifter containing constraints.

tatva.sparse.reduce_sparsity_pattern ¤

reduce_sparsity_pattern(
    pattern: csr_matrix, free_dofs: ArrayLike
) -> sps.csr_matrix

Reduce a sparse matrix pattern in CSR format to only the free dofs.

Parameters:

  • pattern ¤

    (csr_matrix) –

    Sparse matrix pattern in CSR format on the full set of dofs.

  • free_dofs ¤

    (ArrayLike) –

    Array of free dofs as integer indices.

tatva.sparse.distance2_colors ¤

distance2_colors(
    row_ptr: Array, col_idx: Array, n_dofs: int
)

Compute distance-2 coloring based on greedy algorithm.

Parameters:

  • row_ptr ¤

    (Array) –

    CSR row pointer array

  • col_idx ¤

    (Array) –

    CSR column indices array

  • n_dofs ¤

    (int) –

    Number of degrees of freedom (size of the matrix)

Returns: colors: Array of colors assigned to each DOF