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

jinns/loss/_abstract_loss.py

@@ -1,41 +1,77 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Self, Literal, Callable
-from jaxtyping import Array, PyTree, Key
+from typing import Self, Literal, Callable, TypeVar, Generic, Any
+from jaxtyping import PRNGKeyArray, Array, PyTree, Float
 import equinox as eqx
 import jax
 import jax.numpy as jnp
 import optax
-from jinns.loss._loss_weights import AbstractLossWeights
 from jinns.parameters._params import Params
 from jinns.loss._loss_weight_updates import soft_adapt, lr_annealing, ReLoBRaLo
+from jinns.utils._types import AnyLossComponents, AnyBatch, AnyLossWeights
 
-if TYPE_CHECKING:
-    from jinns.utils._types import AnyLossComponents, AnyBatch
+L = TypeVar(
+    "L", bound=AnyLossWeights
+)  # we want to be able to use one of the element of AnyLossWeights
+# that is https://stackoverflow.com/a/79534258 via `bound`
 
+B = TypeVar(
+    "B", bound=AnyBatch
+)  # The above comment also works with Unions (https://docs.python.org/3/library/typing.html#typing.TypeVar)
+# We then do the same TypeVar to be able to use one of the element of AnyBatch
+# in the evaluate_by_terms methods of child classes.
+C = TypeVar(
+    "C", bound=AnyLossComponents[Array | None]
+)  # The above comment also works with Unions (https://docs.python.org/3/library/typing.html#typing.TypeVar)
 
-class AbstractLoss(eqx.Module):
+# In the cases above, without the bound, we could not have covariance on
+# the type because it would break LSP. Note that covariance on the return type
+# is authorized in LSP hence we do not need the same TypeVar instruction for
+# the return types of evaluate_by_terms for example!
+
+
+class AbstractLoss(eqx.Module, Generic[L, B, C]):
     """
     About the call:
     https://github.com/patrick-kidger/equinox/issues/1002 + https://docs.kidger.site/equinox/pattern/
     """
 
-    loss_weights: AbstractLossWeights
+    loss_weights: eqx.AbstractVar[L]
     update_weight_method: Literal["soft_adapt", "lr_annealing", "ReLoBRaLo"] | None = (
         eqx.field(kw_only=True, default=None, static=True)
     )
 
-    @abc.abstractmethod
-    def __call__(self, *_, **__) -> Array:
-        pass
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self.evaluate(*args, **kwargs)
 
     @abc.abstractmethod
-    def evaluate_by_terms(
-        self, params: Params[Array], batch: AnyBatch
-    ) -> tuple[AnyLossComponents, AnyLossComponents]:
+    def evaluate_by_terms(self, params: Params[Array], batch: B) -> tuple[C, C]:
         pass
 
+    def evaluate(self, params: Params[Array], batch: B) -> tuple[Float[Array, " "], C]:
+        """
+        Evaluate the loss function at a batch of points for given parameters.
+
+        We retrieve the total value itself and a PyTree with loss values for each term
+
+        Parameters
+        ---------
+        params
+            Parameters at which the loss is evaluated
+        batch
+            Composed of a batch of points in the
+            domain, a batch of points in the domain
+            border and an optional additional batch of parameters (eg. for
+            metamodeling) and an optional additional batch of observed
+            inputs/outputs/parameters
+        """
+        loss_terms, _ = self.evaluate_by_terms(params, batch)
+
+        loss_val = self.ponderate_and_sum_loss(loss_terms)
+
+        return loss_val, loss_terms
+
     def get_gradients(
         self, fun: Callable[[Params[Array]], Array], params: Params[Array]
     ) -> tuple[Array, Array]:
@@ -48,7 +84,7 @@ class AbstractLoss(eqx.Module):
         loss_val, grads = value_grad_loss(params)
         return loss_val, grads
 
-    def ponderate_and_sum_loss(self, terms):
+    def ponderate_and_sum_loss(self, terms: C) -> Array:
         """
         Get total loss from individual loss terms and weights
 
@@ -58,19 +94,18 @@ class AbstractLoss(eqx.Module):
             self.loss_weights,
             is_leaf=lambda x: eqx.is_inexact_array(x) and x is not None,
         )
-        terms = jax.tree.leaves(
+        terms_list = jax.tree.leaves(
             terms, is_leaf=lambda x: eqx.is_inexact_array(x) and x is not None
         )
-        if len(weights) == len(terms):
-            return jnp.sum(jnp.array(weights) * jnp.array(terms))
-        else:
-            raise ValueError(
-                "The numbers of declared loss weights and "
-                "declared loss terms do not concord "
-                f" got {len(weights)} and {len(terms)}"
-            )
+        if len(weights) == len(terms_list):
+            return jnp.sum(jnp.array(weights) * jnp.array(terms_list))
+        raise ValueError(
+            "The numbers of declared loss weights and "
+            "declared loss terms do not concord "
+            f" got {len(weights)} and {len(terms_list)}"
+        )
 
-    def ponderate_and_sum_gradient(self, terms):
+    def ponderate_and_sum_gradient(self, terms: C) -> C:
         """
         Get total gradients from individual loss gradients and weights
         for each parameter
