Skip to content

Utils¤

tatva.utils.virtual_work_to_residual ¤

virtual_work_to_residual(
    fn=None,
    /,
    *,
    test_arr: ArrayLike | None = None,
    test_shape: tuple[int, ...] | None = None,
    test_size: int | None = None,
    jit: bool = False,
) -> Callable

Convert a virtual-work function that is linear wrt. the test function (first argument) into a residual function.

The returned callable evaluates the Jacobian of fn with respect to its first argument at test_arr. This utility can be used directly or as a decorator. Only one of the arguments test_arr, test_shape, or test_size should be provided.

Parameters:

  • fn ¤

    –

    Function whose first argument is the test array. If None, the function returns a decorator.

  • test_arr ¤

    (ArrayLike | None, default: None ) –

    Test array to use for the Jacobian evaluation.

  • test_shape ¤

    (tuple[int, ...] | None, default: None ) –

    Shape of the test array to use for the Jacobian evaluation. Will use jnp.zeros(test_shape) as the test array.

  • test_size ¤

    (int | None, default: None ) –

    Size of the test array to use for the Jacobian evaluation. Will use jnp.zeros(test_size) as the test array.

  • jit ¤

    (bool, default: False ) –

    Whether to JIT-compile the generated residual function (Default: False).

Returns:

  • Callable ( Callable ) –

    A residual function with the same remaining signature as fn.

tatva.utils.make_project_function ¤

make_project_function(
    nnodes: int,
    colored_matrix: ColoredMatrix | None = None,
    elements: ArrayLike | None = None,
    lifter: Lifter | None = None,
) -> Callable[[Operator, Array], Array]

Factory to create a projection function for a given mesh (defined by nnodes, elements/ColoredMatrix) and optional lifter. The returned function takes an Operator and a field to project, and returns the projected field on the nodes.

Automatically batches the projection if the field has multiple components (e.g., vector or tensor fields) and handles lifting if a lifter is provided.

Parameters:

  • nnodes ¤

    (int) –

    Number of nodes in the mesh.

  • colored_matrix ¤

    (ColoredMatrix | None, default: None ) –

    Optional ColoredMatrix defining the sparsity pattern for the projection. If not provided, it will be created from the elements.

  • elements ¤

    (ArrayLike | None, default: None ) –

    Optional array defining the mesh elements, used to create the ColoredMatrix if not provided.

  • lifter ¤

    (Lifter | None, default: None ) –

    Optional Lifter object to handle lifting of the solution from the reduced system to the full system. If provided, the projection will be performed on the reduced system and then lifted back to the full system.

tatva.utils.create_g2l ¤

create_g2l(l2g: NDArray) -> Callable[[NDArray], NDArray]