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

jinns/utils/_types.py

@@ -1,65 +1,31 @@
-# pragma: exclude file
 from __future__ import (
     annotations,
 )  # https://docs.python.org/3/library/typing.html#constant
 
-from typing import TypeAlias, TYPE_CHECKING, NewType
-from jaxtyping import Int
+from typing import TypeAlias, TYPE_CHECKING, Callable
+from jaxtyping import Float, Array
 
 if TYPE_CHECKING:
-    from jinns.loss._LossPDE import (
-        LossPDEStatio,
-        LossPDENonStatio,
-        SystemLossPDE,
+    from jinns.data._Batchs import ODEBatch, PDEStatioBatch, PDENonStatioBatch
+    from jinns.loss._LossODE import LossODE
+    from jinns.loss._LossPDE import LossPDEStatio, LossPDENonStatio
+    from jinns.loss._loss_components import (
+        ODEComponents,
+        PDEStatioComponents,
+        PDENonStatioComponents,
     )
 
-    from jinns.loss._LossODE import LossODE, SystemLossODE
-    from jinns.parameters._params import Params, ParamsDict
-    from jinns.data._DataGenerators import (
-        DataGeneratorODE,
-        CubicMeshPDEStatio,
-        CubicMeshPDENonStatio,
-        DataGeneratorObservations,
-        DataGeneratorParameter,
-        DataGeneratorObservationsMultiPINNs,
-    )
-
-    from jinns.loss import DynamicLoss
-    from jinns.data._Batchs import *
-    from jinns.nn._pinn import PINN
-    from jinns.nn._hyperpinn import HyperPINN
-    from jinns.nn._spinn_mlp import SPINN
-    from jinns.utils._containers import *
-    from jinns.validation._validation import AbstractValidationModule
-
-    AnyLoss: TypeAlias = (
-        LossPDEStatio | LossPDENonStatio | SystemLossPDE | LossODE | SystemLossODE
-    )
-
-    AnyParams: TypeAlias = Params | ParamsDict
+    # Here we define types available for the whole package
+    BoundaryConditionFun: TypeAlias = Callable[
+        [Float[Array, " dim"] | Float[Array, " dim + 1"]], Float[Array, " dim_solution"]
+    ]
 
-    AnyDataGenerator: TypeAlias = (
-        DataGeneratorODE | CubicMeshPDEStatio | CubicMeshPDENonStatio
-    )
-
-    AnyPINN: TypeAlias = PINN | HyperPINN | SPINN
-
-    AnyBatch: TypeAlias = ODEBatch | PDEStatioBatch | PDENonStatioBatch
-    rar_operands = NewType(
-        "rar_operands", tuple[AnyLoss, AnyParams, AnyDataGenerator, Int]
-    )
+    AnyBatch: TypeAlias = ODEBatch | PDENonStatioBatch | PDEStatioBatch
+    AnyLoss: TypeAlias = LossODE | LossPDEStatio | LossPDENonStatio
 
-    main_carry = NewType(
-        "main_carry",
-        tuple[
-            Int,
-            AnyLoss,
-            OptimizationContainer,
-            OptimizationExtraContainer,
-            DataGeneratorContainer,
-            AbstractValidationModule,
-            LossContainer,
-            StoredObjectContainer,
-            Float[Array, "n_iter"],
-        ],
+    # here we would like a type from 3.12
+    # (https://typing.python.org/en/latest/spec/aliases.html#type-statement) so
+    # that we could have a generic AnyLossComponents
+    AnyLossComponents: TypeAlias = (
+        ODEComponents | PDEStatioComponents | PDENonStatioComponents
     )

jinns/utils/_utils.py

@@ -2,20 +2,13 @@
 Implements various utility functions
 """
 
-from math import prod
 import warnings
 import jax
 import jax.numpy as jnp
-from jaxtyping import PyTree, Array
+from jaxtyping import PyTree, Array, Bool
 
-from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-)
 
-
-def _check_nan_in_pytree(pytree: PyTree) -> bool:
+def _check_nan_in_pytree(pytree: PyTree) -> Bool[Array, " "]:
     """
     Check if there is a NaN value anywhere is the pytree
 
@@ -55,7 +48,7 @@ def get_grid(in_array: Array) -> Array:
 
 
 def _check_shape_and_type(
-    r: Array | int, expected_shape: tuple, cause: str = "", binop: str = ""
+    r: Array | int | float, expected_shape: tuple, cause: str = "", binop: str = ""
 ) -> Array | float:
     """
     Ensures float type and correct shapes for broadcasting when performing a
@@ -90,7 +83,7 @@ def _check_shape_and_type(
 
 
 def _subtract_with_check(
-    a: Array | int, b: Array | int, cause: str = ""
+    a: Array | int | float, b: Array, cause: str = ""
 ) -> Array | float:
     a = _check_shape_and_type(a, b.shape, cause=cause, binop="-")
     return a - b

jinns/validation/__init__.py

@@ -1 +1,3 @@
 from ._validation import AbstractValidationModule, ValidationLoss
+
+__all__ = ["AbstractValidationModule", "ValidationLoss"]

jinns/validation/_validation.py

@@ -7,19 +7,23 @@ from __future__ import (
 )  # https://docs.python.org/3/library/typing.html#constant
 
 import abc
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Array
+from jaxtyping import Array, Float
 
-from jinns.data._DataGenerators import (
+from jinns.data._utils import (
     append_obs_batch,
     append_param_batch,
 )
 
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.data._DataGeneratorParameter import DataGeneratorParameter
+    from jinns.data._DataGeneratorObservations import DataGeneratorObservations
+    from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+    from jinns.parameters._params import Params
+    from jinns.loss._abstract_loss import AbstractLoss
 
 # Using eqx Module for the DataClass + Pytree inheritance
 # Abstract class and abstract/final pattern is used
@@ -40,8 +44,8 @@ class AbstractValidationModule(eqx.Module):
 
     @abc.abstractmethod
     def __call__(
-        self, params: Params | ParamsDict
-    ) -> tuple["AbstractValidationModule", bool, Array, bool]:
+        self, params: Params[Array]
+    ) -> tuple[AbstractValidationModule, bool, Array, Params[Array]]:
         raise NotImplementedError
 
 
@@ -52,24 +56,20 @@ class ValidationLoss(AbstractValidationModule):
     for more complicated validation strategy.
     """
 
-    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(
+    loss: AbstractLoss = eqx.field(kw_only=True)
+    validation_data: AbstractDataGenerator = eqx.field(kw_only=True)
+    validation_param_data: DataGeneratorParameter = eqx.field(
+        kw_only=True, default=None
+    )
+    validation_obs_data: DataGeneratorObservations | None = eqx.field(
         kw_only=True, default=None
     )
-    validation_obs_data: Union[
-        DataGeneratorObservations, DataGeneratorObservationsMultiPINNs, None
-    ] = 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] = eqx.field(kw_only=True, default=10)
+    patience: int = eqx.field(kw_only=True, default=10)
     best_val_loss: Array = eqx.field(
         converter=jnp.asarray, default_factory=lambda: jnp.array(jnp.inf), kw_only=True
     )
@@ -79,10 +79,11 @@ class ValidationLoss(AbstractValidationModule):
     )
 
     def __call__(
-        self, params: AnyParams
-    ) -> tuple["ValidationLoss", bool, float, AnyParams]:
+        self, params: Params[Array]
+    ) -> tuple[ValidationLoss, bool, Float[Array, " "], Params[Array]]:
         # do in-place mutation
 
