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

jinns/utils/_ppinn.py

@@ -0,0 +1,227 @@
+"""
+Implements utility function to create PINNs
+"""
+
+from typing import Callable, Literal
+from dataclasses import InitVar
+import jax
+import jax.numpy as jnp
+import equinox as eqx
+
+from jaxtyping import Array, Key, PyTree, Float
+
+from jinns.parameters._params import Params, ParamsDict
+
+from jinns.utils._pinn import PINN, _MLP
+
+
+class PPINN(PINN):
+    r"""
+    A PPINN (Parallel PINN) object which mimicks the PFNN architecture from
+    DeepXDE. This is in fact a PINN that encompasses several PINNs internally.
+
+    Parameters
+    ----------
+    slice_solution : slice
+        A jnp.s\_ object which indicates which axis of the PINN output is
+        dedicated to the actual equation solution. Default None
+        means that slice_solution = the whole PINN output. This argument is useful
+        when the PINN is also used to output equation parameters for example
+        Note that it must be a slice and not an integer (a preprocessing of the
+        user provided argument takes care of it).
+    eq_type : 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.
+    input_transform : Callable[[Float[Array, "input_dim"], Params], Float[Array, "output_dim"]]
+        A function that will be called before entering the PINN. Its output(s)
+        must match the PINN inputs (except for the parameters).
+        Its inputs are the PINN inputs (`t` and/or `x` concatenated together)
+        and the parameters. Default is no operation.
+    output_transform : Callable[[Float[Array, "input_dim"], Float[Array, "output_dim"], Params], Float[Array, "output_dim"]]
+        A function with arguments begin the same input as the PINN, the PINN
+        output and the parameter. This function will be called after exiting the PINN.
+        Default is no operation.
+    mlp_list : list[eqx.Module]
+        The actual neural networks instanciated as eqx.Modules
+    """
+
+    slice_solution: slice = eqx.field(static=True, kw_only=True)
+    output_slice: slice = eqx.field(static=True, kw_only=True, default=None)
+
+    mlp_list: InitVar[list[eqx.Module]] = eqx.field(kw_only=True)
+
+    params: PyTree = eqx.field(init=False)
+    static: PyTree = eqx.field(init=False, static=True)
+
+    def __post_init__(self, mlp, mlp_list):
+        super().__post_init__(
+            mlp=mlp_list[0],
+        )
+        self.params, self.static = (), ()
+        for mlp in mlp_list:
+            params, static = eqx.partition(mlp, eqx.is_inexact_array)
+            self.params = self.params + (params,)
+            self.static = self.static + (static,)
+
+    @property
+    def init_params(self) -> PyTree:
+        """
+        Returns an initial set of parameters
+        """
+        return self.params
+
+    def __call__(
+        self,
+        inputs: Float[Array, "1"] | Float[Array, "dim"] | Float[Array, "1+dim"],
+        params: PyTree,
+    ) -> Float[Array, "output_dim"]:
+        """
+        Evaluate the PPINN on some inputs with some params.
+        """
+        if len(inputs.shape) == 0:
+            # This can happen often when the user directly provides some
+            # collocation points (eg for plotting, whithout using
+            # DataGenerators)
+            inputs = inputs[None]
+        transformed_inputs = self.input_transform(inputs, params)
+
+        outs = []
+        for params_, static in zip(params.nn_params, self.static):
+            model = eqx.combine(params_, static)
+            outs += [model(transformed_inputs)]
+        # Note that below is then a global output transform
+        res = self.output_transform(inputs, jnp.concatenate(outs, axis=0), params)
+
+        ## force (1,) output for non vectorial solution (consistency)
+        if not res.shape:
+            return jnp.expand_dims(res, axis=-1)
+        return res
+
+
+def create_PPINN(
+    key: Key,
+    eqx_list_list: list[tuple[tuple[Callable, int, int] | Callable, ...]],
+    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+    dim_x: int = 0,
+    input_transform: Callable[
+        [Float[Array, "input_dim"], Params], Float[Array, "output_dim"]
+    ] = None,
+    output_transform: Callable[
+        [Float[Array, "input_dim"], Float[Array, "output_dim"], Params],
+        Float[Array, "output_dim"],
+    ] = None,
+    slice_solution: slice = None,
+) -> tuple[PINN | list[PINN], PyTree | list[PyTree]]:
+    r"""
+    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_list
+        A list of `eqx_list` (see `create_PINN`). The input dimension must be the
+        same for each sub-`eqx_list`. Then the parallel subnetworks can be
+        different. Their respective outputs are concatenated.
+    eq_type
+        A string with three possibilities.
+        "ODE": the PPINN is called with one input `t`.
+        "statio_PDE": the PPINN is called with one input `x`, `x`
+        can be high dimensional.
+        "nonstatio_PDE": the PPINN 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.
+    dim_x
+        An integer. The dimension of `x`. Default `0`.
+    input_transform
+        A function that will be called before entering the PPINN. Its output(s)
+        must match the PPINN inputs (except for the parameters).
+        Its inputs are the PPINN inputs (`t` and/or `x` concatenated together)
+        and the parameters. Default is no operation.
+    output_transform
+        This function will be called after exiting
+        the PPINN, i.e., on the concatenated outputs of all parallel networks
+        Default is no operation.
+    slice_solution
+        A jnp.s\_ object which indicates which axis of the PPINN output is
+        dedicated to the actual equation solution. Default None
+        means that slice_solution = the whole PPINN output. This argument is
+        useful when the PPINN is also used to output equation parameters for
+        example Note that it must be a slice and not an integer (a
+        preprocessing of the user provided argument takes care of it).
+
+
+    Returns
+    -------
+    ppinn
+        A PPINN instance
+    ppinn.init_params
+        An initial set of parameters for the PPINN
+
+    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")
+
+    nb_outputs_declared = 0
+    for eqx_list in eqx_list_list:
+        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]
+
+    if slice_solution is None:
+        slice_solution = jnp.s_[0:nb_outputs_declared]
+    if isinstance(slice_solution, int):
+        # rewrite it as a slice to ensure that axis does not disappear when
+        # indexing
+        slice_solution = jnp.s_[slice_solution : slice_solution + 1]
+
+    if input_transform is None:
+
+        def input_transform(_in, _params):
+            return _in
+
+    if output_transform is None:
+
+        def output_transform(_in_pinn, _out_pinn, _params):
+            return _out_pinn
+
+    mlp_list = []
+    for eqx_list in eqx_list_list:
+        mlp_list.append(_MLP(key=key, eqx_list=eqx_list))
+
+    ppinn = PPINN(
+        mlp=None,
+        mlp_list=mlp_list,
+        slice_solution=slice_solution,
+        eq_type=eq_type,
+        input_transform=input_transform,
+        output_transform=output_transform,
+    )
+    return ppinn, ppinn.init_params

jinns/utils/_save_load.py

@@ -20,18 +20,18 @@ def function_to_string(
     We need this transformation for eqx_list to be pickled
 
     From `((eqx.nn.Linear, 2, 20),
-            (jax.nn.tanh),
+            (jax.nn.tanh,),
             (eqx.nn.Linear, 20, 20),
-            (jax.nn.tanh),
+            (jax.nn.tanh,),
             (eqx.nn.Linear, 20, 20),
-            (jax.nn.tanh),
+            (jax.nn.tanh,),
             (eqx.nn.Linear, 20, 1))` to
     `(("Linear", 2, 20),
-                ("tanh"),
+                ("tanh",),
                 ("Linear", 20, 20),
-                ("tanh"),
+                ("tanh",),
                 ("Linear", 20, 20),
-                ("tanh"),
+                ("tanh",),
                 ("Linear", 20, 1))`
     """
     return jax.tree_util.tree_map(
@@ -210,14 +210,16 @@ def load_pinn(
     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)
+        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)
+        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)
+        u_reloaded_shallow, _ = eqx.filter_eval_shape(
+            create_HYPERPINN, **kwargs_reloaded
+        )
     else:
         raise ValueError(f"{type_} is not valid")
     if key_list_for_paramsdict is None:
@@ -226,15 +228,13 @@ def load_pinn(
         u_reloaded = eqx.tree_deserialise_leaves(
             filename + "-module.eqx", u_reloaded_shallow
         )
-        params = Params(
-            nn_params=u_reloaded.init_params(), eq_params=eq_params_reloaded
-        )
+        params = Params(nn_params=u_reloaded.init_params, eq_params=eq_params_reloaded)
     else:
         nn_params_dict = {}
         for key in key_list_for_paramsdict:
             u_reloaded = eqx.tree_deserialise_leaves(
                 filename + f"-module_{key}.eqx", u_reloaded_shallow
             )
-            nn_params_dict[key] = u_reloaded.init_params()
+            nn_params_dict[key] = u_reloaded.init_params
         params = ParamsDict(nn_params=nn_params_dict, eq_params=eq_params_reloaded)
     return u_reloaded, params

jinns/utils/_spinn.py

@@ -10,6 +10,8 @@ import jax.numpy as jnp
 import equinox as eqx
 from jaxtyping import Key, Array, Float, PyTree
 
+from jinns.parameters._params import Params, ParamsDict
+
 
 class _SPINN(eqx.Module):
     """
@@ -21,7 +23,8 @@ class _SPINN(eqx.Module):
     key : InitVar[Key]
         A jax random key for the layer initializations.
     d : int
-        The number of dimensions to treat separately.
+        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
@@ -62,15 +65,11 @@ class _SPINN(eqx.Module):
             self.separated_mlp.append(self.layers)
 
     def __call__(
-        self, t: Float[Array, "1"], x: Float[Array, "omega_dim"]
+        self, inputs: Float[Array, "dim"] | Float[Array, "dim+1"]
     ) -> Float[Array, "d embed_dim*output_dim"]:
-        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]
+            t_ = inputs[d : d + 1]
             for layer in self.separated_mlp[d]:
                 t_ = layer(t_)
             outputs += [t_]
@@ -82,13 +81,11 @@ class SPINN(eqx.Module):
     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.
+        The number of dimensions to treat separately, including time `t` if
+        used for non-stationnary equations.
 
     """
 
@@ -105,42 +102,28 @@ class SPINN(eqx.Module):
     def __post_init__(self, spinn_mlp):
         self.params, self.static = eqx.partition(spinn_mlp, eqx.is_inexact_array)
 
+    @property
     def init_params(self) -> PyTree:
         """
         Returns an initial set of parameters
         """
         return self.params
 
-    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, 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)
-        if self.eq_type == "nonstatio_PDE":
-            (t, x, params) = args
-            try:
-                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)
-        raise RuntimeError("Wrong parameter value for eq_type")
-
-    def eval_nn(
-        self, res: Float[Array, "d embed_dim*output_dim"]
+    def __call__(
+        self,
+        t_x: Float[Array, "batch_size 1+dim"],
+        params: Params | ParamsDict | PyTree,
     ) -> Float[Array, "output_dim"]:
         """
         Evaluate the SPINN on some inputs with some params.
         """
+        try:
+            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)
+        res = v_model(t_x)
+
         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])])
         res = jnp.stack(
@@ -170,7 +153,7 @@ def create_SPINN(
     eqx_list: tuple[tuple[Callable, int, int] | Callable, ...],
     eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
     m: int = 1,
-) -> SPINN:
+) -> tuple[SPINN, PyTree]:
     """
     Utility function to create a SPINN neural network with the equinox
     library.
@@ -218,16 +201,14 @@ def create_SPINN(
         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
 
 
     Returns
     -------
     spinn
         An instanciated SPINN
+    spinn.init_params
+        The initial set of parameters of the model
 
     Raises
     ------
@@ -265,4 +246,4 @@ def create_SPINN(
     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
+    return spinn, spinn.init_params

jinns/utils/_types.py

@@ -1,3 +1,4 @@
+# pragma: exclude file
 from __future__ import (
     annotations,
 )  # https://docs.python.org/3/library/typing.html#constant

jinns/utils/_utils.py

@@ -2,13 +2,18 @@
 Implements various utility functions
 """
 
