jinns 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

jinns/loss/__init__.py

@@ -1,21 +1,18 @@
 from ._DynamicLossAbstract import DynamicLoss, ODE, PDEStatio, PDENonStatio
-from ._LossODE import LossODE, SystemLossODE
-from ._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
+from ._LossODE import LossODE
+from ._LossPDE import LossPDEStatio, LossPDENonStatio
 from ._DynamicLoss import (
     GeneralizedLotkaVolterra,
     BurgersEquation,
     FPENonStatioLoss2D,
     OU_FPENonStatioLoss2D,
     FisherKPP,
-    MassConservation2DStatio,
-    NavierStokes2DStatio,
+    NavierStokesMassConservation2DStatio,
 )
 from ._loss_weights import (
     LossWeightsODE,
-    LossWeightsODEDict,
     LossWeightsPDENonStatio,
     LossWeightsPDEStatio,
-    LossWeightsPDEDict,
 )
 
 from ._operators import (
@@ -26,3 +23,28 @@ from ._operators import (
     vectorial_laplacian_fwd,
     vectorial_laplacian_rev,
 )
+
+__all__ = [
+    "DynamicLoss",
+    "ODE",
+    "PDEStatio",
+    "PDENonStatio",
+    "LossODE",
+    "LossPDEStatio",
+    "LossPDENonStatio",
+    "GeneralizedLotkaVolterra",
+    "BurgersEquation",
+    "FPENonStatioLoss2D",
+    "OU_FPENonStatioLoss2D",
+    "FisherKPP",
+    "NavierStokesMassConservation2DStatio",
+    "LossWeightsODE",
+    "LossWeightsPDEStatio",
+    "LossWeightsPDENonStatio",
+    "divergence_fwd",
+    "divergence_rev",
+    "laplacian_fwd",
+    "laplacian_rev",
+    "vectorial_laplacian_fwd",
+    "vectorial_laplacian_rev",
+]

jinns/loss/_abstract_loss.py

@@ -0,0 +1,15 @@
+import abc
+from jaxtyping import Array
+import equinox as eqx
+
+
+class AbstractLoss(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/
+    """
+
+    @abc.abstractmethod
+    def __call__(self, *_, **__) -> Array:
+        pass

jinns/loss/_boundary_conditions.py

@@ -7,31 +7,31 @@ from __future__ import (
 )  # https://docs.python.org/3/library/typing.html#constant
 
 from typing import TYPE_CHECKING, Callable
+from jaxtyping import Array, Float
 import jax
 import jax.numpy as jnp
 from jax import vmap, grad
-import equinox as eqx
 from jinns.utils._utils import get_grid, _subtract_with_check
-from jinns.data._Batchs import *
-from jinns.utils._pinn import PINN
-from jinns.utils._spinn import SPINN
+from jinns.data._Batchs import PDEStatioBatch, PDENonStatioBatch
+from jinns.nn._pinn import PINN
+from jinns.nn._spinn import SPINN
 
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.parameters._params import Params
+    from jinns.utils._types import BoundaryConditionFun
+    from jinns.nn._abstract_pinn import AbstractPINN
 
 
 def _compute_boundary_loss(
     boundary_condition_type: str,
-    f: Callable[
-        [Float[Array, "dim"] | Float[Array, "dim + 1"]], Float[Array, "dim_solution"]
-    ],
+    f: BoundaryConditionFun,
     batch: PDEStatioBatch | PDENonStatioBatch,
-    u: eqx.Module,
-    params: AnyParams,
+    u: AbstractPINN,
+    params: Params[Array],
     facet: int,
     dim_to_apply: slice,
     vmap_in_axes: tuple,
-) -> float:
+) -> Float[Array, " "]:
     r"""A generic function that will compute the mini-batch MSE of a
     boundary condition in the stationary case, resp. non-stationary, given by:
 
@@ -67,7 +67,7 @@ def _compute_boundary_loss(
     u
         a PINN
     params
-        Params or ParamsDict
+        Params
     facet
         An integer which represents the id of the facet which is currently
         considered (in the order provided by the DataGenerator which is fixed)
@@ -96,15 +96,15 @@ def _compute_boundary_loss(
 
 def boundary_dirichlet(
     f: Callable[
-        [Float[Array, "dim"] | Float[Array, "dim + 1"]], Float[Array, "dim_solution"]
+        [Float[Array, " dim"] | Float[Array, " dim + 1"]], Float[Array, " dim_solution"]
     ],
     batch: PDEStatioBatch | PDENonStatioBatch,
-    u: eqx.Module,
-    params: Params | ParamsDict,
+    u: AbstractPINN,
+    params: Params[Array],
     facet: int,
     dim_to_apply: slice,
     vmap_in_axes: tuple,
-) -> float:
+) -> Float[Array, " "]:
     r"""
     This omega boundary condition enforces a solution that is equal to `f`
     at `times_batch` x `omega_border` (non stationary case) or at `omega_border`
@@ -135,6 +135,7 @@ def boundary_dirichlet(
     vmap_in_axes
         A tuple object which specifies the in_axes of the vmapping
     """
+    assert batch.border_batch is not None
     batch_array = batch.border_batch
     batch_array = batch_array[..., facet]
 
@@ -168,15 +169,15 @@ def boundary_dirichlet(
 
 def boundary_neumann(
     f: Callable[
-        [Float[Array, "dim"] | Float[Array, "dim + 1"]], Float[Array, "dim_solution"]
+        [Float[Array, " dim"] | Float[Array, " dim + 1"]], Float[Array, " dim_solution"]
     ],
     batch: PDEStatioBatch | PDENonStatioBatch,
-    u: eqx.Module,
-    params: Params | ParamsDict,
+    u: AbstractPINN,
+    params: Params[Array],
     facet: int,
     dim_to_apply: slice,
     vmap_in_axes: tuple,
-) -> float:
+) -> Float[Array, " "]:
     r"""
     This omega boundary condition enforces a solution where $\nabla u\cdot
     n$ is equal to `f` at the cartesian product of `time_batch` x `omega
@@ -208,6 +209,7 @@ def boundary_neumann(
     vmap_in_axes
         A tuple object which specifies the in_axes of the vmapping
     """
+    assert batch.border_batch is not None
     batch_array = batch.border_batch
     batch_array = batch_array[..., facet]
 
@@ -223,7 +225,6 @@ def boundary_neumann(
         n = jnp.array([[-1, 1, 0, 0], [0, 0, -1, 1]])
 
     if isinstance(u, PINN):
-
         u_ = lambda inputs, params: jnp.squeeze(u(inputs, params)[dim_to_apply])
 
         if u.eq_type == "statio_PDE":

jinns/loss/_loss_utils.py

@@ -6,50 +6,53 @@ from __future__ import (
     annotations,
 )  # https://docs.python.org/3/library/typing.html#constant
 
-from typing import TYPE_CHECKING, Callable, Dict
+from typing import TYPE_CHECKING, Callable, TypeGuard
+from types import EllipsisType
 import jax
 import jax.numpy as jnp
 from jax import vmap
-import equinox as eqx
-from jaxtyping import Float, Array, PyTree
+from jaxtyping import Float, Array
 
 from jinns.loss._boundary_conditions import (
     _compute_boundary_loss,
 )
 from jinns.utils._utils import _subtract_with_check, get_grid
-from jinns.data._DataGenerators import append_obs_batch, make_cartesian_product
+from jinns.data._utils import make_cartesian_product
 from jinns.parameters._params import _get_vmap_in_axes_params
-from jinns.utils._pinn import PINN
-from jinns.utils._spinn import SPINN
-from jinns.utils._hyperpinn import HYPERPINN
-from jinns.data._Batchs import *
-from jinns.parameters._params import Params, ParamsDict
+from jinns.nn._pinn import PINN
+from jinns.nn._spinn import SPINN
+from jinns.nn._hyperpinn import HyperPINN
+from jinns.data._Batchs import PDEStatioBatch, PDENonStatioBatch
+from jinns.parameters._params import Params
 
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.utils._types import BoundaryConditionFun
+    from jinns.nn._abstract_pinn import AbstractPINN
 
 
 def dynamic_loss_apply(
-    dyn_loss: DynamicLoss,
-    u: eqx.Module,
+    dyn_loss: Callable,
+    u: AbstractPINN,
     batch: (
-        Float[Array, "batch_size 1"]
-        | Float[Array, "batch_size dim"]
-        | Float[Array, "batch_size 1+dim"]
+        Float[Array, " batch_size 1"]
+        | Float[Array, " batch_size dim"]
+        | Float[Array, " batch_size 1+dim"]
     ),
-    params: Params | ParamsDict,
-    vmap_axes: tuple[int | None, ...],
-    loss_weight: float | Float[Array, "dyn_loss_dimension"],
-    u_type: PINN | HYPERPINN | None = None,
-) -> float:
+    params: Params[Array],
+    vmap_axes: tuple[int, Params[int | None] | None],
+    loss_weight: float | Float[Array, " dyn_loss_dimension"],
+    u_type: PINN | HyperPINN | None = None,
+) -> Float[Array, " "]:
     """
     Sometimes when u is a lambda function a or dict we do not have access to
     its type here, hence the last argument
     """
-    if u_type == PINN or u_type == HYPERPINN or isinstance(u, (PINN, HYPERPINN)):
+    if u_type == PINN or u_type == HyperPINN or isinstance(u, (PINN, HyperPINN)):
         v_dyn_loss = vmap(
             lambda batch, params: dyn_loss(
-                batch, u, params  # we must place the params at the end
+                batch,
+                u,
+                params,  # we must place the params at the end
             ),
             vmap_axes,
             0,
@@ -66,36 +69,38 @@ def dynamic_loss_apply(
 
 
 def normalization_loss_apply(
-    u: eqx.Module,
+    u: AbstractPINN,
     batches: (
-        tuple[Float[Array, "nb_norm_samples dim"]]
+        tuple[Float[Array, " nb_norm_samples dim"]]
         | tuple[
-            Float[Array, "nb_norm_time_slices 1"], Float[Array, "nb_norm_samples dim"]
+            Float[Array, " nb_norm_time_slices 1"], Float[Array, " nb_norm_samples dim"]
         ]
     ),
-    params: Params | ParamsDict,
-    vmap_axes_params: tuple[int | None, ...],
-    int_length: int,
+    params: Params[Array],
+    vmap_axes_params: tuple[Params[int | None] | None],
+    norm_weights: Float[Array, " nb_norm_samples"],
     loss_weight: float,
-) -> float:
+) -> Float[Array, " "]:
     """
     Note the squeezing on each result. We expect unidimensional *PINN since
     they represent probability distributions
     """
-    if isinstance(u, (PINN, HYPERPINN)):
+    if isinstance(u, (PINN, HyperPINN)):
         if len(batches) == 1:
             v_u = vmap(
-                lambda b: u(b)[u.slice_solution],
+                lambda *b: u(*b)[u.slice_solution],
                 (0,) + vmap_axes_params,
                 0,
             )
             res = v_u(*batches, params)
+            assert res.shape[-1] == 1, "norm loss expects unidimensional *PINN"
+            # Monte-Carlo integration using importance sampling
             mse_norm_loss = loss_weight * (
-                jnp.abs(jnp.mean(res.squeeze()) * int_length - 1) ** 2
+                jnp.abs(jnp.mean(res.squeeze() * norm_weights) - 1) ** 2
             )
         else:
             # NOTE this cartesian product is costly
-            batches = make_cartesian_product(
+            batch_cart_prod = make_cartesian_product(
                 batches[0],
                 batches[1],
             ).reshape(batches[0].shape[0], batches[1].shape[0], -1)
@@ -106,21 +111,24 @@ def normalization_loss_apply(
                 ),
                 in_axes=(0,) + vmap_axes_params,
             )
-            res = v_u(batches, params)
-            # Over all the times t, we perform a integration
+            res = v_u(batch_cart_prod, params)
+            assert res.shape[-1] == 1, "norm loss expects unidimensional *PINN"
+            # For all times t, we perform an integration. Then we average the
+            # losses over times.
             mse_norm_loss = loss_weight * jnp.mean(
-                jnp.abs(jnp.mean(res.squeeze(), axis=-1) * int_length - 1) ** 2
+                jnp.abs(jnp.mean(res.squeeze() * norm_weights, axis=-1) - 1) ** 2
             )
     elif isinstance(u, SPINN):
         if len(batches) == 1:
             res = u(*batches, params)
+            assert res.shape[-1] == 1, "norm loss expects unidimensional *SPINN"
             mse_norm_loss = (
                 loss_weight
                 * jnp.abs(
                     jnp.mean(
                         res.squeeze(),
                     )
-                    * int_length
+                    * norm_weights
                     - 1
                 )
                 ** 2
@@ -134,14 +142,15 @@ def normalization_loss_apply(
                 ),
                 params,
             )
+            assert res.shape[-1] == 1, "norm loss expects unidimensional *SPINN"
             # the outer mean() below is for the times stamps
             mse_norm_loss = loss_weight * jnp.mean(
                 jnp.abs(
                     jnp.mean(
                         res.squeeze(),
-                        axis=(d + 1 for d in range(res.ndim - 2)),
+                        axis=list(d + 1 for d in range(res.ndim - 2)),
                     )
-                    * int_length
+                    * norm_weights
                     - 1
                 )
                 ** 2
@@ -153,18 +162,34 @@ def normalization_loss_apply(
 
 
 def boundary_condition_apply(
-    u: eqx.Module,
+    u: AbstractPINN,
     batch: PDEStatioBatch | PDENonStatioBatch,
-    params: Params | ParamsDict,
-    omega_boundary_fun: Callable,
-    omega_boundary_condition: str,
-    omega_boundary_dim: int,
-    loss_weight: float | Float[Array, "boundary_cond_dim"],
-) -> float:
-
+    params: Params[Array],
+    omega_boundary_fun: BoundaryConditionFun | dict[str, BoundaryConditionFun],
+    omega_boundary_condition: str | dict[str, str],
+    omega_boundary_dim: slice | dict[str, slice],
+    loss_weight: float | Float[Array, " boundary_cond_dim"],
+) -> Float[Array, " "]:
+    assert batch.border_batch is not None
     vmap_in_axes = (0,) + _get_vmap_in_axes_params(batch.param_batch_dict, params)
 
-    if isinstance(omega_boundary_fun, dict):
+    def _check_tuple_of_dict(
+        val,
+    ) -> TypeGuard[
+        tuple[
+            dict[str, BoundaryConditionFun],
+            dict[str, BoundaryConditionFun],
+            dict[str, BoundaryConditionFun],
+        ]
+    ]:
+        return all(isinstance(x, dict) for x in val)
+
+    omega_boundary_dicts = (
+        omega_boundary_condition,
+        omega_boundary_fun,
+        omega_boundary_dim,
+    )
+    if _check_tuple_of_dict(omega_boundary_dicts):
         # We must create the facet tree dictionary as we do not have the
         # enumerate from the for loop to pass the id integer
         if batch.border_batch.shape[-1] == 2:
@@ -186,10 +211,10 @@ def boundary_condition_apply(
                     )
                 )
             ),
-            omega_boundary_condition,
-            omega_boundary_fun,
+            omega_boundary_dicts[0],  # omega_boundary_condition,
+            omega_boundary_dicts[1],  # omega_boundary_fun,
             facet_tree,
-            omega_boundary_dim,
+            omega_boundary_dicts[2],  # omega_boundary_dim,
             is_leaf=lambda x: x is None,
         )  # when exploring leaves with None value (no condition) the returned
         # mse is None and we get rid of the None leaves of b_losses_by_facet
@@ -202,13 +227,13 @@ def boundary_condition_apply(
             lambda fa: jnp.mean(
                 loss_weight
                 * _compute_boundary_loss(
-                    omega_boundary_condition,
-                    omega_boundary_fun,
+                    omega_boundary_dicts[0],  # type: ignore -> need TypeIs from 3.13
+                    omega_boundary_dicts[1],  # type: ignore -> need TypeIs from 3.13
                     batch,
                     u,
                     params,
                     fa,
-                    omega_boundary_dim,
+                    omega_boundary_dicts[2],  # type: ignore -> need TypeIs from 3.13
                     vmap_in_axes,
                 )
             ),
@@ -221,22 +246,21 @@ def boundary_condition_apply(
 
 
 def observations_loss_apply(
-    u: eqx.Module,
-    batches: ODEBatch | PDEStatioBatch | PDENonStatioBatch,
-    params: Params | ParamsDict,
-    vmap_axes: tuple[int | None, ...],
-    observed_values: Float[Array, "batch_size observation_dim"],
-    loss_weight: float | Float[Array, "observation_dim"],
-    obs_slice: slice,
-) -> float:
-    # TODO implement for SPINN
-    if isinstance(u, (PINN, HYPERPINN)):
+    u: AbstractPINN,
+    batch: Float[Array, " obs_batch_size input_dim"],
+    params: Params[Array],
+    vmap_axes: tuple[int, Params[int | None] | None],
+    observed_values: Float[Array, " obs_batch_size observation_dim"],
+    loss_weight: float | Float[Array, " observation_dim"],
+    obs_slice: EllipsisType | slice | None,
+) -> Float[Array, " "]:
+    if isinstance(u, (PINN, HyperPINN)):
         v_u = vmap(
             lambda *args: u(*args)[u.slice_solution],
             vmap_axes,
             0,
         )
-        val = v_u(*batches, params)[:, obs_slice]
+        val = v_u(batch, params)[:, obs_slice]
         mse_observation_loss = jnp.mean(
             jnp.sum(
                 loss_weight
@@ -255,16 +279,17 @@ def observations_loss_apply(
 
 
 def initial_condition_apply(
-    u: eqx.Module,
-    omega_batch: Float[Array, "dimension"],
-    params: Params | ParamsDict,
-    vmap_axes: tuple[int | None, ...],
+    u: AbstractPINN,
+    omega_batch: Float[Array, " dimension"],
+    params: Params[Array],
+    vmap_axes: tuple[int, Params[int | None] | None],
     initial_condition_fun: Callable,
-    loss_weight: float | Float[Array, "initial_condition_dimension"],
-) -> float:
+    t0: Float[Array, " 1"],
+    loss_weight: float | Float[Array, " initial_condition_dimension"],
+) -> Float[Array, " "]:
     n = omega_batch.shape[0]
-    t0_omega_batch = jnp.concatenate([jnp.zeros((n, 1)), omega_batch], axis=1)
-    if isinstance(u, (PINN, HYPERPINN)):
+    t0_omega_batch = jnp.concatenate([t0 * jnp.ones((n, 1)), omega_batch], axis=1)
+    if isinstance(u, (PINN, HyperPINN)):
         v_u_t0 = vmap(
             lambda t0_x, params: _subtract_with_check(
                 initial_condition_fun(t0_x[1:]),
@@ -296,103 +321,3 @@ def initial_condition_apply(
     else:
         raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_initial_condition
-
-
-def constraints_system_loss_apply(
-    u_constraints_dict: Dict,
-    batch: ODEBatch | PDEStatioBatch | PDENonStatioBatch,
-    params_dict: ParamsDict,
-    loss_weights: Dict[str, float | Array],
-    loss_weight_struct: PyTree,
-):
-    """
-    Same function for systemlossODE and systemlossPDE!
-    """
-    # Transpose so we have each u_dict as outer structure and the
-    # associated loss_weight as inner structure
-    loss_weights_T = jax.tree_util.tree_transpose(
-        jax.tree_util.tree_structure(loss_weight_struct),
-        jax.tree_util.tree_structure(loss_weights["initial_condition"]),
-        loss_weights,
-    )
-
-    if isinstance(params_dict.nn_params, dict):
-
-        def apply_u_constraint(
-            u_constraint, nn_params, eq_params, loss_weights_for_u, obs_batch_u
-        ):
-            res_dict_for_u = u_constraint.evaluate(
-                Params(
-                    nn_params=nn_params,
-                    eq_params=eq_params,
-                ),
-                append_obs_batch(batch, obs_batch_u),
-            )[1]
-            res_dict_ponderated = jax.tree_util.tree_map(
-                lambda w, l: w * l, res_dict_for_u, loss_weights_for_u
-            )
-            return res_dict_ponderated
-
-        # Note in the case of multiple PINNs, batch.obs_batch_dict is a dict
-        # with keys corresponding to the PINN and value correspondinf to an
-        # original obs_batch_dict. Hence the tree mapping also interates over
-        # batch.obs_batch_dict
-        res_dict = jax.tree_util.tree_map(
-            apply_u_constraint,
-            u_constraints_dict,
-            params_dict.nn_params,
-            (
-                params_dict.eq_params
-                if params_dict.eq_params.keys() == params_dict.nn_params.keys()
-                else {k: params_dict.eq_params for k in params_dict.nn_params.keys()}
-            ),  # this manipulation is needed since we authorize eq_params not to have the same structure as nn_params in ParamsDict
-            loss_weights_T,
-            batch.obs_batch_dict,
-            is_leaf=lambda x: (
-                not isinstance(x, dict)  # to not traverse more than the first
-                # outer dict of the pytrees passed to the function. This will
-                # work because u_constraints_dict is a dict of Losses, and it
-                # thus stops the traversing of other dict too
-            ),
-        )
-    # TODO try to get rid of this condition?
-    else:
-
-        def apply_u_constraint(u_constraint, loss_weights_for_u, obs_batch_u):
-            res_dict_for_u = u_constraint.evaluate(
-                params_dict,
-                append_obs_batch(batch, obs_batch_u),
-            )[1]
-            res_dict_ponderated = jax.tree_util.tree_map(
-                lambda w, l: w * l, res_dict_for_u, loss_weights_for_u
-            )
-            return res_dict_ponderated
-
-        res_dict = jax.tree_util.tree_map(
-            apply_u_constraint, u_constraints_dict, loss_weights_T, batch.obs_batch_dict
-        )
-
-    # Transpose back so we have mses as outer structures and their values
-    # for each u_dict as inner structures. The tree_leaves transforms the
-    # inner structure into a list so we can catch is as leaf it the
-    # tree_map below
-    res_dict = jax.tree_util.tree_transpose(
-        jax.tree_util.tree_structure(
-            jax.tree_util.tree_leaves(loss_weights["initial_condition"])
-        ),
-        jax.tree_util.tree_structure(loss_weight_struct),
-        res_dict,
-    )
-    # For each mse, sum their values on each u_dict
-    res_dict = jax.tree_util.tree_map(
-        lambda mse: jax.tree_util.tree_reduce(
-            lambda x, y: x + y, jax.tree_util.tree_leaves(mse)
-        ),
-        res_dict,
-        is_leaf=lambda x: isinstance(x, list),
-    )
-    # Total loss
-    total_loss = jax.tree_util.tree_reduce(
-        lambda x, y: x + y, jax.tree_util.tree_leaves(res_dict)
-    )
-    return total_loss, res_dict

jinns/loss/_loss_weights.py

@@ -2,58 +2,26 @@
 Formalize the loss weights data structure
 """
 
-from typing import Dict
 from jaxtyping import Array, Float
 import equinox as eqx
 
 
 class LossWeightsODE(eqx.Module):
-
-    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 LossWeightsODEDict(eqx.Module):
-
-    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
-    )
-    observations: Dict[str, Array | Float | None] = eqx.field(
-        kw_only=True, default=None
-    )
+    dyn_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
+    initial_condition: Array | Float = eqx.field(kw_only=True, default=0.0)
+    observations: Array | Float = eqx.field(kw_only=True, default=0.0)
 
 
 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)
+    dyn_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
+    norm_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
+    boundary_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
+    observations: Array | Float = eqx.field(kw_only=True, default=0.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 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
-    )
-    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
-    )
+    dyn_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
+    norm_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
+    boundary_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
+    observations: Array | Float = eqx.field(kw_only=True, default=0.0)
+    initial_condition: Array | Float = eqx.field(kw_only=True, default=0.0)