File size: 7,755 Bytes
78131a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """
Tests for core simulation components:
- DC power flow solver
- Environment lifecycle (reset, step, terminate)
- Grading system (scoring, bounds, reproducibility)
- Baseline heuristic policy
"""
import copy
import unittest
import numpy as np
from src.physics import DCSolver, IslandedException
from src.environment import OpenGridEnv
from src.tasks import TASKS
from src.models import GridAction, BusAdjustment
from src.grader import RobustnessGrader, compute_analytical_ceiling
from src.baseline import heuristic_policy
def task(task_id: str):
"""Get a deep-copied task config to prevent cross-test contamination."""
return copy.deepcopy(TASKS[task_id])
class TestDCSolver(unittest.TestCase):
def setUp(self):
self.num_buses = 3
self.lines = [
{'id': 'L01', 'from': 0, 'to': 1, 'susceptance': 100, 'connected': True},
{'id': 'L12', 'from': 1, 'to': 2, 'susceptance': 50, 'connected': True},
{'id': 'L02', 'from': 0, 'to': 2, 'susceptance': 100, 'connected': True}
]
self.solver = DCSolver(self.num_buses)
self.solver.update_grid(self.lines)
def test_power_flow_balance(self):
"""Slack bus should absorb any generation/load imbalance."""
p_inj = np.array([0.0, 50.0, -100.0])
theta, flows, slack_inj = self.solver.solve(p_inj)
# Check that flows are computed
self.assertIn('L01', flows)
self.assertIn('L02', flows)
def test_islanding_detection(self):
"""Disconnecting lines to island bus 2 should raise IslandedException."""
with self.assertRaises(IslandedException):
broken_lines = [
{'id': 'L01', 'from': 0, 'to': 1, 'susceptance': 100, 'connected': True},
{'id': 'L12', 'from': 1, 'to': 2, 'susceptance': 50, 'connected': False},
{'id': 'L02', 'from': 0, 'to': 2, 'susceptance': 100, 'connected': False}
]
self.solver.update_grid(broken_lines)
def test_slack_injection_returned(self):
"""solve() should return slack bus injection as third element."""
p_inj = np.array([0.0, 50.0, -100.0])
result = self.solver.solve(p_inj)
self.assertEqual(len(result), 3)
theta, flows, slack_inj = result
# Slack should inject ~50 MW to cover the deficit
self.assertAlmostEqual(slack_inj, 50.0, places=0)
def test_solve_before_update_raises(self):
"""Calling solve() on a fresh solver should raise RuntimeError."""
fresh = DCSolver(3)
with self.assertRaises(RuntimeError):
fresh.solve(np.array([0.0, 10.0, -10.0]))
def test_invalid_bus_index_raises(self):
"""Lines referencing out-of-range bus IDs should raise ValueError."""
bad_lines = [
{'id': 'L_bad', 'from': 0, 'to': 99, 'susceptance': 50, 'connected': True},
]
solver = DCSolver(3)
with self.assertRaises(ValueError):
solver.update_grid(bad_lines)
class TestEnvironment(unittest.TestCase):
def test_reset_returns_observation(self):
"""reset() should return a valid GridObservation."""
env = OpenGridEnv(task("task_easy"))
obs = env.reset()
self.assertEqual(obs.timestep, 0)
self.assertGreater(len(obs.buses), 0, "Observation should have buses")
self.assertGreater(len(obs.lines), 0, "Observation should have lines")
def test_step_returns_tuple(self):
"""step() should return (obs, reward, done, info)."""
env = OpenGridEnv(task("task_easy"))
env.reset()
obs, reward, done, info = env.step(GridAction())
self.assertEqual(obs.timestep, 1)
self.assertIsInstance(reward.value, float)
self.assertIsInstance(done, bool)
def test_reproducibility(self):
"""Running the same task twice should produce identical initial observations."""
env1 = OpenGridEnv(task("task_easy"))
obs1 = env1.reset()
env2 = OpenGridEnv(task("task_easy"))
obs2 = env2.reset()
self.assertEqual(obs1.grid_frequency, obs2.grid_frequency)
self.assertEqual(len(obs1.buses), len(obs2.buses))
def test_episode_terminates(self):
"""Episode should end after max_steps."""
config = task("task_easy")
env = OpenGridEnv(config)
env.reset()
done = False
steps = 0
while not done and steps < 100:
_, _, done, _ = env.step(GridAction())
steps += 1
self.assertTrue(done, "Episode should terminate")
self.assertLessEqual(steps, config["max_steps"])
def test_frequency_reasonable(self):
"""Frequency should stay in a reasonable range for do-nothing agent."""
env = OpenGridEnv(task("task_easy"))
obs = env.reset()
for _ in range(10):
obs, _, done, _ = env.step(GridAction())
if done:
break
self.assertGreater(obs.grid_frequency, 40.0,
"Frequency below reasonable minimum")
self.assertLess(obs.grid_frequency, 60.0,
"Frequency above reasonable maximum")
class TestGrader(unittest.TestCase):
def test_grader_score_range(self):
"""Grader should return score strictly in (0, 1) — never 0.0 or 1.0."""
grader = RobustnessGrader(task("task_easy"))
result = grader.evaluate_policy(heuristic_policy, n_episodes=1)
self.assertGreater(result["score"], 0.0)
self.assertLess(result["score"], 1.0)
def test_grader_all_tasks(self):
"""Grader should work on all registered tasks."""
for task_id, config in TASKS.items():
grader = RobustnessGrader(copy.deepcopy(config))
result = grader.evaluate_policy(heuristic_policy, n_episodes=1)
self.assertIn("score", result, f"Missing 'score' for {task_id}")
self.assertIn("avg_raw_reward", result,
f"Missing 'avg_raw_reward' for {task_id}")
class TestBaseline(unittest.TestCase):
def test_heuristic_returns_valid_action(self):
"""Heuristic policy should return a valid GridAction."""
env = OpenGridEnv(task("task_easy"))
obs = env.reset()
action = heuristic_policy(obs)
self.assertIsInstance(action, GridAction)
class TestReproducibility(unittest.TestCase):
def test_floor_deterministic(self):
"""Two calls to _estimate_bounds should produce identical floors (seeded RNG)."""
grader1 = RobustnessGrader(task("task_easy"))
grader1._estimate_bounds(n_samples=3)
grader2 = RobustnessGrader(task("task_easy"))
grader2._estimate_bounds(n_samples=3)
self.assertEqual(grader1.reward_floor, grader2.reward_floor,
"Floor should be deterministic with same seed")
def test_ceiling_is_analytical(self):
"""Ceiling should be max_steps * 1.2, not an empirical estimate."""
config = task("task_easy")
grader = RobustnessGrader(config)
bounds = grader.get_bounds()
expected_ceiling = compute_analytical_ceiling(config["max_steps"])
self.assertEqual(bounds["reward_ceiling"], expected_ceiling,
"Ceiling should match analytical formula")
def test_heuristic_score_below_one(self):
"""With analytical ceiling, heuristic should score < 1.0 (not degenerate)."""
grader = RobustnessGrader(task("task_easy"))
result = grader.evaluate_policy(heuristic_policy, n_episodes=1)
self.assertLess(result["score"], 1.0)
self.assertGreater(result["score"], 0.0)
if __name__ == '__main__':
unittest.main() |