PINN

popinn.PINN

Physics-Informed Neural Network.

A single MLP mapping coordinates to a scalar, with no auxiliary inputs.

See Raissi et al. (2019) for more details on PINNs.

Source code in src/popinn/models.py

class PINN(AbstractModel):
    """Physics-Informed Neural Network.

    A single MLP mapping coordinates to a scalar, with no auxiliary inputs.

    See [Raissi et al. (2019)](https://www.sciencedirect.com/science/article/pii/S0021999118307125) for more details on PINNs.
    """

    mlp: eqx.nn.MLP

    def __init__(
        self,
        key: jr.PRNGKey,
        num_coords: int = 2,
        hidden_dim: int = 64,
        depth: int = 4,
        inner_activation: Callable = jnp.tanh,
        final_activation: Callable = jax.nn.softplus,
        mlp_kwargs: dict = {},
    ):
        """Initialize the PINN as a multi-layer perceptron using `equinox.nn.MLP`.

        Args:
            key (jr.PRNGKey): PRNG key for MLP initialization.
            num_coords (int): Number of coordinate inputs (e.g. 2 for x & t).
            hidden_dim (int): Width of each hidden layer.
            depth (int): Number of hidden layers.
            inner_activation (Callable): Activation between hidden layers.
            final_activation (Callable): Activation on the output.
            mlp_kwargs (dict): Extra keyword arguments forwarded to
                `equinox.nn.MLP`.
        """
        self.mlp = eqx.nn.MLP(
            in_size=num_coords,
            out_size="scalar",
            width_size=hidden_dim,
            depth=depth,
            activation=inner_activation,
            final_activation=final_activation,
            key=key,
            **mlp_kwargs,
        )

    def _eval(
        self,
        coords: Float[Array, "num_coords"],
        aux_inputs: tuple,
    ) -> Float[Array, ""]:
        """Map coordinates to a scalar; aux_inputs is ignored.

        Args:
            coords (Float[Array, 'num_coords']): stacked coordinate array.
            aux_inputs (tuple): accepted for API consistency and not used.

        Returns:
            Float[Array, '']: scalar solution estimate.
        """
        # aux_inputs are not used, included for API consistency
        return self.mlp(coords)

__init__(key, num_coords=2, hidden_dim=64, depth=4, inner_activation=jnp.tanh, final_activation=jax.nn.softplus, mlp_kwargs={})

Initialize the PINN as a multi-layer perceptron using equinox.nn.MLP.

Parameters:

Name	Type	Description	Default
key	jax.random.PRNGKey	PRNG key for MLP initialization.	required
num_coords	int	Number of coordinate inputs (e.g. 2 for x & t).	2
hidden_dim	int	Width of each hidden layer.	64
depth	int	Number of hidden layers.	4
inner_activation	collections.abc.Callable	Activation between hidden layers.	jax.numpy.tanh
final_activation	collections.abc.Callable	Activation on the output.	jax.nn.softplus
mlp_kwargs	dict	Extra keyword arguments forwarded to equinox.nn.MLP.	{}

Source code in src/popinn/models.py

def __init__(
    self,
    key: jr.PRNGKey,
    num_coords: int = 2,
    hidden_dim: int = 64,
    depth: int = 4,
    inner_activation: Callable = jnp.tanh,
    final_activation: Callable = jax.nn.softplus,
    mlp_kwargs: dict = {},
):
    """Initialize the PINN as a multi-layer perceptron using `equinox.nn.MLP`.

    Args:
        key (jr.PRNGKey): PRNG key for MLP initialization.
        num_coords (int): Number of coordinate inputs (e.g. 2 for x & t).
        hidden_dim (int): Width of each hidden layer.
        depth (int): Number of hidden layers.
        inner_activation (Callable): Activation between hidden layers.
        final_activation (Callable): Activation on the output.
        mlp_kwargs (dict): Extra keyword arguments forwarded to
            `equinox.nn.MLP`.
    """
    self.mlp = eqx.nn.MLP(
        in_size=num_coords,
        out_size="scalar",
        width_size=hidden_dim,
        depth=depth,
        activation=inner_activation,
        final_activation=final_activation,
        key=key,
        **mlp_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