DeepONet

popinn.DeepONet

Deep Operator Network consisting of a single trunk network and a branch network for each auxiliary input. The trunk and branch networks are all equinox.nn.MLP.

The final model output is constructed by taking the sum over an element-wise product between the trunk \(tr(\vec{x})\) and branch \(br_j(\vec{\mu_j})\) outputs, plus a learnable bias \(b\). For example, for two branches, the output is:

\[ \sum_i \Big[br_{1, i}(\vec{\mu}_1) * br_{2, i}(\vec{\mu}_2) * tr_i(\vec{x})\Big] + b \]

The trunk and branch outputs must therefore all be the same dimension, which is specified by branch_trunk_output_dim on initialization.

The auxiliary inputs are functions evaluated at a fixed set of sensor points, e.g. the forcing term \(F(t)\) in the PDE:

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

and the initial condition \(u_0(x)\) that may depend on \(\beta\), \(\nu\), and \(\rho\).

Example usage:

import jax.numpy as jnp
import popinn

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

ts = jnp.linspace(0,1,10)
F_t = # some function that depends t
F_t_vals = F_t(ts)

xs = jnp.linspace(0,1,10)
u_0 = # some function that depends on x and may depend on beta, nu, and/or rho
u_0_vals = u_0(xs)

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

# initialize
model = popinn.DeepONet(
            branch_input_dim = (xs.shape[0], ts.shape[0]),
            trunk_input_dim = 2,
            branch_depth = (5, 5)
        )

# evaluate
u = model(x, t, (u_0_vals, F_t_vals))

See Lu et al. (2020) for more details on DeepONets.

Source code in src/popinn/models.py

class DeepONet(AbstractDeepONet):
    """Deep Operator Network consisting of a single trunk network and a branch network for each auxiliary input.
    The trunk and branch networks are all `equinox.nn.MLP`.

    The final model output is constructed by taking the sum over an element-wise product between the
    trunk $tr(\\vec{x})$ and branch $br_j(\\vec{\\mu_j})$ outputs, plus a learnable bias $b$. For example,
    for two branches, the output is:

    $$
    \\sum_i \\Big[br_{1, i}(\\vec{\\mu}_1) * br_{2, i}(\\vec{\\mu}_2) * tr_i(\\vec{x})\\Big] + b
    $$

    The trunk and branch outputs must therefore all be the same dimension, which is specified
    by `branch_trunk_output_dim` on initialization.

    The auxiliary inputs are functions evaluated at a fixed set of sensor points, e.g. the
    forcing term $F(t)$ in the PDE:

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

    and the initial condition $u_0(x)$ that may depend on $\\beta$, $\\nu$, and $\\rho$.

    Example usage:

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

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

    ts = jnp.linspace(0,1,10)
    F_t = # some function that depends t
    F_t_vals = F_t(ts)

    xs = jnp.linspace(0,1,10)
    u_0 = # some function that depends on x and may depend on beta, nu, and/or rho
    u_0_vals = u_0(xs)

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

    # initialize
    model = popinn.DeepONet(
                branch_input_dim = (xs.shape[0], ts.shape[0]),
                trunk_input_dim = 2,
                branch_depth = (5, 5)
            )

    # evaluate
    u = model(x, t, (u_0_vals, F_t_vals))
    ```

    See [Lu et al. (2020)](https://arxiv.org/abs/1910.03193) for more details on DeepONets.
    """

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

    def __init__(
        self,
        key: jr.PRNGKey,
        branch_input_dim: tuple[int],
        trunk_input_dim: int,
        branch_trunk_output_dim: int = 100,
        branch_depth: tuple = (5,),
        trunk_depth: int = 3,
        branch_kwargs: dict = {"activation": jnp.tanh, "final_activation": jnp.tanh},
        trunk_kwargs: dict = {"activation": jnp.tanh, "final_activation": jnp.tanh},
        bias_min: float = -1.0,
        bias_max: float = 1.0,
    ):
        """Build the branch networks, trunk, and bias.

        Args:
            key (jr.PRNGKey): PRNG key.
            branch_input_dim (tuple[int]): Tuple containing the sensor count for each
                branch.
            trunk_input_dim (int): Number of coordinate inputs to the trunk.
            branch_trunk_output_dim (int): The shared hidden and output width
                of the trunk and branch networks.
            branch_depth (tuple): Number of hidden layers per branch; one
                entry per branch (indexed alongside branch_input_dim).
            trunk_depth (int): Number of hidden layers in the trunk.
            branch_kwargs (dict): Extra kwargs forwarded to every branch MLP.
            trunk_kwargs (dict): Extra kwargs forwarded to the trunk MLP.
        """

        key, trunk_key, bias_key = jr.split(key, 3)

        self.bias = jr.uniform(bias_key, minval=bias_min, maxval=bias_max)

        self.trunk = eqx.nn.MLP(
            in_size=trunk_input_dim,
            out_size=branch_trunk_output_dim,
            width_size=branch_trunk_output_dim,
            depth=trunk_depth,
            key=trunk_key,
            **trunk_kwargs,
        )

        self.branches = []
        branch_keys = jr.split(key, len(branch_input_dim))
        for idx in range(len(branch_input_dim)):
            self.branches.append(
                eqx.nn.MLP(
                    in_size=branch_input_dim[idx],
                    out_size=branch_trunk_output_dim,
                    width_size=branch_trunk_output_dim,
                    depth=branch_depth[idx],
                    key=branch_keys[idx],
                    **branch_kwargs,
                )
            )

