"""Strict contracts for a burst-simulation calibration design.""" from __future__ import annotations import math from datetime import datetime from pathlib import Path from typing import Literal from uuid import UUID from pydantic import BaseModel, ConfigDict, Field, model_validator Severity = Literal["light", "medium", "severe"] class StrictModel(BaseModel): model_config = ConfigDict(extra="forbid", allow_inf_nan=False) class ModelExpectations(StrictModel): node_count: int = Field(gt=0) pipe_count: int = Field(gt=0) valve_count: int = Field(gt=0) flow_unit: Literal["LPS"] pressure_unit: Literal["METERS"] pressure_measurement_unit: Literal["m"] flow_measurement_unit: Literal["m3/h"] emitter_exponent: float = Field(gt=0) hydraulic_timestep_seconds: int = Field(gt=0) report_timestep_seconds: int = Field(gt=0) class SeverityLevel(StrictModel): level: Severity cross_section_fraction: float = Field(gt=0, le=0.125) class SelectionRule(StrictModel): primary_pressure_drop_percent: float = Field(gt=0, le=100) corroborating_pressure_drop_percent: float = Field(gt=0, le=100) flow_change_percent: float = Field(gt=0, le=1000) consecutive_samples: int = Field(ge=1, le=20) expected_sample_count: int = Field(ge=2, le=100) relative_change_epsilon: float = Field(gt=0, le=0.001) normal_control_must_not_trigger: Literal[True] negative_monitor_pressure_allowed: Literal[False] choose_smallest_passing_area: Literal[True] class BindingValidation(StrictModel): max_pressure_distance_m: float = Field(gt=0, le=10) max_flow_distance_m: float = Field(gt=0, le=10) class SimulationCase(StrictModel): case_id: str = Field(pattern=r"^B[0-9]{2}$") pipe_id: str = Field(min_length=1) occurred_at: datetime @model_validator(mode="after") def aware_time(self): if self.occurred_at.tzinfo is None: raise ValueError("case time must include a timezone") return self class PressureBinding(StrictModel): device_id: str = Field(min_length=1) node_id: str = Field(min_length=1) class FlowBinding(StrictModel): device_id: str = Field(min_length=1) link_id: str = Field(min_length=1) class PipeCalibration(StrictModel): role: Literal["primary", "secondary"] pipe_id: str = Field(min_length=1) diameter_mm: float = Field(gt=0) length_m: float = Field(gt=0) start_node_id: str = Field(min_length=1) emitter_node_id: str = Field(min_length=1) longitude: float = Field(ge=-180, le=180) latitude: float = Field(gt=-90, lt=90) candidate_orifice_areas_cm2: dict[Severity, float] pressure_bindings: list[PressureBinding] = Field(min_length=1, max_length=20) flow_bindings: list[FlowBinding] = Field(min_length=1, max_length=20) boundary_valve_ids: list[str] = Field(min_length=1, max_length=20) output_node_ids: list[str] = Field(min_length=1, max_length=20) output_link_ids: list[str] = Field(default_factory=list, max_length=20) @model_validator(mode="after") def complete_outputs(self): if set(self.candidate_orifice_areas_cm2) != {"light", "medium", "severe"}: raise ValueError("all three severity areas are required") pressure_device_ids = [value.device_id for value in self.pressure_bindings] pressure_node_ids = [value.node_id for value in self.pressure_bindings] flow_device_ids = [value.device_id for value in self.flow_bindings] flow_link_ids = [value.link_id for value in self.flow_bindings] for values, label in ( (pressure_device_ids, "pressure devices"), (pressure_node_ids, "pressure nodes"), (flow_device_ids, "flow devices"), (flow_link_ids, "flow links"), (self.boundary_valve_ids, "boundary valves"), (self.output_node_ids, "output nodes"), (self.output_link_ids, "output links"), ): if len(values) != len(set(values)): raise ValueError(f"{label} must be distinct") if not {self.start_node_id, self.emitter_node_id, *pressure_node_ids}.issubset( self.output_node_ids ): raise ValueError("pipe endpoints and pressure nodes must be returned") if not {self.pipe_id, *flow_link_ids, *self.boundary_valve_ids}.issubset( self.output_link_ids ): raise ValueError("target, flow and boundary links must be returned") return self class BurstSimulationDesign(StrictModel): version: Literal["burst-simulation-design@1"] suite_id: Literal["tjwater2-v2-burst-calibration"] project_code: Literal["tjwater2_v2"] project_id: UUID model_expectations: ModelExpectations duration_seconds: int = Field(ge=60, le=86400) event_times: list[datetime] = Field(min_length=1, max_length=8) severity_levels: list[SeverityLevel] = Field(min_length=3, max_length=3) calibration_order: list[Severity] = Field(min_length=3, max_length=3) selection_rule: SelectionRule binding_validation: BindingValidation cases: list[SimulationCase] = Field(min_length=1, max_length=20) pipes: list[PipeCalibration] = Field(min_length=1, max_length=4) @model_validator(mode="after") def coherent_design(self): if any(value.tzinfo is None for value in self.event_times): raise ValueError("event times must include a timezone") if len(set(self.event_times)) != len(self.event_times): raise ValueError("event times must be distinct") levels = {entry.level: entry.cross_section_fraction for entry in self.severity_levels} if set(levels) != {"light", "medium", "severe"}: raise ValueError("light, medium and severe levels are required") if not levels["light"] < levels["medium"] < levels["severe"]: raise ValueError("severity fractions must increase") expected_levels = {"light": 1 / 32, "medium": 1 / 16, "severe": 1 / 8} if any( not math.isclose(levels[level], fraction, rel_tol=0, abs_tol=1e-12) for level, fraction in expected_levels.items() ): raise ValueError("severity fractions must remain 1/32, 1/16 and 1/8") if self.calibration_order != ["severe", "medium", "light"]: raise ValueError("calibration order must be severe, medium, light") report_step = self.model_expectations.report_timestep_seconds if self.duration_seconds % report_step: raise ValueError("duration must contain whole report timesteps") expected_samples = self.duration_seconds // report_step + 1 if self.selection_rule.expected_sample_count != expected_samples: raise ValueError("expected sample count does not match duration and report step") pipes = {pipe.pipe_id: pipe for pipe in self.pipes} if len(pipes) != len(self.pipes): raise ValueError("pipe IDs must be distinct") if len({case.case_id for case in self.cases}) != len(self.cases): raise ValueError("case IDs must be distinct") expected_cases = {(pipe_id, at) for pipe_id in pipes for at in self.event_times} actual_cases = {(case.pipe_id, case.occurred_at) for case in self.cases} if actual_cases != expected_cases: raise ValueError("cases must cover every pipe and event time exactly once") for pipe in self.pipes: cross_section_cm2 = math.pi * (pipe.diameter_mm / 10) ** 2 / 4 for level, fraction in levels.items(): expected = cross_section_cm2 * fraction actual = pipe.candidate_orifice_areas_cm2[level] if not math.isclose(actual, expected, rel_tol=0, abs_tol=0.001): raise ValueError( f"{pipe.pipe_id} {level} area does not match its diameter fraction" ) return self def load_simulation_design(path: Path) -> BurstSimulationDesign: return BurstSimulationDesign.model_validate_json(path.read_text(encoding="utf-8"))