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

jinns/data/_DataGenerators.py

@@ -98,6 +98,33 @@ def _reset_or_increment(bend, n_eff, operands):
     )
 
 
+def _check_and_set_rar_parameters(rar_parameters, n, n_start):
+    if rar_parameters is not None and n_start is None:
+        raise ValueError(
+            f"n_start or/and nt_start must be provided in the context of RAR sampling scheme, {n_start} was provided"
+        )
+    if rar_parameters is not None:
+        # Default p is None. However, in the RAR sampling scheme we use 0
+        # probability to specify non-used collocation points (i.e. points
+        # above n_start). Thus, p is a vector of probability of shape (n, 1).
+        p = jnp.zeros((n,))
+        p = p.at[:n_start].set(1 / n_start)
+        # set internal counter for the number of gradient steps since the
+        # last new collocation points have been added
+        # It is not 0 to ensure the first iteration of RAR happens just
+        # after start_iter. See the _proceed_to_rar() function in _rar.py
+        rar_iter_from_last_sampling = rar_parameters["update_every"] - 1
+        # set iternal counter for the number of times collocation points
+        # have been added
+        rar_iter_nb = 0
+    else:
+        p = None
+        rar_iter_from_last_sampling = None
+        rar_iter_nb = None
+
+    return n_start, p, rar_iter_from_last_sampling, rar_iter_nb
+
+
 #####################################################
 # DataGenerator for ODE : only returns time_batches
 #####################################################
@@ -150,10 +177,10 @@ class DataGeneratorODE:
             Default to None: do not use Residual Adaptative Resampling.
             Otherwise a dictionary with keys. `start_iter`: the iteration at
             which we start the RAR sampling scheme (we first have a burn in
-            period). `update_rate`: the number of gradient steps taken between
+            period). `update_every`: the number of gradient steps taken between
             each appending of collocation points in the RAR algo.
-            `sample_size`: the size of the sample from which we will select new
-            collocation points. `selected_sample_size`: the number of selected
+            `sample_size_times`: the size of the sample from which we will select new
+            collocation points. `selected_sample_size_times`: the number of selected
             points from the sample to be added to the current collocation
             points
             "DeepXDE: A deep learning library for solving differential
@@ -179,29 +206,13 @@ class DataGeneratorODE:
         self.method = method
         self.rar_parameters = rar_parameters
 
-        if rar_parameters is not None and nt_start is None:
-            raise ValueError(
-                "nt_start must be provided in the context of RAR sampling scheme"
-            )
-        if rar_parameters is not None:
-            self.nt_start = nt_start
-            # Default p is None. However, in the RAR sampling scheme we use 0
-            # probability to specify non-used collocation points (i.e. points
-            # above nt_start). Thus, p is a vector of probability of shape (nt, 1).
-            self.p = jnp.zeros((self.nt,))
-            self.p = self.p.at[: self.nt_start].set(1 / nt_start)
-            # set internal counter for the number of gradient steps since the
-            # last new collocation points have been added
-            self.rar_iter_from_last_sampling = 0
-            # set iternal counter for the number of times collocation points
-            # have been added
-            self.rar_iter_nb = 0
-
-        if rar_parameters is None or nt_start is None:
-            self.nt_start = self.nt
-            self.p = None
-            self.rar_iter_from_last_sampling = None
-            self.rar_iter_nb = None
+        # Set-up for RAR (if used)
+        (
+            self.nt_start,
+            self.p_times,
+            self.rar_iter_from_last_sampling,
+            self.rar_iter_nb,
+        ) = _check_and_set_rar_parameters(rar_parameters, n=nt, n_start=nt_start)
 
         if not self.data_exists:
             # Useful when using a lax.scan with pytree
@@ -238,7 +249,7 @@ class DataGeneratorODE:
             self.times,
             self.curr_time_idx,
             self.temporal_batch_size,
-            self.p,
+            self.p_times,
         )
 
     def temporal_batch(self):
@@ -253,7 +264,7 @@ class DataGeneratorODE:
         if self.rar_parameters is not None:
             nt_eff = (
                 self.nt_start
-                + self.rar_iter_nb * self.rar_parameters["selected_sample_size"]
+                + self.rar_iter_nb * self.rar_parameters["selected_sample_size_omega"]
             )
         else:
             nt_eff = self.nt
@@ -283,14 +294,19 @@ class DataGeneratorODE:
             self.curr_time_idx,
             self.tmin,
             self.tmax,
-            self.rar_parameters,
-            self.p,
+            self.p_times,
             self.rar_iter_from_last_sampling,
             self.rar_iter_nb,
         )  # arrays / dynamic values
         aux_data = {
             k: vars(self)[k]
-            for k in ["temporal_batch_size", "method", "nt", "nt_start"]
+            for k in [
+                "temporal_batch_size",
+                "method",
+                "nt",
+                "rar_parameters",
+                "nt_start",
+            ]
         }  # static values
         return (children, aux_data)
 
