jinns 0.8.6__py3-none-any.whl → 0.8.8__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/_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

jinns/experimental/_sinuspinn.py

@@ -0,0 +1,135 @@
+import jax
+import equinox as eqx
+import jax.numpy as jnp
+from jinns.utils._pinn import PINN
+
+
+def almost_zero_init(weight: jax.Array, key: jax.random.PRNGKey) -> jax.Array:
+    out, in_ = weight.shape
+    stddev = 1e-2
+    return stddev * jax.random.normal(key, shape=(out, in_))
+
+
+class _SinusPINN(eqx.Module):
+    """
+    A specific PINN whose layers are x_sin2x functions whose frequencies are
+    determined by an other network
+    """
+
+    layers_pinn: list
+    layers_aux_nn: list
+
+    def __init__(self, key, list_layers_pinn, list_layers_aux_nn):
+        """
+        Parameters
+        ----------
+        key
+            A jax random key
+        list_layers_pinn
+            A list as eqx_list in jinns' PINN utility for the main PINN
+        list_layers_aux_nn
+            A list as eqx_list in jinns' PINN utility for the network which outputs
+            the PINN's activation frequencies
+        """
+        self.layers_pinn = []
+        for l in list_layers_pinn:
+            if len(l) == 1:
+                self.layers_pinn.append(l[0])
+            else:
+                key, subkey = jax.random.split(key, 2)
+                self.layers_pinn.append(l[0](*l[1:], key=subkey))
+        self.layers_aux_nn = []
+        for idx, l in enumerate(list_layers_aux_nn):
+            if len(l) == 1:
+                self.layers_aux_nn.append(l[0])
+            else:
+                key, subkey = jax.random.split(key, 2)
+                linear_layer = l[0](*l[1:], key=subkey)
+                key, subkey = jax.random.split(key, 2)
+                linear_layer = eqx.tree_at(
+                    lambda l: l.weight,
+                    linear_layer,
+                    almost_zero_init(linear_layer.weight, subkey),
+                )
+                if (idx == len(list_layers_aux_nn) - 1) or (
+                    idx == len(list_layers_aux_nn) - 2
+                ):
+                    # for the last layer: almost 0 weights and 0.5 bias
+                    linear_layer = eqx.tree_at(
+                        lambda l: l.bias,
+                        linear_layer,
+                        0.5 * jnp.ones(linear_layer.bias.shape),
+                    )
+                else:
+                    # for the other previous layers:
+                    # almost 0 weight and 0 bias
+                    linear_layer = eqx.tree_at(
+                        lambda l: l.bias,
+                        linear_layer,
+                        jnp.zeros(linear_layer.bias.shape),
+                    )
+                self.layers_aux_nn.append(linear_layer)
+
+                ## init to zero the frequency network except last biases
+                # key, subkey = jax.random.split(key, 2)
+                # _pinn = init_linear_weight(_pinn, almost_zero_init, subkey)
+                # key, subkey = jax.random.split(key, 2)
+                # _pinn = init_linear_bias(_pinn, zero_init, subkey)
+                # print(_pinn)
+                # print(jax.tree_util.tree_leaves(_pinn, is_leaf=lambda
+                #    p:not isinstance(p,eqx.nn.Linear))[0].layers_aux_nn[-1].bias)
+                # _pinn = eqx.tree_at(lambda p:_pinn.layers_aux_nn[-1].bias, 0.5 *
+                #        jnp.ones(_pinn.layers_aux_nn[-1].bias.shape))
+                #        #, is_leaf=lambda
+                #        #p:not isinstance(p, eqx.nn.Linear))
+
+    def __call__(self, x):
+        x_ = x.copy()
+        # forward pass in the network which determines the freq
+        for layer in self.layers_aux_nn:
+            x_ = layer(x_)
+        freq_list = jnp.clip(jnp.square(x_), a_min=1e-4, a_max=5)
+        x_ = x.copy()
+        # forward pass through the actual PINN
+        for idx, layer in enumerate(self.layers_pinn):
+            if idx % 2 == 0:
+                # Currently: every two layer we have an activation
+                # requiring a frequency
+                x_ = layer(x_)
+            else:
+                x_ = layer(x_, freq_list[(idx - 1) // 2])
+        return x_
+
+
+class sinusPINN(PINN):
+    """
+    MUST inherit from PINN to pass all the checks
+
+    HOWEVER we dot not bother with reimplementing anything
+    """
+
+    def __init__(self, key, list_layers_pinn, list_layers_aux_nn):
+        super().__init__({}, jnp.s_[...], "statio_PDE", None, None, None)
+        key, subkey = jax.random.split(key, 2)
+        _pinn = _SinusPINN(subkey, list_layers_pinn, list_layers_aux_nn)
+
+        self.params, self.static = eqx.partition(_pinn, eqx.is_inexact_array)
+
+    def init_params(self):
+        return self.params
+
+    def __call__(self, x, params):
+        try:
+            model = eqx.combine(params["nn_params"], self.static)
+        except (KeyError, TypeError) as e:  # give more flexibility
+            model = eqx.combine(params, self.static)
+        res = model(x)
+        if not res.shape:
+            return jnp.expand_dims(res, axis=-1)
+        return res
+
+
+def create_sinusPINN(key, list_layers_pinn, list_layers_aux_nn):
+    """ """
+    u = sinusPINN(key, list_layers_pinn, list_layers_aux_nn)
+    return u

jinns/experimental/_spectralpinn.py

@@ -0,0 +1,87 @@
+import jax
+import equinox as eqx
+import jax.numpy as jnp
+from jinns.utils._pinn import PINN
+
+
+def almost_zero_init(weight: jax.Array, key: jax.random.PRNGKey) -> jax.Array:
+    out, in_ = weight.shape
+    stddev = 1e-2
+    return stddev * jax.random.normal(key, shape=(out, in_))
+
+
+class _SpectralPINN(eqx.Module):
+    """
+    A specific PINN whose acrhitecture is similar to spectral method for simulation of a spatial field
+    (Chilès and Delfiner, 2012) - a single layer with cos() activation function and sum for last layer
+    """
+
+    layers_pinn: list
+    nbands: int
+
+    def __init__(self, key, list_layers_pinn, nbands):
+        """
+        Parameters
+        ----------
+        key
+            A jax random key
+        list_layers_pinn
+            A list as eqx_list in jinns' PINN utility for the main PINN
+        nbands
+            Number of spectral bands (i.e., neurones in the single layer of the PINN)
+        """
+        self.nbands = nbands
+        self.layers_pinn = []
+        for l in list_layers_pinn:
+            if len(l) == 1:
+                self.layers_pinn.append(l[0])
+            else:
+                key, subkey = jax.random.split(key, 2)
+                self.layers_pinn.append(l[0](*l[1:], key=subkey))
+
+    def __call__(self, x):
+        # forward pass through the actual PINN
+        for layer in self.layers_pinn:
+            x = layer(x)
+
+        return jnp.sqrt(2 / self.nbands) * jnp.sum(x)
+
+
+class spectralPINN(PINN):
+    """
+    MUST inherit from PINN to pass all the checks
+
+    HOWEVER we dot not bother with reimplementing anything
+    """
+
+    def __init__(self, key, list_layers_pinn, nbands):
+        super().__init__({}, jnp.s_[...], "statio_PDE", None, None, None)
+        key, subkey = jax.random.split(key, 2)
+        _pinn = _SpectralPINN(subkey, list_layers_pinn, nbands)
+
+        self.params, self.static = eqx.partition(_pinn, eqx.is_inexact_array)
+
+    def init_params(self):
+        return self.params
+
+    def __call__(self, x, params):
+        try:
+            model = eqx.combine(params["nn_params"], self.static)
+        except (KeyError, TypeError) as e:  # give more flexibility
+            model = eqx.combine(params, self.static)
+        # model = eqx.tree_at(lambda m:
+        #        m.layers_pinn[0].bias,
+        #        model,
+        #        model.layers_pinn[0].bias % (2 *
+        #            jnp.pi)
+        #        )
+        res = model(x)
+        if not res.shape:
+            return jnp.expand_dims(res, axis=-1)
+        return res
+
+
+def create_spectralPINN(key, list_layers_pinn, nbands):
+    """ """
+    u = spectralPINN(key, list_layers_pinn, nbands)
+    return u

jinns/loss/_LossODE.py

@@ -19,6 +19,8 @@ from jinns.loss._Losses import (
 )
 from jinns.utils._pinn import PINN
 
+_LOSS_WEIGHT_KEYS_ODE = ["observations", "dyn_loss", "initial_condition"]
+
 
 @register_pytree_node_class
 class LossODE:
@@ -128,6 +130,10 @@ class LossODE:
         if self.obs_slice is None:
             self.obs_slice = jnp.s_[...]
 
+        for k in _LOSS_WEIGHT_KEYS_ODE:
+            if k not in self.loss_weights.keys():
+                self.loss_weights[k] = 0
+
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
 

jinns/loss/_LossPDE.py

@@ -31,6 +31,16 @@ _IMPLEMENTED_BOUNDARY_CONDITIONS = [
     "vonneumann",
 ]
 
+_LOSS_WEIGHT_KEYS_PDESTATIO = [
+    "sobolev",
+    "observations",
+    "norm_loss",
+    "boundary_loss",
+    "dyn_loss",
+]
+
+_LOSS_WEIGHT_KEYS_PDENONSTATIO = _LOSS_WEIGHT_KEYS_PDESTATIO + ["initial_condition"]
+
 
 @register_pytree_node_class
 class LossPDEAbstract:
@@ -269,8 +279,8 @@ class LossPDEStatio(LossPDEAbstract):
             the PINN object
         loss_weights
             a dictionary with values used to ponderate each term in the loss
-            function. Valid keys are `dyn_loss`, `norm_loss`, `boundary_loss`
-            and `observations`.
+            function. Valid keys are `dyn_loss`, `norm_loss`, `boundary_loss`,
+            `observations` and `sobolev`.
             Note that we can have jnp.arrays with the same dimension of
             `u` which then ponderates each output of `u`
         dynamic_loss
@@ -441,18 +451,12 @@ class LossPDEStatio(LossPDEAbstract):
             )  # we return a function, that way
             # the order of sobolev_m is static and the conditional in the recursive
             # function is properly set
-            self.sobolev_m = self.sobolev_m
         else:
             self.sobolev_reg = None
 
-        if self.normalization_loss is None:
-            self.loss_weights["norm_loss"] = 0
-
-        if self.omega_boundary_fun is None:
-            self.loss_weights["boundary_loss"] = 0
-
-        if self.sobolev_reg is None:
-            self.loss_weights["sobolev"] = 0
+        for k in _LOSS_WEIGHT_KEYS_PDESTATIO:
+            if k not in self.loss_weights.keys():
+                self.loss_weights[k] = 0
 
         if (
             isinstance(self.omega_boundary_fun, dict)
@@ -533,7 +537,6 @@ class LossPDEStatio(LossPDEAbstract):
             )
         else:
             mse_norm_loss = jnp.array(0.0)
-            self.loss_weights["norm_loss"] = 0
 
         # boundary part
         params_ = _set_derivatives(params, "boundary_loss", self.derivative_keys)
@@ -567,7 +570,6 @@ class LossPDEStatio(LossPDEAbstract):
             )
         else:
             mse_observation_loss = jnp.array(0.0)
-            self.loss_weights["observations"] = 0
 
         # Sobolev regularization
         params_ = _set_derivatives(params, "sobolev", self.derivative_keys)
@@ -582,7 +584,6 @@ class LossPDEStatio(LossPDEAbstract):
             )
         else:
             mse_sobolev_loss = jnp.array(0.0)
-            self.loss_weights["sobolev"] = 0
 
         # total loss
         total_loss = (
@@ -785,8 +786,9 @@ class LossPDENonStatio(LossPDEStatio):
         else:
             self.sobolev_reg = None
 
-        if self.sobolev_reg is None:
-            self.loss_weights["sobolev"] = 0
+        for k in _LOSS_WEIGHT_KEYS_PDENONSTATIO:
+            if k not in self.loss_weights.keys():
+                self.loss_weights[k] = 0
 
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
@@ -924,7 +926,6 @@ class LossPDENonStatio(LossPDEStatio):
             )
         else:
             mse_observation_loss = jnp.array(0.0)
-            self.loss_weights["observations"] = 0
 
         # Sobolev regularization
         params_ = _set_derivatives(params, "sobolev", self.derivative_keys)
@@ -939,7 +940,6 @@ class LossPDENonStatio(LossPDEStatio):
             )
         else:
             mse_sobolev_loss = jnp.array(0.0)
-            self.loss_weights["sobolev"] = 0.0
 
         # total loss
         total_loss = (