__init__(key, branch_input_dim, trunk_input_dim, branch_trunk_output_dim=100, branch_depth=(5,), trunk_depth=3, branch_kwargs={'activation': jnp.tanh, 'final_activation': jnp.tanh}, trunk_kwargs={'activation': jnp.tanh, 'final_activation': jnp.tanh}, bias_min=-1.0, bias_max=1.0)

Build the branch networks, trunk, and bias.

Parameters:

Name	Type	Description	Default
key	jax.random.PRNGKey	PRNG key.	required
branch_input_dim	tuple[int]	Tuple containing the sensor count for each branch.	required
trunk_input_dim	int	Number of coordinate inputs to the trunk.	required
branch_trunk_output_dim	int	The shared hidden and output width of the trunk and branch networks.	100
branch_depth	tuple	Number of hidden layers per branch; one entry per branch (indexed alongside branch_input_dim).	(5,)
trunk_depth	int	Number of hidden layers in the trunk.	3
branch_kwargs	dict	Extra kwargs forwarded to every branch MLP.	{'activation': jax.numpy.tanh, 'final_activation': jax.numpy.tanh}
trunk_kwargs	dict	Extra kwargs forwarded to the trunk MLP.	{'activation': jax.numpy.tanh, 'final_activation': jax.numpy.tanh}

Source code in src/popinn/models.py

def __init__(
    self,
    key: jr.PRNGKey,
    branch_input_dim: tuple[int],
    trunk_input_dim: int,
    branch_trunk_output_dim: int = 100,
    branch_depth: tuple = (5,),
    trunk_depth: int = 3,
    branch_kwargs: dict = {"activation": jnp.tanh, "final_activation": jnp.tanh},
    trunk_kwargs: dict = {"activation": jnp.tanh, "final_activation": jnp.tanh},
    bias_min: float = -1.0,
    bias_max: float = 1.0,
):
    """Build the branch networks, trunk, and bias.

    Args:
        key (jr.PRNGKey): PRNG key.
        branch_input_dim (tuple[int]): Tuple containing the sensor count for each
            branch.
        trunk_input_dim (int): Number of coordinate inputs to the trunk.
        branch_trunk_output_dim (int): The shared hidden and output width
            of the trunk and branch networks.
        branch_depth (tuple): Number of hidden layers per branch; one
            entry per branch (indexed alongside branch_input_dim).
        trunk_depth (int): Number of hidden layers in the trunk.
        branch_kwargs (dict): Extra kwargs forwarded to every branch MLP.
        trunk_kwargs (dict): Extra kwargs forwarded to the trunk MLP.
    """

    key, trunk_key, bias_key = jr.split(key, 3)

    self.bias = jr.uniform(bias_key, minval=bias_min, maxval=bias_max)

    self.trunk = eqx.nn.MLP(
        in_size=trunk_input_dim,
        out_size=branch_trunk_output_dim,
        width_size=branch_trunk_output_dim,
        depth=trunk_depth,
        key=trunk_key,
        **trunk_kwargs,
    )

    self.branches = []
    branch_keys = jr.split(key, len(branch_input_dim))
    for idx in range(len(branch_input_dim)):
        self.branches.append(
            eqx.nn.MLP(
                in_size=branch_input_dim[idx],
                out_size=branch_trunk_output_dim,
                width_size=branch_trunk_output_dim,
                depth=branch_depth[idx],
                key=branch_keys[idx],
                **branch_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