jinns 0.9.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

jinns/utils/_spinn.py

@@ -3,54 +3,53 @@ Implements utility function to create Separable PINNs
 https://arxiv.org/abs/2211.08761
 """
 
+from dataclasses import InitVar
+from typing import Callable, Literal
 import jax
 import jax.numpy as jnp
 import equinox as eqx
+from jaxtyping import Key, Array, Float, PyTree
 
 
 class _SPINN(eqx.Module):
     """
     Construct a Separable PINN as proposed in
     Cho et al., _Separable Physics-Informed Neural Networks_, NeurIPS, 2023
+
+    Parameters
+    ----------
+    key : InitVar[Key]
+        A jax random key for the layer initializations.
+    d : int
+        The number of dimensions to treat separately.
+    eqx_list : InitVar[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, 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)
+        )`.
     """
 
-    layers: list
-    separated_mlp: list
-    d: int
-    r: int
-    m: int
+    d: int = eqx.field(static=True, kw_only=True)
 
-    def __init__(self, key, d, r, eqx_list, m=1):
-        """
-        Parameters
-        ----------
-        key
-            A jax random key
-        d
-            An integer. The number of dimensions to treat separately
-        r
-            An integer. The dimension of the embedding
-        eqx_list
-            A list of list of successive equinox modules and activation functions to
-            describe *each separable PINN architecture*.
-            The inner lists have the eqx module or
-            axtivation function as first item, other items represents arguments
-            that could be required (eg. the size of the layer).
-            __Note:__ 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]
-            ]`
-        """
-        self.d = d
-        self.r = r
-        self.m = m
+    key: InitVar[Key] = eqx.field(kw_only=True)
+    eqx_list: InitVar[tuple[tuple[Callable, int, int] | Callable, ...]] = eqx.field(
+        kw_only=True
+    )
+
+    layers: list = eqx.field(init=False)
+    separated_mlp: list = eqx.field(init=False)
 
+    def __post_init__(self, key, eqx_list):
         self.separated_mlp = []
         for _ in range(self.d):
             self.layers = []
@@ -62,7 +61,9 @@ class _SPINN(eqx.Module):
                     self.layers.append(l[0](*l[1:], key=subkey))
             self.separated_mlp.append(self.layers)
 
-    def __call__(self, t, x):
+    def __call__(
+        self, t: Float[Array, "1"], x: Float[Array, "omega_dim"]
+    ) -> Float[Array, "d embed_dim*output_dim"]:
         if t is not None:
             dimensions = jnp.concatenate([t, x.flatten()], axis=0)
         else:
@@ -78,53 +79,67 @@ class _SPINN(eqx.Module):
 
 class SPINN(eqx.Module):
     """
-    Basically a wrapper around the `__call__` function to be able to give a type to
-    our former `self.u`
-    The function create_SPINN has the role to population the `__call__` function
+    A SPINN object compatible with the rest of jinns.
+    This is typically created with `create_SPINN`.
 
     **NOTE**: SPINNs with `t` and `x` as inputs are best used with a
     DataGenerator with `self.cartesian_product=False` for memory consideration
+
+    Parameters
+    ----------
+    d : int
+        The number of dimensions to treat separately.
+
     """
 
-    d: int
-    r: int
-    eq_type: str = eqx.field(static=True)
-    m: int
-    params: eqx.Module
-    static: eqx.Module
+    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)
+
+    spinn_mlp: InitVar[eqx.Module] = eqx.field(kw_only=True)
+
+    params: PyTree = eqx.field(init=False)
+    static: PyTree = eqx.field(init=False, static=True)
 
-    def __init__(self, spinn_mlp, d, r, eq_type, m):
-        self.d, self.r, self.m = d, r, m
+    def __post_init__(self, spinn_mlp):
         self.params, self.static = eqx.partition(spinn_mlp, eqx.is_inexact_array)
-        self.eq_type = eq_type
 
-    def init_params(self):
+    def init_params(self) -> PyTree:
+        """
+        Returns an initial set of parameters
+        """
         return self.params
 
-    def __call__(self, *args):
+    def __call__(self, *args) -> Float[Array, "output_dim"]:
+        """
+        Calls `eval_nn` with rearranged arguments
+        """
         if self.eq_type == "statio_PDE":
             (x, params) = args
             try:
-                spinn = eqx.combine(params["nn_params"], self.static)
-            except (KeyError, TypeError) as e:
+                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, (0))
             res = v_model(t=None, x=x)
-            return self._eval_nn(res)
+            return self.eval_nn(res)
         if self.eq_type == "nonstatio_PDE":
             (t, x, params) = args
             try:
-                spinn = eqx.combine(params["nn_params"], self.static)
-            except (KeyError, TypeError) as e:
+                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, ((0, 0)))
             res = v_model(t, x)
-            return self._eval_nn(res)
+            return self.eval_nn(res)
         raise RuntimeError("Wrong parameter value for eq_type")
 
-    def _eval_nn(self, res):
+    def eval_nn(
+        self, res: Float[Array, "d embed_dim*output_dim"]
+    ) -> Float[Array, "output_dim"]:
         """
