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

jinns/loss/_loss_weights.py

@@ -2,58 +2,82 @@
 Formalize the loss weights data structure
 """
 
-from typing import Dict
-from jaxtyping import Array, Float
+from __future__ import annotations
+from dataclasses import fields
+
+from jaxtyping import Array
+import jax.numpy as jnp
 import equinox as eqx
 
 
-class LossWeightsODE(eqx.Module):
+def lw_converter(x):
+    if x is None:
+        return x
+    else:
+        return jnp.asarray(x)
 
-    dyn_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-    initial_condition: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-    observations: Array | Float | None = eqx.field(kw_only=True, default=1.0)
 
+class AbstractLossWeights(eqx.Module):
+    """
+    An abstract class, currently only useful for type hints
 
-class LossWeightsODEDict(eqx.Module):
+    TODO in the future maybe loss weights could be subclasses of
+    XDEComponentsAbstract?
+    """
 
-    dyn_loss: Dict[str, Array | Float | None] = eqx.field(kw_only=True, default=None)
-    initial_condition: Dict[str, Array | Float | None] = eqx.field(
-        kw_only=True, default=None
+    def items(self):
+        """
+        For the dataclass to be iterated like a dictionary.
+        Practical and retrocompatible with old code when loss components were
+        dictionaries
+        """
+        return {
+            field.name: getattr(self, field.name)
+            for field in fields(self)
+            if getattr(self, field.name) is not None
+        }.items()
+
+
+class LossWeightsODE(AbstractLossWeights):
+    dyn_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
     )
-    observations: Dict[str, Array | Float | None] = eqx.field(
-        kw_only=True, default=None
+    initial_condition: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    observations: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
     )
 
 
-class LossWeightsPDEStatio(eqx.Module):
-
-    dyn_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-    norm_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-    boundary_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-    observations: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-
-
-class LossWeightsPDENonStatio(eqx.Module):
-
-    dyn_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-    norm_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-    boundary_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-    observations: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-    initial_condition: Array | Float | None = eqx.field(kw_only=True, default=1.0)
-
+class LossWeightsPDEStatio(AbstractLossWeights):
+    dyn_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    norm_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    boundary_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    observations: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
 
-class LossWeightsPDEDict(eqx.Module):
-    """
-    Only one type of LossWeights data structure for the SystemLossPDE:
-    Include the initial condition always for the code to be more generic
-    """
 
-    dyn_loss: Dict[str, Array | Float | None] = eqx.field(kw_only=True, default=1.0)
-    norm_loss: Dict[str, Array | Float | None] = eqx.field(kw_only=True, default=1.0)
-    boundary_loss: Dict[str, Array | Float | None] = eqx.field(
-        kw_only=True, default=1.0
+class LossWeightsPDENonStatio(AbstractLossWeights):
+    dyn_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    norm_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    boundary_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    observations: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
     )
-    observations: Dict[str, Array | Float | None] = eqx.field(kw_only=True, default=1.0)
-    initial_condition: Dict[str, Array | Float | None] = eqx.field(
-        kw_only=True, default=1.0
+    initial_condition: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
     )

jinns/loss/_operators.py

@@ -2,22 +2,42 @@
 Implements diverse operators for dynamic losses
 """
 
-from typing import Literal
+from __future__ import (
+    annotations,
+)
+
+from typing import Literal, cast, Callable
 
 import jax
 import jax.numpy as jnp
 from jax import grad
-import equinox as eqx
 from jaxtyping import Float, Array
 from jinns.parameters._params import Params
