loadpath 0.1.1__py3-none-any.whl

loadpath/narrowband.py

@@ -0,0 +1,418 @@
+"""Narrow-band sampling of an implicit field for meshing.
+
+A part defined by a field f(P) < 0 only needs f near its surface: deep
+inside or far outside, marching cubes finds nothing. `narrowband_mc` cuts
+the meshing grid into blocks of `block` cells, samples f once per block
+(at its centre) and skips every block whose centre is further from the
+zero level than the field can change across the block (Lipschitz bound
+`L` times the half-diagonal). Only the remaining blocks are sampled at
+full resolution and marched, one block at a time.
+
+The surviving blocks share their boundary samples with their neighbours,
+and marching cubes puts a vertex on an edge from that edge's two samples
+alone, so the blocks' meshes join exactly: after merging the shared
+vertices the result is the same mesh dense marching cubes gives on the
+whole grid (tests/test_narrowband.py checks this), at a fraction of the
+field evaluations and with memory set by one batch of blocks, not by the
+grid. That is what lets a field go to a finer pitch than a dense grid
+fits in memory.
+
+The skip test is only as safe as `L`. For a distance field (shapes,
+booleans, shells, offsets) L = 1; smooth unions, remaps of a distance and
+plain TPMS level sets in fields.py stay close to it, but a warp-graded
+lattice reaches 30 or more. `lipschitz=None` estimates L as the steepest
+gradient at 20k random points or between neighbouring block centres,
+times `margin`, and never below 1. A bound that is still too low leaves
+a crossing in a skipped leaf, which shows as a hole where it meets a
+sampled one: fields.field_mesh checks the mesh is watertight and, if it
+is not, marches again with every block sampled.
+"""
+from __future__ import annotations
+
+import math
+import os
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+
+__all__ = ["narrowband_mc", "march_blocks", "peval", "watertight", "signed_volume", "shell_volumes", "face_cross"]
+
+_BATCH = 2_000_000      # field samples per call
+_CHUNK = 250_000        # field samples per thread task
+
+
+def _threads():
+    env = os.environ.get("LOADPATH_THREADS")
+    if env:
+        return max(1, int(env))
+    try:
+        return max(1, len(os.sched_getaffinity(0)))
+    except AttributeError:
+        return max(1, os.cpu_count() or 1)
+
+
+_POOL = None
+
+
+def peval(field, P, chunk=_CHUNK):
+    """field(P) as float64, in chunks spread over threads. The fields are
+    numpy expressions, which release the GIL, so 4 cores evaluate about
+    3.7x faster, and every chunk is the same computation, so the result is
+    bit for bit the serial one. The first chunk runs alone, so a field that
+    builds a cached resource on first use (a warp, a mesh's distance grid)
+    builds it once. LOADPATH_THREADS=1 turns this off. A field with a
+    `batch` attribute (the torch evaluator, which threads itself or runs
+    on a GPU) is called serially in chunks of that many points instead."""
+    global _POOL
+    n = len(P)
+    out = np.empty(n)
+    nt = _threads()
+    chunk = getattr(field, "chunk", None) or chunk
+    batch = getattr(field, "batch", None)
+    if batch:
+        nt, chunk = 1, batch
+    if nt == 1 or n <= 2 * chunk:
+        for i in range(0, n, chunk):
+            out[i:i + chunk] = np.asarray(field(P[i:i + chunk]), float)
+        return out
+    out[:chunk] = np.asarray(field(P[:chunk]), float)
+    if _POOL is None or _POOL._max_workers != nt:
+        _POOL = ThreadPoolExecutor(nt, thread_name_prefix="loadpath-field")
+
+    def one(i):
+        out[i:i + chunk] = np.asarray(field(P[i:i + chunk]), float)
+
+    for _ in _POOL.map(one, range(chunk, n, chunk)):
+        pass
+    return out
+
+
+def pblock(fn, P, rows, chunk=_CHUNK):
+    """fn(P) for a function giving (rows, len(P)) blocks (a crease tree's
+    leaf values), in chunks spread over threads as peval does: the first
+    chunk alone, then the rest on the pool."""
+    global _POOL
+    n = len(P)
+    out = np.empty((rows, n))
+    nt = _threads()
+    if nt == 1 or n <= 2 * chunk:
+        for i in range(0, n, chunk):
+            out[:, i:i + chunk] = fn(P[i:i + chunk])
+        return out
+    out[:, :chunk] = fn(P[:chunk])
+    if _POOL is None or _POOL._max_workers != nt:
+        _POOL = ThreadPoolExecutor(nt, thread_name_prefix="loadpath-field")
+
+    def one(i):
+        out[:, i:i + chunk] = fn(P[i:i + chunk])
+
+    for _ in _POOL.map(one, range(chunk, n, chunk)):
+        pass
+    return out
+
+
+def _grid_points(lo, pitch, i0, n):
+    """Points of the sub-grid starting at index i0 (3,) with n (3,) samples."""
+    ax = [lo[k] + (i0[k] + np.arange(n[k])) * pitch for k in range(3)]
+    X, Y, Z = np.meshgrid(*ax, indexing="ij")
+    return np.stack([X.ravel(), Y.ravel(), Z.ravel()], axis=1)
+
+
+def _eval(field, P):
+    # grid samples (kept as float32 for marching cubes): lattice trig in
+    # single precision
+    from .lattice import single_precision
+
+    with single_precision():
+        return peval(field, P)
+
+
+def _lipschitz(field, lo, hi, pitch, n=20_000):
+    """The steepest gradient seen at n random points (central differences)."""
+    P = np.random.default_rng(0).uniform(lo, hi, size=(n, 3))
+    h = 0.05 * pitch
+    E = np.eye(3) * h
+    f = _eval(field, np.concatenate([P + e for e in E] + [P - e for e in E])).reshape(6, n)
+    g = (f[:3] - f[3:]) / (2 * h)
+    return float(np.sqrt((g * g).sum(axis=0)).max())
+
+
+_LOCAL = True
+_LOCAL_MARGIN = 3.0     # local slopes come from differences a leaf apart, which smooth over a steep spot
+
+
+def _local_lipschitz(fl, nl, step, margin, floor):
+    """A Lipschitz bound per leaf (blocks, nl**3): the steepest difference
+    between its leaf centre and the next one along each axis (inside its
+    block), over the spacing, times `margin`, never below `floor`. The global
+    bound is the steepest slope anywhere, and a TPMS sheet's distance
+    (value over gradient) climbs steeply round the gradient's zeros, deep
+    in the channels, far from any wall: there it is 20-30, which keeps
+    every leaf of a lattice. Next to a wall the slope is ~1-3, and the
+    field varies over a cell (millimetres), not a leaf, so the local slope
+    is the one that matters for whether the leaf reaches the surface."""
+    F = fl.reshape(-1, nl, nl, nl)
+    G = np.zeros_like(F)
+    with np.errstate(invalid="ignore"):
+        for ax in range(1, 4):
+            d = np.abs(np.diff(F, axis=ax))
+            d = np.where(np.isnan(d), 0.0, d)      # two empty leaves (inf - inf)
+            lo = [slice(None)] * 4
+            hi = [slice(None)] * 4
+            lo[ax] = slice(0, -1)
+            hi[ax] = slice(1, None)
+            G[tuple(lo)] = np.maximum(G[tuple(lo)], d)
+            G[tuple(hi)] = np.maximum(G[tuple(hi)], d)
+    return np.maximum(floor, margin * G / step).reshape(fl.shape)
+
+
+def narrowband_mc(field, lo, dims, pitch, block=32, leaf=4, lipschitz=None, margin=1.5, budget=None,
+                  progress=None):
+    """Marching cubes of field < 0 on the grid lo + i * pitch, i < dims,
+    sampling only near the zero level. The grid's outer layer is forced
+    outside (+pitch), as fields.field_mesh does, so the mesh is closed.
+
+    Two levels: blocks of `block` cells are tested at their centres; in
+    the blocks that pass, leaves of `leaf` cells are tested the same way,
+    and only the samples of the leaves that pass are evaluated. A sample
+    of a block that no passing leaf touches gets the sign of its leaf at
+    a large magnitude: with the Lipschitz bound holding, every edge the
+    surface crosses has both its samples evaluated, so the block's mesh is
+    the dense one. Each block is marched on its own and the shared
+    vertices merged. Raises ValueError, before the fine pass, when the
+    passing leaves need more than `budget` samples. Returns (verts (mm),
+    faces, stats)."""
+    stats = {}
+    verts, faces, nv = [], [], 0
+    for i0, n, v, fa in march_blocks(field, lo, dims, pitch, block, leaf, lipschitz, margin, budget, progress,
+                                     stats):
+        verts.append(v)
+        faces.append(fa + nv)
+        nv += len(v)
+    if not verts:
+        return np.zeros((0, 3)), np.zeros((0, 3), int), stats
+    V, Fc = merge_blocks(np.concatenate(verts), np.concatenate(faces), block)
+    return np.asarray(lo, float) + V * pitch, Fc, stats
+
+
+def merge_blocks(V, F, block):
+    """Merge the vertices the blocks share. A vertex on a block's face comes
+    out of both blocks bit for bit (same edge, same two samples) and every
+    other vertex is unique already, so only the ones on a block plane (a
+    coordinate that is a multiple of `block`) go through the (slow) row
+    unique. Drops the faces the merge collapses."""
+    on = (np.mod(V, block) == 0).any(axis=1)
+    idx = np.flatnonzero(on)
+    remap = np.arange(len(V))
+    if len(idx):
+        _, first, inv = np.unique(V[idx], axis=0, return_index=True, return_inverse=True)
+        remap[idx] = idx[first][inv.reshape(-1)]
+    keep_v = remap == np.arange(len(V))
+    new = np.cumsum(keep_v) - 1
+    F = new[remap[F]]
+    V = V[keep_v]
+    ok = (F[:, 0] != F[:, 1]) & (F[:, 1] != F[:, 2]) & (F[:, 0] != F[:, 2])
+    return V, F[ok]
+
+
+def march_blocks(field, lo, dims, pitch, block=32, leaf=4, lipschitz=None, margin=1.5, budget=None,
+                 progress=None, stats=None):
+    """The narrow-band blocks of narrowband_mc, marched one at a time, in
+    x-major order (all the blocks of one x slab before the next): yields
+    (i0, n, verts, faces) per block that the surface crosses, verts in grid
+    index units (add lo, times pitch, for mm). Fills `stats` once the
+    blocks are planned. The field is evaluated in batches of blocks; memory
+    is one batch, whatever the grid."""
+    from skimage.measure import marching_cubes
+
+    lo = np.asarray(lo, float)
+    dims = np.asarray(dims, int)
+    B, lf = int(block), int(leaf)
+    if B % lf:
+        raise ValueError("block must be a multiple of leaf")
+    ncells = dims - 1
+    nb = np.maximum(1, np.ceil(ncells / B).astype(int))
+
+    def centres(i0, i1):
+        """Centres (mm) and half-diagonals of the boxes [i0, i1) in cells."""
+        return lo + 0.5 * (i0 + i1) * pitch, 0.5 * np.linalg.norm((i1 - i0) * pitch, axis=-1)
+
+    bi = np.stack(np.meshgrid(*[np.arange(n) for n in nb], indexing="ij"), -1).reshape(-1, 3)
+    b0 = bi * B
+    b1 = np.minimum(b0 + B, ncells)
+    C, half = centres(b0, b1)
+    fc = _eval(field, C)
+    if lipschitz is None:
+        L = _lipschitz(field, lo, lo + ncells * pitch, pitch)
+        F3 = fc.reshape(tuple(nb))
+        for k in range(3):
+            if nb[k] > 1:
+                L = max(L, float(np.max(np.abs(np.diff(F3, axis=k)))) / (B * pitch))
+        L = max(1.0, L * margin)
+    else:
+        L = float(lipschitz)
+    on_border = (b0 == 0).any(axis=1) | (b1 == ncells).any(axis=1)
+    act = (np.abs(fc) <= L * half + pitch) | (on_border & (fc < 0))
+    bi, b0, b1 = bi[act], b0[act], b1[act]
+
+    # leaves of the passing blocks
+    nl = B // lf
+    li = np.stack(np.meshgrid(*[np.arange(nl)] * 3, indexing="ij"), -1).reshape(-1, 3)
+    l0 = (b0[:, None, :] + li[None] * lf)
+    l1 = np.minimum(l0 + lf, b1[:, None, :])
+    real = (l0 < b1[:, None, :]).all(axis=2)            # leaves past a short block's end are empty
+    Cl, hl = centres(l0.reshape(-1, 3), l1.reshape(-1, 3))
+    fl = np.full(len(Cl), np.inf)
+    rf = real.ravel()
+    fl[rf] = _eval(field, Cl[rf])
+    fl = fl.reshape(len(b0), -1)
+    Ll = L
+    if lipschitz is None and _LOCAL:
+        Ll = np.minimum(L, _local_lipschitz(fl, nl, lf * pitch, _LOCAL_MARGIN, max(1.0, margin)))
+    lon = ((l0 == 0) | (l1 == ncells)).any(axis=2)
+    lact = real & ((np.abs(fl) <= Ll * hl.reshape(len(b0), -1) + pitch) | (lon & (fl < 0)))
+    keep = lact.any(axis=1)
+    b0, b1, fl, lact = b0[keep], b1[keep], fl[keep], lact[keep]
+
+    def mask(k):
+        """The samples block k needs: every sample of a passing leaf."""
+        n = b1[k] - b0[k] + 1
+        # a leaf's samples run from its first corner to one past its last
+        # cell: its cells, grown by one sample on the high side of each axis
+        c = np.zeros((nl * lf + 1,) * 3, bool)
+        c[:-1, :-1, :-1] = np.repeat(np.repeat(np.repeat(lact[k].reshape(nl, nl, nl), lf, 0), lf, 1), lf, 2)
+        m = c.copy()
+        m[1:] |= c[:-1]
+        m[:, 1:] |= m[:, :-1].copy()
+        m[:, :, 1:] |= m[:, :, :-1].copy()
+        return m[:n[0], :n[1], :n[2]]
+
+    fine = None
+    if budget is not None:
+        fine = sum(int(mask(k).sum()) for k in range(len(b0)))
+        if fine > budget:
+            raise ValueError(f"the surface needs {fine / 1e6:.1f}M samples at pitch {pitch:g}, over the "
+                             f"{budget / 1e6:.0f}M budget: use a coarser pitch")
+    if stats is None:
+        stats = {}
+    stats.update({"blocks": int(np.prod(nb)), "activeBlocks": int(len(b0)), "leaves": int(lact.sum()),
+                  "lipschitz": round(float(L), 3) if np.isfinite(L) else None,
+                  "samples": int(len(C) + rf.sum()), "denseSamples": int(np.prod(dims))})
+
+    big = np.float32(1e3 * pitch)
+    k = 0
+    while k < len(b0):
+        # a batch of blocks, one field call
+        group, P, cnt = [], [], 0
+        while k < len(b0) and (not group or cnt < _BATCH):
+            n = b1[k] - b0[k] + 1
+            ii = np.argwhere(mask(k))
+            P.append(lo + (b0[k] + ii) * pitch)
+            group.append((k, n, ii))
+            cnt += len(ii)
+            k += 1
+        F = _eval(field, np.concatenate(P)).astype(np.float32)
+        del P
+        stats["samples"] += len(F)
+        o = 0
+        for kk, n, ii in group:
+            # unevaluated samples: their leaf's sign, far from zero
+            lix = np.minimum(np.indices(tuple(n)).reshape(3, -1).T // lf, nl - 1)
+            sgn = np.sign(fl[kk][(lix[:, 0] * nl + lix[:, 1]) * nl + lix[:, 2]])
+            f = (sgn * big).astype(np.float32).reshape(tuple(n))
+            f[tuple(ii.T)] = F[o:o + len(ii)]
+            o += len(ii)
+            i0 = b0[kk]
+            for ax in range(3):
+                if i0[ax] == 0:
+                    f[(slice(None),) * ax + (0,)] = pitch
+                if i0[ax] + n[ax] == dims[ax]:
+                    f[(slice(None),) * ax + (-1,)] = pitch
+            f[f == 0.0] = -1e-4 * pitch
+            if f.min() >= 0 or f.max() <= 0:
+                continue
+            v, fa, _, _ = marching_cubes(f, 0.0)
+            yield i0, n, v.astype(np.float64) + i0, fa
+        if progress:
+            progress({"stage": "field", "done": int(k), "of": int(len(b0))})
+
+
+# ---- fast bookkeeping on big triangle meshes (trimesh's takes tens of seconds) ----
+
+def watertight(F, nv=None):
+    """Every edge shared by exactly two faces (trimesh's is_watertight)."""
+    F = np.asarray(F, np.int64)
+    if not len(F):
+        return False
+    nv = int(F.max()) + 1 if nv is None else int(nv)
+    a = F.ravel()
+    b = np.roll(F, -1, axis=1).ravel()
+    e = np.sort(np.minimum(a, b) * nv + np.maximum(a, b))
+    if len(e) % 2:
+        return False
+    if not np.array_equal(e[0::2], e[1::2]):
+        return False
+    return bool(len(e) < 4 or not np.any(e[1:-1:2] == e[2::2]))
+
+
+def face_cross(V, F):
+    """(b - a) x (c - a) for every face (a, b, c): twice its area along its
+    normal (compiled, meshops.py; else column by column in numpy)."""
+    from .meshops import face_cross as fc
+
+    return fc(V, F)
+
+
+def _face_cross_np(V, F):
+    # column by column, ~3x numpy's cross on (n, 3) rows
+    VT = np.ascontiguousarray(np.asarray(V, float).T)
+    a = VT[:, F[:, 0]]
+    b = VT[:, F[:, 1]]
+    c = VT[:, F[:, 2]]
+    b -= a
+    c -= a
+    out = np.empty((len(F), 3))
+    out[:, 0] = b[1] * c[2] - b[2] * c[1]
+    out[:, 1] = b[2] * c[0] - b[0] * c[2]
+    out[:, 2] = b[0] * c[1] - b[1] * c[0]
+    return out
+
+
+def _triple(V, F):
+    """a . (b x c) per face (six times its signed volume from the origin)."""
+    from .meshops import triple
+
+    return triple(V, F)
+
+
+def _triple_np(V, F):
+    VT = np.ascontiguousarray(np.asarray(V, float).T)
+    a = VT[:, F[:, 0]]
+    b = VT[:, F[:, 1]]
+    c = VT[:, F[:, 2]]
+    return a[0] * (b[1] * c[2] - b[2] * c[1]) + a[1] * (b[2] * c[0] - b[0] * c[2]) + a[2] * (b[0] * c[1] - b[1] * c[0])
+
+
+def signed_volume(V, F):
+    """Enclosed volume (divergence theorem), positive with outward normals."""
+    return float(_triple(V, F).sum() / 6.0)
+
+
+def shell_volumes(V, F):
+    """(label per face, signed volume per connected shell): shells joined
+    through shared vertices. A closed body's outside is positive, the wall
+    of a void sealed inside it negative."""
+    import scipy.sparse as sp
+    from scipy.sparse.csgraph import connected_components
+
+    n = len(V)
+    A = sp.coo_matrix((np.ones(2 * len(F), np.int8), (np.r_[F[:, 0], F[:, 1]], np.r_[F[:, 1], F[:, 2]])),
+                      shape=(n, n))
+    _, lab = connected_components(A, directed=False)
+    fl = lab[F[:, 0]]
+    vol = np.bincount(fl, weights=_triple(V, F) / 6.0)
+    # renumber to the shells that have faces
+    used = np.unique(fl)
+    ren = np.full(lab.max() + 1, -1)
+    ren[used] = np.arange(len(used))
+    return ren[fl], vol[used]

loadpath/nlkernel.py

@@ -0,0 +1,241 @@
+"""
+Compiled element kernel for the large-deflection solver (nonlinear.py).
+
+nonlinear._element builds each hexahedron's internal force and tangent from
+stacks of small numpy operations (3x3 inverses, determinants, 4th-order
+einsums, transposed copies), whose overhead, not their arithmetic, was most
+of a large-deflection solve: about 54 us per element per Newton iteration on
+the 1056-cell tendon finger. Here the same element (total Lagrangian hex8,
+2x2x2 Gauss, neo-Hookean or St. Venant-Kirchhoff, optionally with the nine
+incompatible modes condensed out) is one compiled loop over elements, run on
+all cores, and the scatter into the CSR tangent is a second compiled loop.
+
+Same arithmetic as the numpy element, in a different summation order, so
+forces and tangents agree to ~1e-14 relative, not bit for bit.
+LOADPATH_NLKERNEL=0 runs the numpy element instead.
+"""
+from __future__ import annotations
+
+import os
+
+import numpy as np
+
+try:
+    import numba as _nb
+except ImportError:  # pragma: no cover - numba is a dependency, kept optional
+    _nb = None
+
+__all__ = ["available", "elements", "scatter"]
+
+
+def available() -> bool:
+    return _nb is not None and os.environ.get("LOADPATH_NLKERNEL", "1") != "0"
+
+
+if _nb is not None:
+    @_nb.njit(cache=True, nogil=True)
+    def _gauss(ux, G, h, lam, mu, svk, fx, kx):
+        """Accumulate internal force fx (3n) and tangent kx (3n, 3n) of one
+        element with generalized displacements ux (n, 3); returns min det F."""
+        nn = ux.shape[0]
+        w = h ** 3 / 8.0
+        wk = w / (h * h)
+        jmin = np.inf
+        F = np.empty((3, 3))
+        S = np.empty((3, 3))
+        P = np.empty((3, 3))
+        Ci = np.empty((3, 3))
+        A = np.empty((3, 3, 3, 3))
+        FA = np.empty((3, 3, 3, 3))
+        Q = np.empty((3, 3, 3, 3))
+        T = np.empty((nn, 3, 3, 3))
+        for g in range(G.shape[0]):
+            for i in range(3):
+                for j in range(3):
+                    s = 1.0 if i == j else 0.0
+                    for a in range(nn):
+                        s += ux[a, i] * G[g, a, j] / h
+                    F[i, j] = s
+            J = (F[0, 0] * (F[1, 1] * F[2, 2] - F[1, 2] * F[2, 1])
+                 - F[0, 1] * (F[1, 0] * F[2, 2] - F[1, 2] * F[2, 0])
+                 + F[0, 2] * (F[1, 0] * F[2, 1] - F[1, 1] * F[2, 0]))
+            if J < jmin:
+                jmin = J
+            # C = F^T F
+            c00 = F[0, 0] * F[0, 0] + F[1, 0] * F[1, 0] + F[2, 0] * F[2, 0]
+            c11 = F[0, 1] * F[0, 1] + F[1, 1] * F[1, 1] + F[2, 1] * F[2, 1]
+            c22 = F[0, 2] * F[0, 2] + F[1, 2] * F[1, 2] + F[2, 2] * F[2, 2]
+            c01 = F[0, 0] * F[0, 1] + F[1, 0] * F[1, 1] + F[2, 0] * F[2, 1]
+            c12 = F[0, 1] * F[0, 2] + F[1, 1] * F[1, 2] + F[2, 1] * F[2, 2]
+            c02 = F[0, 0] * F[0, 2] + F[1, 0] * F[1, 2] + F[2, 0] * F[2, 2]
+            if svk:
+                e00, e11, e22 = 0.5 * (c00 - 1.0), 0.5 * (c11 - 1.0), 0.5 * (c22 - 1.0)
+                tr = e00 + e11 + e22
+                S[0, 0] = lam * tr + 2 * mu * e00
+                S[1, 1] = lam * tr + 2 * mu * e11
+                S[2, 2] = lam * tr + 2 * mu * e22
+                S[0, 1] = S[1, 0] = mu * c01
+                S[1, 2] = S[2, 1] = mu * c12
+                S[0, 2] = S[2, 0] = mu * c02
+                for i in range(3):
+                    for j in range(3):
+                        for k in range(3):
+                            for l in range(3):
+                                A[i, j, k, l] = (lam * (i == j) * (k == l)
+                                                 + mu * ((i == k) * (j == l) + (i == l) * (j == k)))
+            else:
+                # C^-1 by cofactors (C symmetric); det C = J^2
+                d = c00 * (c11 * c22 - c12 * c12) - c01 * (c01 * c22 - c12 * c02) + c02 * (c01 * c12 - c11 * c02)
+                Ci[0, 0] = (c11 * c22 - c12 * c12) / d
+                Ci[1, 1] = (c00 * c22 - c02 * c02) / d
+                Ci[2, 2] = (c00 * c11 - c01 * c01) / d
+                Ci[0, 1] = Ci[1, 0] = (c02 * c12 - c01 * c22) / d
+                Ci[1, 2] = Ci[2, 1] = (c01 * c02 - c00 * c12) / d
+                Ci[0, 2] = Ci[2, 0] = (c01 * c12 - c02 * c11) / d
+                lnJ = np.log(max(J, 1e-12))
+                for i in range(3):
+                    for j in range(3):
+                        S[i, j] = mu * ((1.0 if i == j else 0.0) - Ci[i, j]) + lam * lnJ * Ci[i, j]
+                cc = mu - lam * lnJ
+                for i in range(3):
+                    for j in range(3):
+                        for k in range(3):
+                            for l in range(3):
+                                A[i, j, k, l] = (lam * Ci[i, j] * Ci[k, l]
+                                                 + cc * (Ci[i, k] * Ci[j, l] + Ci[i, l] * Ci[j, k]))
+            for i in range(3):
+                for l in range(3):
+                    P[i, l] = F[i, 0] * S[0, l] + F[i, 1] * S[1, l] + F[i, 2] * S[2, l]
+            for a in range(nn):
+                for i in range(3):
+                    fx[3 * a + i] += w * (P[i, 0] * G[g, a, 0] + P[i, 1] * G[g, a, 1] + P[i, 2] * G[g, a, 2]) / h
+            # Q_iljn = F_ik A_klpn F_jp + S_ln delta_ij
+            for i in range(3):
+                for l in range(3):
+                    for p in range(3):
+                        for n in range(3):
+                            FA[i, l, p, n] = F[i, 0] * A[0, l, p, n] + F[i, 1] * A[1, l, p, n] + F[i, 2] * A[2, l, p, n]
+            for i in range(3):
+                for l in range(3):
+                    for j in range(3):
+                        for n in range(3):
+                            v = FA[i, l, 0, n] * F[j, 0] + FA[i, l, 1, n] * F[j, 1] + FA[i, l, 2, n] * F[j, 2]
+                            if i == j:
+                                v += S[l, n]
+                            Q[i, l, j, n] = v
+            # K_aibj += wk G_al Q_iljn G_bn, upper triangle in (a, b)
+            for a in range(nn):
+                for i in range(3):
+                    for j in range(3):
+                        for n in range(3):
+                            T[a, i, j, n] = G[g, a, 0] * Q[i, 0, j, n] + G[g, a, 1] * Q[i, 1, j, n] + G[g, a, 2] * Q[i, 2, j, n]
+            for a in range(nn):
+                for b in range(a, nn):
+                    gb0, gb1, gb2 = G[g, b, 0], G[g, b, 1], G[g, b, 2]
+                    for i in range(3):
+                        for j in range(3):
+                            kx[3 * a + i, 3 * b + j] += wk * (T[a, i, j, 0] * gb0 + T[a, i, j, 1] * gb1 + T[a, i, j, 2] * gb2)
+        n3 = 3 * nn
+        for r in range(n3):
+            for c in range(r):
+                if c // 3 < r // 3:
+                    kx[r, c] = kx[c, r]
+        return jmin
+
+    @_nb.njit(cache=True, parallel=True)
+    def _elements(u, edof, alpha, mod, h, lam, mu, svk, tangent, G, im):
+        m = edof.shape[0]
+        nn = G.shape[1]
+        fe = np.zeros((m, 24))
+        ke = np.zeros((m, 24, 24)) if tangent else np.zeros((1, 24, 24))
+        jm = np.empty(m)
+        anew = np.zeros((m, 3, 3)) if im else np.zeros((1, 3, 3))
+        ra = np.zeros((m, 9)) if im else np.zeros((1, 9))
+        cpl = np.zeros((m, 9, 24)) if im else np.zeros((1, 9, 24))
+        for e in _nb.prange(m):
+            ux = np.zeros((nn, 3))
+            for a in range(8):
+                for i in range(3):
+                    ux[a, i] = u[edof[e, 3 * a + i]]
+            if im:
+                for a in range(3):
+                    for i in range(3):
+                        ux[8 + a, i] = alpha[e, a, i]
+            fx = np.zeros(3 * nn)
+            kx = np.zeros((3 * nn, 3 * nn))
+            jm[e] = _gauss(ux, G, h, lam, mu, svk, fx, kx)
+            me = mod[e]
+            if not im:
+                for r in range(24):
+                    fe[e, r] = fx[r] * me
+                if tangent:
+                    for r in range(24):
+                        for c in range(24):
+                            ke[e, r, c] = kx[r, c] * me
+                continue
+            # condense the nine incompatible modes: one local Newton step
+            kaa = kx[24:, 24:].copy()
+            rhs = np.empty((9, 25))
+            for r in range(9):
+                rhs[r, 0] = fx[24 + r]
+                for c in range(24):
+                    rhs[r, 1 + c] = kx[c, 24 + r]          # K_au = K_ua^T
+            sol = np.linalg.solve(kaa, rhs)
+            for r in range(24):
+                s = fx[r]
+                for q in range(9):
+                    s -= kx[r, 24 + q] * sol[q, 0]
+                fe[e, r] = s * me
+            if tangent:
+                for r in range(24):
+                    for c in range(24):
+                        s = kx[r, c]
+                        for q in range(9):
+                            s -= kx[r, 24 + q] * sol[q, 1 + c]
+                        ke[e, r, c] = s * me
+            for a in range(3):
+                for i in range(3):
+                    anew[e, a, i] = alpha[e, a, i] - sol[3 * a + i, 0]
+            for q in range(9):
+                ra[e, q] = fx[24 + q]
+                for c in range(24):
+                    cpl[e, q, c] = sol[q, 1 + c]
+        return fe, ke, jm, anew, ra, cpl
+
+    @_nb.njit(cache=True, nogil=True)
+    def _scatter(data, slots, ke, fint, fe, edof, tangent):
+        m = edof.shape[0]
+        for e in range(m):
+            for r in range(24):
+                fint[edof[e, r]] += fe[e, r]
+            if tangent:
+                for a in range(8):
+                    for b in range(8):
+                        p = 8 * a + b
+                        for i in range(3):
+                            for j in range(3):
+                                data[slots[p, e, 3 * i + j]] += ke[e, 3 * a + i, 3 * b + j]
+
+
+def elements(u, edof, alpha, mod, h, nu, kind, tangent, grads):
+    """All elements' (fint (m, 24), ke (m, 24, 24), det F per element,
+    alpha_new, R_alpha (m, 9), K_aa^-1 K_au (m, 9, 24)), scaled by the
+    moduli (fint, ke) as nonlinear._assemble scatters them. grads: (8, n, 3)
+    unit-cube gradients, n = 8 (plain) or 11 (incompatible modes)."""
+    lam, mu = nu / ((1 + nu) * (1 - 2 * nu)), 1.0 / (2 * (1 + nu))
+    im = alpha is not None
+    a = np.ascontiguousarray(alpha, float) if im else np.zeros((1, 3, 3))
+    return _elements(np.ascontiguousarray(u, float), np.ascontiguousarray(edof), a,
+                     np.ascontiguousarray(mod, float), float(h), lam, mu, kind == "svk", bool(tangent),
+                     np.ascontiguousarray(grads, float), im)
+
+
+def scatter(n_dof, nnz, slots, edof, fe, ke, tangent):
+    """Global internal force and the CSR data array of the tangent."""
+    fint = np.zeros(n_dof)
+    data = np.zeros(nnz) if tangent else np.zeros(1)
+    if ke is None or not tangent:
+        ke = np.zeros((1, 24, 24))
+    _scatter(data, np.ascontiguousarray(slots), np.ascontiguousarray(ke), fint, np.ascontiguousarray(fe),
+             np.ascontiguousarray(edof), bool(tangent))
+    return fint, (data if tangent else None)