P2INN

popinn.P2INN

Parametrized Physics-Informed Neural Network.

Consists of 3 MLPs: a coordinate encoder, parameter encoder, and manifold network. The outputs of the two encoders are stacked together and then passed into the manifold network, which returns the scalar model output.

The parameters can be scalar PDE parameters, e.g. the scalar parameters \(\beta\), \(\nu\), and \(\rho\) in this PDE:

\[ \frac{\partial u}{\partial t} + \beta \frac{\partial u}{\partial x} - \nu \frac{\partial^2 u}{\partial x^2} - \rho u (1 - u) = 0 \]

Example usage:

import jax.numpy as jnp
import popinn

beta = jnp.array(1.)
nu = jnp.array(2.)
rho = jnp.array(3.)

x = jnp.array(1.)
t = jnp.array(1.)

# initialize
model = popinn.P2INN(num_params = 3, num_coords = 2)

# evaluate
u = model(x, t, (beta, nu, rho))

See Cho et al. (2024) for more details about P\(^2\)INNs.

Source code in src/popinn/models.py

class P2INN(AbstractP2INN):
    """Parametrized Physics-Informed Neural Network.

    Consists of 3 MLPs: a coordinate encoder, parameter encoder, and
    manifold network. The outputs of the two encoders are stacked together
    and then passed into the manifold network, which returns the scalar
    model output.

    The parameters can be scalar PDE parameters, e.g. the scalar parameters
    $\\beta$, $\\nu$, and $\\rho$ in this PDE:

    $$
    \\frac{\\partial u}{\\partial t} + \\beta \\frac{\\partial u}{\\partial x} - \\nu \\frac{\\partial^2 u}{\\partial x^2} - \\rho u (1 - u) = 0
    $$

    Example usage:

    ```python
    import jax.numpy as jnp
    import popinn

    beta = jnp.array(1.)
    nu = jnp.array(2.)
    rho = jnp.array(3.)

    x = jnp.array(1.)
    t = jnp.array(1.)

    # initialize
    model = popinn.P2INN(num_params = 3, num_coords = 2)

    # evaluate
    u = model(x, t, (beta, nu, rho))
    ```

    See [Cho et al. (2024)](https://arxiv.org/abs/2408.09446) for more details
    about P$^2$INNs.
    """

    param_encoder: eqx.nn.MLP
    coord_encoder: eqx.nn.MLP
    manifold: eqx.nn.MLP

    def __init__(
        self,
        key: jr.PRNGKey,
        num_params: int = 1,
        num_coords: int = 2,
        param_hidden_dim=150,
        param_depth=4,
        coord_hidden_dim=64,
        coord_depth=3,
        manifold_inner_dim=64,
        manifold_depth=5,
        param_activation=jnp.tanh,
        coord_activation=jax.nn.silu,
        manifold_inner_activation=jnp.tanh,
        manifold_final_activation=jax.nn.softplus,
        param_kwargs: dict = {},
        coord_kwargs: dict = {},
        manifold_kwargs: dict = {},
    ):
        """Initialize the three sub-networks: coordinate encoder, parameter encoder, and manifold network,
        each as an `equinox.nn.MLP`.

        Args:
            key (jr.PRNGKey): PRNG key; split three ways for the encoders and
                manifold.
            num_params (int): Number of scalar PDE parameters; defines the input dimension of the parameter encoder.
            num_coords (int): number of coordinate inputs (coord-encoder
                input size).
            param_hidden_dim (int): width and output size of the parameter
                encoder.
            param_depth (int): number of hidden layers in the parameter
                encoder.
            coord_hidden_dim (int): width and output size of the coordinate
                encoder.
            coord_depth (int): number of hidden layers in the coordinate
                encoder.
            manifold_inner_dim (int): width of the manifold network's hidden
                layers.
            manifold_depth (int): number of hidden layers in the manifold
                network.
            param_activation (Callable): activation for the parameter encoder
                (used inner and final).
            coord_activation (Callable): activation for the coordinate
                encoder (used inner and final).
            manifold_inner_activation (Callable): activation between the
                manifold network's hidden layers.
            manifold_final_activation (Callable): activation on the manifold
                output. The default softplus keeps the solution positive;
                override for sign-changing solutions.
            param_kwargs (dict): extra kwargs forwarded to the parameter
                encoder MLP.
            coord_kwargs (dict): extra kwargs forwarded to the coordinate
                encoder MLP.
            manifold_kwargs (dict): extra kwargs forwarded to the manifold
                MLP.
        """
        k1, k2, k3 = jr.split(key, 3)

        self.param_encoder = eqx.nn.MLP(
            in_size=num_params,
            out_size=param_hidden_dim,
            width_size=param_hidden_dim,
            depth=param_depth,
            activation=param_activation,
            final_activation=param_activation,
            key=k1,
            **param_kwargs,
        )

        self.coord_encoder = eqx.nn.MLP(
            in_size=num_coords,
            out_size=coord_hidden_dim,
            width_size=coord_hidden_dim,
            depth=coord_depth,
            activation=coord_activation,
            final_activation=coord_activation,
            key=k2,
            **coord_kwargs,
        )

        manifold_input_dim = self.param_encoder.out_size + self.coord_encoder.out_size

        self.manifold = eqx.nn.MLP(
            in_size=manifold_input_dim,
            out_size="scalar",  # scalar output so jax.grad / D works
            width_size=manifold_inner_dim,
            depth=manifold_depth,
            activation=manifold_inner_activation,
            final_activation=manifold_final_activation,
            key=k3,
            **manifold_kwargs,
        )

