SystemSimulationApp/PythonModels/registry.py

from __future__ import annotations

from collections.abc import Callable, Mapping
from dataclasses import dataclass
from math import isfinite

from PythonModels.components.cylinder import Cylinder
from PythonModels.components.orifice import Orifice
from PythonModels.components.resistive_pipe import ResistivePipe
from PythonModels.components.tank import Tank
from PythonModels.components.tee import Tee
from PythonModels.core.base import Component
from PythonModels.core.medium import IdealGasMedium
from PythonModels.core.ports import PortDefinition


@dataclass(frozen=True)
class ParameterSpec:
    name: str
    default: float
    minimum: float | None = None
    maximum: float | None = None
    minimum_exclusive: bool = False

    def validation_message(self, value: float) -> str | None:
        if not isfinite(value):
            return "must be finite"
        if self.minimum is not None:
            if self.minimum_exclusive and value <= self.minimum:
                return f"must be greater than {self.minimum:g}"
            if not self.minimum_exclusive and value < self.minimum:
                return f"must be at least {self.minimum:g}"
        if self.maximum is not None and value > self.maximum:
            return f"must be at most {self.maximum:g}"
        return None


ComponentFactory = Callable[
    [str, IdealGasMedium, Mapping[str, float]],
    Component,
]


@dataclass(frozen=True)
class ComponentModelSpec:
    model_type: str
    ports: tuple[PortDefinition, ...]
    parameters: tuple[ParameterSpec, ...]
    factory: ComponentFactory

    @property
    def parameter_by_name(self) -> dict[str, ParameterSpec]:
        return {parameter.name: parameter for parameter in self.parameters}

    def create(
        self,
        name: str,
        medium: IdealGasMedium,
        values: Mapping[str, float],
    ) -> Component:
        resolved = {
            parameter.name: values.get(parameter.name, parameter.default)
            for parameter in self.parameters
        }
        for parameter in self.parameters:
            message = parameter.validation_message(resolved[parameter.name])
            if message is not None:
                raise ValueError(
                    f"Parameter '{parameter.name}' on component '{name}' {message}."
                )
        unknown = sorted(set(values) - set(self.parameter_by_name))
        if unknown:
            raise ValueError(
                f"Component '{name}' contains unsupported parameters: "
                + ", ".join(unknown)
                + "."
            )
        component = self.factory(name, medium, resolved)
        component.model_type = self.model_type
        return component


def _cylinder_factory(
    name: str,
    medium: IdealGasMedium,
    values: Mapping[str, float],
) -> Component:
    return Cylinder(
        name=name,
        medium=medium,
        V=values["volume"],
        p0=values["p0"],
        T0=values["T0"],
    )


def _tank_factory(
    name: str,
    medium: IdealGasMedium,
    values: Mapping[str, float],
) -> Component:
    return Tank(
        name=name,
        medium=medium,
        V=values["volume"],
        p0=values["p0"],
        T0=values["T0"],
    )


def _pipe_factory(
    name: str,
    medium: IdealGasMedium,
    values: Mapping[str, float],
) -> Component:
    return ResistivePipe(
        name=name,
        medium=medium,
        L=values["length"],
        D=values["diameter"],
        lambda_darcy=values["lambda_darcy"],
        p0=values["p0"],
        T0=values["T0"],
    )


def _orifice_factory(
    name: str,
    medium: IdealGasMedium,
    values: Mapping[str, float],
) -> Component:
    return Orifice(name=name, opening=values["opening"], K=values["K"])


def _tee_factory(
    name: str,
    medium: IdealGasMedium,
    values: Mapping[str, float],
) -> Component:
    return Tee(name=name)


COMPONENT_MODEL_REGISTRY: dict[str, ComponentModelSpec] = {
    "cylinder": ComponentModelSpec(
        model_type="cylinder",
        ports=(PortDefinition.pneumatic("port_b", nominal_role="outlet"),),
        parameters=(
            ParameterSpec("volume", 0.01, minimum=0.0, minimum_exclusive=True),
            ParameterSpec("p0", 35e6, minimum=0.0, minimum_exclusive=True),
            ParameterSpec("T0", 300.0, minimum=0.0, minimum_exclusive=True),
        ),
        factory=_cylinder_factory,
    ),
    "tank": ComponentModelSpec(
        model_type="tank",
        ports=(PortDefinition.pneumatic("port_a", nominal_role="inlet"),),
        parameters=(
            ParameterSpec("volume", 0.1, minimum=0.0, minimum_exclusive=True),
            ParameterSpec("p0", 1e5, minimum=0.0, minimum_exclusive=True),
            ParameterSpec("T0", 300.0, minimum=0.0, minimum_exclusive=True),
        ),
        factory=_tank_factory,
    ),
    "pipe": ComponentModelSpec(
        model_type="pipe",
        ports=(
            PortDefinition.pneumatic("port_a", nominal_role="inlet"),
            PortDefinition.pneumatic("port_b", nominal_role="outlet"),
        ),
        parameters=(
            ParameterSpec("length", 5.0, minimum=0.0, minimum_exclusive=True),
            ParameterSpec("diameter", 0.02, minimum=0.0, minimum_exclusive=True),
            ParameterSpec("lambda_darcy", 0.02, minimum=0.0),
            ParameterSpec("p0", 1e5, minimum=0.0, minimum_exclusive=True),
            ParameterSpec("T0", 300.0, minimum=0.0, minimum_exclusive=True),
        ),
        factory=_pipe_factory,
    ),
    "orifice": ComponentModelSpec(
        model_type="orifice",
        ports=(
            PortDefinition.pneumatic("port_a", nominal_role="inlet"),
            PortDefinition.pneumatic("port_b", nominal_role="outlet"),
        ),
        parameters=(
            ParameterSpec("K", 1e-5, minimum=0.0),
            ParameterSpec("opening", 1.0, minimum=0.0, maximum=1.0),
        ),
        factory=_orifice_factory,
    ),
    "tee": ComponentModelSpec(
        model_type="tee",
        ports=(
            PortDefinition.pneumatic("port_in", nominal_role="bidirectional"),
            PortDefinition.pneumatic("port_out1", nominal_role="bidirectional"),
            PortDefinition.pneumatic("port_out2", nominal_role="bidirectional"),
        ),
        parameters=(),
        factory=_tee_factory,
    ),
}


def get_component_model_spec(model_type: str) -> ComponentModelSpec:
    try:
        return COMPONENT_MODEL_REGISTRY[model_type]
    except KeyError as exc:
        raise ValueError(f"Unsupported model type: {model_type}.") from exc