jinns 0.8.10__py3-none-any.whl → 1.0.0__py3-none-any.whl

jinns/utils/_containers.py

@@ -1,57 +1,51 @@
 """
-NamedTuples definition
+equinox Modules used as containers
 """
 
-from typing import Union, NamedTuple
-from jaxtyping import PyTree
-from jax.typing import ArrayLike
-import optax
-import jax.numpy as jnp
-from jinns.loss._LossODE import LossODE, SystemLossODE
-from jinns.loss._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
-from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-    DataGeneratorParameter,
-    DataGeneratorObservations,
-    DataGeneratorObservationsMultiPINNs,
-)
-
-
-class DataGeneratorContainer(NamedTuple):
-    data: Union[DataGeneratorODE, CubicMeshPDEStatio, CubicMeshPDENonStatio]
-    param_data: Union[DataGeneratorParameter, None] = None
-    obs_data: Union[
-        DataGeneratorObservations, DataGeneratorObservationsMultiPINNs, None
-    ] = None
-
-
-class ValidationContainer(NamedTuple):
-    loss: Union[
-        LossODE, SystemLossODE, LossPDEStatio, LossPDENonStatio, SystemLossPDE, None
-    ]
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
+
+from typing import TYPE_CHECKING, Dict
+from jaxtyping import PyTree, Array, Float, Bool
+from optax import OptState
+import equinox as eqx
+
+if TYPE_CHECKING:
+    from jinns.utils._types import *
+
+
+class DataGeneratorContainer(eqx.Module):
+    data: AnyDataGenerator
+    param_data: DataGeneratorParameter | None = None
+    obs_data: DataGeneratorObservations | DataGeneratorObservationsMultiPINNs | None = (
+        None
+    )
+
+
+class ValidationContainer(eqx.Module):
+    loss: AnyLoss | None
     data: DataGeneratorContainer
     hyperparams: PyTree = None
-    loss_values: Union[ArrayLike, None] = None
+    loss_values: Float[Array, "n_iter"] | None = None
 
 
-class OptimizationContainer(NamedTuple):
-    params: dict
-    last_non_nan_params: dict
-    opt_state: optax.OptState
+class OptimizationContainer(eqx.Module):
+    params: Params
+    last_non_nan_params: Params
+    opt_state: OptState
 
 
-class OptimizationExtraContainer(NamedTuple):
+class OptimizationExtraContainer(eqx.Module):
     curr_seq: int
-    seq2seq: Union[dict, None]
-    early_stopping: bool = False
+    best_val_params: Params
+    early_stopping: Bool = False
 
 
-class LossContainer(NamedTuple):
-    stored_loss_terms: dict
-    train_loss_values: ArrayLike
+class LossContainer(eqx.Module):
+    stored_loss_terms: Dict[str, Float[Array, "n_iter"]]
+    train_loss_values: Float[Array, "n_iter"]
 
 
-class StoredObjectContainer(NamedTuple):
-    stored_params: Union[list, None]
+class StoredObjectContainer(eqx.Module):
+    stored_params: list | None

jinns/utils/_hyperpinn.py

@@ -3,90 +3,132 @@ Implements utility function to create HYPERPINNs
 https://arxiv.org/pdf/2111.01008.pdf
 """
 
+import warnings
+from dataclasses import InitVar
+from typing import Callable, Literal
 import copy
 from math import prod
-import numpy as onp
 import jax
 import jax.numpy as jnp
 from jax.tree_util import tree_leaves, tree_map
-from jax.typing import ArrayLike
+from jaxtyping import Array, Float, PyTree, Int, Key
 import equinox as eqx
+import numpy as onp
 
 from jinns.utils._pinn import PINN, _MLP
+from jinns.parameters._params import Params
 
 
-def _get_param_nb(params):
-    """
-    Returns the number of parameters in a equinox module whose parameters
-    are stored in the pytree of parameters params but also the cumulative
-    sum when parsing the pytree
-    In reality, multiply the dimensions of the Arrays in this  pytree and
-    sum everything, using pytree utility functions
+def _get_param_nb(
+    params: Params,
+) -> tuple[Int[onp.ndarray, "1"], Int[onp.ndarray, "n_layers"]]:
+    """Returns the number of parameters in a Params object and also
+    the cumulative sum when parsing the object.
+
+
+    Parameters
+    ----------
+    params :
+        A Params object.
     """
     dim_prod_all_arrays = [
         prod(a.shape)
         for a in tree_leaves(params, is_leaf=lambda x: isinstance(x, jnp.ndarray))
     ]
