jinns 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

jinns/nn/_spinn.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+from typing import Union, Callable, Any, Literal, overload
+from dataclasses import InitVar
+from jaxtyping import PyTree, Float, Array
+import jax
+import jax.numpy as jnp
+import equinox as eqx
+
+from jinns.parameters._params import Params
+from jinns.nn._abstract_pinn import AbstractPINN
+from jinns.nn._utils import _PyTree_to_Params
+
+
+class SPINN(AbstractPINN):
+    """
+    A Separable PINN object compatible with the rest of jinns.
+
+    Parameters
+    ----------
+    d : int
+        The number of dimensions to treat separately, including time `t` if
+        used for non-stationnary equations.
+    r : int
+        An integer. The dimension of the embedding.
+    eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
+        A string with three possibilities.
+        "ODE": the PINN is called with one input `t`.
+        "statio_PDE": the PINN is called with one input `x`, `x`
+        can be high dimensional.
+        "nonstatio_PDE": the PINN is called with two inputs `t` and `x`, `x`
+        can be high dimensional.
+        **Note**: the input dimension as given in eqx_list has to match the sum
+        of the dimension of `t` + the dimension of `x`.
+    m : int
+        The output dimension of the neural network. According to
+        the SPINN article, a total embedding dimension of `r*m` is defined. We
+        then sum groups of `r` embedding dimensions to compute each output.
+        Default is 1.
+    filter_spec : PyTree[Union[bool, Callable[[Any], bool]]]
+        Default is `eqx.is_inexact_array`. This tells Jinns what to consider as
+        a trainable parameter. Quoting from equinox documentation:
+        a PyTree whose structure should be a prefix of the structure of pytree.
+        Each of its leaves should either be 1) True, in which case the leaf or
+        subtree is kept; 2) False, in which case the leaf or subtree is
+        replaced with replace; 3) a callable Leaf -> bool, in which case this is evaluated on the leaf or mapped over the subtree, and the leaf kept or replaced as appropriate.
+    eqx_spinn_network : eqx.Module
+        The actual neural network instanciated as an eqx.Module. It should be
+        an architecture taking `d` inputs and returning `d` times an embedding
+        of dimension `r`*`m`. See the Separable PINN paper for more details.
+
+    """
+
+    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"] = eqx.field(
+        static=True, kw_only=True
+    )
+    d: int = eqx.field(static=True, kw_only=True)
+    r: int = eqx.field(static=True, kw_only=True)
+    m: int = eqx.field(static=True, kw_only=True, default=1)
+    filter_spec: PyTree[Union[bool, Callable[[Any], bool]]] = eqx.field(
+        static=True, kw_only=True, default=None
+    )
+    eqx_spinn_network: InitVar[eqx.Module] = eqx.field(kw_only=True)
+
+    init_params: SPINN = eqx.field(init=False)
+    static: SPINN = eqx.field(init=False, static=True)
+
+    def __post_init__(self, eqx_spinn_network):
+        if self.filter_spec is None:
+            self.filter_spec = eqx.is_inexact_array
+
+        self.init_params, self.static = eqx.partition(
+            eqx_spinn_network, self.filter_spec
+        )
+
+    @overload
+    @_PyTree_to_Params
+    def __call__(
+        self,
+        inputs: Float[Array, " input_dim"],
+        params: PyTree,
+        *args,
+        **kwargs,
+    ) -> Float[Array, " output_dim"]: ...
+
+    @_PyTree_to_Params
+    def __call__(
+        self,
+        t_x: Float[Array, "  batch_size 1+dim"],
+        params: Params[Array],
+    ) -> Float[Array, "  output_dim"]:
+        """
+        Evaluate the SPINN on some inputs with some params.
+
+        Note that that thanks to the decorator, params can also directly be the
+        PyTree (SPINN, PINN_MLP, ...) that we get out of eqx.combine
+        """
+        # try:
+        spinn = eqx.combine(params.nn_params, self.static)
+        # except (KeyError, AttributeError, TypeError) as e:
+        #    spinn = eqx.combine(params, self.static)
+        v_model = jax.vmap(spinn)
+        res = v_model(t_x)  # type: ignore
+
+        a = ", ".join([f"{chr(97 + d)}z" for d in range(res.shape[1])])
+        b = "".join([f"{chr(97 + d)}" for d in range(res.shape[1])])
+        res = jnp.stack(
+            [
+                jnp.einsum(
+                    f"{a} -> {b}",
+                    *(
+                        res[:, d, m * self.r : (m + 1) * self.r]
+                        for d in range(res.shape[1])
+                    ),
+                )
+                for m in range(self.m)
+            ],
+            axis=-1,
+        )  # compute each output dimension
+
+        # force (1,) output for non vectorial solution (consistency)
+        if len(res.shape) == self.d:
+            return jnp.expand_dims(res, axis=-1)
+        return res

