Introducing PorePy model protocols to eliminate excessive type-hinting of mixed-in model members

[vlipovac]

Existing (and newly created) models have an excessive declaration of class members (methods & attributes) for satisfying mypy and helping developers to get the correct mixed-in members.

This drives the numbers of lines of codes up, as well as the maintenance of names, typing of signatures and documentation.
With on-going development and refactoring, we already have multiple inconsistencies of either missing declaration of mixed-in methods, or missing docs of respective declarations.
Also, if a user wants to implement a new model, navigation through the declared members and required mixed-ins is hell.

All that is not required by introducing a Protocol-class, which contains declarations of respective members, but bundled and usable by all models.

As an example for fluid mass balance:
(Motivated by mypy docs)

import porepy as pp

from typing import Callable, Protocol


class PorePyModel(Protocol):

    mobility_discretization: Callable[[list[pp.Grid]], pp.ad.UpwindAd]
    """Description and intra-sphinx link to implementation."""

    advective_flux: Callable[
        [
            list[pp.Grid],
            pp.ad.Operator,
            pp.ad.UpwindAd,
            pp.ad.Operator,
            Callable[[list[pp.MortarGrid]], pp.ad.Operator],
        ],
        pp.ad.Operator,
    ]
    """Description and intra-sphinx link to implementation."""


class MassBalanceEquations(pp.BalanceEquation):

    # HERE we do not need any declaration any more.
    # They are all bundled in the PorePyModel, which can be re-used in other mixins

    def fluid_flux(self: PorePyModel, domains: pp.SubdomainsOrBoundaries) -> pp.ad.Operator:
        ....
        # we have now typed and supported access to the members provided by other model-related classes in porepy
        discr = self.mobility_discretization(domains)
        ....
        flux = self.advective_flux(
            domains, mob_rho, discr, boundary_operator, self.interface_fluid_flux
        )
        flux.set_name("fluid_flux")
        return flux

This example covers covers all models (in the porepy-sense) and I think it is easy to see that it can in principle be applied to every method and attribute, provided by any mixin class.

With Python > 3.12, class typing instead of typing self will be available

class MassBalanceEquations[PorePyModel](pp.BalanceEquation):

    def fluid_flux(self, domains: pp.SubdomainsOrBoundaries) -> pp.ad.Operator:
        ....  # No need to type self
        flux = self.advective_flux(
            domains, mob_rho, discr, boundary_operator, self.interface_fluid_flux
        ) 
        flux.set_name("fluid_flux")
        return flux

Proposal

Since PorePy aims for having a modular multi-physics modelling framework, we can start with implementing a single Protocol-class, which contains a declaration of methods of all constitutive laws.

We can proceed to include also all members provided by single-physics models, like fluid mass balance or energy balance (all-in-one solution).

But we can also design protocols for individual physics. Models using multiple physics, must respectively use combined protocols to have methods from different physics declared as members of their classes.
We would then have a Protocol-class covering the constitutive laws, and a Protocol-class for each physics-block.
Multi-physics models can then be typed with a multi-protocol class. This idea is illustrated by below example:

class ListProtocol(Protocol):

    to_list: Callable[[np.ndarray], list]

class SetProtocol(Protocol):

    to_set: Callable[[np.ndarray], set]


class MultiProtocol(ListProtocol, SetProtocol, Protocol):
    ...


class MyModel:

    def do_something(self: MultiProtocol) -> tuple[list, set]:

        a = np.ones(3)
        return self.to_list(a), self.to_set(a)

[vlipovac]

It is also easy to see that (if required), the constitutive laws can be split by protocols into different families of constitutive laws, relating to different physics.
While it is some additional work to think about the how to define the "families", it will improve the navigation through the framework greatly.

[keileg]

This makes a lot of sense. I think we should explore the possibility of families. This sounds like a user project over summer.