Source code for nekmeshpy.hexmesh.morph

"""Fixed-arity, rung-preserving ``HexMesh`` operations (delta 0)."""

from __future__ import annotations

import logging
from collections.abc import Sequence

import numpy as np

from .._typing import (
    FloatArray,
    IntArray,
    Point,
    PointArray,
    Vec3,
)
from ..core import affine, conform
from ..pointmesh.morph import NodeMap, _mapped
from ..quadmesh import QuadMesh
from ..quadmesh.morph import _affine as quad_affine
from ..quadmesh.morph import _transform_fn as quad_transform_fn
from ..quadmesh.morph import blend as quad_blend
from .hexmesh import HexMesh

_log = logging.getLogger(__name__)




def blend(a: HexMesh, b: HexMesh,
          fractions: FloatArray | Sequence[float]) -> list[HexMesh]:
    """Linearly morph between two conformal blocks ``a`` and ``b`` (identical ``hexes``,
    equal point count), one block per fraction ``t`` with points ``(1-t)*a + t*b`` --
    ``t=0`` reproduces ``a``, ``t=1`` reproduces ``b``.  Each block returned carries
    ``a``'s ``face_tags`` and empty ``element_tags``."""
    A: PointArray = np.asarray(a.points, dtype=float).reshape(-1, 3)
    B: PointArray = np.asarray(b.points, dtype=float).reshape(-1, 3)
    if A.shape[0] != B.shape[0]:
        raise ValueError(
            "blend: blocks must have equal point counts (got %d, %d); build one "
            "from the other's points so they pair by index"
            % (A.shape[0], B.shape[0]))
    if not np.array_equal(a.corners, b.corners):
        raise ValueError(
            "blend: blocks must share identical connectivity (paired by index)")
    if a.order != b.order:
        raise ValueError("blend: blocks must have the same order (got %d, %d)"
                         % (a.order, b.order))
    # identical connectivity => identical edge / face tables, so a's and b's shared
    # edge nodes, shared face nodes and private interiors already pair one-for-one
    # and each morphs with the same lerp the corners get from the blended points.
    # The shared corners, shared edge nodes and shared face nodes *are* the
    # shared-face QuadMesh, so that whole part of the morph is one
    # ``QuadMesh.blend`` of the rung below (which in turn delegates the corners and
    # edge nodes to ``LineMesh.blend``); the result reuses ``a``'s per-hex face
    # indices and D4 codes verbatim -- a blend is a pure point-space morph, so
    # nothing is re-derived.  At order 1 all three tables are empty and this is
    # exactly the plain point blend.
    ho = a.order > 1
    ai, bi = a.interior, b.interior
    fr: FloatArray = np.asarray(fractions, dtype=float).ravel()
    # the face tags ride the blended ``quads`` themselves -- ``QuadMesh.blend`` keeps
    # ``a``'s edge tags but drops its element tags, so they are put back here
    return [HexMesh(QuadMesh(faces.line_mesh, faces.quads, faces.orient, faces.interior,
                             a.quad_mesh.element_tags),
                    a.hexes, a.orient,
                    (1.0 - t) * ai + t * bi if ho else None)
            for t, faces in zip(fr, quad_blend(a.quad_mesh, b.quad_mesh, fr))]



def _affine(mesh: HexMesh, matrix: FloatArray | None, offset: Vec3) -> HexMesh:
    """Map every coordinate of ``mesh`` through the affine pair ``(matrix, offset)``."""
    return HexMesh(quad_affine(mesh.quad_mesh, matrix, offset), mesh.hexes,
                   mesh.orient, affine.apply(mesh.interior, matrix, offset),
                   mesh.element_tags)




def transform(mesh: HexMesh, matrix: FloatArray,
              offset: Vec3 | Sequence[float] = affine.ORIGIN) -> HexMesh:
    """A new block with every node mapped through the affine ``p @ matrix.T + offset``.
    """
    return _affine(mesh, np.asarray(matrix, dtype=float).reshape(3, 3),
                   np.asarray(offset, dtype=float).reshape(3))





def transform_fn(mesh: HexMesh, fn: NodeMap) -> HexMesh:
    """A new block with every node -- the shared-face :class:`QuadMesh
    <nekmeshpy.quadmesh.QuadMesh>` all the way down, and the private per-hex
    ``interior`` -- mapped through an arbitrary ``fn``; :func:`transform` for a warp
    that is not affine.  ``fn`` is called once per node block with an ``(N, 3)`` array
    and must return ``(N, 3)``.  Connectivity and tags ride through verbatim.

    ``fn`` can fold space; the result is checked and a ``logging.warning`` (not an
    error) is emitted when any hex comes out inverted -- sampled on the stored map at
    ``mesh.order``, a flag rather than a certificate, and no substitute for
    :func:`quality_summary <nekmeshpy.hexmesh.query.quality_summary>` at the solver's
    order.  It does **not** re-wind a reflection; use :func:`mirror` for that."""
    from .query import quality_summary

    out = HexMesh(quad_transform_fn(mesh.quad_mesh, fn), mesh.hexes, mesh.orient,
                  _mapped(mesh.interior, fn), mesh.element_tags)
    n = quality_summary(out).n_inverted
    if n:
        _log.warning(
            "hexmesh.transform_fn: fn produced %d inverted hex(es) (scaled Jacobian "
            "sampled at order %d); the warp folds the mesh there", n, out.order)
    return out