-    return sum(dim_prod_all_arrays), onp.cumsum(dim_prod_all_arrays)
+    return onp.asarray(sum(dim_prod_all_arrays)), onp.cumsum(dim_prod_all_arrays)
 
 
 class HYPERPINN(PINN):
     """
-    Composed of a PINN and an hypernetwork
-    """
+    A HYPERPINN object compatible with the rest of jinns.
+    Composed of a PINN and an HYPER network. The HYPERPINN is typically
+    instanciated using with `create_HYPERPINN`. However, a user could directly
+    creates their HYPERPINN using this
+    class by passing an eqx.Module for argument `mlp` (resp. for argument
+    `hyper_mlp`) that plays the role of the NN (resp. hyper NN) and that is
+    already instanciated.
 
-    params_hyper: eqx.Module
-    static_hyper: eqx.Module
+    Parameters
+    ----------
     hyperparams: list = eqx.field(static=True)
+        A list of keys from Params.eq_params that will be considered as
+        hyperparameters for metamodeling.
     hypernet_input_size: int
-    pinn_params_sum: ArrayLike = eqx.field(static=True)
-    pinn_params_cumsum: ArrayLike = eqx.field(static=True)
+        An integer. The input size of the MLP used for the hypernetwork. Must
+        be equal to the flattened concatenations for the array of parameters
+        designated by the `hyperparams` argument.
+    slice_solution : slice
+        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 : str
+        A string with three possibilities.
+        "ODE": the HYPERPINN is called with one input `t`.
+        "statio_PDE": the HYPERPINN is called with one input `x`, `x`
+        can be high dimensional.
+        "nonstatio_PDE": the HYPERPINN 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], 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], 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.
+    output_slice : slice, default=None
+        A jnp.s\_[] to determine the different dimension for the HYPERPINN.
+        See `shared_pinn_outputs` argument of `create_HYPERPINN`.
+    mlp : eqx.Module
+        The actual neural network instanciated as an eqx.Module.
+    hyper_mlp : eqx.Module
+        The actual hyper neural network instanciated as an eqx.Module.
+    """
 
-    def __init__(
-        self,
-        mlp,
-        hyper_mlp,
-        slice_solution,
-        eq_type,
-        input_transform,
-        output_transform,
-        hyperparams,
-        hypernet_input_size,
-        output_slice,
-    ):
-        super().__init__(
+    hyperparams: list[str] = eqx.field(static=True, kw_only=True)
+    hypernet_input_size: int = eqx.field(kw_only=True)
+
+    hyper_mlp: InitVar[eqx.Module] = eqx.field(kw_only=True)
+    mlp: InitVar[eqx.Module] = eqx.field(kw_only=True)
+
+    params_hyper: PyTree = eqx.field(init=False)
+    static_hyper: PyTree = eqx.field(init=False, static=True)
+    pinn_params_sum: Int[onp.ndarray, "1"] = eqx.field(init=False, static=True)
+    pinn_params_cumsum: Int[onp.ndarray, "n_layers"] = eqx.field(
+        init=False, static=True
+    )
+
+    def __post_init__(self, mlp, hyper_mlp):
+        super().__post_init__(
             mlp,
-            slice_solution,
-            eq_type,
-            input_transform,
-            output_transform,
-            output_slice,
         )
         self.params_hyper, self.static_hyper = eqx.partition(
             hyper_mlp, eqx.is_inexact_array
         )
-        self.hyperparams = hyperparams
-        self.hypernet_input_size = hypernet_input_size
         self.pinn_params_sum, self.pinn_params_cumsum = _get_param_nb(self.params)
 
-    def init_params(self):
+    def init_params(self) -> Params:
+        """
+        Returns an initial set of parameters
+        """
         return self.params_hyper
 
-    def hyper_to_pinn(self, hyper_output):
+    def _hyper_to_pinn(self, hyper_output: Float[Array, "output_dim"]) -> PyTree:
         """
         From the output of the hypernetwork we set the well formed
         parameters of the pinn (`self.params`)
         """
         pinn_params_flat = eqx.tree_at(
-            lambda p: tree_leaves(p, is_leaf=lambda x: isinstance(x, jnp.ndarray)),
+            lambda p: tree_leaves(p, is_leaf=eqx.is_array),
             self.params,
-            [hyper_output[0 : self.pinn_params_cumsum[0]]]
-            + [
-                hyper_output[
-                    self.pinn_params_cumsum[i] : self.pinn_params_cumsum[i + 1]
-                ]
-                for i in range(len(self.pinn_params_cumsum) - 1)
-            ],
+            jnp.split(hyper_output, self.pinn_params_cumsum[:-1]),
         )
 
         return tree_map(
@@ -96,26 +138,31 @@ class HYPERPINN(PINN):
             is_leaf=lambda x: isinstance(x, jnp.ndarray),
         )
 
-    def _eval_nn(self, inputs, params, input_transform, output_transform):
+    def eval_nn(
+        self,
+        inputs: Float[Array, "input_dim"],
+        params: Params | PyTree,
+    ) -> Float[Array, "output_dim"]:
         """
