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

jinns/validation/_validation.py

@@ -2,27 +2,24 @@
 Implements some validation functions and their associated hyperparameter
 """
 
-import copy
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
+
 import abc
-from typing import Union
+from typing import TYPE_CHECKING, Union
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Array, Bool, PyTree, Int
-import jinns
-import jinns.data
-from jinns.loss import LossODE, LossPDENonStatio, LossPDEStatio
+from jaxtyping import Array
+
 from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-    DataGeneratorParameter,
-    DataGeneratorObservations,
-    DataGeneratorObservationsMultiPINNs,
     append_obs_batch,
     append_param_batch,
 )
-import jinns.loss
+
+if TYPE_CHECKING:
+    from jinns.utils._types import *
 
 # Using eqx Module for the DataClass + Pytree inheritance
 # Abstract class and abstract/final pattern is used
@@ -35,14 +32,16 @@ class AbstractValidationModule(eqx.Module):
     2. implement a ``__call__`` returning ``(AbstractValidationModule, Bool, Array)``
     """
 
-    call_every: eqx.AbstractVar[Int]  # Mandatory for all validation step,
+    call_every: eqx.AbstractVar[int] = eqx.field(
+        kw_only=True
+    )  # Mandatory for all validation step,
     # it tells that the validation step is performed every call_every
     # iterations.
 
     @abc.abstractmethod
     def __call__(
-        self, params: PyTree
-    ) -> tuple["AbstractValidationModule", Bool, Array, Bool]:
+        self, params: Params | ParamsDict
+    ) -> tuple["AbstractValidationModule", bool, Array, bool]:
         raise NotImplementedError
 
 
