jinns 0.8.6__py3-none-any.whl → 0.8.7__py3-none-any.whl

jinns/loss/_LossODE.py

@@ -19,6 +19,8 @@ from jinns.loss._Losses import (
 )
 from jinns.utils._pinn import PINN
 
+_LOSS_WEIGHT_KEYS_ODE = ["observations", "dyn_loss", "initial_condition"]
+
 
 @register_pytree_node_class
 class LossODE:
@@ -128,6 +130,10 @@ class LossODE:
         if self.obs_slice is None:
             self.obs_slice = jnp.s_[...]
 
+        for k in _LOSS_WEIGHT_KEYS_ODE:
+            if k not in self.loss_weights.keys():
+                self.loss_weights[k] = 0
+
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
 

jinns/loss/_LossPDE.py

@@ -31,6 +31,16 @@ _IMPLEMENTED_BOUNDARY_CONDITIONS = [
     "vonneumann",
 ]
 
+_LOSS_WEIGHT_KEYS_PDESTATIO = [
+    "sobolev",
+    "observations",
+    "norm_loss",
+    "boundary_loss",
+    "dyn_loss",
+]
+
+_LOSS_WEIGHT_KEYS_PDENONSTATIO = _LOSS_WEIGHT_KEYS_PDESTATIO + ["initial_condition"]
+
 
 @register_pytree_node_class
 class LossPDEAbstract:
@@ -269,8 +279,8 @@ class LossPDEStatio(LossPDEAbstract):
             the PINN object
         loss_weights
             a dictionary with values used to ponderate each term in the loss
-            function. Valid keys are `dyn_loss`, `norm_loss`, `boundary_loss`
-            and `observations`.
+            function. Valid keys are `dyn_loss`, `norm_loss`, `boundary_loss`,
+            `observations` and `sobolev`.
             Note that we can have jnp.arrays with the same dimension of
             `u` which then ponderates each output of `u`
         dynamic_loss
@@ -441,18 +451,12 @@ class LossPDEStatio(LossPDEAbstract):
             )  # we return a function, that way
             # the order of sobolev_m is static and the conditional in the recursive
             # function is properly set
-            self.sobolev_m = self.sobolev_m
         else:
             self.sobolev_reg = None
 
-        if self.normalization_loss is None:
-            self.loss_weights["norm_loss"] = 0
-
-        if self.omega_boundary_fun is None:
-            self.loss_weights["boundary_loss"] = 0
-
-        if self.sobolev_reg is None:
-            self.loss_weights["sobolev"] = 0
+        for k in _LOSS_WEIGHT_KEYS_PDESTATIO:
+            if k not in self.loss_weights.keys():
+                self.loss_weights[k] = 0
 
         if (
             isinstance(self.omega_boundary_fun, dict)
@@ -533,7 +537,6 @@ class LossPDEStatio(LossPDEAbstract):
             )
         else:
             mse_norm_loss = jnp.array(0.0)
-            self.loss_weights["norm_loss"] = 0
 
         # boundary part
         params_ = _set_derivatives(params, "boundary_loss", self.derivative_keys)
@@ -567,7 +570,6 @@ class LossPDEStatio(LossPDEAbstract):
             )
         else:
             mse_observation_loss = jnp.array(0.0)
-            self.loss_weights["observations"] = 0
 
         # Sobolev regularization
         params_ = _set_derivatives(params, "sobolev", self.derivative_keys)
@@ -582,7 +584,6 @@ class LossPDEStatio(LossPDEAbstract):
             )
         else:
             mse_sobolev_loss = jnp.array(0.0)
-            self.loss_weights["sobolev"] = 0
 
         # total loss
         total_loss = (
@@ -785,8 +786,9 @@ class LossPDENonStatio(LossPDEStatio):
         else:
             self.sobolev_reg = None
 
-        if self.sobolev_reg is None:
-            self.loss_weights["sobolev"] = 0
+        for k in _LOSS_WEIGHT_KEYS_PDENONSTATIO:
+            if k not in self.loss_weights.keys():
+                self.loss_weights[k] = 0
 
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
@@ -924,7 +926,6 @@ class LossPDENonStatio(LossPDEStatio):
             )
         else:
             mse_observation_loss = jnp.array(0.0)
-            self.loss_weights["observations"] = 0
 
         # Sobolev regularization
         params_ = _set_derivatives(params, "sobolev", self.derivative_keys)