-        inner function to factorize code. apply_fn (which takes varying forms)
-        call _eval_nn which always have the same content.
+        Evaluate the HYPERPINN on some inputs with some params.
         """
         try:
-            hyper = eqx.combine(params["nn_params"], self.static_hyper)
-        except (KeyError, TypeError) as e:  # give more flexibility
+            hyper = eqx.combine(params.nn_params, self.static_hyper)
+        except (KeyError, AttributeError, TypeError) as e:  # give more flexibility
             hyper = eqx.combine(params, self.static_hyper)
 
         eq_params_batch = jnp.concatenate(
-            [params["eq_params"][k].flatten() for k in self.hyperparams], axis=0
+            [params.eq_params[k].flatten() for k in self.hyperparams], axis=0
         )
 
         hyper_output = hyper(eq_params_batch)
 
-        pinn_params = self.hyper_to_pinn(hyper_output)
+        pinn_params = self._hyper_to_pinn(hyper_output)
 
         pinn = eqx.combine(pinn_params, self.static)
-        res = output_transform(inputs, pinn(input_transform(inputs, params)).squeeze())
+        res = self.output_transform(
+            inputs, pinn(self.input_transform(inputs, params)).squeeze(), params
+        )
 
         if self.output_slice is not None:
             res = res[self.output_slice]
@@ -127,18 +174,23 @@ class HYPERPINN(PINN):
 
 
 def create_HYPERPINN(
-    key,
-    eqx_list,
-    eq_type,
-    hyperparams,
-    hypernet_input_size,
-    dim_x=0,
-    input_transform=None,
-    output_transform=None,
-    slice_solution=None,
-    shared_pinn_outputs=None,
-    eqx_list_hyper=None,
-):
+    key: Key,
+    eqx_list: tuple[tuple[Callable, int, int] | Callable, ...],
+    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+    hyperparams: list[str],
+    hypernet_input_size: int,
+    dim_x: int = 0,
+    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,
+    shared_pinn_outputs: slice = None,
+    eqx_list_hyper: tuple[tuple[Callable, int, int] | Callable, ...] = None,
+) -> HYPERPINN | list[HYPERPINN]:
     r"""
     Utility function to create a standard PINN neural network with the equinox
     library.
