jinns 0.8.7__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-0.8.7.dist-info → jinns-0.8.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.7
+Version: 0.8.8
 Summary: Physics Informed Neural Network with JAX
 Author-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
 Maintainer-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-0.8.7.dist-info → jinns-0.8.8.dist-info}/RECORD

@@ -1,9 +1,11 @@
-jinns/__init__.py,sha256=Nw5pdlmDhJwco3bXX3YttkeCF8czX_6m0poh8vu0lDQ,113
+jinns/__init__.py,sha256=T2XlmLbYqcXTumPJL00cJ80W98We5LH8Yg_Lss_exl4,139
 jinns/data/_DataGenerators.py,sha256=N4-U4z3MG46UIzHCbKScv9Z7AN40w1wlLY_VsVNj2sI,62293
 jinns/data/__init__.py,sha256=yBOmoavSD-cABp4XcjQY1zsEVO0mDyIhi2MJ5WNp0l8,326
-jinns/data/_display.py,sha256=6renz4H7kHktutmLY7HM6PmxYH7cBfGHpC7GQa1Fnlk,7778
-jinns/experimental/__init__.py,sha256=3jCIy2R2i_0Erwxg-HwISdH79Nt1XCXhS9yY1F5awiY,208
+jinns/data/_display.py,sha256=vlqggDCgVMEwdGBtjVmZaTQORU6imSfDkssn2XCtITI,10392
+jinns/experimental/__init__.py,sha256=qWbhC7Z8UgLWy0t-zU7RYze6v13-FngiCYXu-2bRVFQ,296
 jinns/experimental/_diffrax_solver.py,sha256=sLT22byqh-6015_fhe1xtMWlFOYcCjzYKET4sLhA9R4,6818
+jinns/experimental/_sinuspinn.py,sha256=hxSzscwMV2LayWOqenIlT1zqEVVrE5Y8CKf7bHX5XFQ,5016
+jinns/experimental/_spectralpinn.py,sha256=-4795pa7AYtRNSE-ugan3gHh64mtu2VdrRG5AS_J9Eg,2654
 jinns/loss/_DynamicLoss.py,sha256=L4CVmmF0rTPbHntgqsLLHlnrlQgLHsetUocpJm7ZYag,27461
 jinns/loss/_DynamicLossAbstract.py,sha256=kTQlhLx7SBuH5dIDmYaE79sVHUZt1nUFa8LxPU5IHhM,8504
 jinns/loss/_LossODE.py,sha256=b9doBHoQwYvlgpqzrNO4dOaTN87LRvjHtHbz9bMoH7E,22119
@@ -27,8 +29,8 @@ jinns/utils/_utils.py,sha256=8dgvWXX9NT7_7-zltWp0C9tG45ZFNwXxueyxPBb4hjo,6740
 jinns/utils/_utils_uspinn.py,sha256=qcKcOw3zrwWSQyGVj6fD8c9GinHt_U6JWN_k0auTtXM,26039
 jinns/validation/__init__.py,sha256=Jv58mzgC3F7cRfXA6caicL1t_U0UAhbwLrmMNVg6E7s,66
 jinns/validation/_validation.py,sha256=KfetbzB0xTNdBcYLwFWjEtP63Tf9wJirlhgqLTJDyy4,6761
-jinns-0.8.7.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.8.7.dist-info/METADATA,sha256=L0P7JvMGKrJHx9OjrtFsmNKEwdKA_RlufAbOBf5l10I,2482
-jinns-0.8.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jinns-0.8.7.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.8.7.dist-info/RECORD,,
+jinns-0.8.8.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.8.8.dist-info/METADATA,sha256=oTs2EJMu4Bwo2n9DLsAPSU5edpbgPtwhNXBuW8YjpOc,2482
+jinns-0.8.8.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jinns-0.8.8.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.8.8.dist-info/RECORD,,