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

jinns/loss/_DynamicLoss.py

@@ -85,9 +85,9 @@ class FisherKPP(PDENonStatio):
             lap = laplacian_rev(t_x, u, params)[..., None]
 
             return du_dt + self.Tmax * (
-                -params.eq_params["D"] * lap
+                -params.eq_params.D * lap
                 - u(t_x, params)
-                * (params.eq_params["r"] - params.eq_params["g"] * u(t_x, params))
+                * (params.eq_params.r - params.eq_params.g * u(t_x, params))
             )
         if isinstance(u, SPINN):
             s = jnp.zeros((1, self.dim_x + 1))
@@ -101,8 +101,8 @@ class FisherKPP(PDENonStatio):
             lap = laplacian_fwd(t_x, u, params)
 
             return du_dt + self.Tmax * (
-                -params.eq_params["D"] * lap
-                - u_tx * (params.eq_params["r"] - params.eq_params["g"] * u_tx)
+                -params.eq_params.D * lap
+                - u_tx * (params.eq_params.r - params.eq_params.g * u_tx)
             )
         raise ValueError("u is not among the recognized types (PINN or SPINN)")
 
@@ -164,17 +164,17 @@ class GeneralizedLotkaVolterra(ODE):
             The parameters in a Params object
         """
         du_dt = jax.jacrev(lambda t: jnp.log(u(t, params)))(t)
-        carrying_term = params.eq_params["carrying_capacities"] * jnp.sum(u(t, params))
+        carrying_term = params.eq_params.carrying_capacities * jnp.sum(u(t, params))
         interactions_terms = jax.tree.map(
             lambda interactions_for_i: jnp.sum(
                 interactions_for_i * u(t, params).squeeze()
             ),
-            params.eq_params["interactions"],
+            params.eq_params.interactions,
             is_leaf=eqx.is_array,
         )
         interactions_terms = jnp.array([*(interactions_terms)])
         return du_dt.squeeze() + self.Tmax * (
-            -params.eq_params["growth_rates"] + interactions_terms + carrying_term
+            -params.eq_params.growth_rates + interactions_terms + carrying_term
         )
 
 
@@ -229,7 +229,7 @@ class BurgersEquation(PDENonStatio):
 
             return du_dtx_values[0:1] + self.Tmax * (
                 u_(t_x) * du_dtx_values[1:2]
-                - params.eq_params["nu"] * d2u_dx_dtx(t_x)[1:2]
+                - params.eq_params.nu * d2u_dx_dtx(t_x)[1:2]
             )
 
         if isinstance(u, SPINN):
@@ -258,7 +258,7 @@ class BurgersEquation(PDENonStatio):
             )[1]
             _, d2u_dx2 = jax.jvp(du_dx_fun, (t_x,), (v1,))
             # Note that ones_like(x) works because x is Bx1 !
-            return du_dt + self.Tmax * (u_tx * du_dx - params.eq_params["nu"] * d2u_dx2)
+            return du_dt + self.Tmax * (u_tx * du_dx - params.eq_params.nu * d2u_dx2)
         raise ValueError("u is not among the recognized types (PINN or SPINN)")
 
 
@@ -458,7 +458,7 @@ class OU_FPENonStatioLoss2D(FPENonStatioLoss2D):
         eq_params
             A dictionary containing the equation parameters
         """
-        return eq_params["alpha"] * (eq_params["mu"] - x)
+        return eq_params.alpha * (eq_params.mu - x)
 
     def sigma_mat(self, x, eq_params):
         r"""
@@ -473,7 +473,7 @@ class OU_FPENonStatioLoss2D(FPENonStatioLoss2D):
             A dictionary containing the equation parameters
         """
 
