jinns 0.8.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

jinns/utils/__init__.py

@@ -7,3 +7,4 @@ from ._pinn import create_PINN
 from ._spinn import create_SPINN
 from ._hyperpinn import create_HYPERPINN
 from ._optim import alternate_optimizer, delayed_optimizer
+from ._save_load import save_pinn, load_pinn

jinns/utils/_hyperpinn.py

@@ -3,13 +3,13 @@ Implements utility function to create HYPERPINNs
 https://arxiv.org/pdf/2111.01008.pdf
 """
 
-from functools import partial
 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
 import equinox as eqx
 
 from jinns.utils._pinn import PINN, _MLP
@@ -35,47 +35,39 @@ class HYPERPINN(PINN):
     Composed of a PINN and an hypernetwork
     """
 
+    params_hyper: eqx.Module
+    static_hyper: eqx.Module
+    hyperparams: list = eqx.field(static=True)
+    hypernet_input_size: int
+    pinn_params_sum: ArrayLike
+    pinn_params_cumsum: ArrayLike
+
     def __init__(
         self,
-        key,
-        eqx_list,
-        eqx_list_hyper,
+        mlp,
+        hyper_mlp,
         slice_solution,
         eq_type,
         input_transform,
         output_transform,
         hyperparams,
         hypernet_input_size,
-        output_slice=None,
+        output_slice,
     ):
-        key, subkey = jax.random.split(key, 2)
         super().__init__(
-            subkey,
-            eqx_list,
+            mlp,
             slice_solution,
             eq_type,
             input_transform,
             output_transform,
             output_slice,
         )
-        self.pinn_params_sum, self.pinn_params_cumsum = _get_param_nb(self.params)
-        # the number of parameters for the pinn will be the number of ouputs
-        # for the hypetnetwork
-        self.hyperparams = hyperparams
-        self.hypernet_input_size = hypernet_input_size
-        key, subkey = jax.random.split(key, 2)
-        try:
-            eqx_list_hyper[-1][2] = self.pinn_params_sum
-        except IndexError:
-            eqx_list_hyper[-2][2] = self.pinn_params_sum
-        try:
-            eqx_list_hyper[0][1] = self.hypernet_input_size
-        except IndexError:
-            eqx_list_hyper[0][1] = self.hypernet_input_size
-        _hyper = _MLP(subkey, eqx_list_hyper)
         self.params_hyper, self.static_hyper = eqx.partition(
-            _hyper, eqx.is_inexact_array
+            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):
         return self.params_hyper
