fluxfem 0.1.4__py3-none-any.whl → 0.2.1__py3-none-any.whl

fluxfem/core/basis.py

@@ -4,8 +4,7 @@ from typing import Protocol
 import jax
 import jax.numpy as jnp
 import numpy as np
-from .dtypes import DEFAULT_DTYPE as _FDTYPE
-# from .dtypes import DEFAULT_DTYPE
+from .dtypes import default_dtype
 
 
 def build_B_matrices(dN_dx: jnp.ndarray) -> jnp.ndarray:
@@ -271,6 +270,7 @@ class TetLinearBasis(SmallStrainBMixin, TotalLagrangeBMixin):
 
     @property
     def ref_node_coords(self) -> jnp.ndarray:
+        dtype = default_dtype()
         return jnp.array(
             [
                 [0.0, 0.0, 0.0],
@@ -278,7 +278,7 @@ class TetLinearBasis(SmallStrainBMixin, TotalLagrangeBMixin):
                 [0.0, 1.0, 0.0],
                 [0.0, 0.0, 1.0],
             ],
-            dtype=_FDTYPE,
+            dtype=dtype,
         )
 
     def shape_functions(self) -> jnp.ndarray:
@@ -294,6 +294,7 @@ class TetLinearBasis(SmallStrainBMixin, TotalLagrangeBMixin):
 
     def shape_grads_ref(self) -> jnp.ndarray:
         # constant gradients in reference tetra
+        dtype = default_dtype()
         dN = jnp.array(
             [
                 [-1.0, -1.0, -1.0],
@@ -301,7 +302,7 @@ class TetLinearBasis(SmallStrainBMixin, TotalLagrangeBMixin):
                 [0.0, 1.0, 0.0],
                 [0.0, 0.0, 1.0],
             ],
-            dtype=_FDTYPE,
+            dtype=dtype,
         )
         dN = jnp.tile(dN[None, :, :], (self.n_q, 1, 1))  # (n_q,4,3)
         return dN
@@ -374,7 +375,7 @@ class TetQuadraticBasis10(SmallStrainBMixin, TotalLagrangeBMixin):
 
         grads = []
         for a, dLa in zip([L1, L2, L3, L4], [dL1, dL2, dL3, dL4]):
-            grads.append((2 * a - 1)[..., None] * dLa[None, :])
+            grads.append((4 * a - 1)[..., None] * dLa[None, :])
 
         dN1 = grads[0]
         dN2 = grads[1]
@@ -420,7 +421,7 @@ class HexTriLinearBasis(SmallStrainBMixin, TotalLagrangeBMixin):
 
     def tree_flatten(self):
         children = (self.quad_points, self.quad_weights)
-        aux_data = {}
+        aux_data: dict[str, object] = {}
         return children, aux_data
 
     @classmethod
