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

jinns/__init__.py

@@ -2,4 +2,5 @@ import jinns.data
 import jinns.loss
 import jinns.solver
 import jinns.utils
+import jinns.experimental
 from jinns.solver._solve import solve

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/data/_display.py

@@ -1,8 +1,13 @@
+"""
+Utility functions for plotting
+"""
+
+from functools import partial
+import warnings
 import matplotlib.pyplot as plt
 import jax.numpy as jnp
 from jax import vmap
 from mpl_toolkits.axes_grid1 import ImageGrid
-from functools import partial
 
 
 def plot2d(
@@ -14,8 +19,10 @@ def plot2d(
     figsize=(7, 7),
     cmap="inferno",
     spinn=False,
+    vmin_vmax=None,
+    ax_for_plot=None,
 ):
-    """Generic function for plotting functions over rectangular 2-D domains
+    r"""Generic function for plotting functions over rectangular 2-D domains
     :math:`\Omega`. It treats both the stationary case :math:`u(x)` or the
     non-stationnary case :math:`u(t, x)`.
 
@@ -40,6 +47,12 @@ def plot2d(
         _description_, by default (7, 7)
     cmap : str, optional
         _description_, by default "inferno"
+    vmin_vmax : tuple, optional
+        The colorbar minimum and maximum value. Defaults None.
+    ax_for_plot : Matplotlib axis, optional
+        If None, jinns triggers the plotting. Otherwise this argument
+        corresponds to the axis which will host the plot. Default is None.
+        NOTE: that this argument will have an effect only if times is None.
 
     Raises
     ------
@@ -59,25 +72,49 @@ def plot2d(
         # Statio case : expect a function of one argument fun(x)
         if not spinn:
             v_fun = vmap(fun, 0, 0)
-            _plot_2D_statio(
-                v_fun, mesh, plot=True, colorbar=True, cmap=cmap, figsize=figsize
+            ret = _plot_2D_statio(
+                v_fun,
+                mesh,
+                plot=not ax_for_plot,
+                colorbar=True,
+                cmap=cmap,
+                figsize=figsize,
+                vmin_vmax=vmin_vmax,
             )
         elif spinn:
             values_grid = jnp.squeeze(
                 fun(jnp.stack([xy_data[0][..., None], xy_data[1][..., None]], axis=1))
             )
-            _plot_2D_statio(
+            ret = _plot_2D_statio(
                 values_grid,
                 mesh,
-                plot=True,
+                plot=not ax_for_plot,
                 colorbar=True,
                 cmap=cmap,
                 spinn=True,
                 figsize=figsize,
+                vmin_vmax=vmin_vmax,
             )
-        plt.title(title)
+        if not ax_for_plot:
+            plt.title(title)
+        else:
+            if vmin_vmax is not None:
+                im = ax_for_plot.pcolormesh(
+                    mesh[0],
+                    mesh[1],
+                    ret[0],
+                    cmap=cmap,
+                    vmin=vmin_vmax[0],
+                    vmax=vmin_vmax[1],
+                )
+            else:
+                im = ax_for_plot.pcolormesh(mesh[0], mesh[1], ret[0], cmap=cmap)
+            ax_for_plot.set_title(title)
+            ax_for_plot.cax.colorbar(im, format="%0.2f")
 
     else:
+        if ax_for_plot is not None:
+            warnings.warn("ax_for_plot is ignored. jinns will plot the figure")
         if not isinstance(times, list):
             try:
                 times = times.tolist()
@@ -101,7 +138,12 @@ def plot2d(
             if not spinn:
                 v_fun_at_t = vmap(lambda x: fun(t=jnp.array([t]), x=x), 0, 0)
                 t_slice, _ = _plot_2D_statio(
-                    v_fun_at_t, mesh, plot=False, colorbar=False, cmap=None
+                    v_fun_at_t,
+                    mesh,
+                    plot=False,
+                    colorbar=False,
+                    cmap=None,
+                    vmin_vmax=vmin_vmax,
                 )
             elif spinn:
                 values_grid = jnp.squeeze(
@@ -113,15 +155,37 @@ def plot2d(
                     )[0]
                 )
                 t_slice, _ = _plot_2D_statio(
-                    values_grid, mesh, plot=False, colorbar=True, spinn=True
+                    values_grid,
+                    mesh,
+                    plot=False,
+                    colorbar=True,
+                    spinn=True,
+                    vmin_vmax=vmin_vmax,
+                )
+            if vmin_vmax is not None:
+                im = ax.pcolormesh(
+                    mesh[0],
+                    mesh[1],
+                    t_slice,
+                    cmap=cmap,
+                    vmin=vmin_vmax[0],
+                    vmax=vmin_vmax[1],
                 )
-            im = ax.pcolormesh(mesh[0], mesh[1], t_slice, cmap=cmap)
+            else:
+                im = ax.pcolormesh(mesh[0], mesh[1], t_slice, cmap=cmap)
             ax.set_title(f"t = {times[idx] * Tmax:.2f}")
             ax.cax.colorbar(im, format="%0.2f")
 
 
 def _plot_2D_statio(
-    v_fun, mesh, plot=True, colorbar=True, cmap="inferno", figsize=(7, 7), spinn=False
+    v_fun,
+    mesh,
+    plot=True,
+    colorbar=True,
+    cmap="inferno",
+    figsize=(7, 7),
+    spinn=False,
+    vmin_vmax=None,
 ):
     """Function that plot the function u(x) with 2-D input x using pcolormesh()
 
@@ -136,6 +200,8 @@ def _plot_2D_statio(
         either show or return the plot, by default True
     colorbar : bool, optional
         add a colorbar, by default True
+    vmin_vmax: tuple, optional
+        The colorbar minimum and maximum value. Defaults None.
 
     Returns
     -------
@@ -153,7 +219,17 @@ def _plot_2D_statio(
 
     if plot:
         fig = plt.figure(figsize=figsize)
-        im = plt.pcolormesh(x_grid, y_grid, values_grid, cmap=cmap)
+        if vmin_vmax is not None:
+            im = plt.pcolormesh(
+                x_grid,
+                y_grid,
+                values_grid,
+                cmap=cmap,
+                vmin=vmin_vmax[0],
+                vmax=vmin_vmax[1],
+            )
+        else:
+            im = plt.pcolormesh(x_grid, y_grid, values_grid, cmap=cmap)
         if colorbar:
             fig.colorbar(im, format="%0.2f")
         # don't plt.show() because it is done in plot2d()
@@ -217,6 +293,7 @@ def plot1d_image(
     colorbar=True,
     cmap="inferno",
     spinn=False,
+    vmin_vmax=None,
 ):
     """Function for plotting the 2-D image of a function :math:`f(t, x)` where
     `t` is time (1-D) and x is space (1-D).
@@ -237,6 +314,8 @@ def plot1d_image(
         , by default ""
     figsize : tuple, optional
         , by default (10, 10)
+    vmin_vmax: tuple
+        The colorbar minimum and maximum value. Defaults None.
     """
 
     mesh = jnp.meshgrid(times, xdata)  # cartesian product
@@ -250,7 +329,17 @@ def plot1d_image(
     elif spinn:
         values_grid = jnp.squeeze(fun((times[..., None]), xdata[..., None]).T)
     fig, ax = plt.subplots(1, 1, figsize=figsize)
-    im = ax.pcolormesh(mesh[0] * Tmax, mesh[1], values_grid, cmap=cmap)
+    if vmin_vmax is not None:
+        im = ax.pcolormesh(
+            mesh[0] * Tmax,
+            mesh[1],
+            values_grid,
+            cmap=cmap,
+            vmin=vmin_vmax[0],
+            vmax=vmin_vmax[1],
+        )
+    else:
+        im = ax.pcolormesh(mesh[0] * Tmax, mesh[1], values_grid, cmap=cmap)
     if colorbar:
         fig.colorbar(im, format="%0.2f")
     ax.set_title(title)

jinns/experimental/__init__.py

@@ -6,3 +6,5 @@ from ._diffrax_solver import (
     neumann_boundary_condition,
     plot_diffrax_solution,
 )
+from ._sinuspinn import create_sinusPINN
+from ._spectralpinn import create_spectralPINN