@@ -308,8 +324,7 @@ class DataGeneratorODE:
             curr_time_idx,
             tmin,
             tmax,
-            rar_parameters,
-            p,
+            p_times,
             rar_iter_from_last_sampling,
             rar_iter_nb,
         ) = children
@@ -318,12 +333,11 @@ class DataGeneratorODE:
             data_exists=True,
             tmin=tmin,
             tmax=tmax,
-            rar_parameters=rar_parameters,
             **aux_data,
         )
         obj.times = times
         obj.curr_time_idx = curr_time_idx
-        obj.p = p
+        obj.p_times = p_times
         obj.rar_iter_from_last_sampling = rar_iter_from_last_sampling
         obj.rar_iter_nb = rar_iter_nb
         return obj
@@ -412,10 +426,10 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             Default to None: do not use Residual Adaptative Resampling.
             Otherwise a dictionary with keys. `start_iter`: the iteration at
             which we start the RAR sampling scheme (we first have a burn in
-            period). `update_rate`: the number of gradient steps taken between
+            period). `update_every`: the number of gradient steps taken between
             each appending of collocation points in the RAR algo.
-            `sample_size`: the size of the sample from which we will select new
-            collocation points. `selected_sample_size`: the number of selected
+            `sample_size_omega`: the size of the sample from which we will select new
+            collocation points. `selected_sample_size_omega`: the number of selected
             points from the sample to be added to the current collocation
             points
             "DeepXDE: A deep learning library for solving differential
@@ -442,30 +456,12 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
         assert dim == len(max_pts) and isinstance(max_pts, tuple)
         self.n = n
         self.rar_parameters = rar_parameters
-
-        if rar_parameters is not None and n_start is None:
-            raise ValueError(
-                "n_start must be provided in the context of RAR sampling scheme"
-            )
-        if rar_parameters is not None:
-            self.n_start = n_start
-            # Default p is None. However, in the RAR sampling scheme we use 0
-            # probability to specify non-used collocation points (i.e. points
-            # above n_start). Thus, p is a vector of probability of shape (n, 1).
-            self.p = jnp.zeros((self.n,))
-            self.p = self.p.at[: self.n_start].set(1 / n_start)
-            # set internal counter for the number of gradient steps since the
-            # last new collocation points have been added
-            self.rar_iter_from_last_sampling = 0
-            # set iternal counter for the number of times collocation points
-            # have been added
-            self.rar_iter_nb = 0
-
-        if rar_parameters is None or n_start is None:
-            self.n_start = self.n
-            self.p = None
-            self.rar_iter_from_last_sampling = None
-            self.rar_iter_nb = None
+        (
+            self.n_start,
+            self.p_omega,
+            self.rar_iter_from_last_sampling,
+            self.rar_iter_nb,
+        ) = _check_and_set_rar_parameters(rar_parameters, n=n, n_start=n_start)
 
         self.p_border = None  # no RAR sampling for border for now
 
@@ -643,7 +639,7 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             self.omega,
             self.curr_omega_idx,
             self.omega_batch_size,
-            self.p,
+            self.p_omega,
         )
 
     def inside_batch(self):
@@ -656,7 +652,7 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
         if self.rar_parameters is not None:
             n_eff = (
                 self.n_start
-                + self.rar_iter_nb * self.rar_parameters["selected_sample_size"]
+                + self.rar_iter_nb * self.rar_parameters["selected_sample_size_omega"]
             )
         else:
             n_eff = self.n
@@ -745,8 +741,7 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             self.curr_omega_border_idx,
             self.min_pts,
             self.max_pts,
-            self.rar_parameters,
-            self.p,
+            self.p_omega,
             self.rar_iter_from_last_sampling,
             self.rar_iter_nb,
         )
@@ -759,6 +754,7 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
                 "omega_border_batch_size",
                 "method",
                 "dim",
+                "rar_parameters",
                 "n_start",
             ]
         }
@@ -780,8 +776,7 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             curr_omega_border_idx,
             min_pts,
             max_pts,
-            rar_parameters,
-            p,
+            p_omega,
             rar_iter_from_last_sampling,
             rar_iter_nb,
         ) = children
@@ -792,14 +787,13 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             data_exists=True,
             min_pts=min_pts,
             max_pts=max_pts,
-            rar_parameters=rar_parameters,
             **aux_data,
         )
         obj.omega = omega
         obj.omega_border = omega_border
         obj.curr_omega_idx = curr_omega_idx
         obj.curr_omega_border_idx = curr_omega_border_idx
