Skip to content

Abstract Models

popinn.AbstractModel ¤

Abstract base class for all models.

The __call__ and D (derivative) methods are concrete and should not be overridden by subclasses.

Subclasses must implement the abstract hook _eval(coords, aux_inputs), which __call__ dispatches to.

Source code in src/popinn/models.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
class AbstractModel(eqx.Module):
    """Abstract base class for all models.

    The `__call__` and `D` (derivative) methods are concrete and should not be overridden by subclasses.

    Subclasses must implement the abstract hook `_eval(coords, aux_inputs)`, which `__call__` dispatches to.
    """

    @abc.abstractmethod
    def _eval(
        self,
        coords: Float[Array, "num_coords"],
        aux_inputs: tuple,
    ) -> Float[Array, ""]:
        """Abstract method.

        Evaluates the model at one grid point.

        Args:
            coords (Float[Array, 'num_coords']): Stacked coordinate array,
                one entry per coordinate axis (e.g. `jnp.array([x, t])`).
            aux_inputs (tuple): Auxiliary inputs (e.g. PDE
                parameters, sensor-sampled functions, etc.). Empty for models
                with no auxiliary inputs (PINN).

        Returns:
            (Float[Array, '']): The model output as a scalar JAX array.
        """
        raise NotImplementedError

    def __call__(self, *args):
        """Evaluate the model at one point.

        Coordinates are passed as individual scalars; auxiliary inputs are
        passed as a single trailing *tuple*:

            model(x, t, (a, b))   # x & t are the coordinates; a & b are the auxiliary inputs.

        For models with no auxiliary inputs, e.g. `popinn.PINN`, the auxiliary tuple
        can be empty, or omitted completely:

            model(x, t, ())       # explicit empty aux -- equivalent to:
            model(x, t)           # no aux

        Note the auxiliary inputs *must be a tuple* specifically, not a list
        or array.

        Args:
            *args: The coordinate scalars, optionally followed by the aux
                tuple.

        Returns:
            (Float[Array, '']): The model output as a scalar JAX array.
        """
        if args and isinstance(args[-1], tuple):
            *coords, aux_inputs = args
        else:
            coords, aux_inputs = args, ()
        return self._eval(jnp.stack(coords), aux_inputs)

    def D(self, *argnums):
        """Build the derivative of the model with respect to one or more coordinates.

        Returns a function with the same call signature as the model whose
        output is the requested derivative.

        !!! Example "Derivative syntax"
            Each entry of `argnums` indexes a coordinate argument (`0` -> first coordinate, `1` -> second, ...);
            chaining differentiates repeatedly. For example, for a model with two coordinates `x` & `t`,
            derivatives are taken & evaluated like:
            ```python
                du_dx  = model.D(0)(x, t, aux)         # du/dx
                d2u_dx2 = model.D(0, 0)(x, t, aux)     # d2/dx2
                du_dt  = model.D(1)(x, t, aux)         # d/dt


            ```

        Args:
            *argnums (int): Coordinate indices to differentiate with respect
                to, applied left to right.

        Returns:
            (Callable): A function with the same signature as the model that returns the
                requested scalar derivative.
        """
        f = self.__call__
        for a in argnums:
            f = jax.grad(f, argnums=a)
        return f
_eval(coords, aux_inputs) ¤

Abstract method.

Evaluates the model at one grid point.

Parameters:

Name Type Description Default
coords jaxtyping.Float[jax.Array, num_coords]

Stacked coordinate array, one entry per coordinate axis (e.g. jnp.array([x, t])).

 required
aux_inputs tuple

Auxiliary inputs (e.g. PDE parameters, sensor-sampled functions, etc.). Empty for models with no auxiliary inputs (PINN).

 required

Returns:

Type Description
jaxtyping.Float[jax.Array, '']

The model output as a scalar JAX array.
Source code in src/popinn/models.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
@abc.abstractmethod
def _eval(
    self,
    coords: Float[Array, "num_coords"],
    aux_inputs: tuple,
) -> Float[Array, ""]:
    """Abstract method.

    Evaluates the model at one grid point.

    Args:
        coords (Float[Array, 'num_coords']): Stacked coordinate array,
            one entry per coordinate axis (e.g. `jnp.array([x, t])`).
        aux_inputs (tuple): Auxiliary inputs (e.g. PDE
            parameters, sensor-sampled functions, etc.). Empty for models
            with no auxiliary inputs (PINN).

    Returns:
        (Float[Array, '']): The model output as a scalar JAX array.
    """
    raise NotImplementedError
__call__(*args) ¤

Evaluate the model at one point.

Coordinates are passed as individual scalars; auxiliary inputs are passed as a single trailing tuple:

model(x, t, (a, b))   # x & t are the coordinates; a & b are the auxiliary inputs.