@@ -53,37 +52,44 @@ class ValidationLoss(AbstractValidationModule):
     for more complicated validation strategy.
     """
 
-    loss: Union[callable, LossODE, LossPDEStatio, LossPDENonStatio] = eqx.field(
-        converter=copy.deepcopy
+    loss: AnyLoss = eqx.field(kw_only=True)  # NOTE that
+    # there used to be a deepcopy here which has been suppressed. 1) No need
+    # because loss are now eqx.Module (immutable) so no risk of in-place
+    # modification. 2) deepcopy is buggy with equinox, InitVar etc. (see issue
+    # #857 on equinox github)
+    validation_data: Union[AnyDataGenerator] = eqx.field(kw_only=True)
+    validation_param_data: Union[DataGeneratorParameter, None] = eqx.field(
+        kw_only=True, default=None
     )
-    validation_data: Union[DataGeneratorODE, CubicMeshPDEStatio, CubicMeshPDENonStatio]
-    validation_param_data: Union[DataGeneratorParameter, None] = None
     validation_obs_data: Union[
         DataGeneratorObservations, DataGeneratorObservationsMultiPINNs, None
-    ] = None
-    call_every: Int = 250  # concrete typing
-    early_stopping: Bool = True  # globally control if early stopping happens
+    ] = eqx.field(kw_only=True, default=None)
+    call_every: int = eqx.field(kw_only=True, default=250)  # concrete typing
+    early_stopping: bool = eqx.field(
+        kw_only=True, default=True
+    )  # globally control if early stopping happens
 
-    patience: Union[Int] = 10
+    patience: Union[int] = eqx.field(kw_only=True, default=10)
     best_val_loss: Array = eqx.field(
-        converter=jnp.asarray, default_factory=lambda: jnp.array(jnp.inf)
+        converter=jnp.asarray, default_factory=lambda: jnp.array(jnp.inf), kw_only=True
     )
 
     counter: Array = eqx.field(
-        converter=jnp.asarray, default_factory=lambda: jnp.array(0.0)
+        converter=jnp.asarray, default_factory=lambda: jnp.array(0.0), kw_only=True
     )
 
-    def __call__(self, params) -> tuple["ValidationLoss", Bool, Array]:
+    def __call__(
+        self, params: AnyParams
+    ) -> tuple["ValidationLoss", bool, float, AnyParams]:
         # do in-place mutation
-        val_batch = self.validation_data.get_batch()
+
+        validation_data, val_batch = self.validation_data.get_batch()
         if self.validation_param_data is not None:
-            val_batch = append_param_batch(
-                val_batch, self.validation_param_data.get_batch()
-            )
+            validation_param_data, param_batch = self.validation_param_data.get_batch()
+            val_batch = append_param_batch(val_batch, param_batch)
         if self.validation_obs_data is not None:
-            val_batch = append_obs_batch(
-                val_batch, self.validation_obs_data.get_batch()
-            )
+            validation_obs_data, obs_batch = self.validation_obs_data.get_batch()
+            val_batch = append_obs_batch(val_batch, obs_batch)
 
         validation_loss_value, _ = self.loss(params, val_batch)
         (counter, best_val_loss, update_best_params) = jax.lax.cond(
@@ -93,9 +99,14 @@ class ValidationLoss(AbstractValidationModule):
             (self.counter, self.best_val_loss),
         )
 
-        # use eqx.tree_at to update attributes
-        # (https://github.com/patrick-kidger/equinox/issues/396)
-        new = eqx.tree_at(lambda t: t.counter, self, counter)
+        new = eqx.tree_at(lambda t: t.validation_data, self, validation_data)
+        if self.validation_param_data is not None:
+            new = eqx.tree_at(
+                lambda t: t.validation_param_data, new, validation_param_data
+            )
+        if self.validation_obs_data is not None:
+            new = eqx.tree_at(lambda t: t.validation_obs_data, new, validation_obs_data)
+        new = eqx.tree_at(lambda t: t.counter, new, counter)
         new = eqx.tree_at(lambda t: t.best_val_loss, new, best_val_loss)
 
         bool_early_stopping = jax.lax.cond(
@@ -109,106 +120,3 @@ class ValidationLoss(AbstractValidationModule):
         )
         # return `new` cause no in-place modification of the eqx.Module
         return (new, bool_early_stopping, validation_loss_value, update_best_params)
-
-
-if __name__ == "__main__":
-    import jax
-    import jax.numpy as jnp
-    import jax.random as random
-    from jinns.loss import BurgerEquation
-
-    key = random.PRNGKey(1)
-    key, subkey = random.split(key)
-
-    n = 50
-    nb = 2 * 2 * 10
-    nt = 10
-    omega_batch_size = 10
-    omega_border_batch_size = 10
-    temporal_batch_size = 4
-    dim = 1
-    xmin = 0
-    xmax = 1
-    tmin, tmax = 0, 1
-    method = "uniform"
-
-    val_data = jinns.data.CubicMeshPDENonStatio(
-        subkey,
-        n,
-        nb,
-        nt,
-        omega_batch_size,
-        omega_border_batch_size,
-        temporal_batch_size,
-        dim,
-        (xmin,),
-        (xmax,),
-        tmin,
-        tmax,
-        method,
-    )
-
-    eqx_list = [
-        [eqx.nn.Linear, 2, 50],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 50, 50],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 50, 50],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 50, 50],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 50, 50],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 50, 2],
-    ]
-
-    key, subkey = random.split(key)
-    u = jinns.utils.create_PINN(
-        subkey, eqx_list, "nonstatio_PDE", 2, slice_solution=jnp.s_[:1]
-    )
-    init_nn_params = u.init_params()
-
-    dyn_loss = BurgerEquation()
-    loss_weights = {"dyn_loss": 1, "boundary_loss": 10, "observations": 10}
-
-    key, subkey = random.split(key)
-    loss = jinns.loss.LossPDENonStatio(
-        u=u,
-        loss_weights=loss_weights,
-        dynamic_loss=dyn_loss,
-        norm_key=subkey,
-        norm_borders=(-1, 1),
-    )
-    print(id(loss))
-    validation = ValidationLoss(
-        call_every=250,
-        early_stopping=True,
-        patience=1000,
-        loss=loss,
-        validation_data=val_data,
-        validation_param_data=None,
-    )
-    print(id(validation.loss) is not id(loss))  # should be True (deepcopy)
-
-    init_params = {"nn_params": init_nn_params, "eq_params": {"nu": 1.0}}
-
-    print(validation.loss is loss)
-    loss.evaluate(init_params, val_data.get_batch())
-    print(loss.norm_key)
-    print("Call validation once")
-    validation, _, _ = validation(init_params)
-    print(validation.loss is loss)
-    print(validation.loss.norm_key == loss.norm_key)
-    print("Crate new pytree from validation and call it once")
-    new_val = eqx.tree_at(lambda t: t.counter, validation, jnp.array(3.0))
-    print(validation.loss is new_val.loss)  # FALSE
-    # test if attribute have been modified
-    new_val, _, _ = new_val(init_params)
-    print(f"{new_val.loss is loss=}")
-    print(f"{loss.norm_key=}")
-    print(f"{validation.loss.norm_key=}")
-    print(f"{new_val.loss.norm_key=}")
-    print(f"{new_val.loss.norm_key == loss.norm_key=}")
-    print(f"{new_val.loss.norm_key == validation.loss.norm_key=}")
-    print(new_val.counter)
-    print(validation.counter)

jinns-1.1.0.dist-info/AUTHORS

@@ -0,0 +1,2 @@
+Hugo Gangloff <hugo.gangloff@inrae.fr>
+Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-0.9.0.dist-info → jinns-1.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.9.0
+Version: 1.1.0
 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>
@@ -10,20 +10,21 @@ Project-URL: Documentation, https://mia_jinns.gitlab.io/jinns/index.html
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Development Status :: 4 - Beta
 Classifier: Programming Language :: Python
-Requires-Python: >=3.7
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
+License-File: AUTHORS
 Requires-Dist: numpy
 Requires-Dist: jax
 Requires-Dist: jaxopt
 Requires-Dist: optax
-Requires-Dist: equinox
+Requires-Dist: equinox >0.11.3
 Requires-Dist: jax-tqdm
 Requires-Dist: diffrax
 Requires-Dist: matplotlib
 Provides-Extra: notebook
 Requires-Dist: jupyter ; extra == 'notebook'
-Requires-Dist: matplotlib ; extra == 'notebook'
+Requires-Dist: seaborn ; extra == 'notebook'
 
 jinns
 =====

jinns-1.1.0.dist-info/RECORD

@@ -0,0 +1,39 @@
+jinns/__init__.py,sha256=5p7V5VJd7PXEINqhqS4mUsnQtXlyPwfctRhL4p0loFg,181
+jinns/data/_Batchs.py,sha256=BLxTDiFb6o9M6Irc2_HKmpr8IgA159u_kJIbCBZ490E,926
+jinns/data/_DataGenerators.py,sha256=TuQDPI8NGx4WnCGtUp8v7o7GcnDfVZvL5E34JI-9Lmw,58455
+jinns/data/__init__.py,sha256=TRCH0Z4-SQZ50MbSf46CUYWBkWVDmXCyez9T-EGiv_8,338
+jinns/experimental/__init__.py,sha256=3jCIy2R2i_0Erwxg-HwISdH79Nt1XCXhS9yY1F5awiY,208
+jinns/experimental/_diffrax_solver.py,sha256=upMr3kTTNrxEiSUO_oLvCXcjS9lPxSjvbB81h3qlhaU,6813
+jinns/loss/_DynamicLoss.py,sha256=WGbAuWnNfsbzUlWEiW_ARd4kI3jmHwdqPjxLC-wCA6s,25753
+jinns/loss/_DynamicLossAbstract.py,sha256=Xyt28Oej_zlhcV3f6cw2vnAKyRJhXBiA63CsdL3PihU,13767
+jinns/loss/_LossODE.py,sha256=ThWPse6Gn5crM3_tzwZCBx-usoD0xWu6y1n0GVl2dpI,23422
+jinns/loss/_LossPDE.py,sha256=R9kNQiaFbFx2eCMdjB7ie7UJ9pJW7PmvHijioNgu-bs,49117
+jinns/loss/__init__.py,sha256=Fm4QAHaVmp0CA7HSwb7KUctwdXnNZ9v5KmTqpeoYPaE,669
+jinns/loss/_boundary_conditions.py,sha256=O0D8eWsFfvNNeO20PQ0rUKBI_MDqaBvqChfXaztZoL4,16679
+jinns/loss/_loss_utils.py,sha256=44J-VF6dxT_o5BcNWFOiLpY40c35YnAxxZkoNtdtcZc,13689
+jinns/loss/_loss_weights.py,sha256=F0Fgji2XpVk3pr9oIryGuXcG1FGQo4Dv6WFgze2BtA0,2201
+jinns/loss/_operators.py,sha256=o-Ljp_9_HXB9Mhm-ANh6ouNw4_PsqLJAha7dFDGl_nQ,10781
+jinns/parameters/__init__.py,sha256=1gxNLoAXUjhUzBWuh86YjU5pYy8SOboCs8TrKcU1wZc,158
+jinns/parameters/_derivative_keys.py,sha256=UyEcgfNF1vwPcGWD2ShAZkZiq4thzRDm_OUJzOfjjiY,21909
+jinns/parameters/_params.py,sha256=wK9ZSqoL9KnjOWqc_ZhJ09ffbsgeUEcttc1Rhme0lLk,3550
+jinns/plot/__init__.py,sha256=Q279h5veYWNLQyttsC8_tDOToqUHh8WaRON90CiWXqk,81
+jinns/plot/_plot.py,sha256=ZGIJdGwEd3NlHRTq_2sOfEH_CtOkvPwdgCMct-nQlJE,11691
+jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jinns/solver/_rar.py,sha256=BAr0qHNtCQ11XRhxRDed8JZPSWBm4E5FT3VOGQPrm4g,22728
+jinns/solver/_solve.py,sha256=iUkvZ9f-0bxO2eJlNv7wLk5oES69LNbDxzOgBAjbHTg,21106
+jinns/utils/__init__.py,sha256=CNxcb_AYzA2aeDsYwLfIuT1zy8NU7LElU160eHCj2oA,174
+jinns/utils/_containers.py,sha256=_PLIkJHY-jE3nigjuAiYE5USJr7rpXCDsNULRKZtqnU,1213
+jinns/utils/_hyperpinn.py,sha256=7fAy6Xv7CkIeIYCw7u9RB0gh4aITMfxz6JCmJ685KB8,16351
+jinns/utils/_pinn.py,sha256=eiw_i72D7-CxZXXOBpgOgDndwGd5sSnFZNk-Rg-6xy8,12977
+jinns/utils/_save_load.py,sha256=kVDeQrpPf7j7kYUy1gWGCr4_QyaBplRbPSitkWTQnQA,8574
+jinns/utils/_spinn.py,sha256=lesBjXwoj3UUzfecPdh1wBFCY0BlA-q0Wb1pDv0RVYA,9211
+jinns/utils/_types.py,sha256=P_dS0odrHbyalYJ0FjS6q0tkXAGr-4GArsiyJYrB1ho,1878
+jinns/utils/_utils.py,sha256=Ow8xB516E7yHDZatokVJHHFNPDu6fXr9-NmraUXjjyw,1819
+jinns/validation/__init__.py,sha256=Jv58mzgC3F7cRfXA6caicL1t_U0UAhbwLrmMNVg6E7s,66
+jinns/validation/_validation.py,sha256=bvqL2poTFJfn9lspWqMqXvQGcQIodKwKrC786QtEZ7A,4700
+jinns-1.1.0.dist-info/AUTHORS,sha256=7NwCj9nU-HNG1asvy4qhQ2w7oZHrn-Lk5_wK_Ve7a3M,80
+jinns-1.1.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-1.1.0.dist-info/METADATA,sha256=3Qk885oguf6S_WPHd9KCFVIWP21nJqX9zFWoS9ZI-T0,2536
+jinns-1.1.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jinns-1.1.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-1.1.0.dist-info/RECORD,,

{jinns-0.9.0.dist-info → jinns-1.1.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (70.2.0)
+Generator: setuptools (75.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

jinns/experimental/_sinuspinn.py

@@ -1,135 +0,0 @@
-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

@@ -1,87 +0,0 @@
-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