fluxfem 0.1.3a0__py3-none-any.whl

fluxfem/core/data.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+
+
+@dataclass(eq=False)
+class MeshData:
+    """Lightweight mesh data container for JAX-friendly serialization."""
+    coords: jnp.ndarray
+    conn: jnp.ndarray
+    cell_tags: jnp.ndarray | None = None
+    node_tags: jnp.ndarray | None = None
+
+    @classmethod
+    def from_mesh(cls, mesh: Any) -> "MeshData":
+        return cls(
+            coords=jnp.asarray(mesh.coords),
+            conn=jnp.asarray(mesh.conn),
+            cell_tags=None if mesh.cell_tags is None else jnp.asarray(mesh.cell_tags),
+            node_tags=None if mesh.node_tags is None else jnp.asarray(mesh.node_tags),
+        )
+
+
+@dataclass(eq=False)
+class BasisData:
+    """Quadrature and basis metadata for reproducible assembly."""
+    quad_points: jnp.ndarray
+    quad_weights: jnp.ndarray
+    dofs_per_node: int
+    kind: str
+
+    @classmethod
+    def from_basis(cls, basis: Any) -> "BasisData":
+        return cls(
+            quad_points=jnp.asarray(basis.quad_points),
+            quad_weights=jnp.asarray(basis.quad_weights),
+            dofs_per_node=int(basis.dofs_per_node),
+            kind=type(basis).__name__,
+        )
+
+
+@dataclass(eq=False)
+class SpaceData:
+    """Snapshot of space-related data used in assembly."""
+    mesh: MeshData
+    basis: BasisData
+    elem_dofs: jnp.ndarray
+    value_dim: int
+    n_dofs: int
+    n_ldofs: int
+
+    @classmethod
+    def from_space(cls, space: Any) -> "SpaceData":
+        return cls(
+            mesh=MeshData.from_mesh(space.mesh),
+            basis=BasisData.from_basis(space.basis),
+            elem_dofs=jnp.asarray(space.elem_dofs),
+            value_dim=int(space.value_dim),
+            n_dofs=int(space.n_dofs),
+            n_ldofs=int(space.n_ldofs),
+        )

fluxfem/core/dtypes.py

@@ -0,0 +1,4 @@
+import jax
+import jax.numpy as jnp
+
+DEFAULT_DTYPE = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32

