Source code for cfdmod.regroup.run

"""Regroup orchestration.

``run_regroup`` is the entry point used by the CLI and by callers that
already have a loaded mesh. It expands any
:class:`cfdmod.regroup.parameters.BySizeRoundedPerComponent` specs into
plain :class:`cfdmod.geometry.grouping.GroupingSpec` chains, runs the
grouping, builds the regrouped LnasFormat, and writes the new geometry
plus the rewritten timeseries.
"""

from __future__ import annotations

import pathlib

import numpy as np
from lnas import LnasFormat

from cfdmod.geometry.grouping import (
    ByDivisionsGrouping,
    ByZoningGrouping,
    GroupingSpec,
)
from cfdmod.geometry.grouping.kinds.by_divisions import _intervals_from_count
from cfdmod.io.geometry.transformation_config import TransformationConfig
from cfdmod.io.mesh import load_mesh
from cfdmod.io.xdmf import (
    write_processing_metadata,
    write_temporal_xdmf,
)
from cfdmod.logger import logger
from cfdmod.regroup.functions import (
    apply_regroup_to_timeseries,
    build_regroup_mapping,
    build_regrouped_mesh,
    build_sliced_regrouped_mesh,
)
from cfdmod.regroup.parameters import (
    BySizeRoundedPerComponent,
    RegroupConfig,
    RegroupSpec,
)

__all__ = ["expand_regroup_chain", "run_regroup"]


def _restricted_centroids(
    mesh: LnasFormat,
    parent_idxs: np.ndarray,
) -> np.ndarray:
    triangles = mesh.geometry.triangle_vertices  # (n_tri, 3, 3)
    return np.mean(triangles[parent_idxs], axis=1)  # (k, 3)


def _rounded_division_count(extent: float, target: float, min_n: int) -> int:
    """``max(min_n, round(extent / target))``; handles the degenerate extent=0 case."""
    if target <= 0:
        raise ValueError(f"target size must be > 0, got {target}")
    if extent <= 0.0:
        return max(min_n, 1)
    # int(x + 0.5) avoids banker's rounding so 0.5 -> 1 (matches user expectation
    # that a "barely there" division still counts as one cell).
    n = int(extent / target + 0.5)
    return max(min_n, n)


_IntervalsTriple = tuple[list[float], list[float], list[float]]


def _expand_one_per_component(
    spec: BySizeRoundedPerComponent,
    mesh: LnasFormat,
    prior_chain: list[GroupingSpec],
    transformation: TransformationConfig | None,
) -> tuple[
    list[GroupingSpec],
    set[str],
    dict[str, _IntervalsTriple],
    dict[str, np.ndarray],
]:
    """Expand a fan-out spec into one ``ByDivisionsGrouping`` per parent group.

    Returns ``(new_specs, consumed_parent_names)``. The consumed names
    are the parent groups whose triangles have been subdivided; the
    caller drops them from the final grouping so the leaf cells (which
    overlap with the parents) are the only output groups.
    """
    if not prior_chain:
        raise ValueError(
            "BySizeRoundedPerComponent requires at least one preceding grouping spec "
            "to define the components to fan out over."
        )

    prior = build_regroup_mapping(mesh, prior_chain, transformation)

    if spec.restrict_to is not None:
        missing = [n for n in spec.restrict_to if n not in prior.groups]
        if missing:
            raise ValueError(
                f"BySizeRoundedPerComponent: restrict_to references unknown "
                f"groups {missing}. Available: {sorted(prior.groups)}"
            )
        parent_names = [n for n in spec.restrict_to]
    else:
        parent_names = list(prior.groups.keys())

    # Compute centroids in the same (transformed) frame the binning will see.
    if transformation is not None:
        mesh_for_binning = mesh.copy()
        mesh_for_binning.geometry.apply_transformation(
            transformation.get_geometry_transformation()
        )
    else:
        mesh_for_binning = mesh

    expanded: list[GroupingSpec] = []
    consumed: set[str] = set()
    parent_intervals: dict[str, _IntervalsTriple] = {}
    parent_triangles: dict[str, np.ndarray] = {}
    for parent_name in parent_names:
        parent_idxs = np.asarray(prior.groups[parent_name], dtype=np.int64)
        if parent_idxs.size == 0:
            continue
        consumed.add(parent_name)
        parent_triangles[parent_name] = parent_idxs
        cents = _restricted_centroids(mesh_for_binning, parent_idxs)
        lo = cents.min(axis=0)
        hi = cents.max(axis=0)

        if spec.target_size_x is not None:
            n_x: int | None = _rounded_division_count(
                float(hi[0] - lo[0]), spec.target_size_x, spec.min_n_div
            )
        else:
            n_x = None
        if spec.target_size_y is not None:
            n_y: int | None = _rounded_division_count(
                float(hi[1] - lo[1]), spec.target_size_y, spec.min_n_div
            )
        else:
            n_y = None
        if spec.target_size_z is not None:
            n_z: int | None = _rounded_division_count(
                float(hi[2] - lo[2]), spec.target_size_z, spec.min_n_div
            )
        else:
            n_z = None

        sub_template = spec.name_template.replace("{parent}", parent_name)
        expanded.append(
            ByDivisionsGrouping(
                n_div_x=n_x,
                n_div_y=n_y,
                n_div_z=n_z,
                name_template=sub_template,
                restrict_to=[parent_name],
            )
        )
        parent_intervals[parent_name] = (
            _intervals_from_count(float(lo[0]), float(hi[0]), n_x),
            _intervals_from_count(float(lo[1]), float(hi[1]), n_y),
            _intervals_from_count(float(lo[2]), float(hi[2]), n_z),
        )

        logger.info(
            f"  expand[{parent_name}]: extents=({hi[0] - lo[0]:.3f}, "
            f"{hi[1] - lo[1]:.3f}, {hi[2] - lo[2]:.3f}) -> "
            f"n_div=({n_x}, {n_y}, {n_z})"
        )

    return expanded, consumed, parent_intervals, parent_triangles