@@ -446,6 +447,7 @@ class HexTriLinearBasis(SmallStrainBMixin, TotalLagrangeBMixin):
         6: ( 1, 1, 1)
         7: (-1, 1, 1)
         """
+        dtype = default_dtype()
         return jnp.array(
             [
                 [-1.0, -1.0, -1.0],
@@ -457,7 +459,7 @@ class HexTriLinearBasis(SmallStrainBMixin, TotalLagrangeBMixin):
                 [ 1.0,  1.0,  1.0],
                 [-1.0,  1.0,  1.0],
             ],
-            dtype=_FDTYPE,
+            dtype=dtype,
         )
 
     # ---------- reference shape functions & gradients ----------
@@ -567,7 +569,7 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
 
     def tree_flatten(self):
         children = (self.quad_points, self.quad_weights)
-        aux_data = {}
+        aux_data: dict[str, object] = {}
         return children, aux_data
 
     @classmethod
@@ -581,6 +583,7 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
 
     @property
     def ref_node_coords(self) -> jnp.ndarray:
+        dtype = default_dtype()
         corners = jnp.array(
             [
                 [-1.0, -1.0, -1.0],
@@ -592,7 +595,7 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
                 [ 1.0,  1.0,  1.0],
                 [-1.0,  1.0,  1.0],
             ],
-            dtype=_FDTYPE,
+            dtype=dtype,
         )
         edges = jnp.array(
             [
@@ -609,7 +612,7 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
                 [ 1.0,  1.0,  0.0],  # 2-6
                 [-1.0,  1.0,  0.0],  # 3-7
             ],
-            dtype=_FDTYPE,
+            dtype=dtype,
         )
         return jnp.concatenate([corners, edges], axis=0)  # (20,3)
 
@@ -620,6 +623,7 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
         zeta = qp[:, 2:3]
 
         # corners
+        dtype = default_dtype()
         s = jnp.array(
             [
                 [-1, -1, -1],
@@ -631,7 +635,7 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
                 [ 1,  1,  1],
                 [-1,  1,  1],
             ],
-            dtype=_FDTYPE,
+            dtype=dtype,
         )
         sx = s[:, 0]
         sy = s[:, 1]
@@ -639,22 +643,30 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
         term = xi * sx + eta * sy + zeta * sz - 2.0  # (n_q,8)
         N_corner = 0.125 * (1 + sx * xi) * (1 + sy * eta) * (1 + sz * zeta) * term  # (n_q,8)
 
-        # edges
-        edges_x = [(-1, -1), (1, -1), (1, 1), (-1, 1)]  # eta, zeta fixed
-        edges_y = [(-1, -1), (1, -1), (1, 1), (-1, 1)]  # xi fixed
-        edges_z = [(-1, -1), (1, -1), (1, 1), (-1, 1)]  # xi, eta fixed
-
-        N_edges = []
-        # along xi (1 - xi^2)
-        for sy, sz in edges_x:
-            N_edges.append(0.25 * (1 - xi * xi) * (1 + sy * eta) * (1 + sz * zeta))
-        # along eta
-        for sx, sz in edges_y:
-            N_edges.append(0.25 * (1 - eta * eta) * (1 + sx * xi) * (1 + sz * zeta))
-        # along zeta
-        for sx, sy in edges_z:
-            N_edges.append(0.25 * (1 - zeta * zeta) * (1 + sx * xi) * (1 + sy * eta))
-
+        # edges: order matches hex20 connectivity (e01, e12, e23, e30, e45, e56, e67, e74, e04, e15, e26, e37)
+        def edge_x(sy, sz):
+            return 0.25 * (1 - xi * xi) * (1 + sy * eta) * (1 + sz * zeta)
+
+        def edge_y(sx, sz):
+            return 0.25 * (1 - eta * eta) * (1 + sx * xi) * (1 + sz * zeta)
+
+        def edge_z(sx, sy):
+            return 0.25 * (1 - zeta * zeta) * (1 + sx * xi) * (1 + sy * eta)
+
+        N_edges = [
+            edge_x(-1, -1),
+            edge_y(1, -1),
+            edge_x(1, -1),
+            edge_y(-1, -1),
+            edge_x(-1, 1),
+            edge_y(1, 1),
+            edge_x(1, 1),
+            edge_y(-1, 1),
+            edge_z(-1, -1),
+            edge_z(1, -1),
+            edge_z(1, 1),
+            edge_z(-1, 1),
+        ]
         N_edges = jnp.concatenate(N_edges, axis=1)  # (n_q, 12)
         return jnp.concatenate([N_corner, N_edges], axis=1)  # (n_q,20)
 
@@ -664,6 +676,7 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
         eta = qp[:, 1:2]
         zeta = qp[:, 2:3]
 
+        dtype = default_dtype()
         s = jnp.array(
             [
                 [-1, -1, -1],
@@ -675,7 +688,7 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
                 [ 1,  1,  1],
                 [-1,  1,  1],
             ],
-            dtype=_FDTYPE,
+            dtype=dtype,
         )
         sx = s[:, 0]
         sy = s[:, 1]
@@ -696,34 +709,38 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
         )  # (n_q,8,3)
 
         # edges derivatives
-        d_list = []
-        # along xi
-        edges_x = [(-1, -1), (1, -1), (1, 1), (-1, 1)]
-        for sy_val, sz_val in edges_x:
-            sy_ = sy_val
-            sz_ = sz_val
+        def d_edge_x(sy_, sz_):
             dxi = -0.5 * xi * (1 + sy_ * eta) * (1 + sz_ * zeta)
             deta = 0.25 * (1 - xi * xi) * sy_ * (1 + sz_ * zeta)
             dzeta = 0.25 * (1 - xi * xi) * (1 + sy_ * eta) * sz_
-            d_list.append(jnp.stack([dxi, deta, dzeta], axis=2))
-        # along eta
-        edges_y = [(-1, -1), (1, -1), (1, 1), (-1, 1)]
-        for sx_val, sz_val in edges_y:
-            sx_ = sx_val
-            sz_ = sz_val
+            return jnp.stack([dxi, deta, dzeta], axis=2)
+
+        def d_edge_y(sx_, sz_):
             dxi = 0.25 * (1 - eta * eta) * sx_ * (1 + sz_ * zeta)
             deta = -0.5 * eta * (1 + sx_ * xi) * (1 + sz_ * zeta)
             dzeta = 0.25 * (1 - eta * eta) * (1 + sx_ * xi) * sz_
-            d_list.append(jnp.stack([dxi, deta, dzeta], axis=2))
-        # along zeta
-        edges_z = [(-1, -1), (1, -1), (1, 1), (-1, 1)]
-        for sx_val, sy_val in edges_z:
-            sx_ = sx_val
-            sy_ = sy_val
+            return jnp.stack([dxi, deta, dzeta], axis=2)
+
+        def d_edge_z(sx_, sy_):
             dxi = 0.25 * (1 - zeta * zeta) * sx_ * (1 + sy_ * eta)
             deta = 0.25 * (1 - zeta * zeta) * (1 + sx_ * xi) * sy_
             dzeta = -0.5 * zeta * (1 + sx_ * xi) * (1 + sy_ * eta)
-            d_list.append(jnp.stack([dxi, deta, dzeta], axis=2))
+            return jnp.stack([dxi, deta, dzeta], axis=2)
+
+        d_list = [
+            d_edge_x(-1, -1),
+            d_edge_y(1, -1),
+            d_edge_x(1, -1),
+            d_edge_y(-1, -1),
+            d_edge_x(-1, 1),
+            d_edge_y(1, 1),
+            d_edge_x(1, 1),
+            d_edge_y(-1, 1),
+            d_edge_z(-1, -1),
+            d_edge_z(1, -1),
+            d_edge_z(1, 1),
+            d_edge_z(-1, 1),
+        ]
 
         d_edges = jnp.concatenate(d_list, axis=1)  # (n_q,12,3)
         return jnp.concatenate([d_corner, d_edges], axis=1)  # (n_q,20,3)
@@ -846,7 +863,8 @@ def _gauss_legendre_1d(order: int) -> tuple[jnp.ndarray, jnp.ndarray]:
     if order <= 0:
         raise ValueError("quadrature order must be positive")
     pts, wts = np.polynomial.legendre.leggauss(order)
-    return jnp.array(pts, dtype=_FDTYPE), jnp.array(wts, dtype=_FDTYPE)
+    dtype = default_dtype()
+    return jnp.array(pts, dtype=dtype), jnp.array(wts, dtype=dtype)
 
 
 def _gl_points_for_degree(degree: int) -> int:
@@ -865,10 +883,12 @@ def _tet_quadrature(degree: int) -> tuple[jnp.ndarray, jnp.ndarray]:
     degree<=1: 1-point; degree<=2: 4-point; degree>=3: 5-point (Stroud T3-5).
     """
     if degree <= 1:
