CelerisLab/src/CelerisLab/body/geometry/base.py

# CelerisLab/body/geometry/base.py
"""
Geometry shape interface: uniform intermediate representation for all shapes.

Design:
    Geometry is an abstract base that defines the interface for shape-specific
    cut-link enumeration, sensor cell discovery, and flag mask generation.
    The output is geometry-agnostic CutLink and SensorCell records -- no
    boundary-method-specific fields leak into the geometry layer.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import List, Optional, Tuple

import numpy as np


@dataclass
class CutLink:
    """One cut link: fluid node -> wall -> solid neighbor.

    Geometry-agnostic intermediate record.  Produced by circle/polygon/mesh
    builders, consumed by the SoA packer.  No Bouzidi-specific fields.

    Fields:
        fluid_idx:  linear index of the fluid-side cell.
        dir:        lattice direction towards wall (1..NQ-1).
        q:          fractional distance along link, in (0, 1].
        body_id:    owning body id.
        hit_x:      wall hit-point x (lattice).
        hit_y:      wall hit-point y.
        hit_z:      wall hit-point z (0 for 2D).
        rx:         hit_x - center_x  (torque lever arm).
        ry:         hit_y - center_y.
        rz:         hit_z - center_z (0 for 2D).
        normal_x:   inward-facing surface normal x at hit point (0 for 2D stub).
        normal_y:   inward-facing surface normal y at hit point (0 for 2D stub).
        fallback:   0 = Bouzidi legal; 1 = half-way bounce-back required.
        scheme_tag: reserved for future boundary-scheme selection (0 = Bouzidi).
        motion_tag: reserved for future motion type (0 = stationary, 1 = rotating, 2 = translating).
    """
    fluid_idx: int
    dir: int
    q: float
    body_id: int
    hit_x: float = 0.0
    hit_y: float = 0.0
    hit_z: float = 0.0
    rx: float = 0.0
    ry: float = 0.0
    rz: float = 0.0
    normal_x: float = 0.0
    normal_y: float = 0.0
    fallback: int = 0
    scheme_tag: int = 0
    motion_tag: int = 0


@dataclass
class SensorCell:
    """One sensor sampling cell."""
    idx: int        # linear cell index
    body_id: int    # owning body id


class Geometry(ABC):
    """Abstract geometry shape.

    Subclasses implement shape-specific grid projection (cut links, sensor cells,
    flag masks).  The output is a list of CutLink / SensorCell records --
    geometry-agnostic, boundary-method-agnostic.
    """

    @property
    @abstractmethod
    def center(self) -> Tuple[float, ...]:
        ...

    @abstractmethod
    def build_cut_links(
        self, nx: int, ny: int,
        domain_flags: Optional[np.ndarray] = None,
    ) -> List[CutLink]:
        """Enumerate cut links for this shape.

        Args:
            nx, ny: Grid dimensions.
            domain_flags: Optional uint16 flags array for donor-cell
                fluid/solid classification. None = all donors assumed fluid.

        Returns:
            List of CutLink records (empty if zero links found).
        """
        ...

    @abstractmethod
    def build_sensor_cells(self, nx: int, ny: int) -> List[SensorCell]:
        """Enumerate sensor cells for this shape.

        Args:
            nx, ny: Grid dimensions.

        Returns:
            List of SensorCell records (empty if zero cells found).
        """
        ...

    @abstractmethod
    def build_flag_mask(self, nx: int, ny: int) -> np.ndarray:
        """Return (nx*ny,) uint16 flag mask with bits set for this shape.

        This variant marks obstacle cells (``SOLID|OBSTACLE|BC_CURVED``).

        Args:
            nx, ny: Grid dimensions.

        Returns:
            uint16 array of shape (nx*ny,) with flag bits set for obstacle cells.
        """
        ...

    @abstractmethod
    def build_sensor_flag_mask(self, nx: int, ny: int) -> np.ndarray:
        """Return (nx*ny,) uint16 flag mask for sensor cells.

        This variant marks sensor cells (``FLUID|SENSOR_FLAG``).

        Args:
            nx, ny: Grid dimensions.

        Returns:
            uint16 array of shape (nx*ny,) with flag bits set for sensor cells.
        """
        ...

    def build_force_region_flag_mask(self, nx: int, ny: int) -> np.ndarray:
        """Return (nx*ny,) uint16 flag mask for force-region cells.

        This variant marks force-region cells (``FLUID|FRC_REGION``).
        Default implementation calls ``build_sensor_flag_mask`` for
        subclasses that treat sensor and force-region identically.

        Args:
            nx, ny: Grid dimensions.

        Returns:
            uint16 array of shape (nx*ny,) with flag bits set for
            force-region cells.
        """
        return self.build_sensor_flag_mask(nx, ny)