fluxfem 0.2.1__tar.gz → 0.2.2__tar.gz

{fluxfem-0.2.1 → fluxfem-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fluxfem
-Version: 0.2.1
+Version: 0.2.2
 Summary: FluxFEM: A weak-form-centric differentiable finite element framework in JAX
 License: Apache-2.0
 Author: Kohei Watanabe
@@ -245,9 +245,8 @@ K_cc = ...
 K_uc = ...
 
 blocks = ff_solver.make_block_matrix(
-    diag=ff_solver.block_diag(u=K_uu, c=K_cc),
+    diag=ff_solver.block_diag(order=("u", "c"), u=K_uu, c=K_cc),
     rel={("u", "c"): K_uc},
-    sizes={"u": K_uu.shape[0], "c": K_cc.shape[0]},
     symmetric=True,
     transpose_rule="T",
 )
@@ -297,8 +296,19 @@ poetry add fluxfem[cuda12]
 
 Optional PETSc-based solvers are available via `petsc4py`. Enable with the extra:
 
+
+```bash
+pip install "fluxfem[petsc]"
+or
+pip install "fluxfem[petsc,cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+```
+
 ```bash
 poetry add fluxfem --extras "petsc"
+or
+poetry add "fluxfem[petsc,cuda12]"
+or
+poetry add fluxfem --extras "petsc" --extras "cuda12"
 ```
 
 Note: newer `petsc4py` expects PETSc builds that include the `PetscRegressor`
@@ -306,6 +316,12 @@ API. If your PETSc build does not have it, `petsc4py` will fail to compile. In
 that case, rebuild PETSc with regressor support or pin `petsc4py` to a version
 compatible with your PETSc build.
 
+Important: you must match the `petsc4py` version to the PETSc version you have
+installed. The current FluxFEM extra pins `petsc4py==3.23.6` (see
+`[project.optional-dependencies]`), so make sure your PETSc install is
+compatible with that `petsc4py` release, or override it to match your PETSc
+build.
+
 GPU note: this repo currently tests CUDA via the `cuda12` extra only. Other CUDA
 versions are not covered by CI and may require manual JAX installation.
 

{fluxfem-0.2.1 → fluxfem-0.2.2}/README.md

@@ -219,9 +219,8 @@ K_cc = ...
 K_uc = ...
 
 blocks = ff_solver.make_block_matrix(
-    diag=ff_solver.block_diag(u=K_uu, c=K_cc),
+    diag=ff_solver.block_diag(order=("u", "c"), u=K_uu, c=K_cc),
     rel={("u", "c"): K_uc},
-    sizes={"u": K_uu.shape[0], "c": K_cc.shape[0]},
     symmetric=True,
     transpose_rule="T",
 )
@@ -271,8 +270,19 @@ poetry add fluxfem[cuda12]
 
 Optional PETSc-based solvers are available via `petsc4py`. Enable with the extra:
 
+
+```bash
+pip install "fluxfem[petsc]"
+or
+pip install "fluxfem[petsc,cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+```
+
 ```bash
 poetry add fluxfem --extras "petsc"
+or
+poetry add "fluxfem[petsc,cuda12]"
+or
+poetry add fluxfem --extras "petsc" --extras "cuda12"
 ```
 
 Note: newer `petsc4py` expects PETSc builds that include the `PetscRegressor`
@@ -280,6 +290,12 @@ API. If your PETSc build does not have it, `petsc4py` will fail to compile. In
 that case, rebuild PETSc with regressor support or pin `petsc4py` to a version
 compatible with your PETSc build.
 
+Important: you must match the `petsc4py` version to the PETSc version you have
+installed. The current FluxFEM extra pins `petsc4py==3.23.6` (see
+`[project.optional-dependencies]`), so make sure your PETSc install is
+compatible with that `petsc4py` release, or override it to match your PETSc
+build.
+
 GPU note: this repo currently tests CUDA via the `cuda12` extra only. Other CUDA
 versions are not covered by CI and may require manual JAX installation.
 

{fluxfem-0.2.1 → fluxfem-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "fluxfem"
-version = "0.2.1"
+version = "0.2.2"
 description = ""
 authors = [
     {name = "Kohei Watanabe",email = "koheitech001@gmail.com"}
@@ -30,7 +30,7 @@ cuda12 = ["jax (==0.8.2)", "jaxlib (==0.8.2)", "jax-cuda12-plugin (==0.8.2)", "j
 
 [tool.poetry]
 name = "fluxfem"
-version = "0.2.1"
+version = "0.2.2"
 description = "FluxFEM: A weak-form-centric differentiable finite element framework in JAX"
 authors = ["Kohei Watanabe <koheitech001@gmail.com>"]
 readme = "README.md"

{fluxfem-0.2.1 → fluxfem-0.2.2}/src/fluxfem/solver/bc.py

@@ -4,6 +4,7 @@ from dataclasses import dataclass
 from typing import Optional, Sequence
 import numpy as np
 import numpy.typing as npt
+import jax
 import jax.numpy as jnp
 
 from ..mesh.surface import SurfaceMesh
@@ -212,6 +213,121 @@ def assemble_surface_linear_form(
     """
     Assemble a linear form over surface facets using a weak-form callback.
     """
+    def _is_jax(x) -> bool:
+        return isinstance(x, jax.Array) or isinstance(x, jax.core.Tracer)
+
+    def _assemble_surface_linear_form_jax():
+        facets = jnp.asarray(surface.conn, dtype=jnp.int32)
+        coords = jnp.asarray(surface.coords)
+        if facets.ndim != 2 or facets.shape[1] not in (3, 4):
+            raise NotImplementedError("JAX surface assembly supports only tri/quad facets.")
+
+        n_nodes = surface.n_nodes if n_total_nodes is None else int(n_total_nodes)
+        n_dofs = n_nodes * dim
+        dtype = coords.dtype if F0 is None else jnp.asarray(F0).dtype
+        F = jnp.zeros(n_dofs, dtype=dtype) if F0 is None else jnp.asarray(F0, dtype=dtype)
+        if F.shape[0] != n_dofs:
+            raise ValueError(f"F length {F.shape[0]} does not match expected {n_dofs}")
+
+        coords_f = coords[facets]
+        facet_ids = jnp.arange(facets.shape[0], dtype=jnp.int32)
+        outward_from = jnp.mean(coords, axis=0)
+
+        def facet_fe_tri(facet_id, nodes):
+            p0, p1, p2 = nodes[0], nodes[1], nodes[2]
+            cross = jnp.cross(p1 - p0, p2 - p0)
+            area = 0.5 * jnp.linalg.norm(cross)
+            centroid = (p0 + p1 + p2) / 3.0
+            norm = jnp.linalg.norm(cross)
+            norm = jnp.where(norm < 1e-12, 1e-12, norm)
+            normal = cross / norm
+            v = centroid - outward_from
+            normal = jnp.where(jnp.dot(normal, v) < 0.0, -normal, normal)
+
+            N = jnp.array([[1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0]], dtype=dtype)
+            w = jnp.array([1.0], dtype=dtype)
+            ctx = SurfaceFormContext(
+                v=SurfaceFormField(N=N, value_dim=dim),
+                x_q=centroid[None, :],
+                w=w,
+                detJ=jnp.array([area], dtype=dtype),
+                facet_id=facet_id,
+                normal=normal,
+            )
+            fe_q = form(ctx, params)
+            if getattr(form, "_includes_measure", False):
+                fe = jnp.einsum("qi->i", fe_q)
+            else:
+                fe = jnp.einsum("qi,q->i", fe_q, w * area)
+            return fe
+
+        def facet_fe_quad(facet_id, nodes):
+            p0, p1, p2, p3 = nodes[0], nodes[1], nodes[2], nodes[3]
+            cross = jnp.cross(p1 - p0, p3 - p0)
+            norm = jnp.linalg.norm(cross)
+            norm = jnp.where(norm < 1e-12, 1e-12, norm)
+            normal = cross / norm
+            centroid = (p0 + p1 + p2 + p3) / 4.0
+            v = centroid - outward_from
+            normal = jnp.where(jnp.dot(normal, v) < 0.0, -normal, normal)
+
+            gp = jnp.array([-1.0 / jnp.sqrt(3.0), 1.0 / jnp.sqrt(3.0)], dtype=dtype)
+            xi, eta = jnp.meshgrid(gp, gp, indexing="xy")
+            xi = xi.reshape(-1)
+            eta = eta.reshape(-1)
+            w = jnp.ones_like(xi)
+
+            N = 0.25 * jnp.stack(
+                [
+                    (1 - xi) * (1 - eta),
+                    (1 + xi) * (1 - eta),
+                    (1 + xi) * (1 + eta),
+                    (1 - xi) * (1 + eta),
+                ],
+                axis=1,
+            )
+            dN_dxi = 0.25 * jnp.stack(
+                [-(1 - eta), (1 - eta), (1 + eta), -(1 + eta)], axis=1
+            )
+            dN_deta = 0.25 * jnp.stack(
+                [-(1 - xi), -(1 + xi), (1 + xi), (1 - xi)], axis=1
+            )
+
+            dx_dxi = jnp.einsum("qa,ai->qi", dN_dxi, nodes)
+            dx_deta = jnp.einsum("qa,ai->qi", dN_deta, nodes)
+            detJ = jnp.linalg.norm(jnp.cross(dx_dxi, dx_deta), axis=1)
+            x_q = jnp.einsum("qa,ai->qi", N, nodes)
+
+            ctx = SurfaceFormContext(
+                v=SurfaceFormField(N=N, value_dim=dim),
+                x_q=x_q,
+                w=w,
+                detJ=detJ,
+                facet_id=facet_id,
+                normal=normal,
+            )
+            fe_q = form(ctx, params)
+            if getattr(form, "_includes_measure", False):
+                fe = jnp.einsum("qi->i", fe_q)
+            else:
+                fe = jnp.einsum("qi,q->i", fe_q, w * detJ)
+            return fe
+
+        if facets.shape[1] == 3:
+            fe = jax.vmap(facet_fe_tri)(facet_ids, coords_f)
+        else:
+            fe = jax.vmap(facet_fe_quad)(facet_ids, coords_f)
+
+        fe = fe.astype(dtype)
+        offsets = jnp.arange(dim, dtype=jnp.int32)
+        dofs = facets[:, :, None] * dim + offsets[None, None, :]
+        dofs = dofs.reshape(facets.shape[0], -1)
+        F = F.at[dofs].add(fe)
+        return F
+
+    if _is_jax(surface.coords) or _is_jax(surface.conn) or _is_jax(F0):
+        return _assemble_surface_linear_form_jax()
+
     facets = np.asarray(surface.conn, dtype=int)
     coords = np.asarray(surface.coords)
     n_nodes = surface.n_nodes if n_total_nodes is None else int(n_total_nodes)

{fluxfem-0.2.1 → fluxfem-0.2.2}/src/fluxfem/solver/block_matrix.py

@@ -18,8 +18,22 @@ FieldKey: TypeAlias = str | int
 BlockMap: TypeAlias = dict[FieldKey, dict[FieldKey, MatrixLike]]
 
 
-def diag(**blocks: MatrixLike) -> dict[str, MatrixLike]:
-    return dict(blocks)
+def diag(*, order: Sequence[FieldKey] | None = None, **blocks: MatrixLike) -> dict[FieldKey, MatrixLike]:
+    """
+    Build a dict of diagonal blocks with an optional explicit field order.
+    """
+    if order is None:
+        return dict(blocks)
+    ordered_blocks: dict[FieldKey, MatrixLike] = {}
+    for name in order:
+        if name not in blocks:
+            raise KeyError(f"Missing block '{name}' in block_diag order")
+        ordered_blocks[name] = blocks[name]
+    extra = set(blocks) - set(order)
+    if extra:
+        extra_list = ", ".join(str(name) for name in sorted(extra, key=str))
+        raise KeyError(f"Unknown block(s) not in order: {extra_list}")
+    return ordered_blocks
 
 
 def _infer_sizes_from_diag(diag_blocks: Mapping[FieldKey, MatrixLike]) -> dict[FieldKey, int]: