CelerisLab/src/CelerisLab/body/jets.py

# CelerisLab/body/jets.py
"""Body-independent jets: spatial region + direction field + profile + amplitude.

A Jet selects cut links purely by the true wall hit point (FP64, before SoA
packing), over every link in the scene; it carries no body association.
Per-link velocity = amplitude * profile(s) * field(hit), where ``s`` in [-1, 1]
is the region's lateral coordinate. Velocities are lattice units of the level
whose links are evaluated; callers convert frames/units before evaluation.
One physical link (receiver, direction) may be driven by at most one jet, and
one jet may drive the links of only one body (checked at assignment, any amplitude).

``condition`` names the prescribed boundary type. Only ``"velocity"`` (the
default: amplitude is a signed lattice velocity, not a flow rate) is
implemented; ``"pressure"`` is reserved and rejected at construction, before
any amplitude validation, so it never reaches assignment or a kernel.

Edge rule: the lateral bound |s| < 1 is exclusive (the profile vanishes
there for parabolic/quartic); Sector r_min/r_max and Tube 0 <= along <= length
are inclusive. Jets are immutable; an amplitude change replaces the Jet.

``surface_link_velocity`` is the single host-only assignment/validation point
for legacy apertures plus jets (no PyCUDA import), shared by dense and fine
levels.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, replace

import numpy as np

CONDITIONS = ("velocity",)
RESERVED_CONDITIONS = ("pressure",)

PROFILES = {
    "uniform": lambda s: np.ones_like(s),
    "parabolic": lambda s: 1 - s*s,
    "quartic": lambda s: (1 - s*s)**2,
}


def _finite(*values):
    if not all(math.isfinite(v) for v in values):
        raise ValueError("Jet parameters must be finite")


def _freeze_origin(spec):
    """Store ``origin`` as a validated float 2-tuple (no caller-owned list/array)."""
    origin = tuple(float(v) for v in spec.origin)
    if len(origin) != 2:
        raise ValueError("Jet origin must be a 2D point")
    _finite(*origin)
    object.__setattr__(spec, "origin", origin)


@dataclass(frozen=True)
class Sector:
    """Annular sector about ``origin``: |angle offset| < half_width, r_min <= r <= r_max."""
    origin: tuple
    angle: float
    half_width: float
    r_max: float
    r_min: float = 0.0

    def __post_init__(self):
        _freeze_origin(self)
        _finite(self.angle, self.half_width, self.r_max, self.r_min)
        if not (0 < self.half_width <= math.pi and 0 <= self.r_min <= self.r_max):
            raise ValueError("Sector requires 0 < half_width <= pi and 0 <= r_min <= r_max")

    def select(self, hits):
        d = hits - np.asarray(self.origin, np.float64)
        r = np.hypot(d[:, 0], d[:, 1])
        delta = np.angle(np.exp(1j*(np.arctan2(d[:, 1], d[:, 0]) - self.angle)))
        return (np.abs(delta) < self.half_width) & (r >= self.r_min) & (r <= self.r_max), \
            delta/self.half_width


@dataclass(frozen=True)
class Tube:
    """Rectangle from ``origin`` along ``axis_angle`` for ``length``, |lateral| < half_width."""
    origin: tuple
    axis_angle: float
    length: float
    half_width: float

    def __post_init__(self):
        _freeze_origin(self)
        _finite(self.axis_angle, self.length, self.half_width)
        if self.length <= 0 or self.half_width <= 0:
            raise ValueError("Tube requires positive length and half_width")

    def select(self, hits):
        d = hits - np.asarray(self.origin, np.float64)
        c, s = math.cos(self.axis_angle), math.sin(self.axis_angle)
        along, lateral = d[:, 0]*c + d[:, 1]*s, -d[:, 0]*s + d[:, 1]*c
        return (along >= 0) & (along <= self.length) & (np.abs(lateral) < self.half_width), \
            lateral/self.half_width


@dataclass(frozen=True)
class Radial:
    """Unit vector along hit - origin (outward) or its reverse (inward)."""
    origin: tuple
    inward: bool = True

    def __post_init__(self):
        _freeze_origin(self)

    def direction(self, hits):
        d = hits - np.asarray(self.origin, np.float64)
        r = np.hypot(d[:, 0], d[:, 1])[:, None]
        if np.any(r <= 0):
            raise ValueError("Radial field undefined at its origin")
        return (-d if self.inward else d)/r


@dataclass(frozen=True)
class Parallel:
    """Uniform unit vector at ``angle`` (radians, counterclockwise from +x)."""
    angle: float

    def __post_init__(self):
        _finite(self.angle)

    def direction(self, hits):
        return np.broadcast_to([math.cos(self.angle), math.sin(self.angle)], hits.shape)


@dataclass(frozen=True)
class Jet:
    """Region (Sector|Tube), field (Radial|Parallel), profile name, signed amplitude,
    prescribed boundary condition (velocity only; pressure reserved).

    Frozen: change amplitude with ``with_amplitude`` (validated, new object);
    the cached per-link shapes stay valid because region/field/profile cannot change.
    """
    region: object
    field: object
    profile: str = "uniform"
    amplitude: float = 0.0
    condition: str = "velocity"

    def __post_init__(self):
        # The condition decides what amplitude and field mean: check it first.
        if self.condition in RESERVED_CONDITIONS:
            raise NotImplementedError(f"{self.condition} jets are reserved but not implemented")
        if self.condition not in CONDITIONS:
            raise ValueError(f"Jet condition must be one of {CONDITIONS + RESERVED_CONDITIONS}")
        if self.profile not in PROFILES:
            raise ValueError(f"Jet profile must be one of {sorted(PROFILES)}")
        _finite(self.amplitude)
        if abs(self.amplitude) >= 1:
            raise ValueError("Jet amplitude must satisfy abs(amplitude) < 1 (lattice units)")

    def with_amplitude(self, amplitude):
        return replace(self, amplitude=float(amplitude))


def jet_link_shapes(jets, hits):
    """Unit-amplitude per-link vectors (n, 2) FP64 and jet index (n,) int32 (-1 = none).

    ``hits`` are true FP64 wall hit points (n, 2) row-aligned with the packed
    links. Raises if a jet selects nothing or two jets select the same link.
    """
    hits = np.asarray(hits, np.float64).reshape(-1, 2)
    if not np.isfinite(hits).all():
        raise ValueError("Jet hit points must be finite")
    shape = np.zeros_like(hits)
    jet_id = np.full(len(hits), -1, np.int32)
    for j, jet in enumerate(jets):
        selected, s = jet.region.select(hits)
        if not selected.any():
            raise ValueError(f"Jet {j} selects no resolved cut links")
        if np.any(jet_id[selected] >= 0):
            raise ValueError(f"Jet {j} overlaps links already driven by another jet")
        jet_id[selected] = j
        shape[selected] = (PROFILES[jet.profile](s[selected])[:, None]
                           * jet.field.direction(hits[selected]))
    return shape, jet_id


def jet_link_velocity(jets, shape, jet_id):
    """FP32 (n, 2) velocity = amplitude[jet_id] * shape; O(n), no geometry work."""
    amplitude = np.array([jet.amplitude for jet in jets] + [0.0], np.float64)
    return (amplitude[jet_id][:, None]*shape).astype(np.float32)


# Two coarse-unit frames coexist (see FixedPatchStepper docs): the patch
# cell-edge frame puts coarse cell I's center at I + 1/2, while Simulation
# geometry uses the integer-node frame with that center at I. Jets are
# evaluated in the frame their region/field were written in; convert hits
# explicitly, never by relabelling.

def fine_hits_to_edge_frame(hits_local, fine_origin, dx):
    """Fine-local cell-index hits -> coarse cell-edge frame: origin + dx*(hit + 1/2)."""
    return np.asarray(fine_origin, np.float64) + dx*(np.asarray(hits_local, np.float64) + .5)


def edge_to_node_frame(points):
    """Coarse cell-edge frame -> Simulation integer-node frame (shift by -1/2)."""
    return np.asarray(points, np.float64) - .5


def surface_link_velocity(objects, cfg, flags, fluid_idx, dir, body_id, rx, ry,
                          jets=(), hits=None, *, _shifted_blocked=False):
    """Per-link prescribed surface velocity (n, 2) FP32 and aperture mask.

    Host-only and storage-neutral: rows follow the given physical cut-link
    order. Returns (None, None) when neither an object carries surface
    actuation nor a jet is given. Jets need true FP64 ``hits`` (n, 2) and may
    not share a link with each other or with a legacy aperture.
    Shifted storage is privately enabled only by the all-blocked stepper.
    """
    enabled = [obj for obj in objects if obj.surface_actuation is not None]
    if not enabled and not jets:
        return None, None
    from ..lbm.descriptors import FLUID, SOLID
    if np.any((flags[fluid_idx] & FLUID) == 0) or np.any(flags[fluid_idx] & SOLID):
        raise ValueError("Surface cut links must have active fluid receivers")
    if len(np.unique(fluid_idx.astype(np.uint64)*cfg.nq+dir)) != len(dir):
        raise ValueError("Surface boundary requires unique cut-link writers")
    surface_velocity = np.zeros((len(dir), 2), np.float32)
    surface_aperture = np.zeros(len(dir), bool)
    for obj in enabled:
        obj.validate_surface_actuation(cfg)
        selected = body_id == obj.obj_id
        velocity, aperture = obj.surface_actuation.project(rx[selected], ry[selected])
        if not aperture.any():
            raise ValueError("Surface aperture hits no resolved cut links")
        surface_velocity[selected] = velocity
        surface_aperture[selected] = aperture
    if jets:
        if hits is None or len(hits) != len(dir):
            raise ValueError("Jets require true hit points row-aligned with the cut links")
        if (cfg is None or not cfg.is_d2q9 or cfg.store_precision != "FP32"
                or cfg.streaming != "double_buffer"
                or (cfg.ddf_shifting and not (_shifted_blocked is True
                    and cfg.collision == "SRT" and not cfg.les_enabled))):
            raise ValueError("Jets require D2Q9/FP32/unshifted/double_buffer")
        shape, jet_id = jet_link_shapes(jets, hits)
        driven = jet_id >= 0
        for j in range(len(jets)):  # scene-wide body identity, independent of storage blocks
            bodies = np.unique(body_id[jet_id == j])
            if len(bodies) > 1:
                raise ValueError(f"Jet {j} selects cut links of bodies {bodies.tolist()}; "
                                 "one jet may drive only one body")
        owners = {obj.obj_id: obj for obj in objects}
        for bid in np.unique(body_id[driven]):
            state = vars(owners[int(bid)].state)
            if any(state.get(k, 0.) != 0 for k in ("vx", "vy", "omega")):
                raise ValueError("Jets may only drive links of stationary bodies")
        if np.any(surface_aperture & driven):
            raise ValueError("Jet overlaps a SurfaceActuation aperture")
        surface_velocity[driven] = jet_link_velocity(jets, shape, jet_id)[driven]
        surface_aperture |= driven
    return surface_velocity, surface_aperture