jinns 1.3.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

jinns/nn/_ppinn.py

@@ -2,17 +2,20 @@
 Implements utility function to create PINNs
 """
 
-from typing import Callable, Literal, Self
+from __future__ import annotations
+
+from typing import Callable, Literal, Self, cast, overload
 from dataclasses import InitVar
 import jax
 import jax.numpy as jnp
 import equinox as eqx
 
-from jaxtyping import Array, Key, PyTree, Float
+from jaxtyping import Array, Key, Float, PyTree
 
-from jinns.parameters._params import Params, ParamsDict
+from jinns.parameters._params import Params
 from jinns.nn._pinn import PINN
 from jinns.nn._mlp import MLP
+from jinns.nn._utils import _PyTree_to_Params
 
 
 class PPINN_MLP(PINN):
@@ -39,12 +42,12 @@ class PPINN_MLP(PINN):
         **Note**: the input dimension as given in eqx_list has to match the sum
         of the dimension of `t` + the dimension of `x` or the output dimension
         after the `input_transform` function.
-    input_transform : Callable[[Float[Array, "input_dim"], Params], Float[Array, "output_dim"]]
+    input_transform : Callable[[Float[Array, " input_dim"], Params[Array]], Float[Array, " output_dim"]]
         A function that will be called before entering the PPINN. Its output(s)
         must match the PPINN inputs (except for the parameters).
         Its inputs are the PPINN inputs (`t` and/or `x` concatenated together)
         and the parameters. Default is no operation.
-    output_transform : Callable[[Float[Array, "input_dim"], Float[Array, "output_dim"], Params], Float[Array, "output_dim"]]
+    output_transform : Callable[[Float[Array, " input_dim"], Float[Array, " output_dim"], Params[Array]], Float[Array, " output_dim"]]
         A function with arguments begin the same input as the PPINN, the PPINN
         output and the parameter. This function will be called after exiting
         the PPINN.
@@ -63,25 +66,46 @@ class PPINN_MLP(PINN):
     """
 
     eqx_network_list: InitVar[list[eqx.Module]] = eqx.field(kw_only=True)
+    init_params: tuple[PINN, ...] = eqx.field(
+        init=False
+    )  # overriding parent attribute type
+    static: tuple[PINN, ...] = eqx.field(
+        init=False, static=True
+    )  # overriding parent attribute type
 
     def __post_init__(self, eqx_network, eqx_network_list):
         super().__post_init__(
             eqx_network=eqx_network_list[0],  # this is not used since it is
             # overwritten just below
         )
-        self.init_params, self.static = (), ()
-        for eqx_network_ in eqx_network_list:
+        params, static = eqx.partition(eqx_network_list[0], self.filter_spec)
+        self.init_params, self.static = (params,), (static,)
+        for eqx_network_ in eqx_network_list[1:]:
             params, static = eqx.partition(eqx_network_, self.filter_spec)
             self.init_params = self.init_params + (params,)
             self.static = self.static + (static,)
 
+    @overload
+    @_PyTree_to_Params
     def __call__(
         self,
-        inputs: Float[Array, "1"] | Float[Array, "dim"] | Float[Array, "1+dim"],
+        inputs: Float[Array, " input_dim"],
         params: PyTree,
-    ) -> Float[Array, "output_dim"]:
+        *args,
+        **kwargs,
+    ) -> Float[Array, " output_dim"]: ...
+
+    @_PyTree_to_Params
+    def __call__(
+        self,
+        inputs: Float[Array, " 1"] | Float[Array, " dim"] | Float[Array, " 1+dim"],
+        params: Params[Array],
+    ) -> Float[Array, " output_dim"]:
         """
         Evaluate the PPINN 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
         """
         if len(inputs.shape) == 0:
             # This can happen often when the user directly provides some
@@ -92,14 +116,14 @@ class PPINN_MLP(PINN):
 
         outs = []
 
-        try:
-            for params_, static in zip(params.nn_params, self.static):
-                model = eqx.combine(params_, static)
-                outs += [model(transformed_inputs)]
-        except (KeyError, AttributeError, TypeError) as e:
-            for params_, static in zip(params, self.static):
-                model = eqx.combine(params_, static)
-                outs += [model(transformed_inputs)]
+        # try:
+        for params_, static in zip(params.nn_params, self.static):
+            model = eqx.combine(params_, static)
+            outs += [model(transformed_inputs)]  # type: ignore
+        # except (KeyError, AttributeError, TypeError) as e:
+        #    for params_, static in zip(params, self.static):
+        #        model = eqx.combine(params_, static)
+        #        outs += [model(transformed_inputs)]
         # Note that below is then a global output transform
         res = self.output_transform(inputs, jnp.concatenate(outs, axis=0), params)
 