jinns/nn/_spinn_mlp.py

@@ -0,0 +1,202 @@
+"""
+Implements utility function to create Separable PINNs
+https://arxiv.org/abs/2211.08761
+"""
+
+from dataclasses import InitVar
+from typing import Callable, Literal, Self, Union, Any, TypeGuard
+import jax
+import jax.numpy as jnp
+import equinox as eqx
+from jaxtyping import Key, Array, Float, PyTree
+
+from jinns.nn._mlp import MLP
+from jinns.nn._spinn import SPINN
+
+
+class SMLP(eqx.Module):
+    """
+    Construct a Separable MLP
+
+    Parameters
+    ----------
+    key : InitVar[Key]
+        A jax random key for the layer initializations.
+    d : int
+        The number of dimensions to treat separately, including time `t` if
+        used for non-stationnary equations.
+    eqx_list : InitVar[tuple[tuple[Callable, int, int] | tuple[Callable], ...]]
+        A tuple of tuples of successive equinox modules and activation functions to
+        describe the PINN architecture. The inner tuples must have the eqx module or
+        activation function as first item, other items represents arguments
+        that could be required (eg. the size of the layer).
+        The `key` argument need not be given.
+        Thus typical example is `eqx_list=
+        ((eqx.nn.Linear, 1, 20),
+            (jax.nn.tanh,),
+            (eqx.nn.Linear, 20, 20),
+            (jax.nn.tanh,),
+            (eqx.nn.Linear, 20, 20),
+            (jax.nn.tanh,),
+            (eqx.nn.Linear, 20, r * m)
+        )`.
+    """
+
+    key: InitVar[Key] = eqx.field(kw_only=True)
+    eqx_list: InitVar[tuple[tuple[Callable, int, int] | tuple[Callable], ...]] = (
+        eqx.field(kw_only=True)
+    )
+    d: int = eqx.field(static=True, kw_only=True)
+
+    separated_mlp: list[MLP] = eqx.field(init=False)
+
+    def __post_init__(self, key, eqx_list):
+        keys = jax.random.split(key, self.d)
+        self.separated_mlp = [
+            MLP(key=keys[d_], eqx_list=eqx_list) for d_ in range(self.d)
+        ]
+
+    def __call__(
+        self, inputs: Float[Array, " dim"] | Float[Array, " dim+1"]
+    ) -> Float[Array, " d embed_dim*output_dim"]:
+        outputs = []
+        for d in range(self.d):
+            x_i = inputs[d : d + 1]
+            outputs += [self.separated_mlp[d](x_i)]
+        return jnp.asarray(outputs)
+
+
+class SPINN_MLP(SPINN):
+    """
+    An implementable SPINN based on a MLP architecture
+    """
+
+    @classmethod
+    def create(
+        cls,
+        key: Key,
+        d: int,
+        r: int,
+        eqx_list: tuple[tuple[Callable, int, int] | tuple[Callable], ...],
+        eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+        m: int = 1,
+        filter_spec: PyTree[Union[bool, Callable[[Any], bool]]] = None,
+    ) -> tuple[Self, SPINN]:
+        """
+        Utility function to create a SPINN neural network with the equinox
+        library.
+
+        *Note* that a SPINN is not vmapped and expects the
+        same batch size for each of its input axis. It directly outputs a
+        solution of shape `(batchsize,) * d`. See the paper for more
+        details.
+
+        Parameters
+        ----------
+        key : Key
+            A JAX random key that will be used to initialize the network parameters
+        d : int
+            The number of dimensions to treat separately.
+        r : int
+            An integer. The dimension of the embedding.
+        eqx_list : tuple[tuple[Callable, int, int] | Callable, ...],
+            A tuple of tuples of successive equinox modules and activation functions to
+            describe the PINN architecture. The inner tuples must have the eqx module or
+            activation function as first item, other items represents arguments
+            that could be required (eg. the size of the layer).
+            The `key` argument need not be given.
+            Thus typical example is
+            `eqx_list=((eqx.nn.Linear, 1, 20),
+                (jax.nn.tanh,),
+                (eqx.nn.Linea)r, 20, 20),
+                (jax.nn.tanh,),
+                (eqx.nn.Linear, 20, 20),
+                (jax.nn.tanh,),
+                (eqx.nn.Linear, 20, r * m)
+            )`.
+        eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
+            A string with three possibilities.
+            "ODE": the PINN is called with one input `t`.
+            "statio_PDE": the PINN is called with one input `x`, `x`
+            can be high dimensional.
+            "nonstatio_PDE": the PINN is called with two inputs `t` and `x`, `x`
+            can be high dimensional.
+            **Note**: the input dimension as given in eqx_list has to match the sum
+            of the dimension of `t` + the dimension of `x`.
+        m : int
+            The output dimension of the neural network. According to
+            the SPINN article, a total embedding dimension of `r*m` is defined. We
+            then sum groups of `r` embedding dimensions to compute each output.
+            Default is 1.
+        filter_spec : PyTree[Union[bool, Callable[[Any], bool]]]
+            Default is None which leads to `eqx.is_inexact_array` in the class
+            instanciation. This tells Jinns what to consider as
+            a trainable parameter. Quoting from equinox documentation:
+            a PyTree whose structure should be a prefix of the structure of pytree.
+            Each of its leaves should either be 1) True, in which case the leaf or
+            subtree is kept; 2) False, in which case the leaf or subtree is
+            replaced with replace; 3) a callable Leaf -> bool, in which case this is evaluated on the leaf or mapped over the subtree, and the leaf kept or replaced as appropriate.
+
+
+
+
+        Returns
+        -------
+        spinn
+            An instanciated SPINN
+        spinn.init_params
+            The initial set of parameters of the model
+
+        Raises
+        ------
+        RuntimeError
+            If the parameter value for eq_type is not in `["ODE", "statio_PDE",
+            "nonstatio_PDE"]` and for various failing checks
+        """
+
+        if eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
+            raise RuntimeError("Wrong parameter value for eq_type")
+
+        def element_is_layer(element: tuple) -> TypeGuard[tuple[Callable, int, int]]:
+            return len(element) > 1
+
+        if element_is_layer(eqx_list[0]):
+            nb_inputs_declared = eqx_list[0][
+                1
+            ]  # normally we look for 2nd ele of 1st layer
+        elif element_is_layer(eqx_list[1]):
+            nb_inputs_declared = eqx_list[1][
+                1
+            ]  # but we can have, eg, a flatten first layer
+        else:
+            nb_inputs_declared = None
+        if nb_inputs_declared is None or nb_inputs_declared != 1:
+            raise ValueError("Input dim must be set to 1 in SPINN!")
+
+        if element_is_layer(eqx_list[-1]):
+            nb_outputs_declared = eqx_list[-1][2]  # normally we look for 3rd ele of
+            # last layer
+        elif element_is_layer(eqx_list[-2]):
+            nb_outputs_declared = eqx_list[-2][2]
+            # but we can have, eg, a `jnp.exp` last layer
+        else:
+            nb_outputs_declared = None
+        if nb_outputs_declared is None or nb_outputs_declared != r * m:
+            raise ValueError("Output dim must be set to r * m in SPINN!")
+
+        if d > 24:
+            raise ValueError(
+                "Too many dimensions, not enough letters available in jnp.einsum"
+            )
+
+        smlp = SMLP(key=key, d=d, eqx_list=eqx_list)
+        spinn = cls(
+            eqx_spinn_network=smlp,
+            d=d,
+            r=r,
+            eq_type=eq_type,
+            m=m,
+            filter_spec=filter_spec,
+        )
+
+        return spinn, spinn.init_params

