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

jinns/nn/_pinn.py

@@ -0,0 +1,204 @@
+"""
+Implement abstract class for PINN architectures
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Union, Any, Literal, overload
+from dataclasses import InitVar
+import equinox as eqx
+from jaxtyping import Float, Array, PyTree
+import jax.numpy as jnp
+from jinns.parameters._params import Params
+from jinns.nn._abstract_pinn import AbstractPINN
+from jinns.nn._utils import _PyTree_to_Params
+
+
+class PINN(AbstractPINN):
+    r"""
+    Base class for PINN objects. It can be seen as a wrapper on
+    an `eqx.Module` which actually implement the NN architectures, with extra
+    arguments handling the "physics-informed" aspect.
+
+    !!! Note
+        We use the `eqx.partition` and `eqx.combine` strategy of Equinox: a
+        `filter_spec` is applied on the PyTree and splits it into two PyTree with
+        the same structure: a static one (invisible to JAX transform such as JIT,
+        grad, etc.) and dynamic one. By convention, anything not static is
+        considered a parameter in Jinns.
+
+    For compatibility with jinns, we require that a `PINN` architecture:
+
+        1) has an eqx.Module (`eqx_network`) InitVar passed to __post_init__
+        representing the network architecture.
+        2) calls `eqx.partition` in __post_init__ in order to store the
+        static part of the model and the initial parameters.
+        3) has a `eq_type` argument, used for handling internal operations in
+        jinns.
+        4) has a `slice_solution` argument. It is a `jnp.s\_` object which
+    indicates which axis of the PINN output is dedicated to the actual equation
+    solution. Default None means that slice_solution = the whole PINN output.
+    For example, this argument is useful when the PINN is also used to output
+    equation parameters. Note that it must be a slice and not an integer (a
+    preprocessing of the user provided argument takes care of it).
+
+    Parameters
+    ----------
+    slice_solution : slice
+        Default is jnp.s\_[...]. A jnp.s\_ object which indicates which axis of the PINN output is
+        dedicated to the actual equation solution. Default None
+        means that slice_solution = the whole PINN output. This argument is useful
+        when the PINN is also used to output equation parameters for example
+        Note that it must be a slice and not an integer (a preprocessing of the
+        user provided argument takes care of it).
+    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` or the output dimension
+        after the `input_transform` function.
+    input_transform : Callable[[Float[Array, " input_dim"], Params[Array]], Float[Array, " output_dim"]]
+        A function that will be called before entering the PINN. Its output(s)
+        must match the PINN inputs (except for the parameters).
+        Its inputs are the PINN 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[Array]], Float[Array, " output_dim"]]
+        A function with arguments begin the same input as the PINN, the PINN
+        output and the parameter. This function will be called after exiting the PINN.
+        Default is no operation.
+    eqx_network : eqx.Module
+        The actual neural network instanciated as an eqx.Module.
+    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.
+
+
+    Raises
+    ------
+    RuntimeError
+        If the parameter value for eq_type is not in `["ODE", "statio_PDE",
+        "nonstatio_PDE"]`
+    """
+
+    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"] = eqx.field(
+        static=True, kw_only=True
+    )
+    slice_solution: slice = eqx.field(static=True, kw_only=True, default=None)
+    input_transform: Callable[
+        [Float[Array, " input_dim"], Params[Array]], Float[Array, " output_dim"]
+    ] = eqx.field(static=True, kw_only=True, default=None)
+    output_transform: Callable[
+        [Float[Array, " input_dim"], Float[Array, " output_dim"], Params[Array]],
+        Float[Array, " output_dim"],
+    ] = eqx.field(static=True, kw_only=True, default=None)
+
+    eqx_network: InitVar[eqx.Module] = eqx.field(kw_only=True)
+    filter_spec: PyTree[Union[bool, Callable[[Any], bool]]] = eqx.field(
+        static=True, kw_only=True, default=eqx.is_inexact_array
+    )
+
+    init_params: PINN = eqx.field(init=False)
+    static: PINN = eqx.field(init=False, static=True)
+
+    def __post_init__(self, eqx_network):
+        if self.eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
+            raise RuntimeError("Wrong parameter value for eq_type")
+        # saving the static part of the model and initial parameters
+
+        if self.filter_spec is None:
+            self.filter_spec = eqx.is_inexact_array
+
+        self.init_params, self.static = eqx.partition(eqx_network, self.filter_spec)
+
+        if self.input_transform is None:
+            self.input_transform = lambda _in, _params: _in
+
+        if self.output_transform is None:
+            self.output_transform = lambda _in_pinn, _out_pinn, _params: _out_pinn
+
+        if self.slice_solution is None:
+            self.slice_solution = jnp.s_[:]
+
+        if isinstance(self.slice_solution, int):
+            # rewrite it as a slice to ensure that axis does not disappear when
+            # indexing
+            self.slice_solution = jnp.s_[self.slice_solution : self.slice_solution + 1]
+
+    def eval(self, network, inputs, *args, **kwargs):
+        """How to call your Equinox module `network`. The purpose of this method
+        is to give more flexibility : user should re-implement `eval`
+        when inheriting from `PINN` if they desire more flexibility on how to
+        evaluate the network.
+
+        Defaults to using `network.__call__(inputs)` but it could be more refined *e.g.* `network.anymethod(inputs)`.
+
+        Parameters
+        ----------
+        network : eqx.Module
+            Your neural network with the parameters set, usually returned by
+            `eqx.combine(self.static, current_params)`.
+        inputs : Array
+            The inputs, evetually transformed by `self.input_transformed` if
+            specified by the user.
+
+        Returns
+        -------
+        Array
+            The output
+        """
+
+        return network(inputs)
+
+    @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,
+        inputs: Float[Array, " input_dim"],
+        params: Params[Array],
+        *args,
+        **kwargs,
+    ) -> Float[Array, " output_dim"]:
+        """
+        A proper __call__ implementation performs an eqx.combine here with
+        `params` and `self.static` to recreate the callable eqx.Module
+        architecture. The rest of the content of this function is dependent on
+        the network.
+
+        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
+            # collocation points (eg for plotting, whithout using
+            # DataGenerators)
+            inputs = inputs[None]
+
+        model = eqx.combine(params.nn_params, self.static)
+
+        # evaluate the model
+        res = self.eval(model, self.input_transform(inputs, params), *args, **kwargs)
+
+        res = self.output_transform(inputs, res.squeeze(), params)
+
+        # force (1,) output for non vectorial solution (consistency)
+        if not res.shape:
+            return jnp.expand_dims(res, axis=-1)
+        return res

jinns/nn/_ppinn.py

@@ -0,0 +1,239 @@
+"""
+Implements utility function to create PINNs
+"""
+
+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, Float, PyTree
+
+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):
+    r"""
+    A PPINN MLP (Parallel PINN with MLPs) object which mimicks the PFNN architecture from
+    DeepXDE. This is in fact a PINN MLP that encompasses several PINN MLPs internally.
+
+    Parameters
+    ----------
+    slice_solution : slice
+        A jnp.s\_ object which indicates which axis of the PPINN output is
+        dedicated to the actual equation solution. Default None
+        means that slice_solution = the whole PPINN output. This argument is useful
+        when the PINN is also used to output equation parameters for example
+        Note that it must be a slice and not an integer (a preprocessing of the
+        user provided argument takes care of it).
+    eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
+        A string with three possibilities.
+        "ODE": the PPINN is called with one input `t`.
+        "statio_PDE": the PPINN is called with one input `x`, `x`
+        can be high dimensional.
+        "nonstatio_PDE": the PPINN 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` or the output dimension
+        after the `input_transform` function.
+    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[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.
+        Default is no operation.
+    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_network_list
+            A list of eqx.nn.MLP objects with same input
+            dimensions. They represent the parallel subnetworks of the PPIN MLP.
+            Their respective outputs are concatenated.
+    """
+
+    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
+        )
+        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, " input_dim"],
+        params: PyTree,
+        *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
+            # collocation points (eg for plotting, whithout using
+            # DataGenerators)
+            inputs = inputs[None]
+        transformed_inputs = self.input_transform(inputs, params)
+
+        outs = []
+
+        # 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)
+
+        ## force (1,) output for non vectorial solution (consistency)
+        if not res.shape:
+            return jnp.expand_dims(res, axis=-1)
+        return res
+
+    @classmethod
+    def create(
+        cls,
+        eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+        eqx_network_list: list[eqx.nn.MLP | MLP] | None = None,
+        key: Key = None,
+        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.
+
+        Parameters
+        ----------
+        eq_type
+            A string with three possibilities.
+            "ODE": the PPINN MLP is called with one input `t`.
+            "statio_PDE": the PPINN MLP is called with one input `x`, `x`
+            can be high dimensional.
+            "nonstatio_PDE": the PPINN MLP 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` or the output dimension
+            after the `input_transform` function.
+        eqx_network_list
+                Default is None. A list of eqx.nn.MLP objects with same input
+                dimensions. They represent the parallel subnetworks of the PPIN MLP.
+                Their respective outputs are concatenated.
+        key
+            Default is None. Must be provided with `eqx_list_list` if
+            `eqx_network_list` is not provided. A JAX random key that will be used
+            to initialize the networks parameters.
+        eqx_list_list
+            Default is None. Must be provided if `eqx_network_list` is not
+            provided. A list of `eqx_list` (see `PINN_MLP.create()`). The input dimension must be the
+            same for each sub-`eqx_list`. Then the parallel subnetworks can be
+            different. Their respective outputs are concatenated.
+        input_transform
+            A function that will be called before entering the PPINN MLP. Its output(s)
+            must match the PPINN MLP inputs (except for the parameters).
+            Its inputs are the PPINN MLP inputs (`t` and/or `x` concatenated together)
+            and the parameters. Default is no operation.
+        output_transform
+            This function will be called after exiting
+            the PPINN MLP, i.e., on the concatenated outputs of all parallel networks
+            Default is no operation.
+        slice_solution
+            A jnp.s\_ object which indicates which axis of the PPINN MLP output is
+            dedicated to the actual equation solution. Default None
+            means that slice_solution = the whole PPINN MLP output. This argument is
+            useful when the PPINN MLP is also used to output equation parameters for
+            example Note that it must be a slice and not an integer (a
+            preprocessing of the user provided argument takes care of it).
+
+
+        Returns
+        -------
+        ppinn
+            A PPINN MLP instance
+        ppinn.init_params
+            An initial set of parameters for the PPINN MLP
+
+        """
+
+        if eqx_network_list is None:
+            if eqx_list_list is None or key is None:
+                raise ValueError(
+                    "If eqx_network_list is None, then key and eqx_list_list"
+                    " must be provided"
+                )
+
+            eqx_network_list = []
+            for eqx_list in eqx_list_list:
+                key, subkey = jax.random.split(key, 2)
+                eqx_network_list.append(MLP(key=subkey, eqx_list=eqx_list))
+
+        ppinn = cls(
+            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,  # type: ignore
+            output_transform=output_transform,  # type: ignore
+        )
+        return ppinn, ppinn.init_params

jinns/{utils → nn}/_save_load.py

@@ -7,14 +7,16 @@ import pickle
 import jax
 import equinox as eqx
 
-from jinns.utils._pinn import create_PINN, PINN
-from jinns.utils._spinn import create_SPINN, SPINN
-from jinns.utils._hyperpinn import create_HYPERPINN, HYPERPINN
-from jinns.parameters._params import Params, ParamsDict
+from jinns.nn._pinn import PINN
+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
 
 
 def function_to_string(
-    eqx_list: tuple[tuple[Callable, int, int] | Callable, ...]
+    eqx_list: tuple[tuple[Callable, int, int] | Callable, ...],
 ) -> tuple[tuple[str, int, int] | str, ...]:
     """
     We need this transformation for eqx_list to be pickled
@@ -40,7 +42,7 @@ def function_to_string(
 
 
 def string_to_function(
-    eqx_list_with_string: tuple[tuple[str, int, int] | str, ...]
+    eqx_list_with_string: tuple[tuple[str, int, int] | str, ...],
 ) -> tuple[tuple[Callable, int, int] | Callable, ...]:
     """
     We need this transformation for eqx_list at the loading ("unpickling")
@@ -84,8 +86,8 @@ def string_to_function(
 
 def save_pinn(
     filename: str,
-    u: PINN | HYPERPINN | SPINN,
-    params: Params | ParamsDict,
+    u: PINN | HyperPINN | SPINN,
+    params: Params,
     kwargs_creation,
 ):
     """
@@ -103,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
     ----------
@@ -120,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.params_hyper, u, params)
-        elif isinstance(u, (PINN, SPINN)):
-            u = eqx.tree_at(lambda m: m.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.params_hyper, u, params_)
-            elif isinstance(u, (PINN, SPINN)):
-                u = eqx.tree_at(lambda m: m.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)
@@ -167,9 +154,8 @@ def save_pinn(
 
 def load_pinn(
     filename: str,
-    type_: Literal["pinn", "hyperpinn", "spinn"],
-    key_list_for_paramsdict: list[str] = None,
-) -> tuple[eqx.Module, Params | ParamsDict]:
+    type_: Literal["pinn_mlp", "hyperpinn", "spinn_mlp"],
+) -> tuple[eqx.Module, Params]:
     """
     Load a PINN model. This function needs to access 3 files :
     `{filename}-module.eqx`, `{filename}-parameters.pkl` and
@@ -187,9 +173,7 @@ def load_pinn(
     filename
         Filename (prefix) without extension.
     type_
-        Type of model to load. Must be in ["pinn", "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.
+        Type of model to load. Must be in ["pinn_mlp", "hyperpinn", "spinn"].
 
     Returns
     -------
@@ -207,34 +191,36 @@ def load_pinn(
         eq_params_reloaded = {}
         print("No pickle file for equation parameters found!")
     kwargs_reloaded["eqx_list"] = string_to_function(kwargs_reloaded["eqx_list"])
-    if type_ == "pinn":
+    if type_ == "pinn_mlp":
         # next line creates a shallow model, the jax arrays are just shapes and
         # not populated, this just recreates the correct pytree structure
-        u_reloaded_shallow, _ = eqx.filter_eval_shape(create_PINN, **kwargs_reloaded)
-    elif type_ == "spinn":
-        u_reloaded_shallow, _ = eqx.filter_eval_shape(create_SPINN, **kwargs_reloaded)
+        u_reloaded_shallow, _ = eqx.filter_eval_shape(
+            PINN_MLP.create, **kwargs_reloaded
+        )
+    elif type_ == "spinn_mlp":
+        u_reloaded_shallow, _ = eqx.filter_eval_shape(
+            SPINN_MLP.create, **kwargs_reloaded
+        )
     elif type_ == "hyperpinn":
         kwargs_reloaded["eqx_list_hyper"] = string_to_function(
             kwargs_reloaded["eqx_list_hyper"]
         )
         u_reloaded_shallow, _ = eqx.filter_eval_shape(
-            create_HYPERPINN, **kwargs_reloaded
+            HyperPINN.create, **kwargs_reloaded
         )
     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
         )
+    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
-            )
-            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