@@ -112,18 +136,31 @@ class PPINN_MLP(PINN):
     def create(
         cls,
         eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
-        eqx_network_list: list[eqx.nn.MLP] = None,
+        eqx_network_list: list[eqx.nn.MLP | MLP] | None = None,
         key: Key = None,
-        eqx_list_list: list[tuple[tuple[Callable, int, int] | Callable, ...]] = None,
-        input_transform: Callable[
-            [Float[Array, "input_dim"], Params], Float[Array, "output_dim"]
-        ] = None,
-        output_transform: Callable[
-            [Float[Array, "input_dim"], Float[Array, "output_dim"], Params],
-            Float[Array, "output_dim"],
-        ] = None,
-        slice_solution: slice = None,
-    ) -> tuple[Self, PyTree]:
+        eqx_list_list: (
+            list[tuple[tuple[Callable, int, int] | tuple[Callable], ...]] | None
+        ) = None,
+        input_transform: (
+            Callable[
+                [Float[Array, " input_dim"], Params[Array]],
+                Float[Array, " output_dim"],
+            ]
+            | None
+        ) = None,
+        output_transform: (
+            Callable[
+                [
+                    Float[Array, " input_dim"],
+                    Float[Array, " output_dim"],
+                    Params[Array],
+                ],
+                Float[Array, " output_dim"],
+            ]
+            | None
+        ) = None,
+        slice_solution: slice | None = None,
+    ) -> tuple[Self, tuple[PINN, ...]]:
         r"""
         Utility function to create a Parrallel PINN neural network for Jinns.
 
@@ -189,15 +226,14 @@ class PPINN_MLP(PINN):
             eqx_network_list = []
             for eqx_list in eqx_list_list:
                 key, subkey = jax.random.split(key, 2)
-                print(subkey)
                 eqx_network_list.append(MLP(key=subkey, eqx_list=eqx_list))
 
         ppinn = cls(
-            eqx_network=None,
-            eqx_network_list=eqx_network_list,
-            slice_solution=slice_solution,
+            eqx_network=None,  # type: ignore
+            eqx_network_list=cast(list[eqx.Module], eqx_network_list),
+            slice_solution=slice_solution,  # type: ignore
             eq_type=eq_type,
-            input_transform=input_transform,
-            output_transform=output_transform,
+            input_transform=input_transform,  # type: ignore
+            output_transform=output_transform,  # type: ignore
         )
         return ppinn, ppinn.init_params

jinns/nn/_save_load.py

@@ -12,7 +12,7 @@ from jinns.nn._spinn import SPINN
 from jinns.nn._mlp import PINN_MLP
 from jinns.nn._spinn_mlp import SPINN_MLP
 from jinns.nn._hyperpinn import HyperPINN
-from jinns.parameters._params import Params, ParamsDict
+from jinns.parameters._params import Params
 
 
 def function_to_string(
@@ -87,7 +87,7 @@ def string_to_function(
 def save_pinn(
     filename: str,
     u: PINN | HyperPINN | SPINN,
-    params: Params | ParamsDict,
+    params: Params,
     kwargs_creation,
 ):
     """