@@ -939,7 +940,6 @@ class LossPDENonStatio(LossPDEStatio):
             )
         else:
             mse_sobolev_loss = jnp.array(0.0)
-            self.loss_weights["sobolev"] = 0.0
 
         # total loss
         total_loss = (

jinns/solver/_solve.py

@@ -3,9 +3,9 @@ This modules implements the main `solve()` function of jinns which
 handles the optimization process
 """
 
+import copy
+from functools import partial
 import optax
-from tqdm import tqdm
-from jax_tqdm import scan_tqdm
 import jax
 from jax import jit
 import jax.numpy as jnp
@@ -19,8 +19,32 @@ from jinns.data._DataGenerators import (
     append_param_batch,
     append_obs_batch,
 )
+from jinns.utils._containers import *
 
-from functools import partial
+
+def check_batch_size(other_data, main_data, attr_name):
+    if (
+        (
+            isinstance(main_data, DataGeneratorODE)
+            and getattr(other_data, attr_name) != main_data.temporal_batch_size
+        )
+        or (
+            isinstance(main_data, CubicMeshPDEStatio)
+            and not isinstance(main_data, CubicMeshPDENonStatio)
+            and getattr(other_data, attr_name) != main_data.omega_batch_size
+        )
+        or (
+            isinstance(main_data, CubicMeshPDENonStatio)
+            and getattr(other_data, attr_name)
+            != main_data.omega_batch_size * main_data.temporal_batch_size
+        )
+    ):
+        raise ValueError(
+            "Optional other_data.param_batch_size must be"
+            " equal to main_data.temporal_batch_size or main_data.omega_batch_size or"
+            " the product of both dependeing on the type of the main"
+            " datagenerator"
+        )
 
 
 def solve(
@@ -35,6 +59,7 @@ def solve(
     tracked_params_key_list=None,
     param_data=None,
     obs_data=None,
+    validation=None,
     obs_batch_sharding=None,
 ):
     """
@@ -88,6 +113,20 @@ def solve(
     obs_data
         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
+        validation strategy. See documentation of :obj:`~jinns.validation.
+        _validation.AbstractValidationModule` for the general interface, and
+        :obj:`~jinns.validation._validation.ValidationLoss` for a practical
+        implementation of a vanilla validation stategy on a validation set of
+        collocation points.
+
+        **Note**: The ``__call__(self, params)`` method should have
+        the latter prescribed signature and return ``(validation [eqx.Module],
+        early_stop [bool], validation_criterion [Array])``. It is called every
+        ``validation.call_every`` iteration. Users are free to design any
+        validation strategy of their choice, and to decide on the early
+        stopping criterion.
     obs_batch_sharding
         Default None. An optional sharding object to constraint the obs_batch.
         Typically, a SingleDeviceSharding(gpu_device) when obs_data has been
@@ -114,39 +153,14 @@ def solve(
         A dictionary. At each key an array of the values of the parameters
         given in tracked_params_key_list is stored
     """
-    params = init_params
-    last_non_nan_params = init_params.copy()
-
     if param_data is not None:
-        if (
-            (
-                isinstance(data, DataGeneratorODE)
-                and param_data.param_batch_size != data.temporal_batch_size
-                and obs_data.obs_batch_size != data.temporal_batch_size
-            )
-            or (
-                isinstance(data, CubicMeshPDEStatio)
-                and not isinstance(data, CubicMeshPDENonStatio)
-                and param_data.param_batch_size != data.omega_batch_size
-                and obs_data.obs_batch_size != data.omega_batch_size
-            )
-            or (
-                isinstance(data, CubicMeshPDENonStatio)
-                and param_data.param_batch_size
-                != data.omega_batch_size * data.temporal_batch_size
-                and obs_data.obs_batch_size
-                != data.omega_batch_size * data.temporal_batch_size
-            )
-        ):
-            raise ValueError(
-                "Optional param_data.param_batch_size must be"
-                " equal to data.temporal_batch_size or data.omega_batch_size or"
-                " the product of both dependeing on the type of the main"
-                " datagenerator"
-            )
+        check_batch_size(param_data, data, "param_batch_size")
+
+    if obs_data is not None:
+        check_batch_size(obs_data, data, "obs_batch_size")
 
     if opt_state is None:
-        opt_state = optimizer.init(params)
+        opt_state = optimizer.init(init_params)
 
     # 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
@@ -160,7 +174,7 @@ def solve(
         ), "data.method must be uniform if using seq2seq learning !"
         data, opt_state = initialize_seq2seq(loss, data, seq2seq, opt_state)
 
-    total_loss_values = jnp.zeros((n_iter))
+    train_loss_values = jnp.zeros((n_iter))
     # depending on obs_batch_sharding we will get the simple get_batch or the
     # get_batch with device_put, the latter is not jittable
     get_batch = get_get_batch(obs_batch_sharding)
@@ -168,68 +182,125 @@ def solve(
     # initialize the dict for stored parameter values
     # we need to get a loss_term to init stuff
     batch_ini, data, param_data, obs_data = get_batch(data, param_data, obs_data)
-    _, loss_terms = loss(params, batch_ini)
+    _, loss_terms = loss(init_params, batch_ini)
     if tracked_params_key_list is None:
         tracked_params_key_list = []
-    tracked_params = _tracked_parameters(params, tracked_params_key_list)
+    tracked_params = _tracked_parameters(init_params, tracked_params_key_list)
     stored_params = jax.tree_util.tree_map(
         lambda tracked_param, param: (
             jnp.zeros((n_iter,) + param.shape) if tracked_param else None
         ),
         tracked_params,
-        params,
+        init_params,
     )
 
     # initialize the dict for stored loss values
     stored_loss_terms = jax.tree_util.tree_map(
-        lambda x: jnp.zeros((n_iter)), loss_terms
+        lambda _: jnp.zeros((n_iter)), loss_terms
+    )
+
+    train_data = DataGeneratorContainer(
+        data=data, param_data=param_data, obs_data=obs_data
+    )
+    optimization = OptimizationContainer(
+        params=init_params, last_non_nan_params=init_params.copy(), opt_state=opt_state
+    )
+    optimization_extra = OptimizationExtraContainer(
+        curr_seq=curr_seq,
+        seq2seq=seq2seq,
     )
+    loss_container = LossContainer(
+        stored_loss_terms=stored_loss_terms,
+        train_loss_values=train_loss_values,
+    )
+    stored_objects = StoredObjectContainer(
+        stored_params=stored_params,
+    )
+
+    if validation is not None:
+        validation_crit_values = jnp.zeros(n_iter)
+    else:
+        validation_crit_values = None
+
+    break_fun = get_break_fun(n_iter)
 
+    iteration = 0
     carry = (
-        init_params,
-        init_params.copy(),
-        data,
-        curr_seq,
-        seq2seq,
-        stored_params,
-        stored_loss_terms,
+        iteration,
         loss,
-        param_data,
-        obs_data,
-        opt_state,
-        total_loss_values,
+        optimization,
+        optimization_extra,
+        train_data,
+        validation,
+        loss_container,
+        stored_objects,
+        validation_crit_values,
     )
 
-    def one_iteration(carry, i):
+    def one_iteration(carry):
         (
-            params,
-            last_non_nan_params,
-            data,
-            curr_seq,
-            seq2seq,
-            stored_params,
-            stored_loss_terms,
+            i,
             loss,
-            param_data,
-            obs_data,
-            opt_state,
-            total_loss_values,
+            optimization,
+            optimization_extra,
+            train_data,
+            validation,
+            loss_container,
+            stored_objects,
+            validation_crit_values,
         ) = carry
-        batch, data, param_data, obs_data = get_batch(data, param_data, obs_data)
 
+        batch, data, param_data, obs_data = get_batch(
+            train_data.data, train_data.param_data, train_data.obs_data
+        )
+
+        # Gradient step
         (
             loss,
-            loss_val,
+            train_loss_value,
             loss_terms,
             params,
             opt_state,
             last_non_nan_params,
         ) = gradient_step(
-            loss, optimizer, batch, params, opt_state, last_non_nan_params
+            loss,
+            optimizer,
+            batch,
+            optimization.params,
+            optimization.opt_state,
+            optimization.last_non_nan_params,
         )
 
-        # Print loss during optimization
-        print_fn(i, loss_val, print_loss_every)
+        # Print train loss value during optimization
+        print_fn(i, train_loss_value, print_loss_every, prefix="[train] ")
+
+        if validation is not None:
+            # there is a jax.lax.cond because we do not necesarily call the
+            # validation step every iteration
+            (
+                validation,  # always return `validation` for in-place mutation
+                early_stopping,
+                validation_criterion,
+            ) = jax.lax.cond(
+                i % validation.call_every == 0,
+                lambda operands: operands[0](*operands[1:]),  # validation.__call__()
+                lambda operands: (
+                    operands[0],
+                    False,
+                    validation_crit_values[i - 1],
+                ),
+                (
+                    validation,  # validation must be in operands
+                    params,
+                ),
+            )
+            # Print validation loss value during optimization
+            print_fn(i, validation_criterion, print_loss_every, prefix="[validation] ")
+            validation_crit_values = validation_crit_values.at[i].set(
+                validation_criterion
+            )
+        else:
+            early_stopping = False
 
         # Trigger RAR
         loss, params, data = trigger_rar(
@@ -238,84 +309,98 @@ def solve(
 
         # Trigger seq2seq
         loss, params, data, opt_state, curr_seq, seq2seq = trigger_seq2seq(
-            i, loss, params, data, opt_state, curr_seq, seq2seq
+            i,
+            loss,
+            params,
+            data,
+            opt_state,
+            optimization_extra.curr_seq,
+            optimization_extra.seq2seq,
         )
 
         # save loss value and selected parameters
-        stored_params, stored_loss_terms, total_loss_values = store_loss_and_params(
+        stored_params, stored_loss_terms, train_loss_values = store_loss_and_params(
             i,
             params,
-            stored_params,
-            stored_loss_terms,
-            total_loss_values,
-            loss_val,
+            stored_objects.stored_params,
+            loss_container.stored_loss_terms,
+            loss_container.train_loss_values,
+            train_loss_value,
             loss_terms,
             tracked_params,
         )
+        i += 1
+
         return (
-            params,
-            last_non_nan_params,
-            data,
-            curr_seq,
-            seq2seq,
-            stored_params,
-            stored_loss_terms,
+            i,
             loss,
-            param_data,
-            obs_data,
-            opt_state,
-            total_loss_values,
-        ), None
+            OptimizationContainer(params, last_non_nan_params, opt_state),
+            OptimizationExtraContainer(curr_seq, seq2seq, early_stopping),
+            DataGeneratorContainer(data, param_data, obs_data),
+            validation,
+            LossContainer(stored_loss_terms, train_loss_values),
+            StoredObjectContainer(stored_params),
+            validation_crit_values,
+        )
 
-    # Main optimization loop. We use the fully scanned (fully jitted) version
-    # if no mixing devices. Otherwise we use the for loop. Here devices only
+    # Main optimization loop. We use the LAX while loop (fully jitted) version
+    # if no mixing devices. Otherwise we use the standard while loop. Here devices only
     # concern obs_batch, but it could lead to more complex scheme in the future
     if obs_batch_sharding is not None:
-        for i in tqdm(range(n_iter)):
-            carry, _ = one_iteration(carry, i)
+        while break_fun(carry):
+            carry = one_iteration(carry)
     else:
-        carry, _ = jax.lax.scan(
-            scan_tqdm(n_iter)(one_iteration),
-            carry,
-            jnp.arange(n_iter),
-        )
+        carry = jax.lax.while_loop(break_fun, one_iteration, carry)
 
     (
-        init_params,
-        last_non_nan_params,
-        data,
-        curr_seq,
-        seq2seq,
-        stored_params,
-        stored_loss_terms,
+        i,
         loss,
-        param_data,
-        obs_data,
-        opt_state,
-        total_loss_values,
+        optimization,
+        optimization_extra,
+        train_data,
+        validation,
+        loss_container,
+        stored_objects,
+        validation_crit_values,
     ) = carry
 
     jax.debug.print(
-        "Iteration {i}: loss value = {total_loss_val}",
-        i=n_iter,
-        total_loss_val=total_loss_values[-1],
+        "Final iteration {i}: train loss value = {train_loss_val}",
+        i=i,
+        train_loss_val=loss_container.train_loss_values[i - 1],
     )
+    if validation is not None:
+        jax.debug.print(
+            "validation loss value = {validation_loss_val}",
+            validation_loss_val=validation_crit_values[i - 1],
+        )
 
+    if validation is None:
+        return (
+            optimization.last_non_nan_params,
+            loss_container.train_loss_values,
+            loss_container.stored_loss_terms,
+            train_data.data,
+            loss,
+            optimization.opt_state,
+            stored_objects.stored_params,
+        )
     return (
-        last_non_nan_params,
-        total_loss_values,
-        stored_loss_terms,
-        data,
+        optimization.last_non_nan_params,
+        loss_container.train_loss_values,
+        loss_container.stored_loss_terms,
+        train_data.data,
         loss,
-        opt_state,
-        stored_params,
+        optimization.opt_state,
+        stored_objects.stored_params,
+        validation_crit_values,
     )
 
 
 @partial(jit, static_argnames=["optimizer"])
 def gradient_step(loss, optimizer, batch, params, opt_state, last_non_nan_params):
     """
-    loss and optimizer cannot be jit-ted.
+    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)
@@ -340,14 +425,14 @@ def gradient_step(loss, optimizer, batch, params, opt_state, last_non_nan_params
     )
 
 
-@jit
-def print_fn(i, loss_val, print_loss_every):
+@partial(jit, static_argnames=["prefix"])
+def print_fn(i, loss_val, print_loss_every, prefix=""):
     # note that if the following is not jitted in the main lor loop, it is
     # super slow
     _ = jax.lax.cond(
         i % print_loss_every == 0,
         lambda _: jax.debug.print(
-            "Iteration {i}: loss value = {loss_val}",
+            prefix + "Iteration {i}: loss value = {loss_val}",
             i=i,
             loss_val=loss_val,
         ),
@@ -362,8 +447,8 @@ def store_loss_and_params(
     params,
     stored_params,
     stored_loss_terms,
-    total_loss_values,
-    loss_val,
+    train_loss_values,
+    train_loss_val,
     loss_terms,
     tracked_params,
 ):
@@ -384,8 +469,66 @@ def store_loss_and_params(
         loss_terms,
     )
 
-    total_loss_values = total_loss_values.at[i].set(loss_val)
-    return stored_params, stored_loss_terms, total_loss_values
+    train_loss_values = train_loss_values.at[i].set(train_loss_val)
+    return (stored_params, stored_loss_terms, train_loss_values)
+
+
+def get_break_fun(n_iter):
+    """
+    Wrapper to get the break_fun with appropriate `n_iter`
+    """
+
+    @jit
+    def break_fun(carry):
+        """
+        Function to break from the main optimization loop
+        We check several conditions
+        """
+
+        def stop_while_loop(msg):
+            """
+            Note that the message is wrapped in the jax.lax.cond because a
+            string is not a valid JAX type that can be fed into the operands
+            """
+            jax.debug.print(f"Stopping main optimization loop, cause: {msg}")
+            return False
+
+        def continue_while_loop(_):
+            return True
+
+        (i, _, optimization, optimization_extra, _, _, _, _, _) = carry
+
+        # Condition 1
+        bool_max_iter = jax.lax.cond(
+            i >= n_iter,
+            lambda _: stop_while_loop("max iteration is reached"),
+            continue_while_loop,
+            None,
+        )
+        # Condition 2
+        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)"
+            ),
+            continue_while_loop,
+            None,
+        )
+        # Condition 3
+        bool_early_stopping = jax.lax.cond(
+            optimization_extra.early_stopping,
+            lambda _: stop_while_loop("early stopping"),
+            continue_while_loop,
+            _,
+        )
+
+        # stop when one of the cond to continue is False
+        return jax.tree_util.tree_reduce(
+            lambda x, y: jnp.logical_and(jnp.array(x), jnp.array(y)),
+            (bool_max_iter, bool_nan_in_params, bool_early_stopping),
+        )
+
+    return break_fun
 
 
 def get_get_batch(obs_batch_sharding):