fluxfem/core/forms.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+from dataclasses import dataclass
+
+from .basis import Basis3D
+
+# FormContext/ScalarFormField/VectorFormField were dataclasses with the default __eq__. JAX ended up calling that
+#   __eq__ during the vmap over residuals, which tries to compare array fields element‑wise and then coerce to a
+#   bool, triggering “truth value of an array is ambiguous.” Setting eq=False (and for ElementVector for consistency)
+#   removes the autogenerated __eq__, so vmap no longer evaluates array equality and the residual/Jacobian assembly
+#   succeeds.
+
+
+@dataclass(eq=False)
+class ElementVector:
+    """
+    Simple vector-valued element wrapper (scikit-fem ElementVector style).
+    dim: dofs per node (e.g., 3 for displacement)
+    """
+    dim: int
+
+    def dof_map(self, conn: jnp.ndarray) -> jnp.ndarray:
+        """
+        Expand scalar connectivity (n_elems, n_nodes_per_elem) to vector dofs.
+        Returns shape (n_elems, n_nodes_per_elem * dim).
+        """
+        base = conn[..., None] * self.dim                     # (n_elems, n_nodes, 1)
+        offsets = jnp.arange(self.dim, dtype=conn.dtype)      # (dim,)
+        dofs = base + offsets                                 # (n_elems, n_nodes, dim)
+        return dofs.reshape(conn.shape[0], -1)
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclass(eq=False)
+class ScalarFormField:
+    """Scalar FE field evaluated on one element."""
+    N: jnp.ndarray              # (n_q, n_nodes)
+    elem_coords: jnp.ndarray    # (n_nodes, 3)
+    basis: Basis3D
+    _gradN: jnp.ndarray | None = None
+    _detJ: jnp.ndarray | None = None
+
+    @property
+    def gradN(self):
+        if self._gradN is None:
+            self._gradN, self._detJ = self.basis.spatial_grads_and_detJ(self.elem_coords)
+        return self._gradN
+
+    @property
+    def detJ(self):
+        if self._detJ is None:
+            self._gradN, self._detJ = self.basis.spatial_grads_and_detJ(self.elem_coords)
+        return self._detJ
+
+    def eval(self, u_elem: jnp.ndarray) -> jnp.ndarray:
+        # u_elem: (n_nodes,)
+        return jnp.einsum("qa,a->q", self.N, u_elem)
+
+    def grad(self, u_elem: jnp.ndarray) -> jnp.ndarray:
+        # returns (n_q, 3)
+        return jnp.einsum("qaj,a->qj", self.gradN, u_elem)
+
+    def tree_flatten(self):
+        children = (self.N, self.elem_coords, self._gradN, self._detJ)
+        aux = {"basis": self.basis}
+        return children, aux
+
+    @classmethod
+    def tree_unflatten(cls, aux, chirdren):
+        N, elem_coords, gradN, detJ = chirdren
+        return cls(N, elem_coords, aux["basis"], gradN, detJ)
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclass(eq=False)
+class VectorFormField:
+    """Vector-valued FE field evaluated on one element."""
+    N: jnp.ndarray
+    elem_coords: jnp.ndarray
+    basis: Basis3D
+    value_dim: int  # ★Python int (static)
+    _gradN: jnp.ndarray | None = None
+    _detJ: jnp.ndarray | None = None
+
+    @property
+    def gradN(self):
+        if self._gradN is None:
+            self._gradN, self._detJ = self.basis.spatial_grads_and_detJ(self.elem_coords)
+        return self._gradN
+
+    @property
+    def detJ(self):
+        if self._detJ is None:
+            self._gradN, self._detJ = self.basis.spatial_grads_and_detJ(self.elem_coords)
+        return self._detJ
+
+    def eval(self, u_elem: jnp.ndarray) -> jnp.ndarray:
+        # u_elem: (value_dim*n_nodes,) expected
+        u_nodes = u_elem.reshape((-1, self.value_dim))  # (n_nodes, vd); vd is a Python int
+        return jnp.einsum("qa,ai->qi", self.N, u_nodes)  # (n_q, vd)
+
+    def grad(self, u_elem: jnp.ndarray) -> jnp.ndarray:
+        u_nodes = u_elem.reshape((-1, self.value_dim))
+        return jnp.einsum("qaj,ai->qij", self.gradN, u_nodes)  # (n_q, vd, 3)
+
+    def tree_flatten(self):
+        children = (self.N, self.elem_coords, self._gradN, self._detJ)
+        aux = {
+            "basis": self.basis,
+            "value_dim": int(self.value_dim)
+        }
+        return children, aux
+
+    @classmethod
+    def tree_unflatten(cls, aux, children):
+        N, elem_coords, gradN, detJ = children
+        return cls(N, elem_coords, aux["basis"], aux["value_dim"], gradN, detJ)
+
+
+FormFieldLike = ScalarFormField | VectorFormField
+
+
+def vector_load_form(field: FormFieldLike, load_vec: jnp.ndarray) -> jnp.ndarray:
+    """
+    Build vector linear form values from a FormField and a load vector.
+    """
+    lv = jnp.asarray(load_vec)
+    if lv.ndim == 1:
+        lv = lv[None, :]
+    elif lv.ndim != 2:
+        raise ValueError("load_vec must be shape (dim,) or (n_q, dim)")
+    if lv.shape[0] == 1:
+        lv = jnp.broadcast_to(lv, (field.N.shape[0], lv.shape[1]))
+    elif lv.shape[0] != field.N.shape[0]:
+        raise ValueError("load_vec must be shape (dim,) or (n_q, dim)")
+    load = field.N[..., None] * lv[:, None, :]
+    return load.reshape(load.shape[0], -1)
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclass(eq=False)
+class FormContext:
+    """Bundle test/trial fields and quadrature data for element assembly."""
+    test: FormFieldLike
+    trial: FormFieldLike
+    x_q: jnp.ndarray       # (n_q, 3)
+    w: jnp.ndarray         # (n_q,)
+    elem_id: jnp.ndarray | int = 0
+
+    @property
+    def u(self) -> FormFieldLike:
+        return self.trial
+
+    @property
+    def v(self) -> FormFieldLike:
+        return self.test
+
+    def tree_flatten(self):
+        children = (
+            self.test,
+            self.trial,
+            self.x_q,
+            self.w,
+            self.elem_id,
+        )
+        return children, {}
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        (
+            test,
+            trial,
+            x_q,
+            w,
+            elem_id,
+        ) = children
+        return cls(
+            test,
+            trial,
+            x_q,
+            w,
+            elem_id,
+        )
+
+
+@dataclass(eq=False)
+class FieldPair:
+    """Named test/trial/unknown grouping for mixed formulations."""
+    test: FormFieldLike
+    trial: FormFieldLike
+    unknown: FormFieldLike | None = None
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclass(eq=False)
+class MixedFormContext:
+    """FormContext for mixed formulations keyed by field name."""
+    fields: dict[str, FieldPair]
+    x_q: jnp.ndarray       # (n_q, 3)
+    w: jnp.ndarray         # (n_q,)
+    elem_id: jnp.ndarray | int = 0
+    unknown: FormFieldLike | None = None
+    trial_fields: dict[str, FormFieldLike] | None = None
+    test_fields: dict[str, FormFieldLike] | None = None
+    unknown_fields: dict[str, FormFieldLike] | None = None
+
+    def tree_flatten(self):
+        children = (
+            self.fields,
+            self.x_q,
+            self.w,
+            self.elem_id,
+            self.unknown,
+            self.trial_fields,
+            self.test_fields,
+            self.unknown_fields,
+        )
+        return children, {}
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        (
+            fields,
+            x_q,
+            w,
+            elem_id,
+            unknown,
+            trial_fields,
+            test_fields,
+            unknown_fields,
+        ) = children
+        return cls(fields, x_q, w, elem_id, unknown, trial_fields, test_fields, unknown_fields)

