jinns 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

jinns/loss/_boundary_conditions.py

@@ -1,6 +1,9 @@
 import jax
 import jax.numpy as jnp
 from jax import vmap, grad
+from jinns.utils._utils import _get_grid, _check_user_func_return
+from jinns.utils._pinn import PINN
+from jinns.utils._spinn import SPINN
 
 
 def _compute_boundary_loss_statio(
@@ -127,18 +130,32 @@ def boundary_dirichlet_statio(f, border_batch, u, params):
         dictionaries: `eq_params` and `nn_params``, respectively the
         differential equation parameters and the neural network parameter
     """
-    v_u_boundary = vmap(
-        lambda dx: u(
-            dx,
-            u_params=params["nn_params"],
-            eq_params=jax.lax.stop_gradient(params["eq_params"]),
+    if isinstance(u, PINN):
+        v_u_boundary = vmap(
+            lambda dx: u(
+                dx,
+                u_params=params["nn_params"],
+                eq_params=jax.lax.stop_gradient(params["eq_params"]),
+            )
+            - f(dx),
+            (0),
+            0,
         )
-        - f(dx),
-        (0),
-        0,
-    )
 
-    mse_u_boundary = jnp.mean((v_u_boundary(border_batch)) ** 2, axis=0)
+        mse_u_boundary = jnp.sum((v_u_boundary(border_batch)) ** 2, axis=-1)
+    elif isinstance(u, SPINN):
+        values = u(
+            border_batch,
+            params["nn_params"],
+            jax.lax.stop_gradient(params["eq_params"]),
+        )
+        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,
+        )
     return mse_u_boundary
 
 
@@ -177,17 +194,69 @@ def boundary_neumann_statio(f, border_batch, u, params, facet):
         # border_batch shape (batch_size, ndim, nfacets)
         n = jnp.array([[-1, 1, 0, 0], [0, 0, -1, 1]])
 
-    v_neumann = vmap(
-        lambda dx: jnp.dot(
-            grad(u, 0)(
-                dx, params["nn_params"], jax.lax.stop_gradient(params["eq_params"])
-            ),
-            n[..., facet],
+    if isinstance(u, PINN):
+        u_ = lambda x, nn, eq: u(t, x, nn, eq)[0]
+        v_neumann = vmap(
+            lambda dx: jnp.dot(
+                grad(u_, 0)(
+                    dx, params["nn_params"], jax.lax.stop_gradient(params["eq_params"])
+                ),
+                n[..., facet],
+            )
+            - f(dx),
+            0,
         )
-        - f(dx),
-        0,
-    )
-    mse_u_boundary = jnp.mean((v_neumann(border_batch)) ** 2, axis=0)
+        mse_u_boundary = jnp.sum((v_neumann(border_batch)) ** 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["nn_params"],
+                    jax.lax.stop_gradient(params["eq_params"]),
+                ),
+                (omega_border_batch,),
+                (jnp.ones_like(x),),
+            )
+            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(
+                    x,
+                    params["nn_params"],
+                    jax.lax.stop_gradient(params["eq_params"]),
+                ),
+                (omega_border_batch,),
+                (tangent_vec_0,),
+            )
+            _, du_dx2 = jax.jvp(
+                lambda x: u(
+                    x,
+                    params["nn_params"],
+                    jax.lax.stop_gradient(params["eq_params"]),
+                ),
+                (omega_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)
     return mse_u_boundary
 
 
@@ -212,30 +281,56 @@ def boundary_dirichlet_nonstatio(f, times_batch, omega_border_batch, u, params):
         dictionaries: `eq_params` and `nn_params``, respectively the
         differential equation parameters and the neural network parameter
     """
-    tile_omega_border_batch = jnp.tile(
-        omega_border_batch, reps=(times_batch.shape[0], 1)
-    )
-
-    def rep_times(k):
-        return jnp.repeat(times_batch, k, axis=0)
-
-    v_u_boundary = vmap(
-        lambda t, dx: u(
-            t,
-            dx,
-            u_params=params["nn_params"],
-            eq_params=jax.lax.stop_gradient(params["eq_params"]),
+    if isinstance(u, PINN):
+        tile_omega_border_batch = jnp.tile(
+            omega_border_batch, reps=(times_batch.shape[0], 1)
+        )
+
+        def rep_times(k):
+            return jnp.repeat(times_batch, k, axis=0)
+
+        v_u_boundary = vmap(
+            lambda t, dx: u(
+                t,
+                dx,
+                u_params=params["nn_params"],
+                eq_params=jax.lax.stop_gradient(params["eq_params"]),
+            )
+            - f(t, dx),
+            (0, 0),
+            0,
+        )
+        res = v_u_boundary(
+            rep_times(omega_border_batch.shape[0]), tile_omega_border_batch
+        )  # TODO check if this cartesian product is always relevant
+        mse_u_boundary = jnp.sum(
+            res**2,
+            axis=-1,
+        )
+    elif isinstance(u, SPINN):
+        tile_omega_border_batch = jnp.tile(
+            omega_border_batch, reps=(times_batch.shape[0], 1)
+        )
+
+        if omega_border_batch.shape[0] == 1:
+            omega_border_batch = jnp.tile(
+                omega_border_batch, reps=(times_batch.shape[0], 1)
+            )
+        # otherwise we require batches to have same shape and we do not need
+        # this operation
+
+        values = u(
+            times_batch,
+            tile_omega_border_batch,
+            params["nn_params"],
+            jax.lax.stop_gradient(params["eq_params"]),
+        )
+        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
         )
-        - f(t, dx),
-        (0, 0),
-        0,
-    )
-
-    mse_u_boundary = jnp.mean(
-        (v_u_boundary(rep_times(omega_border_batch.shape[0]), tile_omega_border_batch))
-        ** 2,
-        axis=0,
-    )
+        res = values - boundaries
+        mse_u_boundary = jnp.sum(res**2, axis=-1)
     return mse_u_boundary
 
 
@@ -264,13 +359,6 @@ def boundary_neumann_nonstatio(f, times_batch, omega_border_batch, u, params, fa
         An integer which represents the id of the facet which is currently
         considered (in the order provided wy the DataGenerator which is fixed)
     """
-    tile_omega_border_batch = jnp.tile(
-        omega_border_batch, reps=(times_batch.shape[0], 1)
-    )
-
-    def rep_times(k):
-        return jnp.repeat(times_batch, k, axis=0)
-
     # 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)
@@ -282,21 +370,99 @@ def boundary_neumann_nonstatio(f, times_batch, omega_border_batch, u, params, fa
         # border_batch shape (batch_size, ndim, nfacets)
         n = jnp.array([[-1, 1, 0, 0], [0, 0, -1, 1]])
 
-    v_neumann = vmap(
-        lambda t, dx: jnp.dot(
-            grad(u, 1)(
-                t, dx, params["nn_params"], jax.lax.stop_gradient(params["eq_params"])
-            ),
-            n[..., facet],
+    if isinstance(u, PINN):
+        tile_omega_border_batch = jnp.tile(
+            omega_border_batch, reps=(times_batch.shape[0], 1)
         )
-        - f(t, dx),
-        0,
-        0,
-    )
-    mse_u_boundary = jnp.mean(
-        (v_neumann(rep_times(omega_border_batch.shape[0]), tile_omega_border_batch))
-        ** 2,
-        axis=0,
-    )
 
+        def rep_times(k):
+            return jnp.repeat(times_batch, k, axis=0)
+
+        u_ = lambda t, x, nn, eq: u(t, x, nn, eq)[0]
+        v_neumann = vmap(
+            lambda t, dx: jnp.dot(
+                grad(u_, 1)(
+                    t,
+                    dx,
+                    params["nn_params"],
+                    jax.lax.stop_gradient(params["eq_params"]),
+                ),
+                n[..., facet],
+            )
+            - f(t, dx),
+            0,
+            0,
+        )
+        mse_u_boundary = jnp.sum(
+            (v_neumann(rep_times(omega_border_batch.shape[0]), tile_omega_border_batch))
+            ** 2,
+            axis=-1,
+        )  # TODO check if this cartesian product is always relevant
+
+    elif isinstance(u, SPINN):
+        if omega_border_batch.shape[0] == 1:
+            omega_border_batch = jnp.tile(
+                omega_border_batch, reps=(times_batch.shape[0], 1)
+            )
+            # ie case 1D
+        # otherwise we require batches to have same shape and we do not need
+        # this operation
+
+        # 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 omega_border_batch.shape[0] == 1:  # i.e. case 1D
+            _, du_dx = jax.jvp(
+                lambda x: u(
+                    times_batch,
+                    x,
+                    params["nn_params"],
+                    jax.lax.stop_gradient(params["eq_params"]),
+                ),
+                (omega_border_batch,),
+                (jnp.ones_like(x),),
+            )
+            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["nn_params"],
+                    jax.lax.stop_gradient(params["eq_params"]),
+                ),
+                (omega_border_batch,),
+                (tangent_vec_0,),
+            )
+            _, du_dx2 = jax.jvp(
+                lambda x: u(
+                    times_batch,
+                    x,
+                    params["nn_params"],
+                    jax.lax.stop_gradient(params["eq_params"]),
+                ),
+                (omega_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")
+
+        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
+        mse_u_boundary = jnp.sum(
+            res**2,
+            axis=-1,
+        )
     return mse_u_boundary

jinns/loss/_operators.py

@@ -2,13 +2,16 @@ import jax
 import jax.numpy as jnp
 from jax import grad
 from functools import partial
+from jinns.utils._pinn import PINN
+from jinns.utils._spinn import SPINN
 
 
-def _div(u, nn_params, eq_params, x, t=None):
+def _div_rev(u, nn_params, eq_params, x, t=None):
     r"""
-    Compute the divergence of a vector field :math:`\mathbf{u}` ie
-    :math:`\nabla \cdot u(x)` with u a vector field from :math:`\mathbb{R}^n`
-    to :math:`\mathbb{R}^n`
+    Compute the divergence of a vector field :math:`\mathbf{u}`, i.e.,
+    :math:`\nabla \cdot \mathbf{u}(\mathbf{x})` with :math:`\mathbf{u}` a vector
+    field from :math:`\mathbb{R}^d` to :math:`\mathbb{R}^d`.
+    The computation is done using backward AD
     """
 
     def scan_fun(_, i):
@@ -26,42 +29,131 @@ def _div(u, nn_params, eq_params, x, t=None):
     return jnp.sum(accu)
 
 
-def _laplacian(u, nn_params, eq_params, x, t=None):
+def _div_fwd(u, nn_params, eq_params, x, t=None):
     r"""
-    Compute the Laplacian of a scalar field u (from :math:`\mathbb{R}^n`
-    to :math:`\mathbb{R}`) for x of arbitrary dimension ie
-    :math:`\Delta u(x)=\nabla\cdot\nabla u(x)`
-    For computational reason we do not compute the trace of the Hessian but
-    we explicitly call the gradient twice
+    Compute the divergence of a **batched** vector field :math:`\mathbf{u}`, i.e.,
+    :math:`\nabla \cdot \mathbf{u}(\mathbf{x})` with :math:`\mathbf{u}` a vector
+    field from :math:`\mathbb{R}^{b \times d}` to :math:`\mathbb{R}^{b \times b
+    \times d}`. The result is then in :math:`\mathbb{R}^{b\times b}`.
+    Because of the embedding that happens in SPINNs the
+    computation is most efficient with forward AD. This is the idea behind Separable PINNs.
+    This function is to be used in the context of SPINNs only.
     """
 
     def scan_fun(_, i):
+        tangent_vec = jnp.repeat(
+            jax.nn.one_hot(i, x.shape[-1])[None], x.shape[0], axis=0
+        )
         if t is None:
-            d2u_dxi2 = grad(
-                lambda x, nn_params, eq_params: grad(u, 0)(x, nn_params, eq_params)[i],
-                0,
-            )(x, nn_params, eq_params)[i]
+            __, du_dxi = jax.jvp(
+                lambda x: u(x, nn_params, eq_params)[..., i], (x,), (tangent_vec,)
+            )
         else:
-            d2u_dxi2 = grad(
-                lambda t, x, nn_params, eq_params: grad(u, 1)(
-                    t, x, nn_params, eq_params
-                )[i],
-                1,
-            )(t, x, nn_params, eq_params)[i]
+            __, du_dxi = jax.jvp(
+                lambda x: u(t, x, nn_params, eq_params)[..., i], (x,), (tangent_vec,)
+            )
+        return _, du_dxi
+
+    _, accu = jax.lax.scan(scan_fun, {}, jnp.arange(x.shape[1]))
+    return jnp.sum(accu, axis=0)
+
+
+def _laplacian_rev(u, nn_params, eq_params, x, t=None):
+    r"""
+    Compute the Laplacian of a scalar field :math:`u` (from :math:`\mathbb{R}^d`
+    to :math:`\mathbb{R}`) for :math:`\mathbf{x}` of arbitrary dimension, i.e.,
+    :math:`\Delta u(\mathbf{x})=\nabla\cdot\nabla u(\mathbf{x})`.
+    The computation is done using backward AD.
+    """
+
+    # Note that the last dim of u is nec. 1
+    if t is None:
+        u_ = lambda x: u(x, nn_params, eq_params)[0]
+    else:
+        u_ = lambda t, x: u(t, x, nn_params, eq_params)[0]
+
+    if t is None:
+        return jnp.trace(jax.hessian(u_)(x))
+    else:
+        return jnp.trace(jax.hessian(u_, argnums=1)(t, x))
+
+    # For a small d, we found out that trace of the Hessian is faster, but the
+    # trick below for taking directly the diagonal elements might prove useful
+    # in higher dimensions?
+
+    # def scan_fun(_, i):
+    #    if t is None:
+    #        d2u_dxi2 = grad(
+    #            lambda x: grad(u_, 0)(x)[i],
+    #            0,
+    #        )(
+    #            x
+    #        )[i]
+    #    else:
+    #        d2u_dxi2 = grad(
+    #            lambda t, x: grad(u_, 1)(t, x)[i],
+    #            1,
+    #        )(
+    #            t, x
+    #        )[i]
+    #    return _, d2u_dxi2
+
+    # _, trace_hessian = jax.lax.scan(scan_fun, {}, jnp.arange(x.shape[0]))
+    # return jnp.sum(trace_hessian)
+
+
+def _laplacian_fwd(u, nn_params, eq_params, x, t=None):
+    r"""
+    Compute the Laplacian of a **batched** scalar field :math:`u`
+    (from :math:`\mathbb{R}^{b\times d}` to :math:`\mathbb{R}^{b\times b}`)
+    for :math:`\mathbf{x}` of arbitrary dimension :math:`d` with batch
+    dimension :math:`b`.
+    Because of the embedding that happens in SPINNs the
+    computation is most efficient with forward AD. This is the idea behind Separable PINNs.
+    This function is to be used in the context of SPINNs only.
+    """
+
+    def scan_fun(_, i):
+        tangent_vec = jnp.repeat(
+            jax.nn.one_hot(i, x.shape[-1])[None], x.shape[0], axis=0
+        )
+
+        if t is None:
+            du_dxi_fun = lambda x: jax.jvp(
+                lambda x: u(x, nn_params, eq_params)[..., 0], (x,), (tangent_vec,)
+            )[
+                1
+            ]  # Note the indexing [..., 0]
+            __, d2u_dxi2 = jax.jvp(du_dxi_fun, (x,), (tangent_vec,))
+        else:
+            du_dxi_fun = lambda x: jax.jvp(
+                lambda x: u(t, x, nn_params, eq_params)[..., 0], (x,), (tangent_vec,)
+            )[
+                1
+            ]  # Note the indexing [..., 0]
+            __, d2u_dxi2 = jax.jvp(du_dxi_fun, (x,), (tangent_vec,))
         return _, d2u_dxi2
 
-    _, trace_hessian = jax.lax.scan(scan_fun, {}, jnp.arange(x.shape[0]))
-    return jnp.sum(trace_hessian)
+    _, trace_hessian = jax.lax.scan(scan_fun, {}, jnp.arange(x.shape[1]))
+    return jnp.sum(trace_hessian, axis=0)
 
 
 def _vectorial_laplacian(u, nn_params, eq_params, x, t=None, u_vec_ndim=None):
     r"""
-    Compute the vectorial Laplacian of a vector field u (from :math:`\mathbb{R}^m`
-    to :math:`\mathbb{R}^n`) for x of arbitrary dimension ie
-    :math:`\Delta \mathbf{u}(x)=\nabla\cdot\nabla \mathbf{u}(x)`
+    Compute the vectorial Laplacian of a vector field :math:`\mathbf{u}` (from
+    :math:`\mathbb{R}^d`
+    to :math:`\mathbb{R}^n`) for :math:`\mathbf{x}` of arbitrary dimension, i.e.,
+    :math:`\Delta \mathbf{u}(\mathbf{x})=\nabla\cdot\nabla
+    \mathbf{u}(\mathbf{x})`.
+
+    **Note:** We need to provide `u_vec_ndim` the dimension of the vector
+    :math:`\mathbf{u}(\mathbf{x})` if it is different than that of
+    :math:`\mathbf{x}`.
 
-    **Note:** We need to provide in u_vec_ndim the dimension of the vector
-    :math:`\mathbf{u}(x)` if it is different than that of x
+    **Note:** `u` can be a SPINN, in this case, it corresponds to a vector
+    field from (from :math:`\mathbb{R}^{b\times d}` to
+    :math:`\mathbb{R}^{b\times b\times n}`) and forward mode AD is used.
+    Technically, the return is of dimension :math:`n\times b \times b`.
     """
     if u_vec_ndim is None:
         u_vec_ndim = x.shape[0]
@@ -69,25 +161,42 @@ def _vectorial_laplacian(u, nn_params, eq_params, x, t=None, u_vec_ndim=None):
     def scan_fun(_, j):
         # The loop over the components of u(x). We compute one Laplacian for
         # each of these components
-        if t is None:
-            uj = lambda x, nn_params, eq_params: u(x, nn_params, eq_params)[j]
-        else:
-            uj = lambda t, x, nn_params, eq_params: u(t, x, nn_params, eq_params)[j]
-        lap_on_j = _laplacian(uj, nn_params, eq_params, x, t)
+        # Note the expand_dims
+        if isinstance(u, PINN):
+            if t is None:
+                uj = lambda x, nn_params, eq_params: jnp.expand_dims(
+                    u(x, nn_params, eq_params)[j], axis=-1
+                )
+            else:
+                uj = lambda t, x, nn_params, eq_params: jnp.expand_dims(
+                    u(t, x, nn_params, eq_params)[j], axis=-1
+                )
+            lap_on_j = _laplacian_rev(uj, nn_params, eq_params, x, t)
+        elif isinstance(u, SPINN):
+            if t is None:
+                uj = lambda x, nn_params, eq_params: jnp.expand_dims(
+                    u(x, nn_params, eq_params)[..., j], axis=-1
+                )
+            else:
+                uj = lambda t, x, nn_params, eq_params: jnp.expand_dims(
+                    u(t, x, nn_params, eq_params)[..., j], axis=-1
+                )
+            lap_on_j = _laplacian_fwd(uj, nn_params, eq_params, x, t)
+
         return _, lap_on_j
 
     _, vec_lap = jax.lax.scan(scan_fun, {}, jnp.arange(u_vec_ndim))
     return vec_lap
 
 
-def _u_dot_nabla_times_u(u, nn_params, eq_params, x, t=None):
+def _u_dot_nabla_times_u_rev(u, nn_params, eq_params, x, t=None):
     r"""
-    Implement :math:`((\mathbf{u}\cdot\nabla)\mathbf{u})(x)` for x of arbitrary
-    dimension. Note that :math:`\mathbf{u}` is a vector field from :math:`\mathbb{R}^n`
-    to :math:`\mathbb{R}^n`
-    Currently for `x.ndim=2`
-
-    **Note:** We do not use loops but code explicitly the expression to avoid
+    Implement :math:`((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})` for
+    :math:`\mathbf{x}` of arbitrary
+    dimension. :math:`\mathbf{u}` is a vector field from :math:`\mathbb{R}^n`
+    to :math:`\mathbb{R}^n`. **Currently for** `x.ndim=2` **only**.
+    The computation is done using backward AD.
+    We do not use loops but code explicitly the expression to avoid
     computing twice some terms
     """
     if x.shape[0] == 2:
@@ -127,17 +236,64 @@ def _u_dot_nabla_times_u(u, nn_params, eq_params, x, t=None):
         raise NotImplementedError("x.ndim must be 2")
 
 
+def _u_dot_nabla_times_u_fwd(u, nn_params, eq_params, x, t=None):
+    r"""
+    Implement :math:`((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})` for
+    :math:`\mathbf{x}` of arbitrary dimension **with a batch dimension**.
+    I.e., :math:`\mathbf{u}` is a vector field from :math:`\mathbb{R}^{b\times
+    b}`
+    to :math:`\mathbb{R}^{b\times b \times d}`. **Currently for** :math:`d=2`
+    **only**.
+    We do not use loops but code explicitly the expression to avoid
+    computing twice some terms.
+    Because of the embedding that happens in SPINNs the
+    computation is most efficient with forward AD. This is the idea behind Separable PINNs.
+    This function is to be used in the context of SPINNs only.
+    """
+    if x.shape[-1] == 2:
+        tangent_vec_0 = jnp.repeat(jnp.array([1.0, 0.0])[None], x.shape[0], axis=0)
+        tangent_vec_1 = jnp.repeat(jnp.array([0.0, 1.0])[None], x.shape[0], axis=0)
+        if t is None:
+            u_at_x, du_dx = jax.jvp(
+                lambda x: u(x, nn_params, eq_params), (x,), (tangent_vec_0,)
+            )  # thanks to forward AD this gets dux_dx and duy_dx in a vector
+            # ie the derivatives of both components of u wrt x
+            # this also gets the vector of u evaluated at x
+            u_at_x, du_dy = jax.jvp(
+                lambda x: u(x, nn_params, eq_params), (x,), (tangent_vec_1,)
+            )  # thanks to forward AD this gets dux_dy and duy_dy in a vector
+            # ie the derivatives of both components of u wrt y
+
+        else:
+            u_at_x, du_dx = jax.jvp(
+                lambda x: u(t, x, nn_params, eq_params), (x,), (tangent_vec_0,)
+            )
+            u_at_x, du_dy = jax.jvp(
+                lambda x: u(t, x, nn_params, eq_params), (x,), (tangent_vec_1,)
+            )
+
+        return jnp.stack(
+            [
+                u_at_x[..., 0] * du_dx[..., 0] + u_at_x[..., 1] * du_dy[..., 0],
+                u_at_x[..., 0] * du_dx[..., 1] + u_at_x[..., 1] * du_dy[..., 1],
+            ],
+            axis=-1,
+        )
+    else:
+        raise NotImplementedError("x.ndim must be 2")
+
+
 def _sobolev(u, m, statio=True):
     r"""
-    Compute the Sobolev regularization of order m
-    of a scalar field u (from :math:`\mathbb{R}^d1` to :math:`\mathbb{R}`)
-    for x of arbitrary dimension i.e.
+    Compute the Sobolev regularization of order :math:`m`
+    of a scalar field :math:`u` (from :math:`\mathbb{R}^{d}` to :math:`\mathbb{R}`)
+    for :math:`\mathbf{x}` of arbitrary dimension :math:`d`, i.e.,
     :math:`\frac{1}{n_l}\sum_{l=1}^{n_l}\sum_{|\alpha|=1}^{m+1} ||\partial^{\alpha} u(x_l)||_2^2` where
-    :math:`m\geq\max(d_1 // 2, K)` with `K` the order of the differential
+    :math:`m\geq\max(d_1 // 2, K)` with :math:`K` the order of the differential
     operator.
 
-    This regularization is proposed in _Convergence and error analysis of
-    PINNs_, Doumeche et al., 2023, https://arxiv.org/pdf/2305.01240.pdf
+    This regularization is proposed in *Convergence and error analysis of
+    PINNs*, Doumeche et al., 2023, https://arxiv.org/pdf/2305.01240.pdf
     """
 
     def jac_recursive(u, order, start):

jinns/utils/__init__.py

@@ -3,5 +3,6 @@ from ._utils import (
     euler_maruyama_density,
     log_euler_maruyama_density,
     alternate_optax_solver,
-    create_PINN,
 )
+from ._pinn import create_PINN
+from ._spinn import create_SPINN