def translate(mesh: HexMesh, vector: Vec3 | Sequence[float]) -> HexMesh:
    """A new block shifted rigidly by ``vector`` ``(3,)``.  Bit-exact: the offset is
    added without a matmul, so translating by ``0`` returns the identical
    coordinates."""
    return _affine(mesh, *affine.translation(vector))





def rotate(mesh: HexMesh, angle: float,
           axis: Vec3 | Sequence[float] = affine.Z_AXIS,
           center: Point | Sequence[float] = affine.ORIGIN) -> HexMesh:
    """A new block rotated by ``angle`` **radians** about the line through ``center``
    along ``axis`` (right-handed, ``axis`` need not be normalized)."""
    return _affine(mesh, *affine.rotation(angle, axis, center))





def scale(mesh: HexMesh, factor: float | Vec3 | Sequence[float],
          center: Point | Sequence[float] = affine.ORIGIN) -> HexMesh:
    """A new block scaled about ``center`` by ``factor`` -- a scalar (uniform) or a
    ``(3,)`` per-axis vector.  Every factor must be positive."""
    return _affine(mesh, *affine.scaling(factor, center))


#: Nek corner order under a re-winding: swap the bottom and top quads.  That is the
#: reference-space reflection ``t -> 1 - t``, and the cheapest determinant flip that
#: keeps all six faces the same corner sets they already were.
_REWIND_CORNERS: IntArray = np.array([4, 5, 6, 7, 0, 1, 2, 3], dtype=np.int64)

#: The face columns that follow from :data:`_REWIND_CORNERS`: faces 1-4 keep their
#: corner sets (their frames turn over in place) and faces 5 / 6 trade places.
_REWIND_FACES: IntArray = np.array([0, 1, 2, 3, 5, 4], dtype=np.int64)


def _rewind(mesh: HexMesh) -> HexMesh:
    """The same block with every hex turned inside out -- the corner order permuted by
    :data:`_REWIND_CORNERS`, which negates the Jacobian.

    The shared faces are left exactly as stored: a face row is canonical and belongs to
    both its hexes, so re-winding is entirely a matter of each element's *own* view of
    it.  That view is the ``hex`` face columns plus the ``face_orient`` D4 codes, and
    the codes are refitted against the unchanged canonical rows rather than reasoned
    about -- the same call ``merge`` makes when a weld hands an element a face in a
    frame it did not choose."""
    hexes: IntArray = mesh.corners[:, _REWIND_CORNERS]
    elem_faces: IntArray = mesh.hexes[:, _REWIND_FACES]
    orient: IntArray = conform.face_frame_code(
        hexes[:, conform._LOCAL_FACES], mesh.quad_mesh.corners[elem_faces])
    n = mesh.order - 1
    perm: IntArray = (np.arange(n ** 3, dtype=np.int64).reshape(n, n, n)[::-1].ravel()
                      if n else np.zeros(0, dtype=np.int64))
    return HexMesh(mesh.quad_mesh, elem_faces, orient, mesh.interior[:, perm, :],
                   mesh.element_tags)




def mirror(mesh: HexMesh, normal: Vec3 | Sequence[float],
           point: Point | Sequence[float] = affine.ORIGIN) -> HexMesh:
    """A new block reflected through the plane with ``normal`` through ``point``, with
    every hex **re-wound** so it comes back the right way out.

    A reflection has determinant ``-1`` and inverts every element: a bare ``transform``
    by a reflection matrix hands back a block whose scaled Jacobian is negative
    throughout and whose :func:`volume <nekmeshpy.hexmesh.query.volume>` reads negative
    -- a mesh no solver will take.  ``mirror`` pairs the coordinate map with the
    winding fix, so the result measures exactly as its original did.

    The symmetric-domain idiom -- mesh the half the geometry lets you, then recover the
    whole::

        full = hexmesh.merge([half, hexmesh.mirror(half, normal=(1, 0, 0))])

    ``merge`` welds the two along the symmetry plane they now share, so the seam comes
    out conformal rather than doubled.  The faces that were tagged on the symmetry plane
    are *interior* afterwards and still carry their tag: drop those rows, or leave the
    plane untagged on the half in the first place."""
    return _rewind(_affine(mesh, *affine.reflection(normal, point)))



__all__ = [
    "blend",
    "mirror",
    "rotate",
    "scale",
    "transform",
    "transform_fn",
    "translate",
]