@@ -276,14 +268,32 @@ def create_HYPERPINN(
         def output_transform(_in_pinn, _out_pinn):
             return _out_pinn
 
+    key, subkey = jax.random.split(key, 2)
+    mlp = _MLP(subkey, 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
+    key, subkey = jax.random.split(key, 2)
+    hyper_mlp = _MLP(subkey, eqx_list_hyper)
+
     if shared_pinn_outputs is not None:
         hyperpinns = []
         static = None
         for output_slice in shared_pinn_outputs:
             hyperpinn = HYPERPINN(
-                key,
-                eqx_list,
-                eqx_list_hyper,
+                mlp,
+                hyper_mlp,
                 slice_solution,
                 eq_type,
                 input_transform,
@@ -300,14 +310,14 @@ def create_HYPERPINN(
             hyperpinns.append(hyperpinn)
         return hyperpinns
     hyperpinn = HYPERPINN(
-        key,
-        eqx_list,
-        eqx_list_hyper,
+        mlp,
+        hyper_mlp,
         slice_solution,
         eq_type,
         input_transform,
         output_transform,
         hyperparams,
         hypernet_input_size,
+        None,
     )
     return hyperpinn

jinns/utils/_pinn.py

@@ -2,8 +2,10 @@
 Implements utility function to create PINNs
 """
 
+from typing import Callable
 import jax
 import jax.numpy as jnp
+from jax.typing import ArrayLike
 import equinox as eqx
 
 
@@ -39,16 +41,11 @@ class _MLP(eqx.Module):
         """
 
         self.layers = []
-        # TODO we are limited currently in the number of layer type we can
-        # parse and we lack some safety checks
         for l in eqx_list:
             if len(l) == 1:
                 self.layers.append(l[0])
             else:
-                # By default we append a random key at the end of the
-                # arguments fed into a layer module call
                 key, subkey = jax.random.split(key, 2)
-                # the argument key is keyword only
                 self.layers.append(l[0](*l[1:], key=subkey))
 
     def __call__(self, t):
@@ -57,25 +54,31 @@ class _MLP(eqx.Module):
         return t
 
 
-class PINN:
+class PINN(eqx.Module):
     """
     Basically a wrapper around the `__call__` function to be able to give a type to
     our former `self.u`
     The function create_PINN has the role to population the `__call__` function
     """
 
+    slice_solution: ArrayLike
+    eq_type: str = eqx.field(static=True)
+    input_transform: Callable = eqx.field(static=True)
+    output_transform: Callable = eqx.field(static=True)
+    output_slice: ArrayLike
+    params: eqx.Module
+    static: eqx.Module
+
     def __init__(
         self,
-        key,
-        eqx_list,
+        mlp,
         slice_solution,
         eq_type,
         input_transform,
         output_transform,
-        output_slice=None,
+        output_slice,
     ):
-        _pinn = _MLP(key, eqx_list)
-        self.params, self.static = eqx.partition(_pinn, eqx.is_inexact_array)
+        self.params, self.static = eqx.partition(mlp, eqx.is_inexact_array)
         self.slice_solution = slice_solution
         self.eq_type = eq_type
         self.input_transform = input_transform
@@ -134,7 +137,7 @@ def create_PINN(
     shared_pinn_outputs=None,
     slice_solution=None,
 ):
-    """
+    r"""
     Utility function to create a standard PINN neural network with the equinox
     library.
 
@@ -246,13 +249,14 @@ def create_PINN(
         def output_transform(_in_pinn, _out_pinn):
             return _out_pinn
 
+    mlp = _MLP(key, eqx_list)
+
     if shared_pinn_outputs is not None:
         pinns = []
         static = None
         for output_slice in shared_pinn_outputs:
             pinn = PINN(
-                key,
-                eqx_list,
+                mlp,
                 slice_solution,
                 eq_type,
                 input_transform,
@@ -266,7 +270,5 @@ def create_PINN(
                 pinn.static = static
             pinns.append(pinn)
         return pinns
-    pinn = PINN(
-        key, eqx_list, slice_solution, eq_type, input_transform, output_transform
-    )
+    pinn = PINN(mlp, slice_solution, eq_type, input_transform, output_transform, None)
     return pinn

jinns/utils/_save_load.py

@@ -0,0 +1,175 @@
+"""
+Implements save and load functions
+"""
+
+import pickle
+import jax
+import equinox as eqx
+
+from jinns.utils._pinn import create_PINN
+from jinns.utils._spinn import create_SPINN
+from jinns.utils._hyperpinn import create_HYPERPINN
+
+
+def function_to_string(eqx_list):
+    """
+    We need this transformation for eqx_list to be pickled
+
+    From `[[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]` to
+    `[["Linear", 2, 20],
+                ["tanh"],
+                ["Linear", 20, 20],
+                ["tanh"],
+                ["Linear", 20, 20],
+                ["tanh"],
+                ["Linear", 20, 1]`
+    """
+    return jax.tree_util.tree_map(
+        lambda x: x.__name__ if hasattr(x, "__call__") else x, eqx_list
+    )
+
+
+def string_to_function(eqx_list_with_string):
+    """
+    We need this transformation for eqx_list at the loading ("unpickling")
+    operation.
+
+    From `[["Linear", 2, 20],
+                ["tanh"],
+                ["Linear", 20, 20],
+                ["tanh"],
+                ["Linear", 20, 20],
+                ["tanh"],
+                ["Linear", 20, 1]`
+    to  `[[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]` to
+    """
+
+    def _str_to_fun(l):
+        try:
+            try:
+                try:
+                    return getattr(jax.nn, l)
+                except AttributeError:
+                    return getattr(jax.numpy, l)
+            except AttributeError:
+                return getattr(eqx.nn, l)
+        except AttributeError as exc:
+            raise ValueError(
+                "Activation functions must be from jax.nn or jax.numpy,"
+                + "or layers must be eqx.nn layers"
+            ) from exc
+
+    return jax.tree_util.tree_map(
+        lambda x: _str_to_fun(x) if isinstance(x, str) else x, eqx_list_with_string
+    )
+
+
+def save_pinn(filename, u, params, kwargs_creation):
+    """
+    Save a PINN / HyperPINN / SPINN model
+    This function creates 3 files, beggining by `filename`
+
+     1. an eqx file to save the eqx.Module (the PINN, HyperPINN, ...)
+     2. a pickle file for the parameters
+     3. a pickle file for the arguments that have been used at PINN
+
+     creation and that we need to reconstruct the eqx.module later on.
+
+    Parameters
+    ----------
+    filename
+        Filename (prefix) without extension
+    u
+        The PINN
+    params
+        The dictionary of parameters of the model.
+        Typically, it is a dictionary of
+        dictionaries: `eq_params` and `nn_params`, respectively the
+        differential equation parameters and the neural network parameter
+    kwargs_creation
+        The dictionary of arguments that were used to create the PINN, e.g.
+        the layers list, O/PDE type, etc.
+    """
+    eqx.tree_serialise_leaves(filename + "-module.eqx", u)
+    with open(filename + "-parameters.pkl", "wb") as f:
+        pickle.dump(params, f)
+    kwargs_creation = kwargs_creation.copy()  # avoid side-effect that would be
+    # very probably harmless anyway
+
+    # we now need to transform the functions in eqx_list into strings otherwise
+    # it could not be pickled
+    kwargs_creation["eqx_list"] = function_to_string(kwargs_creation["eqx_list"])
+
+    # same thing if there is an hypernetwork:
+    try:
+        kwargs_creation["eqx_list_hyper"] = function_to_string(
+            kwargs_creation["eqx_list_hyper"]
+        )
+    except KeyError:
+        pass
+
+    with open(filename + "-arguments.pkl", "wb") as f:
+        pickle.dump(kwargs_creation, f)
+
+
+def load_pinn(filename, type_):
+    """
+    Load a PINN model. This function needs to access 3 files :
+    `{filename}-module.eqx`, `{filename}-parameters.pkl` and
+    `{filename}-arguments.pkl`.
+
+    These files are created by `jinns.utils.save_pinn`.
+
+    Note that this requires equinox v0.11.3 (currently latest version) for the
+    `eqx.filter_eval_shape` to work.
+
+    Parameters
+    ----------
+    filename
+        Filename (prefix) without extension.
+    type_
+        Type of model to load. Must be in ["pinn", "hyperpinn", "spinn"].
+
+    Returns
+    -------
+    u_reloaded
+        The reloaded PINN
+    params_reloaded
+        The reloaded parameters
+    """
+    with open(filename + "-arguments.pkl", "rb") as f:
+        kwargs_reloaded = pickle.load(f)
+    with open(filename + "-parameters.pkl", "rb") as f:
+        params_reloaded = pickle.load(f)
+    kwargs_reloaded["eqx_list"] = string_to_function(kwargs_reloaded["eqx_list"])
+    if type_ == "pinn":
+        # 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)
+    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)
+    else:
+        raise ValueError(f"{type_} is not valid")
+    # 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
+    )
+    return u_reloaded, params_reloaded

jinns/utils/_spinn.py

@@ -76,17 +76,23 @@ class _SPINN(eqx.Module):
         return jnp.asarray(outputs)
 
 
-class SPINN:
+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
     """
 
-    def __init__(self, key, d, r, eqx_list, eq_type, m=1):
+    d: int
+    r: int
+    eq_type: str = eqx.field(static=True)
+    m: int
+    params: eqx.Module
+    static: eqx.Module
+
+    def __init__(self, spinn_mlp, d, r, eq_type, m):
         self.d, self.r, self.m = d, r, m
-        _spinn = _SPINN(key, d, r, eqx_list, m)
-        self.params, self.static = eqx.partition(_spinn, eqx.is_inexact_array)
+        self.params, self.static = eqx.partition(spinn_mlp, eqx.is_inexact_array)
         self.eq_type = eq_type
 
     def init_params(self):
@@ -229,6 +235,7 @@ def create_SPINN(key, d, r, eqx_list, eq_type, m=1):
             "Too many dimensions, not enough letters available in jnp.einsum"
         )
 
-    spinn = SPINN(key, d, r, eqx_list, eq_type, m)
+    spinn_mlp = _SPINN(key, d, r, eqx_list, m)
+    spinn = SPINN(spinn_mlp, d, r, eq_type, m)
 
     return spinn

jinns/utils/_utils.py

@@ -7,7 +7,6 @@ from operator import getitem
 import numpy as np
 import jax
 import jax.numpy as jnp
-import optax
 
 
 def _check_nan_in_pytree(pytree):

jinns/utils/_utils_uspinn.py

@@ -0,0 +1,727 @@
+import numpy as np
+import jax
+import jax.numpy as jnp
+import optax
+import equinox as eqx
+from functools import reduce
+from operator import getitem
+
+
+def _check_nan_in_pytree(pytree):
+    """
+    Check if there is a NaN value anywhere is the pytree
+
+    Parameters
+    ----------
+    pytree
+        A pytree
+
+    Returns
+    -------
+    res
+        A boolean. True if any of the pytree content is NaN
+    """
+    return jnp.any(
+        jnp.array(
+            [
+                value
+                for value in 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
+
+
+class _MLP(eqx.Module):
+    """
+    Class to construct an equinox module from a key and a eqx_list. To be used
+    in pair with the function `create_PINN`
+    """
+
+    layers: list
+
+    def __init__(self, key, eqx_list):
+        """
+        Parameters
+        ----------
+        key
+            A jax random key
+        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
+            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, 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]
+            ]`
+        """
+
+        self.layers = []
+        # TODO we are limited currently in the number of layer type we can
+        # parse and we lack some safety checks
+        for l in eqx_list:
+            if len(l) == 1:
+                self.layers.append(l[0])
+            else:
+                # By default we append a random key at the end of the
+                # arguments fed into a layer module call
+                key, subkey = jax.random.split(key, 2)
+                # the argument key is keyword only
+                self.layers.append(l[0](*l[1:], key=subkey))
+
+    def __call__(self, t):
+        for layer in self.layers:
+            t = layer(t)
+        return t
+
+
+class PINN:
+    """
+    Basically a wrapper around the `__call__` function to be able to give a type to
+    our former `self.u`
+    The function create_PINN has the role to population the `__call__` function
+    """
+
+    def __init__(self, key, eqx_list, output_slice=None):
+        _pinn = _MLP(key, eqx_list)
+        self.params, self.static = eqx.partition(_pinn, eqx.is_inexact_array)
+        self.output_slice = output_slice
+
+    def init_params(self):
+        return self.params
+
+    def __call__(self, *args, **kwargs):
+        return self.apply_fn(self, *args, **kwargs)
+
+
+def create_PINN(
+    key,
+    eqx_list,
+    eq_type,
+    dim_x=0,
+    with_eq_params=None,
+    input_transform=None,
+    output_transform=None,
+    shared_pinn_outputs=None,
+):
+    """
+    Utility function to create a standard PINN neural network with the equinox
+    library.
+
+    Parameters
+    ----------
+    key
+        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
+        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, 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`
+        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` + the number of
+        parameters in `eq_params` if with_eq_params is `True` (see below)**
+    dim_x
+        An integer. The dimension of `x`. Default `0`
+    with_eq_params
+        Default is None. Otherwise a list of keys from the dict `eq_params`
+        that  the network also takes as inputs.
+        the equation parameters (`eq_params`).
+        **If some keys are provided, the input dimension
+        as given in eqx_list must take into account the number of such provided
+        keys (i.e., the input dimension is the addition of the dimension of ``t``
+        + the dimension of ``x`` + the number of ``eq_params``)**
+    input_transform
+        A function that will be called before entering the PINN. Its output(s)
+        must mathc the PINN inputs.
+    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
+    shared_pinn_outputs
+        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. Default is None, we only return one PINN.
+
+
+    Returns
+    -------
+    init_fn
+        A function which (re-)initializes the PINN parameters with the provided
+        jax random key
+    apply_fn
+        A function to apply the neural network on given inputs for given
+        parameters. A typical call will be of the form `u(t, nn_params)` for
+        ODE or `u(t, x, nn_params)` for nD PDEs (`x` being multidimensional)
+        or even `u(t, x, nn_params, eq_params)` if with_eq_params is `True`
+
+    Raises
+    ------
+    RuntimeError
+        If the parameter value for eq_type is not in `["ODE", "statio_PDE",
+        "nonstatio_PDE"]`
+    RuntimeError
+        If we have a `dim_x > 0` and `eq_type == "ODE"`
+        or if we have a `dim_x = 0` and `eq_type != "ODE"`
+    """
+    if eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
+        raise RuntimeError("Wrong parameter value for eq_type")
+
+    if eq_type == "ODE" and dim_x != 0:
+        raise RuntimeError("Wrong parameter combination eq_type and dim_x")
+
+    if eq_type != "ODE" and dim_x == 0:
+        raise RuntimeError("Wrong parameter combination eq_type and dim_x")
+
+    dim_t = 0 if eq_type == "statio_PDE" else 1
+    dim_in_params = len(with_eq_params) if with_eq_params is not None else 0
+    try:
+        nb_inputs_declared = eqx_list[0][1]  # normally we look for 2nd ele of 1st layer
+    except IndexError:
+        nb_inputs_declared = eqx_list[1][1]
+        # but we can have, eg, a flatten first layer
+
+    try:
+        nb_outputs_declared = eqx_list[-1][2]  # normally we look for 3rd ele of
+        # last layer
+    except IndexError:
+        nb_outputs_declared = eqx_list[-2][2]
+        # but we can have, eg, a `jnp.exp` last layer
+
+    # NOTE Currently the check below is disabled because we added
+    # input_transform
+    # if dim_t + dim_x + dim_in_params != nb_inputs_declared:
+    #    raise RuntimeError("Error in the declarations of the number of parameters")
+
+    if eq_type == "ODE":
+        if with_eq_params is None:
+
+            def apply_fn(self, t, u_params, eq_params=None):
+                model = eqx.combine(u_params, self.static)
+                t = t[
+                    None
+                ]  # Note that we added a dimension to t which is lacking for the ODE batches
+                if output_transform is None:
+                    if input_transform is not None:
+                        res = model(input_transform(t)).squeeze()
+                    else:
+                        res = model(t).squeeze()
+                else:
+                    if input_transform is not None:
+                        res = output_transform(t, model(input_transform(t)).squeeze())
+                    else:
+                        res = output_transform(t, model(t).squeeze())
+                if self.output_slice is not None:
+                    return res[self.output_slice]
+                else:
+                    return res
+
+        else:
+
+            def apply_fn(self, t, u_params, eq_params):
+                model = eqx.combine(u_params, self.static)
+                t = t[
+                    None
+                ]  # We added a dimension to t which is lacking for the ODE batches
+                eq_params_flatten = jnp.concatenate(
+                    [e.ravel() for k, e in eq_params.items() if k in with_eq_params]
+                )
+                t_eq_params = jnp.concatenate([t, eq_params_flatten], axis=-1)
+
+                if output_transform is None:
+                    if input_transform is not None:
+                        res = model(input_transform(t_eq_params)).squeeze()
+                    else:
+                        res = model(t_eq_params).squeeze()
+                else:
+                    if input_transform is not None:
+                        res = output_transform(
+                            t_eq_params,
+                            model(input_transform(t_eq_params)).squeeze(),
+                        )
+                    else:
+                        res = output_transform(
+                            t_eq_params, model(t_eq_params).squeeze()
+                        )
+
+                if self.output_slice is not None:
+                    return res[self.output_slice]
+                else:
+                    return res
+
+    elif eq_type == "statio_PDE":
+        # Here we add an argument `x` which can be high dimensional
+        if with_eq_params is None:
+
+            def apply_fn(self, x, u_params, eq_params=None):
+                model = eqx.combine(u_params, self.static)
+
+                if output_transform is None:
+                    if input_transform is not None:
+                        res = model(input_transform(x)).squeeze()
+                    else:
+                        res = model(x).squeeze()
+                else:
+                    if input_transform is not None:
+                        res = output_transform(x, model(input_transform(x)).squeeze())
+                    else:
+                        res = output_transform(x, model(x).squeeze()).squeeze()
+
+                if self.output_slice is not None:
+                    res = res[self.output_slice]
+
+                # force (1,) output for non vectorial solution (consistency)
+                if not res.shape:
+                    return jnp.expand_dims(res, axis=-1)
+                else:
+                    return res
+
+        else:
+
+            def apply_fn(self, x, u_params, eq_params):
+                model = eqx.combine(u_params, self.static)
+                eq_params_flatten = jnp.concatenate(
+                    [e.ravel() for k, e in eq_params.items() if k in with_eq_params]
+                )
+                x_eq_params = jnp.concatenate([x, eq_params_flatten], axis=-1)
+
+                if output_transform is None:
+                    if input_transform is not None:
+                        res = model(input_transform(x_eq_params)).squeeze()
+                    else:
+                        res = model(x_eq_params).squeeze()
+                else:
+                    if input_transform is not None:
+                        res = output_transform(
+                            x_eq_params,
+                            model(input_transform(x_eq_params)).squeeze(),
+                        )
+                    else:
+                        res = output_transform(
+                            x_eq_params, model(x_eq_params).squeeze()
+                        )
+
+                if self.output_slice is not None:
+                    res = res[self.output_slice]
+
+                # force (1,) output for non vectorial solution (consistency)
+                if not res.shape:
+                    return jnp.expand_dims(res, axis=-1)
+                else:
+                    return res
+
+    elif eq_type == "nonstatio_PDE":
+        # Here we add an argument `x` which can be high dimensional
+        if with_eq_params is None:
+
+            def apply_fn(self, t, x, u_params, eq_params=None):
+                model = eqx.combine(u_params, self.static)
+                t_x = jnp.concatenate([t, x], axis=-1)
+
+                if output_transform is None:
+                    if input_transform is not None:
+                        res = model(input_transform(t_x)).squeeze()
+                    else:
+                        res = model(t_x).squeeze()
+                else:
+                    if input_transform is not None:
+                        res = output_transform(
+                            t_x, model(input_transform(t_x)).squeeze()
+                        )
+                    else:
+                        res = output_transform(t_x, model(t_x).squeeze())
+
+                if self.output_slice is not None:
+                    res = res[self.output_slice]
+
+                ## force (1,) output for non vectorial solution (consistency)
+                if not res.shape:
+                    return jnp.expand_dims(res, axis=-1)
+                else:
+                    return res
+
+        else:
+
+            def apply_fn(self, t, x, u_params, eq_params):
+                model = eqx.combine(u_params, self.static)
+                t_x = jnp.concatenate([t, x], axis=-1)
+                eq_params_flatten = jnp.concatenate(
+                    [e.ravel() for k, e in eq_params.items() if k in with_eq_params]
+                )
+                t_x_eq_params = jnp.concatenate([t_x, eq_params_flatten], axis=-1)
+
+                if output_transform is None:
+                    if input_transform is not None:
+                        res = model(input_transform(t_x_eq_params)).squeeze()
+                    else:
+                        res = model(t_x_eq_params).squeeze()
+                else:
+                    if input_transform is not None:
+                        res = output_transform(
+                            t_x_eq_params,
+                            model(input_transform(t_x_eq_params)).squeeze(),
+                        )
+                    else:
+                        res = output_transform(
+                            t_x_eq_params,
+                            model(input_transform(t_x_eq_params)).squeeze(),
+                        )
+
+                if self.output_slice is not None:
+                    res = res[self.output_slice]
+
+                # force (1,) output for non vectorial solution (consistency)
+                if not res.shape:
+                    return jnp.expand_dims(res, axis=-1)
+                else:
+                    return res
+
+    else:
+        raise RuntimeError("Wrong parameter value for eq_type")
+
+    if shared_pinn_outputs is not None:
+        pinns = []
+        static = None
+        for output_slice in shared_pinn_outputs:
+            pinn = PINN(key, eqx_list, output_slice)
+            pinn.apply_fn = apply_fn
+            # all the pinns are in fact the same so we share the same static
+            if static is None:
+                static = pinn.static
+            else:
+                pinn.static = static
+            pinns.append(pinn)
+        return pinns
+    else:
+        pinn = PINN(key, eqx_list)
+        pinn.apply_fn = apply_fn
+        return pinn
+
+
+class _SPINN(eqx.Module):
+    """
+    Construct a Separable PINN as proposed in
+    Cho et al., _Separable Physics-Informed Neural Networks_, NeurIPS, 2023
+    """
+
+    layers: list
+    separated_mlp: list
+    d: int
+    r: int
+    m: int
+
+    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, d, 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]
+            ]`
+        """
+        keys = jax.random.split(key, 8)
+
+        self.d = d
+        self.r = r
+        self.m = m
+
+        self.separated_mlp = []
+        for d in range(self.d):
+            self.layers = []
+            for l in eqx_list:
+                if len(l) == 1:
+                    self.layers.append(l[0])
+                else:
+                    key, subkey = jax.random.split(key, 2)
+                    self.layers.append(l[0](*l[1:], key=subkey))
+            self.separated_mlp.append(self.layers)
+
+    def __call__(self, t, x):
+        if t is not None:
+            dimensions = jnp.concatenate([t, x.flatten()], axis=0)
+        else:
+            dimensions = jnp.concatenate([x.flatten()], axis=0)
+        outputs = []
+        for d in range(self.d):
+            t_ = dimensions[d][None]
+            for layer in self.separated_mlp[d]:
+                t_ = layer(t_)
+            outputs += [t_]
+        return jnp.asarray(outputs)
+
+
+def _get_grid(in_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
+    of values
+    """
+    if in_array.shape[-1] > 1 or in_array.ndim > 1:
+        return jnp.stack(
+            jnp.meshgrid(
+                *(in_array[..., d] for d in range(in_array.shape[-1])), indexing="ij"
+            ),
+            axis=-1,
+        )
+    else:
+        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,)
+    else:
+        # 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 = (
+            {
+                "eq_params": {
+                    k: (0 if k in eq_params_batch_dict.keys() else None)
+                    for k in params["eq_params"].keys()
+                },
+                "nn_params": None,
+            },
+        )
+        return vmap_in_axes_params
+
+
+def _check_user_func_return(r, shape):
+    """
+    Correctly handles the result from a user defined function (eg a boundary
+    condition) to get the correct broadcast
+    """
+    if isinstance(r, int) or isinstance(r, float):
+        # if we have a scalar cast it to float
+        return float(r)
+    if r.shape == () or len(r.shape) == 1:
+        # if we have a scalar (or a vector, but no batch dim) inside an array
+        return r.astype(float)
+    else:
+        # if we have an array of the shape of the batch dimension(s) check that
+        # we have the correct broadcast
+        # the reshape below avoids a missing (1,) ending dimension
+        # depending on how the user has coded the inital function
+        return r.reshape(shape)
+
+
+def alternate_optax_solver(
+    steps, parameters_set1, parameters_set2, lr_set1, lr_set2, label_fn=None
+):
+    """
+    This function creates an optax optimizer that alternates the optimization
+    between two set of parameters (ie. when some parameters are update to a
+    given learning rates, others are not updated (learning rate = 0)
+    The optimizers are scaled by adam parameters.
+
+    __Note:__ The alternating pattern relies on
+    `optax.piecewise_constant_schedule` which __multiplies__ learning rates of
+    previous steps (current included) to set the new learning rate. Hence, our
+    strategy used here is to relying on potentially cancelling power of tens to
+    create the alternating scheme.
+
+    Parameters
+    ----------
+    steps
+        An array which describes the epochis number at which we alternate the
+        optimization: the parameter_set that is being updated now stops
+        updating, the other parameter_set starts updating.
+        __Note:__ The step 0 should not be included
+    parameters_set1
+        A list of leaf level keys which must be found in the general `params` dict. The
+        parameters in this `set1` will be the parameters which are updated
+        first in the alternating scheme.
+    parameters_set2
+        A list of leaf level keys which must be found in the general `params` dict. The
+        parameters in this `set2` will be the parameters which are not updated
+        first in the alternating scheme.
+    lr_set1
+        A float. The learning rate of updates for set1.
+    lr_set2
+        A float. The learning rate of updates for set2.
+    label_fn
+        The same function as the label_fn function passed in an optax
+        `multi_transform`
+        [https://optax.readthedocs.io/en/latest/api.html#optax.multi_transform](see
+        here)
+        Default None, ie, we already internally provide the default one (as
+        proposed in the optax documentation) which may suit many use cases
+
+    Returns
+    -------
+    tx
+        The optax optimizer object
+    """
+
+    def map_nested_fn(fn):
+        """
+        Recursively apply `fn` to the key-value pairs of a nested dict
+        We follow the example from
+        https://optax.readthedocs.io/en/latest/api.html#optax.multi_transform
+        for different learning rates
+        """
+
+        def map_fn(nested_dict):
+            return {
+                k: (map_fn(v) if isinstance(v, dict) else fn(k, v))
+                for k, v in nested_dict.items()
+            }
+
+        return map_fn
+
+    label_fn = map_nested_fn(lambda k, _: k)
+
+    power_to_0 = 1e-25  # power of ten used to force a learning rate to 0
+    power_to_lr = 1 / power_to_0  # power of ten used to force a learning rate to lr
+    nn_params_scheduler = optax.piecewise_constant_schedule(
+        init_value=lr_set1,
+        boundaries_and_scales={
+            k: (
+                power_to_0
+                if even_odd % 2 == 0  # set lr to 0 eg if even_odd is even ie at
+                # first step
+                else power_to_lr
+            )
+            for even_odd, k in enumerate(steps)
+        },
+    )
+    eq_params_scheduler = optax.piecewise_constant_schedule(
+        init_value=power_to_0 * lr_set2,  # so normal learning rate is 1e-3
+        boundaries_and_scales={
+            k: (power_to_lr if even_odd % 2 == 0 else power_to_0)
+            for even_odd, k in enumerate(steps)
+        },
+    )
+
+    # the scheduler for set1 is called nn_chain because we usually start by
+    # updating the NN parameters
+    nn_chain = optax.chain(
+        optax.scale_by_adam(),
+        optax.scale_by_schedule(nn_params_scheduler),
+        optax.scale(-1.0),
+    )
+    eq_chain = optax.chain(
+        optax.scale_by_adam(),
+        optax.scale_by_schedule(eq_params_scheduler),
+        optax.scale(-1.0),
+    )
+    dict_params_set1 = {p: nn_chain for p in parameters_set1}
+    dict_params_set2 = {p: eq_chain for p in parameters_set2}
+    tx = optax.multi_transform(
+        {**dict_params_set1, **dict_params_set2},
+        label_fn,
+    )
+
+    return tx
+
+
+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 i 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)