For models with no auxiliary inputs, e.g. popinn.PINN, the auxiliary tuple can be empty, or omitted completely:

model(x, t, ())       # explicit empty aux -- equivalent to:
model(x, t)           # no aux

Note the auxiliary inputs must be a tuple specifically, not a list or array.

Parameters:

Name Type Description Default
*args

The coordinate scalars, optionally followed by the aux tuple.

 ()

Returns:

Type Description
jaxtyping.Float[jax.Array, '']

The model output as a scalar JAX array.
Source code in src/popinn/models.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def __call__(self, *args):
    """Evaluate the model at one point.

    Coordinates are passed as individual scalars; auxiliary inputs are
    passed as a single trailing *tuple*:

        model(x, t, (a, b))   # x & t are the coordinates; a & b are the auxiliary inputs.

    For models with no auxiliary inputs, e.g. `popinn.PINN`, the auxiliary tuple
    can be empty, or omitted completely:

        model(x, t, ())       # explicit empty aux -- equivalent to:
        model(x, t)           # no aux

    Note the auxiliary inputs *must be a tuple* specifically, not a list
    or array.

    Args:
        *args: The coordinate scalars, optionally followed by the aux
            tuple.

    Returns:
        (Float[Array, '']): The model output as a scalar JAX array.
    """
    if args and isinstance(args[-1], tuple):
        *coords, aux_inputs = args
    else:
        coords, aux_inputs = args, ()
    return self._eval(jnp.stack(coords), aux_inputs)
D(*argnums) ¤

Build the derivative of the model with respect to one or more coordinates.

Returns a function with the same call signature as the model whose output is the requested derivative.

Derivative syntax

Each entry of argnums indexes a coordinate argument (0 -> first coordinate, 1 -> second, ...); chaining differentiates repeatedly. For example, for a model with two coordinates x & t, derivatives are taken & evaluated like: 

    du_dx  = model.D(0)(x, t, aux)         # du/dx
    d2u_dx2 = model.D(0, 0)(x, t, aux)     # d2/dx2
    du_dt  = model.D(1)(x, t, aux)         # d/dt

Parameters:

Name Type Description Default
*argnums int

Coordinate indices to differentiate with respect to, applied left to right.

 ()

Returns:

Type Description
collections.abc.Callable

A function with the same signature as the model that returns the requested scalar derivative.
Source code in src/popinn/models.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def D(self, *argnums):
    """Build the derivative of the model with respect to one or more coordinates.

    Returns a function with the same call signature as the model whose
    output is the requested derivative.

    !!! Example "Derivative syntax"
        Each entry of `argnums` indexes a coordinate argument (`0` -> first coordinate, `1` -> second, ...);
        chaining differentiates repeatedly. For example, for a model with two coordinates `x` & `t`,
        derivatives are taken & evaluated like:
        ```python
            du_dx  = model.D(0)(x, t, aux)         # du/dx
            d2u_dx2 = model.D(0, 0)(x, t, aux)     # d2/dx2
            du_dt  = model.D(1)(x, t, aux)         # d/dt


        ```

    Args:
        *argnums (int): Coordinate indices to differentiate with respect
            to, applied left to right.

    Returns:
        (Callable): A function with the same signature as the model that returns the
            requested scalar derivative.
    """
    f = self.__call__
    for a in argnums:
        f = jax.grad(f, argnums=a)
    return f

popinn.AbstractP2INN ¤

Abstract Parameterized Physics-Informed Neural Network.

Composes a parameter encoder, a coordinate encoder, and a manifold network: the coordinates and parameters are encoded separately, their embeddings concatenated, and the manifold network maps to a scalar. Subclasses must provide the three sub-networks as fields.

Source code in src/popinn/models.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
class AbstractP2INN(AbstractModel):
    """Abstract Parameterized Physics-Informed Neural Network.

    Composes a parameter encoder, a coordinate encoder, and a manifold
    network: the coordinates and parameters are encoded separately, their
    embeddings concatenated, and the manifold network maps to a
    scalar. Subclasses must provide the three sub-networks as fields.
    """

    param_encoder: eqx.AbstractVar[eqx.nn.MLP]
    coord_encoder: eqx.AbstractVar[eqx.nn.MLP]
    manifold: eqx.AbstractVar[eqx.nn.MLP]

    def _eval(self, coords: Float[Array, "num_coords"], aux_inputs: tuple) -> Float[Array, ""]:
        """Concatenate coordinate and parameter embeddings, then pass through manifold network.

        Args:
            coords (Float[Array, 'num_coords']): Array of coordinate values, one coordinate per axis, e.g. `jnp.array([x,t])` for scalars `x` and `t`.
            aux_inputs (tuple): Tuple of scalar PDE parameters, e.g. `(a,b)` for scalar values `a` and `b`.

        Returns:
            (Float[Array, '']): The model output as a scalar JAX array.
        """

        h_param = self.param_encoder(jnp.stack(aux_inputs))
        h_coord = self.coord_encoder(coords)
        h_concat = jnp.concatenate([h_param, h_coord])
        return self.manifold(h_concat)