+        # pylint / pyright complains below when using the self attributes see: https://github.com/patrick-kidger/equinox/issues/1013
         validation_data, val_batch = self.validation_data.get_batch()
         if self.validation_param_data is not None:
             validation_param_data, param_batch = self.validation_param_data.get_batch()

{jinns-1.3.0.dist-info → jinns-1.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: jinns
-Version: 1.3.0
+Version: 1.5.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>
@@ -25,6 +25,7 @@ Requires-Dist: matplotlib
 Provides-Extra: notebook
 Requires-Dist: jupyter; extra == "notebook"
 Requires-Dist: seaborn; extra == "notebook"
+Dynamic: license-file
 
 jinns
 =====
@@ -53,6 +54,9 @@ It can also be used for forward problems and hybrid-modeling.
 
     - [Hyper PINNs](https://arxiv.org/pdf/2111.01008.pdf): useful for meta-modeling
 
+- Other
+    - Adaptative Loss Weights are now implemented. Some SoftAdapt, LRAnnealing and ReLoBRaLo are available and users can implement their own strategy. See the [tutorial](https://mia_jinns.gitlab.io/jinns/Notebooks/Tutorials/implementing_your_own_PDE_problem/)
+
 
 - **Get started**: check out our various notebooks on the [documentation](https://mia_jinns.gitlab.io/jinns/index.html).
 
@@ -99,7 +103,7 @@ Here are the contributors guidelines:
 pip install -e .
 ```
 
-3. Install pre-commit and run it.
+3. Install pre-commit and run it. Our pre-commit hooks consist in `ruff format` and `ruff check`. You can install `ruff` simply by `pip install ruff`. We highly recommend you to check the code type hints with `pyright` even though we currently have no rule concerning type checking in the pipeline.
 
 ```bash
 pip install pre-commit
@@ -112,7 +116,7 @@ pre-commit install
 
 Don't hesitate to contribute and get your name on the list here !
 
-**List of contributors:** Hugo Gangloff, Nicolas Jouvin, Lucia Clarotto, Inass Soukarieh
+**List of contributors:** Hugo Gangloff, Nicolas Jouvin, Lucia Clarotto, Inass Soukarieh, Mohamed Badi
 
 # Cite us
 

jinns-1.5.0.dist-info/RECORD

@@ -0,0 +1,55 @@
+jinns/__init__.py,sha256=hyh3QKO2cQGK5cmvFYP0MrXb-tK_DM2T9CwLwO-sEX8,500
+jinns/data/_AbstractDataGenerator.py,sha256=O61TBOyeOFKwf1xqKzFD4KwCWRDnm2XgyJ-kKY9fmB4,557
+jinns/data/_Batchs.py,sha256=-DlD6Qag3zs5QbKtKAOvOzV7JOpNOqAm_P8cwo1dIZg,1574
+jinns/data/_CubicMeshPDENonStatio.py,sha256=c_8czJpxSoEvgZ8LDpL2sqtF9dcW4ELNO4juEFMOxog,16400
+jinns/data/_CubicMeshPDEStatio.py,sha256=stZ0Kbb7_VwFmWUSPs0P6a6qRj2Tu67p7sxEfb1Ajks,17865
+jinns/data/_DataGeneratorODE.py,sha256=5RzUbQFEsooAZsocDw4wRgA_w5lJmDMuY4M6u79K-1c,7260
+jinns/data/_DataGeneratorObservations.py,sha256=jknepLsJatSJHFq5lLMD-fFHkPGj5q286LEjE-vH24k,7738
+jinns/data/_DataGeneratorParameter.py,sha256=IedX3jcOj7ZDW_18IAcRR75KVzQzo85z9SICIKDBJl4,8539
+jinns/data/__init__.py,sha256=4b4eVsoGHV89m2kGDiAOHsrGialZQ6j5ja575qWwQHs,677
+jinns/data/_utils.py,sha256=XxaLIg_HIgcB7ACBIhTpHbCT1HXKcDaY1NABncAYX1c,5223
+jinns/experimental/__init__.py,sha256=DT9e57zbjfzPeRnXemGUqnGd--MhV77FspChT0z4YrE,410
+jinns/experimental/_diffrax_solver.py,sha256=upMr3kTTNrxEiSUO_oLvCXcjS9lPxSjvbB81h3qlhaU,6813
+jinns/loss/_DynamicLoss.py,sha256=4mb7OCP-cGZ_mG2MQ-AniddDcuBT78p4bQI7rZpwte4,22722
+jinns/loss/_DynamicLossAbstract.py,sha256=QhHRgvtcT-ifHlOxTyXbjDtHk9UfPN2Si8s3v9nEQm4,12672
+jinns/loss/_LossODE.py,sha256=DeejnU2ytgrOxUnwuVkQDWWRKJAgNQyjacTx-jT0xPA,13796
+jinns/loss/_LossPDE.py,sha256=ycjWJ99SuXe9DV5nROSWyq--xcp2JJ2PGWxsdWyZZog,36942
+jinns/loss/__init__.py,sha256=z5xYgBipNFf66__5BqQc6R_8r4F6A3TXL60YjsM8Osk,1287
+jinns/loss/_abstract_loss.py,sha256=DMxn0SQe9PW-pq3p5Oqvb0YK3_ulLDOnoIXzK219GV4,4576
+jinns/loss/_boundary_conditions.py,sha256=9HGw1cGLfmEilP4V4B2T0zl0YP1kNtrtXVLQNiBmWgc,12464
+jinns/loss/_loss_components.py,sha256=MMzaGlaRqESPjRzT0j0WU9HAqWQSbIXpGAqM1xQCZHw,1106
+jinns/loss/_loss_utils.py,sha256=R6PffBAtg6z9M8x1DFXmmqZpC095b9gZ_DB1phQxSuY,11168
+jinns/loss/_loss_weight_updates.py,sha256=9Bwouh7shLyc_wrdzN6CYL0ZuQH81uEs-L6wCeiYFx8,6817
+jinns/loss/_loss_weights.py,sha256=kII5WddORgeommFTudT3CSvhICpo6nSe47LclUgu_78,2429
+jinns/loss/_operators.py,sha256=Ds5yRH7hu-jaGRp7PYbt821BgYuEvgWHufWhYgdMjw0,22909
+jinns/nn/__init__.py,sha256=gwE48oqB_FsSIE-hUvCLz0jPaqX350LBxzH6ueFWYk4,456
+jinns/nn/_abstract_pinn.py,sha256=JUFjlV_nyheZw-max_tAUgFh6SspIbD5we_4bn70V6k,671
+jinns/nn/_hyperpinn.py,sha256=hF7HRLMMVBPT9CTQC2DjpDRcQDJCrT9cAj8wfApT_WE,19412
+jinns/nn/_mlp.py,sha256=Xmi-mG6uakN67R2S2UsBazdXIJVaGsD2B6TeJM1QjGY,8881
+jinns/nn/_pinn.py,sha256=4pvgUPQdQaO3cPBuEU7W4UaLV7lodqcR3pVR1sC0ni4,8774
+jinns/nn/_ppinn.py,sha256=LtjGQaLozdA4Kwutn8Pyerbu9yOc0t3_b701yfMb1ac,10392
+jinns/nn/_save_load.py,sha256=UqVy2oBzvIeBy6XB9tb61x3-x8i4dNCXJHC5_-bko-I,7477
+jinns/nn/_spinn.py,sha256=u5YG2FXcrg8p_uS2QFGmWoeFXYLxXnyV2e6BUHpo4xk,4774
+jinns/nn/_spinn_mlp.py,sha256=uCL454sF0Tfj7KT-fdXPnvKJYRQOuq60N0r2b2VAB8Q,7606
+jinns/nn/_utils.py,sha256=9UXz73iHKHVQYPBPIEitrHYJzJ14dspRwPfLA8avx0c,1120
+jinns/parameters/__init__.py,sha256=O0n7y6R1LRmFzzugCxMFCMS2pgsuWSh-XHjfFViN_eg,265
+jinns/parameters/_derivative_keys.py,sha256=YlLDX49PfYhr2Tj--t3praiD8JOUTZU6PTmjbNZsbMc,19173
+jinns/parameters/_params.py,sha256=qn4IGMJhD9lDBqOWmGEMy4gXt5a6KHfirkYZwHO7Vwk,2633
+jinns/plot/__init__.py,sha256=KPHX0Um4FbciZO1yD8kjZbkaT8tT964Y6SE2xCQ4eDU,135
+jinns/plot/_plot.py,sha256=-A5auNeElaz2_8UzVQJQE4143ZFg0zgMjStU7kwttEY,11565
+jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jinns/solver/_rar.py,sha256=vSVTnCGCusI1vTZCvIkP2_G8we44G_42yZHx2sOK9DE,10291
+jinns/solver/_solve.py,sha256=oVHnuc7Z0V2-ZYgZtCx7xdFd7TpB9w-6AwafX-kgBE4,28379
+jinns/solver/_utils.py,sha256=sM2UbVzYyjw24l4QSIR3IlynJTPGD_S08r8v0lXMxA8,5876
+jinns/utils/__init__.py,sha256=OEYWLCw8pKE7xoQREbd6SHvCjuw2QZHuVA6YwDcsBE8,53
+jinns/utils/_containers.py,sha256=YShcrPKfj5_I9mn3NMAS4Ea9MhhyL7fjv0e3MRbITHg,1837
+jinns/utils/_types.py,sha256=jl_91HtcrtE6UHbdTrRI8iUmr2kBUL0oP0UNIKhAXYw,1170
+jinns/utils/_utils.py,sha256=M7NXX9ok-BkH5o_xo74PB1_Cc8XiDipSl51rq82dTH4,2821
+jinns/validation/__init__.py,sha256=FTyUO-v1b8Tv-FDSQsntrH7zl9E0ENexqKMT_dFRkYo,124
+jinns/validation/_validation.py,sha256=8p6sMKiBAvA6JNm65hjkMj0997LJ0BkyCREEh0AnPVE,4803
+jinns-1.5.0.dist-info/licenses/AUTHORS,sha256=7NwCj9nU-HNG1asvy4qhQ2w7oZHrn-Lk5_wK_Ve7a3M,80
+jinns-1.5.0.dist-info/licenses/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-1.5.0.dist-info/METADATA,sha256=jEp__DP39B1HiTYVhtVcWKPmzS22kSUD6jNVSmHFh8g,5314
+jinns-1.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+jinns-1.5.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-1.5.0.dist-info/RECORD,,

{jinns-1.3.0.dist-info → jinns-1.5.0.dist-info}/WHEEL

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