{jinns-0.8.0.dist-info → jinns-0.8.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.0
+Version: 0.8.1
 Summary: Physics Informed Neural Network with JAX
 Author-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
 Maintainer-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-0.8.0.dist-info → jinns-0.8.1.dist-info}/RECORD

@@ -16,14 +16,16 @@ jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jinns/solver/_rar.py,sha256=K-0y1-ofOAo1n_Ea3QShSGCGKVYTwiaE_Bz9-DZMJm8,14525
 jinns/solver/_seq2seq.py,sha256=FL-42hTgmVl7O3hHh1ccFVw2bT8bW82hvlDRz971Chk,5620
 jinns/solver/_solve.py,sha256=6kWFWpJ33uOUzZKn7gIOM7yQsVZUwSuorOWPojVeMQY,13721
-jinns/utils/__init__.py,sha256=RxWVc4GUHkCRVtATwF-tJZvEF9VqU0t6I6vvFI8pvzY,268
-jinns/utils/_hyperpinn.py,sha256=oPQrzqDhZK1plukRFPgNPOeBWTqrezrmgKFTwI1I0eU,11007
+jinns/utils/__init__.py,sha256=44ms5UR6vMw3Nf6u4RCAzPFs4fom_YbBnH9mfne8m6k,313
+jinns/utils/_hyperpinn.py,sha256=nuy_V6qIXNxvLKvQAY6aZ_PLVroiiXQZ7RkHlF3GG60,11320
 jinns/utils/_optim.py,sha256=550kxH75TL30o1iKx1swJyP0KqyUPsJ7-imL1w65Qd0,4444
