feat: add local environment scaffold and baselines

baselines/__init__.py

@@ -0,0 +1 @@
+"""Random and heuristic baselines for the stellarator design environment."""

baselines/compare.py

@@ -0,0 +1,47 @@
+"""Run both baselines and print a comparison summary."""
+
+from __future__ import annotations
+
+import sys
+
+from baselines.heuristic_agent import heuristic_episode
+from baselines.random_agent import random_episode
+from server.environment import StellaratorEnvironment
+
+
+def main(n_episodes: int = 20) -> None:
+    env = StellaratorEnvironment()
+
+    random_rewards: list[float] = []
+    heuristic_rewards: list[float] = []
+    random_best_qs: list[float] = []
+    heuristic_best_qs: list[float] = []
+
+    for i in range(n_episodes):
+        rr, rt = random_episode(env, seed=i)
+        random_rewards.append(rr)
+        random_best_qs.append(rt[-1]["best_qs"])
+
+        hr, ht = heuristic_episode(env, seed=i)
+        heuristic_rewards.append(hr)
+        heuristic_best_qs.append(ht[-1]["best_qs"])
+
+    r_mean = sum(random_rewards) / len(random_rewards)
+    h_mean = sum(heuristic_rewards) / len(heuristic_rewards)
+    r_qs = sum(random_best_qs) / len(random_best_qs)
+    h_qs = sum(heuristic_best_qs) / len(heuristic_best_qs)
+
+    print(f"{'Metric':<25} {'Random':>12} {'Heuristic':>12}")
+    print("-" * 51)
+    print(f"{'Mean reward':<25} {r_mean:>+12.4f} {h_mean:>+12.4f}")
+    print(f"{'Mean best QS residual':<25} {r_qs:>12.6f} {h_qs:>12.6f}")
+    print(f"{'Episodes':<25} {n_episodes:>12d} {n_episodes:>12d}")
+    print()
+
+    wins = sum(1 for h, r in zip(heuristic_rewards, random_rewards) if h > r)
+    print(f"Heuristic wins: {wins}/{n_episodes} episodes ({100 * wins / n_episodes:.0f}%)")
+
+
+if __name__ == "__main__":
+    n = int(sys.argv[1]) if len(sys.argv) > 1 else 20
+    main(n)

baselines/heuristic_agent.py

@@ -0,0 +1,110 @@
+"""Heuristic baseline agent for the stellarator design environment.
+
+Strategy: guided perturbations informed by domain knowledge.
+1. Probe the most sensitive coefficient (zs12) first with a small move.
+2. Apply medium perturbations in directions that typically improve QS.
+3. Use restore_best to recover from any worsening.
+4. Submit before exhausting budget.
+"""
+
+from __future__ import annotations
+
+import sys
+
+from fusion_lab.models import StellaratorAction
+from server.environment import StellaratorEnvironment
+
+STRATEGY: list[tuple[str, str, str, str]] = [
+    ("tune_zs12", "decrease", "small", "hot"),
+    ("tune_zs12", "decrease", "medium", "hot"),
+    ("tune_rc11", "increase", "small", "hot"),
+    ("tune_rc10", "increase", "medium", "hot"),
+    ("tune_zs11", "decrease", "small", "hot"),
+]
+
+
+def heuristic_episode(
+    env: StellaratorEnvironment, seed: int | None = None
+) -> tuple[float, list[dict[str, object]]]:
+    obs = env.reset(seed=seed)
+    total_reward = 0.0
+    trace: list[dict[str, object]] = [{"step": 0, "qs": obs.quasi_symmetry_residual}]
+    prev_best = obs.best_qs_residual
+
+    for operator, direction, magnitude, restart in STRATEGY:
+        if obs.done or obs.budget_remaining <= 1:
+            break
+
+        action = StellaratorAction(
+            intent="run",
+            operator=operator,
+            direction=direction,
+            magnitude=magnitude,
+            restart=restart,
+        )
+        obs = env.step(action)
+        total_reward += obs.reward or 0.0
+        trace.append(
+            {
+                "step": len(trace),
+                "action": f"{operator} {direction} {magnitude}",
+                "qs": obs.quasi_symmetry_residual,
+                "best_qs": obs.best_qs_residual,
+                "reward": obs.reward,
+            }
+        )
+
+        if obs.best_qs_residual > prev_best and obs.budget_remaining > 1:
+            restore = StellaratorAction(intent="restore_best")
+            obs = env.step(restore)
+            total_reward += obs.reward or 0.0
+            trace.append(
+                {
+                    "step": len(trace),
+                    "action": "restore_best",
+                    "qs": obs.quasi_symmetry_residual,
+                    "best_qs": obs.best_qs_residual,
+                    "reward": obs.reward,
+                }
+            )
+
+        prev_best = obs.best_qs_residual
+
+    if not obs.done:
+        submit = StellaratorAction(intent="submit")
+        obs = env.step(submit)
+        total_reward += obs.reward or 0.0
+        trace.append(
+            {
+                "step": len(trace),
+                "action": "submit",
+                "qs": obs.quasi_symmetry_residual,
+                "best_qs": obs.best_qs_residual,
+                "reward": obs.reward,
+            }
+        )
+
+    return total_reward, trace
+
+
+def main(n_episodes: int = 20) -> None:
+    env = StellaratorEnvironment()
+    rewards: list[float] = []
+
+    for i in range(n_episodes):
+        total_reward, trace = heuristic_episode(env, seed=i)
+        final = trace[-1]
+        rewards.append(total_reward)
+        print(
+            f"Episode {i:3d}: steps={len(trace) - 1}  "
+            f"final_qs={final['qs']:.6f}  best_qs={final['best_qs']:.6f}  "
+            f"reward={total_reward:+.4f}"
+        )
+
+    mean_reward = sum(rewards) / len(rewards)
+    print(f"\nHeuristic baseline ({n_episodes} episodes): mean_reward={mean_reward:+.4f}")
+
+
+if __name__ == "__main__":
+    n = int(sys.argv[1]) if len(sys.argv) > 1 else 20
+    main(n)

