jinns 0.8.6__py3-none-any.whl → 0.8.8__py3-none-any.whl

jinns/validation/_validation.py

@@ -0,0 +1,214 @@
+"""
+Implements some validation functions and their associated hyperparameter
+"""
+
+import copy
+import abc
+from typing import 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 jinns.data._DataGenerators import (
+    DataGeneratorODE,
+    CubicMeshPDEStatio,
+    CubicMeshPDENonStatio,
+    DataGeneratorParameter,
+    DataGeneratorObservations,
+    DataGeneratorObservationsMultiPINNs,
+    append_obs_batch,
+    append_param_batch,
+)
+import jinns.loss
+
+# Using eqx Module for the DataClass + Pytree inheritance
+# Abstract class and abstract/final pattern is used
+# see : https://docs.kidger.site/equinox/pattern/
+
+
+class AbstractValidationModule(eqx.Module):
+    """Abstract class representing interface for any validation module. It must
+    1. have a ``call_every`` attribute.
+    2. implement a ``__call__`` returning ``(AbstractValidationModule, Bool, Array)``
+    """
+
+    call_every: eqx.AbstractVar[Int]  # 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]:
+        raise NotImplementedError
+
+
+class ValidationLoss(AbstractValidationModule):
+    """
+    Implementation of a vanilla validation module returning the PINN loss
+    on a validation set of collocation points. This can be used as a baseline
+    for more complicated validation strategy.
+    """
+
+    loss: Union[callable, LossODE, LossPDEStatio, LossPDENonStatio] = eqx.field(
+        converter=copy.deepcopy
+    )
+    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
+
+    patience: Union[Int] = 10
+    best_val_loss: Array = eqx.field(
+        converter=jnp.asarray, default_factory=lambda: jnp.array(jnp.inf)
+    )
+
+    counter: Array = eqx.field(
+        converter=jnp.asarray, default_factory=lambda: jnp.array(0.0)
+    )
+
+    def __call__(self, params) -> tuple["ValidationLoss", Bool, Array]:
+        # do in-place mutation
+        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()
+            )
+        if self.validation_obs_data is not None:
+            val_batch = append_obs_batch(
+                val_batch, self.validation_obs_data.get_batch()
+            )
+
+        validation_loss_value, _ = self.loss(params, val_batch)
+        (counter, best_val_loss) = jax.lax.cond(
+            validation_loss_value < self.best_val_loss,
+            lambda _: (jnp.array(0.0), validation_loss_value),  # reset
+            lambda operands: (operands[0] + 1, operands[1]),  # increment
+            (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.best_val_loss, new, best_val_loss)
+
+        bool_early_stopping = jax.lax.cond(
+            jnp.logical_and(
+                jnp.array(self.counter == self.patience),
+                jnp.array(self.early_stopping),
+            ),
+            lambda _: True,
+            lambda _: False,
+            None,
+        )
+        # return `new` cause no in-place modification of the eqx.Module
+        return (new, bool_early_stopping, validation_loss_value)
+
+
+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-0.8.6.dist-info → jinns-0.8.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.6
+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.6.dist-info → jinns-0.8.8.dist-info}/RECORD

@@ -1,13 +1,15 @@
-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=sxpgiDR6mfoREuc-qe0AkirOe5K_5oblaYCnodTNxoI,21912
-jinns/loss/_LossPDE.py,sha256=_yX3R-FrAScTn9_QfVC8PfDYRE4UQ5lnzITUYgNFitA,61766
+jinns/loss/_LossODE.py,sha256=b9doBHoQwYvlgpqzrNO4dOaTN87LRvjHtHbz9bMoH7E,22119
+jinns/loss/_LossPDE.py,sha256=purAEtc0e71kv9XnZUT-a7MrkDAkM_3tTI4xJPu6fH4,61629
 jinns/loss/_Losses.py,sha256=XOL3MFiKEd3ndsc78Qnpi1vbgR0B2HaAWOGGW2meDM8,11190
 jinns/loss/__init__.py,sha256=pFNYUxns-NPXBFdqrEVSiXkQLfCtKw-t2trlhvLzpYE,355
 jinns/loss/_boundary_conditions.py,sha256=YfSnLZ25hXqQ5KWAuxOrWSKkf_oBqAc9GQV4z7MjWyQ,17434
@@ -15,8 +17,9 @@ jinns/loss/_operators.py,sha256=zDGJqYqeYH7xd-4dtGX9PS-pf0uSOpUUXGo5SVjIJ4o,1106
 jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jinns/solver/_rar.py,sha256=K-0y1-ofOAo1n_Ea3QShSGCGKVYTwiaE_Bz9-DZMJm8,14525
 jinns/solver/_seq2seq.py,sha256=FL-42hTgmVl7O3hHh1ccFVw2bT8bW82hvlDRz971Chk,5620
-jinns/solver/_solve.py,sha256=r4jn6hx7_t-Y2rBWA2npUmWWnDg4iRbgYBHZDNn9tmY,13745
+jinns/solver/_solve.py,sha256=mGi0zaT_fK_QpBjTxof5Ix4mmfmnPi66CNJ3GQFZuo4,19099
 jinns/utils/__init__.py,sha256=44ms5UR6vMw3Nf6u4RCAzPFs4fom_YbBnH9mfne8m6k,313
+jinns/utils/_containers.py,sha256=eYD277fO7X4EfX7PUFCCl69r3JBfh1sCfq8LkL5gd6o,1495
 jinns/utils/_hyperpinn.py,sha256=93hbiATdp5W4l1cu9Oe6O2c45o-ZF_z2u6FzNLyjnm4,10878
 jinns/utils/_optim.py,sha256=550kxH75TL30o1iKx1swJyP0KqyUPsJ7-imL1w65Qd0,4444
 jinns/utils/_pinn.py,sha256=mhA4-3PazyQTbWIx9oLaNwL0QDe8ZIBhbiy5J3kwa4I,9471
@@ -24,8 +27,10 @@ jinns/utils/_save_load.py,sha256=qgZ23nUcB8-B5IZ2guuUWC4M7r5Lxd_Ms3staScdyJo,566
 jinns/utils/_spinn.py,sha256=SzOUt1KHtB9QOpghpvitnXN-KEqXUXbvabC5k0TnKEo,7793
 jinns/utils/_utils.py,sha256=8dgvWXX9NT7_7-zltWp0C9tG45ZFNwXxueyxPBb4hjo,6740
 jinns/utils/_utils_uspinn.py,sha256=qcKcOw3zrwWSQyGVj6fD8c9GinHt_U6JWN_k0auTtXM,26039
-jinns-0.8.6.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.8.6.dist-info/METADATA,sha256=3Ml6PCA-569v9-1FgyPDySX09RQas0zPOVEV_gqy9lk,2482
-jinns-0.8.6.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jinns-0.8.6.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.8.6.dist-info/RECORD,,
+jinns/validation/__init__.py,sha256=Jv58mzgC3F7cRfXA6caicL1t_U0UAhbwLrmMNVg6E7s,66
+jinns/validation/_validation.py,sha256=KfetbzB0xTNdBcYLwFWjEtP63Tf9wJirlhgqLTJDyy4,6761
+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,,