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

jinns/plot/_plot.py

@@ -21,8 +21,7 @@ def plot2d(
     figsize: tuple = (7, 7),
     cmap: str = "inferno",
     spinn: bool = False,
-    vmin_vmax: tuple[float, float] | None = None,
-    ax_for_plot: plt.Axes | None = None,
+    vmin_vmax: tuple[float, float] = [None, None],
 ):
     r"""Generic function for plotting functions over rectangular 2-D domains
     $\Omega$. It handles both the
@@ -55,11 +54,7 @@ def plot2d(
     vmin_vmax :
         The colorbar minimum and maximum value. Defaults None.
     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.
+        True if the function is a `SPINN` object.
 
     Raises
     ------
@@ -73,55 +68,48 @@ def plot2d(
             "xy_data must be a list of length 2 containing"
             "jnp.array of shape (nx,) and (ny,)."
         )
+
     mesh = jnp.meshgrid(xy_data[0], xy_data[1])  # cartesian product
 
     if times is None:
         # Statio case : expect a function of one argument fun(x)
+
         if not spinn:
             v_fun = vmap(fun, 0, 0)
             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))
-            )
+            values_grid = jnp.squeeze(fun(jnp.stack([xy_data[0], xy_data[1]], axis=1)))
             ret = _plot_2D_statio(
                 values_grid,
                 mesh,
-                plot=not ax_for_plot,
                 colorbar=True,
                 cmap=cmap,
-                spinn=True,
                 figsize=figsize,
                 vmin_vmax=vmin_vmax,
             )
-        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")
+            fig, ax = plt.subplots(1, 1, figsize=figsize)
+
+            im = ax.pcolormesh(
+                mesh[0],
+                mesh[1],
+                ret[0],
+                cmap=cmap,
+                vmin=vmin_vmax[0],
+                vmax=vmin_vmax[1],
+            )
+
+            ax.set_title(title)
+            fig.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()
@@ -143,56 +131,64 @@ def plot2d(
 
         for idx, (t, ax) in enumerate(zip(times, grid)):
             if not spinn:
-                v_fun_at_t = vmap(lambda x: fun(t=jnp.array([t]), x=x), 0, 0)
+                x_grid, y_grid = mesh
+                v_fun_at_t = vmap(fun)(
+                    jnp.concatenate(
+                        [
+                            t
+                            * jnp.ones((xy_data[0].shape[0] * xy_data[1].shape[0], 1)),
+                            jnp.vstack([x_grid.flatten(), y_grid.flatten()]).T,
+                        ],
+                        axis=-1,
+                    )
+                )
                 t_slice, _ = _plot_2D_statio(
                     v_fun_at_t,
                     mesh,
-                    plot=False,
+                    plot=False,  # only use to compute t_slice
                     colorbar=False,
                     cmap=None,
                     vmin_vmax=vmin_vmax,
                 )
             elif spinn:
-                values_grid = jnp.squeeze(
-                    fun(
+                t_x = jnp.concatenate(
+                    [
                         t * jnp.ones((xy_data[0].shape[0], 1)),
-                        jnp.stack(
-                            [xy_data[0][..., None], xy_data[1][..., None]], axis=1
+                        jnp.concatenate(
+                            [xy_data[0][..., None], xy_data[1][..., None]], axis=-1
                         ),
-                    )[0]
+                    ],
+                    axis=-1,
                 )
+                values_grid = jnp.squeeze(fun(t_x)[0]).T
                 t_slice, _ = _plot_2D_statio(
                     values_grid,
                     mesh,
-                    plot=False,
+                    plot=False,  # only use to compute t_slice
                     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],
-                )
-            else:
-                im = ax.pcolormesh(mesh[0], mesh[1], t_slice, cmap=cmap)
+
+            im = ax.pcolormesh(
+                mesh[0],
+                mesh[1],
+                t_slice,
+                cmap=cmap,
+                vmin=vmin_vmax[0],
+                vmax=vmin_vmax[1],
+            )
             ax.set_title(f"t = {times[idx] * Tmax:.2f}")
             ax.cax.colorbar(im, format="%0.2f")
 
 
 def _plot_2D_statio(
-    v_fun,
+    v_fun: Callable | Float[Array, "(nx*ny)^2 1"],
     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,
+    vmin_vmax: tuple[float, float] = [None, None],
 ):
     """Function that plot the function u(x) with 2-D input x using pcolormesh()
 