[docs] def expand_regroup_chain( chain: list[RegroupSpec], mesh: LnasFormat, transformation: TransformationConfig | None, ) -> tuple[ list[GroupingSpec], set[str], dict[str, _IntervalsTriple], dict[str, np.ndarray], ]: """Resolve any regroup-local specs into plain ``GroupingSpec`` entries. Returns ``(expanded_specs, consumed_group_names, parent_intervals, parent_triangles)``. ``consumed`` names are intermediate parent groups that ``BySizeRoundedPerComponent`` has fanned out over; ``parent_intervals`` and ``parent_triangles`` carry the per-parent cut planes and triangle indices needed to drive the ``"sliced"`` aggregation mode (empty dicts when no fan-out happened). """ expanded: list[GroupingSpec] = [] consumed: set[str] = set() parent_intervals: dict[str, _IntervalsTriple] = {} parent_triangles: dict[str, np.ndarray] = {} for spec in chain: if isinstance(spec, BySizeRoundedPerComponent): new_specs, new_consumed, new_intervals, new_triangles = _expand_one_per_component( spec, mesh, expanded, transformation ) expanded.extend(new_specs) consumed |= new_consumed parent_intervals.update(new_intervals) parent_triangles.update(new_triangles) else: expanded.append(spec) return expanded, consumed, parent_intervals, parent_triangles
def _resolve_global_intervals( expanded: list[GroupingSpec], grouping, mesh: LnasFormat, ) -> tuple[_IntervalsTriple, np.ndarray] | None: """For a chain whose only zoning is a single trailing ByZoning/Divisions, return ``((x, y, z), all_grouped_triangle_idxs)``; else None. Used by sliced mode when the user wrote a one-shot zoning config (not a BySizeRoundedPerComponent fan-out). Picks the last zoning-like spec. """ last_zoning: ByZoningGrouping | None = None for spec in expanded: if isinstance(spec, ByZoningGrouping): last_zoning = spec elif isinstance(spec, ByDivisionsGrouping): # ByDivisionsGrouping resolves to a ByZoningGrouping; bbox-derived. cents = mesh.geometry.triangle_vertices.mean(axis=1) allowed = ( np.concatenate([grouping.groups[n] for n in (spec.restrict_to or [])]) if spec.restrict_to else np.arange(cents.shape[0], dtype=np.int64) ) if allowed.size == 0: continue cand = cents[allowed] lo = cand.min(axis=0) hi = cand.max(axis=0) last_zoning = ByZoningGrouping( x_intervals=_intervals_from_count(float(lo[0]), float(hi[0]), spec.n_div_x), y_intervals=_intervals_from_count(float(lo[1]), float(hi[1]), spec.n_div_y), z_intervals=_intervals_from_count(float(lo[2]), float(hi[2]), spec.n_div_z), ) if last_zoning is None: return None all_idxs = ( np.unique( np.concatenate([np.asarray(idxs, dtype=np.int64) for idxs in grouping.groups.values()]) ) if grouping.groups else np.array([], dtype=np.int64) ) return ( ( list(last_zoning.x_intervals), list(last_zoning.y_intervals), list(last_zoning.z_intervals), ), all_idxs, ) def _record_processing_metadata( output_h5: pathlib.Path, group: str, cfg: RegroupConfig, expanded: list[GroupingSpec], *, input_geometry: pathlib.Path | None, input_timeseries: pathlib.Path, ) -> None: config_dict = { "regroup": { "groupings": [s.model_dump(mode="python") for s in cfg.groupings], "expanded_groupings": [s.model_dump(mode="python") for s in expanded], "aggregation": cfg.aggregation, "timeseries_group": cfg.timeseries_group, "transformation": ( cfg.transformation.model_dump(mode="python") if cfg.transformation is not None else None ), "unassigned_policy": cfg.unassigned_policy, } } extra = {"input_timeseries": str(input_timeseries)} if input_geometry is not None: extra["input_geometry"] = str(input_geometry) write_processing_metadata(output_h5, group, config=config_dict, extra=extra)
[docs] def run_regroup( cfg: RegroupConfig, geometry: pathlib.Path | LnasFormat, timeseries: pathlib.Path, output_dir: pathlib.Path, ) -> None: """Run the full regroup pipeline and write outputs to ``output_dir``. Outputs: - ``geometry.lnas`` (always); ``geometry.stl`` if ``cfg.output_geometry_format == "lnas_and_stl"``. - ``{cfg.timeseries_group}.regrouped.h5`` and a sibling ``.xdmf``. """ output_dir = pathlib.Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) timeseries = pathlib.Path(timeseries) input_geometry_path = pathlib.Path(geometry) if not isinstance(geometry, LnasFormat) else None mesh = load_mesh(geometry) expanded, consumed, parent_intervals, parent_triangles = expand_regroup_chain( cfg.groupings, mesh, cfg.transformation ) if not expanded: raise ValueError( "regroup: chain expanded to zero specs (no work to do); check the config." ) grouping = build_regroup_mapping(mesh, expanded, cfg.transformation) if consumed: kept = {n: idxs for n, idxs in grouping.groups.items() if n not in consumed} from cfdmod.geometry.grouping import GroupingResult grouping = GroupingResult(parent_n_triangles=grouping.parent_n_triangles, groups=kept) logger.info( f"regroup: dropped {len(consumed)} consumed parent group(s): {sorted(consumed)}" ) if not grouping.groups and cfg.unassigned_policy == "drop": raise ValueError( "regroup: chain produced zero non-empty groups and unassigned_policy='drop'." ) if cfg.aggregation == "sliced": if not parent_intervals: # No fan-out happened. Try to extract a single trailing zoning to # use as the global cut. Fall back to error otherwise. global_resolved = _resolve_global_intervals(expanded, grouping, mesh) if global_resolved is None: raise ValueError( "regroup (sliced): could not resolve cut intervals from the chain. " "Use a chain with BySizeRoundedPerComponent or a trailing " "ByZoningGrouping / ByDivisionsGrouping." ) global_intervals, global_idxs = global_resolved parent_intervals = {"*": global_intervals} parent_triangles = {"*": global_idxs} logger.info( f"regroup (sliced): slicing {len(parent_intervals)} parent " f"component(s) along axis-aligned cut planes..." ) new_lnas, regroup_index = build_sliced_regrouped_mesh( mesh, grouping, parent_intervals=parent_intervals, parent_triangles=parent_triangles, unassigned_policy=cfg.unassigned_policy, ) else: new_lnas, regroup_index = build_regrouped_mesh( mesh, grouping, aggregation=cfg.aggregation, unassigned_policy=cfg.unassigned_policy, ) out_lnas = output_dir / "geometry.lnas" new_lnas.to_file(out_lnas) if cfg.output_geometry_format == "lnas_and_stl": new_lnas.geometry.export_stl(output_dir / "geometry.stl") logger.info( f"regroup: wrote geometry to {out_lnas} ({new_lnas.geometry.triangles.shape[0]} " f"triangle(s), {len(new_lnas.surfaces)} surface(s))" ) out_h5 = output_dir / f"{cfg.timeseries_group}.regrouped.h5" apply_regroup_to_timeseries( timeseries, out_h5, group=cfg.timeseries_group, regroup_index=regroup_index, new_triangles=new_lnas.geometry.triangles, new_vertices=new_lnas.geometry.vertices, ) out_xdmf = out_h5.with_suffix(".xdmf") write_temporal_xdmf(out_h5, out_xdmf, cfg.timeseries_group) _record_processing_metadata( out_h5, cfg.timeseries_group, cfg, expanded, input_geometry=input_geometry_path, input_timeseries=timeseries, ) logger.info(f"regroup: wrote timeseries to {out_h5} (+ sibling XDMF)")