jinns 0.8.10__py3-none-any.whl → 0.9.0__py3-none-any.whl

jinns/data/_DataGenerators.py

@@ -17,9 +17,8 @@ class ODEBatch(NamedTuple):
 
 
 class PDENonStatioBatch(NamedTuple):
-    inside_batch: ArrayLike
-    border_batch: ArrayLike
-    temporal_batch: ArrayLike
+    times_x_inside_batch: ArrayLike
+    times_x_border_batch: ArrayLike
     param_batch_dict: dict = None
     obs_batch_dict: dict = None
 
@@ -47,6 +46,18 @@ def append_obs_batch(batch, obs_batch_dict):
     return batch._replace(obs_batch_dict=obs_batch_dict)
 
 
+def make_cartesian_product(b1, b2):
+    """
+    Create the cartesian product of a time and a border omega batches
+    by tiling and repeating
+    """
+    n1 = b1.shape[0]
+    n2 = b2.shape[0]
+    b1 = jnp.repeat(b1, n2, axis=0)
+    b2 = jnp.tile(b2, reps=(n1,) + tuple(1 for i in b2.shape[1:]))
+    return jnp.concatenate([b1, b2], axis=1)
+
+
 def _reset_batch_idx_and_permute(operands):
     key, domain, curr_idx, _, p = operands
     # resetting counter
@@ -476,10 +487,9 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             # always set to 2.
             self.nb = 2
             self.omega_border_batch_size = 2
-        # warnings.warn("We are in 1-D case => omega_border_batch_size is "
-        #               "ignored since borders of Omega are singletons."
-        #               " self.border_batch() will return [xmin, xmax]"
-        #               )
+            # We are in 1-D case => omega_border_batch_size is
+            # ignored since borders of Omega are singletons.
+            #  self.border_batch() will return [xmin, xmax]
         else:
             if nb % (2 * self.dim) != 0 or nb < 2 * self.dim:
                 raise ValueError(
@@ -829,6 +839,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         rar_parameters=None,
         n_start=None,
         nt_start=None,
+        cartesian_product=True,
         data_exists=False,
     ):
         r"""
@@ -899,6 +910,10 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             Defaults to None. A RAR hyper-parameter. Same as ``n_start`` but
             for times collocation point. See also ``DataGeneratorODE``
             documentation.
+        cartesian_product
+            Defaults to True. Whether we return the cartesian product of the
+            temporal batch with the inside and border batches. If False we just
+            return their concatenation.
         data_exists
             Must be left to `False` when created by the user. Avoids the
             regeneration of :math:`\Omega`, :math:`\partial\Omega` and
@@ -923,6 +938,30 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         self.tmax = tmax
         self.nt = nt
 
+        self.cartesian_product = cartesian_product
+        if not self.cartesian_product:
+            if self.temporal_batch_size != self.omega_batch_size:
+                raise ValueError(
+                    "If stacking is requested between the time and "
+                    "inside batches of collocation points, self.temporal_batch_size "
+                    "must then be equal to self.omega_batch_size"
+                )
+            if (
+                self.dim > 1
+                and self.omega_border_batch_size is not None
+                and self.temporal_batch_size != self.omega_border_batch_size
+            ):
+                raise ValueError(
+                    "If dim > 1 and stacking is requested between the time and "
+                    "inside batches of collocation points, self.temporal_batch_size "
+                    "must then be equal to self.omega_border_batch_size"
+                )
+            # Note if self.dim == 1:
+            #    print(
+            #        "Cartesian product is not requested but will be "
+            #        "executed anyway since dim=1"
+            #    )
+
         # Set-up for timewise RAR (some quantity are already set-up by super())
         (
             self.nt_start,
@@ -1003,11 +1042,26 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         Generic method to return a batch. Here we call `self.inside_batch()`,
         `self.border_batch()` and `self.temporal_batch()`
         """
-        return PDENonStatioBatch(
-            inside_batch=self.inside_batch(),
-            border_batch=self.border_batch(),
-            temporal_batch=self.temporal_batch(),
-        )
+        x = self.inside_batch()
+        dx = self.border_batch()
+        t = self.temporal_batch().reshape(self.temporal_batch_size, 1)
+
+        if self.cartesian_product:
+            t_x = make_cartesian_product(t, x)
+        else:
+            t_x = jnp.concatenate([t, x], axis=1)
+
+        if dx is not None:
+            t_ = t.reshape(self.temporal_batch_size, 1, 1)
+            t_ = jnp.repeat(t_, dx.shape[-1], axis=2)
+            if self.cartesian_product or self.dim == 1:
+                t_dx = make_cartesian_product(t_, dx)
+            else:
+                t_dx = jnp.concatenate([t_, dx], axis=1)
+        else:
+            t_dx = None
+
+        return PDENonStatioBatch(times_x_inside_batch=t_x, times_x_border_batch=t_dx)
 
     def tree_flatten(self):
         children = (
@@ -1041,6 +1095,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
                 "rar_parameters",
                 "n_start",
                 "nt_start",
+                "cartesian_product",
             ]
         }
         return (children, aux_data)