@@ -146,59 +198,61 @@ def create_HYPERPINN(
     Parameters
     ----------
     key
-        A jax random key that will be used to initialize the network parameters
+        A JAX random key that will be used to initialize the network
+        parameters.
     eqx_list
-        A list of list of successive equinox modules and activation functions to
-        describe the PINN architecture. The inner lists have the eqx module or
-        axtivation function as first item, other items represents arguments
+        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 represent arguments
         that could be required (eg. the size of the layer).
-        __Note:__ the `key` argument need not be given.
+        The `key` argument need not be given.
         Thus typical example is `eqx_list=
-        [[eqx.nn.Linear, 2, 20],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 20, 20],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 20, 20],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 20, 1]
-        ]`
+        ((eqx.nn.Linear, 2, 20),
+            jax.nn.tanh,
+            (eqx.nn.Linear, 20, 20),
+            jax.nn.tanh,
+            (eqx.nn.Linear, 20, 20),
+            jax.nn.tanh,
+            (eqx.nn.Linear, 20, 1)
+        )`.
     eq_type
         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`
+        "ODE": the HYPERPINN is called with one input `t`.
+        "statio_PDE": the HYPERPINN is called with one input `x`, `x`
         can be high dimensional.
-        "nonstatio_PDE": the PINN is called with two inputs `t` and `x`, `x`
+        "nonstatio_PDE": the HYPERPINN 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
     hyperparams
-        A list of keys from params["eq_params"] that will be considered as
-        hyperparameters for metamodeling
+        A list of keys from Params.eq_params that will be considered as
+        hyperparameters for metamodeling.
     hypernet_input_size
         An integer. The input size of the MLP used for the hypernetwork. Must
         be equal to the flattened concatenations for the array of parameters
-        designated by the `hyperparams` argument
+        designated by the `hyperparams` argument.
     dim_x
-        An integer. The dimension of `x`. Default `0`
+        An integer. The dimension of `x`. Default `0`.
     input_transform
         A function that will be called before entering the PINN. Its output(s)
-        must match the PINN inputs. Its inputs are the PINN inputs (`t` and/or
-        `x` concatenated together and the parameters). Default is the No operation
+        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
-        A function with arguments the same input(s) as the PINN AND the PINN
-        output that will be called after exiting the PINN. Default is the No
-        operation
+        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.
     slice_solution
         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)
+        user provided argument takes care of it).
     shared_pinn_outputs
         Default is None, for a stantard PINN.
-        A tuple of jnp.s_[] (slices) to determine the different output for each
+        A tuple of jnp.s\_[] (slices) to determine the different output for each
         network. In this case we return a list of PINNs, one for each output in
         shared_pinn_outputs. This is useful to create PINNs that share the
         same network and same parameters; **the user must then use the same