-        obj.p = p
+        obj.p_omega = p_omega
         obj.rar_iter_from_last_sampling = rar_iter_from_last_sampling
         obj.rar_iter_nb = rar_iter_nb
         return obj
@@ -834,6 +828,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         method="grid",
         rar_parameters=None,
         n_start=None,
+        nt_start=None,
         data_exists=False,
     ):
         r"""
@@ -887,27 +882,23 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             Default to None: do not use Residual Adaptative Resampling.
             Otherwise a dictionary with keys. `start_iter`: the iteration at
             which we start the RAR sampling scheme (we first have a burn in
-            period). `update_rate`: the number of gradient steps taken between
+            period). `update_every`: the number of gradient steps taken between
             each appending of collocation points in the RAR algo.
-            `sample_size`: the size of the sample from which we will select new
-            collocation points. `selected_sample_size`: the number of selected
+            `sample_size_omega`: the size of the sample from which we will select new
+            collocation points. `selected_sample_size_omega`: the number of selected
             points from the sample to be added to the current collocation
             points.
-            __Note:__ that if RAR sampling is chosen it will currently affect both
-            self.times and self.omega with the same hyperparameters
-            (rar_parameters and n_start)
-            "DeepXDE: A deep learning library for solving differential
-            equations", L. Lu, SIAM Review, 2021
         n_start
             Defaults to None. The effective size of n used at start time.
             This value must be
             provided when rar_parameters is not None. Otherwise we set internally
             n_start = n and this is hidden from the user.
             In RAR, n_start
-            then corresponds to the initial number of points we train the PINN.
-            __Note:__ that if RAR sampling is chosen it will currently affect both
-            self.times and self.omega with the same hyperparameters
-            (rar_parameters and n_start)
+            then corresponds to the initial number of omega points we train the PINN.
+        nt_start
+            Defaults to None. A RAR hyper-parameter. Same as ``n_start`` but
+            for times collocation point. See also ``DataGeneratorODE``
+            documentation.
         data_exists
             Must be left to `False` when created by the user. Avoids the
             regeneration of :math:`\Omega`, :math:`\partial\Omega` and
@@ -931,6 +922,15 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         self.tmin = tmin
         self.tmax = tmax
         self.nt = nt
+
+        # Set-up for timewise RAR (some quantity are already set-up by super())
+        (
+            self.nt_start,
+            self.p_times,
+            _,
+            _,
+        ) = _check_and_set_rar_parameters(rar_parameters, n=nt, n_start=nt_start)
+
         if not self.data_exists:
             # Useful when using a lax.scan with pytree
             # Optionally can tell JAX not to re-generate data
@@ -950,7 +950,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             self.times,
             self.curr_time_idx,
             self.temporal_batch_size,
-            self.p,
+            self.p_times,
         )
 
     def generate_data_nonstatio(self):
@@ -980,7 +980,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         if self.rar_parameters is not None:
             nt_eff = (
                 self.n_start
-                + self.rar_iter_nb * self.rar_parameters["selected_sample_size"]
+                + self.rar_iter_nb * self.rar_parameters["selected_sample_size_times"]
             )
         else:
             nt_eff = self.nt
@@ -1022,8 +1022,8 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             self.max_pts,
             self.tmin,
             self.tmax,
-            self.rar_parameters,
-            self.p,
+            self.p_times,
+            self.p_omega,
             self.rar_iter_from_last_sampling,
             self.rar_iter_nb,
         )
@@ -1038,7 +1038,9 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
                 "temporal_batch_size",
                 "method",
                 "dim",
+                "rar_parameters",
                 "n_start",
+                "nt_start",
             ]
         }
         return (children, aux_data)
@@ -1063,8 +1065,8 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             max_pts,
             tmin,
             tmax,
-            rar_parameters,
-            p,
+            p_times,
+            p_omega,
             rar_iter_from_last_sampling,
             rar_iter_nb,
         ) = children
@@ -1075,7 +1077,6 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             max_pts=max_pts,
             tmin=tmin,
             tmax=tmax,
-            rar_parameters=rar_parameters,
             **aux_data,
         )
         obj.omega = omega
@@ -1084,9 +1085,11 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         obj.curr_omega_idx = curr_omega_idx
         obj.curr_omega_border_idx = curr_omega_border_idx
         obj.curr_time_idx = curr_time_idx
-        obj.p = p
+        obj.p_times = p_times
+        obj.p_omega = p_omega
         obj.rar_iter_from_last_sampling = rar_iter_from_last_sampling
         obj.rar_iter_nb = rar_iter_nb
+
         return obj
 
 

jinns/loss/_LossODE.py

@@ -237,6 +237,7 @@ class LossODE:
             "u": self.u,
             "dynamic_loss": self.dynamic_loss,
             "obs_slice": self.obs_slice,
+            "derivative_keys": self.derivative_keys,
         }
         return (children, aux_data)
 

jinns/solver/_rar.py

@@ -11,6 +11,41 @@ from jinns.loss._LossODE import LossODE, SystemLossODE
 from jinns.loss._DynamicLossAbstract import PDEStatio
 
 from functools import partial
+from jinns.utils._hyperpinn import HYPERPINN
+from jinns.utils._spinn import SPINN
+
+
+def _proceed_to_rar(data, i):
+    """Utilility function with various check to ensure we can proceed with the rar_step.
+    Return True if yes, and False otherwise"""
+
+    # Overall checks (universal for any data generator)
+    check_list = [
+        # check if burn-in period has ended
+        data.rar_parameters["start_iter"] <= i,
+        # check if enough iterations since last points added
+        (data.rar_parameters["update_every"] - 1) == data.rar_iter_from_last_sampling,
+    ]
+
+    # Memory allocation checks (depends on the type of DataGenerator)
+    # check if we still have room to append new collocation points in the
+    # allocated jnp.array (can concern `data.p_times` or `p_omega`)
+    if isinstance(data, DataGeneratorODE) or isinstance(data, CubicMeshPDENonStatio):
+        check_list.append(
+            data.rar_parameters["selected_sample_size_times"]
+            <= jnp.count_nonzero(data.p_times == 0),
+        )
+
+    if isinstance(data, CubicMeshPDEStatio) or isinstance(data, CubicMeshPDENonStatio):
+        # for now the above check are redundants but there may be a time when
+        # we drop inheritence
+        check_list.append(
+            data.rar_parameters["selected_sample_size_omega"]
+            <= jnp.count_nonzero(data.p_omega == 0),
+        )
+
+    proceed = jnp.all(jnp.array(check_list))
+    return proceed
 
 
 @partial(jax.jit, static_argnames=["_rar_step_true", "_rar_step_false"])
@@ -22,21 +57,7 @@ def trigger_rar(i, loss, params, data, _rar_step_true, _rar_step_false):
     else:
         # update `data` according to rar scheme.
         data = jax.lax.cond(
-            jnp.all(
-                jnp.array(
-                    [
-                        # check if enough it since last points added
-                        data.rar_parameters["update_rate"]
-                        == data.rar_iter_from_last_sampling,
-                        # check if burn in period has ended
-                        data.rar_parameters["start_iter"] < i,
-                        # check if we still have room to append new
-                        # collocation points in the allocated jnp array
-                        data.rar_parameters["selected_sample_size"]
-                        <= jnp.count_nonzero(data.p == 0),
-                    ]
-                )
-            ),
+            _proceed_to_rar(data, i),
             _rar_step_true,
             _rar_step_false,
             (loss, params, data, i),
@@ -49,13 +70,37 @@ def init_rar(data):
     Separated from the main rar, because the initialization to get _true and
     _false cannot be jit-ted.
     """