-        return jnp.diag(eq_params["sigma"])
+        return jnp.diag(eq_params.sigma)
 
     def diffusion(self, x, eq_params, i=None, j=None):
         r"""
@@ -587,15 +587,15 @@ class NavierStokesMassConservation2DStatio(PDEStatio):
             # dynamic loss on x axis
             result_x = (
                 u_dot_nabla_x_u[0]
-                + 1 / params.eq_params["rho"] * jac_p[0, 0]
-                - params.eq_params["nu"] * vec_laplacian_u[0]
+                + 1 / params.eq_params.rho * jac_p[0, 0]
+                - params.eq_params.nu * vec_laplacian_u[0]
             )
 
             # dynamic loss on y axis
             result_y = (
                 u_dot_nabla_x_u[1]
-                + 1 / params.eq_params["rho"] * jac_p[0, 1]
-                - params.eq_params["nu"] * vec_laplacian_u[1]
+                + 1 / params.eq_params.rho * jac_p[0, 1]
+                - params.eq_params.nu * vec_laplacian_u[1]
             )
 
             # MASS CONVERVATION
@@ -605,7 +605,8 @@ class NavierStokesMassConservation2DStatio(PDEStatio):
             # output is 3D
             if mc.ndim == 0 and not result_x.ndim == 0:
                 mc = mc[None]
-            return jnp.stack([result_x, result_y, mc], axis=-1)
+
+            return jnp.stack([result_x, result_y, mc], axis=-1).squeeze()
 
         if isinstance(u_p, SPINN):
             u = lambda x, params: u_p(x, params)[..., 0:2]
@@ -626,14 +627,14 @@ class NavierStokesMassConservation2DStatio(PDEStatio):
             # dynamic loss on x axis
             result_x = (
                 u_dot_nabla_x_u[..., 0]
-                + 1 / params.eq_params["rho"] * dp_dx.squeeze()
-                - params.eq_params["nu"] * vec_laplacian_u[..., 0]
+                + 1 / params.eq_params.rho * dp_dx.squeeze()
+                - params.eq_params.nu * vec_laplacian_u[..., 0]
             )
             # dynamic loss on y axis
             result_y = (
                 u_dot_nabla_x_u[..., 1]
-                + 1 / params.eq_params["rho"] * dp_dy.squeeze()
-                - params.eq_params["nu"] * vec_laplacian_u[..., 1]
+                + 1 / params.eq_params.rho * dp_dy.squeeze()
+                - params.eq_params.nu * vec_laplacian_u[..., 1]
             )
 
             # MASS CONVERVATION

jinns/loss/_DynamicLossAbstract.py

@@ -9,10 +9,13 @@ from __future__ import (
 import warnings
 import abc
 from functools import partial
-from typing import Callable, TYPE_CHECKING, ClassVar, Generic, TypeVar
+from dataclasses import InitVar
+from typing import Callable, TYPE_CHECKING, ClassVar, Generic, TypeVar, Any
 import equinox as eqx
-from jaxtyping import Float, Array
+from jaxtyping import Float, Array, PyTree
+import jax
 import jax.numpy as jnp
+from jinns.parameters._params import EqParams
 
 
 # See : https://docs.kidger.site/equinox/api/module/advanced_fields/#equinox.AbstractClassVar--known-issues
@@ -59,7 +62,7 @@ class DynamicLoss(eqx.Module, Generic[InputDim]):
         Tmax needs to be given when the PINN time input is normalized in
         [0, 1], ie. we have performed renormalization of the differential
         equation
-    eq_params_heterogeneity : dict[str, Callable | None], default=None
+    eq_params_heterogeneity : dict[str, Callable[[InputDim, AbstractPINN, Params[Array]], Array] | None], default=None
         A dict with the same keys as eq_params and the value being either None
         (no heterogeneity) or a function which encodes for the spatio-temporal
         heterogeneity of the parameter.
@@ -71,7 +74,10 @@ class DynamicLoss(eqx.Module, Generic[InputDim]):
         `eq_params` too.
         A value can be missing, in this case there is no heterogeneity (=None).
         Default None, meaning there is no heterogeneity in the equation
-        parameters.
+        parameters. Note that since 1.6.0, this is handled inernally as a
+        `PyTree[Callable[[InputDim, AbstractPINN, Params[Array]], Array] |
+        None] | None` (`Params.eq_params` is not a dict
+        anymore).
     vectorial_dyn_loss_ponderation : Float[Array, " dim"], default=None
         Add a different ponderation weight to each of the dimension to the
         dynamic loss. This array must have the same dimension as the output of
@@ -86,40 +92,49 @@ class DynamicLoss(eqx.Module, Generic[InputDim]):
 
     _eq_type = AbstractClassVar[str]  # class variable denoting the type of
     # differential equation
-    Tmax: Float = eqx.field(kw_only=True, default=1)
-    eq_params_heterogeneity: dict[str, Callable | None] | None = eqx.field(
-        kw_only=True, default=None, static=True
-    )
+    Tmax: float = eqx.field(kw_only=True, default=1)
+    eq_params_heterogeneity: (
+        PyTree[Callable[[InputDim, AbstractPINN, Params[Array]], Array] | None] | None
+    ) = eqx.field(kw_only=True, default=None, static=True)
     vectorial_dyn_loss_ponderation: Float[Array, " dim"] | None = eqx.field(
-        kw_only=True, default=None
+        kw_only=True, default_factory=lambda: jnp.array(1.0)
     )
-
-    def __post_init__(self):
-        if self.vectorial_dyn_loss_ponderation is None:
-            self.vectorial_dyn_loss_ponderation = jnp.array(1.0)
+    params: InitVar[Params[Array]] = eqx.field(default=None)
+
+    def __post_init__(self, params: Params[Array] | None = None):
+        if isinstance(self.eq_params_heterogeneity, dict):  # type: ignore
+            # we cannot use the same converter as in Params.eq_params
+            # we don't want to create a new type but use the same type as
+            # Params.eq_params which already exists.
+            if params is None:
+                raise ValueError(
+                    "When `self.eq_params_heterogeneity` is "
+                    "provided, `params` must be specified at init"
+                )
+            self.eq_params_heterogeneity = EqParams(
+                self.eq_params_heterogeneity,
+                "EqParams",  # type: ignore
+            )
 
     def _eval_heterogeneous_parameters(
         self,
         inputs: InputDim,
         u: AbstractPINN,
         params: Params[Array],
-        eq_params_heterogeneity: dict[str, Callable | None] | None = None,
-    ) -> dict[str, Array]:
-        eq_params_ = {}
+        eq_params_heterogeneity: PyTree[
+            Callable[[InputDim, AbstractPINN, Params[Array]], Array] | None
+        ]
+        | None = None,
+    ) -> PyTree[Array]:
         if eq_params_heterogeneity is None:
             return params.eq_params
-
-        for k, p in params.eq_params.items():
-            try:
-                if eq_params_heterogeneity[k] is not None:
-                    eq_params_[k] = eq_params_heterogeneity[k](inputs, u, params)  # type: ignore don't know why pyright says
-                    # eq_params_heterogeneity[k] can be None here
-                else:
-                    eq_params_[k] = p
-            except KeyError:
-                # we authorize missing eq_params_heterogeneity key
-                # if its heterogeneity is None anyway
-                eq_params_[k] = p
+        eq_params_ = jax.tree.map(
+            lambda p, fun: (  # type: ignore
+                fun(inputs, u, params) if fun is not None else p
+            ),
+            params.eq_params,
+            eq_params_heterogeneity,
+        )
         return eq_params_
 
     @partial(_decorator_heteregeneous_params)
@@ -128,7 +143,7 @@ class DynamicLoss(eqx.Module, Generic[InputDim]):
         inputs: InputDim,
         u: AbstractPINN,
         params: Params[Array],
-    ) -> float:
+    ) -> Float[Array, " eq_dim"]:
         evaluation = self.vectorial_dyn_loss_ponderation * self.equation(
             inputs, u, params
         )
@@ -147,7 +162,7 @@ class DynamicLoss(eqx.Module, Generic[InputDim]):
         return evaluation
 
     @abc.abstractmethod
-    def equation(self, *args, **kwargs):
+    def equation(self, *args: Any, **kwargs: Any) -> Float[Array, " eq_dim"]:
         # TO IMPLEMENT
         # Point-wise evaluation of the differential equation N[u](.)
         raise NotImplementedError("You should implement your equation.")
@@ -183,7 +198,7 @@ class ODE(DynamicLoss[Float[Array, " 1"]]):
     @abc.abstractmethod
     def equation(
         self, t: Float[Array, " 1"], u: AbstractPINN, params: Params[Array]
-    ) -> float:
+    ) -> Float[Array, " eq_dim"]:
         r"""
         The differential operator defining the ODE.
 