@@ -105,15 +105,12 @@ def save_pinn(
     tree_serialise_leaves`).
 
     Equation parameters are saved apart because the initial type of attribute
-    `params` in PINN / HyperPINN / SPINN is not `Params` nor `ParamsDict`
+    `params` in PINN / HyperPINN / SPINN is not `Params`
     but `PyTree` as inherited from `eqx.partition`.
     Therefore, if we want to ensure a proper serialization/deserialization:
     - we cannot save a `Params` object at this
       attribute field ; the `Params` object must be split into `Params.nn_params`
       (type `PyTree`) and `Params.eq_params` (type `dict`).
-    - in the case of a `ParamsDict` we cannot save `ParamsDict.nn_params` at
-      the attribute field `params` because it is not a `PyTree` (as expected in
-      the PINN / HyperPINN / SPINN signature) but it is still a dictionary.
 
     Parameters
     ----------
@@ -122,28 +119,16 @@ def save_pinn(
     u
         The PINN
     params
-        Params or ParamsDict to be save
+        Params to be saved
     kwargs_creation
         The dictionary of arguments that were used to create the PINN, e.g.
         the layers list, O/PDE type, etc.
     """
-    if isinstance(params, Params):
-        if isinstance(u, HyperPINN):
-            u = eqx.tree_at(lambda m: m.init_params_hyper, u, params)
-        elif isinstance(u, (PINN, SPINN)):
-            u = eqx.tree_at(lambda m: m.init_params, u, params)
-        eqx.tree_serialise_leaves(filename + "-module.eqx", u)
-
-    elif isinstance(params, ParamsDict):
-        for key, params_ in params.nn_params.items():
-            if isinstance(u, HyperPINN):
-                u = eqx.tree_at(lambda m: m.init_params_hyper, u, params_)
-            elif isinstance(u, (PINN, SPINN)):
-                u = eqx.tree_at(lambda m: m.init_params, u, params_)
-            eqx.tree_serialise_leaves(filename + f"-module_{key}.eqx", u)
-
-    else:
-        raise ValueError("The parameters to be saved must be a Params or a ParamsDict")
+    if isinstance(u, HyperPINN):
+        u = eqx.tree_at(lambda m: m.init_params_hyper, u, params)
+    elif isinstance(u, (PINN, SPINN)):
+        u = eqx.tree_at(lambda m: m.init_params, u, params)
+    eqx.tree_serialise_leaves(filename + "-module.eqx", u)
 
     with open(filename + "-eq_params.pkl", "wb") as f:
         pickle.dump(params.eq_params, f)
@@ -170,8 +155,7 @@ def save_pinn(
 def load_pinn(
     filename: str,
     type_: Literal["pinn_mlp", "hyperpinn", "spinn_mlp"],
-    key_list_for_paramsdict: list[str] = None,
-) -> tuple[eqx.Module, Params | ParamsDict]:
+) -> tuple[eqx.Module, Params]:
     """
     Load a PINN model. This function needs to access 3 files :
     `{filename}-module.eqx`, `{filename}-parameters.pkl` and
@@ -190,8 +174,6 @@ def load_pinn(
         Filename (prefix) without extension.
     type_
         Type of model to load. Must be in ["pinn_mlp", "hyperpinn", "spinn"].
-    key_list_for_paramsdict
-        Pass the name of the keys of the dictionnary `ParamsDict.nn_params`. Default is None. In this case, we expect to retrieve a ParamsDict.
 
     Returns
     -------
@@ -228,29 +210,17 @@ def load_pinn(
         )
     else:
         raise ValueError(f"{type_} is not valid")
-    if key_list_for_paramsdict is None:
-        # now the empty structure is populated with the actual saved array values
-        # stored in the eqx file
-        u_reloaded = eqx.tree_deserialise_leaves(
-            filename + "-module.eqx", u_reloaded_shallow
+    # now the empty structure is populated with the actual saved array values
+    # stored in the eqx file
+    u_reloaded = eqx.tree_deserialise_leaves(
+        filename + "-module.eqx", u_reloaded_shallow
+    )
+    if isinstance(u_reloaded, HyperPINN):
+        params = Params(
+            nn_params=u_reloaded.init_params_hyper, eq_params=eq_params_reloaded
         )
-        if isinstance(u_reloaded, HyperPINN):
-            params = Params(
-                nn_params=u_reloaded.init_params_hyper, eq_params=eq_params_reloaded
-            )
-        elif isinstance(u_reloaded, (PINN, SPINN)):
-            params = Params(
-                nn_params=u_reloaded.init_params, eq_params=eq_params_reloaded
-            )
+    elif isinstance(u_reloaded, (PINN, SPINN)):
+        params = Params(nn_params=u_reloaded.init_params, eq_params=eq_params_reloaded)
     else:
-        nn_params_dict = {}
-        for key in key_list_for_paramsdict:
-            u_reloaded = eqx.tree_deserialise_leaves(
-                filename + f"-module_{key}.eqx", u_reloaded_shallow
-            )
-            if isinstance(u_reloaded, HyperPINN):
-                nn_params_dict[key] = u_reloaded.init_params_hyper
-            elif isinstance(u_reloaded, (PINN, SPINN)):
-                nn_params_dict[key] = u_reloaded.init_params
-        params = ParamsDict(nn_params=nn_params_dict, eq_params=eq_params_reloaded)
+        raise ValueError("Wrong type for u_reloaded")
     return u_reloaded, params

jinns/nn/_spinn.py

@@ -1,14 +1,17 @@
-from typing import Union, Callable, Any
+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, ParamsDict
+from jinns.parameters._params import Params
+from jinns.nn._abstract_pinn import AbstractPINN
+from jinns.nn._utils import _PyTree_to_Params
 
 
-class SPINN(eqx.Module):
+class SPINN(AbstractPINN):
     """
     A Separable PINN object compatible with the rest of jinns.
 
@@ -47,21 +50,21 @@ class SPINN(eqx.Module):
 
     """
 
+    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)
-    eq_type: str = 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: PyTree = eqx.field(init=False)
-    static: PyTree = eqx.field(init=False, static=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
 
@@ -69,20 +72,34 @@ class SPINN(eqx.Module):
             eqx_spinn_network, self.filter_spec
         )
 
+    @overload
+    @_PyTree_to_Params
     def __call__(
         self,
-        t_x: Float[Array, "batch_size 1+dim"],
-        params: Params | ParamsDict | PyTree,
-    ) -> Float[Array, "output_dim"]:
+        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)
+        # 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)
+        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])])

jinns/nn/_spinn_mlp.py

@@ -4,13 +4,12 @@ https://arxiv.org/abs/2211.08761
 """
 
 from dataclasses import InitVar
-from typing import Callable, Literal, Self, Union, Any
+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.parameters._params import Params, ParamsDict
 from jinns.nn._mlp import MLP
 from jinns.nn._spinn import SPINN
 
@@ -26,7 +25,7 @@ class SMLP(eqx.Module):
     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] | Callable, ...]]
+    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
@@ -34,18 +33,18 @@ class SMLP(eqx.Module):
         The `key` argument need not be given.
         Thus typical example is `eqx_list=
         ((eqx.nn.Linear, 1, 20),
-            jax.nn.tanh,
+            (jax.nn.tanh,),
             (eqx.nn.Linear, 20, 20),
-            jax.nn.tanh,
+            (jax.nn.tanh,),
             (eqx.nn.Linear, 20, 20),
-            jax.nn.tanh,
+            (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] | Callable, ...]] = 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)
 
