import random from typing import Dict, Union, Any, Tuple, List from enum import Enum from flask import Flask, request, jsonify from nucon import Nucon, ParameterEnum, PumpStatus, PumpDryStatus, PumpOverloadStatus, BreakerStatus import threading import torch from nucon.model import ReactorDynamicsModel class OperatingState(Enum): # Tuple indicates a range of values, while list indicates a set of possible values OFFLINE = { 'CORE_TEMP': (18.0, 22.0), 'CORE_TEMP_OPERATIVE': [306.0], 'CORE_TEMP_MAX': [1000.0], 'CORE_TEMP_MIN': [0.0], 'CORE_TEMP_RESIDUAL': [False], 'CORE_PRESSURE': (0.9, 1.1), 'CORE_PRESSURE_MAX': [1.5], 'CORE_PRESSURE_OPERATIVE': [1.0], 'CORE_INTEGRITY': [100.0], 'CORE_WEAR': [0.0], 'CORE_STATE': ['OFFLINE'], 'CORE_STATE_CRITICALITY': [0.0], 'CORE_CRITICAL_MASS_REACHED': [False], 'CORE_CRITICAL_MASS_REACHED_COUNTER': [0], 'CORE_IMMINENT_FUSION': [False], 'CORE_READY_FOR_START': [False], 'CORE_STEAM_PRESENT': [False], 'CORE_HIGH_STEAM_PRESENT': [False], 'TIME': ['00:00:00'], 'TIME_STAMP': ['1970-01-01 00:00:00'], 'COOLANT_CORE_STATE': ['INACTIVE'], 'COOLANT_CORE_PRESSURE': (0.9, 1.1), 'COOLANT_CORE_MAX_PRESSURE': [1.5], 'COOLANT_CORE_VESSEL_TEMPERATURE': (18.0, 22.0), 'COOLANT_CORE_QUANTITY_IN_VESSEL': [0.0], 'COOLANT_CORE_PRIMARY_LOOP_LEVEL': [0.0], 'COOLANT_CORE_FLOW_SPEED': [0.0], 'COOLANT_CORE_FLOW_ORDERED_SPEED': [0.0], 'COOLANT_CORE_FLOW_REACHED_SPEED': [False], 'COOLANT_CORE_QUANTITY_CIRCULATION_PUMPS_PRESENT': [3], 'COOLANT_CORE_QUANTITY_FREIGHT_PUMPS_PRESENT': [2], 'COOLANT_CORE_CIRCULATION_PUMP_0_STATUS': [PumpStatus.INACTIVE], 'COOLANT_CORE_CIRCULATION_PUMP_1_STATUS': [PumpStatus.INACTIVE], 'COOLANT_CORE_CIRCULATION_PUMP_2_STATUS': [PumpStatus.INACTIVE], 'COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID], 'COOLANT_CORE_CIRCULATION_PUMP_1_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID], 'COOLANT_CORE_CIRCULATION_PUMP_2_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID], 'COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD], 'COOLANT_CORE_CIRCULATION_PUMP_1_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD], 'COOLANT_CORE_CIRCULATION_PUMP_2_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD], 'COOLANT_CORE_CIRCULATION_PUMP_0_ORDERED_SPEED': [0.0], 'COOLANT_CORE_CIRCULATION_PUMP_1_ORDERED_SPEED': [0.0], 'COOLANT_CORE_CIRCULATION_PUMP_2_ORDERED_SPEED': [0.0], 'COOLANT_CORE_CIRCULATION_PUMP_0_SPEED': [0.0], 'COOLANT_CORE_CIRCULATION_PUMP_1_SPEED': [0.0], 'COOLANT_CORE_CIRCULATION_PUMP_2_SPEED': [0.0], 'RODS_STATUS': ['INACTIVE'], 'RODS_MOVEMENT_SPEED': [0.0], 'RODS_MOVEMENT_SPEED_DECREASED_HIGH_TEMPERATURE': [False], 'RODS_DEFORMED': [False], 'RODS_TEMPERATURE': (18.0, 22.0), 'RODS_MAX_TEMPERATURE': [1000.0], 'RODS_POS_ORDERED': [0.0], 'RODS_POS_ACTUAL': [0.0], 'RODS_POS_REACHED': [True], 'RODS_QUANTITY': [100], 'RODS_ALIGNED': [True], 'GENERATOR_0_KW': [0.0], 'GENERATOR_1_KW': [0.0], 'GENERATOR_2_KW': [0.0], 'GENERATOR_0_V': [0.0], 'GENERATOR_1_V': [0.0], 'GENERATOR_2_V': [0.0], 'GENERATOR_0_A': [0.0], 'GENERATOR_1_A': [0.0], 'GENERATOR_2_A': [0.0], 'GENERATOR_0_HERTZ': [0.0], 'GENERATOR_1_HERTZ': [0.0], 'GENERATOR_2_HERTZ': [0.0], 'GENERATOR_0_BREAKER': [BreakerStatus.OPEN], 'GENERATOR_1_BREAKER': [BreakerStatus.OPEN], 'GENERATOR_2_BREAKER': [BreakerStatus.OPEN], 'STEAM_TURBINE_0_RPM': [0.0], 'STEAM_TURBINE_1_RPM': [0.0], 'STEAM_TURBINE_2_RPM': [0.0], 'STEAM_TURBINE_0_TEMPERATURE': (18.0, 22.0), 'STEAM_TURBINE_1_TEMPERATURE': (18.0, 22.0), 'STEAM_TURBINE_2_TEMPERATURE': (18.0, 22.0), 'STEAM_TURBINE_0_PRESSURE': (0.9, 1.1), 'STEAM_TURBINE_1_PRESSURE': (0.9, 1.1), 'STEAM_TURBINE_2_PRESSURE': (0.9, 1.1), } NOMINAL = { 'CORE_TEMP': (290.0, 370.0), 'CORE_TEMP_OPERATIVE': [306.0], 'CORE_TEMP_MAX': [1000.0], 'CORE_TEMP_MIN': [0.0], 'CORE_TEMP_RESIDUAL': [False], 'CORE_PRESSURE': (14.5, 15.5), 'CORE_PRESSURE_MAX': [16.0], 'CORE_PRESSURE_OPERATIVE': [15.0], 'CORE_INTEGRITY': (99.0, 100.0), 'CORE_WEAR': (0.0, 1.0), 'CORE_STATE': ['ACTIVE'], 'CORE_STATE_CRITICALITY': (0.9, 1.1), 'CORE_CRITICAL_MASS_REACHED': [True], 'CORE_CRITICAL_MASS_REACHED_COUNTER': [1], 'CORE_IMMINENT_FUSION': [False], 'CORE_READY_FOR_START': [True], 'CORE_STEAM_PRESENT': [True], 'CORE_HIGH_STEAM_PRESENT': [False], 'TIME': ['00:00:00'], 'TIME_STAMP': ['1970-01-01 00:00:00'], 'COOLANT_CORE_STATE': ['ACTIVE'], 'COOLANT_CORE_PRESSURE': (13.5, 14.5), 'COOLANT_CORE_MAX_PRESSURE': [15.0], 'COOLANT_CORE_VESSEL_TEMPERATURE': (270.0, 290.0), 'COOLANT_CORE_QUANTITY_IN_VESSEL': (95.0, 100.0), 'COOLANT_CORE_PRIMARY_LOOP_LEVEL': (95.0, 100.0), 'COOLANT_CORE_FLOW_SPEED': (9.5, 10.5), 'COOLANT_CORE_FLOW_ORDERED_SPEED': [10.0], 'COOLANT_CORE_FLOW_REACHED_SPEED': [True], 'COOLANT_CORE_QUANTITY_CIRCULATION_PUMPS_PRESENT': [3], 'COOLANT_CORE_QUANTITY_FREIGHT_PUMPS_PRESENT': [3], 'COOLANT_CORE_CIRCULATION_PUMP_0_STATUS': [PumpStatus.ACTIVE_SPEED_REACHED], 'COOLANT_CORE_CIRCULATION_PUMP_1_STATUS': [PumpStatus.ACTIVE_SPEED_REACHED], 'COOLANT_CORE_CIRCULATION_PUMP_2_STATUS': [PumpStatus.ACTIVE_SPEED_REACHED], 'COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID], 'COOLANT_CORE_CIRCULATION_PUMP_1_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID], 'COOLANT_CORE_CIRCULATION_PUMP_2_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID], 'COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD], 'COOLANT_CORE_CIRCULATION_PUMP_1_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD], 'COOLANT_CORE_CIRCULATION_PUMP_2_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD], 'COOLANT_CORE_CIRCULATION_PUMP_0_ORDERED_SPEED': (95.0, 100.0), 'COOLANT_CORE_CIRCULATION_PUMP_1_ORDERED_SPEED': (95.0, 100.0), 'COOLANT_CORE_CIRCULATION_PUMP_2_ORDERED_SPEED': (95.0, 100.0), 'COOLANT_CORE_CIRCULATION_PUMP_0_SPEED': (95.0, 100.0), 'COOLANT_CORE_CIRCULATION_PUMP_1_SPEED': (95.0, 100.0), 'COOLANT_CORE_CIRCULATION_PUMP_2_SPEED': (95.0, 100.0), 'RODS_STATUS': ['ACTIVE'], 'RODS_MOVEMENT_SPEED': (0.9, 1.1), 'RODS_MOVEMENT_SPEED_DECREASED_HIGH_TEMPERATURE': [False], 'RODS_DEFORMED': [False], 'RODS_TEMPERATURE': (290.0, 310.0), 'RODS_MAX_TEMPERATURE': [1000.0], 'RODS_POS_ORDERED': [50.0], 'RODS_POS_ACTUAL': (49.5, 50.5), 'RODS_POS_REACHED': [True], 'RODS_QUANTITY': [100], 'RODS_ALIGNED': [True], 'GENERATOR_0_KW': (950.0, 1050.0), 'GENERATOR_1_KW': (950.0, 1050.0), 'GENERATOR_2_KW': (950.0, 1050.0), 'GENERATOR_0_V': (380.0, 420.0), 'GENERATOR_1_V': (380.0, 420.0), 'GENERATOR_2_V': (380.0, 420.0), 'GENERATOR_0_A': (2375.0, 2625.0), 'GENERATOR_1_A': (2375.0, 2625.0), 'GENERATOR_2_A': (2375.0, 2625.0), 'GENERATOR_0_HERTZ': (49.5, 50.5), 'GENERATOR_1_HERTZ': (49.5, 50.5), 'GENERATOR_2_HERTZ': (49.5, 50.5), 'GENERATOR_0_BREAKER': [BreakerStatus.CLOSED], 'GENERATOR_1_BREAKER': [BreakerStatus.CLOSED], 'GENERATOR_2_BREAKER': [BreakerStatus.CLOSED], 'STEAM_TURBINE_0_RPM': (2950.0, 3050.0), 'STEAM_TURBINE_1_RPM': (2950.0, 3050.0), 'STEAM_TURBINE_2_RPM': (2950.0, 3050.0), 'STEAM_TURBINE_0_TEMPERATURE': (270.0, 290.0), 'STEAM_TURBINE_1_TEMPERATURE': (270.0, 290.0), 'STEAM_TURBINE_2_TEMPERATURE': (270.0, 290.0), 'STEAM_TURBINE_0_PRESSURE': (13.5, 14.5), 'STEAM_TURBINE_1_PRESSURE': (13.5, 14.5), 'STEAM_TURBINE_2_PRESSURE': (13.5, 14.5), } class NuconSimulator: class Parameters: def __init__(self, nucon: Nucon): for param_name in nucon.get_all_readable(): setattr(self, param_name, None) def __init__(self, host: str = 'localhost', port: int = 8786): self._nucon = Nucon() self.parameters = self.Parameters(self._nucon) self.time = 0.0 self.allow_all_writes = False self.set_state(OperatingState.OFFLINE) self.model = None self.readable_params = list(self._nucon.get_all_readable().keys()) self.non_writable_params = [name for name, param in self._nucon.get_all_readable().items() if not param.is_writable] self._run(host, port) def get(self, parameter: Union[str, Any]) -> Any: if isinstance(parameter, str): return getattr(self.parameters, parameter) return getattr(self.parameters, parameter.id) def set(self, parameter: Union[str, Any], value: Any, force: bool = False) -> None: if isinstance(parameter, str): param_obj = self._nucon[parameter] else: param_obj = parameter if not param_obj.is_writable and not force and not self.allow_all_writes: raise ValueError(f"Parameter {param_obj.id} is not writable") # Convert value to the correct type try: if param_obj.enum_type: if isinstance(value, str): value = param_obj.enum_type[value.upper()] elif isinstance(value, int): value = param_obj.enum_type(value) elif not isinstance(value, param_obj.enum_type): raise ValueError(f"Invalid enum value for {param_obj.id}") else: value = param_obj.param_type(value) except (ValueError, KeyError): raise ValueError(f"Invalid type for parameter {param_obj.id}. Expected {param_obj.param_type}") # Check range if not forced if not force and param_obj.min_val is not None and param_obj.max_val is not None: if not param_obj.min_val <= value <= param_obj.max_val: raise ValueError(f"Value {value} is out of range for parameter {param_obj.id}. " f"Valid range: [{param_obj.min_val}, {param_obj.max_val}]") setattr(self.parameters, param_obj.id, value) def set_allow_all_writes(self, allow: bool) -> None: self.allow_all_writes = allow def update(self, time_step: float) -> None: self._update_reactor_state(time_step) self.time += time_step def load_model(self, model_path: str) -> None: try: self.model = ReactorDynamicsModel(self.readable_params, self.non_writable_params) self.model.load_state_dict(torch.load(model_path)) self.model.eval() # Set the model to evaluation mode print(f"Model loaded successfully from {model_path}") except Exception as e: print(f"Error loading model: {str(e)}") self.model = None def _update_reactor_state(self, time_step: float) -> None: if not self.model: raise ValueError("Model not set. Please load a model using load_model() method.") state = {} for param in self.readable_params: value = self.get(param) if isinstance(value, Enum): value = value.value state[param] = value # Use the model to predict the next state with torch.no_grad(): next_state = self.model(state, time_step) # Update the simulator's state for param, value in next_state.items(): self.set(param, value) def set_state(self, state: OperatingState) -> None: self._sample_parameters_from_state(state) def _sample_parameters_from_state(self, state: OperatingState) -> None: for param_name, value_spec in state.value.items(): param = self._nucon.get_all_readable()[param_name] if isinstance(value_spec, tuple): value = random.uniform(*value_spec) elif isinstance(value_spec, list): value = random.choice(value_spec) else: raise ValueError(f"Invalid value specification for parameter {param_name}") self.set(param, value, force=True) def _run(self, port: int = 8786, host: str = 'localhost', debug: bool = False): app = Flask(__name__) @app.route('/', methods=['GET']) def get_parameter(): variable = request.args.get('variable') if not variable: return jsonify({"error": "No variable specified"}), 400 try: value = self.get(variable) return str(value), 200 except KeyError: return jsonify({"error": f"Unknown variable: {variable}"}), 404 @app.route('/', methods=['POST']) def set_parameter(): variable = request.args.get('variable') value = request.args.get('value') if not variable or value is None: return jsonify({"error": "Both variable and value must be specified"}), 400 try: self.set(variable, value) return "OK", 200 except ValueError as e: return jsonify({"error": str(e)}), 400 except KeyError: return jsonify({"error": f"Unknown variable: {variable}"}), 404 def run_simulator(host='localhost', port=8786, debug=False): app.run(host=host, port=port, debug=debug) threading.Thread(target=run_simulator, args=(host, port, debug), daemon=True).start()