gri-plot 0.2.0.post1__py3-none-any.whl

gri_plot/surfaces/mesh.py

@@ -0,0 +1,603 @@
+"""Core mesh utilities for 3D surface visualization.
+
+This module provides the fundamental functions for creating Plotly Mesh3d traces
+from vertices and faces, and for extracting mesh geometry from scalar fields.
+"""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import plotly.graph_objects as go
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from numpy.typing import NDArray
+
+# Default mesh resolution for field_to_mesh and related functions
+DEFAULT_MESH_RESOLUTION = 100
+
+
+def vertices_to_mesh3d(  # noqa: PLR0913 - mesh rendering needs many options
+    vertices: NDArray[np.floating],
+    faces: NDArray[np.integer],
+    *,
+    intensity_fn: Callable[[NDArray[np.floating]], NDArray[np.floating]] | None = None,
+    color: str | None = None,
+    opacity: float = 0.7,
+    name: str | None = None,
+    colorscale: str | list = "Viridis",
+    colorbar_title: str | None = None,
+    showscale: bool = False,
+    **kwargs,
+) -> go.Mesh3d:
+    """Create a Plotly Mesh3d trace from vertices and faces.
+
+    This is the core rendering function. All shape and surface classes should
+    use this to generate their Plotly traces.
+
+    Args:
+        vertices: Vertex positions, shape (N, 3).
+        faces: Triangle indices, shape (M, 3).
+        intensity_fn: Optional function to compute vertex intensities for coloring.
+            Takes vertices (N, 3) and returns intensities (N,). If None and color
+            is None, defaults to distance from mesh centroid (works well in any
+            coordinate frame).
+        color: Uniform surface color. Mutually exclusive with intensity_fn.
+        opacity: Surface opacity (0-1).
+        name: Trace name for legend.
+        colorscale: Plotly colorscale name or custom colorscale list
+            (used with intensity_fn).
+        colorbar_title: Title for colorbar (used with intensity_fn).
+        showscale: Whether to show the colorbar.
+        **kwargs: Additional arguments for go.Mesh3d.
+
+    Returns:
+        Plotly Mesh3d trace.
+    """
+    vertices = np.asarray(vertices)
+    faces = np.asarray(faces)
+
+    trace_kwargs: dict = {
+        "x": vertices[:, 0],
+        "y": vertices[:, 1],
+        "z": vertices[:, 2],
+        "i": faces[:, 0],
+        "j": faces[:, 1],
+        "k": faces[:, 2],
+        "opacity": opacity,
+    }
+
+    if color is not None:
+        # Uniform color mode
+        trace_kwargs["color"] = color
+    else:
+        # Intensity-based coloring
+        if intensity_fn is not None:
+            intensity = intensity_fn(vertices)
+        else:
+            # Default to distance from mesh centroid - works in any coordinate frame
+            centroid = vertices.mean(axis=0)
+            intensity = np.linalg.norm(vertices - centroid, axis=-1)
+
+        trace_kwargs["intensity"] = intensity
+        trace_kwargs["colorscale"] = colorscale
+        trace_kwargs["showscale"] = showscale
+        if colorbar_title is not None:
+            trace_kwargs["colorbar"] = {"title": colorbar_title}
+
+    if name is not None:
+        trace_kwargs["name"] = name
+        trace_kwargs["showlegend"] = True
+
+    trace_kwargs.update(kwargs)
+    return go.Mesh3d(**trace_kwargs)
+
+
+def _remove_degenerate_faces(
+    vertices: NDArray[np.floating],
+    faces: NDArray[np.integer],
+    min_area_ratio: float = 0.001,
+    max_dihedral_angle: float = 1.0,  # ~57 degrees
+) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+    """Remove degenerate triangles from mesh.
+
+    Degenerate faces can appear after projecting vertices onto an isosurface,
+    especially for thin tube-like surfaces where opposite sides can fold
+    onto each other.
+
+    Args:
+        vertices: Vertex positions, shape (N, 3).
+        faces: Triangle indices, shape (M, 3).
+        min_area_ratio: Remove faces with area less than this fraction of
+            median face area.
+        max_dihedral_angle: Remove faces that have an edge with dihedral
+            angle exceeding this value (radians). Default ~57 degrees.
+
+    Returns:
+        Tuple of (vertices, filtered_faces). Vertices are unchanged.
+    """
+    if len(faces) == 0:
+        return vertices, faces
+
+    # Compute face areas
+    v0 = vertices[faces[:, 0]]
+    v1 = vertices[faces[:, 1]]
+    v2 = vertices[faces[:, 2]]
+    edge1 = v1 - v0
+    edge2 = v2 - v0
+    areas = 0.5 * np.linalg.norm(np.cross(edge1, edge2), axis=1)
+
+    # Find faces to keep based on area
+    median_area = np.median(areas)
+    min_area = median_area * min_area_ratio
+    area_ok = areas >= min_area
+
+    # Find edges with extreme dihedral angles
+    dihedral_angles = compute_edge_dihedral_angles(vertices, faces)
+    bad_edges = {
+        edge for edge, angle in dihedral_angles.items() if angle > max_dihedral_angle
+    }
+
+    # Mark faces that contain bad edges
+    angle_ok = np.ones(len(faces), dtype=bool)
+    if bad_edges:
+        for face_idx, face in enumerate(faces):
+            for i in range(3):
+                v_a, v_b = int(face[i]), int(face[(i + 1) % 3])
+                edge = (min(v_a, v_b), max(v_a, v_b))
+                if edge in bad_edges:
+                    angle_ok[face_idx] = False
+                    break
+
+    # Keep faces that pass both criteria
+    keep_mask = area_ok & angle_ok
+    filtered_faces = faces[keep_mask]
+
+    return vertices, filtered_faces
+
+
+def field_to_mesh(  # noqa: PLR0913 - marching cubes has many parameters
+    residual_fn: Callable[[NDArray[np.floating]], NDArray[np.floating]],
+    bounds: tuple[NDArray[np.floating], NDArray[np.floating]],
+    resolution: int = DEFAULT_MESH_RESOLUTION,
+    isovalue: float = 0.0,
+    *,
+    refine: bool = True,
+    refine_angle_threshold: float = 0.15,
+    refine_iterations: int = 5,
+) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+    """Extract mesh geometry from a scalar field using marching cubes.
+
+    The surface is extracted where residual_fn(xyz) = isovalue.
+
+    Args:
+        residual_fn: Function (xyz) -> scalar where xyz has shape (..., 3)
+            and output has shape (...).
+        bounds: Tuple of (min_corner, max_corner), each shape (3,).
+        resolution: Number of samples per dimension. Default 100.
+        isovalue: Value at which to extract the surface.
+        refine: If True (default), apply adaptive mesh refinement to subdivide
+            high-curvature regions after marching cubes.
+        refine_angle_threshold: Dihedral angle threshold in radians for
+            refinement. Edges exceeding this angle are subdivided.
+            Default 0.15 (~8.5 degrees).
+        refine_iterations: Number of refinement passes. Default 5.
+
+    Returns:
+        Tuple of (vertices, faces) where:
+            - vertices: shape (N, 3) in world coordinates
+            - faces: shape (M, 3) triangle indices
+
+    Raises:
+        ImportError: If scikit-image is not available.
+        RuntimeError: If marching cubes fails or finds no surface.
+    """
+    try:
+        from skimage.measure import marching_cubes  # noqa: PLC0415
+    except ImportError as e:
+        raise ImportError(
+            "scikit-image is required for field_to_mesh. "
+            "Install with: pip install scikit-image",
+        ) from e
+
+    min_corner = np.asarray(bounds[0], dtype=np.float64)
+    max_corner = np.asarray(bounds[1], dtype=np.float64)
+
+    # Create 3D grid
+    x = np.linspace(min_corner[0], max_corner[0], resolution)
+    y = np.linspace(min_corner[1], max_corner[1], resolution)
+    z = np.linspace(min_corner[2], max_corner[2], resolution)
+
+    spacing = (
+        (max_corner[0] - min_corner[0]) / (resolution - 1),
+        (max_corner[1] - min_corner[1]) / (resolution - 1),
+        (max_corner[2] - min_corner[2]) / (resolution - 1),
+    )
+
+    # Create meshgrid and evaluate field
+    x_grid, y_grid, z_grid = np.meshgrid(x, y, z, indexing="ij")
+    xyz = np.stack([x_grid, y_grid, z_grid], axis=-1)
+    values = residual_fn(xyz)
+
+    # Extract surface
+    try:
+        verts, faces, _, _ = marching_cubes(values, level=isovalue, spacing=spacing)
+    except Exception as e:
+        raise RuntimeError(f"Marching cubes failed: {e}") from e
+
+    if len(verts) == 0:
+        raise RuntimeError("No surface found at specified isovalue")
+
+    # Convert to world coordinates
+    verts_world = verts + min_corner
+
+    # Project marching cubes vertices onto the true isosurface
+    # Marching cubes uses linear interpolation which leaves vertices off-surface
+    verts_world = _project_to_surface(verts_world, residual_fn)
+
+    # Apply adaptive refinement if requested
+    if refine:
+        verts_world, faces = refine_mesh(
+            verts_world,
+            faces,
+            residual_fn=residual_fn,
+            angle_threshold=refine_angle_threshold,
+            max_iterations=refine_iterations,
+        )
+
+    return verts_world, faces
+
+
+def grid_to_mesh(
+    x: NDArray[np.floating],
+    y: NDArray[np.floating],
+    z: NDArray[np.floating],
+) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+    """Convert a 2D grid of 3D points to vertices and triangular faces.
+
+    Used for parametric surfaces where x, y, z are 2D arrays defining
+    a surface patch.
+
+    Args:
+        x: X coordinates, 2D array of shape (M, N).
+        y: Y coordinates, 2D array of shape (M, N).
+        z: Z coordinates, 2D array of shape (M, N).
+
+    Returns:
+        Tuple of (vertices, faces) where:
+            - vertices: shape (M*N, 3)
+            - faces: shape ((M-1)*(N-1)*2, 3)
+    """
+    m, n = x.shape
+    vertices = np.column_stack([x.ravel(), y.ravel(), z.ravel()])
+
+    # Create triangular faces from the grid
+    faces = []
+    for i in range(m - 1):
+        for j in range(n - 1):
+            # Indices of the four corners of this grid cell
+            idx00 = i * n + j
+            idx01 = i * n + (j + 1)
+            idx10 = (i + 1) * n + j
+            idx11 = (i + 1) * n + (j + 1)
+
+            # Two triangles per cell
+            faces.append([idx00, idx10, idx11])
+            faces.append([idx00, idx11, idx01])
+
+    return vertices, np.array(faces, dtype=np.int64)
+
+
+def compute_face_normals(
+    vertices: NDArray[np.floating],
+    faces: NDArray[np.integer],
+) -> NDArray[np.floating]:
+    """Compute unit normal vectors for each face.
+
+    Args:
+        vertices: Vertex positions, shape (N, 3).
+        faces: Triangle indices, shape (M, 3).
+
+    Returns:
+        Unit normals, shape (M, 3).
+    """
+    v0 = vertices[faces[:, 0]]
+    v1 = vertices[faces[:, 1]]
+    v2 = vertices[faces[:, 2]]
+
+    # Cross product of two edges
+    edge1 = v1 - v0
+    edge2 = v2 - v0
+    normals = np.cross(edge1, edge2)
+
+    # Normalize
+    lengths = np.linalg.norm(normals, axis=1, keepdims=True)
+    lengths = np.maximum(lengths, 1e-12)  # Avoid division by zero
+    return normals / lengths
+
+
+def compute_edge_dihedral_angles(
+    vertices: NDArray[np.floating],
+    faces: NDArray[np.integer],
+) -> dict[tuple[int, int], float]:
+    """Compute dihedral angles at each edge.
+
+    The dihedral angle is the angle between normals of adjacent faces.
+    Edges with high dihedral angles indicate high curvature.
+
+    Args:
+        vertices: Vertex positions, shape (N, 3).
+        faces: Triangle indices, shape (M, 3).
+
+    Returns:
+        Dictionary mapping (min_idx, max_idx) edge tuples to dihedral angles
+        in radians. Boundary edges (single adjacent face) are not included.
+    """
+    normals = compute_face_normals(vertices, faces)
+
+    # Build edge to face adjacency
+    edge_to_faces: dict[tuple[int, int], list[int]] = {}
+    for face_idx, face in enumerate(faces):
+        for i in range(3):
+            v0, v1 = int(face[i]), int(face[(i + 1) % 3])
+            edge = (min(v0, v1), max(v0, v1))
+            if edge not in edge_to_faces:
+                edge_to_faces[edge] = []
+            edge_to_faces[edge].append(face_idx)
+
+    # Compute dihedral angles for interior edges
+    dihedral_angles: dict[tuple[int, int], float] = {}
+    for edge, face_indices in edge_to_faces.items():
+        if len(face_indices) == 2:  # noqa: PLR2004 - interior edge has exactly 2 faces
+            n1 = normals[face_indices[0]]
+            n2 = normals[face_indices[1]]
+            # Angle between normals
+            dot = np.clip(np.dot(n1, n2), -1.0, 1.0)
+            angle = np.arccos(dot)
+            dihedral_angles[edge] = float(angle)
+
+    return dihedral_angles
+
+
+def _project_to_surface(
+    points: NDArray[np.floating],
+    residual_fn: Callable[[NDArray[np.floating]], NDArray[np.floating]],
+    max_iterations: int = 20,
+    tolerance: float = 1e-8,
+) -> NDArray[np.floating]:
+    """Project points onto the isosurface using Newton-Raphson along gradient.
+
+    Uses a scale-invariant approach: step = residual / |grad|^2 * grad,
+    which gives correct step sizes regardless of residual normalization.
+
+    Args:
+        points: Points to project, shape (N, 3).
+        residual_fn: Implicit surface function.
+        max_iterations: Maximum Newton iterations.
+        tolerance: Convergence tolerance for residual.
+
+    Returns:
+        Projected points, shape (N, 3).
+    """
+    points = np.asarray(points, dtype=np.float64).copy()
+    h = 1e-6  # Finite difference step
+
+    for _ in range(max_iterations):
+        residuals = residual_fn(points)
+
+        # Check convergence
+        if np.all(np.abs(residuals) < tolerance):
+            break
+
+        # Compute gradient via finite differences
+        grad = np.zeros_like(points)
+        for dim in range(3):
+            points_plus = points.copy()
+            points_plus[:, dim] += h
+            grad[:, dim] = (residual_fn(points_plus) - residuals) / h
+
+        # Newton-Raphson step: move by residual / |grad|^2 * grad
+        # This is scale-invariant: if residual is normalized by k, grad is also
+        # scaled by 1/k, so residual / |grad|^2 has units of length
+        grad_sq = np.sum(grad * grad, axis=1, keepdims=True)
+        grad_sq = np.maximum(grad_sq, 1e-20)
+
+        # Step along gradient direction by residual / |grad|^2
+        points -= (residuals[:, np.newaxis] / grad_sq) * grad
+
+    return points
+
+
+def _subdivide_face(
+    face: NDArray[np.integer],
+    edge_to_new_vertex: dict[tuple[int, int], int],
+) -> list[list[int]]:
+    """Subdivide a single face based on available edge midpoints.
+
+    Args:
+        face: Triangle vertex indices, shape (3,).
+        edge_to_new_vertex: Mapping from edge tuples to midpoint vertex indices.
+
+    Returns:
+        List of new face vertex index lists.
+    """
+    v0, v1, v2 = int(face[0]), int(face[1]), int(face[2])
+    edges = [
+        (min(v0, v1), max(v0, v1)),
+        (min(v1, v2), max(v1, v2)),
+        (min(v2, v0), max(v2, v0)),
+    ]
+    midpoints = [edge_to_new_vertex.get(e) for e in edges]
+    split_count = sum(m is not None for m in midpoints)
+
+    if split_count == 3:  # noqa: PLR2004 - all 3 edges of triangle split
+        # All three edges split: 1-to-4 subdivision
+        m01, m12, m20 = midpoints[0], midpoints[1], midpoints[2]
+        if m01 is None or m12 is None or m20 is None:
+            return [[v0, v1, v2]]  # Should not happen when split_count == 3
+        return [[v0, m01, m20], [v1, m12, m01], [v2, m20, m12], [m01, m12, m20]]
+
+    if split_count == 2:  # noqa: PLR2004 - 2 of 3 edges split
+        return _subdivide_face_two_edges(v0, v1, v2, midpoints)
+
+    if split_count == 1:
+        return _subdivide_face_one_edge(v0, v1, v2, midpoints)
+
+    return [[v0, v1, v2]]
+
+
+def _subdivide_face_two_edges(
+    v0: int,
+    v1: int,
+    v2: int,
+    midpoints: list[int | None],
+) -> list[list[int]]:
+    """Handle 1-to-3 subdivision when two edges are split."""
+    if midpoints[0] is not None and midpoints[1] is not None:
+        m01, m12 = midpoints[0], midpoints[1]
+        return [[v0, m01, v2], [m01, v1, m12], [m01, m12, v2]]
+    if midpoints[1] is not None and midpoints[2] is not None:
+        m12, m20 = midpoints[1], midpoints[2]
+        return [[v0, v1, m20], [v1, m12, m20], [m12, v2, m20]]
+    # midpoints[0] and midpoints[2] must be set
+    m01, m20 = midpoints[0], midpoints[2]
+    if m01 is None or m20 is None:
+        return [[v0, v1, v2]]  # Should not happen
+    return [[v0, m01, m20], [m01, v1, v2], [m01, v2, m20]]
+
+
+def _subdivide_face_one_edge(
+    v0: int,
+    v1: int,
+    v2: int,
+    midpoints: list[int | None],
+) -> list[list[int]]:
+    """Handle 1-to-2 subdivision when one edge is split."""
+    if midpoints[0] is not None:
+        m01 = midpoints[0]
+        return [[v0, m01, v2], [m01, v1, v2]]
+    if midpoints[1] is not None:
+        m12 = midpoints[1]
+        return [[v0, v1, m12], [v0, m12, v2]]
+    m20 = midpoints[2]
+    if m20 is None:
+        return [[v0, v1, v2]]  # Should not happen
+    return [[v0, v1, m20], [m20, v1, v2]]
+
+
+def _find_faces_to_split(
+    faces: NDArray[np.integer],
+    edges_to_split: set[tuple[int, int]],
+) -> set[int]:
+    """Find face indices that have at least one edge marked for splitting."""
+    faces_to_split = set()
+    for face_idx, face in enumerate(faces):
+        for i in range(3):
+            v0, v1 = int(face[i]), int(face[(i + 1) % 3])
+            edge = (min(v0, v1), max(v0, v1))
+            if edge in edges_to_split:
+                faces_to_split.add(face_idx)
+                break
+    return faces_to_split
+
+
+def _create_midpoint_vertices(
+    vertices: NDArray[np.floating],
+    edges_to_split: set[tuple[int, int]],
+) -> tuple[NDArray[np.floating], dict[tuple[int, int], int]]:
+    """Create new vertices at edge midpoints.
+
+    Returns:
+        Tuple of (new_vertices_array, edge_to_new_vertex_mapping).
+    """
+    edge_to_new_vertex: dict[tuple[int, int], int] = {}
+    new_vertices_list = [vertices]
+    next_vertex_idx = len(vertices)
+
+    for edge in edges_to_split:
+        midpoint = (vertices[edge[0]] + vertices[edge[1]]) / 2
+        edge_to_new_vertex[edge] = next_vertex_idx
+        new_vertices_list.append(midpoint.reshape(1, 3))
+        next_vertex_idx += 1
+
+    return np.vstack(new_vertices_list), edge_to_new_vertex
+
+
+def refine_mesh(
+    vertices: NDArray[np.floating],
+    faces: NDArray[np.integer],
+    residual_fn: Callable[[NDArray[np.floating]], NDArray[np.floating]] | None = None,
+    angle_threshold: float = 0.3,
+    max_iterations: int = 3,
+) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+    """Refine mesh by subdividing high-curvature triangles.
+
+    Triangles with edges that have high dihedral angles (indicating high
+    curvature) are subdivided using a 1-to-4 split. New vertices are placed
+    at edge midpoints and optionally projected onto the implicit surface.
+
+    Args:
+        vertices: Vertex positions, shape (N, 3).
+        faces: Triangle indices, shape (M, 3).
+        residual_fn: Optional implicit surface function for projecting new
+            vertices onto the surface. If None, new vertices stay at midpoints.
+        angle_threshold: Dihedral angle threshold in radians. Edges with
+            angles above this are marked for subdivision. Default 0.3 (~17 deg).
+        max_iterations: Number of refinement passes. Each pass subdivides
+            triangles that exceed the threshold.
+
+    Returns:
+        Tuple of (refined_vertices, refined_faces).
+    """
+    vertices = np.asarray(vertices, dtype=np.float64)
+    faces = np.asarray(faces, dtype=np.int64)
+
+    for _ in range(max_iterations):
+        # Compute dihedral angles and find edges to subdivide
+        dihedral_angles = compute_edge_dihedral_angles(vertices, faces)
+        edges_to_split = {
+            edge for edge, angle in dihedral_angles.items() if angle > angle_threshold
+        }
+
+        if not edges_to_split:
+            break
+
+        faces_to_split = _find_faces_to_split(faces, edges_to_split)
+        if not faces_to_split:
+            break
+
+        # Create midpoint vertices
+        vertices, edge_to_new_vertex = _create_midpoint_vertices(
+            vertices,
+            edges_to_split,
+        )
+
+        # Project new vertices onto surface if residual_fn provided
+        if residual_fn is not None and edge_to_new_vertex:
+            new_vertex_indices = list(edge_to_new_vertex.values())
+            projected = _project_to_surface(vertices[new_vertex_indices], residual_fn)
+            vertices[new_vertex_indices] = projected
+
+        # Build new face list
+        new_faces = []
+        for face_idx, face in enumerate(faces):
+            if face_idx in faces_to_split:
+                new_faces.extend(_subdivide_face(face, edge_to_new_vertex))
+            else:
+                new_faces.append(face.tolist())
+
+        faces = np.array(new_faces, dtype=np.int64)
+
+    return vertices, faces
+
+
+__all__ = [
+    "DEFAULT_MESH_RESOLUTION",
+    "compute_edge_dihedral_angles",
+    "compute_face_normals",
+    "field_to_mesh",
+    "grid_to_mesh",
+    "refine_mesh",
+    "vertices_to_mesh3d",
+]