jinns/nn/_utils.py

@@ -0,0 +1,38 @@
+from typing import Any, ParamSpec, Callable, Concatenate
+from jaxtyping import PyTree, Array
+from jinns.parameters._params import Params
+
+
+P = ParamSpec("P")
+
+
+def _PyTree_to_Params(
+    call_fun: Callable[
+        Concatenate[Any, Any, PyTree | Params[Array], P],
+        Any,
+    ],
+) -> Callable[
+    Concatenate[Any, Any, PyTree | Params[Array], P],
+    Any,
+]:
+    """
+    Decorator to be used around __call__ functions of PINNs, SPINNs, etc. It
+    authorizes the __call__ with `params` being directly be the
+    PyTree (SPINN, PINN_MLP, ...) that we get out of `eqx.combine`
+
+    This generic approach enables to cleanly handle type hints, up to the small
+    effort required to understand type hints for decorators (ie ParamSpec).
+    """
+
+    def wrapper(
+        self: Any,
+        inputs: Any,
+        params: PyTree | Params[Array],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ):
+        if isinstance(params, PyTree) and not isinstance(params, Params):
+            params = Params(nn_params=params, eq_params={})
+        return call_fun(self, inputs, params, *args, **kwargs)
+
+    return wrapper

jinns/parameters/__init__.py

@@ -1,6 +1,13 @@
-from ._params import Params, ParamsDict
+from ._params import Params
 from ._derivative_keys import (
     DerivativeKeysODE,
     DerivativeKeysPDEStatio,
     DerivativeKeysPDENonStatio,
 )
+
+__all__ = [
+    "Params",
+    "DerivativeKeysODE",
+    "DerivativeKeysPDEStatio",
+    "DerivativeKeysPDENonStatio",
+]