jinns/utils/_containers.py

@@ -0,0 +1,57 @@
+"""
+NamedTuples definition
+"""
+
+from typing import Union, NamedTuple
+from jaxtyping import PyTree
+from jax.typing import ArrayLike
+import optax
+import jax.numpy as jnp
+from jinns.loss._LossODE import LossODE, SystemLossODE
+from jinns.loss._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
+from jinns.data._DataGenerators import (
+    DataGeneratorODE,
+    CubicMeshPDEStatio,
+    CubicMeshPDENonStatio,
+    DataGeneratorParameter,
+    DataGeneratorObservations,
+    DataGeneratorObservationsMultiPINNs,
+)
+
+
+class DataGeneratorContainer(NamedTuple):
+    data: Union[DataGeneratorODE, CubicMeshPDEStatio, CubicMeshPDENonStatio]
+    param_data: Union[DataGeneratorParameter, None] = None
+    obs_data: Union[
+        DataGeneratorObservations, DataGeneratorObservationsMultiPINNs, None
+    ] = None
+
+
+class ValidationContainer(NamedTuple):
+    loss: Union[
+        LossODE, SystemLossODE, LossPDEStatio, LossPDENonStatio, SystemLossPDE, None
+    ]
+    data: DataGeneratorContainer
+    hyperparams: PyTree = None
+    loss_values: Union[ArrayLike, None] = None
+
+
+class OptimizationContainer(NamedTuple):
+    params: dict
+    last_non_nan_params: dict
+    opt_state: optax.OptState
+
+
+class OptimizationExtraContainer(NamedTuple):
+    curr_seq: int
+    seq2seq: Union[dict, None]
+    early_stopping: bool = False
+
+
+class LossContainer(NamedTuple):
+    stored_loss_terms: dict
+    train_loss_values: ArrayLike
+
+
+class StoredObjectContainer(NamedTuple):
+    stored_params: Union[list, None]