fluxfem/core/interp.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import numpy as np
+
+from .basis import HexTriLinearBasis
+from .space import FESpace
+
+
+def eval_shape_functions_hex8(xi_eta_zeta: np.ndarray) -> np.ndarray:
+    """
+    Evaluate trilinear Hex8 shape functions at given local coords (xi, eta, zeta) in [-1,1]^3.
+    Returns N with shape (n_q, 8).
+    """
+    pts = np.atleast_2d(np.asarray(xi_eta_zeta, dtype=float))
+    xi, eta, zeta = pts[:, 0], pts[:, 1], pts[:, 2]
+    N = np.stack(
+        [
+            0.125 * (1 - xi) * (1 - eta) * (1 - zeta),
+            0.125 * (1 + xi) * (1 - eta) * (1 - zeta),
+            0.125 * (1 + xi) * (1 + eta) * (1 - zeta),
+            0.125 * (1 - xi) * (1 + eta) * (1 - zeta),
+            0.125 * (1 - xi) * (1 - eta) * (1 + zeta),
+            0.125 * (1 + xi) * (1 - eta) * (1 + zeta),
+            0.125 * (1 + xi) * (1 + eta) * (1 + zeta),
+            0.125 * (1 - xi) * (1 + eta) * (1 + zeta),
+        ],
+        axis=1,
+    )
+    return N
+
+
+def interpolate_field_at_element_points(space: FESpace, u: np.ndarray, xi_eta_zeta: np.ndarray) -> np.ndarray:
+    """
+    Interpolate vector field u (3 dof/node ordering) at given local points for all elements.
+    - xi_eta_zeta: (m,3) local coords in [-1,1]^3
+    Returns: (n_elem, m, 3)
+    """
+    if not isinstance(space.basis, HexTriLinearBasis):
+        raise NotImplementedError("interpolate_field_at_element_points currently supports Hex8 (trilinear) only.")
+    N = eval_shape_functions_hex8(xi_eta_zeta)  # (m,8)
+    u_arr = np.asarray(u)
+    n_nodes = space.mesh.coords.shape[0]
+    if u_arr.shape[0] != 3 * n_nodes:
+        raise ValueError(f"Expected 3 dof/node; got {u_arr.shape[0]} for {n_nodes} nodes")
+    u_nodes = u_arr.reshape(n_nodes, 3)
+    conn = np.asarray(space.elem_dofs) // 3  # node indices
+    elem_u = u_nodes[conn]  # (n_elem,8,3)
+    vals = np.einsum("pq,eqr->epr", N, elem_u)  # (n_elem, m, 3)
+    return vals
+
+
+__all__ = [
+    "eval_shape_functions_hex8",
+    "interpolate_field_at_element_points",
+]

fluxfem/core/solver.py

