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

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