@@ -216,7 +270,11 @@ def create_HYPERPINN(
 
     Returns
     -------
-    `u`, a :class:`.HyperPINN` object which inherits from `eqx.Module` (hence callable).
+    hyperpinn
+        A HYPERPINN instance or, when `shared_pinn_ouput` is not None,
+        a list of HYPERPINN instances with the same structure is returned,
+        only differing by there final slicing of the network output.
+
 
     Raises
     ------
@@ -259,53 +317,94 @@ def create_HYPERPINN(
 
     if output_transform is None:
 
-        def output_transform(_in_pinn, _out_pinn):
+        def output_transform(_in_pinn, _out_pinn, _params):
             return _out_pinn
 
     key, subkey = jax.random.split(key, 2)
-    mlp = _MLP(subkey, eqx_list)
+    mlp = _MLP(key=subkey, eqx_list=eqx_list)
     # quick partitioning to get the params to get the correct number of neurons
     # for the last layer of hyper network
     params_mlp, _ = eqx.partition(mlp, eqx.is_inexact_array)
     pinn_params_sum, _ = _get_param_nb(params_mlp)
     # the number of parameters for the pinn will be the number of ouputs
     # for the hyper network
-    try:
-        eqx_list_hyper[-1][2] = pinn_params_sum
-    except IndexError:
-        eqx_list_hyper[-2][2] = pinn_params_sum
-    try:
-        eqx_list_hyper[0][1] = hypernet_input_size
-    except IndexError:
-        eqx_list_hyper[1][1] = hypernet_input_size
+    if len(eqx_list_hyper[-1]) > 1:
+        eqx_list_hyper = eqx_list_hyper[:-1] + (
+            (eqx_list_hyper[-1][:2] + (pinn_params_sum,)),
+        )
+    else:
+        eqx_list_hyper = (
+            eqx_list_hyper[:-2]
+            + ((eqx_list_hyper[-2][:2] + (pinn_params_sum,)),)
+            + eqx_list_hyper[-1]
+        )
+    if len(eqx_list_hyper[0]) > 1:
+        eqx_list_hyper = (
+            (
+                (eqx_list_hyper[0][0],)
+                + (hypernet_input_size,)
+                + (eqx_list_hyper[0][2],)
+            ),
+        ) + eqx_list_hyper[1:]
+    else:
+        eqx_list_hyper = (
+            eqx_list_hyper[0]
+            + (
+                (
+                    (eqx_list_hyper[1][0],)
+                    + (hypernet_input_size,)
+                    + (eqx_list_hyper[1][2],)
+                ),
+            )
+            + eqx_list_hyper[2:]
+        )
     key, subkey = jax.random.split(key, 2)
-    hyper_mlp = _MLP(subkey, eqx_list_hyper)
+
+    with warnings.catch_warnings():
+        # TODO check why this warning is raised here and not in the PINN
+        # context ?
+        warnings.filterwarnings("ignore", message="A JAX array is being set as static!")
+        hyper_mlp = _MLP(key=subkey, eqx_list=eqx_list_hyper)
 
     if shared_pinn_outputs is not None:
         hyperpinns = []
         for output_slice in shared_pinn_outputs:
-            hyperpinn = HYPERPINN(
-                mlp,
-                hyper_mlp,
-                slice_solution,
-                eq_type,
-                input_transform,
-                output_transform,
-                hyperparams,
-                hypernet_input_size,
-                output_slice,
-            )
+            with warnings.catch_warnings():
+                # Catch the equinox warning because we put the number of
+                # parameters as static while being jnp.Array. This this time
+                # this is correct to do so, because they are used as indices
+                # and will never be modified
+                warnings.filterwarnings(
+                    "ignore", message="A JAX array is being set as static!"
+                )
+                hyperpinn = HYPERPINN(
+                    mlp=mlp,
+                    hyper_mlp=hyper_mlp,
+                    slice_solution=slice_solution,
+                    eq_type=eq_type,
+                    input_transform=input_transform,
+                    output_transform=output_transform,
+                    hyperparams=hyperparams,
+                    hypernet_input_size=hypernet_input_size,
+                    output_slice=output_slice,
+                )
             hyperpinns.append(hyperpinn)
         return hyperpinns
-    hyperpinn = HYPERPINN(
-        mlp,
-        hyper_mlp,
-        slice_solution,
-        eq_type,
-        input_transform,
-        output_transform,
-        hyperparams,
-        hypernet_input_size,
-        None,
-    )
+    with warnings.catch_warnings():
+        # Catch the equinox warning because we put the number of
+        # parameters as static while being jnp.Array. This this time
+        # this is correct to do so, because they are used as indices
+        # and will never be modified
+        warnings.filterwarnings("ignore", message="A JAX array is being set as static!")
+        hyperpinn = HYPERPINN(
+            mlp=mlp,
+            hyper_mlp=hyper_mlp,
+            slice_solution=slice_solution,
+            eq_type=eq_type,
+            input_transform=input_transform,
+            output_transform=output_transform,
+            hyperparams=hyperparams,
+            hypernet_input_size=hypernet_input_size,
+            output_slice=None,
+        )
     return hyperpinn