-        qp = jnp.array([[0.25, 0.25, 0.25]], dtype=_FDTYPE)
-        qw = jnp.array([1.0 / 6.0], dtype=_FDTYPE)
+        dtype = default_dtype()
+        qp = jnp.array([[0.25, 0.25, 0.25]], dtype=dtype)
+        qw = jnp.array([1.0 / 6.0], dtype=dtype)
         return qp, qw
     if degree <= 2:
+        dtype = default_dtype()
         qp = jnp.array(
             [
                 [0.58541020, 0.13819660, 0.13819660],
@@ -876,11 +896,12 @@ def _tet_quadrature(degree: int) -> tuple[jnp.ndarray, jnp.ndarray]:
                 [0.13819660, 0.13819660, 0.58541020],
                 [0.13819660, 0.13819660, 0.13819660],
             ],
-            dtype=_FDTYPE,
+            dtype=dtype,
         )
-        qw = jnp.full((4,), (1.0 / 24.0), dtype=_FDTYPE)
+        qw = jnp.full((4,), (1.0 / 24.0), dtype=dtype)
         return qp, qw
     # degree 3 rule: centroid + 4 symmetric points
+    dtype = default_dtype()
     qp = jnp.array(
         [
             [0.25, 0.25, 0.25],
@@ -889,11 +910,11 @@ def _tet_quadrature(degree: int) -> tuple[jnp.ndarray, jnp.ndarray]:
             [1.0 / 6.0, 1.0 / 6.0, 0.50],
             [1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0],
         ],
-        dtype=_FDTYPE,
+        dtype=dtype,
     )
     qw = jnp.array(
         [-2.0 / 15.0, 3.0 / 40.0, 3.0 / 40.0, 3.0 / 40.0, 3.0 / 40.0],
-        dtype=_FDTYPE,
+        dtype=dtype,
     )
     return qp, qw
 