@@ -0,0 +1,113 @@
+"""
+Helper to bridge JAX-assembled matrices back to NumPy/SciPy and solve.
+"""
+from __future__ import annotations
+
+import numpy as np
+import jax.numpy as jnp
+from typing import Any
+
+try:
+    import scipy.sparse as sp
+    from scipy.sparse.linalg import spsolve
+except Exception as exc:  # pragma: no cover
+    raise ImportError("scipy is required for spsolve utilities") from exc
+
+
+def coo_to_csr(rows: Any, cols: Any, data: Any, n_dofs: int):
+    """
+    Convert COO triplets to SciPy CSR matrix.
+    """
+    r = np.asarray(rows, dtype=np.int64)
+    c = np.asarray(cols, dtype=np.int64)
+    d = np.asarray(data)
+    return sp.csr_matrix((d, (r, c)), shape=(n_dofs, n_dofs))
+
+
+def spdirect_solve_cpu(K: Any, F: jnp.ndarray, *, use_jax: bool = False) -> np.ndarray:
+    """
+    Convert JAX arrays to NumPy/SciPy and solve K u = F with sparse solver.
+    If use_jax=True, dispatch to JAX's experimental sparse spsolve.
+
+    Parameters
+    ----------
+    K : jnp.ndarray
+        Global stiffness matrix (n_dofs, n_dofs), dense or symmetric.
+    F : jnp.ndarray
+        Load vector (n_dofs,) or multiple RHS (n_dofs, n_rhs)
+
+    Returns
+    -------
+    np.ndarray
+        Solution vector u (n_dofs,) or (n_dofs, n_rhs)
+    """
+    if use_jax:
+        try:
+            return spdirect_solve_jax(K, F)
+        except Exception:
+            pass
+
+    if hasattr(K, "to_csr"):
+        K_csr = K.to_csr()
+    elif isinstance(K, tuple) and len(K) == 4:
+        K_csr = coo_to_csr(*K)
+    elif sp.issparse(K):
+        K_csr = K.tocsr()
+    else:
+        K_np = np.asarray(K)
+        K_csr = sp.csr_matrix(K_np)
+
+    F_np = np.asarray(F)
+    u = spsolve(K_csr, F_np)
+    return np.asarray(u)
+
+
+def spdirect_solve_jax(K: Any, F: jnp.ndarray) -> np.ndarray:
+    """
+    Direct sparse solve in JAX via jax.experimental.sparse.linalg.spsolve.
+    Accepts FluxSparseMatrix or jax.experimental.sparse.BCOO.
+    """
+    try:
+        import jax
+        if jax.default_backend() == "cpu":
+            # JAX spsolve falls back to SciPy on CPU and can hit read-only buffers.
+            return spdirect_solve_cpu(K, F, use_jax=False)
+    except Exception:
+        pass
+    try:
+        from jax.experimental.sparse.linalg import spsolve as jspsolve
+        from jax.experimental import sparse as jsparse
+    except Exception as exc:  # pragma: no cover
+        raise ImportError("jax.experimental.sparse is required for spdirect_solve_jax") from exc
+
+    if sp.issparse(K):
+        data = jnp.asarray(K.data)
+        indices = jnp.asarray(K.indices)
+        indptr = jnp.asarray(K.indptr)
+        F_arr = jnp.asarray(F)
+        if F_arr.ndim == 1:
+            return np.asarray(jspsolve(data, indices, indptr, F_arr))
+        return np.asarray(jnp.stack([jspsolve(data, indices, indptr, F_arr[:, i]) for i in range(F_arr.shape[1])], axis=1))
+
+    if isinstance(K, tuple) and len(K) == 4:
+        rows, cols, data, n_dofs = K
+        idx = jnp.stack([jnp.asarray(rows), jnp.asarray(cols)], axis=-1)
+        bcoo = jsparse.BCOO((jnp.asarray(data), idx), shape=(int(n_dofs), int(n_dofs)))
+    elif isinstance(K, jsparse.BCOO):
+        bcoo = K
+    elif hasattr(K, "to_bcoo"):
+        bcoo = K.to_bcoo()
+    else:
+        raise TypeError("spdirect_solve_jax expects FluxSparseMatrix, BCOO, CSR, or COO tuple")
+
+    bcsr = jsparse.BCSR.from_bcoo(bcoo)
+    F_arr = jnp.asarray(F)
+    if F_arr.ndim == 1:
+        return np.asarray(jspsolve(bcsr.data, bcsr.indices, bcsr.indptr, F_arr))
+    return np.asarray(jnp.stack([jspsolve(bcsr.data, bcsr.indices, bcsr.indptr, F_arr[:, i]) for i in range(F_arr.shape[1])], axis=1))
+
+def spdirect_solve_gpu(K: Any, F: jnp.ndarray) -> np.ndarray:
+    """
+    GPU direct sparse solve via JAX experimental sparse solver.
+    """
+    return spdirect_solve_jax(K, F)