+    # NOTE if a user misspell some entry of ``rar_parameters`` the error
+    # risks to be a bit obscure but it should be ok.
     if data.rar_parameters is None:
         _rar_step_true, _rar_step_false = None, None
     else:
-        _rar_step_true, _rar_step_false = _rar_step_init(
-            data.rar_parameters["sample_size"],
-            data.rar_parameters["selected_sample_size"],
-        )
+        if isinstance(data, DataGeneratorODE):
+            # In this case we only need rar parameters related to `times`
+            _rar_step_true, _rar_step_false = _rar_step_init(
+                data.rar_parameters["sample_size_times"],
+                data.rar_parameters["selected_sample_size_times"],
+            )
+        elif isinstance(data, CubicMeshPDENonStatio):
+            # In this case we need rar parameters related to both `times`
+            # and`omega`
+            _rar_step_true, _rar_step_false = _rar_step_init(
+                (
+                    data.rar_parameters["sample_size_times"],
+                    data.rar_parameters["sample_size_omega"],
+                ),
+                (
+                    data.rar_parameters["selected_sample_size_times"],
+                    data.rar_parameters["selected_sample_size_omega"],
+                ),
+            )
+        elif isinstance(data, CubicMeshPDEStatio):
+            # In this case we only need rar parameters related to `omega`
+            _rar_step_true, _rar_step_false = _rar_step_init(
+                data.rar_parameters["sample_size_omega"],
+                data.rar_parameters["selected_sample_size_omega"],
+            )
+
         data.rar_parameters["iter_from_last_sampling"] = 0
 
     return data, _rar_step_true, _rar_step_false
