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

jinns/{data/_display.py → plot/_plot.py}

@@ -1,5 +1,5 @@
 """
-Utility functions for plotting
+Utility functions for plotting in 1D and 2D, with and without time.
 """
 
 from functools import partial
@@ -8,48 +8,55 @@ import matplotlib.pyplot as plt
 import jax.numpy as jnp
 from jax import vmap
 from mpl_toolkits.axes_grid1 import ImageGrid
+from typing import Callable, List
+from jaxtyping import Array, Float, Bool
 
 
 def plot2d(
-    fun,
-    xy_data,
-    times=None,
-    Tmax=1,
-    title="",
-    figsize=(7, 7),
-    cmap="inferno",
-    spinn=False,
-    vmin_vmax=None,
-    ax_for_plot=None,
+    fun: Callable,
+    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,
+    ax_for_plot: plt.Axes | None = None,
 ):
     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)`.
+    $\Omega$. It handles both the
 
-    When in the non-stationnary case, the `times` argument gives the time
-    slices :math:`t_i` at which to plot :math:`u(t_i, x)`.
+     1. the stationary case $u(x)$
+     2. the non-stationnary case $u(t, x)$
+
+    In the non-stationnary case, the `times` argument gives the time
+    slices $t_i$ at which to plot $u(t_i, x)$.
 
 
     Parameters
     ----------
-    fun : _type_
-        _description_
-    xy_data : _type_
-        _description_
-    times : _type_, optional
-        _description_, by default None
-    Tmax : float, only in non-stationary cases
-        Useful if you used time re-scaling in the differential equation, by
-        default 1
-    title : str, optional
+    fun :
+        the function $u$ to plot on the meshgrid, and eventually the time
+        slices. It's suppose to have signature `u(x)` in the stationnary case,, and `u(t, x)` in the non-stationnary case. Use `partial` or `lambda to freeze / reorder any other arguments.
+    xy_data :
+        A list of 2 `jnp.Array` providing grid values for meshgrid creation
+    times :
+        list or Array of time slices where to plot the function. Use Tmax if
+        you trained with time-rescaling.
+    Tmax :
+        Useful if you used time rescaling in the differential equation for training, default to 1 (no rescaling).
+    title :
         plot title, by default ""
-    figsize : tuple, optional
-        _description_, by default (7, 7)
-    cmap : str, optional
-        _description_, by default "inferno"
-    vmin_vmax : tuple, optional
+    figsize :
+        By default (7, 7)
+    cmap :
+        the matplotlib color map used in the ImageGrid.
+    vmin_vmax :
         The colorbar minimum and maximum value. Defaults None.
-    ax_for_plot : Matplotlib axis, optional
+    spinn :
+        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+    ax_for_plot :
         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.
@@ -57,7 +64,7 @@ def plot2d(
     Raises
     ------
     ValueError
-        if xy_data is not a list of type 2
+        if xy_data is not a list of length 2
     """
 
     # if not isinstance(xy_data, jnp.ndarray) and not xy_data.shape[-1] == 2:
@@ -179,13 +186,13 @@ def plot2d(
 
 def _plot_2D_statio(
     v_fun,
-    mesh,
-    plot=True,
-    colorbar=True,
-    cmap="inferno",
-    figsize=(7, 7),
-    spinn=False,
-    vmin_vmax=None,
+    mesh: Float[Array, "nx*ny nx*ny"],
+    plot: Bool = True,
+    colorbar: Bool = True,
+    cmap: str = "inferno",
+    figsize: tuple[int, int] = (7, 7),
+    spinn: Bool = False,
+    vmin_vmax: tuple[float, float] = None,
 ):
     """Function that plot the function u(x) with 2-D input x using pcolormesh()
 
@@ -200,6 +207,12 @@ def _plot_2D_statio(
         either show or return the plot, by default True
     colorbar : bool, optional
         add a colorbar, by default True
+    cmap :
+        the matplotlib color map used in the ImageGrid.
+    figsize :
+        By default (7, 7)
+    spinn :
+        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
     vmin_vmax: tuple, optional
         The colorbar minimum and maximum value. Defaults None.
 
@@ -238,33 +251,37 @@ def _plot_2D_statio(
 
 
 def plot1d_slice(
-    fun,
-    xdata,
-    time_slices=jnp.array([0]),
-    Tmax=1,
-    title="",
-    figsize=(10, 10),
-    spinn=False,
+    fun: Callable[[float, float], float],
+    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,
 ):
     """Function for plotting time slices of a function :math:`f(t_i, x)` where
-    `t` is time (1-D) and x is 1-D
+    `t_i` is time (1-D) and x is 1-D
 
     Parameters
     ----------
-    fun : callable with two arguments `t` and `x`
+    fun
         f(t, x)
-    xdata : jnp.array
+    xdata
         the discretization of space
-    time_slices : list, optional
-        the time slices :math:`t_i` at which to plot, by default [0]
-    Tmax : int, optional
+    time_slices
+        the time slices :math:`t_i` at which to plot.
+    Tmax
         Useful if you used time re-scaling in the differential equation, by
         default 1
-    title : str, optional
+    title
         title of the plot, by default ""
-    figsize : tuple, optional
+    spinn :
+        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+    figsize
         size of the figure, by default (10, 10)
     """
+    if time_slices is None:
+        time_slices = jnp.array([0])
     plt.figure(figsize=figsize)
     for t in time_slices:
         if not spinn:
@@ -284,16 +301,16 @@ def plot1d_slice(
 
 
 def plot1d_image(
-    fun,
-    xdata,
-    times,
-    Tmax=1,
-    title="",
-    figsize=(10, 10),
-    colorbar=True,
-    cmap="inferno",
-    spinn=False,
-    vmin_vmax=None,
+    fun: Callable[[float, float], float],
+    xdata: Float[Array, "nx"],
+    times: Float[Array, "nt"],
+    Tmax: float = 1.0,
+    title: str = "",
+    figsize: tuple[int, int] = (10, 10),
+    colorbar: Bool = True,
+    cmap: str = "inferno",
+    spinn: Bool = False,
+    vmin_vmax: tuple[float, float] = 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).
@@ -302,19 +319,25 @@ def plot1d_image(
 
     Parameters
     ----------
-    fun : callable with two arguments t and x
-        the function to plot
-    xdata : jnp.array
+    fun :
+        callable with two arguments t and x the function to plot
+    xdata :
         the discretization of space
-    times : jnp.array
+    times :
         the discretization of time
-    Tmax : int, optional
-        _description_, by default 1
-    title : str, optional
-        , by default ""
-    figsize : tuple, optional
-        , by default (10, 10)
-    vmin_vmax: tuple
+    Tmax :
+        by default 1
+    title :
+        by default ""
+    figsize :
+        by default (10, 10)
+    colorbar :
+        Whether to add a colobar
+    cmap :
+        the matplotlib color map used in the ImageGrid.
+    spinn :
+        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+    vmin_vmax:
         The colorbar minimum and maximum value. Defaults None.
     """