-jinns/utils/_pinn.py,sha256=U2cgxFG0Eaa94xe_Niifc_l6JXg0Ft5jTTs5nvS_258,9764
-jinns/utils/_spinn.py,sha256=C0z3d19cpj3sHPfUX8bv2GwtAgGfRR7m89LAI4j8bec,7932
-jinns/utils/_utils.py,sha256=8U0AcczbMbSEUS1J__06zI1B3TWfZhx7E0BDnHWPcPQ,6753
-jinns-0.8.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.8.0.dist-info/METADATA,sha256=QmD157p0PpAFV7OHm7tznSCUFVReSis-VamGhWbi8qQ,2482
-jinns-0.8.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jinns-0.8.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.8.0.dist-info/RECORD,,
+jinns/utils/_pinn.py,sha256=N8LuB9Ql472O01USghkJkEOmx67DTjc279T8Lj-Lwd4,9722
+jinns/utils/_save_load.py,sha256=qgZ23nUcB8-B5IZ2guuUWC4M7r5Lxd_Ms3staScdyJo,5668
+jinns/utils/_spinn.py,sha256=aeIC3DBY7f_N8HABjvBNv375dMyjll3zt6KjY2bEIkM,8058
+jinns/utils/_utils.py,sha256=8dgvWXX9NT7_7-zltWp0C9tG45ZFNwXxueyxPBb4hjo,6740
+jinns/utils/_utils_uspinn.py,sha256=qcKcOw3zrwWSQyGVj6fD8c9GinHt_U6JWN_k0auTtXM,26039
+jinns-0.8.1.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.8.1.dist-info/METADATA,sha256=AKJ921rioUOvCh8PRFt-_KbvuUMPt73i-i9QV2Orrrg,2482
+jinns-0.8.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jinns-0.8.1.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.8.1.dist-info/RECORD,,