@@ -102,7 +137,7 @@ class AbstractLoss(eqx.Module):
         loss_terms: PyTree,
         stored_loss_terms: PyTree,
         grad_terms: PyTree,
-        key: Key,
+        key: PRNGKeyArray,
     ) -> Self:
         """
         Update the loss weights according to a predefined scheme

jinns/loss/_loss_components.py

@@ -1,38 +1,17 @@
 from typing import TypeVar, Generic
-from dataclasses import fields
-import equinox as eqx
+
+from jinns.utils._ItemizableModule import ItemizableModule
 
 T = TypeVar("T")
 
 
-class XDEComponentsAbstract(eqx.Module, Generic[T]):
-    """
-    Provides a template for ODE components with generic types.
-    One can inherit to specialize and add methods and attributes
-    We do not enforce keyword only to avoid being to verbose (this then can
-    work like a tuple)
-    """
-
-    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 ODEComponents(XDEComponentsAbstract[T]):
+class ODEComponents(ItemizableModule, Generic[T]):
     dyn_loss: T
     initial_condition: T
     observations: T
 
 
-class PDEStatioComponents(XDEComponentsAbstract[T]):
+class PDEStatioComponents(ItemizableModule, Generic[T]):
     dyn_loss: T
     norm_loss: T
     boundary_loss: T

jinns/loss/_loss_utils.py

@@ -308,3 +308,26 @@ def initial_condition_apply(
     else:
         raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_initial_condition
+
+
+def initial_condition_check(x, dim_size=None):
+    """
+    Make a (dim_size,) jnp array from an int, a float or a 0D jnp array
+
+    """
+    if isinstance(x, Array):
+        if not x.shape:  # e.g. user input: jnp.array(0.)
+            x = jnp.array([x])
+        if dim_size is not None:  # we check for the required dims_ize
+            if x.shape != (dim_size,):
+                raise ValueError(
+                    f"Wrong dim_size. It should be({dim_size},). Got shape: {x.shape}"
+                )
+
+    elif isinstance(x, float):  # e.g. user input: 0.
+        x = jnp.array([x])
+    elif isinstance(x, int):  # e.g. user input: 0
+        x = jnp.array([float(x)])
+    else:
+        raise ValueError(f"Wrong value, expected Array, float or int, got {type(x)}")
+    return x

jinns/loss/_loss_weight_updates.py

@@ -4,18 +4,17 @@ A collection of specific weight update schemes in jinns
 
 from __future__ import annotations
 from typing import TYPE_CHECKING
-from jaxtyping import Array, Key
+from jaxtyping import Array, PRNGKeyArray
 import jax.numpy as jnp
 import jax
 import equinox as eqx
 
 if TYPE_CHECKING:
-    from jinns.loss._loss_weights import AbstractLossWeights
-    from jinns.utils._types import AnyLossComponents
+    from jinns.utils._types import AnyLossComponents, AnyLossWeights
 
 
 def soft_adapt(
-    loss_weights: AbstractLossWeights,
+    loss_weights: AnyLossWeights,
     iteration_nb: int,
     loss_terms: AnyLossComponents,
     stored_loss_terms: AnyLossComponents,
@@ -58,11 +57,11 @@ def soft_adapt(
 
 
 def ReLoBRaLo(
-    loss_weights: AbstractLossWeights,
+    loss_weights: AnyLossWeights,
     iteration_nb: int,
     loss_terms: AnyLossComponents,
     stored_loss_terms: AnyLossComponents,
-    key: Key,
+    key: PRNGKeyArray,
     decay_factor: float = 0.9,
     tau: float = 1,  ## referred to as temperature in the article
     p: float = 0.9,
@@ -146,7 +145,7 @@ def ReLoBRaLo(
 
 
 def lr_annealing(
-    loss_weights: AbstractLossWeights,
+    loss_weights: AnyLossWeights,
     grad_terms: AnyLossComponents,
     decay_factor: float = 0.9,  # 0.9 is the recommended value from the article
     eps: float = 1e-6,

jinns/loss/_loss_weights.py

@@ -3,81 +3,80 @@ Formalize the loss weights data structure
 """
 
 from __future__ import annotations
