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, ReactorKNNModel import pickle class HighUncertaintyError(Exception): """Raised by NuconSimulator when the dynamics model uncertainty exceeds the threshold. Caught by NuconEnv/NuconGoalEnv and returned as truncated=True so the RL algorithm bootstraps the value rather than treating it as a terminal state. """ def __init__(self, uncertainty: float, threshold: float): self.uncertainty = uncertainty self.threshold = threshold super().__init__(f"Model uncertainty {uncertainty:.3f} exceeded threshold {threshold:.3f}") 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'], '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'], '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, uncertainty_threshold: float = None): self._nucon = Nucon() self.parameters = self.Parameters(self._nucon) self.host = host self.port = port self.time = 0.0 self.allow_all_writes = False self.uncertainty_threshold = uncertainty_threshold 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 set_model(self, model) -> None: """Set a pre-loaded ReactorDynamicsModel or ReactorKNNModel directly.""" self.model = model if isinstance(model, ReactorDynamicsModel): self.model.eval() def load_model(self, model_path: str) -> None: """Load a model from a file. .pkl → ReactorKNNModel, otherwise → ReactorDynamicsModel (torch).""" try: if model_path.endswith('.pkl'): with open(model_path, 'rb') as f: self.model = pickle.load(f) print(f"kNN model loaded from {model_path}") else: # Reconstruct shell from the saved state dict; input/output params # are stored inside the checkpoint. checkpoint = torch.load(model_path, weights_only=False) if isinstance(checkpoint, dict) and 'input_params' in checkpoint: self.model = ReactorDynamicsModel(checkpoint['input_params'], checkpoint['output_params']) self.model.load_state_dict(checkpoint['state_dict']) else: # Legacy: plain state dict — fall back using sim readable/non-writable lists self.model = ReactorDynamicsModel(self.readable_params, self.non_writable_params) self.model.load_state_dict(checkpoint) self.model.eval() print(f"NN model loaded 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() or set_model().") # Build state dict using only the params the model knows about state = {} for param_id in self.model.input_params: value = getattr(self.parameters, param_id, None) if isinstance(value, Enum): value = value.value if value is None: value = 0.0 # fallback for params not initialised in sim state state[param_id] = value # Forward pass — same interface for both NN and kNN if isinstance(self.model, ReactorDynamicsModel): with torch.no_grad(): next_state = self.model.forward(state, time_step) else: if self.uncertainty_threshold is not None: next_state, uncertainty = self.model.forward_with_uncertainty(state, time_step) if uncertainty > self.uncertainty_threshold: raise HighUncertaintyError(uncertainty, self.uncertainty_threshold) else: next_state = self.model.forward(state, time_step) # Update only the output params the model predicts for param_id, value in next_state.items(): try: self.set(param_id, value, force=True) except (ValueError, KeyError): pass # ignore params that can't be set (type mismatch, unknown) 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, host: str = 'localhost', port: int = 8786, debug: bool = False): app = Flask(__name__) @app.route('/', methods=['GET']) def get_parameter(): variable = request.args.get('variable', '').upper() if not variable: return jsonify({"error": "No variable specified"}), 400 if variable == 'WEBSERVER_LIST_VARIABLES': return '\n'.join(self._nucon._parameters.keys()), 200 if variable == 'WEBSERVER_BATCH_GET': value_arg = request.args.get('value', '') names = [n.strip().upper() for n in value_arg.split(',') if n.strip()] if not names: names = list(self._nucon._parameters.keys()) result = {} for name in names: try: result[name] = str(self.get(name)) except (KeyError, AttributeError): pass return jsonify(result), 200 try: value = self.get(variable) return str(value), 200 except (KeyError, AttributeError): 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()