jinns/validation/__init__.py

@@ -0,0 +1 @@
+from ._validation import AbstractValidationModule, ValidationLoss

jinns/validation/_validation.py

@@ -0,0 +1,214 @@
+"""
+Implements some validation functions and their associated hyperparameter
+"""
+
+import copy
+import abc
+from typing import Union
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, Bool, PyTree, Int
+import jinns
+import jinns.data
+from jinns.loss import LossODE, LossPDENonStatio, LossPDEStatio
+from jinns.data._DataGenerators import (
+    DataGeneratorODE,
+    CubicMeshPDEStatio,
+    CubicMeshPDENonStatio,
+    DataGeneratorParameter,
+    DataGeneratorObservations,
+    DataGeneratorObservationsMultiPINNs,
+    append_obs_batch,
+    append_param_batch,
+)
+import jinns.loss
+
+# Using eqx Module for the DataClass + Pytree inheritance
+# Abstract class and abstract/final pattern is used
+# see : https://docs.kidger.site/equinox/pattern/
+
+
+class AbstractValidationModule(eqx.Module):
+    """Abstract class representing interface for any validation module. It must
+    1. have a ``call_every`` attribute.
+    2. implement a ``__call__`` returning ``(AbstractValidationModule, Bool, Array)``
+    """
+
+    call_every: eqx.AbstractVar[Int]  # Mandatory for all validation step,
+    # it tells that the validation step is performed every call_every
+    # iterations.
+
+    @abc.abstractmethod
+    def __call__(
+        self, params: PyTree
+    ) -> tuple["AbstractValidationModule", Bool, Array]:
+        raise NotImplementedError
+
+
+class ValidationLoss(AbstractValidationModule):
+    """
+    Implementation of a vanilla validation module returning the PINN loss
+    on a validation set of collocation points. This can be used as a baseline
+    for more complicated validation strategy.
+    """
+
+    loss: Union[callable, LossODE, LossPDEStatio, LossPDENonStatio] = eqx.field(
+        converter=copy.deepcopy
+    )
+    validation_data: Union[DataGeneratorODE, CubicMeshPDEStatio, CubicMeshPDENonStatio]
+    validation_param_data: Union[DataGeneratorParameter, None] = None
+    validation_obs_data: Union[
+        DataGeneratorObservations, DataGeneratorObservationsMultiPINNs, None
+    ] = None
+    call_every: Int = 250  # concrete typing
+    early_stopping: Bool = True  # globally control if early stopping happens
+
+    patience: Union[Int] = 10
+    best_val_loss: Array = eqx.field(
+        converter=jnp.asarray, default_factory=lambda: jnp.array(jnp.inf)
+    )
+
+    counter: Array = eqx.field(
+        converter=jnp.asarray, default_factory=lambda: jnp.array(0.0)
+    )
+
+    def __call__(self, params) -> tuple["ValidationLoss", Bool, Array]:
+        # do in-place mutation
+        val_batch = self.validation_data.get_batch()
+        if self.validation_param_data is not None:
+            val_batch = append_param_batch(
+                val_batch, self.validation_param_data.get_batch()
+            )
+        if self.validation_obs_data is not None:
+            val_batch = append_obs_batch(
+                val_batch, self.validation_obs_data.get_batch()
+            )
+
+        validation_loss_value, _ = self.loss(params, val_batch)
+        (counter, best_val_loss) = jax.lax.cond(
+            validation_loss_value < self.best_val_loss,
+            lambda _: (jnp.array(0.0), validation_loss_value),  # reset
+            lambda operands: (operands[0] + 1, operands[1]),  # increment
+            (self.counter, self.best_val_loss),
+        )
+
+        # use eqx.tree_at to update attributes
+        # (https://github.com/patrick-kidger/equinox/issues/396)
+        new = eqx.tree_at(lambda t: t.counter, self, counter)
+        new = eqx.tree_at(lambda t: t.best_val_loss, new, best_val_loss)
+
+        bool_early_stopping = jax.lax.cond(
+            jnp.logical_and(
+                jnp.array(self.counter == self.patience),
+                jnp.array(self.early_stopping),
+            ),
+            lambda _: True,
+            lambda _: False,
+            None,
+        )
+        # return `new` cause no in-place modification of the eqx.Module
+        return (new, bool_early_stopping, validation_loss_value)
+
+
+if __name__ == "__main__":
+    import jax
+    import jax.numpy as jnp
+    import jax.random as random
+    from jinns.loss import BurgerEquation
+
+    key = random.PRNGKey(1)
+    key, subkey = random.split(key)
+
+    n = 50
+    nb = 2 * 2 * 10
+    nt = 10
+    omega_batch_size = 10
+    omega_border_batch_size = 10
+    temporal_batch_size = 4
+    dim = 1
+    xmin = 0
+    xmax = 1
+    tmin, tmax = 0, 1
+    method = "uniform"
+
+    val_data = jinns.data.CubicMeshPDENonStatio(
+        subkey,
+        n,
+        nb,
+        nt,
+        omega_batch_size,
+        omega_border_batch_size,
+        temporal_batch_size,
+        dim,
+        (xmin,),
+        (xmax,),
+        tmin,
+        tmax,
+        method,
+    )
+
+    eqx_list = [
+        [eqx.nn.Linear, 2, 50],
+        [jax.nn.tanh],
+        [eqx.nn.Linear, 50, 50],
+        [jax.nn.tanh],
+        [eqx.nn.Linear, 50, 50],
+        [jax.nn.tanh],
+        [eqx.nn.Linear, 50, 50],
+        [jax.nn.tanh],
+        [eqx.nn.Linear, 50, 50],
+        [jax.nn.tanh],
+        [eqx.nn.Linear, 50, 2],
+    ]
+
+    key, subkey = random.split(key)
+    u = jinns.utils.create_PINN(
+        subkey, eqx_list, "nonstatio_PDE", 2, slice_solution=jnp.s_[:1]
+    )
+    init_nn_params = u.init_params()
+
+    dyn_loss = BurgerEquation()
+    loss_weights = {"dyn_loss": 1, "boundary_loss": 10, "observations": 10}
+
+    key, subkey = random.split(key)
+    loss = jinns.loss.LossPDENonStatio(
+        u=u,
+        loss_weights=loss_weights,
+        dynamic_loss=dyn_loss,
+        norm_key=subkey,
+        norm_borders=(-1, 1),
+    )
+    print(id(loss))
+    validation = ValidationLoss(
+        call_every=250,
+        early_stopping=True,
+        patience=1000,
+        loss=loss,
+        validation_data=val_data,
+        validation_param_data=None,
+    )
+    print(id(validation.loss) is not id(loss))  # should be True (deepcopy)
+
+    init_params = {"nn_params": init_nn_params, "eq_params": {"nu": 1.0}}
+
+    print(validation.loss is loss)
+    loss.evaluate(init_params, val_data.get_batch())
+    print(loss.norm_key)
+    print("Call validation once")
+    validation, _, _ = validation(init_params)
+    print(validation.loss is loss)
+    print(validation.loss.norm_key == loss.norm_key)
+    print("Crate new pytree from validation and call it once")
+    new_val = eqx.tree_at(lambda t: t.counter, validation, jnp.array(3.0))
+    print(validation.loss is new_val.loss)  # FALSE
+    # test if attribute have been modified
+    new_val, _, _ = new_val(init_params)
+    print(f"{new_val.loss is loss=}")
+    print(f"{loss.norm_key=}")
+    print(f"{validation.loss.norm_key=}")
+    print(f"{new_val.loss.norm_key=}")
+    print(f"{new_val.loss.norm_key == loss.norm_key=}")
+    print(f"{new_val.loss.norm_key == validation.loss.norm_key=}")
+    print(new_val.counter)
+    print(validation.counter)

