fields.py

"""Derived per-frame fields for visualisation.

Nothing here feeds the solver. Every quantity is a *read* of a state the solver has already
accepted, computed once per exported frame and thrown at a file -- so the rules that govern
the solver's own kernels (fixed shapes, no host sync, differentiability) do not apply, and
these are written for clarity instead.

The one rule that does carry over is that the deformation gradient is computed **once** per
frame and shared. ``jacobian`` and ``energy_density`` both need it, and at six million tets
an accidental second ``vmap`` over ``inv(Dm)`` is a real cost paid on every single frame.

Field *selection* is the scale story. A million-vertex mesh writing displacement, velocity,
speed, J and energy density costs ~76 MB per frame; the same mesh writing points and J costs
~28 MB. Which fields exist is therefore a caller's decision, not a default, and everything
below is addressed by name from a registry rather than computed eagerly.
"""

from __future__ import annotations

import jax
import jax.numpy as jnp
import numpy as np

from diff_vbd.model import SimulationProblem, SimulationState
from diff_vbd.solver.kinematics import deformation_gradient, tet_volume
from diff_vbd.solver.materials import stable_neo_hookean_energy_density

POINT_FIELDS = ("displacement", "velocity", "speed")
CELL_FIELDS = ("jacobian", "energy_density", "energy", "volume_ratio")

#: What a caller gets if it does not choose. Deliberately frugal: the surface stream is
#: written every frame, so its default is the two things you cannot reconstruct afterwards
#: (velocity is gone once the frame is written; speed is what you actually colour by).
DEFAULT_SURFACE_POINT_FIELDS = ("velocity", "speed")
#: The volume stream exists to answer "is any element inverted, and where is the energy",
#: so J earns its place. Everything else there is opt-in.
DEFAULT_VOLUME_CELL_FIELDS = ("jacobian",)


def to_numpy(value) -> np.ndarray:
    return np.asarray(jax.device_get(value))


@jax.jit
def deformation_gradients(
    rest_positions: jnp.ndarray, positions: jnp.ndarray, tets: jnp.ndarray
) -> jnp.ndarray:
    """Return F for every tetrahedron, shape ``(M, 3, 3)``."""
    return jax.vmap(deformation_gradient)(rest_positions[tets], positions[tets])


@jax.jit
def rest_volumes(rest_positions: jnp.ndarray, tets: jnp.ndarray) -> jnp.ndarray:
    return jax.vmap(tet_volume)(rest_positions[tets])


@jax.jit
def energy_densities(material, gradients: jnp.ndarray) -> jnp.ndarray:
    return jax.vmap(stable_neo_hookean_energy_density, in_axes=(None, 0))(
        material, gradients
    )


def compute_point_fields(
    problem: SimulationProblem,
    state: SimulationState,
    names: tuple[str, ...],
    dtype=np.float32,
    indices: jnp.ndarray | None = None,
) -> dict[str, np.ndarray]:
    """Compute the requested per-vertex fields for one frame.

    ``indices`` subsets **before** the device transfer, not after. The surface stream wants
    perhaps 4% of the vertices of a large mesh, and gathering on the host means moving the
    other 96% across the bus every frame in order to throw it away.
    """
    unknown = set(names) - set(POINT_FIELDS)
    if unknown:
        raise ValueError(
            f"unknown point field(s) {sorted(unknown)}; available: {list(POINT_FIELDS)}"
        )

    def gather(array: jnp.ndarray) -> jnp.ndarray:
        return array if indices is None else array[indices]

    fields: dict[str, np.ndarray] = {}
    if "displacement" in names:
        fields["displacement"] = to_numpy(
            gather(state.position) - gather(problem.mesh.rest_positions)
        ).astype(dtype)
    if "velocity" in names:
        fields["velocity"] = to_numpy(gather(state.velocity)).astype(dtype)
    if "speed" in names:
        fields["speed"] = to_numpy(
            jnp.linalg.norm(gather(state.velocity), axis=-1)
        ).astype(dtype)
    return fields