@@ -58,8 +57,8 @@ class SMLP(eqx.Module):
         ]
 
     def __call__(
-        self, inputs: Float[Array, "dim"] | Float[Array, "dim+1"]
-    ) -> Float[Array, "d embed_dim*output_dim"]:
+        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]
@@ -78,11 +77,11 @@ class SPINN_MLP(SPINN):
         key: Key,
         d: int,
         r: int,
-        eqx_list: tuple[tuple[Callable, int, int] | Callable, ...],
+        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, PyTree]:
+    ) -> tuple[Self, SPINN]:
         """
         Utility function to create a SPINN neural network with the equinox
         library.
@@ -108,11 +107,11 @@ class SPINN_MLP(SPINN):
             The `key` argument need not be given.
             Thus typical example is
             `eqx_list=((eqx.nn.Linear, 1, 20),
-                jax.nn.tanh,
+                (jax.nn.tanh,),
+                (eqx.nn.Linea)r, 20, 20),
+                (jax.nn.tanh,),
                 (eqx.nn.Linear, 20, 20),
-                jax.nn.tanh,
-                (eqx.nn.Linear, 20, 20),
-                jax.nn.tanh,
+                (jax.nn.tanh,),
                 (eqx.nn.Linear, 20, r * m)
             )`.
         eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
@@ -158,24 +157,31 @@ class SPINN_MLP(SPINN):
         if eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
             raise RuntimeError("Wrong parameter value for eq_type")
 
-        try:
+        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
-        except IndexError:
+        elif element_is_layer(eqx_list[1]):
             nb_inputs_declared = eqx_list[1][
                 1
             ]  # but we can have, eg, a flatten first layer
-        if nb_inputs_declared != 1:
+        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!")
 
-        try:
+        if element_is_layer(eqx_list[-1]):
             nb_outputs_declared = eqx_list[-1][2]  # normally we look for 3rd ele of
             # last layer
-        except IndexError:
+        elif element_is_layer(eqx_list[-2]):
             nb_outputs_declared = eqx_list[-2][2]
             # but we can have, eg, a `jnp.exp` last layer
-        if nb_outputs_declared != r * m:
+        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:

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",
+]