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

jinns/loss/_boundary_conditions.py

@@ -11,11 +11,7 @@ import jax
 import jax.numpy as jnp
 from jax import vmap, grad
 import equinox as eqx
-from jinns.utils._utils import (
-    _get_grid,
-    _check_user_func_return,
-)
-from jinns.parameters._params import _get_vmap_in_axes_params
+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
@@ -26,12 +22,15 @@ if TYPE_CHECKING:
 
 def _compute_boundary_loss(
     boundary_condition_type: str,
-    f: Callable,
+    f: Callable[
+        [Float[Array, "dim"] | Float[Array, "dim + 1"]], Float[Array, "dim_solution"]
+    ],
     batch: PDEStatioBatch | PDENonStatioBatch,
     u: eqx.Module,
-    params: Params | ParamsDict,
+    params: AnyParams,
     facet: int,
     dim_to_apply: slice,
+    vmap_in_axes: tuple,
 ) -> float:
     r"""A generic function that will compute the mini-batch MSE of a
     boundary condition in the stationary case, resp. non-stationary, given by:
@@ -62,9 +61,9 @@ def _compute_boundary_loss(
         unitary outgoing vector normal to $\partial\Omega$
     f
         the function to be matched in the boundary condition. It should have
-        one or two arguments only (other are ignored).
+        one argument only (for `t`, `x` or `t_x`) (other are ignored).
     batch
-        a PDEStatioBatch or PDENonStatioBatch
+        the batch
     u
         a PINN
     params
@@ -75,233 +74,54 @@ def _compute_boundary_loss(
     dim_to_apply
         A `jnp.s_` object which indicates which dimension(s) of u will be forced
         to match the boundary condition
+    vmap_in_axes
+        A tuple object which specifies the in_axes of the vmapping
 
     Returns
     -------
     scalar
         the MSE computed on `batch`
     """
-    mse = None
-    if isinstance(batch, PDEStatioBatch):
-        if boundary_condition_type.lower() in "dirichlet":
-            mse = boundary_dirichlet_statio(f, batch, u, params, facet, dim_to_apply)
-        elif any(
-            boundary_condition_type.lower() in s
-            for s in ["von neumann", "vn", "vonneumann"]
-        ):
-            mse = boundary_neumann_statio(f, batch, u, params, facet, dim_to_apply)
-    elif isinstance(batch, PDENonStatioBatch):
-        if boundary_condition_type.lower() in "dirichlet":
-            mse = boundary_dirichlet_nonstatio(f, batch, u, params, facet, dim_to_apply)
-        elif any(
-            boundary_condition_type.lower() in s
-            for s in ["von neumann", "vn", "vonneumann"]
-        ):
-            mse = boundary_neumann_nonstatio(f, batch, u, params, facet, dim_to_apply)
+    if boundary_condition_type.lower() in "dirichlet":
+        mse = boundary_dirichlet(f, batch, u, params, facet, dim_to_apply, vmap_in_axes)
+    elif any(
+        boundary_condition_type.lower() in s
+        for s in ["von neumann", "vn", "vonneumann"]
+    ):
+        mse = boundary_neumann(f, batch, u, params, facet, dim_to_apply, vmap_in_axes)
     else:
-        raise ValueError("Wrong type of batch")
+        raise ValueError("Wrong type of initial condition")
     return mse
 
 