fluxfem/core/context_types.py

@@ -1,36 +1,60 @@
 from __future__ import annotations
 
-from typing import Any, Protocol, TypeAlias, runtime_checkable
+from typing import Any, Mapping, Protocol, TYPE_CHECKING, TypeAlias, runtime_checkable
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
 
 
 @runtime_checkable
 class VolumeContext(Protocol):
     """Minimum interface for volume weak-form evaluation."""
 
-    test: Any
-    trial: Any
-    w: Any
+    test: "FormFieldLike"
+    trial: "FormFieldLike"
+    w: ArrayLike
 
 
 @runtime_checkable
 class SurfaceContext(Protocol):
     """Minimum interface for surface weak-form evaluation."""
 
-    v: Any
-    w: Any
-    detJ: Any
-    normal: Any
+    v: "FormFieldLike"
+    w: ArrayLike
+    detJ: ArrayLike
+    normal: ArrayLike
 
 
 @runtime_checkable
 class FormFieldLike(Protocol):
     """Minimum interface for form fields used in weak-form evaluation."""
 
-    N: Any
-    gradN: Any
-    detJ: Any
+    N: ArrayLike
+    gradN: ArrayLike
+    detJ: ArrayLike
     value_dim: int
+    basis: Any
+
+
+@runtime_checkable
+class WeakFormContext(Protocol):
+    """Context interface used when resolving field references."""
+
+    test: FormFieldLike
+    trial: FormFieldLike
+    v: FormFieldLike
+    unknown: FormFieldLike | None
+    fields: Mapping[str, Any] | None
+    test_fields: Mapping[str, FormFieldLike] | None
+    trial_fields: Mapping[str, FormFieldLike] | None
+    unknown_fields: Mapping[str, FormFieldLike] | None
 
 
-UElement: TypeAlias = Any
+UElement: TypeAlias = ArrayLike | Mapping[str, ArrayLike]
 ParamsLike: TypeAlias = Any

fluxfem/core/dtypes.py

@@ -1,4 +1,12 @@
 import jax
 import jax.numpy as jnp