@@ -64,7 +109,7 @@ def init_rar(data):
 def _rar_step_init(sample_size, selected_sample_size):
     """
     This is a wrapper because the sampling size and
-    selected_sample_size, must be treated static
+    selected_sample_size, must be treated as static
     in order to slice. So they must be set before jitting and not with the jitted
     dictionary values rar["test_points_nb"] and rar["added_points_nb"]
 
@@ -72,16 +117,10 @@ def _rar_step_init(sample_size, selected_sample_size):
     """
 
     def rar_step_true(operands):
-        """
-        Note: in all generality, we would need a stop gradient operator around
-        these dynamic_loss evaluations that follow which produce weights for
-        sampling. However, they appear through a argsort and sampling
-        operations which definitly kills gradient flows
-        """
         loss, params, data, i = operands
 
         if isinstance(data, DataGeneratorODE):
-            s = data.sample_in_time_domain(sample_size)
+            new_omega_samples = data.sample_in_time_domain(sample_size)
 
             # We can have different types of Loss
             if isinstance(loss, LossODE):
@@ -90,7 +129,7 @@ def _rar_step_init(sample_size, selected_sample_size):
                     (0),
                     0,
                 )
-                dyn_on_s = v_dyn_loss(s)
+                dyn_on_s = v_dyn_loss(new_omega_samples)
                 if dyn_on_s.ndim > 1:
                     mse_on_s = (jnp.linalg.norm(dyn_on_s, axis=-1) ** 2).flatten()
                 else:
@@ -106,7 +145,7 @@ def _rar_step_init(sample_size, selected_sample_size):
                         (0),
                         0,
                     )
-                    dyn_on_s = v_dyn_loss(s)
+                    dyn_on_s = v_dyn_loss(new_omega_samples)
                     if dyn_on_s.ndim > 1:
                         mse_on_s += (jnp.linalg.norm(dyn_on_s, axis=-1) ** 2).flatten()
                     else:
@@ -118,9 +157,7 @@ def _rar_step_init(sample_size, selected_sample_size):
                 (mse_on_s.shape[0] - selected_sample_size,),
                 (selected_sample_size,),
             )
-            higher_residual_points = s[higher_residual_idx]
-
-            data.rar_iter_from_last_sampling = 0
+            higher_residual_points = new_omega_samples[higher_residual_idx]
 
             ## add the new points in times
             # start indices of update can be dynamic but the the shape (length)
@@ -135,7 +172,7 @@ def _rar_step_init(sample_size, selected_sample_size):
             ## points are non-zero
             new_proba = 1 / (data.nt_start + data.rar_iter_nb * selected_sample_size)
             # the next work because nt_start is static
-            data.p = data.p.at[: data.nt_start].set(new_proba)
+            data.p_times = data.p_times.at[: data.nt_start].set(new_proba)
 
             # the next requires a fori_loop because the range is dynamic
             def update_slices(i, p):
@@ -147,16 +184,14 @@ def _rar_step_init(sample_size, selected_sample_size):
 
             data.rar_iter_nb += 1
 
-            data.p = jax.lax.fori_loop(0, data.rar_iter_nb, update_slices, data.p)
-
-            # NOTE must return data to be correctly updated because we cannot
-            # have side effects in this function that will be jitted
-            return data
+            data.p_times = jax.lax.fori_loop(
+                0, data.rar_iter_nb, update_slices, data.p_times
+            )
 
         elif isinstance(data, CubicMeshPDEStatio) and not isinstance(
             data, CubicMeshPDENonStatio
         ):
-            s = data.sample_in_omega_domain(sample_size)
+            new_omega_samples = data.sample_in_omega_domain(sample_size)
 
             # We can have different types of Loss
             if isinstance(loss, LossPDEStatio):
@@ -169,7 +204,7 @@ def _rar_step_init(sample_size, selected_sample_size):
                     (0),
                     0,
                 )
-                dyn_on_s = v_dyn_loss(s)
+                dyn_on_s = v_dyn_loss(new_omega_samples)
                 if dyn_on_s.ndim > 1:
                     mse_on_s = (jnp.linalg.norm(dyn_on_s, axis=-1) ** 2).flatten()
                 else:
@@ -185,7 +220,7 @@ def _rar_step_init(sample_size, selected_sample_size):
                         0,
                         0,
                     )
-                    dyn_on_s = v_dyn_loss(s)
+                    dyn_on_s = v_dyn_loss(new_omega_samples)
                     if dyn_on_s.ndim > 1:
                         mse_on_s += (jnp.linalg.norm(dyn_on_s, axis=-1) ** 2).flatten()
                     else:
@@ -197,12 +232,10 @@ def _rar_step_init(sample_size, selected_sample_size):
                 (mse_on_s.shape[0] - selected_sample_size,),
                 (selected_sample_size,),
             )
-            higher_residual_points = s[higher_residual_idx]
+            higher_residual_points = new_omega_samples[higher_residual_idx]
 
