Reimplementation of Ad projection operators

[keileg]

Background

In the Ad operator framework, the classes SubdomainProjection, MortarProjection, Trace and InvTrace all represent mappings betweeen grid entities of various kinds. These are all implemented as projection matrices that, during parsing of an operator tree, are left multiplied with arrays and matrices. This is unsatisfactory for two reasons:

1. Constructing the projection matrices is slow and adds significantly to the computational cost of working with the Ad operators. As a partial remedy, some if not all of the projection classes construct the projection matrices during initializationof the projection object, but this to a large degree just moves the computational burden. The ahead-of-time construction will also be hard to make compatible with the caching outlined in #1181.
2. Projection operations are in reality slicing (for a global-to-local restriction) or filling in of selected rows in large arrays and matrices (for local-to-global prolongation). Representing this as matrix-matrix / matrix-vector products likely adds significantly to the computational cost.

Suggested change

The projection methods, say SubdomainProjection.face_restriction(), now return pp.ad.SparseArrayss. Instead, it should return a new slicer-object that roughly looks like this:

class RestrictionSlicer:

    self._target_indices: np.ndarray
    # EK note to self: In the future, this may also be a PETSc IS object

    def __init__(self, int: dim, mdg: MixedDimensionalGrid, domain: list[pp.Grid]) -> None:
        # self._target_indices is computed here.
        # What to do with mortar projections is not clear, but similar logic should apply

    def __matmul__(self, other: pp.ad.AdArray | np.ndarray) -> pp.ad.AdArray | np.ndarray:
        return other[self._target_ind]

class ProjectionSlicer:

    self._target_indices: np.ndarray
    self._target_num_rows: int

    def __init__(self, int: dim, mdg: MixedDimensionalGrid, domain: list[pp.Grid]) -> None:
        # self._target_indices and self._target_num_rows are computed here

    def __matmul__(self, other: pp.ad.AdArray | np.ndarray) -> pp.ad.AdArray | np.ndarray:
        if isinstance(other, np.ndarray):
            result = np.ndarray(self._target_size, dtype=other.dtype)
            result[self._target_ind] = other
        else:  # pp.ad.AdArray
            res_val = np.ndarray([self._target_size, dtype=other.val.dtype)
            res_jac = sps.csr_matrix((self._target_size, other.jac.shape[1]))
            res_val[self._target_ind] = other.val
            res_jac[self._target_ind] = other.jac
            return pp.ad.AdArray(res_val, res_jac)

Comments:

1. It is not clear whether we need one slicer each for SubdomainProjection, MortarProjection etc.
2. The implementation should change the Ad backend only, no changes in the computational grid.
3. Testing will be needed.

Regarding the divergence class

The class pp.ad.Divergence represents a differential operator rather than a mapping. This cannot readily be implemented as slicing, and although improvements should be possible there as well, it will not be part of this issue.

Dependencies

This should be undertaken after #1181.

[keileg]

Some of the functionality in po.models.geometry may also be suited for this approach.

[IvarStefansson]

Some of the functionality in po.models.geometry may also be suited for this approach.

That's an interesting comment. Is there any chance this approach can help us with representation of tensors and their products?

[keileg]

Perhaps, it depends on what you mean. I suggest we discuss in person.