__init__(key, num_params=1, num_coords=2, param_hidden_dim=150, param_depth=4, coord_hidden_dim=64, coord_depth=3, manifold_inner_dim=64, manifold_depth=5, param_activation=jnp.tanh, coord_activation=jax.nn.silu, manifold_inner_activation=jnp.tanh, manifold_final_activation=jax.nn.softplus, param_kwargs={}, coord_kwargs={}, manifold_kwargs={})

Initialize the three sub-networks: coordinate encoder, parameter encoder, and manifold network, each as an equinox.nn.MLP.

Parameters:

Name	Type	Description	Default
key	jax.random.PRNGKey	PRNG key; split three ways for the encoders and manifold.	required
num_params	int	Number of scalar PDE parameters; defines the input dimension of the parameter encoder.	1
num_coords	int	number of coordinate inputs (coord-encoder input size).	2
param_hidden_dim	int	width and output size of the parameter encoder.	150
param_depth	int	number of hidden layers in the parameter encoder.	4
coord_hidden_dim	int	width and output size of the coordinate encoder.	64
coord_depth	int	number of hidden layers in the coordinate encoder.	3
manifold_inner_dim	int	width of the manifold network's hidden layers.	64
manifold_depth	int	number of hidden layers in the manifold network.	5
param_activation	collections.abc.Callable	activation for the parameter encoder (used inner and final).	jax.numpy.tanh
coord_activation	collections.abc.Callable	activation for the coordinate encoder (used inner and final).	jax.nn.silu
manifold_inner_activation	collections.abc.Callable	activation between the manifold network's hidden layers.	jax.numpy.tanh
manifold_final_activation	collections.abc.Callable	activation on the manifold output. The default softplus keeps the solution positive; override for sign-changing solutions.	jax.nn.softplus
param_kwargs	dict	extra kwargs forwarded to the parameter encoder MLP.	{}
coord_kwargs	dict	extra kwargs forwarded to the coordinate encoder MLP.	{}
manifold_kwargs	dict	extra kwargs forwarded to the manifold MLP.	{}

Source code in src/popinn/models.py

def __init__(
    self,
    key: jr.PRNGKey,
    num_params: int = 1,
    num_coords: int = 2,
    param_hidden_dim=150,
    param_depth=4,
    coord_hidden_dim=64,
    coord_depth=3,
    manifold_inner_dim=64,
    manifold_depth=5,
    param_activation=jnp.tanh,
    coord_activation=jax.nn.silu,
    manifold_inner_activation=jnp.tanh,
    manifold_final_activation=jax.nn.softplus,
    param_kwargs: dict = {},
    coord_kwargs: dict = {},
    manifold_kwargs: dict = {},
):
    """Initialize the three sub-networks: coordinate encoder, parameter encoder, and manifold network,
    each as an `equinox.nn.MLP`.

    Args:
        key (jr.PRNGKey): PRNG key; split three ways for the encoders and
            manifold.
        num_params (int): Number of scalar PDE parameters; defines the input dimension of the parameter encoder.
        num_coords (int): number of coordinate inputs (coord-encoder
            input size).
        param_hidden_dim (int): width and output size of the parameter
            encoder.
        param_depth (int): number of hidden layers in the parameter
            encoder.
        coord_hidden_dim (int): width and output size of the coordinate
            encoder.
        coord_depth (int): number of hidden layers in the coordinate
            encoder.
        manifold_inner_dim (int): width of the manifold network's hidden
            layers.
        manifold_depth (int): number of hidden layers in the manifold
            network.
        param_activation (Callable): activation for the parameter encoder
            (used inner and final).
        coord_activation (Callable): activation for the coordinate
            encoder (used inner and final).
        manifold_inner_activation (Callable): activation between the
            manifold network's hidden layers.
        manifold_final_activation (Callable): activation on the manifold
            output. The default softplus keeps the solution positive;
            override for sign-changing solutions.
        param_kwargs (dict): extra kwargs forwarded to the parameter
            encoder MLP.
        coord_kwargs (dict): extra kwargs forwarded to the coordinate
            encoder MLP.
        manifold_kwargs (dict): extra kwargs forwarded to the manifold
            MLP.
    """
    k1, k2, k3 = jr.split(key, 3)

    self.param_encoder = eqx.nn.MLP(
        in_size=num_params,
        out_size=param_hidden_dim,
        width_size=param_hidden_dim,
        depth=param_depth,
        activation=param_activation,
        final_activation=param_activation,
        key=k1,
        **param_kwargs,
    )

    self.coord_encoder = eqx.nn.MLP(
        in_size=num_coords,
        out_size=coord_hidden_dim,
        width_size=coord_hidden_dim,
        depth=coord_depth,
        activation=coord_activation,
        final_activation=coord_activation,
        key=k2,
        **coord_kwargs,
    )

    manifold_input_dim = self.param_encoder.out_size + self.coord_encoder.out_size

    self.manifold = eqx.nn.MLP(
        in_size=manifold_input_dim,
        out_size="scalar",  # scalar output so jax.grad / D works
        width_size=manifold_inner_dim,
        depth=manifold_depth,
        activation=manifold_inner_activation,
        final_activation=manifold_final_activation,
        key=k3,
        **manifold_kwargs,
    )

__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

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

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