SystemSimulationApp/PythonModels/registry.py

from __future__ import annotations

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

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.metadata import ParameterDefinition
from PythonModels.core.medium import IdealGasMedium
from PythonModels.core.ports import PortDefinition


ParameterSpec = ParameterDefinition


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


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

    @property
    def parameter_by_name(self) -> dict[str, ParameterDefinition]:
        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
        if component.port_definitions != self.ports:
            raise ValueError(
                f"Component implementation {self.model_type} does not match "
                "its declared ports."
            )
        if component.parameter_values != resolved:
            raise ValueError(
                f"Component implementation {self.model_type} did not preserve its parameters."
            )
        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.MODEL_TYPE,
        ports=Cylinder.PORTS,
        parameters=Cylinder.PARAMETERS,
        factory=_cylinder_factory,
    ),
    "tank": ComponentModelSpec(
        model_type=Tank.MODEL_TYPE,
        ports=Tank.PORTS,
        parameters=Tank.PARAMETERS,
        factory=_tank_factory,
    ),
    "pipe": ComponentModelSpec(
        model_type=ResistivePipe.MODEL_TYPE,
        ports=ResistivePipe.PORTS,
        parameters=ResistivePipe.PARAMETERS,
        factory=_pipe_factory,
    ),
    "orifice": ComponentModelSpec(
        model_type=Orifice.MODEL_TYPE,
        ports=Orifice.PORTS,
        parameters=Orifice.PARAMETERS,
        factory=_orifice_factory,
    ),
    "tee": ComponentModelSpec(
        model_type=Tee.MODEL_TYPE,
        ports=Tee.PORTS,
        parameters=Tee.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