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

jinns/solver/_solve.py

@@ -8,55 +8,71 @@ from __future__ import (
 )  # https://docs.python.org/3/library/typing.html#constant
 
 import time
-from typing import TYPE_CHECKING, NamedTuple, Dict, Union
+from typing import TYPE_CHECKING, Any, TypeAlias, Callable
 from functools import partial
 import optax
 import jax
 from jax import jit
 import jax.numpy as jnp
-from jaxtyping import Int, Bool, Float, Array
+from jaxtyping import Float, Array
 from jinns.solver._rar import init_rar, trigger_rar
 from jinns.utils._utils import _check_nan_in_pytree
 from jinns.solver._utils import _check_batch_size
-from jinns.utils._containers import *
-from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-    append_obs_batch,
-    append_param_batch,
+from jinns.utils._containers import (
+    DataGeneratorContainer,
+    OptimizationContainer,
+    OptimizationExtraContainer,
+    LossContainer,
+    StoredObjectContainer,
 )
+from jinns.data._utils import append_param_batch, append_obs_batch
 
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.parameters._params import Params
+    from jinns.utils._types import AnyLoss, AnyBatch
+    from jinns.validation._validation import AbstractValidationModule
+    from jinns.data._DataGeneratorParameter import DataGeneratorParameter
+    from jinns.data._DataGeneratorObservations import DataGeneratorObservations
+    from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+
+    main_carry: TypeAlias = tuple[
+        int,
+        AnyLoss,
+        OptimizationContainer,
+        OptimizationExtraContainer,
+        DataGeneratorContainer,
+        AbstractValidationModule | None,
+        LossContainer,
+        StoredObjectContainer,
+        Float[Array, " n_iter"] | None,
+    ]
 
 
 def solve(
-    n_iter: Int,
-    init_params: AnyParams,
-    data: AnyDataGenerator,
+    n_iter: int,
+    init_params: Params[Array],
+    data: AbstractDataGenerator,
     loss: AnyLoss,
     optimizer: optax.GradientTransformation,
-    print_loss_every: Int = 1000,
-    opt_state: Union[NamedTuple, None] = None,
-    tracked_params: Params | ParamsDict | None = None,
+    print_loss_every: int = 1000,
+    opt_state: optax.OptState | None = None,
+    tracked_params: Params[Any | None] | None = None,
     param_data: DataGeneratorParameter | None = None,
-    obs_data: (
-        DataGeneratorObservations | DataGeneratorObservationsMultiPINNs | None
-    ) = None,
+    obs_data: DataGeneratorObservations | None = None,
     validation: AbstractValidationModule | None = None,
     obs_batch_sharding: jax.sharding.Sharding | None = None,
-    verbose: Bool = True,
+    verbose: bool = True,
+    ahead_of_time: bool = True,
 ) -> tuple[
-    Params | ParamsDict,
-    Float[Array, "n_iter"],
-    Dict[str, Float[Array, "n_iter"]],
-    AnyDataGenerator,
+    Params[Array],
+    Float[Array, " n_iter"],
+    dict[str, Float[Array, " n_iter"]],
+    AbstractDataGenerator,
     AnyLoss,
-    NamedTuple,
-    AnyParams,
-    Float[Array, "n_iter"],
-    AnyParams,
+    optax.OptState,
+    Params[Array | None],
+    Float[Array, " n_iter"] | None,
+    Params[Array],
 ]:
     """
     Performs the optimization process via stochastic gradient descent
@@ -93,8 +109,7 @@ def solve(
         Default None. A DataGeneratorParameter object which can be used to
         sample equation parameters.
     obs_data
-        Default None. A DataGeneratorObservations or
-        DataGeneratorObservationsMultiPINNs
+        Default None. A DataGeneratorObservations
         object which can be used to sample minibatches of observations.
     validation
         Default None. Otherwise, a callable ``eqx.Module`` which implements a
@@ -118,6 +133,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
     -------
@@ -162,11 +185,21 @@ def solve(
             _check_batch_size(obs_data, param_data, "n")
 
     if opt_state is None:
-        opt_state = optimizer.init(init_params)
+        opt_state = optimizer.init(init_params)  # type: ignore
+        # our Params are eqx.Module (dataclass + PyTree), PyTree is
+        # compatible with optax transform but not dataclass, this leads to a
+        # type hint error: we could prevent this by ensuring with the eqx.filter that
+        # we have only floating points optimizable params given to optax
+        # see https://docs.kidger.site/equinox/faq/#optax-throwing-a-typeerror
+        # opt_state = optimizer.init(
+        #    eqx.filter(init_params, eqx.is_inexact_array)
+        # )
+        # but this seems like a hack and there is no better way
+        # https://github.com/google-deepmind/optax/issues/384
 
     # RAR sampling init (ouside scanned function to avoid dynamic slice error)
     # If RAR is not used the _rar_step_*() are juste None and data is unchanged
-    data, _rar_step_true, _rar_step_false = init_rar(data)
+    data, _rar_step_true, _rar_step_false = init_rar(data)  # type: ignore
 
     # Seq2seq
     curr_seq = 0
@@ -283,7 +316,7 @@ def solve(
         if verbose:
             _print_fn(i, train_loss_value, print_loss_every, prefix="[train] ")
 
-        if validation is not None:
+        if validation is not None and validation_crit_values is not None:
             # there is a jax.lax.cond because we do not necesarily call the
             # validation step every iteration
             (
@@ -297,7 +330,7 @@ def solve(
                 lambda operands: (
                     operands[0],
                     False,
-                    validation_crit_values[i - 1],
+                    validation_crit_values[i - 1],  # type: ignore don't know why it can still be None
                     False,
                 ),
                 (
@@ -384,16 +417,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,
@@ -417,7 +455,7 @@ def solve(
     # get ready to return the parameters at last iteration...
     # (by default arbitrary choice, this could be None)
     validation_parameters = optimization.last_non_nan_params
-    if validation is not None:
+    if validation is not None and validation_crit_values is not None:
         jax.debug.print(
             "validation loss value = {validation_loss_val}",
             validation_loss_val=validation_crit_values[i - 1],
@@ -452,24 +490,28 @@ def _gradient_step(
     loss: AnyLoss,
     optimizer: optax.GradientTransformation,
     batch: AnyBatch,
-    params: AnyParams,
-    opt_state: NamedTuple,
-    last_non_nan_params: AnyParams,
+    params: Params[Array],
+    opt_state: optax.OptState,
+    last_non_nan_params: Params[Array],
 ) -> tuple[
     AnyLoss,
     float,
-    Dict[str, float],
-    AnyParams,
-    NamedTuple,
-    AnyParams,
+    dict[str, float],
+    Params[Array],
+    optax.OptState,
+    Params[Array],
 ]:
     """
     optimizer cannot be jit-ted.
     """
     value_grad_loss = jax.value_and_grad(loss, has_aux=True)
     (loss_val, loss_terms), grads = value_grad_loss(params, batch)
-    updates, opt_state = optimizer.update(grads, opt_state, params)
-    params = optax.apply_updates(params, updates)
+    updates, opt_state = optimizer.update(
+        grads,
+        opt_state,
+        params,  # type: ignore
+    )  # see optimizer.init for explaination
+    params = optax.apply_updates(params, updates)  # type: ignore
 
     # check if any of the parameters is NaN
     last_non_nan_params = jax.lax.cond(
@@ -490,7 +532,7 @@ def _gradient_step(
 
 
 @partial(jit, static_argnames=["prefix"])
-def _print_fn(i: Int, loss_val: Float, print_loss_every: Int, prefix: str = ""):
+def _print_fn(i: int, loss_val: Float, print_loss_every: int, prefix: str = ""):
     # note that if the following is not jitted in the main lor loop, it is
     # super slow
     _ = jax.lax.cond(
@@ -507,17 +549,15 @@ def _print_fn(i: Int, loss_val: Float, print_loss_every: Int, prefix: str = ""):
 
 @jit
 def _store_loss_and_params(
-    i: Int,
-    params: AnyParams,
-    stored_params: AnyParams,
-    stored_loss_terms: Dict[str, Float[Array, "n_iter"]],
-    train_loss_values: Float[Array, "n_iter"],
+    i: int,
+    params: Params[Array],
+    stored_params: Params[Array],
+    stored_loss_terms: dict[str, Float[Array, " n_iter"]],
+    train_loss_values: Float[Array, " n_iter"],
     train_loss_val: float,
-    loss_terms: Dict[str, float],
-    tracked_params: AnyParams,
-) -> tuple[
-    Params | ParamsDict, Dict[str, Float[Array, "n_iter"]], Float[Array, "n_iter"]
-]:
+    loss_terms: dict[str, float],
+    tracked_params: Params,
+) -> tuple[Params, dict[str, Float[Array, " n_iter"]], Float[Array, " n_iter"]]:
     stored_params = jax.tree_util.tree_map(
         lambda stored_value, param, tracked_param: (
             None
@@ -544,7 +584,7 @@ def _store_loss_and_params(
     return (stored_params, stored_loss_terms, train_loss_values)
 
 
-def _get_break_fun(n_iter: Int, verbose: Bool) -> Callable[[main_carry], Bool]:
+def _get_break_fun(n_iter: int, verbose: bool) -> Callable[[main_carry], bool]:
     """
     Wrapper to get the break_fun with appropriate `n_iter`.
     The verbose argument is here to control printing (or not) when exiting