_eval(coords, aux_inputs) ¤

Concatenate coordinate and parameter embeddings, then pass through manifold network.

Parameters:

Name Type Description Default
coords jaxtyping.Float[jax.Array, num_coords]

Array of coordinate values, one coordinate per axis, e.g. jnp.array([x,t]) for scalars x and t.

 required
aux_inputs tuple

Tuple of scalar PDE parameters, e.g. (a,b) for scalar values a and b.

 required

Returns:

Type Description
jaxtyping.Float[jax.Array, '']

The model output as a scalar JAX array.
Source code in src/popinn/models.py 
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def _eval(self, coords: Float[Array, "num_coords"], aux_inputs: tuple) -> Float[Array, ""]:
    """Concatenate coordinate and parameter embeddings, then pass through manifold network.

    Args:
        coords (Float[Array, 'num_coords']): Array of coordinate values, one coordinate per axis, e.g. `jnp.array([x,t])` for scalars `x` and `t`.
        aux_inputs (tuple): Tuple of scalar PDE parameters, e.g. `(a,b)` for scalar values `a` and `b`.

    Returns:
        (Float[Array, '']): The model output as a scalar JAX array.
    """

    h_param = self.param_encoder(jnp.stack(aux_inputs))
    h_coord = self.coord_encoder(coords)
    h_concat = jnp.concatenate([h_param, h_coord])
    return self.manifold(h_concat)

popinn.AbstractDeepONet ¤

Abstract Deep Operator Network.

Evaluates a sum over a product of branch and trunk embeddings: the trunk encodes the coordinates, each branch encodes one auxiliary input (e.g., a sensor-sampled function), and their elementwise product is summed and offset by a bias. Subclasses must provide the branch list, trunk, and bias as fields.

Source code in src/popinn/models.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
class AbstractDeepONet(AbstractModel):
    """Abstract Deep Operator Network.

    Evaluates a sum over a product of branch and trunk embeddings: the trunk
    encodes the coordinates, each branch encodes one auxiliary input (e.g., a
    sensor-sampled function), and their elementwise product is summed and
    offset by a bias. Subclasses must provide the branch list, trunk, and
    bias as fields.
    """

    branches: eqx.AbstractVar[list[eqx.nn.MLP]]
    trunk: eqx.AbstractVar[eqx.nn.MLP]
    bias: eqx.AbstractVar[Array]

    def _eval(self, coords: Float[Array, "num_coords"], aux_inputs: tuple) -> Float[Array, ""]:
        """Sum element-wise product of trunk (coordinate) and branch (auxiliary) embeddings and add bias.

        Args:
            coords (Float[Array, 'num_coords']): Stacked coordinate array,
                shape (num_coords,), fed to the trunk.
            aux_inputs (tuple): Tuple containing arrays (sensor-sampled function) and/or scalars (PDE parameters).
                Each top level element of the tuple is mapped to its own branch network.

        Returns:
            (Float[Array, '']): The model output as a scalar JAX array.
        """
        h = self.trunk(coords)
        for idx, branch in enumerate(self.branches):
            h = h * branch(jnp.atleast_1d(aux_inputs[idx]))
        return jnp.sum(h) + self.bias
_eval(coords, aux_inputs) ¤

Sum element-wise product of trunk (coordinate) and branch (auxiliary) embeddings and add bias.

Parameters:

Name Type Description Default
coords jaxtyping.Float[jax.Array, num_coords]

Stacked coordinate array, shape (num_coords,), fed to the trunk.

 required
aux_inputs tuple

Tuple containing arrays (sensor-sampled function) and/or scalars (PDE parameters). Each top level element of the tuple is mapped to its own branch network.

 required

Returns:

Type Description
jaxtyping.Float[jax.Array, '']

The model output as a scalar JAX array.
Source code in src/popinn/models.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def _eval(self, coords: Float[Array, "num_coords"], aux_inputs: tuple) -> Float[Array, ""]:
    """Sum element-wise product of trunk (coordinate) and branch (auxiliary) embeddings and add bias.

    Args:
        coords (Float[Array, 'num_coords']): Stacked coordinate array,
            shape (num_coords,), fed to the trunk.
        aux_inputs (tuple): Tuple containing arrays (sensor-sampled function) and/or scalars (PDE parameters).
            Each top level element of the tuple is mapped to its own branch network.

    Returns:
        (Float[Array, '']): The model output as a scalar JAX array.
    """
    h = self.trunk(coords)
    for idx, branch in enumerate(self.branches):
        h = h * branch(jnp.atleast_1d(aux_inputs[idx]))
    return jnp.sum(h) + self.bias