Source code for cfdmod.core.grouping

"""Per-:class:`DataSource` element grouping.

Elements may belong to *one* group per grouping. A data source can
carry many groupings simultaneously (surface name, planar selection,
S1 separation, ...); each lives at a key under
``DataSource.groupings``.

This module provides:

- :class:`Grouping` -- a frozen view over the per-element index plus an
  optional ``id_to_label`` mapping for human-readable group names.
- helpers (:func:`groups_in`, :func:`elements_in_group`) shared by ops.

The :class:`GroupsDataSource` (in ``data_source.py``) is a *separate*
concept: it carries one row per group rather than one row per element,
and stores fields that are aggregations over the original elements.
"""

from __future__ import annotations

__all__ = [
    "Grouping",
    "groups_in",
    "elements_in_group",
    "AggregationKind",
    "aggregate_rows",
]

from typing import Any, Literal

import numpy as np
from pydantic import BaseModel, ConfigDict, field_validator

AggregationKind = Literal["mean", "sum", "max", "min", "area_weighted_mean"]


[docs] class Grouping(BaseModel): """One grouping over a data source's element axis. Each element gets exactly one group id under this grouping; the sentinel ``-1`` marks "ungrouped". Groups are addressed by integer id; an optional :attr:`id_to_label` maps id to a human label (mirroring the legacy ``"{idx}-{body}"`` convention used in cfdmod's pressure pipeline). Attributes: name: Grouping name (e.g. ``"surface"``, ``"zoning"``). indices: ``(n_elements,)`` int32 array of group ids. ``-1`` = ungrouped. id_to_label: Optional dict mapping group id -> string label. """ model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) name: str indices: np.ndarray id_to_label: dict[int, str] | None = None @field_validator("indices", mode="before") @classmethod def _coerce(cls, v: Any) -> np.ndarray: arr = np.asarray(v, dtype=np.int32) if arr.ndim != 1: raise ValueError(f"Grouping indices must be 1-D; got shape {arr.shape}") return arr @property def n_elements(self) -> int: return int(self.indices.shape[0])
[docs] def label(self, group_id: int) -> str: """Resolve a group id to a label. Falls back to ``str(group_id)`` if no mapping is registered or the id is missing from the mapping. """ if self.id_to_label is None: return str(group_id) return self.id_to_label.get(int(group_id), str(group_id))
def groups_in(grouping: Grouping, *, include_ungrouped: bool = False) -> np.ndarray: """Sorted unique group ids in a grouping. Args: grouping: The grouping to inspect. include_ungrouped: If False (default), ``-1`` is dropped from the result. """ uniq = np.unique(grouping.indices) if not include_ungrouped: uniq = uniq[uniq != -1] return uniq def elements_in_group(grouping: Grouping, group_id: int) -> np.ndarray: """Indices of every element belonging to ``group_id``.""" return np.flatnonzero(grouping.indices == int(group_id)) def aggregate_rows( arr: np.ndarray, members: np.ndarray, agg: AggregationKind, weights: np.ndarray | None = None, ) -> np.ndarray: """Reduce ``arr[members]`` along axis 0 using ``agg``. Shared by ``field_series_for_groups`` and ``regroup_topology``. ``area_weighted_mean`` requires ``weights`` (per-element areas) and raises if the selected members have non-positive total weight (e.g. an empty group). """ sub = arr[members] if agg == "mean": return sub.mean(axis=0) if agg == "sum": return sub.sum(axis=0) if agg == "max": return sub.max(axis=0) if agg == "min": return sub.min(axis=0) if agg == "area_weighted_mean": if weights is None: raise ValueError("area_weighted_mean requires per-element weights (areas)") w = weights[members] total = w.sum() if total <= 0: raise ValueError("area_weighted_mean: total area is non-positive") if sub.ndim == 1: return float((sub * w).sum() / total) return (sub * w[:, None]).sum(axis=0) / total raise ValueError(f"unknown aggregation {agg!r}")