Implement an Ad parser class

[keileg]

Background

In the current code, the evaluation of a pp.ad.Operator is implemented in the Operator class itself. The parsing of a tree of an operator (which is really a computational graph) is implemented as a depth-first traversal which identifies leaves in the tree, parses these to numerical values, and then propagates the values through the graph in a bottom-up fashion. While this works, there are shortcomings with the approach:

1. The graph traversal method is tailored for evaluation of numerical values. While other relevant graph operations, such as compatibility checks between operator sizes, or consistency checks of units, would require a similar traversal, the current code would necessitate the writing of a new, but very similar, traversal function for each new operation.
2. There is no way in which we can identify identical nodes, or subgraphs, and cache evaluated values to save computational time.
3. The top-level Operator class is now responsible both for representing combinations of operations and the evaluation of computational graphs, thus violating the single responsibility principle. While this is not a major problem by itself, it is a signal that the current design is not ideal.

Suggested changes

The following large-scale changes are suggested - specifications are needed before implementation:

1. Implement a function for traversal of a computational graph. To EK, it seems natural to convert the tree into a priority queue, but the choice of data structure may not be a primary concern. A rough signature for the function wolud be def graph_traversal(list[pp.ad.Operator]) -> Iterator[pp.ad.Operator]. The motivation for accepting a list of operators is given in the outline of the Ad parser class below.
2. Implement a class, say AdParser, which takes over responsibility for parsing from the Operator class. This will use the graph traversal function to iterate over the computational graph and convert it to numerical values. Moreover, the parser should contain functionality for chaching of operator values to reduce evaluation time.

Thoughts on the AdParser class

Below is an outline of the class:

class AdParser:

    mdg: MixedDimensionalGrid
   # Parsing of Operator leaves will require access to the grid. This is currently obtained
   # by passing the EquationSystem, but direct access is preferable.

    _cache_table_for_current_value: dict[pp.ad.Operator, np.ndarray]
    _cache_table_for_current_value_and_jacobian: dict[pp.ad.Operator, pp.ad.AdArray]
    # Temporarily store computed values here during evaluation. Not sure if we need separate
    # tables for numpy and Ad arrays

     _longer_term_caching_of_values: dict[pp.ad.Operator, np.ndarray]
    ...
    # Operators that will not (are not likely) to change between evaluation steps, and which are
    # costly to compute, can be stored here inbetween evaluations. This should not be part of
    # the initial implementation.

    def value(self, list[pp.ad.Operator]) -> list[np.ndarray]:
        # Fetch a joint iterator for the operators. Evaluation has two steps: First check if this operator
        # has previously been evaluated (check if it is found in the cache tables, which will use the new
        # hash functions), compute if not available.
        # The signature is non-trivial: We need to accept a list to utilize caching between evaluation of
        # multiple equations, but when iterating, we also need to keep track of which operator tree the
        # individual operators are associated with. Achiving this may require tweaks to the outlined implementation.
        # The need to cache between equations also show that parsing must be moved out of the
        # individual operators.

    def value_and_jacobian():
        # Similar to above

Access to the AdParser

While it is tempting to make the AdParser a singleton accessible as pp.AdParser or similar, there is no clear reason to do so. Instead, it is suggested that all models are equipped with an AdParser, just as they have an EquationSystem. It should also be possible to parse outside a model, by a call pp.ad.AdParser(mdg).value_and_jacobian([operator]), though such operations will not be very common.

The role of the EquationSystem

While the functionality of the AdParser could have been integrated into the EquationSystem, it is preferable not to add more functionality into this (arguably already overloaded) class. Instead, the EquationSystem can be given the model-wide AdParser as an attribute and call this for evaluation of equations.

Dependencies

This should not be started until completion of #1179 and #1180

[IvarStefansson]

Will there be functionality reuse between value and value_and_jacobian? If yes, how do you envision this? Also, will the AdParser be the natural candidate to do consistency checks on dimension compatibility?

[keileg]

Note to self: Operator evaluation will also be needed as part of the solution process (test case: the line search methods being introduced in #1208). We need to make sure that the AdParser class can be invoked also in these cases.

[keileg]

NOTE: For debugging purposes, it should always be possible to switch off all caching.