-            data.rar_iter_from_last_sampling = 0
-
-            ## add the new points in times
-            # start indices of update can be dynamic but the the shape (length)
+            ## add the new points in omega
+            # start indices of update can be dynamic but not the shape (length)
             # of the slice
             data.omega = jax.lax.dynamic_update_slice(
                 data.omega,
@@ -214,7 +247,7 @@ def _rar_step_init(sample_size, selected_sample_size):
             ## points are non-zero
             new_proba = 1 / (data.n_start + data.rar_iter_nb * selected_sample_size)
             # the next work because n_start is static
-            data.p = data.p.at[: data.n_start].set(new_proba)
+            data.p_omega = data.p_omega.at[: data.n_start].set(new_proba)
 
             # the next requires a fori_loop because the range is dynamic
             def update_slices(i, p):
@@ -226,145 +259,169 @@ def _rar_step_init(sample_size, selected_sample_size):
 
             data.rar_iter_nb += 1
 
-            data.p = jax.lax.fori_loop(0, data.rar_iter_nb, update_slices, data.p)
-
-            # NOTE must return data to be correctly updated because we cannot
-            # have side effects in this function that will be jitted
-            return data
+            data.p_omega = jax.lax.fori_loop(
+                0, data.rar_iter_nb, update_slices, data.p_omega
+            )
 
         elif isinstance(data, CubicMeshPDENonStatio):
-            st = data.sample_in_time_domain(sample_size)
-            sx = data.sample_in_omega_domain(sample_size)
-
-            # According to the Loss type we have different syntax to call the
-            # dynamic_loss evaluate function
-            if isinstance(loss, LossPDEStatio) and not isinstance(
-                loss, LossPDENonStatio
-            ):
-                # This case might not happen very often...
-                v_dyn_loss = vmap(
-                    lambda x: loss.dynamic_loss.evaluate(
-                        x,
-                        loss.u,
-                        params,
-                    ),
-                    (0),
-                    0,
-                )
-                dyn_on_s = v_dyn_loss(sx)
-                if dyn_on_s.ndim > 1:
-                    mse_on_s = (jnp.linalg.norm(dyn_on_s, axis=-1) ** 2).flatten()
-                else:
-                    mse_on_s = dyn_on_s**2
-            elif isinstance(loss, LossPDENonStatio):
+            # 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 = (
+                selected_sample_size
+            )
+            sample_size_times, sample_size_omega = 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")
+            else:
+                # do cartesian product on new points
+                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)[
+                    ..., None
+                ]  # it is repeated + add an axis
+
+            if isinstance(loss, LossPDENonStatio):
                 v_dyn_loss = vmap(
                     lambda t, x: loss.dynamic_loss.evaluate(t, x, loss.u, params),
                     (0, 0),
                     0,
                 )
-                dyn_on_s = v_dyn_loss(st[..., None], sx)
-                if dyn_on_s.ndim > 1:
-                    mse_on_s = (jnp.linalg.norm(dyn_on_s, axis=-1) ** 2).flatten()
-                else:
-                    mse_on_s = dyn_on_s**2
+                dyn_on_s = v_dyn_loss(repeat_times, tile_omega).reshape(
+                    (sample_size_times, sample_size_omega)
+                )
+                mse_on_s = dyn_on_s**2
             elif isinstance(loss, SystemLossPDE):
-                mse_on_s = 0
+                dyn_on_s = jnp.zeros((sample_size_times, sample_size_omega))
                 for i in loss.dynamic_loss_dict.keys():
-                    if isinstance(loss.dynamic_loss_dict[i], PDEStatio):
-                        v_dyn_loss = vmap(
-                            lambda x: loss.dynamic_loss_dict[i].evaluate(
-                                x, loss.u_dict, params
-                            ),
-                            0,
-                            0,
-                        )
-                        dyn_on_s = v_dyn_loss(sx)
-                        if dyn_on_s.ndim > 1:
-                            mse_on_s += (
-                                jnp.linalg.norm(dyn_on_s, axis=-1) ** 2
-                            ).flatten()
-                        else:
-                            mse_on_s += dyn_on_s**2
-                    else:
-                        v_dyn_loss = vmap(
-                            lambda t, x: loss.dynamic_loss_dict[i].evaluate(
-                                t, x, loss.u_dict, params
-                            ),
-                            (0, 0),
-                            0,
-                        )
-                        dyn_on_s = v_dyn_loss(st[..., None], sx)
-                        if dyn_on_s.ndim > 1:
-                            mse_on_s += (
-                                jnp.linalg.norm(dyn_on_s, axis=-1) ** 2
-                            ).flatten()
-                        else:
-                            mse_on_s += dyn_on_s**2
-
-            ## Now that we have the residuals, select the m points
-            # with higher dynamic loss (residuals)
-            higher_residual_idx = jax.lax.dynamic_slice(
-                jnp.argsort(mse_on_s),
-                (mse_on_s.shape[0] - selected_sample_size,),
-                (selected_sample_size,),
-            )
-            higher_residual_points_st = st[higher_residual_idx]
-            higher_residual_points_sx = sx[higher_residual_idx]
+                    v_dyn_loss = vmap(
+                        lambda t, x: loss.dynamic_loss_dict[i].evaluate(
+                            t, x, loss.u_dict, params
+                        ),
+                        (0, 0),
+                        0,
+                    )
+                    dyn_on_s += v_dyn_loss(repeat_times, tile_omega).reshape(
+                        (sample_size_times, sample_size_omega)
+                    )
 