-from functools import reduce
-from operator import getitem
-import numpy as np
+from math import prod
+import warnings
 import jax
 import jax.numpy as jnp
 from jaxtyping import PyTree, Array
 
+from jinns.data._DataGenerators import (
+    DataGeneratorODE,
+    CubicMeshPDEStatio,
+    CubicMeshPDENonStatio,
+)
+
 
 def _check_nan_in_pytree(pytree: PyTree) -> bool:
     """
@@ -33,7 +38,7 @@ def _check_nan_in_pytree(pytree: PyTree) -> bool:
     )
 
 
-def _get_grid(in_array: Array) -> 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
@@ -49,10 +54,14 @@ def _get_grid(in_array: Array) -> Array:
     return in_array
 
 
-def _check_user_func_return(r: Array | int, shape: tuple) -> Array | int:
+def _check_shape_and_type(
+    r: Array | int, expected_shape: tuple, cause: str = "", binop: str = ""
+) -> Array | float:
     """
-    Correctly handles the result from a user defined function (eg a boundary
-    condition) to get the correct broadcast
+    Ensures float type and correct shapes for broadcasting when performing a
+    binary operation (like -, + or *) between two arrays.
+    First array is a custom user (observation data or output of initial/BC
+    functions), the expected shape is the same as the PINN's.
     """
     if isinstance(r, (int, float)):
         # if we have a scalar cast it to float
@@ -60,9 +69,28 @@ def _check_user_func_return(r: Array | int, shape: tuple) -> Array | int:
     if r.shape == ():
         # if we have a scalar inside a ndarray
         return r.astype(float)
-    if r.shape[-1] == shape[-1]:
-        # the broadcast will be OK
+    if r.shape[-1] == expected_shape[-1]:
+        # broadcasting will be OK
         return r.astype(float)
-    # the reshape below avoids a missing (1,) ending dimension
-    # depending on how the user has coded the inital function
-    return r.reshape(shape)
+
+    if r.shape != expected_shape:
+        # Usually, the reshape below  adds a missing (1,) final axis to ensure # the PINN output and the other function (initial/boundary condition)
+        # have the correct shape, depending on how the user has coded the
+        # initial/boundary condition.
+        warnings.warn(
+            f"[{cause}] Performing operation `{binop}` between arrays"
+            f" of different shapes: got {r.shape} for the custom array and"
+            f" {expected_shape} for the PINN."
+            f" This can cause unexpected and wrong broadcasting."
+            f" Reshaping {r.shape} into {expected_shape}. Reshape your"
+            f" custom array to math the {expected_shape=} to prevent this"
+            f" warning."
+        )
+    return r.reshape(expected_shape)
+
+
+def _subtract_with_check(
+    a: Array | int, b: Array | int, cause: str = ""
+) -> Array | float:
+    a = _check_shape_and_type(a, b.shape, cause=cause, binop="-")
+    return a - b

jinns-1.2.0.dist-info/AUTHORS

@@ -0,0 +1,2 @@
+Hugo Gangloff <hugo.gangloff@inrae.fr>
+Nicolas Jouvin <nicolas.jouvin@inrae.fr>

jinns-1.2.0.dist-info/METADATA

@@ -0,0 +1,127 @@
+Metadata-Version: 2.1
+Name: jinns
+Version: 1.2.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>
+License: Apache License 2.0
+Project-URL: Repository, https://gitlab.com/mia_jinns/jinns
+Project-URL: Documentation, https://mia_jinns.gitlab.io/jinns/index.html
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Development Status :: 4 - Beta
+Classifier: Programming Language :: Python
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: AUTHORS
+Requires-Dist: numpy
+Requires-Dist: jax
+Requires-Dist: jaxopt
+Requires-Dist: optax
+Requires-Dist: equinox>0.11.3
+Requires-Dist: jax-tqdm
+Requires-Dist: diffrax
+Requires-Dist: matplotlib
+Provides-Extra: notebook
+Requires-Dist: jupyter; extra == "notebook"
+Requires-Dist: seaborn; extra == "notebook"
+
+jinns
+=====
+
+![status](https://gitlab.com/mia_jinns/jinns/badges/main/pipeline.svg) ![coverage](https://gitlab.com/mia_jinns/jinns/badges/main/coverage.svg)
+
+Physics Informed Neural Networks with JAX. **jinns** is developed to estimate solutions of ODE and PDE problems using neural networks, with a strong focus on
+
+ 1. inverse problems: find equation parameters given noisy/indirect observations
+ 2. meta-modeling: solve for a parametric family of differential equations
+
+It can also be used for forward problems and hybrid-modeling.
+
+**jinns** specific points:
+
+- **jinns uses JAX** - It is directed to JAX users: forward and backward autodiff, vmapping, jitting and more! No reinventing the wheel: it relies on the JAX ecosystem whenever possible, such as [equinox](https://github.com/patrick-kidger/equinox/) for neural networks or [optax](https://optax.readthedocs.io/) for optimization.
+
+- **jinns is highly modular** - It gives users maximum control for defining their problems, and extending the package. The maths and computations are visible and not hidden behind layers of code!
+
+- **jinns is efficient** - It compares favorably to other existing Python package for PINNs on the [PINNacle benchmarks](https://github.com/i207M/PINNacle/), as demonstrated in the table below. For more details on the benchmarks, checkout the [PINN multi-library benchmark](https://gitlab.com/mia_jinns/pinn-multi-library-benchmark)
+
+- Implemented PINN architectures
+    - Vanilla Multi-Layer Perceptron popular accross the PINNs litterature.
+
+    - [Separable PINNs](https://openreview.net/pdf?id=dEySGIcDnI): allows to leverage forward-mode autodiff for computational speed.
+
+    - [Hyper PINNs](https://arxiv.org/pdf/2111.01008.pdf): useful for meta-modeling
+
+
+- **Get started**: check out our various notebooks on the [documentation](https://mia_jinns.gitlab.io/jinns/index.html).
+
+|  | jinns | DeepXDE - JAX | DeepXDE - Pytorch | PINA | Nvidia Modulus |
+|---|:---:|:---:|:---:|:---:|:---:|
+| Burgers1D | **445** | 723 | 671 | 1977 | 646 |
+| NS2d-C | **265** | 278 | 441 | 1600 | 275 |
+| PInv | 149 | 218 | *CC* | 1509 | **135** |
+| Diffusion-Reaction-Inv | **284** | *NI* | 3424 | 4061 | 2541 |
+| Navier-Stokes-Inv | **175** | *NI* | 1511 | 1403 | 498 |
+
+*Training time in seconds on an Nvidia T600  GPU. NI means problem cannot be implemented in the backend, CC means the code crashed.*
+
+![A diagram of jinns workflow](img/jinns-diagram.png)
+
+
+# Installation
+
+Install the latest version with pip
+
+```bash
+pip install jinns
+```
+
+# Documentation
+
+The project's documentation is hosted on Gitlab page and available at [https://mia_jinns.gitlab.io/jinns/index.html](https://mia_jinns.gitlab.io/jinns/index.html).
+
+
+# Found a bug / want a feature ?
+
+Open an issue on the [Gitlab repo](https://gitlab.com/mia_jinns/jinns/-/issues).
+
+
+# Contributing
+
+Here are the contributors guidelines:
+
+1. First fork the library on Gitlab.
+
+2. Then clone and install the library in development mode with
+
+```bash
+pip install -e .
+```
+
+3. Install pre-commit and run it.
+
+```bash
+pip install pre-commit
+pre-commit install
+```
+
+ 4. Open a merge request once you are done with your changes, the review will be done via Gitlab.
+
+# Contributors
+
+Don't hesitate to contribute and get your name on the list here !
+
+**List of contributors:** Hugo Gangloff, Nicolas Jouvin
+
+# Cite us
+
+Please consider citing our work if you found it useful to yours, using the following lines
+```
+@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}
+}
+```