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

jinns/nn/_spinn_mlp.py

@@ -0,0 +1,196 @@
+"""
+Implements utility function to create Separable PINNs
+https://arxiv.org/abs/2211.08761
+"""
+
+from dataclasses import InitVar
+from typing import Callable, Literal, Self, Union, Any
+import jax
+import jax.numpy as jnp
+import equinox as eqx
+from jaxtyping import Key, Array, Float, PyTree
+
+from jinns.parameters._params import Params, ParamsDict
+from jinns.nn._mlp import MLP
+from jinns.nn._spinn import SPINN
+
+
+class SMLP(eqx.Module):
+    """
+    Construct a Separable MLP
+
+    Parameters
+    ----------
+    key : InitVar[Key]
+        A jax random key for the layer initializations.
+    d : int
+        The number of dimensions to treat separately, including time `t` if
+        used for non-stationnary equations.
+    eqx_list : InitVar[tuple[tuple[Callable, int, int] | Callable, ...]]
+        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, 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 * m)
+        )`.
+    """
+
+    key: InitVar[Key] = eqx.field(kw_only=True)
+    eqx_list: InitVar[tuple[tuple[Callable, int, int] | Callable, ...]] = eqx.field(
+        kw_only=True
+    )
+    d: int = eqx.field(static=True, kw_only=True)
+
+    separated_mlp: list[MLP] = eqx.field(init=False)
+
+    def __post_init__(self, key, eqx_list):
+        keys = jax.random.split(key, self.d)
+        self.separated_mlp = [
+            MLP(key=keys[d_], eqx_list=eqx_list) for d_ in range(self.d)
+        ]
+
+    def __call__(
+        self, inputs: Float[Array, "dim"] | Float[Array, "dim+1"]
+    ) -> Float[Array, "d embed_dim*output_dim"]:
+        outputs = []
+        for d in range(self.d):
+            x_i = inputs[d : d + 1]
+            outputs += [self.separated_mlp[d](x_i)]
+        return jnp.asarray(outputs)
+
+
+class SPINN_MLP(SPINN):
+    """
+    An implementable SPINN based on a MLP architecture
+    """
+
+    @classmethod
+    def create(
+        cls,
+        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,
+        filter_spec: PyTree[Union[bool, Callable[[Any], bool]]] = None,
+    ) -> tuple[Self, PyTree]:
+        """
+        Utility function to create a SPINN neural network with the equinox
+        library.
+
+        *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,) * d`. See the paper for more
+        details.
+
+        Parameters
+        ----------
+        key : Key
+            A JAX random key that will be used to initialize the network parameters
+        d : int
+            The number of dimensions to treat separately.
+        r : int
+            An integer. The dimension of the embedding.
+        eqx_list : 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, 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 * m)
+            )`.
+        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`.
+        m : int
+            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.
+        filter_spec : PyTree[Union[bool, Callable[[Any], bool]]]
+            Default is None which leads to `eqx.is_inexact_array` in the class
+            instanciation. 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.
+
+
+
+
+        Returns
+        -------
+        spinn
+            An instanciated SPINN
+        spinn.init_params
+            The initial set of parameters of the model
+
+        Raises
+        ------
+        RuntimeError
+            If the parameter value for eq_type is not in `["ODE", "statio_PDE",
+            "nonstatio_PDE"]` and for various failing checks
+        """
+
+        if eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
+            raise RuntimeError("Wrong parameter value for eq_type")
+
+        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
+        if nb_inputs_declared != 1:
+            raise ValueError("Input dim must be set to 1 in SPINN!")
+
+        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
+        if nb_outputs_declared != r * m:
+            raise ValueError("Output dim must be set to r * m in SPINN!")
+
+        if d > 24:
+            raise ValueError(
+                "Too many dimensions, not enough letters available in jnp.einsum"
+            )
+
+        smlp = SMLP(key=key, d=d, eqx_list=eqx_list)
+        spinn = cls(
+            eqx_spinn_network=smlp,
+            d=d,
+            r=r,
+            eq_type=eq_type,
+            m=m,
+            filter_spec=filter_spec,
+        )
+
+        return spinn, spinn.init_params

jinns/plot/_plot.py

