jinns 1.3.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

jinns/plot/_plot.py

@@ -2,26 +2,24 @@
 Utility functions for plotting in 1D and 2D, with and without time.
 """
 
-from functools import partial
-import warnings
-import matplotlib.pyplot as plt
+from typing import Callable
 import jax.numpy as jnp
 from jax import vmap
+import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1 import ImageGrid
-from typing import Callable, List
-from jaxtyping import Array, Float, Bool
+from jaxtyping import Array, Float
 
 
 def plot2d(
     fun: Callable,
-    xy_data: tuple[Float[Array, "nx"], Float[Array, "ny"]],
-    times: Float[Array, "nt"] | List[float] | None = None,
+    xy_data: tuple[Float[Array, " nx"], Float[Array, " ny"]],
+    times: Float[Array, " nt"] | list[float] | None = None,
     Tmax: float = 1,
     title: str = "",
     figsize: tuple = (7, 7),
     cmap: str = "inferno",
     spinn: bool = False,
-    vmin_vmax: tuple[float, float] = [None, None],
+    vmin_vmax: tuple[float | None, float | None] | None = None,
 ):
     r"""Generic function for plotting functions over rectangular 2-D domains
     $\Omega$. It handles both the
@@ -61,6 +59,8 @@ def plot2d(
     ValueError
         if xy_data is not a list of length 2
     """
+    if vmin_vmax is None:
+        vmin_vmax = (None, None)
 
     # if not isinstance(xy_data, jnp.ndarray) and not xy_data.shape[-1] == 2:
     if not isinstance(xy_data, list) and not len(xy_data) == 2:
@@ -100,7 +100,7 @@ def plot2d(
             im = ax.pcolormesh(
                 mesh[0],
                 mesh[1],
-                ret[0],
+                ret,
                 cmap=cmap,
                 vmin=vmin_vmax[0],
                 vmax=vmin_vmax[1],
@@ -113,7 +113,7 @@ def plot2d(
         if not isinstance(times, list):
             try:
                 times = times.tolist()
-            except:
+            except AttributeError:
                 raise ValueError("times must be a list or an array")
 
         fig = plt.figure(figsize=figsize)
@@ -129,7 +129,7 @@ def plot2d(
             cbar_pad=0.4,
         )
 
-        for idx, (t, ax) in enumerate(zip(times, grid)):
+        for idx, (t, ax) in enumerate(zip(times, iter(grid))):
             if not spinn:
                 x_grid, y_grid = mesh
                 v_fun_at_t = vmap(fun)(
@@ -142,12 +142,11 @@ def plot2d(
                         axis=-1,
                     )
                 )
-                t_slice, _ = _plot_2D_statio(
+                t_slice = _plot_2D_statio(
                     v_fun_at_t,
                     mesh,
                     plot=False,  # only use to compute t_slice
                     colorbar=False,
-                    cmap=None,
                     vmin_vmax=vmin_vmax,
                 )
             elif spinn:
@@ -161,7 +160,7 @@ def plot2d(
                     axis=-1,
                 )
                 values_grid = jnp.squeeze(fun(t_x)[0]).T
-                t_slice, _ = _plot_2D_statio(
+                t_slice = _plot_2D_statio(
                     values_grid,
                     mesh,
                     plot=False,  # only use to compute t_slice
@@ -182,14 +181,14 @@ def plot2d(
 
 
 def _plot_2D_statio(
-    v_fun: Callable | Float[Array, "(nx*ny)^2 1"],
-    mesh: Float[Array, "nx*ny nx*ny"],
-    plot: Bool = True,
-    colorbar: Bool = True,
+    v_fun: Callable | Float[Array, " (nx*ny)^2 1"],
+    mesh: list[Float[Array, " nx*ny nx*ny"]],
+    plot: bool = True,
+    colorbar: bool = True,
     cmap: str = "inferno",
     figsize: tuple[int, int] = (7, 7),
-    vmin_vmax: tuple[float, float] = [None, None],
-):
+    vmin_vmax: tuple[float | None, float | None] | None = None,
+) -> Array | None:
     """Function that plot the function u(x) with 2-D input x using pcolormesh()
 
 
@@ -217,6 +216,8 @@ def _plot_2D_statio(
     Either None or the values of u() over the meshgrid and the current plt axis
 
     """
+    if vmin_vmax is None:
+        vmin_vmax = (None, None)
 
     x_grid, y_grid = mesh
     if callable(v_fun):
@@ -238,19 +239,18 @@ def _plot_2D_statio(
 
         if colorbar:
             fig.colorbar(im, format="%0.2f")
-        # don't plt.show() because it is done in plot2d()
     else:
-        return values_grid, plt.gca()
+        return values_grid
 
 
 def plot1d_slice(
-    fun: Callable[[float, float], float],
-    xdata: Float[Array, "nx"],
-    time_slices: Float[Array, "nt"] | None = None,
+    fun: Callable[[Float[Array, " "]], Float[Array, " "]],
+    xdata: Float[Array, " nx"],
+    time_slices: Float[Array, " nt"] | None = None,
     Tmax: float = 1.0,
     title: str = "",
     figsize: tuple[int, int] = (10, 10),
-    spinn: Bool = False,
+    spinn: bool = False,
     ax=None,
 ):
     """Function for plotting time slices of a function :math:`f(t_i, x)` where
@@ -284,7 +284,7 @@ def plot1d_slice(
     if time_slices is None:
         time_slices = jnp.array([0])
     if ax is None:
-        fig, ax = plt.subplots(figsize=figsize)
+        _, ax = plt.subplots(figsize=figsize)
 
     for t in time_slices:
         t_xdata = jnp.concatenate(
@@ -306,16 +306,16 @@ def plot1d_slice(
 
 
 def plot1d_image(
-    fun: Callable[[float, float], float],
-    xdata: Float[Array, "nx"],
-    times: Float[Array, "nt"],
+    fun: Callable[[Float[Array, " "]], Float[Array, " "]],
+    xdata: Float[Array, " nx"],
+    times: Float[Array, " nt"],
     Tmax: float = 1.0,
     title: str = "",
     figsize: tuple[int, int] = (10, 10),
-    colorbar: Bool = True,
+    colorbar: bool = True,
     cmap: str = "inferno",
-    spinn: Bool = False,
-    vmin_vmax: tuple[float, float] = [None, None],
+    spinn: bool = False,
+    vmin_vmax: tuple[float | None, float | None] | None = 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).
@@ -350,7 +350,8 @@ def plot1d_image(
     fig, ax
         A `matplotlib` `Figure` and `Axes` objects with the figure.
     """
-
+    if vmin_vmax is None:
+        vmin_vmax = (None, None)
     mesh = jnp.meshgrid(times, xdata)  # cartesian product
     if not spinn:
         # the trick is to use _plot2Dstatio

jinns/solver/_rar.py

@@ -2,41 +2,73 @@ from __future__ import (
     annotations,
 )  # https://docs.python.org/3/library/typing.html#constant
 
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Callable, TypeAlias, Any, TypedDict
 from functools import partial
+from jaxtyping import Float, Array, Bool
 import jax
 from jax import vmap
 import jax.numpy as jnp
 import equinox as eqx
-from jaxtyping import Int, Bool
-
-from jinns.data._Batchs import *
-from jinns.loss._LossODE import LossODE, SystemLossODE
-from jinns.loss._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
-from jinns.loss._loss_utils import dynamic_loss_apply
-from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-)
+
+from jinns.data._DataGeneratorODE import DataGeneratorODE
+from jinns.data._CubicMeshPDEStatio import CubicMeshPDEStatio
+from jinns.data._CubicMeshPDENonStatio import CubicMeshPDENonStatio
 from jinns.nn._hyperpinn import HyperPINN
 from jinns.nn._spinn import SPINN
 
 
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+    from jinns.utils._types import AnyLoss
+    from jinns.parameters._params import Params
+
+    class DataGeneratorWithRAR(AbstractDataGenerator):
+        """
+        Add the required RAR operands for type checks
+        """
+
+        rar_parameters: RarParameterDict
+        n_start: int
+        rar_iter_from_last_sampling: int
+        rar_iter_nb: int
+        p: Float[Array, " n 1"]
+
+    rar_operands: TypeAlias = tuple[Any, Params, DataGeneratorWithRAR, int]
+
+
+class RarParameterDict(TypedDict):
+    """
+    TypedDict to specify the Residual Adaptative Resampling procedure
+    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_every`: the number of gradient steps taken between
+    each update 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
+    points from the sample to be added to the current collocation
+    points.
+    """
 
+    start_iter: int
+    update_every: int
+    sample_size: int
+    selected_sample_size: int
 
-def _proceed_to_rar(data: AnyDataGenerator, i: Int) -> Bool:
+
+def _proceed_to_rar(data: DataGeneratorWithRAR, i: int) -> Bool[Array, " "]:
     """Utilility function with various check to ensure we can proceed with the rar_step.
     Return True if yes, and False otherwise"""
 
     # Overall checks
     check_list = [
         # check if burn-in period has ended
-        data.rar_parameters["start_iter"] <= i,
+        jnp.asarray(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,
+        jnp.asarray(
+            (data.rar_parameters["update_every"] - 1)
+            == data.rar_iter_from_last_sampling
+        ),
     ]
 
     # Memory allocation checks
@@ -52,14 +84,13 @@ def _proceed_to_rar(data: AnyDataGenerator, i: Int) -> Bool:
 
 @partial(jax.jit, static_argnames=["_rar_step_true", "_rar_step_false"])
 def trigger_rar(
-    i: Int,
+    i: int,
     loss: AnyLoss,
-    params: AnyParams,
-    data: AnyDataGenerator,
-    _rar_step_true: Callable[[rar_operands], AnyDataGenerator],
-    _rar_step_false: Callable[[rar_operands], AnyDataGenerator],
-) -> tuple[AnyLoss, AnyParams, AnyDataGenerator]:
-
+    params: Params,
+    data: DataGeneratorWithRAR,
+    _rar_step_true: Callable[[rar_operands], DataGeneratorWithRAR],
+    _rar_step_false: Callable[[rar_operands], DataGeneratorWithRAR],
+) -> tuple[AnyLoss, Params, DataGeneratorWithRAR]:
     if data.rar_parameters is None:
         # do nothing.
         return loss, params, data
@@ -75,11 +106,11 @@ def trigger_rar(
 
 
 def init_rar(
-    data: AnyDataGenerator,
+    data: DataGeneratorWithRAR,
 ) -> tuple[
-    AnyDataGenerator,
-    Callable[[rar_operands], AnyDataGenerator],
-    Callable[[rar_operands], AnyDataGenerator],
+    DataGeneratorWithRAR,
+    Callable[[rar_operands], DataGeneratorWithRAR] | None,
+    Callable[[rar_operands], DataGeneratorWithRAR] | None,
 ]:
     """
     Separated from the main rar, because the initialization to get _true and
@@ -100,9 +131,11 @@ def init_rar(
     return data, _rar_step_true, _rar_step_false
 
 
-def _rar_step_init(sample_size: Int, selected_sample_size: Int) -> tuple[
-    Callable[[rar_operands], AnyDataGenerator],
-    Callable[[rar_operands], AnyDataGenerator],
+def _rar_step_init(
+    sample_size: int, selected_sample_size: int
+) -> tuple[
+    Callable[[rar_operands], DataGeneratorWithRAR],
+    Callable[[rar_operands], DataGeneratorWithRAR],
 ]:
     """
     This is a wrapper because the sampling size and
@@ -113,13 +146,12 @@ def _rar_step_init(sample_size: Int, selected_sample_size: Int) -> tuple[
     This is a kind of manual declaration of static argnums
     """
 
-    def rar_step_true(operands: rar_operands) -> AnyDataGenerator:
-        loss, params, data, i = operands
+    def rar_step_true(operands: rar_operands) -> DataGeneratorWithRAR:
+        loss, params, data, _ = operands
         if isinstance(loss.u, HyperPINN) or isinstance(loss.u, SPINN):
             raise NotImplementedError("RAR not implemented for hyperPINN and SPINN")
 
         if isinstance(data, DataGeneratorODE):
-
             new_key, subkey = jax.random.split(data.key)
             new_samples = data.sample_in_time_domain(subkey, sample_size)
             data = eqx.tree_at(lambda m: m.key, data, new_key)
@@ -145,33 +177,15 @@ def _rar_step_init(sample_size: Int, selected_sample_size: Int) -> tuple[
 
             data = eqx.tree_at(lambda m: m.key, data, new_key)
 
-        # We can have different types of Loss
-        if isinstance(loss, (LossODE, LossPDEStatio, LossPDENonStatio)):
-            v_dyn_loss = vmap(
-                lambda inputs: loss.dynamic_loss.evaluate(inputs, loss.u, params),
-            )
-            dyn_on_s = v_dyn_loss(new_samples)
+        v_dyn_loss = vmap(
+            lambda inputs: loss.dynamic_loss.evaluate(inputs, loss.u, params),
+        )
+        dyn_on_s = v_dyn_loss(new_samples)
 
-            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, SystemLossODE, SystemLossPDE):
-            mse_on_s = 0
-
-            for i in loss.dynamic_loss_dict.keys():
-                v_dyn_loss = vmap(
-                    lambda inputs: loss.dynamic_loss_dict[i].evaluate(
-                        inputs, loss.u_dict, params
-                    ),
-                    (0),
-                    0,
-                )
-                dyn_on_s = v_dyn_loss(new_samples)
-                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
+        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
 
         ## Select the m points with higher dynamic loss
         higher_residual_idx = jax.lax.dynamic_slice(
@@ -188,7 +202,7 @@ def _rar_step_init(sample_size: Int, selected_sample_size: Int) -> tuple[
             new_times = jax.lax.dynamic_update_slice(
                 data.times,
                 higher_residual_points,
-                (data.n_start + data.rar_iter_nb * selected_sample_size,),
+                (data.n_start + data.rar_iter_nb * selected_sample_size,),  # type: ignore
             )
 
             data = eqx.tree_at(lambda m: m.times, data, new_times)
@@ -198,7 +212,7 @@ def _rar_step_init(sample_size: Int, selected_sample_size: Int) -> tuple[
             new_omega = jax.lax.dynamic_update_slice(
                 data.omega,
                 higher_residual_points,
-                (data.n_start + data.rar_iter_nb * selected_sample_size, data.dim),
+                (data.n_start + data.rar_iter_nb * selected_sample_size, data.dim),  # type: ignore
             )
 
             data = eqx.tree_at(lambda m: m.omega, data, new_omega)
@@ -207,7 +221,10 @@ def _rar_step_init(sample_size: Int, selected_sample_size: Int) -> tuple[
             new_domain = jax.lax.dynamic_update_slice(
                 data.domain,
                 higher_residual_points,
-                (data.n_start + data.rar_iter_nb * selected_sample_size, 1 + data.dim),
+                (
+                    data.n_start + data.rar_iter_nb * selected_sample_size,  # type: ignore
+                    1 + data.dim,
+                ),
             )
 
             data = eqx.tree_at(lambda m: m.domain, data, new_domain)
@@ -246,7 +263,7 @@ def _rar_step_init(sample_size: Int, selected_sample_size: Int) -> tuple[
         # have side effects in this function that will be jitted
         return data
 
-    def rar_step_false(operands: rar_operands) -> AnyDataGenerator:
+    def rar_step_false(operands: rar_operands) -> DataGeneratorWithRAR:
         _, _, data, i = operands
 
         # Add 1 only if we are after the burn in period