Sparse

tatva.sparse

Functions:

• create_sparsity_pattern –

  Create a sparsity pattern for a given set of elements and constraints.

• create_sparsity_pattern_KKT –

  Create a sparsity pattern for the KKT system.

• reduce_sparsity_pattern –

  Reduce a sparse matrix pattern to only the free dofs (for K_ff).

• get_bc_indices –

  Get the indices of the fixed degrees of freedom.

• distance2_colors –

  Compute distance-2 coloring based on greedy algorithm.

• largest_degree_first_distance2_colors –

  Compute distance-2 coloring based on Largest Degree First (LDF) heuristic.

• smallest_last_distance2_colors –

  Compute distance-2 coloring based on Smallest-Last (SL) ordering.

• jacfwd –

  Compute the sparse Jacobian using forward-mode automatic differentiation and graph coloring.

• jacfwd_with_batch –

  Compute the sparse Jacobian using forward-mode automatic differentiation and graph coloring and provided seeds.

create_sparsity_pattern

create_sparsity_pattern(mesh: Mesh, n_dofs_per_node: int, K_shape: Optional[Tuple[int, int]] = None, constraint_elements: Optional[Array] = None)

Create a sparsity pattern for a given set of elements and constraints.

Parameters:

• mesh

  (Mesh) –

  Mesh object

• n_dofs_per_node

  (int) –

  Number of degrees of freedom per node

• constraint_elements

  (Optional[Array], default: None ) –

  Optional array of constraint elements. If provided, the sparsity pattern will be created for the constraint elements.

Returns:

• sparsity_pattern –

  jax.experimental.sparse.BCOO

create_sparsity_pattern_KKT

create_sparsity_pattern_KKT(mesh: Mesh, n_dofs_per_node: int, B: Array)

Create a sparsity pattern for the KKT system. Args: mesh: Mesh object n_dofs_per_node: Number of degrees of freedom per node B: Constraint matrix (nb_cons, n_dofs) Returns: sparsity_pattern_KKT: jax.experimental.sparse.BCOO

reduce_sparsity_pattern

reduce_sparsity_pattern(pattern: BCOO, free_dofs: Array) -> BCOO

Reduce a sparse matrix pattern to only the free dofs (for K_ff).

Parameters:

• pattern

  (BCOO) –

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

• free_dofs

  (Array) –

  Array of free dofs as integer indices.

Returns:

• BCOO ( BCOO ) –

  Reduced sparse matrix pattern with rows and columns remapped to the reduced indexing of free dofs.

get_bc_indices

get_bc_indices(sparsity_pattern: BCOO, fixed_dofs: Array)

Get the indices of the fixed degrees of freedom. Args: sparsity_pattern: jax.experimental.sparse.BCOO fixed_dofs: (num_fixed_dofs,) Returns: zero_indices: (num_zero_indices,) one_indices: (num_one_indices,)

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

largest_degree_first_distance2_colors

largest_degree_first_distance2_colors(row_ptr: Array, col_idx: Array, n_dofs: int) -> Array

Compute distance-2 coloring based on Largest Degree First (LDF) heuristic. Args: row_ptr: CSR row pointer array col_idx: CSR column indices array n_dofs: Number of degrees of freedom (size of the matrix) Returns: colors: Array of colors assigned to each DOF

smallest_last_distance2_colors

smallest_last_distance2_colors(row_ptr: Array, col_idx: Array, n_dofs: int) -> Array

Compute distance-2 coloring based on Smallest-Last (SL) ordering. Args: row_ptr: CSR row pointer array col_idx: CSR column indices array n_dofs: Number of degrees of freedom (size of the matrix) Returns: colors: Array of colors assigned to each DOF

jacfwd

jacfwd(gradient: Callable, row_ptr: Array, col_indices: Array, colors: Array, has_aux_args: bool = False) -> Callable

Compute the sparse Jacobian using forward-mode automatic differentiation and graph coloring. The seeds are reconstructed on-the-fly to save memory. Use jax.lax.scan to iterate over colors without materializing the entire compressed Jacobian in memory at once.

Parameters:

• gradient

  (Callable) –

  Function whose Jacobian is to be computed

• row_ptr, col_indices

  –

  The sparsity pattern of the matrix

• colors

  (Array) –

  The array of colors used for compression

• has_aux_args

  (bool, default: False ) –

  Whether the gradient function has auxiliary arguments. If True, the returned function will expect additional arguments after 'u'.

Returns: A function that computes the sparse Jacobian in BCOO format

jacfwd_with_batch

jacfwd_with_batch(gradient: Callable, row_ptr: Array, col_indices: Array, colors: Array, color_batch_size: int = 1, has_aux_args: bool = False) -> Callable

Compute the sparse Jacobian using forward-mode automatic differentiation and graph coloring and provided seeds. Uses jax.lax.map to iterate over colors without materializing the entire compressed Jacobian in memory at once.

Parameters:

• gradient

  (Callable) –

  Function whose Jacobian is to be computed

• row_ptr, col_indices

  –

  The sparsity pattern of the matrix

• seeds

  –

  List of seed vectors for each color

• colors

  (Array) –

  The array of colors used for compression

• color_batch_size

  (int, default: 1 ) –

  Number of colors to process in each batch (1 for minimal memory usage, >1 for faster computation but higher memory)

• has_aux_args

  (bool, default: False ) –

  Whether the gradient function has auxiliary arguments. If True, the returned function will expect additional arguments after 'u'.

Returns: A function that computes the sparse Jacobian in BCOO format