@@ -208,7 +208,7 @@ def _plot_2D_statio(
     figsize :
         By default (7, 7)
     spinn :
-        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+        True if a SPINN is to be plotted. False for PINNs and HyperPINNs
     vmin_vmax: list, optional
         The colorbar minimum and maximum value. Defaults None.
 
@@ -272,7 +272,7 @@ def plot1d_slice(
     figsize
         size of the figure, by default (10, 10)
     spinn
-        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+        True if a SPINN is to be plotted. False for PINNs and HyperPINNs
     ax
         A pre-defined `matplotlib.Axes` where you want to plot.
 
@@ -341,7 +341,7 @@ def plot1d_image(
     cmap :
         the matplotlib color map used in the ImageGrid.
     spinn :
-        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+        True if a SPINN is to be plotted. False for PINNs and HyperPINNs
     vmin_vmax:
         The colorbar minimum and maximum value. Defaults None.
 

jinns/solver/_rar.py

@@ -19,8 +19,8 @@ from jinns.data._DataGenerators import (
     CubicMeshPDEStatio,
     CubicMeshPDENonStatio,
 )
-from jinns.utils._hyperpinn import HYPERPINN
-from jinns.utils._spinn import SPINN
+from jinns.nn._hyperpinn import HyperPINN
+from jinns.nn._spinn import SPINN
 
 
 if TYPE_CHECKING:
@@ -115,7 +115,7 @@ def _rar_step_init(sample_size: Int, selected_sample_size: Int) -> tuple[
 
     def rar_step_true(operands: rar_operands) -> AnyDataGenerator:
         loss, params, data, i = operands
-        if isinstance(loss.u, HYPERPINN) or isinstance(loss.u, SPINN):
+        if isinstance(loss.u, HyperPINN) or isinstance(loss.u, SPINN):
             raise NotImplementedError("RAR not implemented for hyperPINN and SPINN")
 
         if isinstance(data, DataGeneratorODE):

jinns/solver/_solve.py

@@ -47,6 +47,7 @@ def solve(
     validation: AbstractValidationModule | None = None,
     obs_batch_sharding: jax.sharding.Sharding | None = None,
     verbose: Bool = True,
+    ahead_of_time: Bool = True,
 ) -> tuple[
     Params | ParamsDict,
     Float[Array, "n_iter"],
@@ -118,6 +119,14 @@ def solve(
     verbose
         Default True. If False, no std output (loss or cause of
         exiting the optimization loop) will be produced.
+    ahead_of_time
+        Default True. Separate the compilation of the main training loop from
+        the execution to get both timings. You might need to avoid this
+        behaviour if you need to perform JAX transforms over chunks of code
+        containing `jinns.solve()` since AOT-compiled functions cannot be JAX
+        transformed (see https://jax.readthedocs.io/en/latest/aot.html#aot-compiled-functions-cannot-be-transformed).
+        When False, jinns does not provide any timing information (which would
+        be nonsense in a JIT transformed `solve()` function).
 
     Returns
     -------
@@ -384,16 +393,21 @@ def solve(
         def train_fun(carry):
             return jax.lax.while_loop(break_fun, _one_iteration, carry)
 
-        start = time.time()
-        compiled_train_fun = jax.jit(train_fun).lower(carry).compile()
-        end = time.time()
-        print("\nCompilation took\n", end - start, "\n")
+        if ahead_of_time:
+            start = time.time()
+            compiled_train_fun = jax.jit(train_fun).lower(carry).compile()
+            end = time.time()
+            if verbose:
+                print("\nCompilation took\n", end - start, "\n")
 
-        start = time.time()
-        carry = compiled_train_fun(carry)
-        jax.block_until_ready(carry)
-        end = time.time()
-        print("\nTraining took\n", end - start, "\n")
+            start = time.time()
+            carry = compiled_train_fun(carry)
+            jax.block_until_ready(carry)
+            end = time.time()
+            if verbose:
+                print("\nTraining took\n", end - start, "\n")
+        else:
+            carry = train_fun(carry)
 
     (
         i,

jinns/utils/__init__.py

@@ -1,6 +1 @@
-from ._pinn import create_PINN, PINN
-from ._ppinn import create_PPINN, PPINN
-from ._spinn import create_SPINN, SPINN
-from ._hyperpinn import create_HYPERPINN, HYPERPINN
-from ._save_load import save_pinn, load_pinn
 from ._utils import get_grid

jinns/utils/_types.py

@@ -26,9 +26,9 @@ if TYPE_CHECKING:
 
     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.nn._pinn import PINN
+    from jinns.nn._hyperpinn import HyperPINN
+    from jinns.nn._spinn_mlp import SPINN
     from jinns.utils._containers import *
     from jinns.validation._validation import AbstractValidationModule
 
@@ -42,7 +42,7 @@ if TYPE_CHECKING:
         DataGeneratorODE | CubicMeshPDEStatio | CubicMeshPDENonStatio
     )
 
-    AnyPINN: TypeAlias = PINN | HYPERPINN | SPINN
+    AnyPINN: TypeAlias = PINN | HyperPINN | SPINN
 
     AnyBatch: TypeAlias = ODEBatch | PDEStatioBatch | PDENonStatioBatch
     rar_operands = NewType(

{jinns-1.2.0.dist-info → jinns-1.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: jinns
-Version: 1.2.0
+Version: 1.3.0
 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>
@@ -112,16 +112,16 @@ pre-commit install
 
 Don't hesitate to contribute and get your name on the list here !
 
-**List of contributors:** Hugo Gangloff, Nicolas Jouvin
+**List of contributors:** Hugo Gangloff, Nicolas Jouvin, Lucia Clarotto, Inass Soukarieh
 
 # Cite us
 
-Please consider citing our work if you found it useful to yours, using the following lines
+Please consider citing our work if you found it useful to yours, using this [ArXiV preprint](https://arxiv.org/abs/2412.14132)
 ```
-@software{jinns2024,
-    title={\texttt{jinns}: Physics-Informed Neural Networks with JAX},
-    author={Gangloff, Hugo and Jouvin, Nicolas},
-    url={https://gitlab.com/mia_jinns},
-    year={2024}
+@article{gangloff_jouvin2024jinns,
+  title={jinns: a JAX Library for Physics-Informed Neural Networks},
+  author={Gangloff, Hugo and Jouvin, Nicolas},
+  journal={arXiv preprint arXiv:2412.14132},
+  year={2024}
 }
 ```

jinns-1.3.0.dist-info/RECORD

@@ -0,0 +1,44 @@
+jinns/__init__.py,sha256=5p7V5VJd7PXEINqhqS4mUsnQtXlyPwfctRhL4p0loFg,181
+jinns/data/_Batchs.py,sha256=oc7-N1wEbsEvbe9fjVFKG2OPoZJVEjzPm8uj_icACf4,817
+jinns/data/_DataGenerators.py,sha256=3pyUqzQ12AUBqOV-yqpt4X6K_7CqTFtUKMjg-gJE6KA,65101
+jinns/data/__init__.py,sha256=TRCH0Z4-SQZ50MbSf46CUYWBkWVDmXCyez9T-EGiv_8,338
+jinns/experimental/__init__.py,sha256=3jCIy2R2i_0Erwxg-HwISdH79Nt1XCXhS9yY1F5awiY,208
+jinns/experimental/_diffrax_solver.py,sha256=upMr3kTTNrxEiSUO_oLvCXcjS9lPxSjvbB81h3qlhaU,6813
+jinns/loss/_DynamicLoss.py,sha256=lUpFl37_TfwxSREpoVKqUOpQEVqD3hrFXqwP2GZReWw,25817
+jinns/loss/_DynamicLossAbstract.py,sha256=bqmPxyrcvZh_dL74DTpj-TGiFxchvG8qC6KhuGeyOoA,12006
+jinns/loss/_LossODE.py,sha256=QhhSyJpDbcyW4TdShX0HkxbvJQWXvnYg8lik8_wyOg4,23415
+jinns/loss/_LossPDE.py,sha256=DZPinl7KYV2vp_CdjnhaR9M_gE-WOvyi4s8VSDEgti0,51046
+jinns/loss/__init__.py,sha256=PRiJV9fd2GSwaCBVCPyh6pFc6pdA40jfb_T1YvO8ERc,712
+jinns/loss/_boundary_conditions.py,sha256=kxHwNFSMsNzFso6nvAewcAdzW50yTi7IX-5Pthe65XY,12271
+jinns/loss/_loss_utils.py,sha256=IkZAWmBumNWwk3hzeO0dh5RjHKZpt_hL4XnG5-Gpfr8,14690
+jinns/loss/_loss_weights.py,sha256=F0Fgji2XpVk3pr9oIryGuXcG1FGQo4Dv6WFgze2BtA0,2201
+jinns/loss/_operators.py,sha256=qaRxwqgnZzlE_zTyUvafZGnUH5EZY1lpgjT9Vb7QJAQ,21718
+jinns/nn/__init__.py,sha256=k9guJSKmKlHEadAjU-0HlYXJe55Tt783QrkZz6EYyO8,231
+jinns/nn/_hyperpinn.py,sha256=nH8c9DeiiAujprEd7CVKU1chWn-kcSAY-fYLzd8_ikY,18049
+jinns/nn/_mlp.py,sha256=AbbFLF85ayJcQ6kVwfSNdAvjP69UWBP6Z3V-1De-pI4,8028
+jinns/nn/_pinn.py,sha256=45lXgrZQHv-7PQ3EDWWIoo8FlXRnjL1nl7mALTSJ45o,8391
+jinns/nn/_ppinn.py,sha256=vqIH_v1DF3LoHyl3pJ1qhfnGMRMfvbfNK6m9s5LC21k,9212
+jinns/nn/_save_load.py,sha256=VaO9LtR6dajEfo8iP7FgOvyLdQxT2IawazC2sxs97lc,9139
+jinns/nn/_spinn.py,sha256=QmKhDZ0-ToJk3_glQ9BQWgoC0d-EEAWxMrDeHfB2slw,4191
+jinns/nn/_spinn_mlp.py,sha256=9iU_-TIUFMVBcYv0nQmsa07ZwApIKqnXm7v4CY87PTo,7224
+jinns/parameters/__init__.py,sha256=1gxNLoAXUjhUzBWuh86YjU5pYy8SOboCs8TrKcU1wZc,158
+jinns/parameters/_derivative_keys.py,sha256=UyEcgfNF1vwPcGWD2ShAZkZiq4thzRDm_OUJzOfjjiY,21909
+jinns/parameters/_params.py,sha256=wK9ZSqoL9KnjOWqc_ZhJ09ffbsgeUEcttc1Rhme0lLk,3550
+jinns/plot/__init__.py,sha256=Q279h5veYWNLQyttsC8_tDOToqUHh8WaRON90CiWXqk,81
+jinns/plot/_plot.py,sha256=6OqCNvOeqbat3dViOtehILbRfGIS3pnTmNRfbZYaVTA,11433
+jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jinns/solver/_rar.py,sha256=JU4FgWt5w3tzgn2mNyftGi8Erxn5N0Za60-lRaL2poI,9724
+jinns/solver/_solve.py,sha256=Bh7uplfcInJEQj1wmMquisN_vvUghARgX_uaYf7NUpw,23423
+jinns/solver/_utils.py,sha256=b2zYvwZY_fU0NMNWvUEMvHez9s7hwcxfpGzQlz5F6HA,5762
+jinns/utils/__init__.py,sha256=uw3I-lWT3wLabo6-H8FbKpSXI2xobzSs2W-Xno280g0,29
+jinns/utils/_containers.py,sha256=a7A-iUApnjc1YVc7bdt9tKUvHHPDOKMB9OfdrDZGWN8,1450
+jinns/utils/_types.py,sha256=4Qgsg6r9UPGpRwmERv4Cx2nU5ZIweehDlZQPo-FuR4Y,1896
+jinns/utils/_utils.py,sha256=hoRcJqcTuQi_Ip40oI4EbxW46E1rp2C01_HfuCpwKRM,2932
+jinns/validation/__init__.py,sha256=Jv58mzgC3F7cRfXA6caicL1t_U0UAhbwLrmMNVg6E7s,66
+jinns/validation/_validation.py,sha256=bvqL2poTFJfn9lspWqMqXvQGcQIodKwKrC786QtEZ7A,4700
+jinns-1.3.0.dist-info/AUTHORS,sha256=7NwCj9nU-HNG1asvy4qhQ2w7oZHrn-Lk5_wK_Ve7a3M,80
+jinns-1.3.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-1.3.0.dist-info/METADATA,sha256=PM3iLQFd-vHDU697ECGjD2vQpgxo1vo1GTFl5AdIWoo,4744
+jinns-1.3.0.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+jinns-1.3.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-1.3.0.dist-info/RECORD,,

{jinns-1.2.0.dist-info → jinns-1.3.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
+Generator: setuptools (75.8.2)
 Root-Is-Purelib: true
 Tag: py3-none-any