{jinns-0.8.6.dist-info → jinns-0.8.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.6
+Version: 0.8.7
 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.6.dist-info → jinns-0.8.7.dist-info}/RECORD

@@ -6,8 +6,8 @@ jinns/experimental/__init__.py,sha256=3jCIy2R2i_0Erwxg-HwISdH79Nt1XCXhS9yY1F5awi
 jinns/experimental/_diffrax_solver.py,sha256=sLT22byqh-6015_fhe1xtMWlFOYcCjzYKET4sLhA9R4,6818
 jinns/loss/_DynamicLoss.py,sha256=L4CVmmF0rTPbHntgqsLLHlnrlQgLHsetUocpJm7ZYag,27461
 jinns/loss/_DynamicLossAbstract.py,sha256=kTQlhLx7SBuH5dIDmYaE79sVHUZt1nUFa8LxPU5IHhM,8504
-jinns/loss/_LossODE.py,sha256=sxpgiDR6mfoREuc-qe0AkirOe5K_5oblaYCnodTNxoI,21912
-jinns/loss/_LossPDE.py,sha256=_yX3R-FrAScTn9_QfVC8PfDYRE4UQ5lnzITUYgNFitA,61766
+jinns/loss/_LossODE.py,sha256=b9doBHoQwYvlgpqzrNO4dOaTN87LRvjHtHbz9bMoH7E,22119
+jinns/loss/_LossPDE.py,sha256=purAEtc0e71kv9XnZUT-a7MrkDAkM_3tTI4xJPu6fH4,61629
 jinns/loss/_Losses.py,sha256=XOL3MFiKEd3ndsc78Qnpi1vbgR0B2HaAWOGGW2meDM8,11190
 jinns/loss/__init__.py,sha256=pFNYUxns-NPXBFdqrEVSiXkQLfCtKw-t2trlhvLzpYE,355
 jinns/loss/_boundary_conditions.py,sha256=YfSnLZ25hXqQ5KWAuxOrWSKkf_oBqAc9GQV4z7MjWyQ,17434