@@ -585,7 +625,7 @@ def _get_break_fun(n_iter: Int, verbose: Bool) -> Callable[[main_carry], Bool]:
         bool_nan_in_params = jax.lax.cond(
             _check_nan_in_pytree(optimization.params),
             lambda _: stop_while_loop(
-                "NaN values in parameters " "(returning last non NaN values)"
+                "NaN values in parameters (returning last non NaN values)"
             ),
             continue_while_loop,
             None,
@@ -608,18 +648,18 @@ def _get_break_fun(n_iter: Int, verbose: Bool) -> Callable[[main_carry], Bool]:
 
 
 def _get_get_batch(
-    obs_batch_sharding: jax.sharding.Sharding,
+    obs_batch_sharding: jax.sharding.Sharding | None,
 ) -> Callable[
     [
-        AnyDataGenerator,
+        AbstractDataGenerator,
         DataGeneratorParameter | None,
-        DataGeneratorObservations | DataGeneratorObservationsMultiPINNs | None,
+        DataGeneratorObservations | None,
     ],
     tuple[
         AnyBatch,
-        AnyDataGenerator,
+        AbstractDataGenerator,
         DataGeneratorParameter | None,
-        DataGeneratorObservations | DataGeneratorObservationsMultiPINNs | None,
+        DataGeneratorObservations | None,
     ],
 ]:
     """

jinns/solver/_utils.py

@@ -1,9 +1,7 @@
-from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-    DataGeneratorParameter,
-)
+from jinns.data._DataGeneratorODE import DataGeneratorODE
+from jinns.data._CubicMeshPDEStatio import CubicMeshPDEStatio
+from jinns.data._CubicMeshPDENonStatio import CubicMeshPDENonStatio
+from jinns.data._DataGeneratorParameter import DataGeneratorParameter
 
 
 def _check_batch_size(other_data, main_data, attr_name):

jinns/utils/__init__.py

@@ -1,6 +1,3 @@
-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
+
+__all__ = ["get_grid"]

jinns/utils/_containers.py

@@ -11,23 +11,26 @@ from jaxtyping import PyTree, Array, Float, Bool
 from optax import OptState
 import equinox as eqx
 
+from jinns.parameters._params import Params
+
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+    from jinns.data._DataGeneratorParameter import DataGeneratorParameter
+    from jinns.data._DataGeneratorObservations import DataGeneratorObservations
+    from jinns.utils._types import AnyLoss
 
 
 class DataGeneratorContainer(eqx.Module):
-    data: AnyDataGenerator
+    data: AbstractDataGenerator
     param_data: DataGeneratorParameter | None = None
-    obs_data: DataGeneratorObservations | DataGeneratorObservationsMultiPINNs | None = (
-        None
-    )
+    obs_data: DataGeneratorObservations | None = None
 
 
 class ValidationContainer(eqx.Module):
     loss: AnyLoss | None
     data: DataGeneratorContainer
     hyperparams: PyTree = None
-    loss_values: Float[Array, "n_iter"] | None = None
+    loss_values: Float[Array, " n_iter"] | None = None
 
 
 class OptimizationContainer(eqx.Module):
@@ -45,9 +48,9 @@ class OptimizationExtraContainer(eqx.Module):
 
 
 class LossContainer(eqx.Module):
-    stored_loss_terms: Dict[str, Float[Array, "n_iter"]]
-    train_loss_values: Float[Array, "n_iter"]
+    stored_loss_terms: Dict[str, Float[Array, " n_iter"]]
+    train_loss_values: Float[Array, " n_iter"]
 
 
 class StoredObjectContainer(eqx.Module):
-    stored_params: list | None
+    stored_params: Params[Array | None]

jinns/utils/_types.py

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

jinns/utils/_utils.py

@@ -2,20 +2,13 @@
 Implements various utility functions
 """
 
-from math import prod
 import warnings
 import jax
 import jax.numpy as jnp
-from jaxtyping import PyTree, Array
+from jaxtyping import PyTree, Array, Bool
 
-from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-)
 