+import numpy as np
 
-DEFAULT_DTYPE = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
+
+def default_dtype() -> jnp.dtype:
+    return jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
+
+
+DEFAULT_DTYPE = default_dtype()
+INDEX_DTYPE = jnp.int64
+NP_INDEX_DTYPE = np.int64

fluxfem/core/forms.py

@@ -189,6 +189,7 @@ class FormContext:
         )
 
 
+@jax.tree_util.register_pytree_node_class
 @dataclass(eq=False)
 class FieldPair:
     """Named test/trial/unknown grouping for mixed formulations."""
@@ -196,6 +197,15 @@ class FieldPair:
     trial: FormFieldLike
     unknown: FormFieldLike | None = None
 
+    def tree_flatten(self):
+        children = (self.test, self.trial, self.unknown)
+        return children, {}
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        test, trial, unknown = children
+        return cls(test=test, trial=trial, unknown=unknown)
+
 
 @jax.tree_util.register_pytree_node_class
 @dataclass(eq=False)

fluxfem/core/mixed_assembly.py

@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+from typing import Mapping
+
+import jax
+import jax.numpy as jnp
+
+from .dtypes import INDEX_DTYPE
+from .assembly import element_residual, make_sparsity_pattern, chunk_pad_stats, _maybe_trace_pad
+
+
+def _coerce_mixed_u(space, u):
+    if isinstance(u, Mapping):
+        return space.pack_fields(u)
+    return jnp.asarray(u)
+
+
+def _split_elem_vec(field_names, elem_slices, u_elem_vec):
+    return {name: u_elem_vec[elem_slices[name]] for name in field_names}
+
+
+def _concat_residuals(field_names, res_dict):
+    return jnp.concatenate([res_dict[name] for name in field_names], axis=0)
+
+
+def make_element_mixed_residual_kernel(res_form, params, field_names, elem_slices):
+    """Jitted element residual kernel for mixed systems."""
+
+    def per_element(ctx, u_elem_vec):
+        u_elem = _split_elem_vec(field_names, elem_slices, u_elem_vec)
+        res_dict = element_residual(res_form, ctx, u_elem, params)
+        return _concat_residuals(field_names, res_dict)
+
+    return jax.jit(per_element)
+
+
+def make_element_mixed_jacobian_kernel(res_form, params, field_names, elem_slices):
+    """Jitted element Jacobian kernel for mixed systems."""
+    res_kernel = make_element_mixed_residual_kernel(res_form, params, field_names, elem_slices)
+
+    def fe_fun(u_elem_vec, ctx):
+        return res_kernel(ctx, u_elem_vec)
+
+    return jax.jit(jax.jacrev(fe_fun, argnums=0))
+
+
+def assemble_mixed_residual_scatter(
+    space,
+    res_form,
+    u,
+    params,
+    *,
+    sparse: bool = False,
+    kernel=None,
+    n_chunks: int | None = None,
+    pad_trace: bool = False,
+):
+    """Assemble mixed residual using jitted element kernels + scatter_add."""
+    u_vec = _coerce_mixed_u(space, u)
+    ctxs = space.build_form_contexts()
+    ker = kernel if kernel is not None else make_element_mixed_residual_kernel(
+        res_form, params, space.field_names, space.elem_slices
+    )
+
+    u_elems = u_vec[space.elem_dofs]
+    if n_chunks is None:
+        elem_res = jax.vmap(ker)(ctxs, u_elems)
+    else:
+        n_elems = int(u_elems.shape[0])
+        if n_chunks <= 0:
+            raise ValueError("n_chunks must be a positive integer.")
+        n_chunks = min(int(n_chunks), int(n_elems))
+        chunk_size = (n_elems + n_chunks - 1) // n_chunks
+        stats = chunk_pad_stats(n_elems, n_chunks)
+        _maybe_trace_pad(stats, n_chunks=n_chunks, pad_trace=pad_trace)
+        pad = (-n_elems) % chunk_size
+        if pad:
+            ctxs_pad = jax.tree_util.tree_map(
+                lambda x: jnp.concatenate([x, jnp.repeat(x[-1:], pad, axis=0)], axis=0),
+                ctxs,
+            )
+            u_elems_pad = jnp.concatenate([u_elems, jnp.repeat(u_elems[-1:], pad, axis=0)], axis=0)
+        else:
+            ctxs_pad = ctxs
+            u_elems_pad = u_elems
+
+        n_pad = n_elems + pad
+        n_chunks = n_pad // chunk_size
+
+        def _slice_first_dim(x, start, size):
+            start_idx = (start,) + (0,) * (x.ndim - 1)
+            slice_sizes = (size,) + x.shape[1:]
+            return jax.lax.dynamic_slice(x, start_idx, slice_sizes)
+
+        def chunk_fn(i):
+            start = i * chunk_size
+            ctx_chunk = jax.tree_util.tree_map(
+                lambda x: _slice_first_dim(x, start, chunk_size),
+                ctxs_pad,
+            )
+            u_chunk = _slice_first_dim(u_elems_pad, start, chunk_size)
+            res_chunk = jax.vmap(ker)(ctx_chunk, u_chunk)
+            return res_chunk.reshape(-1)
+
+        data_chunks = jax.vmap(chunk_fn)(jnp.arange(n_chunks))
+        elem_res = data_chunks.reshape(-1)[: n_elems * int(space.n_ldofs)].reshape(n_elems, -1)
+    rows = space.elem_dofs.reshape(-1)
+    data = elem_res.reshape(-1)
+
+    if sparse:
+        return rows, data, space.n_dofs
+
+    sdn = jax.lax.ScatterDimensionNumbers(
+        update_window_dims=(),
+        inserted_window_dims=(0,),
+        scatter_dims_to_operand_dims=(0,),
+    )
+    F = jnp.zeros((space.n_dofs,), dtype=data.dtype)
+    F = jax.lax.scatter_add(F, rows[:, None], data, sdn)
+    return F
+
+
+def assemble_mixed_jacobian_values(
+    space, res_form, u, params, *, kernel=None, n_chunks: int | None = None, pad_trace: bool = False
+):
+    """Assemble numeric values for mixed Jacobian (pattern-free)."""
+    u_vec = _coerce_mixed_u(space, u)
+    ctxs = space.build_form_contexts()
+    ker = kernel if kernel is not None else make_element_mixed_jacobian_kernel(
+        res_form, params, space.field_names, space.elem_slices
+    )
+
+    u_elems = u_vec[space.elem_dofs]
+    if n_chunks is None:
+        J_e_all = jax.vmap(ker)(u_elems, ctxs)
+        return J_e_all.reshape(-1)
+
+    n_elems = int(u_elems.shape[0])
+    if n_chunks <= 0:
+        raise ValueError("n_chunks must be a positive integer.")
+    n_chunks = min(int(n_chunks), int(n_elems))
+    chunk_size = (n_elems + n_chunks - 1) // n_chunks
+    stats = chunk_pad_stats(n_elems, n_chunks)
+    _maybe_trace_pad(stats, n_chunks=n_chunks, pad_trace=pad_trace)
+    pad = (-n_elems) % chunk_size
+    if pad:
+        ctxs_pad = jax.tree_util.tree_map(
+            lambda x: jnp.concatenate([x, jnp.repeat(x[-1:], pad, axis=0)], axis=0),
+            ctxs,
+        )
+        u_elems_pad = jnp.concatenate([u_elems, jnp.repeat(u_elems[-1:], pad, axis=0)], axis=0)
+    else:
+        ctxs_pad = ctxs
+        u_elems_pad = u_elems
+
+    n_pad = n_elems + pad
+    n_chunks = n_pad // chunk_size
+    m = int(space.n_ldofs)
+
+    def _slice_first_dim(x, start, size):
+        start_idx = (start,) + (0,) * (x.ndim - 1)
+        slice_sizes = (size,) + x.shape[1:]
+        return jax.lax.dynamic_slice(x, start_idx, slice_sizes)
+
+    def chunk_fn(i):
+        start = i * chunk_size
+        ctx_chunk = jax.tree_util.tree_map(
+            lambda x: _slice_first_dim(x, start, chunk_size),
+            ctxs_pad,
+        )
+        u_chunk = _slice_first_dim(u_elems_pad, start, chunk_size)
+        J_e = jax.vmap(ker)(u_chunk, ctx_chunk)
+        return J_e.reshape(-1)
+
+    data_chunks = jax.vmap(chunk_fn)(jnp.arange(n_chunks))
+    return data_chunks.reshape(-1)[: n_elems * m * m]
+
+
+def assemble_mixed_jacobian_scatter(
+    space,
+    res_form,
+    u,
+    params,
+    *,
+    kernel=None,
+    sparse: bool = True,
+    return_flux_matrix: bool = False,
+    pattern=None,
+    n_chunks: int | None = None,
+    pad_trace: bool = False,
+):
+    """Assemble mixed Jacobian using jitted element kernels + scatter_add."""
+    from ..solver import FluxSparseMatrix  # local import to avoid circular
+
+    pat = pattern if pattern is not None else make_sparsity_pattern(space, with_idx=not sparse)
+    data = assemble_mixed_jacobian_values(
+        space, res_form, u, params, kernel=kernel, n_chunks=n_chunks, pad_trace=pad_trace
+    )
+
+    if sparse:
+        if return_flux_matrix:
+            return FluxSparseMatrix(pat, data)
+        return pat.rows, pat.cols, data, pat.n_dofs
+
+    idx = pat.idx
+    if idx is None:
+        idx = (pat.rows.astype(jnp.int64) * int(pat.n_dofs) + pat.cols.astype(jnp.int64)).astype(INDEX_DTYPE)
+
+    n_entries = pat.n_dofs * pat.n_dofs
+    sdn = jax.lax.ScatterDimensionNumbers(
+        update_window_dims=(),
+        inserted_window_dims=(0,),
+        scatter_dims_to_operand_dims=(0,),
+    )
+    K_flat = jnp.zeros(n_entries, dtype=data.dtype)
+    K_flat = jax.lax.scatter_add(K_flat, idx[:, None], data, sdn)
+    return K_flat.reshape(pat.n_dofs, pat.n_dofs)
+
+
+def assemble_mixed_residual(
+    space, res_form, u, params, *, sparse: bool = False, n_chunks: int | None = None, pad_trace: bool = False
+):
+    """Assemble the global mixed residual vector."""
+    return assemble_mixed_residual_scatter(
+        space, res_form, u, params, sparse=sparse, n_chunks=n_chunks, pad_trace=pad_trace
+    )
+
+
+def assemble_mixed_jacobian(
+    space,
+    res_form,
+    u,
+    params,
+    *,
+    sparse: bool = True,
+    return_flux_matrix: bool = False,
+    pattern=None,
+    n_chunks: int | None = None,
+    pad_trace: bool = False,
+):
+    """Assemble the global mixed Jacobian."""
+    return assemble_mixed_jacobian_scatter(
+        space,
+        res_form,
+        u,
+        params,
+        sparse=sparse,
+        return_flux_matrix=return_flux_matrix,
+        pattern=pattern,
+        n_chunks=n_chunks,
+        pad_trace=pad_trace,
+    )
+
+
+__all__ = [
+    "make_element_mixed_residual_kernel",
+    "make_element_mixed_jacobian_kernel",
+    "assemble_mixed_residual",
+    "assemble_mixed_jacobian",
+    "assemble_mixed_residual_scatter",
+    "assemble_mixed_jacobian_scatter",
+    "assemble_mixed_jacobian_values",
+]