@@ -15,8 +15,9 @@ jinns/loss/_operators.py,sha256=zDGJqYqeYH7xd-4dtGX9PS-pf0uSOpUUXGo5SVjIJ4o,1106
 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=r4jn6hx7_t-Y2rBWA2npUmWWnDg4iRbgYBHZDNn9tmY,13745
+jinns/solver/_solve.py,sha256=mGi0zaT_fK_QpBjTxof5Ix4mmfmnPi66CNJ3GQFZuo4,19099
 jinns/utils/__init__.py,sha256=44ms5UR6vMw3Nf6u4RCAzPFs4fom_YbBnH9mfne8m6k,313
+jinns/utils/_containers.py,sha256=eYD277fO7X4EfX7PUFCCl69r3JBfh1sCfq8LkL5gd6o,1495
 jinns/utils/_hyperpinn.py,sha256=93hbiATdp5W4l1cu9Oe6O2c45o-ZF_z2u6FzNLyjnm4,10878
 jinns/utils/_optim.py,sha256=550kxH75TL30o1iKx1swJyP0KqyUPsJ7-imL1w65Qd0,4444
 jinns/utils/_pinn.py,sha256=mhA4-3PazyQTbWIx9oLaNwL0QDe8ZIBhbiy5J3kwa4I,9471