def compute_cell_fields(
    problem: SimulationProblem,
    state: SimulationState,
    names: tuple[str, ...],
    dtype=np.float32,
) -> dict[str, np.ndarray]:
    """Compute the requested per-tet fields for one frame, sharing one F across them."""
    unknown = set(names) - set(CELL_FIELDS)
    if unknown:
        raise ValueError(
            f"unknown cell field(s) {sorted(unknown)}; available: {list(CELL_FIELDS)}"
        )
    if not names:
        return {}

    gradients = deformation_gradients(
        problem.mesh.rest_positions, state.position, problem.mesh.tets
    )
    fields: dict[str, np.ndarray] = {}

    # J = det F. Signed on purpose: stable Neo-Hookean is finite and smooth through
    # inversion, so an inverted element produces no error, no warning and no NaN -- it just
    # quietly keeps simulating. Thresholding J < 0 in the viewer is the only thing that will
    # ever tell you it happened.
    if {"jacobian", "volume_ratio"} & set(names):
        jacobians = jnp.linalg.det(gradients)
        if "jacobian" in names:
            fields["jacobian"] = to_numpy(jacobians).astype(dtype)
        if "volume_ratio" in names:
            fields["volume_ratio"] = to_numpy(jnp.abs(jacobians)).astype(dtype)

    if {"energy_density", "energy"} & set(names):
        densities = energy_densities(problem.material, gradients)
        if "energy_density" in names:
            fields["energy_density"] = to_numpy(densities).astype(dtype)
        if "energy" in names:
            volumes = rest_volumes(problem.mesh.rest_positions, problem.mesh.tets)
            fields["energy"] = to_numpy(densities * volumes).astype(dtype)

    return fields


def build_body_ids(tets: np.ndarray, num_vertices: int) -> np.ndarray:
    """Label each vertex with the connected component of the tet mesh it belongs to.

    A scene like ``two_blocks`` is one mesh with two disjoint components, and without a label
    there is no way to select one of them in a viewer -- the tets interleave in index order,
    so no index range picks out a body. With the label it is one Threshold filter.

    Connecting each tet's corners to its zeroth corner is enough: the remaining corner pairs
    are connected transitively through it, so three edges per tet suffice instead of six.
    """
    tets = np.asarray(tets, dtype=np.int64)
    rows = np.repeat(tets[:, 0], 3)
    cols = tets[:, 1:].reshape(-1)

    try:
        from scipy.sparse import coo_matrix
        from scipy.sparse.csgraph import connected_components
    except ImportError:
        return _body_ids_union_find(rows, cols, num_vertices)

    graph = coo_matrix(
        (np.ones(rows.shape[0], dtype=np.int8), (rows, cols)),
        shape=(num_vertices, num_vertices),
    )
    _, labels = connected_components(graph, directed=False)
    return labels.astype(np.int32)


def _body_ids_union_find(
    rows: np.ndarray, cols: np.ndarray, num_vertices: int
) -> np.ndarray:
    """Fallback for environments without scipy. Correct, but interpreter-bound."""
    parent = list(range(num_vertices))

    def find(node: int) -> int:
        root = node
        while parent[root] != root:
            root = parent[root]
        while parent[node] != root:  # path compression, iterative to avoid recursion limits
            parent[node], node = root, parent[node]
        return root

    for source, target in zip(rows.tolist(), cols.tolist()):
        root_a, root_b = find(source), find(target)
        if root_a != root_b:
            parent[root_b] = root_a

    roots = np.array([find(v) for v in range(num_vertices)], dtype=np.int64)
    _, labels = np.unique(roots, return_inverse=True)
    return labels.astype(np.int32)


def static_point_fields(
    problem: SimulationProblem, dtype=np.float32
) -> dict[str, np.ndarray]:
    """Per-vertex quantities that do not change over the run, so are written once.

    The masks are ``int32`` rather than bool, because a bool array is not something XDMF can
    describe -- and int32 rather than the int8 that would obviously suffice, because XDMF
    would then declare ``Precision="1"`` and single-byte integer attributes are the corner of
    the format readers are least consistent about. This costs three bytes a vertex on data
    written exactly once, which is not a trade worth thinking about twice.
    """
    boundary = problem.boundary_conditions
    tets = to_numpy(problem.mesh.tets)
    num_vertices = int(to_numpy(problem.mesh.rest_positions).shape[0])
    return {
        "dirichlet": to_numpy(boundary.dirichlet_mask).astype(np.int32),
        "rigid": to_numpy(boundary.rigid_mask).astype(np.int32),
        "mass": to_numpy(problem.topology.mass).astype(dtype),
        "body_id": build_body_ids(tets, num_vertices),
    }