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 Implement an Ad parser class #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. Similar thoughts may apply to the geometry basis functions pp.models.geometry.basis, research is needed.
3. The implementation should change the Ad backend only, no changes in the computational grid.
4. 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 can be done independently of #1179, #1181, but should not be done in parallel.

Task (in roughly prioritized order)

Beta Give feedback

1. Implement general slicer and prolongation classes, following the above outline. These should take matrix sizes and indices as input.
2. In the SubdomainProjection class, replace the current projection matrices with the new slicer and projection classes. This requires defining indices for the operations.
3. Define new mortar projections through the new slicers.
4. Try to make slicers work with the basis functions in the model geometry mixin

[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.