@@ -1121,12 +1176,17 @@ class DataGeneratorParameter:
             once during 1 epoch.
         param_batch_size
             An integer. The size of the batch of randomly selected points among
-            the `n` points.  `param_batch_size` will be the same for all the
-            additional batch(es) of parameter(s). `param_batch_size` must be
-            equal to `temporal_batch_size` or `omega_batch_size` or the product
-            of both whether the present DataGeneratorParameter instance
-            complements and ODEBatch, a PDEStatioBatch or a PDENonStatioBatch,
-            respectively.
+            the `n` points. `param_batch_size` will be the same for all
+            additional batch of parameter.
+            NOTE: no check is done BUT users should be careful that
+            `param_batch_size` must be equal to `temporal_batch_size` or
+            `omega_batch_size` or the product of both. In the first case, the
+            present DataGeneratorParameter instance complements an ODEBatch, a
+            PDEStatioBatch or a PDENonStatioBatch (with self.cartesian_product
+            = False). In the second case, `param_batch_size` =
+            `temporal_batch_size * omega_batch_size` if the present
+            DataGeneratorParameter complements a PDENonStatioBatch
+            with self.cartesian_product = True
         param_ranges
             A dict. A dict of tuples (min, max), which
             reprensents the range of real numbers where to sample batches (of
@@ -1335,13 +1395,18 @@ class DataGeneratorObservations:
         key
             Jax random key to sample new time points and to shuffle batches
         obs_batch_size
-            An integer. The size of the batch of randomly selected observations
-            `obs_batch_size` will be the same for all the
-            elements of the obs dict. `obs_batch_size` must be
-            equal to `temporal_batch_size` or `omega_batch_size` or the product
-            of both whether the present DataGeneratorParameter instance
-            complements and ODEBatch, a PDEStatioBatch or a PDENonStatioBatch,
-            respectively.
+            An integer. The size of the batch of randomly selected points among
+            the `n` points. `obs_batch_size` will be the same for all
+            elements of the obs dict.
+            NOTE: no check is done BUT users should be careful that
+            `obs_batch_size` must be equal to `temporal_batch_size` or
+            `omega_batch_size` or the product of both. In the first case, the
+            present DataGeneratorObservations instance complements an ODEBatch,
+            PDEStatioBatch or a PDENonStatioBatch (with self.cartesian_product
+            = False). In the second case, `obs_batch_size` =
+            `temporal_batch_size * omega_batch_size` if the present
+            DataGeneratorParameter complements a PDENonStatioBatch
+            with self.cartesian_product = True
         observed_pinn_in
             A jnp.array with 2 dimensions.
             Observed values corresponding to the input of the PINN
@@ -1549,11 +1614,16 @@ class DataGeneratorObservationsMultiPINNs:
         obs_batch_size
             An integer. The size of the batch of randomly selected observations
             `obs_batch_size` will be the same for all the
-            elements of the obs dict. `obs_batch_size` must be
-            equal to `temporal_batch_size` or `omega_batch_size` or the product
-            of both whether the present DataGeneratorParameter instance
-            complements and ODEBatch, a PDEStatioBatch or a PDENonStatioBatch,
-            respectively.
+            elements of the obs dict.
+            NOTE: no check is done BUT users should be careful that
+            `obs_batch_size` must be equal to `temporal_batch_size` or
+            `omega_batch_size` or the product of both. In the first case, the
+            present DataGeneratorObservations instance complements an ODEBatch,
+            PDEStatioBatch or a PDENonStatioBatch (with self.cartesian_product
+            = False). In the second case, `obs_batch_size` =
+            `temporal_batch_size * omega_batch_size` if the present
+            DataGeneratorParameter complements a PDENonStatioBatch
+            with self.cartesian_product = True
         observed_pinn_in_dict
             A dict of observed_pinn_in as defined in DataGeneratorObservations.
             Keys must be that of `u_dict`.

jinns/loss/_LossODE.py

@@ -459,7 +459,7 @@ class SystemLossODE:
 
         Parameters
         ---------
-        params
+        params_dict
             A dictionary of dictionaries of parameters of the model.
             Typically, it is a dictionary of dictionaries of
             dictionaries: `eq_params` and `nn_params``, respectively the
@@ -489,7 +489,7 @@ class SystemLossODE:
         # and update vmap_in_axes
         if batch.param_batch_dict is not None:
             # update params with the batches of generated params
-            params = _update_eq_params_dict(params, batch.param_batch_dict)
+            params_dict = _update_eq_params_dict(params_dict, batch.param_batch_dict)
 
         vmap_in_axes_params = _get_vmap_in_axes_params(
             batch.param_batch_dict, params_dict

jinns/loss/_LossPDE.py

@@ -818,17 +818,9 @@ class LossPDENonStatio(LossPDEStatio):
             inputs/outputs/parameters
         """
 
-        omega_batch, omega_border_batch, times_batch = (
-            batch.inside_batch,
-            batch.border_batch,
-            batch.temporal_batch,
-        )
+        times_batch = batch.times_x_inside_batch[:, 0:1]
+        omega_batch = batch.times_x_inside_batch[:, 1:]
         n = omega_batch.shape[0]
-        nt = times_batch.shape[0]
-        times_batch = times_batch.reshape(nt, 1)
-
-        def rep_times(k):
-            return jnp.repeat(times_batch, k, axis=0)
 
         vmap_in_axes_x_t = (0, 0)
 
@@ -844,10 +836,6 @@ class LossPDENonStatio(LossPDEStatio):
 
         vmap_in_axes_params = _get_vmap_in_axes_params(batch.param_batch_dict, params)
 
-        if isinstance(self.u, PINN):
-            omega_batch = jnp.tile(omega_batch, reps=(nt, 1))  # it is tiled
-            times_batch = rep_times(n)  # it is repeated
-
         # dynamic part
         params_ = _set_derivatives(params, "dyn_loss", self.derivative_keys)
         if self.dynamic_loss is not None:
@@ -1372,27 +1360,12 @@ class SystemLossPDE:
 
         if isinstance(batch, PDEStatioBatch):
             omega_batch, _ = batch.inside_batch, batch.border_batch
-            n = omega_batch.shape[0]
             vmap_in_axes_x_or_x_t = (0,)
 
             batches = (omega_batch,)
         elif isinstance(batch, PDENonStatioBatch):
-            omega_batch, _, times_batch = (
-                batch.inside_batch,
-                batch.border_batch,
-                batch.temporal_batch,
-            )
-            n = omega_batch.shape[0]
-            nt = times_batch.shape[0]
-            times_batch = times_batch.reshape(nt, 1)
-
-            def rep_times(k):
-                return jnp.repeat(times_batch, k, axis=0)
-
-            # Moreover...
-            if isinstance(list(self.u_dict.values())[0], PINN):
-                omega_batch = jnp.tile(omega_batch, reps=(nt, 1))  # it is tiled
-                times_batch = rep_times(n)  # it is repeated
+            times_batch = batch.times_x_inside_batch[:, 0:1]
+            omega_batch = batch.times_x_inside_batch[:, 1:]
 
             batches = (omega_batch, times_batch)
             vmap_in_axes_x_or_x_t = (0, 0)

jinns/loss/_Losses.py

@@ -120,16 +120,23 @@ def boundary_condition_apply(
         else:
             raise ValueError("Other border batches are not implemented")
         b_losses_by_facet = jax.tree_util.tree_map(
-            lambda c, f, fa, d: jnp.mean(
-                loss_weight * _compute_boundary_loss(c, f, batch, u, params, fa, d)
+            lambda c, f, fa, d: (
+                None
+                if c is None
+                else jnp.mean(
+                    loss_weight * _compute_boundary_loss(c, f, batch, u, params, fa, d)
+                )
             ),
             omega_boundary_condition,
             omega_boundary_fun,
             facet_tree,
             omega_boundary_dim,
+            is_leaf=lambda x: x is None,
         )  # when exploring leaves with None value (no condition) the returned
         # mse is None and we get rid of the None leaves of b_losses_by_facet
         # with the tree_leaves below
+        # Note that to keep the behaviour given in the comment above we neede
+        # to specify is_leaf according to the note in the release of 0.4.29
     else:
         facet_tuple = tuple(f for f in range(batch[1].shape[-1]))
         b_losses_by_facet = jax.tree_util.tree_map(

jinns/loss/_boundary_conditions.py

@@ -279,23 +279,10 @@ def boundary_dirichlet_nonstatio(f, batch, u, params, facet, dim_to_apply):
     dim_to_apply
         A jnp.s\_ object. The dimension of u on which to apply the boundary condition
     """
-    _, omega_border_batch, times_batch = (
-        batch.inside_batch,
-        batch.border_batch,
-        batch.temporal_batch,
-    )
-    nt = times_batch.shape[0]
-    times_batch = times_batch.reshape(nt, 1)
-    omega_border_batch = omega_border_batch[..., facet]
+    times_batch = batch.times_x_border_batch[:, 0:1, facet]
+    omega_border_batch = batch.times_x_border_batch[:, 1:, facet]
 
     if isinstance(u, PINN):
-        tile_omega_border_batch = jnp.tile(
-            omega_border_batch, reps=(times_batch.shape[0], 1)
-        )
-
-        def rep_times(k):
-            return jnp.repeat(times_batch, k, axis=0)
-
         vmap_in_axes_params = _get_vmap_in_axes_params(batch.param_batch_dict, params)
         vmap_in_axes_x_t = (0, 0)
 
@@ -309,26 +296,13 @@ def boundary_dirichlet_nonstatio(f, batch, u, params, facet, dim_to_apply):
             vmap_in_axes_x_t + vmap_in_axes_params,
             0,
         )
-        res = v_u_boundary(
-            rep_times(omega_border_batch.shape[0]), tile_omega_border_batch, params
-        )
+        res = v_u_boundary(times_batch, omega_border_batch, params)
         mse_u_boundary = jnp.sum(
             res**2,
             axis=-1,
         )
     elif isinstance(u, SPINN):
-        tile_omega_border_batch = jnp.tile(
-            omega_border_batch, reps=(times_batch.shape[0], 1)
-        )
-
-        if omega_border_batch.shape[0] == 1:
-            omega_border_batch = jnp.tile(
-                omega_border_batch, reps=(times_batch.shape[0], 1)
-            )
-        # otherwise we require batches to have same shape and we do not need
-        # this operation
-
-        values = u(times_batch, tile_omega_border_batch, params)[..., dim_to_apply]
+        values = u(times_batch, omega_border_batch, params)[..., dim_to_apply]
         tx_grid = _get_grid(jnp.concatenate([times_batch, omega_border_batch], axis=-1))
         boundaries = _check_user_func_return(
             f(tx_grid[..., 0:1], tx_grid[..., 1:]), values.shape
@@ -367,14 +341,8 @@ def boundary_neumann_nonstatio(f, batch, u, params, facet, dim_to_apply):
     dim_to_apply
         A jnp.s\_ object. The dimension of u on which to apply the boundary condition
     """
-    _, omega_border_batch, times_batch = (
-        batch.inside_batch,
-        batch.border_batch,
-        batch.temporal_batch,
-    )
-    nt = times_batch.shape[0]
-    times_batch = times_batch.reshape(nt, 1)
-    omega_border_batch = omega_border_batch[..., facet]
+    times_batch = batch.times_x_border_batch[:, 0:1, facet]
+    omega_border_batch = batch.times_x_border_batch[:, 1:, facet]
 
     # We resort to the shape of the border_batch to determine the dimension as
     # described in the border_batch function
@@ -388,13 +356,6 @@ def boundary_neumann_nonstatio(f, batch, u, params, facet, dim_to_apply):
         n = jnp.array([[-1, 1, 0, 0], [0, 0, -1, 1]])
 
     if isinstance(u, PINN):
-        tile_omega_border_batch = jnp.tile(
-            omega_border_batch, reps=(times_batch.shape[0], 1)
-        )
-
-        def rep_times(k):
-            return jnp.repeat(times_batch, k, axis=0)
-
         vmap_in_axes_params = _get_vmap_in_axes_params(batch.param_batch_dict, params)
         vmap_in_axes_x_t = (0, 0)
 
@@ -411,8 +372,8 @@ def boundary_neumann_nonstatio(f, batch, u, params, facet, dim_to_apply):
         mse_u_boundary = jnp.sum(
             (
                 v_neumann(
-                    rep_times(omega_border_batch.shape[0]),
-                    tile_omega_border_batch,
+                    times_batch,
+                    omega_border_batch,
                     params,
                 )
             )
@@ -421,14 +382,6 @@ def boundary_neumann_nonstatio(f, batch, u, params, facet, dim_to_apply):
         )
 
     elif isinstance(u, SPINN):
-        if omega_border_batch.shape[0] == 1:
-            omega_border_batch = jnp.tile(
-                omega_border_batch, reps=(times_batch.shape[0], 1)
-            )
-            # ie case 1D
-        # otherwise we require batches to have same shape and we do not need
-        # this operation
-
         # the gradient we see in the PINN case can get gradients wrt to x
         # dimensions at once. But it would be very inefficient in SPINN because
         # of the high dim output of u. So we do 2 explicit forward AD, handling all the

jinns/solver/_rar.py

@@ -264,6 +264,9 @@ def _rar_step_init(sample_size, selected_sample_size):
             )
 
         elif isinstance(data, CubicMeshPDENonStatio):
+            if isinstance(loss.u, HYPERPINN) or isinstance(loss.u, SPINN):
+                raise NotImplementedError("RAR not implemented for hyperPINN and SPINN")
+
             # NOTE in this case sample_size and selected_sample_size
             # are tuples (times, omega) => we unpack them for clarity
             selected_sample_size_times, selected_sample_size_omega = (
@@ -274,14 +277,15 @@ def _rar_step_init(sample_size, selected_sample_size):
             new_times_samples = data.sample_in_time_domain(sample_size_times)
             new_omega_samples = data.sample_in_omega_domain(sample_size_omega)
 
-            if isinstance(loss.u, HYPERPINN) or isinstance(loss.u, SPINN):
-                raise NotImplementedError("RAR not implemented for hyperPINN and SPINN")
+            if not data.cartesian_product:
+                times = new_times_samples
+                omega = new_omega_samples
             else:
                 # do cartesian product on new points
-                tile_omega = jnp.tile(
+                omega = jnp.tile(
                     new_omega_samples, reps=(sample_size_times, 1)
                 )  # it is tiled
-                repeat_times = jnp.repeat(new_times_samples, sample_size_omega, axis=0)[
+                times = jnp.repeat(new_times_samples, sample_size_omega, axis=0)[
                     ..., None
                 ]  # it is repeated + add an axis
 
@@ -291,7 +295,7 @@ def _rar_step_init(sample_size, selected_sample_size):
                     (0, 0),
                     0,
                 )
-                dyn_on_s = v_dyn_loss(repeat_times, tile_omega).reshape(
+                dyn_on_s = v_dyn_loss(times, omega).reshape(
                     (sample_size_times, sample_size_omega)
                 )
                 mse_on_s = dyn_on_s**2
@@ -305,7 +309,7 @@ def _rar_step_init(sample_size, selected_sample_size):
                         (0, 0),
                         0,
                     )
-                    dyn_on_s += v_dyn_loss(repeat_times, tile_omega).reshape(
+                    dyn_on_s += v_dyn_loss(times, omega).reshape(
                         (sample_size_times, sample_size_omega)
                     )
 

jinns/solver/_solve.py

@@ -61,6 +61,7 @@ def solve(
     obs_data=None,
     validation=None,
     obs_batch_sharding=None,
+    verbose=True,
 ):
     """
     Performs the optimization process via stochastic gradient descent
@@ -132,6 +133,9 @@ def solve(
         Typically, a SingleDeviceSharding(gpu_device) when obs_data has been
         created with sharding_device=SingleDeviceSharding(cpu_device) to avoid
         loading on GPU huge datasets of observations
+    verbose:
+        Boolean, default True. If False, no std output (loss or cause of
+        exiting the optimization loop) will be produced.
 
     Returns
     -------
@@ -203,11 +207,14 @@ def solve(
         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
+        params=init_params,
+        last_non_nan_params=init_params.copy(),
+        opt_state=opt_state,
     )
     optimization_extra = OptimizationExtraContainer(
         curr_seq=curr_seq,
         seq2seq=seq2seq,
+        best_val_params=init_params.copy(),
     )
     loss_container = LossContainer(
         stored_loss_terms=stored_loss_terms,
@@ -222,7 +229,7 @@ def solve(
     else:
         validation_crit_values = None
 
-    break_fun = get_break_fun(n_iter)
+    break_fun = get_break_fun(n_iter, verbose)
 
     iteration = 0
     carry = (
@@ -272,7 +279,8 @@ def solve(
         )
 
         # Print train loss value during optimization
-        print_fn(i, train_loss_value, print_loss_every, prefix="[train] ")
+        if verbose:
+            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
@@ -281,6 +289,7 @@ def solve(
                 validation,  # always return `validation` for in-place mutation
                 early_stopping,
                 validation_criterion,
+                update_best_params,
             ) = jax.lax.cond(
                 i % validation.call_every == 0,
                 lambda operands: operands[0](*operands[1:]),  # validation.__call__()
@@ -288,6 +297,7 @@ def solve(
                     operands[0],
                     False,
                     validation_crit_values[i - 1],
+                    False,
                 ),
                 (
                     validation,  # validation must be in operands
@@ -295,12 +305,24 @@ def solve(
                 ),
             )
             # Print validation loss value during optimization
-            print_fn(i, validation_criterion, print_loss_every, prefix="[validation] ")
+            if verbose:
+                print_fn(
+                    i, validation_criterion, print_loss_every, prefix="[validation] "
+                )
             validation_crit_values = validation_crit_values.at[i].set(
                 validation_criterion
             )
+
+            # update best_val_params w.r.t val_loss if needed
+            best_val_params = jax.lax.cond(
+                update_best_params,
+                lambda _: params,  # update with current value
+                lambda operands: operands[0].best_val_params,  # unchanged
+                (optimization_extra,),
+            )
         else:
             early_stopping = False
+            best_val_params = params
 
         # Trigger RAR
         loss, params, data = trigger_rar(
@@ -329,13 +351,17 @@ def solve(
             loss_terms,
             tracked_params,
         )
+
+        # increment iteration number
         i += 1
 
         return (
             i,
             loss,
             OptimizationContainer(params, last_non_nan_params, opt_state),
-            OptimizationExtraContainer(curr_seq, seq2seq, early_stopping),
+            OptimizationExtraContainer(
+                curr_seq, seq2seq, best_val_params, early_stopping
+            ),
             DataGeneratorContainer(data, param_data, obs_data),
             validation,
             LossContainer(stored_loss_terms, train_loss_values),
@@ -364,36 +390,28 @@ def solve(
         validation_crit_values,
     ) = carry
 
-    jax.debug.print(
-        "Final iteration {i}: train loss value = {train_loss_val}",
-        i=i,
-        train_loss_val=loss_container.train_loss_values[i - 1],
-    )
+    if verbose:
+        jax.debug.print(
+            "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 (
         optimization.last_non_nan_params,
         loss_container.train_loss_values,
         loss_container.stored_loss_terms,
-        train_data.data,
-        loss,
+        train_data.data,  # return the DataGenerator if needed (no in-place modif)
+        loss,  # return the Loss if needed (no-inplace modif)
         optimization.opt_state,
         stored_objects.stored_params,
-        validation_crit_values,
+        validation_crit_values if validation is not None else None,
+        optimization_extra.best_val_params if validation is not None else None,
     )
 
 
@@ -453,15 +471,20 @@ def store_loss_and_params(
     tracked_params,
 ):
     stored_params = jax.tree_util.tree_map(
-        lambda stored_value, param, tracked_param: jax.lax.cond(
-            tracked_param,
-            lambda ope: ope[0].at[i].set(ope[1]),
-            lambda ope: ope[0],
-            (stored_value, param),
+        lambda stored_value, param, tracked_param: (
+            None
+            if stored_value is None
+            else jax.lax.cond(
+                tracked_param,
+                lambda ope: ope[0].at[i].set(ope[1]),
+                lambda ope: ope[0],
+                (stored_value, param),
+            )
         ),
         stored_params,
         params,
         tracked_params,
+        is_leaf=lambda x: x is None,
     )
     stored_loss_terms = jax.tree_util.tree_map(
         lambda stored_term, loss_term: stored_term.at[i].set(loss_term),
@@ -473,16 +496,20 @@ def store_loss_and_params(
     return (stored_params, stored_loss_terms, train_loss_values)
 
 
-def get_break_fun(n_iter):
+def get_break_fun(n_iter, verbose: str):
     """
-    Wrapper to get the break_fun with appropriate `n_iter`
+    Wrapper to get the break_fun with appropriate `n_iter`.
+    The verbose argument is here to control printing (or not) when exiting
+    the optimisation loop. It can be convenient is jinns.solve is itself
+    called in a loop and user want to avoid std output.
     """
 
     @jit
-    def break_fun(carry):
+    def break_fun(carry: tuple):
         """
         Function to break from the main optimization loop
-        We check several conditions
+        We the following conditions : maximum number of iterations, NaN
+        appearing in the parameters, and early stopping criterion.
         """
 
         def stop_while_loop(msg):
@@ -490,7 +517,8 @@ def get_break_fun(n_iter):
             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}")
+            if verbose:
+                jax.debug.print(f"Stopping main optimization loop, cause: {msg}")
             return False
 
         def continue_while_loop(_):

jinns/utils/_containers.py

@@ -45,6 +45,7 @@ class OptimizationContainer(NamedTuple):
 class OptimizationExtraContainer(NamedTuple):
     curr_seq: int
     seq2seq: Union[dict, None]
+    best_val_params: dict
     early_stopping: bool = False
 
 

jinns/utils/_hyperpinn.py

@@ -78,15 +78,9 @@ class HYPERPINN(PINN):
         parameters of the pinn (`self.params`)
         """
         pinn_params_flat = eqx.tree_at(
-            lambda p: tree_leaves(p, is_leaf=lambda x: isinstance(x, jnp.ndarray)),
+            lambda p: tree_leaves(p, is_leaf=eqx.is_array),
             self.params,
-            [hyper_output[0 : self.pinn_params_cumsum[0]]]
-            + [
-                hyper_output[
-                    self.pinn_params_cumsum[i] : self.pinn_params_cumsum[i + 1]
-                ]
-                for i in range(len(self.pinn_params_cumsum) - 1)
-            ],
+            jnp.split(hyper_output, self.pinn_params_cumsum[:-1]),
         )
 
         return tree_map(

jinns/utils/_spinn.py

@@ -81,6 +81,9 @@ class SPINN(eqx.Module):
     Basically a wrapper around the `__call__` function to be able to give a type to
     our former `self.u`
     The function create_SPINN has the role to population the `__call__` function
+
+    **NOTE**: SPINNs with `t` and `x` as inputs are best used with a
+    DataGenerator with `self.cartesian_product=False` for memory consideration
     """
 
     d: int
@@ -191,6 +194,9 @@ def create_SPINN(key, d, r, eqx_list, eq_type, m=1):
         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
     -------

jinns/validation/_validation.py

@@ -42,7 +42,7 @@ class AbstractValidationModule(eqx.Module):
     @abc.abstractmethod
     def __call__(
         self, params: PyTree
-    ) -> tuple["AbstractValidationModule", Bool, Array]:
+    ) -> tuple["AbstractValidationModule", Bool, Array, Bool]:
         raise NotImplementedError
 
 
@@ -86,10 +86,10 @@ class ValidationLoss(AbstractValidationModule):
             )
 
         validation_loss_value, _ = self.loss(params, val_batch)
-        (counter, best_val_loss) = jax.lax.cond(
+        (counter, best_val_loss, update_best_params) = 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
+            lambda _: (jnp.array(0.0), validation_loss_value, True),  # reset
+            lambda operands: (operands[0] + 1, operands[1], False),  # increment
             (self.counter, self.best_val_loss),
         )
 
@@ -108,7 +108,7 @@ class ValidationLoss(AbstractValidationModule):
             None,
         )
         # return `new` cause no in-place modification of the eqx.Module
-        return (new, bool_early_stopping, validation_loss_value)
+        return (new, bool_early_stopping, validation_loss_value, update_best_params)
 
 
 if __name__ == "__main__":

{jinns-0.8.10.dist-info → jinns-0.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.10
+Version: 0.9.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>
@@ -20,6 +20,7 @@ Requires-Dist: optax
 Requires-Dist: equinox
 Requires-Dist: jax-tqdm
 Requires-Dist: diffrax
+Requires-Dist: matplotlib
 Provides-Extra: notebook
 Requires-Dist: jupyter ; extra == 'notebook'
 Requires-Dist: matplotlib ; extra == 'notebook'

{jinns-0.8.10.dist-info → jinns-0.9.0.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jinns/__init__.py,sha256=T2XlmLbYqcXTumPJL00cJ80W98We5LH8Yg_Lss_exl4,139
-jinns/data/_DataGenerators.py,sha256=_um-giHQ8mCILUOJHX231njHTHZp4S7EcGrUs7R1dUs,61829
+jinns/data/_DataGenerators.py,sha256=aR-L2ZTHnogU0CFRemtcg1PHlQCp_UZFPnXvxE8hBj0,65069
 jinns/data/__init__.py,sha256=yBOmoavSD-cABp4XcjQY1zsEVO0mDyIhi2MJ5WNp0l8,326
 jinns/data/_display.py,sha256=vlqggDCgVMEwdGBtjVmZaTQORU6imSfDkssn2XCtITI,10392
 jinns/experimental/__init__.py,sha256=qWbhC7Z8UgLWy0t-zU7RYze6v13-FngiCYXu-2bRVFQ,296
@@ -8,29 +8,29 @@ jinns/experimental/_sinuspinn.py,sha256=hxSzscwMV2LayWOqenIlT1zqEVVrE5Y8CKf7bHX5
 jinns/experimental/_spectralpinn.py,sha256=-4795pa7AYtRNSE-ugan3gHh64mtu2VdrRG5AS_J9Eg,2654
 jinns/loss/_DynamicLoss.py,sha256=L4CVmmF0rTPbHntgqsLLHlnrlQgLHsetUocpJm7ZYag,27461
 jinns/loss/_DynamicLossAbstract.py,sha256=kTQlhLx7SBuH5dIDmYaE79sVHUZt1nUFa8LxPU5IHhM,8504
-jinns/loss/_LossODE.py,sha256=Y1mxryPVFf7ruqw_mGNACLExfx4iQT4R2bZP3s5rg4c,22172
-jinns/loss/_LossPDE.py,sha256=purAEtc0e71kv9XnZUT-a7MrkDAkM_3tTI4xJPu6fH4,61629
-jinns/loss/_Losses.py,sha256=XOL3MFiKEd3ndsc78Qnpi1vbgR0B2HaAWOGGW2meDM8,11190
+jinns/loss/_LossODE.py,sha256=Ava8kcAzj4-GiUkHJfZyhq15vGw7ABXk8Zpa3ynMlmY,22187
+jinns/loss/_LossPDE.py,sha256=zsm5gxIU2aRq_UdBOkQa6CyVXU6lqu740W8iAghDHV8,60707
+jinns/loss/_Losses.py,sha256=4tRnOgT31ZMDmdAKRuJv1e4Ob5zS5a54YnIVOQr0-uc,11480
 jinns/loss/__init__.py,sha256=pFNYUxns-NPXBFdqrEVSiXkQLfCtKw-t2trlhvLzpYE,355
-jinns/loss/_boundary_conditions.py,sha256=YfSnLZ25hXqQ5KWAuxOrWSKkf_oBqAc9GQV4z7MjWyQ,17434
+jinns/loss/_boundary_conditions.py,sha256=lqva0LKyPWls0zSTD-UgWfh_Aul_Q88NTN2Lz7KMM1M,15978
 jinns/loss/_operators.py,sha256=zDGJqYqeYH7xd-4dtGX9PS-pf0uSOpUUXGo5SVjIJ4o,11069
 jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jinns/solver/_rar.py,sha256=IYP-jdbM0rbjBtxislrBYBuj49p9_QDOqejZKCHrKg8,17072
+jinns/solver/_rar.py,sha256=wxoq9yNGd3UR6669FhqXJgD-LeKPBJ01t38alXculN8,17164
 jinns/solver/_seq2seq.py,sha256=S6IPfsXpS_fbqIqAy01eUM7GBSBSkRzURan_J-iXXzI,5632
-jinns/solver/_solve.py,sha256=mGi0zaT_fK_QpBjTxof5Ix4mmfmnPi66CNJ3GQFZuo4,19099
+jinns/solver/_solve.py,sha256=QGaiH9EXgk8t19250qB1KhYMNuIMlv0YL6RCM710K58,20376
 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/_containers.py,sha256=nH67kpgN3Ir6M31rqApNVR7xHV7bWzFtQsrDAYuQHDo,1521
+jinns/utils/_hyperpinn.py,sha256=gzT6keKvlLPZ5VnBxvW-IBDyLuAtPLgNdhtPulosmfc,10638
 jinns/utils/_optim.py,sha256=550kxH75TL30o1iKx1swJyP0KqyUPsJ7-imL1w65Qd0,4444
 jinns/utils/_pinn.py,sha256=mhA4-3PazyQTbWIx9oLaNwL0QDe8ZIBhbiy5J3kwa4I,9471
 jinns/utils/_save_load.py,sha256=qgZ23nUcB8-B5IZ2guuUWC4M7r5Lxd_Ms3staScdyJo,5668
-jinns/utils/_spinn.py,sha256=SzOUt1KHtB9QOpghpvitnXN-KEqXUXbvabC5k0TnKEo,7793
+jinns/utils/_spinn.py,sha256=2CKTQv2PCvprJJiUlKv3eeo6SIZ5ZCTdajL0D1sul90,8093
 jinns/utils/_utils.py,sha256=8dgvWXX9NT7_7-zltWp0C9tG45ZFNwXxueyxPBb4hjo,6740
 jinns/utils/_utils_uspinn.py,sha256=qcKcOw3zrwWSQyGVj6fD8c9GinHt_U6JWN_k0auTtXM,26039
 jinns/validation/__init__.py,sha256=Jv58mzgC3F7cRfXA6caicL1t_U0UAhbwLrmMNVg6E7s,66
-jinns/validation/_validation.py,sha256=KfetbzB0xTNdBcYLwFWjEtP63Tf9wJirlhgqLTJDyy4,6761
-jinns-0.8.10.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.8.10.dist-info/METADATA,sha256=5lGoyi2W9MRamQdHVgZnYflJtp__zWDGyTiYgmfGc6g,2483
-jinns-0.8.10.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jinns-0.8.10.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.8.10.dist-info/RECORD,,
+jinns/validation/_validation.py,sha256=ob2b9l2txJDujM8QR3RrS_eUrrz2RPsSadM07zXJJPs,6820
+jinns-0.9.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.9.0.dist-info/METADATA,sha256=e-K6ITCYu6DogJd5P6jg5C1vc7MNgHEoxdVmY4H-LcQ,2508
+jinns-0.9.0.dist-info/WHEEL,sha256=y4mX-SOX4fYIkonsAGA5N0Oy-8_gI4FXw5HNI1xqvWg,91
+jinns-0.9.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.9.0.dist-info/RECORD,,

{jinns-0.8.10.dist-info → jinns-0.9.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: setuptools (70.2.0)
 Root-Is-Purelib: true
 Tag: py3-none-any