-            data.rar_iter_from_last_sampling = 0
+            mse_on_s = dyn_on_s**2
+            # -- Select the m points with highest average residuals on time and
+            # -- space (times in rows / omega in columns)
+            # mean_times = mse_on_s.mean(axis=1)
+            # mean_omega = mse_on_s.mean(axis=0)
+            # times_idx = jax.lax.dynamic_slice(
+            #     jnp.argsort(mean_times),
+            #     (mse_on_s.shape[0] - selected_sample_size_times,),
+            #     (selected_sample_size_times,),
+            # )
+            # omega_idx = jax.lax.dynamic_slice(
+            #     jnp.argsort(mean_omega),
+            #     (mse_on_s.shape[1] - selected_sample_size_omega,),
+            #     (selected_sample_size_omega,),
+            # )
+
+            # -- Select the m worst points (t, x) with highest residuals
+            n_select = max(selected_sample_size_times, selected_sample_size_omega)
+            _, idx = jax.lax.top_k(mse_on_s.flatten(), k=n_select)
+            arr_idx = jnp.unravel_index(idx, mse_on_s.shape)
+            times_idx = arr_idx[0][:selected_sample_size_times]
+            omega_idx = arr_idx[1][:selected_sample_size_omega]
+
+            higher_residual_points_times = new_times_samples[times_idx]
+            higher_residual_points_omega = new_omega_samples[omega_idx]
 
             ## add the new points in times
-            # start indices of update can be dynamic but the the shape (length)
+            # start indices of update can be dynamic but not the shape (length)
             # of the slice
             data.times = jax.lax.dynamic_update_slice(
                 data.times,
-                higher_residual_points_st,
-                (data.n_start + data.rar_iter_nb * selected_sample_size,),
+                higher_residual_points_times,
+                (data.n_start + data.rar_iter_nb * selected_sample_size_times,),
             )
 
             ## add the new points in omega
             data.omega = jax.lax.dynamic_update_slice(
                 data.omega,
-                higher_residual_points_sx,
+                higher_residual_points_omega,
                 (
-                    data.n_start + data.rar_iter_nb * selected_sample_size,
+                    data.n_start + data.rar_iter_nb * selected_sample_size_omega,
                     data.dim,
                 ),
             )
 
             ## rearrange probabilities so that the probabilities of the new
             ## points are non-zero
-            new_proba = 1 / (data.n_start + data.rar_iter_nb * selected_sample_size)
+            new_p_times = 1 / (
+                data.nt_start + data.rar_iter_nb * selected_sample_size_times
+            )
             # the next work because nt_start is static
-            data.p = data.p.at[: data.n_start].set(new_proba)
+            data.p_times = data.p_times.at[: data.nt_start].set(new_p_times)
 
-            # the next requires a fori_loop because the range is dynamic
-            def update_slices(i, p):
-                return jax.lax.dynamic_update_slice(
-                    p,
-                    1 / new_proba * jnp.ones((selected_sample_size,)),
-                    ((data.n_start + i * selected_sample_size),),
-                )
+            # same for p_omega (work because n_start is static)
+            new_p_omega = 1 / (
+                data.n_start + data.rar_iter_nb * selected_sample_size_omega
+            )
+            data.p_omega = data.p_omega.at[: data.n_start].set(new_p_omega)
+
+            # the part of data.p_* after n_start requires a fori_loop because
+            # the range is dynamic
+            def create_update_slices(new_val, selected_sample_size):
+                def update_slices(i, p):
+                    new_p = jax.lax.dynamic_update_slice(
+                        p,
+                        new_val * jnp.ones((selected_sample_size,)),
+                        ((data.n_start + i * selected_sample_size),),
+                    )
+                    return new_p
+
+                return update_slices
 
             data.rar_iter_nb += 1
 
