99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
"""Reactor simulation harness and CLI."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import json
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from threading import Event
|
|
import time
|
|
from typing import Callable, Iterable, Optional
|
|
|
|
from .commands import ReactorCommand
|
|
from .logging_utils import configure_logging
|
|
from .reactor import Reactor
|
|
from .state import PlantState
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
|
|
CommandProvider = Callable[[float, PlantState], Optional[ReactorCommand]]
|
|
|
|
|
|
@dataclass
|
|
class ReactorSimulation:
|
|
reactor: Reactor
|
|
timestep: float = 1.0
|
|
duration: float | None = 3600.0
|
|
command_provider: CommandProvider | None = None
|
|
realtime: bool = False
|
|
stop_event: Event = field(default_factory=Event)
|
|
start_state: PlantState | None = None
|
|
last_state: PlantState | None = field(default=None, init=False)
|
|
|
|
def run(self) -> Iterable[PlantState]:
|
|
state = copy.deepcopy(self.start_state) if self.start_state else self.reactor.initial_state()
|
|
elapsed = 0.0
|
|
last_step_wall = time.time()
|
|
while self.duration is None or elapsed < self.duration:
|
|
if self.stop_event.is_set():
|
|
LOGGER.info("Stop signal received, terminating simulation loop")
|
|
break
|
|
snapshot = copy.deepcopy(state)
|
|
yield snapshot
|
|
command = self.command_provider(elapsed, snapshot) if self.command_provider else None
|
|
self.reactor.step(state, self.timestep, command)
|
|
elapsed += self.timestep
|
|
if self.realtime:
|
|
wall_elapsed = time.time() - last_step_wall
|
|
sleep_time = self.timestep - wall_elapsed
|
|
if sleep_time > 0:
|
|
time.sleep(sleep_time)
|
|
last_step_wall = time.time()
|
|
self.last_state = state
|
|
LOGGER.info("Simulation complete, %.0fs simulated", elapsed)
|
|
|
|
def log(self) -> list[dict[str, float]]:
|
|
return [snapshot for snapshot in (s.snapshot() for s in self.run())]
|
|
|
|
def stop(self) -> None:
|
|
self.stop_event.set()
|
|
|
|
|
|
def main() -> None:
|
|
log_level = os.getenv("FISSION_LOG_LEVEL", "INFO")
|
|
log_file = os.getenv("FISSION_LOG_FILE")
|
|
configure_logging(log_level, log_file)
|
|
realtime = os.getenv("FISSION_REALTIME", "0") == "1"
|
|
duration_env = os.getenv("FISSION_SIM_DURATION")
|
|
if duration_env:
|
|
duration = None if duration_env.lower() in {"none", "infinite"} else float(duration_env)
|
|
else:
|
|
duration = None if realtime else 600.0
|
|
reactor = Reactor.default()
|
|
sim = ReactorSimulation(reactor, timestep=5.0, duration=duration, realtime=realtime)
|
|
load_path = os.getenv("FISSION_LOAD_STATE")
|
|
save_path = os.getenv("FISSION_SAVE_STATE")
|
|
if load_path:
|
|
sim.start_state = reactor.load_state(load_path)
|
|
try:
|
|
if realtime:
|
|
LOGGER.info("Running in real-time mode (Ctrl+C to stop)...")
|
|
for _ in sim.run():
|
|
pass
|
|
else:
|
|
snapshots = sim.log()
|
|
LOGGER.info("Captured %d snapshots", len(snapshots))
|
|
print(json.dumps(snapshots[-5:], indent=2))
|
|
except KeyboardInterrupt:
|
|
sim.stop()
|
|
LOGGER.warning("Simulation interrupted by user")
|
|
finally:
|
|
if save_path and sim.last_state:
|
|
reactor.save_state(save_path, sim.last_state)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|