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

jinns/nn/_spinn_mlp.py

@@ -0,0 +1,196 @@
+"""
+Implements utility function to create Separable PINNs
+https://arxiv.org/abs/2211.08761
+"""
+
+from dataclasses import InitVar
+from typing import Callable, Literal, Self, Union, Any
+import jax
+import jax.numpy as jnp
+import equinox as eqx
+from jaxtyping import Key, Array, Float, PyTree
+
+from jinns.parameters._params import Params, ParamsDict
+from jinns.nn._mlp import MLP
+from jinns.nn._spinn import SPINN
+
+
+class SMLP(eqx.Module):
+    """
+    Construct a Separable MLP
+
+    Parameters
+    ----------
+    key : InitVar[Key]
+        A jax random key for the layer initializations.
+    d : int
+        The number of dimensions to treat separately, including time `t` if
+        used for non-stationnary equations.
+    eqx_list : InitVar[tuple[tuple[Callable, int, int] | Callable, ...]]
+        A tuple of tuples of successive equinox modules and activation functions to
+        describe the PINN architecture. The inner tuples must have the eqx module or
+        activation function as first item, other items represents arguments
+        that could be required (eg. the size of the layer).
+        The `key` argument need not be given.
+        Thus typical example is `eqx_list=
+        ((eqx.nn.Linear, 1, 20),
+            jax.nn.tanh,
+            (eqx.nn.Linear, 20, 20),
+            jax.nn.tanh,
+            (eqx.nn.Linear, 20, 20),
+            jax.nn.tanh,
+            (eqx.nn.Linear, 20, r * m)
+        )`.
+    """
+
+    key: InitVar[Key] = eqx.field(kw_only=True)
+    eqx_list: InitVar[tuple[tuple[Callable, int, int] | Callable, ...]] = eqx.field(
+        kw_only=True
+    )
+    d: int = eqx.field(static=True, kw_only=True)
+
+    separated_mlp: list[MLP] = eqx.field(init=False)
+
+    def __post_init__(self, key, eqx_list):
+        keys = jax.random.split(key, self.d)
+        self.separated_mlp = [
+            MLP(key=keys[d_], eqx_list=eqx_list) for d_ in range(self.d)
+        ]
+
+    def __call__(
+        self, inputs: Float[Array, "dim"] | Float[Array, "dim+1"]
+    ) -> Float[Array, "d embed_dim*output_dim"]:
+        outputs = []
+        for d in range(self.d):
+            x_i = inputs[d : d + 1]
+            outputs += [self.separated_mlp[d](x_i)]
+        return jnp.asarray(outputs)
+
+
+class SPINN_MLP(SPINN):
+    """
+    An implementable SPINN based on a MLP architecture
+    """
+
+    @classmethod
+    def create(
+        cls,
+        key: Key,
+        d: int,
+        r: int,
+        eqx_list: tuple[tuple[Callable, int, int] | Callable, ...],
+        eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+        m: int = 1,
+        filter_spec: PyTree[Union[bool, Callable[[Any], bool]]] = None,
+    ) -> tuple[Self, PyTree]:
+        """
+        Utility function to create a SPINN neural network with the equinox
+        library.
+
+        *Note* that a SPINN is not vmapped and expects the
+        same batch size for each of its input axis. It directly outputs a
+        solution of shape `(batchsize,) * d`. See the paper for more
+        details.
+
+        Parameters
+        ----------
+        key : Key
+            A JAX random key that will be used to initialize the network parameters
+        d : int
+            The number of dimensions to treat separately.
+        r : int
+            An integer. The dimension of the embedding.
+        eqx_list : tuple[tuple[Callable, int, int] | Callable, ...],
+            A tuple of tuples of successive equinox modules and activation functions to
+            describe the PINN architecture. The inner tuples must have the eqx module or
+            activation function as first item, other items represents arguments
+            that could be required (eg. the size of the layer).
+            The `key` argument need not be given.
+            Thus typical example is
+            `eqx_list=((eqx.nn.Linear, 1, 20),
+                jax.nn.tanh,
+                (eqx.nn.Linear, 20, 20),
+                jax.nn.tanh,
+                (eqx.nn.Linear, 20, 20),
+                jax.nn.tanh,
+                (eqx.nn.Linear, 20, r * m)
+            )`.
+        eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
+            A string with three possibilities.
+            "ODE": the PINN is called with one input `t`.
+            "statio_PDE": the PINN is called with one input `x`, `x`
+            can be high dimensional.
+            "nonstatio_PDE": the PINN is called with two inputs `t` and `x`, `x`
+            can be high dimensional.
+            **Note**: the input dimension as given in eqx_list has to match the sum
+            of the dimension of `t` + the dimension of `x`.
+        m : int
+            The output dimension of the neural network. According to
+            the SPINN article, a total embedding dimension of `r*m` is defined. We
+            then sum groups of `r` embedding dimensions to compute each output.
+            Default is 1.
+        filter_spec : PyTree[Union[bool, Callable[[Any], bool]]]
+            Default is None which leads to `eqx.is_inexact_array` in the class
+            instanciation. This tells Jinns what to consider as
+            a trainable parameter. Quoting from equinox documentation:
+            a PyTree whose structure should be a prefix of the structure of pytree.
+            Each of its leaves should either be 1) True, in which case the leaf or
+            subtree is kept; 2) False, in which case the leaf or subtree is
+            replaced with replace; 3) a callable Leaf -> bool, in which case this is evaluated on the leaf or mapped over the subtree, and the leaf kept or replaced as appropriate.
+
+
+
+
+        Returns
+        -------
+        spinn
+            An instanciated SPINN
+        spinn.init_params
+            The initial set of parameters of the model
+
+        Raises
+        ------
+        RuntimeError
+            If the parameter value for eq_type is not in `["ODE", "statio_PDE",
+            "nonstatio_PDE"]` and for various failing checks
+        """
+
+        if eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
+            raise RuntimeError("Wrong parameter value for eq_type")
+
+        try:
+            nb_inputs_declared = eqx_list[0][
+                1
+            ]  # normally we look for 2nd ele of 1st layer
+        except IndexError:
+            nb_inputs_declared = eqx_list[1][
+                1
+            ]  # but we can have, eg, a flatten first layer
+        if nb_inputs_declared != 1:
+            raise ValueError("Input dim must be set to 1 in SPINN!")
+
+        try:
+            nb_outputs_declared = eqx_list[-1][2]  # normally we look for 3rd ele of
+            # last layer
+        except IndexError:
+            nb_outputs_declared = eqx_list[-2][2]
+            # but we can have, eg, a `jnp.exp` last layer
+        if nb_outputs_declared != r * m:
+            raise ValueError("Output dim must be set to r * m in SPINN!")
+
+        if d > 24:
+            raise ValueError(
+                "Too many dimensions, not enough letters available in jnp.einsum"
+            )
+
+        smlp = SMLP(key=key, d=d, eqx_list=eqx_list)
+        spinn = cls(
+            eqx_spinn_network=smlp,
+            d=d,
+            r=r,
+            eq_type=eq_type,
+            m=m,
+            filter_spec=filter_spec,
+        )
+
+        return spinn, spinn.init_params

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 :
@@ -212,8 +208,8 @@ def _plot_2D_statio(
     figsize :
         By default (7, 7)
     spinn :
-        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
-    vmin_vmax: tuple, optional
+        True if a SPINN is to be plotted. False for PINNs and HyperPINNs
+    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).
@@ -336,33 +341,41 @@ def plot1d_image(
     cmap :
         the matplotlib color map used in the ImageGrid.
     spinn :
-        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+        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