jinns 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

jinns/utils/_utils.py

@@ -57,314 +57,21 @@ def _tracked_parameters(params, tracked_params_key_list):
     return tracked_params
 
 
-class _MLP(eqx.Module):
+def _get_grid(in_array):
     """
-    Class to construct an equinox module from a key and a eqx_list. To be used
-    in pair with the function `create_PINN`
+    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
     """
-
-    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
-
-
-def create_PINN(
-    key,
-    eqx_list,
-    eq_type,
-    dim_x=0,
-    with_eq_params=None,
-    input_transform=None,
-    output_transform=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
-
-
-    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
-
-    # 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")
-
-    def make_mlp(key, eqx_list):
-        mlp = _MLP(key, eqx_list)
-        params, static = eqx.partition(mlp, eqx.is_inexact_array)
-
-        def init_fn():
-            return params
-
-        if eq_type == "ODE":
-            if with_eq_params is None:
-
-                def apply_fn(t, u_params, eq_params=None):
-                    model = eqx.combine(u_params, 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:
-                            return model(input_transform(t)).squeeze()
-                        else:
-                            return model(t).squeeze()
-                    else:
-                        if input_transform is not None:
-                            return output_transform(
-                                t, model(input_transform(t)).squeeze()
-                            )
-                        else:
-                            return output_transform(t, model(t).squeeze())
-
-            else:
-
-                def apply_fn(t, u_params, eq_params):
-                    model = eqx.combine(u_params, 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:
-                            return model(input_transform(t_eq_params)).squeeze()
-                        else:
-                            return model(t_eq_params).squeeze()
-                    else:
-                        if input_transform is not None:
-                            return output_transform(
-                                t_eq_params,
-                                model(input_transform(t_eq_params)).squeeze(),
-                            )
-                        else:
-                            return output_transform(
-                                t_eq_params, model(t_eq_params).squeeze()
-                            )
-
-        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(x, u_params, eq_params=None):
-                    model = eqx.combine(u_params, static)
-
-                    if output_transform is None:
-                        if input_transform is not None:
-                            return model(input_transform(x)).squeeze()
-                        else:
-                            return model(x).squeeze()
-                    else:
-                        if input_transform is not None:
-                            return output_transform(
-                                x, model(input_transform(x)).squeeze()
-                            )
-                        else:
-                            return output_transform(x, model(x).squeeze())
-
-            else:
-
-                def apply_fn(x, u_params, eq_params):
-                    model = eqx.combine(u_params, 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:
-                            return model(input_transform(x_eq_params)).squeeze()
-                        else:
-                            return model(x_eq_params).squeeze()
-                    else:
-                        if input_transform is not None:
-                            return output_transform(
-                                x_eq_params,
-                                model(input_transform(x_eq_params)).squeeze(),
-                            )
-                        else:
-                            return output_transform(
-                                x_eq_params, model(x_eq_params).squeeze()
-                            )
-
-        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(t, x, u_params, eq_params=None):
-                    model = eqx.combine(u_params, static)
-                    t_x = jnp.concatenate([t, x], axis=-1)
-
-                    if output_transform is None:
-                        if input_transform is not None:
-                            return model(input_transform(t_x)).squeeze()
-                        else:
-                            return model(t_x).squeeze()
-                    else:
-                        if input_transform is not None:
-                            return output_transform(
-                                t_x, model(input_transform(t_x)).squeeze()
-                            )
-                        else:
-                            return output_transform(t_x, model(t_x).squeeze())
-
-            else:
-
-                def apply_fn(t, x, u_params, eq_params):
-                    model = eqx.combine(u_params, 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:
-                            return model(input_transform(t_x_eq_params)).squeeze()
-                        else:
-                            return model(t_x_eq_params).squeeze()
-                    else:
-                        if input_transform is not None:
-                            return output_transform(
-                                t_x_eq_params,
-                                model(input_transform(t_x_eq_params)).squeeze(),
-                            )
-                        else:
-                            return output_transform(
-                                t_x_eq_params,
-                                model(input_transform(t_x_eq_params)).squeeze(),
-                            )
-
-        else:
-            raise RuntimeError("Wrong parameter value for eq_type")
-
-        return init_fn, apply_fn
-
-    return make_mlp(key, eqx_list)
+    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):
@@ -392,6 +99,25 @@ def _get_vmap_in_axes_params(eq_params_batch_dict, params):
         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
 ):

{jinns-0.4.2.dist-info → jinns-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.4.2
+Version: 0.5.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>
@@ -29,6 +29,18 @@ jinns
 Physics Informed Neural Networks with JAX. **jinns** has been developed to estimate solutions to your ODE et PDE problems using neural networks.
 **jinns** is built on JAX.
 
+**jinns** specific points:
+
+- **jinns** is coded with JAX as a backend: forward and backward autodiff, vmapping, jitting and more!
+
+- We focus the development towards inverse problems and inference in mecanistic-statistical models
+
+- [Separable PINN](https://openreview.net/pdf?id=dEySGIcDnI) are implemented
+
+- Check out our various notebooks to get started with `jinns`
+
+For more information, open an issue or contact us!
+
 # Installation
 
 Install the latest version with pip
@@ -41,6 +53,8 @@ pip install jinns
 
 The project's documentation is available at [https://mia_jinns.gitlab.io/jinns/index.html](https://mia_jinns.gitlab.io/jinns/index.html)
 
+Note that all the tests were performed on a rather small Nvidia T600 GPU, expect a substancial performance gain on bigger devices.
+
 # Contributing
 
 * First fork the library on Gitlab.
@@ -56,7 +70,6 @@ pip install -e .
 ```bash
 pip install pre-commit
 pre-commit install
-pre-commit install --hook-type pre-push
 ```
 
 * Open a merge request once you are done with your changes.

jinns-0.5.0.dist-info/RECORD

@@ -0,0 +1,24 @@
+jinns/__init__.py,sha256=Nw5pdlmDhJwco3bXX3YttkeCF8czX_6m0poh8vu0lDQ,113
+jinns/data/_DataGenerators.py,sha256=nIuKtkX4V4ckfT4-g0bjlY7BLkgcok5JbI9OzJn73mA,44461
+jinns/data/__init__.py,sha256=S13J59Fxuph4uNJ542fP_Mj8U72ilhb5t_UQ-c1k3nY,232
+jinns/data/_display.py,sha256=NfINLJAGmQSPz30cWVaeQFpabzCXprp4RNH6Iycx-VU,7722
+jinns/loss/_DynamicLoss.py,sha256=AuW71-Kyt7tj_aatxroym9BcvqzTfsBI7qMdEF3021w,36550
+jinns/loss/_DynamicLossAbstract.py,sha256=V9NJvMmqnSC06yccu9bkFCWLF3cjyr8ze8qf0LRykjo,7718
+jinns/loss/_LossODE.py,sha256=FHTKQPqLSoMh18j_RYUoR7tUGg_ljd0JKL8xf2HiF5M,17541
+jinns/loss/_LossPDE.py,sha256=Pegyr4jZsI9hZMUGjRjohYeVVKLuSU53rc-epD0KZGI,64328
+jinns/loss/__init__.py,sha256=p9pXrttAgsx_V25Yt8lT8Guoc3z9OAfRRvzl9OIc3Gs,402
+jinns/loss/_boundary_conditions.py,sha256=OE7lDwCBKUJaXkOP5uJZyakT1YWb1D7YogbvFT0RaJo,16633
+jinns/loss/_operators.py,sha256=l9nrLGluD4Y-0L-yydYS3--mH-nKXLRymhm3JCapaZ0,12067
+jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jinns/solver/_rar.py,sha256=V9y07F6objmP6rPA305dIJ82h7kwP4AWBAktZ68b-38,13894
+jinns/solver/_seq2seq.py,sha256=XNL9e0fBj85Q86XfGDzq9dzxIkPPMwoJF38C8doNYtM,6032
+jinns/solver/_solve.py,sha256=Yz_asD0ZuYN923E6ysxtPemN-36Zt-BtlSPkTX-BfA8,10047
+jinns/utils/__init__.py,sha256=ClGy9Ppye1z75daH_3Ngb0406K-AJcZHgj4gM5vfX_8,196
+jinns/utils/_pinn.py,sha256=-UmdSRMzkqe_TEsuVHCtu8b_9Na3wz-5YDoKSAeF8ZM,11307
+jinns/utils/_spinn.py,sha256=SYz8eatM_AkIPYhNf2sOKoGGhtAWKZX47W8xDXR0aK0,8027
+jinns/utils/_utils.py,sha256=zsSPoqQi9OQKTRiCRfeINCwCfzo7MHFCc4zPCIU2UpY,8790
+jinns-0.5.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.5.0.dist-info/METADATA,sha256=C8fZOd7PqroIKIwS80YUIvwbaLbNVuYPxGCv0JG4PkM,2338
+jinns-0.5.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+jinns-0.5.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.5.0.dist-info/RECORD,,

jinns-0.4.2.dist-info/RECORD

@@ -1,22 +0,0 @@
-jinns/__init__.py,sha256=Nw5pdlmDhJwco3bXX3YttkeCF8czX_6m0poh8vu0lDQ,113
-jinns/data/_DataGenerators.py,sha256=nIuKtkX4V4ckfT4-g0bjlY7BLkgcok5JbI9OzJn73mA,44461
-jinns/data/__init__.py,sha256=S13J59Fxuph4uNJ542fP_Mj8U72ilhb5t_UQ-c1k3nY,232
-jinns/data/_display.py,sha256=Xnfo6_PH1g-ZFpWJcbF6CF6Pp12wJtNQb1W1bADuQrA,6134
-jinns/loss/_DynamicLoss.py,sha256=VyoyWdkoxRPeP2vs4ZZBK_T9xWgwkcDuaFrjUSid3Zo,52975
-jinns/loss/_DynamicLossAbstract.py,sha256=V9NJvMmqnSC06yccu9bkFCWLF3cjyr8ze8qf0LRykjo,7718
-jinns/loss/_LossODE.py,sha256=FHTKQPqLSoMh18j_RYUoR7tUGg_ljd0JKL8xf2HiF5M,17541
-jinns/loss/_LossPDE.py,sha256=8wEEPnkibNOdqCXtcuU5nieRsPoffSyV8wcUMjcUkvg,57145
-jinns/loss/__init__.py,sha256=4JxMHHVMxTMsVZmV8mRSIyMVAEp3QIR8QtZvnwvj96Q,560
-jinns/loss/_boundary_conditions.py,sha256=qe1fHnMduVM1h89feHOc-O2--Zgd4ZYM4XjqHYUjuvw,9970
-jinns/loss/_operators.py,sha256=HYGDq3K_K7lztDp6al88Q78F6o-hDRo3ODhw8UMRdC8,5690
-jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jinns/solver/_rar.py,sha256=V9y07F6objmP6rPA305dIJ82h7kwP4AWBAktZ68b-38,13894
-jinns/solver/_seq2seq.py,sha256=XNL9e0fBj85Q86XfGDzq9dzxIkPPMwoJF38C8doNYtM,6032
-jinns/solver/_solve.py,sha256=Yz_asD0ZuYN923E6ysxtPemN-36Zt-BtlSPkTX-BfA8,10047
-jinns/utils/__init__.py,sha256=-jDlwCjyEzWweswKdwLal3OhaUU3FVzK_Ge2S-7KHXs,149
-jinns/utils/_utils.py,sha256=bQm6z_xPKJj9BMCr2tXc44IA8JyGNN-PR5LNRhZ1fD8,20085
-jinns-0.4.2.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.4.2.dist-info/METADATA,sha256=maaTIojnCdHhTIPd2kCJa8QGfX36mWAt1DN7q8Pd_3o,1821
-jinns-0.4.2.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-jinns-0.4.2.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.4.2.dist-info/RECORD,,