-def boundary_dirichlet_statio(
-    f: Callable,
-    batch: PDEStatioBatch,
-    u: eqx.Module,
-    params: Params | ParamsDict,
-    facet: int,
-    dim_to_apply: slice,
-) -> float:
-    r"""
-    This omega boundary condition enforces a solution that is equal to f on
-    border batch.
-
-    __Note__: if using a batch.param_batch_dict, we need to resolve the
-    vmapping axes here however params["eq_params"] has already been fed with
-    the batch in the `evaluate()` of `LossPDEStatio`.
-
-    Parameters
-    ----------
-    f
-        the constraint function
-    batch
-        A PDEStatioBatch object.
-    u
-        The PINN
-    params
-        Params or ParamsDict
-    dim_to_apply
-        A jnp.s\_ object. The dimension of u on which to apply the boundary condition
-    """
-    _, border_batch = batch.inside_batch, batch.border_batch
-    border_batch = border_batch[..., facet]
-
-    if isinstance(u, PINN):
-        vmap_in_axes_params = _get_vmap_in_axes_params(batch.param_batch_dict, params)
-        vmap_in_axes_x = (0,)
-
-        v_u_boundary = vmap(
-            lambda dx, params: u(dx, params)[dim_to_apply] - f(dx),
-            vmap_in_axes_x + vmap_in_axes_params,
-            0,
-        )
-
-        mse_u_boundary = jnp.sum((v_u_boundary(border_batch, params)) ** 2, axis=-1)
-    elif isinstance(u, SPINN):
-        values = u(border_batch, params)[..., dim_to_apply]
-        x_grid = _get_grid(border_batch)
-        boundaries = _check_user_func_return(f(x_grid), values.shape)
-        res = values - boundaries
-        mse_u_boundary = jnp.sum(
-            res**2,
-            axis=-1,
-        )
-    else:
-        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
-    return mse_u_boundary
-
-
-def boundary_neumann_statio(
-    f: Callable,
-    batch: PDEStatioBatch,
-    u: eqx.Module,
-    params: Params | ParamsDict,
-    facet: int,
-    dim_to_apply: slice,
-) -> float:
-    r"""
-    This omega boundary condition enforces a solution where $\nabla u\cdot
-    n$ is equal to `f` on omega borders. $n$ is the unitary
-    outgoing vector normal at border $\partial\Omega$.
-
-    __Note__: if using a batch.param_batch_dict, we need to resolve the
-    vmapping axes here however params["eq_params"] has already been fed with
-    the batch in the `evaluate()` of `LossPDEStatio`.
-
-    Parameters
-    ----------
-    f
-        the constraint function
-    batch
-        A PDEStatioBatch object.
-    u
-        The PINN
-    params
-        The dictionary of parameters of the model.
-        Typically, it is a dictionary of
-        dictionaries: `eq_params` and `nn_params``, respectively the
-        differential equation parameters and the neural network parameter
-    facet
-        An integer which represents the id of the facet which is currently
-        considered (in the order provided wy the DataGenerator which is fixed)
-    dim_to_apply
-        A jnp.s\_ object. The dimension of u on which to apply the boundary condition
-    """
-    _, border_batch = batch.inside_batch, batch.border_batch
-    border_batch = border_batch[..., facet]
-
-    # We resort to the shape of the border_batch to determine the dimension as
-    # described in the border_batch function
-    if jnp.squeeze(border_batch).ndim == 0:  # case 1D borders (just a scalar)
-        n = jnp.array([1, -1])  # the unit vectors normal to the two borders
-
-    else:  # case 2D borders (because 3D borders are not supported yet)
-        # they are in the order: left, right, bottom, top so we give the normal
-        # outgoing vectors accordingly with shape in concordance with
-        # border_batch shape (batch_size, ndim, nfacets)
-        n = jnp.array([[-1, 1, 0, 0], [0, 0, -1, 1]])
-
-    if isinstance(u, PINN):
-        u_ = lambda x, params: jnp.squeeze(u(x, params)[dim_to_apply])
-
-        vmap_in_axes_params = _get_vmap_in_axes_params(batch.param_batch_dict, params)
-        vmap_in_axes_x = (0,)
-
-        v_neumann = vmap(
-            lambda dx, params: jnp.dot(
-                grad(u_, 0)(dx, params),
-                n[..., facet],
-            )
-            - f(dx),
-            vmap_in_axes_x + vmap_in_axes_params,
-            0,
-        )
-        mse_u_boundary = jnp.sum((v_neumann(border_batch, params)) ** 2, axis=-1)
-    elif isinstance(u, SPINN):
-        # the gradient we see in the PINN case can get gradients wrt to x
-        # dimensions at once. But it would be very inefficient in SPINN because
-        # of the high dim output of u. So we do 2 explicit forward AD, handling all the
-        # high dim output at once
-        if border_batch.shape[0] == 1:  # i.e. case 1D
-            _, du_dx = jax.jvp(
-                lambda x: u(
-                    x,
-                    params,
-                )[..., dim_to_apply],
-                (border_batch,),
-                (jnp.ones_like(border_batch),),
-            )
-            values = du_dx * n[facet]
-        elif border_batch.shape[-1] == 2:
-            tangent_vec_0 = jnp.repeat(
-                jnp.array([1.0, 0.0])[None], border_batch.shape[0], axis=0
-            )
-            tangent_vec_1 = jnp.repeat(
-                jnp.array([0.0, 1.0])[None], border_batch.shape[0], axis=0
-            )
-            _, du_dx1 = jax.jvp(
-                lambda x: u(
-                    x,
-                    params,
-                ),
-                (border_batch,),
-                (tangent_vec_0,),
-            )
-            _, du_dx2 = jax.jvp(
-                lambda x: u(
-                    x,
-                    params,
-                ),
-                (border_batch,),
-                (tangent_vec_1,),
-            )
-            values = du_dx1 * n[0, facet] + du_dx2 * n[1, facet]  # dot product
-            # explicitly written
-        else:
-            raise ValueError("Not implemented, we'll do that with a loop")
-
-        x_grid = _get_grid(border_batch)
-        boundaries = _check_user_func_return(f(x_grid), values.shape)
-        res = values - boundaries
-        mse_u_boundary = jnp.sum(res**2, axis=-1)
-    else:
-        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
-    return mse_u_boundary
-
-
-def boundary_dirichlet_nonstatio(
-    f: Callable,
-    batch: PDENonStatioBatch,
+def boundary_dirichlet(
+    f: Callable[
+        [Float[Array, "dim"] | Float[Array, "dim + 1"]], Float[Array, "dim_solution"]
+    ],
+    batch: PDEStatioBatch | PDENonStatioBatch,
     u: eqx.Module,
     params: Params | ParamsDict,
     facet: int,
     dim_to_apply: slice,
+    vmap_in_axes: tuple,
 ) -> float:
     r"""
     This omega boundary condition enforces a solution that is equal to `f`
-    at `times_batch` x `omega borders`
+    at `times_batch` x `omega_border` (non stationary case) or at `omega_border`
+    (stationary case)
 
     __Note__: if using a batch.param_batch_dict, we need to resolve the
     vmapping axes here however params["eq_params"] has already been fed with
-    the batch in the `evaluate()` of `LossPDENonStatio`.
+    the batch in the `evaluate()` of `LossPDE*`.
 
     Parameters
     ----------
     f
         the constraint function
     batch
-        A PDENonStatioBatch object.
+        The batch
     u
-        The PINN
+        The PINN or SPINN
     params
         The dictionary of parameters of the model.
         Typically, it is a dictionary of
@@ -312,49 +132,50 @@ def boundary_dirichlet_nonstatio(
         considered (in the order provided wy the DataGenerator which is fixed)
     dim_to_apply
         A jnp.s\_ object. The dimension of u on which to apply the boundary condition
+    vmap_in_axes
+        A tuple object which specifies the in_axes of the vmapping
     """
-    times_batch = batch.times_x_border_batch[:, 0:1, facet]
-    omega_border_batch = batch.times_x_border_batch[:, 1:, facet]
+    batch_array = batch.border_batch
+    batch_array = batch_array[..., facet]
 
     if isinstance(u, PINN):
-        vmap_in_axes_params = _get_vmap_in_axes_params(batch.param_batch_dict, params)
-        vmap_in_axes_x_t = (0, 0)
-
         v_u_boundary = vmap(
-            lambda t, dx, params: u(
-                t,
-                dx,
-                params,
-            )[dim_to_apply]
-            - f(t, dx),
-            vmap_in_axes_x_t + vmap_in_axes_params,
+            lambda inputs, params: _subtract_with_check(
+                f(inputs),
+                u(
+                    inputs,
+                    params,
+                )[dim_to_apply],
+                cause="boundary condition fun",
+            ),
+            vmap_in_axes,
             0,
         )
-        res = v_u_boundary(times_batch, omega_border_batch, params)
+        res = v_u_boundary(batch_array, params)
         mse_u_boundary = jnp.sum(
             res**2,
             axis=-1,
         )
     elif isinstance(u, SPINN):
-        values = u(times_batch, omega_border_batch, params)[..., dim_to_apply]
-        tx_grid = _get_grid(jnp.concatenate([times_batch, omega_border_batch], axis=-1))
-        boundaries = _check_user_func_return(
-            f(tx_grid[..., 0:1], tx_grid[..., 1:]), values.shape
-        )
-        res = values - boundaries
+        values = u(batch_array, params)[..., dim_to_apply]
+        grid = get_grid(batch_array)
+        res = _subtract_with_check(f(grid), values, cause="boundary condition fun")
         mse_u_boundary = jnp.sum(res**2, axis=-1)
     else:
         raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_u_boundary
 
 
-def boundary_neumann_nonstatio(
-    f: Callable,
-    batch: PDENonStatioBatch,
+def boundary_neumann(
+    f: Callable[
+        [Float[Array, "dim"] | Float[Array, "dim + 1"]], Float[Array, "dim_solution"]
+    ],
+    batch: PDEStatioBatch | PDENonStatioBatch,
     u: eqx.Module,
     params: Params | ParamsDict,
     facet: int,
     dim_to_apply: slice,
+    vmap_in_axes: tuple,
 ) -> float:
     r"""
     This omega boundary condition enforces a solution where $\nabla u\cdot
@@ -371,7 +192,7 @@ def boundary_neumann_nonstatio(
     f:
         the constraint function
     batch
-        A PDENonStatioBatch object.
+        The batch
     u
         The PINN
     params
@@ -384,13 +205,15 @@ def boundary_neumann_nonstatio(
         considered (in the order provided wy the DataGenerator which is fixed)
     dim_to_apply
         A jnp.s\_ object. The dimension of u on which to apply the boundary condition
+    vmap_in_axes
+        A tuple object which specifies the in_axes of the vmapping
     """
-    times_batch = batch.times_x_border_batch[:, 0:1, facet]
-    omega_border_batch = batch.times_x_border_batch[:, 1:, facet]
+    batch_array = batch.border_batch
+    batch_array = batch_array[..., facet]
 
     # We resort to the shape of the border_batch to determine the dimension as
     # described in the border_batch function
-    if jnp.squeeze(omega_border_batch).ndim == 0:  # case 1D borders (just a scalar)
+    if jnp.squeeze(batch_array).ndim == 0:  # case 1D borders (just a scalar)
         n = jnp.array([1, -1])  # the unit vectors normal to the two borders
 
     else:  # case 2D borders (because 3D borders are not supported yet)
@@ -400,24 +223,41 @@ def boundary_neumann_nonstatio(
         n = jnp.array([[-1, 1, 0, 0], [0, 0, -1, 1]])
 
     if isinstance(u, PINN):
-        vmap_in_axes_params = _get_vmap_in_axes_params(batch.param_batch_dict, params)
-        vmap_in_axes_x_t = (0, 0)
 
-        u_ = lambda t, x, params: jnp.squeeze(u(t, x, params)[dim_to_apply])
-        v_neumann = vmap(
-            lambda t, dx, params: jnp.dot(
-                grad(u_, 1)(t, dx, params),
-                n[..., facet],
+        u_ = lambda inputs, params: jnp.squeeze(u(inputs, params)[dim_to_apply])
+
+        if u.eq_type == "statio_PDE":
+            v_neumann = vmap(
+                lambda inputs, params: _subtract_with_check(
+                    f(inputs),
+                    jnp.dot(
+                        grad(u_, 0)(inputs, params),
+                        n[..., facet],
+                    ),
+                    cause="boundary condition fun",
+                ),
+                vmap_in_axes,
+                0,
             )
-            - f(t, dx),
-            vmap_in_axes_x_t + vmap_in_axes_params,
-            0,
-        )
+        elif u.eq_type == "nonstatio_PDE":
+            v_neumann = vmap(
+                lambda inputs, params: _subtract_with_check(
+                    f(inputs),
+                    jnp.dot(
+                        grad(u_, 0)(inputs, params)[1:],  # get rid of time dim
+                        n[..., facet],
+                    ),
+                    cause="boundary condition fun",
+                ),
+                vmap_in_axes,
+                0,
+            )
+        else:
+            raise ValueError("Wrong u.eq_type")
         mse_u_boundary = jnp.sum(
             (
                 v_neumann(
-                    times_batch,
-                    omega_border_batch,
+                    batch_array,
                     params,
                 )
             )
@@ -430,40 +270,69 @@ def boundary_neumann_nonstatio(
         # dimensions at once. But it would be very inefficient in SPINN because
         # of the high dim output of u. So we do 2 explicit forward AD, handling all the
         # high dim output at once
-        if omega_border_batch.shape[0] == 1:  # i.e. case 1D
-            _, du_dx = jax.jvp(
-                lambda x: u(times_batch, x, params)[..., dim_to_apply],
-                (omega_border_batch,),
-                (jnp.ones_like(omega_border_batch),),
-            )
-            values = du_dx * n[facet]
-        elif omega_border_batch.shape[-1] == 2:
-            tangent_vec_0 = jnp.repeat(
-                jnp.array([1.0, 0.0])[None], omega_border_batch.shape[0], axis=0
-            )
-            tangent_vec_1 = jnp.repeat(
-                jnp.array([0.0, 1.0])[None], omega_border_batch.shape[0], axis=0
-            )
-            _, du_dx1 = jax.jvp(
-                lambda x: u(times_batch, x, params)[..., dim_to_apply],
-                (omega_border_batch,),
-                (tangent_vec_0,),
-            )
-            _, du_dx2 = jax.jvp(
-                lambda x: u(times_batch, x, params)[..., dim_to_apply],
-                (omega_border_batch,),
-                (tangent_vec_1,),
-            )
-            values = du_dx1 * n[0, facet] + du_dx2 * n[1, facet]  # dot product
-            # explicitly written
+        if (batch_array.shape[0] == 1 and isinstance(batch, PDEStatioBatch)) or (
+            batch_array.shape[-1] == 2 and isinstance(batch, PDENonStatioBatch)
+        ):
+            if u.eq_type == "statio_PDE":
+                _, du_dx = jax.jvp(
+                    lambda inputs: u(inputs, params)[..., dim_to_apply],
+                    (batch_array,),
+                    (jnp.ones_like(batch_array),),
+                )
+                values = du_dx * n[facet]
+            if u.eq_type == "nonstatio_PDE":
+                _, du_dx = jax.jvp(
+                    lambda inputs: u(inputs, params)[..., dim_to_apply],
+                    (batch_array,),
+                    (jnp.ones_like(batch_array),),
+                )
+                values = du_dx[..., 1] * n[facet]
+        elif (batch_array.shape[-1] == 2 and isinstance(batch, PDEStatioBatch)) or (
+            batch_array.shape[-1] == 3 and isinstance(batch, PDENonStatioBatch)
+        ):
+            if u.eq_type == "statio_PDE":
+                tangent_vec_0 = jnp.repeat(
+                    jnp.array([1.0, 0.0])[None], batch_array.shape[0], axis=0
+                )
+                tangent_vec_1 = jnp.repeat(
+                    jnp.array([0.0, 1.0])[None], batch_array.shape[0], axis=0
+                )
+                _, du_dx1 = jax.jvp(
+                    lambda inputs: u(inputs, params)[..., dim_to_apply],
+                    (batch_array,),
+                    (tangent_vec_0,),
+                )
+                _, du_dx2 = jax.jvp(
+                    lambda inputs: u(inputs, params)[..., dim_to_apply],
+                    (batch_array,),
+                    (tangent_vec_1,),
+                )
+                values = du_dx1 * n[0, facet] + du_dx2 * n[1, facet]  # dot product
+            if u.eq_type == "nonstatio_PDE":
+                tangent_vec_0 = jnp.repeat(
+                    jnp.array([0.0, 1.0, 0.0])[None], batch_array.shape[0], axis=0
+                )
+                tangent_vec_1 = jnp.repeat(
+                    jnp.array([0.0, 0.0, 1.0])[None], batch_array.shape[0], axis=0
+                )
+                _, du_dx1 = jax.jvp(
+                    lambda inputs: u(inputs, params)[..., dim_to_apply],
+                    (batch_array,),
+                    (tangent_vec_0,),
+                )
+                _, du_dx2 = jax.jvp(
+                    lambda inputs: u(inputs, params)[..., dim_to_apply],
+                    (batch_array,),
+                    (tangent_vec_1,),
+                )
+                values = (
+                    du_dx1.squeeze() * n[0, facet] + du_dx2.squeeze() * n[1, facet]
+                )  # dot product
         else:
             raise ValueError("Not implemented, we'll do that with a loop")
 
-        tx_grid = _get_grid(jnp.concatenate([times_batch, omega_border_batch], axis=-1))
-        boundaries = _check_user_func_return(
-            f(tx_grid[..., 0:1], tx_grid[..., 1:]), values.shape
-        )
-        res = values - boundaries
+        grid = get_grid(batch_array)
+        res = _subtract_with_check(f(grid), values, cause="boundary condition fun")
         mse_u_boundary = jnp.sum(
             res**2,
             axis=-1,