-from dataclasses import fields
 
 from jaxtyping import Array
 import jax.numpy as jnp
 import equinox as eqx
 
+from jinns.loss._loss_components import (
+    ODEComponents,
+    PDEStatioComponents,
+    PDENonStatioComponents,
+)
 
-def lw_converter(x):
+
+def lw_converter(x: Array | None) -> Array | None:
     if x is None:
         return x
     else:
         return jnp.asarray(x)
 
 
-class AbstractLossWeights(eqx.Module):
+class LossWeightsODE(ODEComponents[Array | None]):
     """
-    An abstract class, currently only useful for type hints
-
-    TODO in the future maybe loss weights could be subclasses of
-    XDEComponentsAbstract?
+    Value given at initialization is converted to a jnp.array orunmodified if None.
+    This means that at initialization, the user can pass a float or int
     """
 
-    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(
+    dyn_loss: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
-    initial_condition: Array | float | None = eqx.field(
+    initial_condition: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
-    observations: Array | float | None = eqx.field(
+    observations: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
 
 
-class LossWeightsPDEStatio(AbstractLossWeights):
-    dyn_loss: Array | float | None = eqx.field(
+class LossWeightsPDEStatio(PDEStatioComponents[Array | None]):
+    """
+    Value given at initialization is converted to a jnp.array orunmodified if None.
+    This means that at initialization, the user can pass a float or int
+    """
+
+    dyn_loss: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
-    norm_loss: Array | float | None = eqx.field(
+    norm_loss: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
-    boundary_loss: Array | float | None = eqx.field(
+    boundary_loss: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
-    observations: Array | float | None = eqx.field(
+    observations: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
 
 
-class LossWeightsPDENonStatio(AbstractLossWeights):
-    dyn_loss: Array | float | None = eqx.field(
+class LossWeightsPDENonStatio(PDENonStatioComponents[Array | None]):
+    """
+    Value given at initialization is converted to a jnp.array orunmodified if None.
+    This means that at initialization, the user can pass a float or int
+    """
+
+    dyn_loss: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
-    norm_loss: Array | float | None = eqx.field(
+    norm_loss: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
-    boundary_loss: Array | float | None = eqx.field(
+    boundary_loss: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
-    observations: Array | float | None = eqx.field(
+    observations: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )
-    initial_condition: Array | float | None = eqx.field(
+    initial_condition: Array | None = eqx.field(
         kw_only=True, default=None, converter=lw_converter
     )

jinns/nn/_abstract_pinn.py

@@ -3,7 +3,6 @@ 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
 
 
@@ -17,6 +16,5 @@ class AbstractPINN(eqx.Module):
     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

jinns/nn/_hyperpinn.py

@@ -7,18 +7,17 @@ from __future__ import annotations
 
 import warnings
 from dataclasses import InitVar
-from typing import Callable, Literal, Self, Union, Any, cast, overload
+from typing import Callable, Literal, Self, Union, Any, cast
 from math import prod
 import jax
 import jax.numpy as jnp
-from jaxtyping import Array, Float, PyTree, Key
+from jaxtyping import PRNGKeyArray, Array, Float, PyTree
 import equinox as eqx
 import numpy as onp
 
 from jinns.nn._pinn import PINN
 from jinns.nn._mlp import MLP
 from jinns.parameters._params import Params
-from jinns.nn._utils import _PyTree_to_Params
 
 
 def _get_param_nb(
@@ -138,6 +137,32 @@ class HyperPINN(PINN):
             jnp.split(hyper_output, self.pinn_params_cumsum[:-1]),
         )
 
+        # For the record. We exhibited that the jnp.split was a serious time
+        # bottleneck. However none of the approaches below improved the speed.
+        # Moreover, this operation is not well implemented by a triton kernel
+        # apparently so such an optim is not an option.
+        # 1)
+        # pinn_params_flat = jax.tree.unflatten(self.pinn_params_struct,
+        #    jnp.split(hyper_output, self.pinn_params_cumsum[:-1]),
+        # )
+        # 2)
+        # pinn_params_flat = jax.tree.unflatten(self.pinn_params_struct,
+        #    [jax.lax.slice(hyper_output, (s,), (e,)).reshape(r) for s, e, r in
+        #     zip(self.pinn_params_cumsum_start, self.pinn_params_cumsum,
+        #         self.pinn_params_shapes)]
+        # )
+        # 3)
+        # pinn_params_flat = jax.tree.unflatten(self.pinn_params_struct,
+        #    [hyper_output[s:e].reshape(r) for s, e, r in
+        #     zip(self.pinn_params_cumsum_start, self.pinn_params_cumsum,
+        #         self.pinn_params_shapes)]
+        # )
+        # 4)
+        # pinn_params_flat = jax.tree.unflatten(self.pinn_params_struct,
+        #    [jax.lax.dynamic_slice(hyper_output, (s,), (size,)) for s, size in
+        #     zip(self.pinn_params_cumsum_start, self.pinn_params_cumsum_size)]
+        # )
+
         return jax.tree.map(
             lambda a, b: a.reshape(b.shape),
             pinn_params_flat,
@@ -145,17 +170,6 @@ class HyperPINN(PINN):
             is_leaf=lambda x: isinstance(x, jnp.ndarray),
         )
 
-    @overload
-    @_PyTree_to_Params
-    def __call__(
-        self,
-        inputs: Float[Array, " input_dim"],
-        params: PyTree,
-        *args,
-        **kwargs,
-    ) -> Float[Array, " output_dim"]: ...
-
-    @_PyTree_to_Params
     def __call__(
         self,
         inputs: Float[Array, "  input_dim"],
@@ -175,13 +189,10 @@ class HyperPINN(PINN):
             # DataGenerators)
             inputs = inputs[None]
 
-        # try:
         hyper = eqx.combine(params.nn_params, self.static_hyper)
-        # except (KeyError, AttributeError, TypeError) as e:  # give more flexibility
-        #    hyper = eqx.combine(params, self.static_hyper)
 
         eq_params_batch = jnp.concatenate(
-            [params.eq_params[k].flatten() for k in self.hyperparams],
+            [getattr(params.eq_params, k).flatten() for k in self.hyperparams],
             axis=0,
         )
 
@@ -202,12 +213,13 @@ class HyperPINN(PINN):
     @classmethod
     def create(
         cls,
+        *,
         eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
         hyperparams: list[str],
         hypernet_input_size: int,
+        key: PRNGKeyArray | None = None,
         eqx_network: eqx.nn.MLP | MLP | None = None,
         eqx_hyper_network: eqx.nn.MLP | MLP | None = None,
-        key: Key = None,
         eqx_list: tuple[tuple[Callable, int, int] | tuple[Callable], ...] | None = None,
         eqx_list_hyper: (
             tuple[tuple[Callable, int, int] | tuple[Callable], ...] | None
@@ -359,10 +371,10 @@ class HyperPINN(PINN):
 
             ### Now we finetune the hypernetwork architecture
 
-            key, subkey = jax.random.split(key, 2)
+            subkey1, subkey2 = jax.random.split(key, 2)
             # with warnings.catch_warnings():
             #    warnings.filterwarnings("ignore", message="A JAX array is being set as static!")
-            eqx_network = MLP(key=subkey, eqx_list=eqx_list)
+            eqx_network = MLP(key=subkey1, eqx_list=eqx_list)
             # quick partitioning to get the params to get the correct number of neurons
             # for the last layer of hyper network
             params_mlp, _ = eqx.partition(eqx_network, eqx.is_inexact_array)
@@ -405,10 +417,9 @@ class HyperPINN(PINN):
                         + eqx_list_hyper[2:]
                     ),
                 )
-            key, subkey = jax.random.split(key, 2)
             # with warnings.catch_warnings():
             #    warnings.filterwarnings("ignore", message="A JAX array is being set as static!")
-            eqx_hyper_network = cast(MLP, MLP(key=subkey, eqx_list=eqx_list_hyper))
+            eqx_hyper_network = cast(MLP, MLP(key=subkey2, eqx_list=eqx_list_hyper))
 
             ### End of finetuning the hypernetwork architecture
 

jinns/nn/_mlp.py

@@ -9,7 +9,7 @@ from dataclasses import InitVar
 import jax
 import equinox as eqx
 from typing import Protocol
-from jaxtyping import Array, Key, PyTree, Float
+from jaxtyping import Array, PRNGKeyArray, PyTree, Float
 
 from jinns.parameters._params import Params
 from jinns.nn._pinn import PINN
@@ -33,7 +33,7 @@ class MLP(eqx.Module):
 
     Parameters
     ----------
-    key : InitVar[Key]
+    key : InitVar[PRNGKeyArray]
         A jax random key for the layer initializations.
     eqx_list : InitVar[tuple[tuple[Callable, int, int] | tuple[Callable], ...]]
         A tuple of tuples of successive equinox modules and activation functions to
@@ -52,7 +52,7 @@ class MLP(eqx.Module):
         )`.
     """
 
-    key: InitVar[Key] = eqx.field(kw_only=True)
+    key: InitVar[PRNGKeyArray] = eqx.field(kw_only=True)
     eqx_list: InitVar[tuple[tuple[Callable, int, int] | tuple[Callable], ...]] = (
         eqx.field(kw_only=True)
     )
@@ -94,9 +94,10 @@ class PINN_MLP(PINN):
     @classmethod
     def create(
         cls,
+        *,
         eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+        key: PRNGKeyArray | None = None,
         eqx_network: eqx.nn.MLP | MLP | None = None,
-        key: Key = None,
         eqx_list: tuple[tuple[Callable, int, int] | tuple[Callable], ...] | None = None,
         input_transform: (
             Callable[

jinns/nn/_pinn.py

@@ -4,14 +4,13 @@ Implement abstract class for PINN architectures
 
 from __future__ import annotations
 
-from typing import Callable, Union, Any, Literal, overload
+from typing import Callable, Union, Any, Literal
 from dataclasses import InitVar
 import equinox as eqx
 from jaxtyping import Float, Array, PyTree
 import jax.numpy as jnp
 from jinns.parameters._params import Params
 from jinns.nn._abstract_pinn import AbstractPINN
-from jinns.nn._utils import _PyTree_to_Params
 
 
 class PINN(AbstractPINN):
@@ -157,17 +156,6 @@ class PINN(AbstractPINN):
 
         return network(inputs)
 
-    @overload
-    @_PyTree_to_Params
-    def __call__(
-        self,
-        inputs: Float[Array, " input_dim"],
-        params: PyTree,
-        *args,
-        **kwargs,
-    ) -> Float[Array, " output_dim"]: ...
-
-    @_PyTree_to_Params
     def __call__(
         self,
         inputs: Float[Array, " input_dim"],
@@ -180,9 +168,6 @@ class PINN(AbstractPINN):
         `params` and `self.static` to recreate the callable eqx.Module
         architecture. The rest of the content of this function is dependent on
         the network.
-
-        Note that that thanks to the decorator, params can also directly be the
-        PyTree (SPINN, PINN_MLP, ...) that we get out of eqx.combine
         """
 
         if len(inputs.shape) == 0:

jinns/nn/_ppinn.py

@@ -4,18 +4,17 @@ Implements utility function to create PINNs
 
 from __future__ import annotations
 
-from typing import Callable, Literal, Self, cast, overload
+from typing import Callable, Literal, Self, cast
 from dataclasses import InitVar
 import jax
 import jax.numpy as jnp
 import equinox as eqx
 
-from jaxtyping import Array, Key, Float, PyTree
+from jaxtyping import Array, Float, PRNGKeyArray
 
 from jinns.parameters._params import Params
 from jinns.nn._pinn import PINN
 from jinns.nn._mlp import MLP
-from jinns.nn._utils import _PyTree_to_Params
 
 
 class PPINN_MLP(PINN):
@@ -85,17 +84,6 @@ class PPINN_MLP(PINN):
             self.init_params = self.init_params + (params,)
             self.static = self.static + (static,)
 
-    @overload
-    @_PyTree_to_Params
-    def __call__(
-        self,
-        inputs: Float[Array, " input_dim"],
-        params: PyTree,
-        *args,
-        **kwargs,
-    ) -> Float[Array, " output_dim"]: ...
-
-    @_PyTree_to_Params
     def __call__(
         self,
         inputs: Float[Array, " 1"] | Float[Array, " dim"] | Float[Array, " 1+dim"],
@@ -135,9 +123,10 @@ class PPINN_MLP(PINN):
     @classmethod
     def create(
         cls,
+        *,
+        key: PRNGKeyArray | None = None,
         eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
         eqx_network_list: list[eqx.nn.MLP | MLP] | None = None,
-        key: Key = None,
         eqx_list_list: (
             list[tuple[tuple[Callable, int, int] | tuple[Callable], ...]] | None
         ) = None,
@@ -225,7 +214,7 @@ class PPINN_MLP(PINN):
 
             eqx_network_list = []
             for eqx_list in eqx_list_list:
-                key, subkey = jax.random.split(key, 2)
+                key, subkey = jax.random.split(key, 2)  # type: ignore
                 eqx_network_list.append(MLP(key=subkey, eqx_list=eqx_list))
 
         ppinn = cls(