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.pattern_from_energy ¤

pattern_from_energy(
    energy_fn: Callable[Concatenate[Array, P], Array],
    n_dofs: int,
    *static_args,
) -> sps.csr_matrix

Return the sparsity pattern of d²E/du² as a symmetric CSR matrix for a scalar energy function E(u) where u has n_dofs degrees of freedom.

Parameters:

  • energy_fn ¤

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

    scalar JAX array energy function E(u, *static_args) as a function of input variable u and optional static arguments

  • n_dofs ¤

    (int) –

    number of DOFs (integer size of flattened input array u)

  • static_args ¤

    –

    extra args passed to energy_fn, treated as constants

Returns:

  • csr_matrix –

    A symmetric CSR matrix of shape (n_dofs, n_dofs) with binary entries indicating

  • csr_matrix –

    the sparsity pattern of the Hessian d²E/du².

tatva.sparse.pattern_from_virtual_work ¤

pattern_from_virtual_work(
    virtual_work_fn: Callable[Concatenate[Array, P], Array],
    n_dofs: int,
    trial_arg: str,
    test_arg: str,
    *static_args,
) -> sps.csr_matrix

Return the sparsity pattern of the tangent stiffness matrix K = dR/du = d²G/dvdu for a virtual work function virtual_work_fn(*args) as a CSR matrix.

Parameters:

  • virtual_work_fn ¤

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

    scalar JAX array (virtual work) as a function of trial and test variables (e.g., G(u, v, *static_args))

  • n_dofs ¤

    (int) –

    number of DOFs (integer size of flattened input arrays u and v)

  • trial_arg ¤

    (str) –

    parameter name of the trial function in virtual_work_fn

  • test_arg ¤

    (str) –

    parameter name of the test function in virtual_work_fn

  • static_args ¤

    –

    extra arguments (e.g., mesh coordinates, parameters) passed to virtual_work_fn, treated as constants

Returns:

  • csr_matrix –

    A CSR matrix of shape (n_dofs, n_dofs) with binary entries indicating the sparsity

  • csr_matrix –

    pattern of the tangent stiffness matrix K = dR/du = d²G/dvdu, where G is the

  • csr_matrix –

    virtual work and u,v are the trial and test functions respectively.

tatva.sparse.pattern_from_compound ¤

pattern_from_compound(
    compound_cls: type[Compound], block_wise: bool = False
) -> sps.csr_matrix | list[list[sps.csr_matrix]]

Create a sparsity pattern automatically from a Compound class and its attached mesh.

Nodal fields are fully coupled within elements. All other fields (Local, Shared) are only connected to themselves (diagonal entries).

Parameters:

  • compound_cls ¤

    (type[Compound]) –

    The Compound class defining the state layout.

  • block_wise ¤

    (bool, default: False ) –

    If True, return the pattern as a list of lists of sparse matrices corresponding to the compound fields/blocks. Stacked fields are one block.

tatva.sparse.pattern_from_mesh ¤

pattern_from_mesh(
    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