-            data.p = jax.lax.fori_loop(0, data.rar_iter_nb, update_slices, data.p)
+            ## update rest of p_times
+            update_slices_times = create_update_slices(
+                new_p_times, selected_sample_size_times
+            )
+            data.p_times = jax.lax.fori_loop(
+                0,
+                data.rar_iter_nb,
+                update_slices_times,
+                data.p_times,
+            )
+            ## update rest of p_omega
+            update_slices_omega = create_update_slices(
+                new_p_omega, selected_sample_size_omega
+            )
+            data.p_omega = jax.lax.fori_loop(
+                0,
+                data.rar_iter_nb,
+                update_slices_omega,
+                data.p_omega,
+            )
 
-            # NOTE must return data to be correctly updated because we cannot
-            # have side effects in this function that will be jitted
-            return data
+        # update RAR parameters for all cases
+        data.rar_iter_from_last_sampling = 0
+
+        # NOTE must return data to be correctly updated because we cannot
+        # have side effects in this function that will be jitted
+        return data
 
     def rar_step_false(operands):
         _, _, data, i = operands
 
         # Add 1 only if we are after the burn in period
-        data.rar_iter_from_last_sampling = jax.lax.cond(
-            i < data.rar_parameters["start_iter"],
-            lambda operand: 0,
-            lambda operand: operand + 1,
-            (data.rar_iter_from_last_sampling),
+        increment = jax.lax.cond(
+            i <= data.rar_parameters["start_iter"],
+            lambda: 0,
+            lambda: 1,
         )
 
+        data.rar_iter_from_last_sampling += increment
         return data
 
     return rar_step_true, rar_step_false

jinns/solver/_seq2seq.py

@@ -88,7 +88,7 @@ def initialize_seq2seq(loss, data, seq2seq, opt_state):
         data.curr_omega_idx = 0
         data.generate_time_data()
         data._key, data.times, _ = _reset_batch_idx_and_permute(
-            (data._key, data.times, data.curr_omega_idx, None, data.p)
+            (data._key, data.times, data.curr_omega_idx, None, data.p_times)
         )
         opt_state.hyperparams["learning_rate"] = seq2seq["learning_rate"][curr_seq]
 
@@ -145,7 +145,7 @@ def _update_seq2seq_SystemLossODE(operands):
     data.curr_omega_idx = 0
     data.generate_time_data()
     data._key, data.times, _ = _reset_batch_idx_and_permute(
-        (data._key, data.times, data.curr_omega_idx, None, data.p)
+        (data._key, data.times, data.curr_omega_idx, None, data.p_times)
     )
 
     opt_state.hyperparams["learning_rate"] = seq2seq["learning_rate"][curr_seq]

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

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

@@ -1,5 +1,5 @@
 jinns/__init__.py,sha256=T2XlmLbYqcXTumPJL00cJ80W98We5LH8Yg_Lss_exl4,139
-jinns/data/_DataGenerators.py,sha256=N4-U4z3MG46UIzHCbKScv9Z7AN40w1wlLY_VsVNj2sI,62293
+jinns/data/_DataGenerators.py,sha256=_um-giHQ8mCILUOJHX231njHTHZp4S7EcGrUs7R1dUs,61829
 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,15 +8,15 @@ 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=b9doBHoQwYvlgpqzrNO4dOaTN87LRvjHtHbz9bMoH7E,22119
+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/__init__.py,sha256=pFNYUxns-NPXBFdqrEVSiXkQLfCtKw-t2trlhvLzpYE,355
 jinns/loss/_boundary_conditions.py,sha256=YfSnLZ25hXqQ5KWAuxOrWSKkf_oBqAc9GQV4z7MjWyQ,17434
 jinns/loss/_operators.py,sha256=zDGJqYqeYH7xd-4dtGX9PS-pf0uSOpUUXGo5SVjIJ4o,11069
 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/_rar.py,sha256=IYP-jdbM0rbjBtxislrBYBuj49p9_QDOqejZKCHrKg8,17072
+jinns/solver/_seq2seq.py,sha256=S6IPfsXpS_fbqIqAy01eUM7GBSBSkRzURan_J-iXXzI,5632
 jinns/solver/_solve.py,sha256=mGi0zaT_fK_QpBjTxof5Ix4mmfmnPi66CNJ3GQFZuo4,19099
 jinns/utils/__init__.py,sha256=44ms5UR6vMw3Nf6u4RCAzPFs4fom_YbBnH9mfne8m6k,313
 jinns/utils/_containers.py,sha256=eYD277fO7X4EfX7PUFCCl69r3JBfh1sCfq8LkL5gd6o,1495
@@ -29,8 +29,8 @@ 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.8.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.8.8.dist-info/METADATA,sha256=oTs2EJMu4Bwo2n9DLsAPSU5edpbgPtwhNXBuW8YjpOc,2482
-jinns-0.8.8.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jinns-0.8.8.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.8.8.dist-info/RECORD,,
+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,,