+from jinns.nn._abstract_pinn import AbstractPINN
+
+
+def _get_eq_type(
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None,
+) -> Literal["nonstatio_PDE", "statio_PDE"]:
+    """
+    But we filter out ODE from eq_type because we only have operators that does
+    not work with ODEs so far
+    """
+    if isinstance(u, AbstractPINN):
+        assert u.eq_type != "ODE", "Cannot compute the operator for ODE PINNs"
+        return u.eq_type
+    if eq_type is None:
+        raise ValueError("eq_type could not be set!")
+    return eq_type
 
 
 def divergence_rev(
-    inputs: Float[Array, "dim"] | Float[Array, "1+dim"],
-    u: eqx.Module,
-    params: Params,
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] = None,
-) -> float:
+    inputs: Float[Array, " dim"] | Float[Array, " 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " "]:
     r"""
     Compute the divergence of a vector field $\mathbf{u}$, i.e.,
     $\nabla_\mathbf{x} \cdot \mathbf{u}(\mathrm{inputs})$ with $\mathbf{u}$ a vector
@@ -41,13 +61,7 @@ def divergence_rev(
         can know that by inspecting the `u` argument (PINN object). But if `u` is
         a function, we must set this attribute.
     """
-
-    try:
-        eq_type = u.eq_type
-    except AttributeError:
-        pass  # use the value passed as argument
-    if eq_type is None:
-        raise ValueError("eq_type could not be set!")
+    eq_type = _get_eq_type(u, eq_type)
 
     def scan_fun(_, i):
         if eq_type == "nonstatio_PDE":
@@ -70,11 +84,11 @@ def divergence_rev(
 
 
 def divergence_fwd(
-    inputs: Float[Array, "batch_size dim"] | Float[Array, "batch_size 1+dim"],
-    u: eqx.Module,
-    params: Params,
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] = None,
-) -> Float[Array, "batch_size * (1+dim) 1"] | Float[Array, "batch_size * (dim) 1"]:
+    inputs: Float[Array, " batch_size dim"] | Float[Array, " batch_size 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " batch_size * (1+dim) 1"] | Float[Array, " batch_size * (dim) 1"]:
     r"""
     Compute the divergence of a **batched** vector field $\mathbf{u}$, i.e.,
     $\nabla_\mathbf{x} \cdot \mathbf{u}(\mathbf{x})$ with $\mathbf{u}$ a vector
@@ -103,13 +117,7 @@ def divergence_fwd(
         can know that by inspecting the `u` argument (PINN object). But if `u` is
         a function, we must set this attribute.
     """
-
-    try:
-        eq_type = u.eq_type
-    except AttributeError:
-        pass  # use the value passed as argument
-    if eq_type is None:
-        raise ValueError("eq_type could not be set!")
+    eq_type = _get_eq_type(u, eq_type)
 
     def scan_fun(_, i):
         if eq_type == "nonstatio_PDE":
@@ -142,12 +150,12 @@ def divergence_fwd(
 
 
 def laplacian_rev(
-    inputs: Float[Array, "dim"] | Float[Array, "1+dim"],
-    u: eqx.Module,
-    params: Params,
+    inputs: Float[Array, " dim"] | Float[Array, " 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
     method: Literal["trace_hessian_x", "trace_hessian_t_x", "loop"] = "trace_hessian_x",
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] = None,
-) -> float:
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " "]:
     r"""
     Compute the Laplacian of a scalar field $u$ from $\mathbb{R}^d$
     to $\mathbb{R}$ or from $\mathbb{R}^{1+d}$ to $\mathbb{R}$, i.e., this
@@ -180,13 +188,7 @@ def laplacian_rev(
         can know that by inspecting the `u` argument (PINN object). But if `u` is
         a function, we must set this attribute.
     """
-
-    try:
-        eq_type = u.eq_type
-    except AttributeError:
-        pass  # use the value passed as argument
-    if eq_type is None:
-        raise ValueError("eq_type could not be set!")
+    eq_type = _get_eq_type(u, eq_type)
 
     if method == "trace_hessian_x":
         # NOTE we afford a concatenate here to avoid computing Hessian elements for
@@ -226,16 +228,12 @@ def laplacian_rev(
             if eq_type == "nonstatio_PDE":
                 d2u_dxi2 = grad(
                     lambda inputs: grad(u_)(inputs)[1 + i],
-                )(
-                    inputs
-                )[1 + i]
+                )(inputs)[1 + i]
             else:
                 d2u_dxi2 = grad(
                     lambda inputs: grad(u_, 0)(inputs)[i],
                     0,
-                )(
-                    inputs
-                )[i]
+                )(inputs)[i]
             return _, d2u_dxi2
 
         if eq_type == "nonstatio_PDE":
@@ -251,12 +249,12 @@ def laplacian_rev(
 
 
 def laplacian_fwd(
-    inputs: Float[Array, "batch_size 1+dim"] | Float[Array, "batch_size dim"],
-    u: eqx.Module,
-    params: Params,
+    inputs: Float[Array, " batch_size 1+dim"] | Float[Array, " batch_size dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
     method: Literal["trace_hessian_t_x", "trace_hessian_x", "loop"] = "loop",
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] = None,
-) -> Float[Array, "batch_size * (1+dim) 1"] | Float[Array, "batch_size * (dim) 1"]:
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " batch_size * (1+dim) 1"] | Float[Array, " batch_size * (dim) 1"]:
     r"""
     Compute the Laplacian of a **batched** scalar field $u$
     from $\mathbb{R}^{b\times d}$ to $\mathbb{R}^{b\times b}$ or
@@ -299,13 +297,7 @@ def laplacian_fwd(
         can know that by inspecting the `u` argument (PINN object). But if `u` is
         a function, we must set this attribute.
     """
-
-    try:
-        eq_type = u.eq_type
-    except AttributeError:
-        pass  # use the value passed as argument
-    if eq_type is None:
-        raise ValueError("eq_type could not be set!")
+    eq_type = _get_eq_type(u, eq_type)
 
     if method == "loop":
 
@@ -398,11 +390,12 @@ def laplacian_fwd(
 
 
 def vectorial_laplacian_rev(
-    inputs: Float[Array, "dim"] | Float[Array, "1+dim"],
-    u: eqx.Module,
-    params: Params,
-    dim_out: int = None,
-) -> Float[Array, "dim_out"]:
+    inputs: Float[Array, " dim"] | Float[Array, " 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+    dim_out: int | None = None,
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " dim_out"]:
     r"""
     Compute the vectorial Laplacian of a vector field $\mathbf{u}$ from
     $\mathbb{R}^d$ to $\mathbb{R}^n$ or from $\mathbb{R}^{1+d}$ to
@@ -426,7 +419,12 @@ def vectorial_laplacian_rev(
     dim_out
         Dimension of the vector $\mathbf{u}(\mathrm{inputs})$. This needs to be
         provided if it is different than that of $\mathrm{inputs}$.
+    eq_type
+        whether we consider a stationary or non stationary PINN. Most often we
+        can know that by inspecting the `u` argument (PINN object). But if `u` is
+        a function, we must set this attribute.
     """
+    eq_type = _get_eq_type(u, eq_type)
     if dim_out is None:
         dim_out = inputs.shape[0]
 
@@ -435,7 +433,9 @@ def vectorial_laplacian_rev(
         # each of these components
         # Note the jnp.expand_dims call
         uj = lambda inputs, params: jnp.expand_dims(u(inputs, params)[j], axis=-1)
-        lap_on_j = laplacian_rev(inputs, uj, params, eq_type=u.eq_type)
+        lap_on_j = laplacian_rev(
+            inputs, cast(AbstractPINN, uj), params, eq_type=eq_type
+        )
 
         return _, lap_on_j
 
@@ -444,11 +444,12 @@ def vectorial_laplacian_rev(
 
 
 def vectorial_laplacian_fwd(
-    inputs: Float[Array, "batch_size dim"] | Float[Array, "batch_size 1+dim"],
-    u: eqx.Module,
-    params: Params,
-    dim_out: int = None,
-) -> Float[Array, "batch_size * (1+dim) n"] | Float[Array, "batch_size * (dim) n"]:
+    inputs: Float[Array, " batch_size dim"] | Float[Array, " batch_size 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+    dim_out: int | None = None,
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " batch_size * (1+dim) n"] | Float[Array, " batch_size * (dim) n"]:
     r"""
     Compute the vectorial Laplacian of a vector field $\mathbf{u}$ when
     `u` is a SPINN, in this case, it corresponds to a vector
@@ -474,7 +475,12 @@ def vectorial_laplacian_fwd(
     dim_out
         the value of the output dimension ($n$ in the formula above). Must be
         set if different from $d$.
+    eq_type
+        whether we consider a stationary or non stationary PINN. Most often we
+        can know that by inspecting the `u` argument (PINN object). But if `u` is
+        a function, we must set this attribute.
     """
+    eq_type = _get_eq_type(u, eq_type)
     if dim_out is None:
         dim_out = inputs.shape[0]
 
@@ -483,7 +489,9 @@ def vectorial_laplacian_fwd(
         # each of these components
         # Note the expand_dims
         uj = lambda inputs, params: jnp.expand_dims(u(inputs, params)[..., j], axis=-1)
-        lap_on_j = laplacian_fwd(inputs, uj, params, eq_type=u.eq_type)
+        lap_on_j = laplacian_fwd(
+            inputs, cast(AbstractPINN, uj), params, eq_type=eq_type
+        )
 
         return _, lap_on_j
 
@@ -492,8 +500,10 @@ def vectorial_laplacian_fwd(
 
 
 def _u_dot_nabla_times_u_rev(
-    x: Float[Array, "2"], u: eqx.Module, params: Params
-) -> Float[Array, "2"]:
+    x: Float[Array, " 2"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+) -> Float[Array, " 2"]:
     r"""
     Implement $((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})$ for
     $\mathbf{x}$ of arbitrary
@@ -522,10 +532,10 @@ def _u_dot_nabla_times_u_rev(
 
 
 def _u_dot_nabla_times_u_fwd(
-    x: Float[Array, "batch_size 2"],
-    u: eqx.Module,
-    params: Params,
-) -> Float[Array, "batch_size batch_size 2"]:
+    x: Float[Array, " batch_size 2"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+) -> Float[Array, " batch_size batch_size 2"]:
     r"""
     Implement :math:`((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})` for
     :math:`\mathbf{x}` of arbitrary dimension **with a batch dimension**.

jinns/nn/__init__.py

@@ -1,7 +1,22 @@
 from ._save_load import save_pinn, load_pinn
+from ._abstract_pinn import AbstractPINN
 from ._pinn import PINN
 from ._spinn import SPINN
 from ._mlp import PINN_MLP, MLP
 from ._spinn_mlp import SPINN_MLP, SMLP
 from ._hyperpinn import HyperPINN
 from ._ppinn import PPINN_MLP
+
+__all__ = [
+    "save_pinn",
+    "load_pinn",
+    "AbstractPINN",
+    "PINN",
+    "SPINN",
+    "PINN_MLP",
+    "MLP",
+    "SPINN_MLP",
+    "SMLP",
+    "HyperPINN",
+    "PPINN_MLP",
+]

jinns/nn/_abstract_pinn.py

@@ -0,0 +1,22 @@
+import abc
+from typing import Literal, Any
+from jaxtyping import Array
+import equinox as eqx
+
+from jinns.nn._utils import _PyTree_to_Params
+from jinns.parameters._params import Params
+
+
+class AbstractPINN(eqx.Module):
+    """
+    Basically just a way to add a __call__ to an eqx.Module.
+    The way to go for correct type hints apparently
+    https://github.com/patrick-kidger/equinox/issues/1002 + https://docs.kidger.site/equinox/pattern/
+    """
+
+    eq_type: eqx.AbstractVar[Literal["ODE", "statio_PDE", "nonstatio_PDE"]]
+
+    @abc.abstractmethod
+    @_PyTree_to_Params
+    def __call__(self, inputs: Any, params: Params[Array], *args, **kwargs) -> Any:
+        pass