-        common content of apply_fn put here in order to factorize code
+        Evaluate the SPINN on some inputs with some params.
         """
         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])])
@@ -148,59 +163,71 @@ class SPINN(eqx.Module):
         return res
 
 
-def create_SPINN(key, d, r, eqx_list, eq_type, m=1):
+def create_SPINN(
+    key: Key,
+    d: int,
+    r: int,
+    eqx_list: tuple[tuple[Callable, int, int] | Callable, ...],
+    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+    m: int = 1,
+) -> SPINN:
     """
     Utility function to create a SPINN neural network with the equinox
     library.
 
-    *Note* that a SPINN is not vmapped from the outside and expects batch of the
-    same size for each input. It outputs directly a solution of shape
-    (batchsize, batchsize). See the paper for more details.
+    *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, batchsize)`. See the paper for more details.
 
     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
     d
-        An integer. The number of dimensions to treat separately
+        The number of dimensions to treat separately.
     r
-        An integer. The dimension of the embedding
+        An integer. The dimension of the embedding.
     eqx_list
-        A list of list of successive equinox modules and activation functions to
-        describe *each separable 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 represents arguments
         that could be required (eg. the size of the layer).
-        __Note:__ 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]
-        ]`
-    eq_type
+        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)
+        )`.
+    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.
     m
-        An integer. The output dimension of the neural network. According to
+        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.
 
-    **NOTE**: SPINNs with `t` and `x` as inputs are best used with a
-    DataGenerator with `self.cartesian_product=False` for memory consideration
+    !!! note
+        SPINNs with `t` and `x` as inputs are best used with a
+        DataGenerator with `self.cartesian_product=False` for memory
+        consideration
 
 
     Returns
     -------
-    `u`, a :class:`.SPINN` object which inherits from `eqx.Module` (hence callable).
+    spinn
+        An instanciated SPINN
 
     Raises
     ------
@@ -235,7 +262,7 @@ def create_SPINN(key, d, r, eqx_list, eq_type, m=1):
             "Too many dimensions, not enough letters available in jnp.einsum"
         )
 
-    spinn_mlp = _SPINN(key, d, r, eqx_list, m)
-    spinn = SPINN(spinn_mlp, d, r, eq_type, m)
+    spinn_mlp = _SPINN(key=key, d=d, eqx_list=eqx_list)
+    spinn = SPINN(spinn_mlp=spinn_mlp, d=d, r=r, eq_type=eq_type, m=m)
 
     return spinn

jinns/utils/_types.py

@@ -0,0 +1,64 @@
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
+
+from typing import TypeAlias, TYPE_CHECKING, NewType
+from jaxtyping import Int
+
+if TYPE_CHECKING:
+    from jinns.loss._LossPDE import (
+        LossPDEStatio,
+        LossPDENonStatio,
+        SystemLossPDE,
+    )
+
+    from jinns.loss._LossODE import LossODE, SystemLossODE
+    from jinns.parameters._params import Params, ParamsDict
+    from jinns.data._DataGenerators import (
+        DataGeneratorODE,
+        CubicMeshPDEStatio,
+        CubicMeshPDENonStatio,
+        DataGeneratorObservations,
+        DataGeneratorParameter,
+        DataGeneratorObservationsMultiPINNs,
+    )
+
+    from jinns.loss import DynamicLoss
+    from jinns.data._Batchs import *
+    from jinns.utils._pinn import PINN
+    from jinns.utils._hyperpinn import HYPERPINN
+    from jinns.utils._spinn import SPINN
+    from jinns.utils._containers import *
+    from jinns.validation._validation import AbstractValidationModule
+
+    AnyLoss: TypeAlias = (
+        LossPDEStatio | LossPDENonStatio | SystemLossPDE | LossODE | SystemLossODE
+    )
+
+    AnyParams: TypeAlias = Params | ParamsDict
+
+    AnyDataGenerator: TypeAlias = (
+        DataGeneratorODE | CubicMeshPDEStatio | CubicMeshPDENonStatio
+    )
+
+    AnyPINN: TypeAlias = PINN | HYPERPINN | SPINN
+
+    AnyBatch: TypeAlias = ODEBatch | PDEStatioBatch | PDENonStatioBatch
+    rar_operands = NewType(
+        "rar_operands", tuple[AnyLoss, AnyParams, AnyDataGenerator, Int]
+    )
+
+    main_carry = NewType(
+        "main_carry",
+        tuple[
+            Int,
+            AnyLoss,
+            OptimizationContainer,
+            OptimizationExtraContainer,
+            DataGeneratorContainer,
+            AbstractValidationModule,
+            LossContainer,
+            StoredObjectContainer,
+            Float[Array, "n_iter"],
+        ],
+    )

jinns/utils/_utils.py

@@ -7,9 +7,10 @@ from operator import getitem
 import numpy as np
 import jax
 import jax.numpy as jnp
+from jaxtyping import PyTree, Array
 
 
-def _check_nan_in_pytree(pytree):
+def _check_nan_in_pytree(pytree: PyTree) -> bool:
     """
     Check if there is a NaN value anywhere is the pytree
 
@@ -25,40 +26,14 @@ def _check_nan_in_pytree(pytree):
     """
     return jnp.any(
         jnp.array(
-            list(
-                jax.tree_util.tree_leaves(
-                    jax.tree_util.tree_map(lambda x: jnp.any(jnp.isnan(x)), pytree)
-                )
+            jax.tree_util.tree_leaves(
+                jax.tree_util.tree_map(lambda x: jnp.any(jnp.isnan(x)), pytree)
             )
         )
     )
 
 
-def _tracked_parameters(params, tracked_params_key_list):
-    """
-    Returns a pytree with the same structure as params with True is the
-    parameter is tracked False otherwise
-    """
-
-    def set_nested_item(dataDict, mapList, val):
-        """
-        Set item in nested dictionary
-        https://stackoverflow.com/questions/54137991/how-to-update-values-in-nested-dictionary-if-keys-are-in-a-list
-        """
-        reduce(getitem, mapList[:-1], dataDict)[mapList[-1]] = val
-        return dataDict
-
-    tracked_params = jax.tree_util.tree_map(
-        lambda x: False, params
-    )  # init with all False
-
-    for key_list in tracked_params_key_list:
-        tracked_params = set_nested_item(tracked_params, key_list, True)
-
-    return tracked_params
-
-
-def _get_grid(in_array):
+def _get_grid(in_array: Array) -> Array:
     """
     From an array of shape (B, D), D > 1, get the grid array, i.e., an array of
     shape (B, B, ...(D times)..., B, D): along the last axis we have the array
@@ -74,31 +49,7 @@ def _get_grid(in_array):
     return in_array
 
 
-def _get_vmap_in_axes_params(eq_params_batch_dict, params):
-    """
-    Return the input vmap axes when there is batch(es) of parameters to vmap
-    over. The latter are designated by keys in eq_params_batch_dict
-    If eq_params_batch_dict (ie no additional parameter batch), we return None
-    """
-    if eq_params_batch_dict is None:
-        return (None,)
-    # We use pytree indexing of vmapped axes and vmap on axis
-    # 0 of the eq_parameters for which we have a batch
-    # this is for a fine-grained vmaping
-    # scheme over the params
-    vmap_in_axes_params = (
-        {
-            "nn_params": None,
-            "eq_params": {
-                k: (0 if k in eq_params_batch_dict.keys() else None)
-                for k in params["eq_params"].keys()
-            },
-        },
-    )
-    return vmap_in_axes_params
-
-
-def _check_user_func_return(r, shape):
+def _check_user_func_return(r: Array | int, shape: tuple) -> Array | int:
     """
     Correctly handles the result from a user defined function (eg a boundary
     condition) to get the correct broadcast
@@ -115,108 +66,3 @@ def _check_user_func_return(r, shape):
     # the reshape below avoids a missing (1,) ending dimension
     # depending on how the user has coded the inital function
     return r.reshape(shape)
-
-
-def _set_derivatives(params, loss_term, derivative_keys):
-    """
-    Given derivative_keys, the parameters wrt which we want to compute
-    gradients in the loss, we set stop_gradient operators to not take the
-    derivatives with respect to the others. Note that we only operator at
-    top level
-    """
-    try:
-        params = {
-            k: (
-                value
-                if k in derivative_keys[loss_term]
-                else jax.lax.stop_gradient(value)
-            )
-            for k, value in params.items()
-        }
-    except KeyError:  # if the loss_term key has not been specified we
-        # only take gradients wrt "nn_params", all the other entries have
-        # stopped gradient
-        params = {
-            k: value if k in ["nn_params"] else jax.lax.stop_gradient(value)
-            for k, value in params.items()
-        }
-
-    return params
-
-
-def _extract_nn_params(params_dict, nn_key):
-    """
-    Given a params_dict for system loss (ie "nn_params" and "eq_params" as main
-    keys which contain dicts for each PINN (the nn_keys)) we extract the
-    corresponding "nn_params" for `nn_key` and reform a dict with "nn_params"
-    as main key as expected by the PINN/SPINN apply_fn
-    """
-    try:
-        return {
-            "nn_params": params_dict["nn_params"][nn_key],
-            "eq_params": params_dict["eq_params"][nn_key],
-        }
-    except (KeyError, IndexError) as e:
-        return {
-            "nn_params": params_dict["nn_params"][nn_key],
-            "eq_params": params_dict["eq_params"],
-        }
-
-
-def euler_maruyama_density(t, x, s, y, params, Tmax=1):
-    eps = 1e-6
-    delta = jnp.abs(t - s) * Tmax
-    mu = params["alpha_sde"] * (params["mu_sde"] - y) * delta
-    var = params["sigma_sde"] ** 2 * delta
-    return (
-        1 / jnp.sqrt(2 * jnp.pi * var) * jnp.exp(-0.5 * ((x - y) - mu) ** 2 / var) + eps
-    )
-
-
-def log_euler_maruyama_density(t, x, s, y, params):
-    eps = 1e-6
-    delta = jnp.abs(t - s)
-    mu = params["alpha_sde"] * (params["mu_sde"] - y) * delta
-    logvar = params["logvar_sde"]
-    return (
-        -0.5
-        * (jnp.log(2 * jnp.pi * delta) + logvar + ((x - y) - mu) ** 2 / jnp.exp(logvar))
-        + eps
-    )
-
-
-def euler_maruyama(x0, alpha, mu, sigma, T, N):
-    """
-    Simulate 1D diffusion process with simple parametrization using the Euler
-    Maruyama method in the interval [0, T]
-    """
-    path = [np.array([x0])]
-
-    time_steps, step_size = np.linspace(0, T, N, retstep=True)
-    for _ in time_steps[1:]:
-        path.append(
-            path[-1]
-            + step_size * (alpha * (mu - path[-1]))
-            + sigma * np.random.normal(loc=0.0, scale=np.sqrt(step_size))
-        )
-
-    return time_steps, np.stack(path)
-
-
-def _update_eq_params_dict(params, param_batch_dict):
-    # update params["eq_params"] with a batch of eq_params
-    # we avoid side_effect by recreating the dict `params`
-    # TODO transform `params` in a NamedTuple to be able to use _replace
-    # see Issue #1
-    param_batch_dict_ = param_batch_dict | {
-        k: None for k in set(params["eq_params"].keys()) - set(param_batch_dict.keys())
-    }
-    params = {"nn_params": params["nn_params"]} | {
-        "eq_params": jax.tree_util.tree_map(
-            lambda p, q: q if q is not None else p,
-            params["eq_params"],
-            param_batch_dict_,
-        )
-    }
-
-    return params