@@ -24,8 +25,10 @@ jinns/utils/_save_load.py,sha256=qgZ23nUcB8-B5IZ2guuUWC4M7r5Lxd_Ms3staScdyJo,566
 jinns/utils/_spinn.py,sha256=SzOUt1KHtB9QOpghpvitnXN-KEqXUXbvabC5k0TnKEo,7793
 jinns/utils/_utils.py,sha256=8dgvWXX9NT7_7-zltWp0C9tG45ZFNwXxueyxPBb4hjo,6740
 jinns/utils/_utils_uspinn.py,sha256=qcKcOw3zrwWSQyGVj6fD8c9GinHt_U6JWN_k0auTtXM,26039
-jinns-0.8.6.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.8.6.dist-info/METADATA,sha256=3Ml6PCA-569v9-1FgyPDySX09RQas0zPOVEV_gqy9lk,2482
-jinns-0.8.6.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jinns-0.8.6.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.8.6.dist-info/RECORD,,
+jinns/validation/__init__.py,sha256=Jv58mzgC3F7cRfXA6caicL1t_U0UAhbwLrmMNVg6E7s,66
+jinns/validation/_validation.py,sha256=KfetbzB0xTNdBcYLwFWjEtP63Tf9wJirlhgqLTJDyy4,6761
+jinns-0.8.7.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.8.7.dist-info/METADATA,sha256=L0P7JvMGKrJHx9OjrtFsmNKEwdKA_RlufAbOBf5l10I,2482
+jinns-0.8.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jinns-0.8.7.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.8.7.dist-info/RECORD,,