baselines/random_agent.py

@@ -0,0 +1,71 @@
+"""Random baseline agent for the stellarator design environment."""
+
+from __future__ import annotations
+
+import random
+import sys
+
+from fusion_lab.models import StellaratorAction
+from server.environment import StellaratorEnvironment
+
+OPERATORS = ["tune_rc10", "tune_rc11", "tune_zs11", "tune_zs12"]
+DIRECTIONS = ["increase", "decrease"]
+MAGNITUDES = ["small", "medium", "large"]
+RESTARTS = ["hot", "cold"]
+
+
+def random_episode(
+    env: StellaratorEnvironment, seed: int | None = None
+) -> tuple[float, list[dict[str, object]]]:
+    rng = random.Random(seed)
+    obs = env.reset(seed=seed)
+    total_reward = 0.0
+    trace: list[dict[str, object]] = [{"step": 0, "qs": obs.quasi_symmetry_residual}]
+
+    while not obs.done:
+        if obs.budget_remaining <= 0:
+            action = StellaratorAction(intent="submit")
+        else:
+            action = StellaratorAction(
+                intent="run",
+                operator=rng.choice(OPERATORS),
+                direction=rng.choice(DIRECTIONS),
+                magnitude=rng.choice(MAGNITUDES),
+                restart=rng.choice(RESTARTS),
+            )
+        obs = env.step(action)
+        total_reward += obs.reward or 0.0
+        trace.append(
+            {
+                "step": len(trace),
+                "action": action.intent,
+                "qs": obs.quasi_symmetry_residual,
+                "best_qs": obs.best_qs_residual,
+                "reward": obs.reward,
+            }
+        )
+
+    return total_reward, trace
+
+
+def main(n_episodes: int = 20) -> None:
+    env = StellaratorEnvironment()
+    rewards: list[float] = []
+
+    for i in range(n_episodes):
+        total_reward, trace = random_episode(env, seed=i)
+        final = trace[-1]
+        rewards.append(total_reward)
+        print(
+            f"Episode {i:3d}: steps={len(trace) - 1}  "
+            f"final_qs={final['qs']:.6f}  best_qs={final['best_qs']:.6f}  "
+            f"reward={total_reward:+.4f}"
+        )
+
+    mean_reward = sum(rewards) / len(rewards)
+    print(f"\nRandom baseline ({n_episodes} episodes): mean_reward={mean_reward:+.4f}")
+
+
+if __name__ == "__main__":
+    n = int(sys.argv[1]) if len(sys.argv) > 1 else 20
+    main(n)