@@ -200,11 +196,11 @@ def _plot_2D_statio(
     Parameters
     ----------
     v_fun :
-        a vmapped function over jnp.array of shape (*, 2)
+        a vmapped function over jnp.array of shape (*, 2) OR a precomputed array of function values with shape compatible with `mesh`.
     mesh :
         a tuple of size 2, containing the x and y meshgrid.
     plot : bool, optional
-        either show or return the plot, by default True
+        either displays the plot, or silently returns the grid of values `v_fun(mesh)`.
     colorbar : bool, optional
         add a colorbar, by default True
     cmap :
@@ -213,7 +209,7 @@ def _plot_2D_statio(
         By default (7, 7)
     spinn :
         True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
-    vmin_vmax: tuple, optional
+    vmin_vmax: list, optional
         The colorbar minimum and maximum value. Defaults None.
 
     Returns
@@ -223,26 +219,23 @@ def _plot_2D_statio(
     """
 
     x_grid, y_grid = mesh
-    if not spinn:
+    if callable(v_fun):
         values = v_fun(jnp.vstack([x_grid.flatten(), y_grid.flatten()]).T)
         values_grid = values.reshape(x_grid.shape)
-    elif spinn:
-        # in this case v_fun is directly the values :)
-        values_grid = v_fun.T
+    else:
+        values_grid = v_fun.reshape(x_grid.shape)
 
     if plot:
         fig = plt.figure(figsize=figsize)
-        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)
+        im = plt.pcolormesh(
+            x_grid,
+            y_grid,
+            values_grid,
+            cmap=cmap,
+            vmin=vmin_vmax[0],
+            vmax=vmin_vmax[1],
+        )
+
         if colorbar:
             fig.colorbar(im, format="%0.2f")
         # don't plt.show() because it is done in plot2d()
@@ -258,6 +251,7 @@ def plot1d_slice(
     title: str = "",
     figsize: tuple[int, int] = (10, 10),
     spinn: Bool = False,
+    ax=None,
 ):
     """Function for plotting time slices of a function :math:`f(t_i, x)` where
     `t_i` is time (1-D) and x is 1-D
@@ -275,29 +269,40 @@ def plot1d_slice(
         default 1
     title
         title of the plot, by default ""
-    spinn :
-        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
     figsize
         size of the figure, by default (10, 10)
+    spinn
+        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+    ax
+        A pre-defined `matplotlib.Axes` where you want to plot.
+
+    Returns
+    -------
+    ax
+        A `matplotlib.Axes` object
     """
     if time_slices is None:
         time_slices = jnp.array([0])
-    plt.figure(figsize=figsize)
+    if ax is None:
+        fig, ax = plt.subplots(figsize=figsize)
+
     for t in time_slices:
+        t_xdata = jnp.concatenate(
+            [t * jnp.ones((xdata.shape[0], 1)), xdata[:, None]], axis=1
+        )
         if not spinn:
             # fix t with partial : shape is (1,)
-            v_u_tfixed = vmap(partial(fun, t=t * jnp.ones((1,))), 0, 0)
+            v_u_tfixed = vmap(fun)
             # add an axis to xdata for the concatenate function in the neural net
-            values = v_u_tfixed(x=xdata[:, None])
+            values = v_u_tfixed(t_xdata)
         elif spinn:
-            values = jnp.squeeze(
-                fun(t * jnp.ones((xdata.shape[0], 1)), xdata[..., None])[0]
-            )
-        plt.plot(xdata, values, label=f"$t_i={t * Tmax:.2f}$")
-    plt.xlabel("x")
-    plt.ylabel(r"$u(t_i, x)$")
-    plt.legend()
-    plt.title(title)
+            values = jnp.squeeze(fun(t_xdata)[0])
+        ax.plot(xdata, values, label=f"$t_i={t * Tmax:.2f}$")
+    ax.set_xlabel("x")
+    ax.set_ylabel(r"$u(t_i, x)$")
+    ax.legend()
+    ax.set_title(title)
+    return ax
 
 
 def plot1d_image(
@@ -310,7 +315,7 @@ def plot1d_image(
     colorbar: Bool = True,
     cmap: str = "inferno",
     spinn: Bool = False,
-    vmin_vmax: tuple[float, float] = None,
+    vmin_vmax: tuple[float, float] = [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).
@@ -339,30 +344,38 @@ def plot1d_image(
         True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
     vmin_vmax:
         The colorbar minimum and maximum value. Defaults None.
+
+    Returns
+    -------
+    fig, ax
+        A `matplotlib` `Figure` and `Axes` objects with the figure.
     """
 
     mesh = jnp.meshgrid(times, xdata)  # cartesian product
     if not spinn:
         # the trick is to use _plot2Dstatio
-        v_fun = vmap(lambda tx: fun(t=tx[0, None], x=tx[1, None]), 0, 0)
+        v_fun = vmap(fun)  # lambda tx: fun(t=tx[0, None], x=tx[1, None]), 0, 0)
         t_grid, x_grid = mesh
         values_grid = v_fun(jnp.vstack([t_grid.flatten(), x_grid.flatten()]).T).reshape(
             t_grid.shape
         )
     elif spinn:
-        values_grid = jnp.squeeze(fun((times[..., None]), xdata[..., None]).T)
+        values_grid = jnp.squeeze(
+            fun(jnp.concatenate([times[..., None], xdata[..., None]], axis=-1))
+        ).T
+
     fig, ax = plt.subplots(1, 1, figsize=figsize)
-    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)
+    im = ax.pcolormesh(
+        mesh[0] * Tmax,
+        mesh[1],
+        values_grid,
+        cmap=cmap,
+        vmin=vmin_vmax[0],
+        vmax=vmin_vmax[1],
+    )
+
     if colorbar:
         fig.colorbar(im, format="%0.2f")
     ax.set_title(title)
+
+    return fig, ax