openFOAM-RANS-to-GPU/python/src/foam_stepper/__init__.py

"""Python-driven OpenFOAM observability stepper."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Iterable

from ._runtime import configure_openfoam_environment

configure_openfoam_environment()

from . import _foam_stepper as _native
from .types import (
    PatchFieldView,
    PatchView,
    RaggedIntArray,
    FieldRegistryView,
    FieldView,
    MatrixView,
    MeshView,
    SolveResult,
    SourceLocation,
    TransformResult,
    _wrap_value,
)
from .state import (
    SolverStateValidationError,
    describe_matrix_state,
    describe_solver_state,
    export_matrix_state,
    export_solver_state,
    validate_matrix_state,
    validate_solver_state,
)

OpenFoamError = _native.OpenFoamError


def version() -> str:
    return _native.version()


class Case:
    """OpenFOAM case handle that creates Python-controlled steppers."""

    def __init__(self, path: str | Path, solver: str = "incompressibleFluid", time: str | None = None, libs: list[str] | None = None):
        if time is not None:
            raise NotImplementedError("Explicit time selection is not implemented yet; OpenFOAM selects startTime")
        if libs:
            raise NotImplementedError("Additional library loading is not implemented yet")
        self._path = str(path)
        self._solver = solver
        self._native = _native.Case(self._path, solver)

    @property
    def path(self) -> str:
        return self._path

    @property
    def solver_name(self) -> str:
        return self._solver

    def times(self) -> list[str]:
        root = Path(self._path)
        times: list[tuple[float, str]] = []
        for child in root.iterdir():
            if child.is_dir() and _looks_like_time(child.name):
                times.append((float(child.name), child.name))
        return [name for _, name in sorted(times)]

    def control_dict(self) -> str:
        return self.make_stepper().control_dict()

    def fv_schemes(self) -> str:
        return self.make_stepper().fv_schemes()

    def fv_solution(self) -> str:
        return self.make_stepper().fv_solution()

    def make_stepper(self) -> "SimpleStepper":
        return SimpleStepper(self._native.make_stepper())


class SimpleStepper:
    """Python facade over the C++ OpenFOAM observability stepper."""

    def __init__(self, native: Any):
        self._native = native

    @property
    def case_path(self) -> str:
        return self._native.case_path

    @property
    def solver_name(self) -> str:
        return self._native.solver_name

    def state(self) -> dict[str, Any]:
        return dict(self._native.state())

    def export_state(self, *, required_fields: Iterable[str] = ()) -> dict[str, Any]:
        return export_solver_state(self.mesh(), self.fields(), required_fields=required_fields)

    def state_summary(self, *, required_fields: Iterable[str] = ()) -> dict[str, Any]:
        return describe_solver_state(self.export_state(required_fields=required_fields))

    def control_dict(self) -> str:
        return self._native.control_dict()

    def fv_schemes(self) -> str:
        return self._native.fv_schemes()

    def fv_solution(self) -> str:
        return self._native.fv_solution()

    def mesh(self) -> MeshView:
        return MeshView.from_dict(self._native.mesh())

    def fields(self) -> FieldRegistryView:
        return FieldRegistryView.from_dict(self._native.fields())

    def pre_solve(self) -> TransformResult:
        return _transform(self._native.pre_solve())

    def advance_time(self) -> TransformResult:
        return _transform(self._native.advance_time())

    def begin_pimple_iteration(self) -> TransformResult:
        return _transform(self._native.begin_pimple_iteration())

    def end_pimple_iteration(self) -> TransformResult:
        return _transform(self._native.end_pimple_iteration())

    def post_solve(self, write: bool = False) -> TransformResult:
        return _transform(self._native.post_solve(write))

    def fv_models_correct(self) -> TransformResult:
        return _transform(self._native.fv_models_correct())

    def pre_predictor(self) -> TransformResult:
        return _transform(self._native.pre_predictor())

    def momentum_transport_predictor(self) -> TransformResult:
        return _transform(self._native.momentum_transport_predictor())

    def momentum_transport_corrector(self) -> TransformResult:
        return _transform(self._native.momentum_transport_corrector())

    def assemble_momentum_terms(self) -> TransformResult:
        return _transform(self._native.assemble_momentum_terms())

    def assemble_momentum_matrix(self, terms: Any | None = None) -> TransformResult:
        if terms is not None:
            # C++ recomputes the OpenFOAM expression to preserve native invariants.
            pass
        return _transform(self._native.assemble_momentum_matrix())

    def relax_matrix(self, matrix: Any | None = None) -> TransformResult:
        return _transform(self._native.relax_matrix())

    def constrain_matrix(self, matrix: Any | None = None) -> TransformResult:
        return _transform(self._native.constrain_matrix())

    def solve_momentum(self, matrix: Any | None = None) -> TransformResult:
        return _transform(self._native.solve_momentum())

    def compute_pressure_inputs(self, UEqn: Any | None = None) -> TransformResult:
        return _transform(self._native.compute_pressure_inputs())

    def assemble_pressure_matrix(self, inputs: Any | None = None) -> TransformResult:
        return _transform(self._native.assemble_pressure_matrix())

    def solve_pressure(self, matrix: Any | None = None) -> TransformResult:
        return _transform(self._native.solve_pressure())

    def correct_velocity_pressure_flux(self, inputs: Any | None = None, pEqn: Any | None = None) -> TransformResult:
        return _transform(self._native.correct_velocity_pressure_flux())

    def run_one_pimple_iteration(self) -> TransformResult:
        return _transform(self._native.run_one_pimple_iteration())

    def run_until(self, max_steps: int | None = None, max_iterations: int | None = None) -> list[TransformResult]:
        limit = max_steps if max_steps is not None else max_iterations
        results: list[TransformResult] = []
        while limit is None or len(results) < limit:
            result = self.run_one_pimple_iteration()
            results.append(result)
            if not result.outputs.get("time_loop_active", False):
                break
        return results


def _transform(raw: Any) -> TransformResult:
    return TransformResult.from_dict(raw)


def _looks_like_time(name: str) -> bool:
    try:
        float(name)
    except ValueError:
        return False
    return True


__all__ = [
    "Case",
    "SimpleStepper",
    "OpenFoamError",
    "PatchFieldView",
    "PatchView",
    "RaggedIntArray",
    "FieldRegistryView",
    "FieldView",
    "MatrixView",
    "MeshView",
    "SolveResult",
    "SourceLocation",
    "TransformResult",
    "SolverStateValidationError",
    "describe_matrix_state",
    "describe_solver_state",
    "export_matrix_state",
    "export_solver_state",
    "validate_matrix_state",
    "validate_solver_state",
    "version",
]