fusion_lab/client.py

@@ -7,13 +7,13 @@ from fusion_lab.models import StellaratorAction, StellaratorObservation, Stellar
 
 
 class FusionLabClient(EnvClient[StellaratorAction, StellaratorObservation, StellaratorState]):
-    """Thin typed client wrapper for the remote OpenEnv environment."""
+    """Typed client wrapper for the remote Fusion Design Lab environment."""
 
     def _step_payload(self, action: StellaratorAction) -> dict[str, object]:
         return action.model_dump(exclude_none=True)
 
     def _parse_result(self, payload: dict[str, object]) -> StepResult[StellaratorObservation]:
-        observation = StellaratorObservation(**payload)
+        observation = StellaratorObservation.model_validate(payload)
         return StepResult(
             observation=observation,
             reward=observation.reward,
@@ -21,4 +21,4 @@ class FusionLabClient(EnvClient[StellaratorAction, StellaratorObservation, Stell
         )
 
     def _parse_state(self, payload: dict[str, object]) -> StellaratorState:
-        return StellaratorState(**payload)
+        return StellaratorState.model_validate(payload)

fusion_lab/models.py

@@ -2,8 +2,8 @@ from __future__ import annotations
 
 from typing import Literal
 
-from pydantic import BaseModel, Field
-
+from openenv.core import Action, Observation, State
+from pydantic import Field
 
 ActionIntent = Literal["run", "submit", "restore_best"]
 OperatorName = Literal["tune_rc10", "tune_rc11", "tune_zs11", "tune_zs12"]
@@ -12,7 +12,7 @@ MagnitudeName = Literal["small", "medium", "large"]
 RestartMode = Literal["hot", "cold"]
 
 
-class StellaratorAction(BaseModel):
+class StellaratorAction(Action):
     intent: ActionIntent
     operator: OperatorName | None = None
     direction: DirectionName | None = None
@@ -21,26 +21,23 @@ class StellaratorAction(BaseModel):
     reasoning: str = ""
 
 
-class StellaratorObservation(BaseModel):
-    diagnostics_text: str
-    quasi_symmetry_residual: float
-    aspect_ratio: float
-    rotational_transform_axis: float
-    rotational_transform_edge: float
-    magnetic_well_depth: float
-    volume: float
-    vmec_converged: bool
-    step_number: int
-    budget_remaining: int
-    best_qs_residual: float
-    constraints_satisfied: bool
-    target_spec: str
-    reward: float | None = None
-    done: bool = False
-
-
-class StellaratorState(BaseModel):
-    step_count: int = 0
+class StellaratorObservation(Observation):
+    diagnostics_text: str = ""
+    quasi_symmetry_residual: float = 0.0
+    aspect_ratio: float = 0.0
+    rotational_transform_axis: float = 0.0
+    rotational_transform_edge: float = 0.0
+    magnetic_well_depth: float = 0.0
+    volume: float = 0.0
+    vmec_converged: bool = True
+    step_number: int = 0
+    budget_remaining: int = 6
+    best_qs_residual: float = float("inf")
+    constraints_satisfied: bool = True
+    target_spec: str = ""
+
+
+class StellaratorState(State):
     initial_qs: float = 0.0
     current_qs: float = 0.0
     prev_qs: float = 0.0

server/app.py

@@ -1,17 +1,46 @@
 from __future__ import annotations
 
-from fastapi import FastAPI
+from openenv.core import create_fastapi_app
 
-from server.environment import TASK, environment_status
+from fusion_lab.models import StellaratorAction, StellaratorObservation
+from server.environment import (
+    ASPECT_RATIO_RANGE,
+    BUDGET,
+    IOTA_EDGE_RANGE,
+    VOLUME_MIN,
+    StellaratorEnvironment,
+)
 
-app = FastAPI(title="Fusion Design Lab")
-
-
-@app.get("/healthz")
-def healthcheck() -> dict[str, str]:
-    return {"status": "ok", "environment": environment_status()}
+app = create_fastapi_app(
+    env=StellaratorEnvironment,
+    action_cls=StellaratorAction,
+    observation_cls=StellaratorObservation,
+)
 
 
 @app.get("/task")
 def task_summary() -> dict[str, object]:
-    return TASK
+    return {
+        "description": "Minimize quasi-symmetry error for a 2-period quasi-helical stellarator.",
+        "constraints": {
+            "aspect_ratio": list(ASPECT_RATIO_RANGE),
+            "rotational_transform_edge": list(IOTA_EDGE_RANGE),
+            "volume_min": VOLUME_MIN,
+        },
+        "budget": BUDGET,
+        "actions": ["run", "submit", "restore_best"],
+        "operators": ["tune_rc10", "tune_rc11", "tune_zs11", "tune_zs12"],
+        "directions": ["increase", "decrease"],
+        "magnitudes": ["small", "medium", "large"],
+        "restart_modes": ["hot", "cold"],
+    }
+
+
+def main() -> None:
+    import uvicorn
+
+    uvicorn.run("server.app:app", host="0.0.0.0", port=8000, reload=True)
+
+
+if __name__ == "__main__":
+    main()

server/environment.py

@@ -1,19 +1,262 @@
 from __future__ import annotations
 
-from typing import Final
-
-TASK: Final[dict[str, object]] = {
-    "description": "Minimize quasi-symmetry error for a 2-period quasi-helical stellarator.",
-    "constraints": {
-        "aspect_ratio": [4.5, 7.0],
-        "rotational_transform_edge": [0.3, 0.6],
-        "volume_min": 0.5,
-    },
-    "budget": 6,
-    "baseline_input": "server/data/input.QH_baseline",
-}
-
-
-def environment_status() -> str:
-    """Return a simple status string until the full environment is implemented."""
-    return "scaffolded"
+from typing import Any, Final, Optional
+
+from openenv.core import Environment as BaseEnvironment
+
+from fusion_lab.models import (
+    StellaratorAction,
+    StellaratorObservation,
+    StellaratorState,
+)
+from server.physics import Diagnostics, PhysicsEngine
+
+BUDGET: Final[int] = 6
+
+ASPECT_RATIO_RANGE: Final[tuple[float, float]] = (4.5, 7.0)
+IOTA_EDGE_RANGE: Final[tuple[float, float]] = (0.3, 0.6)
+VOLUME_MIN: Final[float] = 0.5
+
+TARGET_SPEC: Final[str] = (
+    "Minimize quasi-symmetry residual for a 2-period quasi-helical stellarator. "
+    "Constraints: aspect ratio in [4.5, 7.0], edge iota in [0.3, 0.6], volume > 0.5 m³. "
+    "Budget: 6 evaluations."
+)
+
+
+def check_constraints(diag: Diagnostics) -> bool:
+    ar_lo, ar_hi = ASPECT_RATIO_RANGE
+    iota_lo, iota_hi = IOTA_EDGE_RANGE
+    return (
+        ar_lo <= diag.aspect_ratio <= ar_hi
+        and iota_lo <= diag.iota_edge <= iota_hi
+        and diag.volume >= VOLUME_MIN
+    )
+
+
+class StellaratorEnvironment(
+    BaseEnvironment[StellaratorAction, StellaratorObservation, StellaratorState]
+):
+    def __init__(self) -> None:
+        super().__init__()
+        self._engine = PhysicsEngine()
+        self._state = StellaratorState()
+        self._last_diag: Diagnostics | None = None
+
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        episode_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> StellaratorObservation:
+        diag = self._engine.reset(seed)
+        satisfied = check_constraints(diag)
+        self._state = StellaratorState(
+            episode_id=episode_id,
+            step_count=0,
+            initial_qs=diag.qs_residual,
+            current_qs=diag.qs_residual,
+            prev_qs=diag.qs_residual,
+            best_qs=diag.qs_residual,
+            budget_total=BUDGET,
+            budget_remaining=BUDGET,
+            constraints_satisfied=satisfied,
+        )
+        self._last_diag = diag
+        return self._build_observation(
+            diag, satisfied, action_summary="Episode started. Baseline design loaded."
+        )
+
+    def step(
+        self,
+        action: StellaratorAction,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> StellaratorObservation:
+        self._state.prev_qs = self._state.current_qs
+        self._state.step_count += 1
+
+        if action.intent == "submit":
+            return self._handle_submit()
+        if action.intent == "restore_best":
+            return self._handle_restore()
+        return self._handle_run(action)
+
+    @property
+    def state(self) -> StellaratorState:
+        return self._state
+
+    # ------------------------------------------------------------------
+    # Action handlers
+    # ------------------------------------------------------------------
+
+    def _handle_run(self, action: StellaratorAction) -> StellaratorObservation:
+        if not all([action.operator, action.direction, action.magnitude]):
+            return self._handle_invalid_run()
+
+        self._state.budget_remaining -= 1
+
+        diag = self._engine.modify_and_run(
+            operator=action.operator,
+            direction=action.direction,
+            magnitude=action.magnitude,
+            restart=action.restart or "hot",
+        )
+
+        satisfied = check_constraints(diag) if diag.converged else self._state.constraints_satisfied
+
+        if diag.converged:
+            self._state.current_qs = diag.qs_residual
+            if diag.qs_residual < self._state.best_qs:
+                self._state.best_qs = diag.qs_residual
+            self._state.constraints_satisfied = satisfied
+
+        done = self._state.budget_remaining <= 0
+        reward = self._compute_reward(diag, action.intent, done)
+        summary = self._summary_run(action, diag)
+        self._state.history.append(summary)
+        self._last_diag = diag
+
+        return self._build_observation(
+            diag, satisfied, action_summary=summary, reward=reward, done=done
+        )
+
+    def _handle_submit(self) -> StellaratorObservation:
+        diag = self._last_diag or self._engine.restore_best()
+        satisfied = check_constraints(diag)
+        reward = self._compute_reward(diag, "submit", done=True)
+        summary = self._summary_submit(satisfied)
+        self._state.history.append(summary)
+
+        return self._build_observation(
+            diag, satisfied, action_summary=summary, reward=reward, done=True
+        )
+
+    def _handle_restore(self) -> StellaratorObservation:
+        self._state.budget_remaining -= 1
+
+        diag = self._engine.restore_best()
+        self._state.current_qs = diag.qs_residual
+        satisfied = check_constraints(diag)
+        self._state.constraints_satisfied = satisfied
+
+        done = self._state.budget_remaining <= 0
+        reward = self._compute_reward(diag, "restore_best", done)
+        summary = f"Restored best design. QS residual: {diag.qs_residual:.6f}."
+        self._state.history.append(summary)
+        self._last_diag = diag
+
+        return self._build_observation(
+            diag, satisfied, action_summary=summary, reward=reward, done=done
+        )
+
+    def _handle_invalid_run(self) -> StellaratorObservation:
+        self._state.budget_remaining -= 1
+        diag = self._last_diag or self._engine.restore_best()
+        satisfied = check_constraints(diag)
+        done = self._state.budget_remaining <= 0
+        summary = "Invalid run action: operator, direction, and magnitude are required."
+        self._state.history.append(summary)
+        return self._build_observation(
+            diag, satisfied, action_summary=summary, reward=-1.0, done=done
+        )
+
+    # ------------------------------------------------------------------
+    # Reward V0
+    # ------------------------------------------------------------------
+
+    def _compute_reward(self, diag: Diagnostics, intent: str, done: bool) -> float:
+        reward = 0.0
+
+        if diag.converged and self._state.prev_qs < float("inf"):
+            improvement = self._state.prev_qs - diag.qs_residual
+            reward += improvement * 500.0
+
+        if diag.converged and not check_constraints(diag):
+            reward -= 2.0
+
+        if not diag.converged:
+            reward -= 1.5
+
+        if intent != "submit":
+            reward -= 0.1
+
+        if intent == "submit":
+            if self._state.best_qs < self._state.initial_qs:
+                ratio = 1.0 - (self._state.best_qs / max(self._state.initial_qs, 1e-9))
+                reward += 5.0 * ratio
+                reward += 1.0 * (self._state.budget_remaining / self._state.budget_total)
+            else:
+                reward -= 1.0
+
+        if done and intent != "submit":
+            if self._state.best_qs < self._state.initial_qs:
+                ratio = 1.0 - (self._state.best_qs / max(self._state.initial_qs, 1e-9))
+                reward += 2.0 * ratio
+
+        return round(reward, 4)
+
+    # ------------------------------------------------------------------
+    # Observation builders
+    # ------------------------------------------------------------------
+
+    def _build_observation(
+        self,
+        diag: Diagnostics,
+        satisfied: bool,
+        action_summary: str,
+        reward: float | None = None,
+        done: bool = False,
+    ) -> StellaratorObservation:
+        text_lines = [
+            action_summary,
+            "",
+            f"QS Residual: {diag.qs_residual:.6f}  |  Best: {self._state.best_qs:.6f}",
+            f"Aspect Ratio: {diag.aspect_ratio:.4f}  [4.5, 7.0]",
+            f"Edge Iota: {diag.iota_edge:.4f}  [0.3, 0.6]",
+            f"Volume: {diag.volume:.4f} m³  (min 0.5)",
+            f"Magnetic Well: {diag.magnetic_well_depth:.4f}",
+            f"VMEC Converged: {diag.converged}",
+            f"Constraints: {'SATISFIED' if satisfied else 'VIOLATED'}",
+            f"Step: {self._state.step_count}  |  Budget: {self._state.budget_remaining}/{self._state.budget_total}",
+        ]
+
+        return StellaratorObservation(
+            diagnostics_text="\n".join(text_lines),
+            quasi_symmetry_residual=diag.qs_residual,
+            aspect_ratio=diag.aspect_ratio,
+            rotational_transform_axis=diag.iota_axis,
+            rotational_transform_edge=diag.iota_edge,
+            magnetic_well_depth=diag.magnetic_well_depth,
+            volume=diag.volume,
+            vmec_converged=diag.converged,
+            step_number=self._state.step_count,
+            budget_remaining=self._state.budget_remaining,
+            best_qs_residual=self._state.best_qs,
+            constraints_satisfied=satisfied,
+            target_spec=TARGET_SPEC,
+            reward=reward,
+            done=done,
+        )
+
+    # ------------------------------------------------------------------
+    # Action summaries
+    # ------------------------------------------------------------------
+
+    def _summary_run(self, action: StellaratorAction, diag: Diagnostics) -> str:
+        restart_note = f" ({action.restart} restart)" if action.restart else ""
+        header = f"Applied {action.operator} {action.direction} {action.magnitude}{restart_note}."
+
+        if diag.converged:
+            delta = self._state.prev_qs - diag.qs_residual
+            direction = "improved" if delta > 0 else "worsened" if delta < 0 else "unchanged"
+            return f"{header} VMEC converged. QS {direction}: {self._state.prev_qs:.6f} -> {diag.qs_residual:.6f}."
+        return f"{header} VMEC failed to converge. Change reverted."
+
+    def _summary_submit(self, satisfied: bool) -> str:
+        status = "Constraints satisfied." if satisfied else "Constraints VIOLATED."
+        improvement = self._state.initial_qs - self._state.best_qs
+        return (
+            f"Design submitted. Best QS residual: {self._state.best_qs:.6f} "
+            f"(improved by {improvement:.6f} from initial). {status}"
+        )

server/physics.py

@@ -1,20 +1,141 @@
 from __future__ import annotations
 
+import math
+import random
+from dataclasses import dataclass, field
+from typing import Final
 
+NFP: Final[int] = 2
+
+BASELINE_COEFFS: Final[dict[str, float]] = {
+    "rc10": 1.0,
+    "rc11": 0.12,
+    "zs11": 0.12,
+    "zs12": -0.02,
+}
+
+OPTIMAL_COEFFS: Final[dict[str, float]] = {
+    "rc10": 1.02,
+    "rc11": 0.135,
+    "zs11": 0.115,
+    "zs12": -0.035,
+}
+
+MAGNITUDE_DELTAS: Final[dict[str, float]] = {
+    "small": 0.005,
+    "medium": 0.02,
+    "large": 0.05,
+}
+
+
+@dataclass(frozen=True)
+class Diagnostics:
+    qs_residual: float
+    aspect_ratio: float
+    iota_axis: float
+    iota_edge: float
+    volume: float
+    magnetic_well_depth: float
+    converged: bool
+
+
+@dataclass
 class PhysicsEngine:
-    """Placeholder for the VMEC-backed physics loop.
-
-    The next implementation step should make this the single place that:
-    - loads the baseline input
-    - applies discrete coefficient updates
-    - runs the solver
-    - computes diagnostics
-    - tracks best-known designs
-    """
-
-    def __init__(self) -> None:
-        self._status = "unimplemented"
-
-    @property
-    def status(self) -> str:
-        return self._status
+    coeffs: dict[str, float] = field(default_factory=lambda: dict(BASELINE_COEFFS))
+    best_coeffs: dict[str, float] = field(default_factory=lambda: dict(BASELINE_COEFFS))
+    best_qs: float = float("inf")
+    _rng: random.Random = field(default_factory=random.Random)
+
+    def reset(self, seed: int | None = None) -> Diagnostics:
+        self.coeffs = dict(BASELINE_COEFFS)
+        self._rng = random.Random(seed)
+        if seed is not None:
+            for key in self.coeffs:
+                self.coeffs[key] += self._rng.gauss(0, 0.002)
+        self.best_coeffs = dict(self.coeffs)
+        diag = self._compute_diagnostics(converged=True)
+        self.best_qs = diag.qs_residual
+        return diag
+
+    def modify_and_run(
+        self,
+        operator: str,
+        direction: str,
+        magnitude: str,
+        restart: str,
+    ) -> Diagnostics:
+        coeff_key = operator.removeprefix("tune_")
+        delta = MAGNITUDE_DELTAS[magnitude]
+        if direction == "decrease":
+            delta = -delta
+
+        prev_value = self.coeffs[coeff_key]
+        self.coeffs[coeff_key] = prev_value + delta
+
+        converged = self._simulate_convergence(magnitude, restart)
+        if not converged:
+            self.coeffs[coeff_key] = prev_value
+            return self._compute_diagnostics(converged=False)
+
+        diag = self._compute_diagnostics(converged=True)
+        if diag.qs_residual < self.best_qs:
+            self.best_qs = diag.qs_residual
+            self.best_coeffs = dict(self.coeffs)
+        return diag
+
+    def restore_best(self) -> Diagnostics:
+        self.coeffs = dict(self.best_coeffs)
+        return self._compute_diagnostics(converged=True)
+
+    def _compute_diagnostics(self, *, converged: bool) -> Diagnostics:
+        rc10 = self.coeffs["rc10"]
+        rc11 = self.coeffs["rc11"]
+        zs11 = self.coeffs["zs11"]
+        zs12 = self.coeffs["zs12"]
+
+        r_minor = math.sqrt(rc11**2 + zs11**2)
+        aspect_ratio = rc10 / max(r_minor, 1e-6)
+        volume = 2.0 * math.pi**2 * rc10 * r_minor**2
+
+        helical_excursion = abs(zs11 / max(abs(rc11), 1e-6))
+        iota_axis = 0.35 + 0.15 * helical_excursion + 0.5 * abs(zs12)
+        shear = 0.04 + 0.02 * abs(rc10 - 1.0)
+        iota_edge = iota_axis + shear
+
+        magnetic_well = 0.02 + 0.01 * (rc11 / max(abs(zs11), 1e-6) - 1.0)
+
+        qs_residual = self._compute_qs_residual() if converged else float("inf")
+
+        return Diagnostics(
+            qs_residual=round(qs_residual, 6),
+            aspect_ratio=round(aspect_ratio, 4),
+            iota_axis=round(iota_axis, 4),
+            iota_edge=round(iota_edge, 4),
+            volume=round(volume, 4),
+            magnetic_well_depth=round(magnetic_well, 4),
+            converged=converged,
+        )
+
+    def _compute_qs_residual(self) -> float:
+        d = {k: self.coeffs[k] - OPTIMAL_COEFFS[k] for k in OPTIMAL_COEFFS}
+        quadratic = (
+            2.0 * d["rc10"] ** 2
+            + 8.0 * d["rc11"] ** 2
+            + 8.0 * d["zs11"] ** 2
+            + 15.0 * d["zs12"] ** 2
+        )
+        cross = 4.0 * d["rc11"] * d["zs11"] - 3.0 * d["rc10"] * d["zs12"]
+        noise = self._rng.gauss(0, 0.0003)
+        return max(quadratic + cross + 0.002 + noise, 0.001)
+
+    def _simulate_convergence(self, magnitude: str, restart: str) -> bool:
+        fail_prob = {"small": 0.02, "medium": 0.08, "large": 0.20}[magnitude]
+        if restart == "hot":
+            fail_prob *= 0.5
+        for key, val in self.coeffs.items():
+            deviation = abs(val - BASELINE_COEFFS[key])
+            if deviation > 0.1:
+                fail_prob += 0.15
+            elif deviation > 0.05:
+                fail_prob += 0.05
+        return self._rng.random() > min(fail_prob, 0.8)