@@ -202,7 +217,7 @@ class ODE(DynamicLoss[Float[Array, " 1"]]):
 
         Returns
         -------
-        float
+        Float[Array, "eq_dim"]
             The residual, *i.e.* the differential operator $\mathcal{N}_\theta[u_\nu](t)$ evaluated at point `t`.
 
         Raises
@@ -242,7 +257,7 @@ class PDEStatio(DynamicLoss[Float[Array, " dim"]]):
     @abc.abstractmethod
     def equation(
         self, x: Float[Array, " dim"], u: AbstractPINN, params: Params[Array]
-    ) -> float:
+    ) -> Float[Array, " eq_dim"]:
         r"""The differential operator defining the stationnary PDE.
 
         !!! warning
@@ -260,7 +275,7 @@ class PDEStatio(DynamicLoss[Float[Array, " dim"]]):
 
         Returns
         -------
-        float
+        Float[Array, "eq_dim"]
             The residual, *i.e.* the differential operator $\mathcal{N}_\theta[u_\nu](x)$ evaluated at point `x`.
 
         Raises
@@ -303,7 +318,7 @@ class PDENonStatio(DynamicLoss[Float[Array, " 1 + dim"]]):
         t_x: Float[Array, " 1 + dim"],
         u: AbstractPINN,
         params: Params[Array],
-    ) -> float:
+    ) -> Float[Array, " eq_dim"]:
         r"""The differential operator defining the non-stationnary PDE.
 
         !!! warning
@@ -320,7 +335,7 @@ class PDENonStatio(DynamicLoss[Float[Array, " 1 + dim"]]):
             The parameters of the equation and the networks, $\theta$ and $\nu$ respectively.
         Returns
         -------
-        float
+        Float[Array, "eq_dim"]
             The residual, *i.e.* the differential operator $\mathcal{N}_\theta[u_\nu](t, x)$ evaluated at point `(t, x)`.