-
-def _check_nan_in_pytree(pytree: PyTree) -> bool:
+def _check_nan_in_pytree(pytree: PyTree) -> Bool[Array, " "]:
     """
     Check if there is a NaN value anywhere is the pytree
 
@@ -55,7 +48,7 @@ def get_grid(in_array: Array) -> Array:
 
 
 def _check_shape_and_type(
-    r: Array | int, expected_shape: tuple, cause: str = "", binop: str = ""
+    r: Array | int | float, expected_shape: tuple, cause: str = "", binop: str = ""
 ) -> Array | float:
     """
     Ensures float type and correct shapes for broadcasting when performing a
@@ -90,7 +83,7 @@ def _check_shape_and_type(
 
 
 def _subtract_with_check(
-    a: Array | int, b: Array | int, cause: str = ""
+    a: Array | int | float, b: Array, cause: str = ""
 ) -> Array | float:
     a = _check_shape_and_type(a, b.shape, cause=cause, binop="-")
     return a - b

jinns/validation/__init__.py

@@ -1 +1,3 @@
 from ._validation import AbstractValidationModule, ValidationLoss
+
+__all__ = ["AbstractValidationModule", "ValidationLoss"]

jinns/validation/_validation.py

@@ -7,19 +7,23 @@ from __future__ import (
 )  # https://docs.python.org/3/library/typing.html#constant
 
 import abc
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Array
+from jaxtyping import Array, Float
 
-from jinns.data._DataGenerators import (
+from jinns.data._utils import (
     append_obs_batch,
     append_param_batch,
 )
 
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.data._DataGeneratorParameter import DataGeneratorParameter
+    from jinns.data._DataGeneratorObservations import DataGeneratorObservations
+    from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+    from jinns.parameters._params import Params
+    from jinns.loss._abstract_loss import AbstractLoss
 
 # Using eqx Module for the DataClass + Pytree inheritance
 # Abstract class and abstract/final pattern is used
@@ -40,8 +44,8 @@ class AbstractValidationModule(eqx.Module):
 
     @abc.abstractmethod
     def __call__(
-        self, params: Params | ParamsDict
-    ) -> tuple["AbstractValidationModule", bool, Array, bool]:
+        self, params: Params[Array]
+    ) -> tuple[AbstractValidationModule, bool, Array, Params[Array]]:
         raise NotImplementedError
 
 
@@ -52,24 +56,20 @@ class ValidationLoss(AbstractValidationModule):
     for more complicated validation strategy.
     """
 
-    loss: AnyLoss = eqx.field(kw_only=True)  # NOTE that
-    # there used to be a deepcopy here which has been suppressed. 1) No need
-    # because loss are now eqx.Module (immutable) so no risk of in-place
-    # modification. 2) deepcopy is buggy with equinox, InitVar etc. (see issue
-    # #857 on equinox github)
-    validation_data: Union[AnyDataGenerator] = eqx.field(kw_only=True)
-    validation_param_data: Union[DataGeneratorParameter, None] = eqx.field(
+    loss: AbstractLoss = eqx.field(kw_only=True)
+    validation_data: AbstractDataGenerator = eqx.field(kw_only=True)
+    validation_param_data: DataGeneratorParameter = eqx.field(
+        kw_only=True, default=None
+    )
+    validation_obs_data: DataGeneratorObservations | None = eqx.field(
         kw_only=True, default=None
     )
-    validation_obs_data: Union[
-        DataGeneratorObservations, DataGeneratorObservationsMultiPINNs, None
-    ] = eqx.field(kw_only=True, default=None)
     call_every: int = eqx.field(kw_only=True, default=250)  # concrete typing
     early_stopping: bool = eqx.field(
         kw_only=True, default=True
     )  # globally control if early stopping happens
 
-    patience: Union[int] = eqx.field(kw_only=True, default=10)
+    patience: int = eqx.field(kw_only=True, default=10)
     best_val_loss: Array = eqx.field(
         converter=jnp.asarray, default_factory=lambda: jnp.array(jnp.inf), kw_only=True
     )
@@ -79,10 +79,11 @@ class ValidationLoss(AbstractValidationModule):
     )
 
     def __call__(
-        self, params: AnyParams
-    ) -> tuple["ValidationLoss", bool, float, AnyParams]:
+        self, params: Params[Array]
+    ) -> tuple[ValidationLoss, bool, Float[Array, " "], Params[Array]]:
         # do in-place mutation
 
+        # pylint / pyright complains below when using the self attributes see: https://github.com/patrick-kidger/equinox/issues/1013
         validation_data, val_batch = self.validation_data.get_batch()
         if self.validation_param_data is not None:
             validation_param_data, param_batch = self.validation_param_data.get_batch()

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: jinns
-Version: 1.2.0
+Version: 1.4.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>
@@ -25,6 +25,7 @@ Requires-Dist: matplotlib
 Provides-Extra: notebook
 Requires-Dist: jupyter; extra == "notebook"
 Requires-Dist: seaborn; extra == "notebook"
+Dynamic: license-file
 
 jinns
 =====
@@ -99,7 +100,7 @@ Here are the contributors guidelines:
 pip install -e .
 ```
 
-3. Install pre-commit and run it.
+3. Install pre-commit and run it. Our pre-commit hooks consist in `ruff format` and `ruff check`. You can install `ruff` simply by `pip install ruff`. We highly recommend you to check the code type hints with `pyright` even though we currently have no rule concerning type checking in the pipeline.
 
 ```bash
 pip install pre-commit
@@ -112,16 +113,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}
 }
 ```