assets.py

"""Loading author-supplied validation geometry.

Validation geometry is a CAD artifact, not something this module invents. Each case owns a
directory under ``assets/validation/``::

    thick_cylinder/
      surface.stl        the CAD export, and the source of truth for the solid
      selector_*.stl     optional boundary-condition selector volumes
      h0.msh h1.msh ...  a refinement family, coarsest first
      manifest.json      which fTetWild invocation produced each level

Everything for one case sits together because a convergence study needs the *family*, not a
mesh: the whole point is to refine the same geometry and watch the error fall at a rate, so
the levels have to be enumerable and their provenance recorded. ``manifest.json`` stores the
fTetWild flags and the pinned submodule commit; everything measurable -- vertex and element
counts, the representative element size -- is computed from the mesh instead of stored, so it
cannot drift away from the file it describes.

This module also has to **verify** what it loads. An external mesh can arrive with inverted
tetrahedra, with dimensions that drifted from nominal, or exported in a format this repo
cannot read, and each of those would otherwise surface much later disguised as a solver
defect. An inverted rest element in particular flips the sign of ``det F`` for that element
permanently, and would be reported as an inversion by the very diagnostic built to detect
real ones.

Format is not negotiable: ``setup.mesh_io.parse_gmsh22_binary_tets`` reads **Gmsh 2.2,
binary, little-endian** and rejects everything else, including Gmsh 4.x and ASCII 2.2.
fTetWild writes exactly that by default, so ``--no-binary`` must not be passed and
``FLOAT_TETWILD_USE_FLOAT`` must stay ``OFF`` -- float output does not raise, it misparses
silently against the parser's hardcoded 8-byte layout.

``REQUIRED_ASSETS`` is the authoritative specification for every case, quoted verbatim into
the error message when one is missing, so the requirements cannot drift away from the code
that consumes them.
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from pathlib import Path

import jax.numpy as jnp
import numpy as np

from diff_vbd.setup.mesh_io import parse_gmsh22_binary_tets

#: Overrides asset discovery, for installed packages or a relocated asset tree.
ASSET_ROOT_ENV = "DIFF_VBD_VALIDATION_ASSETS"


class MissingValidationAsset(FileNotFoundError):
    """Raised when required validation geometry has not been authored yet.

    Carries the specification rather than just the path, because the fix is to produce a mesh
    to particular requirements and a bare "file not found" does not say what they are.
    """


@dataclass(frozen=True)
class MeshAsset:
    """One refinement level of one validation case."""

    name: str
    level: int
    path: Path
    positions: jnp.ndarray
    tets: jnp.ndarray
    flags: tuple[str, ...] = ()

    @property
    def extent(self) -> tuple[float, float, float]:
        points = np.asarray(self.positions)
        return tuple(float(v) for v in (points.max(axis=0) - points.min(axis=0)))

    @property
    def num_vertices(self) -> int:
        return int(self.positions.shape[0])

    @property
    def num_tets(self) -> int:
        return int(self.tets.shape[0])

    @property
    def surface_path(self) -> Path:
        return self.path.parent / "surface.stl"

    @property
    def element_size(self) -> float:
        """Representative element size ``h``, measured from the mesh.

        The cube root of the mean tetrahedron volume. This is what a convergence plot needs
        on its abscissa -- *not* the ``-a`` that was requested, because fTetWild treats that
        as a target for an optimisation and the delivered mesh routinely differs from it. A
        study plotted against the requested value would report the wrong slope.
        """
        points = np.asarray(self.positions)[np.asarray(self.tets)]
        volumes = np.abs(np.linalg.det(points[:, 1:] - points[:, :1])) / 6.0
        return float(np.cbrt(volumes.mean()))

    def describe(self) -> str:
        length, height, depth = self.extent
        return (
            f"{self.name}[h{self.level}]: {self.num_vertices} verts, {self.num_tets} tets, "
            f"h={self.element_size:.4g}, extent {length:.4g} x {height:.4g} x {depth:.4g}"
        )


def asset_root() -> Path:
    """Directory holding validation cases.

    ``$DIFF_VBD_VALIDATION_ASSETS`` wins if set; otherwise ``assets/validation`` relative to
    the repository root, resolved from this file rather than the working directory.
    """
    override = os.environ.get(ASSET_ROOT_ENV)
    if override:
        return Path(override).expanduser()
    return Path(__file__).resolve().parents[3] / "assets" / "validation"


def case_root(name: str) -> Path:
    return asset_root() / name


def load_manifest(name: str) -> dict:
    path = case_root(name) / "manifest.json"
    if not path.is_file():
        raise MissingValidationAsset(_missing_message(name, path))
    return json.loads(path.read_text())


def refinement_levels(name: str) -> list[int]:
    """Indices of the meshed refinement levels for a case, coarsest first."""
    directory = case_root(name)
    if not directory.is_dir():
        return []
    levels = []
    for path in directory.glob("h*.msh"):
        try:
            levels.append(int(path.stem[1:]))
        except ValueError:  # pragma: no cover -- ignore anything not h<N>.msh
            continue
    return sorted(levels)


def available_assets() -> list[str]:
    """Cases with at least one meshed level."""
    root = asset_root()
    if not root.is_dir():
        return []
    return sorted(d.name for d in root.iterdir() if d.is_dir() and refinement_levels(d.name))


def load_mesh(name: str, level: int = 0, *, dtype=jnp.float64) -> MeshAsset:
    """Load one refinement level. ``level=0`` is the coarsest.

    Coarsest by default deliberately: a test that only needs *a* discretisation should pay
    for the cheapest one, and a convergence study asks for the family explicitly.
    """
    path = case_root(name) / f"h{level}.msh"
    if not path.is_file():
        raise MissingValidationAsset(_missing_message(name, path))
    positions, tets = parse_gmsh22_binary_tets(path, dtype=dtype)
    flags = ()
    manifest_path = case_root(name) / "manifest.json"
    if manifest_path.is_file():
        for entry in json.loads(manifest_path.read_text()).get("levels", []):
            if entry.get("file") == path.name:
                flags = tuple(entry.get("flags", ()))
    return MeshAsset(
        name=name, level=level, path=path, positions=positions, tets=tets, flags=flags
    )


def load_refinement_family(name: str, *, dtype=jnp.float64) -> list[MeshAsset]:
    """Every meshed level of a case, coarsest first, with its provenance verified.

    Raises if fewer than two levels exist: a convergence study needs at least two points to
    have a slope at all, and silently reporting a rate from one mesh would be worse than
    refusing.

    The three consistency checks below exist because a family is easy to corrupt without
    noticing. Two mesher runs racing over the same directory -- one killed but not actually
    dead -- will happily leave a mix of levels from different parameter sets alongside a
    manifest describing neither, and the result still looks like a valid family: N files,
    N entries, decreasing sizes. A convergence rate measured from that is meaningless, and
    nothing else in the pipeline would question it.
    """
    levels = refinement_levels(name)
    if len(levels) < 2:
        raise MissingValidationAsset(
            f"case {name!r} has {len(levels)} meshed level(s); a convergence study needs at "
            f"least two. Mesh the same surface at additional element sizes and record them "
            f"in {case_root(name) / 'manifest.json'}."
        )

    declared = [entry.get("file") for entry in load_manifest(name).get("levels", [])]
    on_disk = [f"h{level}.msh" for level in levels]
    if declared != on_disk:
        raise MissingValidationAsset(
            f"case {name!r} is inconsistent: manifest declares {declared} but the directory "
            f"holds {on_disk}. Re-mesh the family so every level's provenance is recorded; "
            f"a convergence study over meshes of unknown parameters measures nothing."
        )

    family = [load_mesh(name, level, dtype=dtype) for level in levels]
    sizes = [mesh.element_size for mesh in family]
    if any(fine >= coarse for coarse, fine in zip(sizes, sizes[1:])):
        raise MissingValidationAsset(
            f"case {name!r} is not monotonically refining: measured element sizes are "
            f"{[round(s, 4) for s in sizes]}, coarsest first. Levels must strictly decrease, "
            f"or the abscissa of the convergence plot is not ordered."
        )
    return family


def selector_path(name: str, role: str) -> Path:
    path = case_root(name) / f"selector_{role}.stl"
    if not path.is_file():
        raise MissingValidationAsset(
            f"selector {role!r} for case {name!r} not found at {path}"
        )
    return path


def _missing_message(name: str, path: Path) -> str:
    spec = REQUIRED_ASSETS.get(name)
    lines = [f"validation geometry for {name!r} not found at {path}"]
    if spec is not None:
        lines.append(f"  purpose : {spec['purpose']}")
        lines.append(f"  geometry: {spec['geometry']}")
        lines.append(f"  meshing : {spec['meshing']}")
        lines.append(f"  export  : {spec['export']}")
        surface = case_root(name) / "surface.stl"
        if surface.is_file():
            lines.append(f"  surface : {surface} (mesh this with fTetWild)")
    lines.append(f"  levels  : {refinement_levels(name) or 'none meshed'}")
    lines.append(f"  cases   : {available_assets() or 'no meshed cases in the asset root'}")
    return "\n".join(lines)


_GMSH_EXPORT = "Gmsh 2.2, binary, little-endian (Gmsh 4.x and ASCII are rejected)"

#: The cases the validation suite expects, and what each must satisfy -- the authoritative
#: specification, quoted verbatim into the error message when one is missing, so the
#: requirements cannot drift away from the code that consumes them.
REQUIRED_ASSETS: dict[str, dict[str, str]] = {
    "block_dynamics": {
        "purpose": (
            "free fall, single-DOF oscillator, momentum conservation, and block-on-incline "
            "Coulomb friction -- the dynamic tier, none of which is blocked by the solver's "
            "static-convergence problem"
        ),
        "geometry": (
            "box 4 x 4 x 2, centred in x and y, bottom face on z = 0. Wide and short so it "
            "does not tip on an inclined plane; the flat bottom face is what contacts the "
            "plane collider"
        ),
        "meshing": "linear tets, -a 0.5. Coarse on purpose: these cases run thousands of "
        "timesteps and need a body, not resolution",
        "export": _GMSH_EXPORT,
    },
    "bar_axial": {
        "purpose": (
            "uniaxial bar hanging under self-weight, against the exact 1-D solution. Loads "
            "an axial mode, which converges far faster under Gauss-Seidel than the bending "
            "mode a cantilever excites -- so it localises the convergence fault"
        ),
        "geometry": (
            "prism 5 x 5 x 50, long axis along +Z spanning z in [0, 50], centred in x and y. "
            "The clamp is the TOP face at z = 50 (hanging, pure tension); selector_top.stl "
            "encloses it"
        ),
        "meshing": "linear tets, -a 1.25 -e 1e-4. The tightened envelope matters for slender "
        "parts: -e scales with the bounding-box DIAGONAL, so a long thin solid gets a loose "
        "envelope on its small cross-section and loses volume",
        "export": _GMSH_EXPORT,
    },
    "halfspace_block": {
        "purpose": "shared specimen for the Hertz, flat-punch and cone indentation cases",
        "geometry": (
            "box 8 x 8 x 4, top face on z = 0 with the body in z in [-4, 0], centred in x "
            "and y. Contact patch at the origin on the top face"
        ),
        "meshing": (
            "GRADED, not uniform: ~0.15 within radius 1.5 of the origin relaxing to ~0.8 at "
            "the boundary. Use fTetWild --bg-mesh with a per-vertex sizing field. A uniform "
            "mesh resolving the contact patch over this block runs to millions of elements, "
            "which the solver cannot step"
        ),
        "export": _GMSH_EXPORT,
    },
    "indenter_sphere": {
        "purpose": (
            "Hertz sphere indentation, and the mesh-mesh half of the differential test "
            "against the analytic sphere collider (same closed form, disjoint code paths)"
        ),
        "geometry": (
            "spherical cap on a cylindrical backing, NOT a full sphere: cap of sphere radius "
            "4 and height 0.6 (base radius ~2.11), backed by a cylinder giving a flat face "
            "to drive. Pole at the origin, axis +Z, body above"
        ),
        "meshing": "linear tets. Rigid and Dirichlet-driven, so only the contact surface "
        "matters -- the interior can be coarse",
        "export": _GMSH_EXPORT,
    },
    "indenter_flat_punch": {
        "purpose": "Sneddon flat punch: load is LINEAR in indentation depth, a different "
        "signature from Hertz's delta^1.5",
        "geometry": "cylinder radius 1.0, height 2.0, flat face on z = 0, axis +Z, body above",
        "meshing": (
            "linear tets with a TIGHTENED envelope -e 1e-4. The sharp circular rim is the "
            "solution -- Sneddon's result depends on the stress singularity there, and a "
            "default envelope rounds it off"
        ),
        "export": _GMSH_EXPORT,
    },
    "indenter_cone": {
        "purpose": "Sneddon cone: load goes as delta^2, the third distinct exponent",
        "geometry": (
            "cone of semi-angle 70 degrees from the axis (blunt), height 1.0, base radius "
            "~2.75, tip rounded to radius 0.05 so it is meshable. Tip at the origin, axis +Z"
        ),
        "meshing": "linear tets. The rounded tip changes the solution locally, so validate "
        "the delta^2 trend rather than tip stress",
        "export": _GMSH_EXPORT,
    },
    "thick_cylinder": {
        "purpose": (
            "Lame thick-walled cylinder under internal pressure: the pointwise-FIELD "
            "validation case (u_r, sigma_rr, sigma_theta have exact closed forms, so the "
            "error is meaningful everywhere). NOT the convergence-rate case: rates are "
            "measured by the MMS study on subdivided Kuhn-box families, because an "
            "fTetWild family is an uncontrolled abscissa (this one spans 2.0x in h with "
            "two levels 6.6% apart -- a 15.6x noise amplifier on that pair's order) and "
            "a curved domain freezes geometry error under subdivision"
        ),
        "geometry": (
            "hollow cylinder, inner radius 5, outer radius 10, height 10, axis +Z spanning "
            "z in [0, 10]. Driven by prescribed radial displacement on the bore (there is no "
            "traction API); constrain z on both end faces for plane strain"
        ),
        "meshing": "linear tets, resolved enough for percent-level stress profiles "
        "(the committed h2 is what the tests read)",
        "export": _GMSH_EXPORT,
    },
    "cantilever_slender": {
        "purpose": (
            "cantilever tip deflection under self-weight, compared against Timoshenko beam "
            "theory"
        ),
        "geometry": (
            "rectangular prism with slenderness L/h >= 20, long axis along +X, one end face "
            "at x = 0 (that face is the clamp, enclosed by selector_clamp.stl). The case "
            "must also satisfy BeamReference.violations: small deflection and nu below 0.40"
        ),
        "meshing": "linear tets, >= 4 elements through the section, -e 1e-4 for the same "
        "slender-part reason as bar_axial",
        "export": _GMSH_EXPORT,
    },
}