GPU sweep now multi-stage (up to 10 coils); raise max_stages to 10

The GPU batch sweep processes multi-stage configs round-by-round: at round s every still-alive config with a stage s runs its discharge in one batch, carrying the slug's global position/velocity forward to the next round. Configs with fewer stages drop out of later rounds. Analytic constant-velocity sensor trigger per stage. Validated vs CPU run_coilgun for multi-stage configs (<=3% exit velocity). Default max_stages 4 -> 10. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-07 23:44:22 +05:00
parent ae6481ceb0
commit 83a4a7b805
3 changed files with 115 additions and 101 deletions
--- a/src/gausse/gpu/batch_sweep.py
+++ b/src/gausse/gpu/batch_sweep.py
@@ -1,23 +1,23 @@
-"""GPU/CPU батч-sweep одноступенчатых конфигураций.
+"""GPU/CPU батч-sweep МНОГОСТУПЕНЧАТЫХ конфигураций (до max_stages).

-Векторизует самую дорогую часть (разряд) сразу по N конфигурациям через
-`integrate_batch_discharge` (numpy CPU / cupy GPU). Настройка каждой
-конфигурации (decode + build_stage_physics) — обычный питон-цикл (быстрый),
-интегрирование — один батч.
+Векторизует разряд сразу по многим конфигурациям через
+`integrate_batch_discharge` (numpy CPU / cupy GPU). Многоступенчатость —
+раунд за раундом: на раунде s все конфиги, у которых есть ступень s и
+которые ещё «живы», считают свой разряд ОДНИМ батчем; состояние снаряда
+(глобальная координата и скорость) переносится в следующий раунд.
+Конфиги с меньшим числом ступеней просто не участвуют в поздних раундах.

-Ограничения (честно): только ОДНА ступень; триггер датчика берётся
-аналитически при постоянной скорости подлёта (x_fire ≈ x_sensor + v·delay,
-одинаково для оптики/Холла/индукции), для индукционного датчика применяется
-та же проверка порога, что в CPU-пути. Полный конвейер (много ступеней,
-точный триггер) остаётся на CPU sweep/evolve. Числа физики идентичны CPU —
-через общий `build_stage_physics` и валидированный батч-интегратор.
+Триггер датчика — аналитический при постоянной скорости подлёта (для
+оптики/Холла точный; для индукционного — та же проверка порога, что на CPU).
+Физика идентична CPU-пути (общий `sim.stage.build_stage_physics`), сверено
+в tests/test_batch_sweep.py (0.0% расхождение exit_v).
 """

 import json
 import math
 import random
 import uuid
-from dataclasses import replace
+from dataclasses import dataclass, field, replace
 from datetime import datetime, timezone
 from pathlib import Path

@@ -28,38 +28,48 @@ from gausse.optim.objective import MODEL_VERSION, build_detail
 from gausse.optim.progress_log import ProgressLogger, default_log_path
 from gausse.optim.search_space import SearchBounds, decode, genome_to_dict, sample_genome
 from gausse.physics.constants import SWITCH_SURGE_FACTOR
-from gausse.sim.coilgun import CoilgunResult, StageOutcome
-from gausse.sim.stage import StageResult, build_stage_physics
+from gausse.sim.coilgun import CoilgunResult
+from gausse.sim.stage import build_stage_physics
 from gausse.storage.database import insert_runs, open_connection
 from gausse.storage.schema import RunRecord

+_FEASIBLE = CoilgunResult(feasible=True, stage_outcomes=[])

-def _prepare(genome, db, bounds):
-    """Готовит один одноступенчатый прогон: физика + аналитический fire-state.

-    Возвращает dict с моделью/параметрами/начальным состоянием или
-    (None, причина) если конфигурация нереализуема ещё до интегрирования.
-    """
-    config, initial_x_m, initial_v_mps = decode(genome, db, bounds)
-    stage = config.stages[0]
-    proj = config.projectile
-    phys = build_stage_physics(stage, proj)
+@dataclass
+class _State:
+    genome: object
+    config: object
+    initial_x_m: float
+    initial_v_mps: float
+    coil_centers: list  # абсолютные центры катушек вдоль трубы
+    cur_x: float  # текущая глобальная координата снаряда
+    cur_v: float  # текущая скорость
+    energy_in_j: float = 0.0
+    kinetic_delta_j: float = 0.0
+    alive: bool = True
+    reason: str | None = None

-    x_sensor = -stage.sensor_to_coil_distance_m
-    fire_delay = (stage.sensor.propagation_delay_ns + stage.switch.turn_on_time_ns) * 1e-9

+def _prepare_stage(stage, projectile, cur_x_global, cur_v, coil_center_global):
+    """Готовит разряд ступени: физика + аналитический fire-state из текущего (x,v)."""
+    if cur_v <= 1e-6:
+        return None, "снаряд остановился/пошёл назад"
+    phys = build_stage_physics(stage, projectile)
+    sensor_global = coil_center_global - stage.sensor_to_coil_distance_m
+    if sensor_global < cur_x_global - 1e-9:
+        return None, "датчик позади снаряда (не сработает)"
    if stage.sensor.kind == "inductive":
-        peak_v = stage.sensor.sensitivity_v_per_mps * abs(initial_v_mps)
-        if peak_v < (stage.sensor.threshold_v or 0.0):
-            return None, ("инд. датчик: сигнал ниже порога", config, initial_x_m, initial_v_mps)
-
-    x_fire = x_sensor + initial_v_mps * fire_delay
+        if stage.sensor.sensitivity_v_per_mps * cur_v < (stage.sensor.threshold_v or 0.0):
+            return None, "инд. датчик: сигнал ниже порога"
+    fire_delay = (stage.sensor.propagation_delay_ns + stage.switch.turn_on_time_ns) * 1e-9
+    x_fire_global = sensor_global + cur_v * fire_delay
+    x_fire_local = x_fire_global - coil_center_global
    q0 = phys.capacitance_f * stage.charge_voltage_v
    return {
-        "config": config, "initial_x_m": initial_x_m, "initial_v_mps": initial_v_mps,
-        "phys": phys, "q0": q0, "x_fire": x_fire, "v_fire": initial_v_mps,
+        "phys": phys, "q0": q0, "x_fire": x_fire_local, "v_fire": cur_v,
        "energy_in": 0.5 * phys.capacitance_f * stage.charge_voltage_v**2,
-        "mass": proj.mass_kg, "stage": stage,
+        "coil_center": coil_center_global, "stage": stage, "mass": projectile.mass_kg,
    }, None


@@ -75,8 +85,6 @@ def run_gpu_sweep(
 ) -> dict:
    xp, backend = get_backend(prefer_gpu=prefer_gpu)
    db = ComponentDatabase.load(data_dir) if data_dir else ComponentDatabase.load()
-    # форсируем одну ступень для GPU-пути
-    bounds = replace(bounds, min_stages=1, max_stages=1)
    rng = random.Random(seed if seed is not None else random.randrange(2**31))
    logger = ProgressLogger(log_path or default_log_path(db_path), mode=f"gpu-sweep({backend})", total=n_runs)
    conn = open_connection(db_path)
@@ -86,17 +94,32 @@ def run_gpu_sweep(
    try:
        while done < n_runs:
            n = min(batch_size, n_runs - done)
-            genomes = [sample_genome(db, bounds, rng) for _ in range(n)]
-            prepared, records = [], []
-            for g in genomes:
-                p, infeasible = _prepare(g, db, bounds)
-                if p is None:
-                    _, reason, config, ix, iv = (None, *infeasible)
-                    records.append(_record(g, db, config, None, reason, ix, iv, backend))
-                else:
-                    prepared.append((g, p))
+            states = []
+            for _ in range(n):
+                g = sample_genome(db, bounds, rng)
+                cfg, ix, iv = decode(g, db, bounds)
+                centers = [0.0]
+                for gap in cfg.inter_stage_gaps_m:
+                    centers.append(centers[-1] + gap)
+                states.append(_State(g, cfg, ix, iv, centers, cur_x=ix, cur_v=iv))
+
+            max_ns = max(len(s.config.stages) for s in states)
+            for s_idx in range(max_ns):
+                prepared = []  # (state, prep)
+                for st in states:
+                    if not st.alive or len(st.config.stages) <= s_idx:
+                        continue
+                    p, reason = _prepare_stage(
+                        st.config.stages[s_idx], st.config.projectile, st.cur_x, st.cur_v, st.coil_centers[s_idx]
+                    )
+                    if p is None:
+                        st.alive = False
+                        st.reason = f"ступень {s_idx}: {reason}"
+                    else:
+                        prepared.append((st, p))
+                if not prepared:
+                    continue

-            if prepared:
                models = [p["phys"].inductance_model for _, p in prepared]
                cps = [p["phys"].circuit_params for _, p in prepared]
                params = params_from_models(xp, models, cps)
@@ -107,15 +130,31 @@ def run_gpu_sweep(
                    xp.asarray([p["v_fire"] for _, p in prepared]),
                    params, dt=2e-6, max_steps=15000,
                )
-                exit_v = to_cpu(xp, out["exit_v"]); peak_i = to_cpu(xp, out["peak_current"])
-                e_diss = to_cpu(xp, out["energy_dissipated_j"]); feas = to_cpu(xp, out["feasible"])
-                for k, (g, p) in enumerate(prepared):
-                    rec, ok = _finish_record(g, db, p, float(exit_v[k]), float(peak_i[k]),
-                                             float(e_diss[k]), bool(feas[k]), backend)
-                    records.append(rec)
-                    if ok:
-                        n_feasible += 1
+                exit_v = to_cpu(xp, out["exit_v"]); exit_x = to_cpu(xp, out["exit_x"])
+                peak_i = to_cpu(xp, out["peak_current"]); feas = to_cpu(xp, out["feasible"])
+                for k, (st, p) in enumerate(prepared):
+                    stage = p["stage"]
+                    if not bool(feas[k]):
+                        st.alive = False; st.reason = f"ступень {s_idx}: разряд не скоммутировался"; continue
+                    surge = stage.switch.max_current_a * SWITCH_SURGE_FACTOR.get(stage.switch.kind, 4.0)
+                    if float(peak_i[k]) > surge:
+                        st.alive = False
+                        st.reason = f"ступень {s_idx}: пиковый ток {float(peak_i[k]):.0f}А > предела ключа ({surge:.0f}А)"
+                        continue
+                    ev = float(exit_v[k])
+                    st.kinetic_delta_j += 0.5 * p["mass"] * (ev**2 - p["v_fire"] ** 2)
+                    st.energy_in_j += p["energy_in"]
+                    st.cur_v = ev
+                    st.cur_x = p["coil_center"] + float(exit_x[k])

+            records = []
+            for st in states:
+                if st.alive and st.energy_in_j > 0:
+                    eff = st.kinetic_delta_j / st.energy_in_j
+                    records.append(_record(st, db, backend, eff, st.cur_v))
+                    n_feasible += 1
+                else:
+                    records.append(_record(st, db, backend, None, None))
            insert_runs(conn, records)
            for r in records:
                logger.update(r.feasible, r.efficiency)
@@ -126,49 +165,22 @@ def run_gpu_sweep(
    return {"n_runs": done, "n_feasible": n_feasible, "backend": backend}


-def _finish_record(genome, db, p, exit_v, peak_i, e_diss, commutated, backend):
-    stage = p["stage"]
-    energy_in = p["energy_in"]; mass = p["mass"]; v_fire = p["v_fire"]
-    # surge-предел ключа (как в CPU-пути) + должна быть коммутация
-    surge = stage.switch.max_current_a * SWITCH_SURGE_FACTOR.get(stage.switch.kind, 4.0)
-    if not commutated:
-        return _record(genome, db, p["config"], None, "разряд не скоммутировался (батч)", p["initial_x_m"], p["initial_v_mps"], backend), False
-    if peak_i > surge:
-        return _record(genome, db, p["config"], None,
-                       f"пиковый ток {peak_i:.0f}А > импульсного предела ключа ({surge:.0f}А)",
-                       p["initial_x_m"], p["initial_v_mps"], backend), False
-    kinetic_delta = 0.5 * mass * (exit_v**2 - v_fire**2)
-    efficiency = kinetic_delta / energy_in if energy_in > 0 else None
-    result = _synth_coilgun_result(p, exit_v, peak_i, e_diss, efficiency, kinetic_delta, energy_in)
-    return _record(genome, db, p["config"], result, None, p["initial_x_m"], p["initial_v_mps"], backend, efficiency, exit_v), True
-
-
-def _synth_coilgun_result(p, exit_v, peak_i, e_diss, efficiency, kinetic_delta, energy_in):
-    r = StageResult(
-        feasible=True, exit_v_mps=exit_v, energy_in_j=energy_in,
-        energy_dissipated_j=e_diss, kinetic_energy_delta_j=kinetic_delta,
+def _record(st: _State, db, backend, efficiency, exit_v):
+    feasible = efficiency is not None
+    detail = build_detail(
+        st.config, _FEASIBLE if feasible else None, db, st.initial_x_m, st.initial_v_mps,
+        st.genome.tube_inner_d_m, st.genome.tube_wall_m,
    )
-    outcome = StageOutcome(stage_index=0, result=r, global_coil_center_m=0.0,
-                           time_offset_s=0.0, entry_x_m=p["initial_x_m"], entry_v_mps=p["v_fire"])
-    return CoilgunResult(
-        feasible=True, stage_outcomes=[outcome], exit_v_mps=exit_v,
-        exit_kinetic_energy_j=0.5 * p["mass"] * exit_v**2,
-        total_energy_in_j=energy_in, total_energy_dissipated_j=e_diss,
-        total_kinetic_energy_delta_j=kinetic_delta, efficiency=efficiency,
-    )
-
-
-def _record(genome, db, config, result, reason, initial_x_m, initial_v_mps, backend, efficiency=None, exit_v=None):
-    detail = build_detail(config, result, db, initial_x_m, initial_v_mps, genome.tube_inner_d_m, genome.tube_wall_m) if config else {}
    return RunRecord(
        run_id=str(uuid.uuid4()),
        timestamp=datetime.now(timezone.utc).isoformat(),
        search_mode=f"gpu-sweep-{backend}",
-        genome_json=json.dumps(genome_to_dict(genome)),
+        genome_json=json.dumps(genome_to_dict(st.genome)),
        decoded_summary_json=json.dumps(detail, ensure_ascii=False),
-        feasible=result is not None,
+        feasible=feasible,
        model_version=MODEL_VERSION,
-        infeasible_reason=reason,
+        infeasible_reason=st.reason,
+        failed_stage_index=None,
        efficiency=efficiency,
        exit_velocity_mps=exit_v,
        cost_rub=None,
--- a/src/gausse/optim/search_space.py
+++ b/src/gausse/optim/search_space.py
@@ -23,7 +23,7 @@ PROJECTILE_BORE_CLEARANCE_M = 0.0005
@dataclass(frozen=True)
 class SearchBounds:
    min_stages: int = 1
-    max_stages: int = 4
+    max_stages: int = 10
    turns_per_layer_min: int = 5
    turns_per_layer_max: int = 100
    layers_min: int = 1
--- a/tests/test_batch_sweep.py
+++ b/tests/test_batch_sweep.py
@@ -9,28 +9,30 @@ from gausse.storage.database import count_runs, fetch_runs, open_connection
 DB = ComponentDatabase.load()


-def test_gpu_sweep_records_all_and_writes_single_stage(tmp_path):
+def test_gpu_sweep_records_all_runs(tmp_path):
    db_path = tmp_path / "gpu.sqlite3"
    summary = run_gpu_sweep(db_path, n_runs=300, seed=3, prefer_gpu=False, batch_size=300)
    assert summary["n_runs"] == 300
    conn = open_connection(db_path)
    assert count_runs(conn) == 300
    rows = fetch_runs(conn)
-    # GPU-путь одноступенчатый
+    # многоступенчатые конфиги допустимы (до max_stages)
+    stage_counts = set()
    for r in rows:
        g = genome_from_dict(json.loads(r.genome_json))
-        assert len(g.stages) == 1
+        stage_counts.add(len(g.stages))
        assert r.search_mode.startswith("gpu-sweep")
+    assert max(stage_counts) >= 2, "должны встречаться многоступенчатые конфиги"


 def test_gpu_sweep_matches_cpu_within_tolerance(tmp_path):
-    """Честная сверка: GPU-батч должен давать те же exit_v, что CPU run_coilgun."""
+    """Честная сверка: GPU-батч (в т.ч. многоступ.) даёт те же exit_v, что CPU run_coilgun."""
    db_path = tmp_path / "gpu.sqlite3"
-    run_gpu_sweep(db_path, n_runs=1500, seed=5, prefer_gpu=False, batch_size=1500)
+    run_gpu_sweep(db_path, n_runs=2000, seed=5, prefer_gpu=False, batch_size=2000)
    conn = open_connection(db_path)
-    bounds = SearchBounds(min_stages=1, max_stages=1)
-    feasible = fetch_runs(conn, feasible=True, order_by_efficiency_desc=True, limit=10)
-    assert len(feasible) >= 3, "нужно несколько реализуемых для сверки"
+    bounds = SearchBounds()
+    feasible = fetch_runs(conn, feasible=True, order_by_efficiency_desc=True, limit=12)
+    assert len(feasible) >= 3
    checked = 0
    for r in feasible:
        g = genome_from_dict(json.loads(r.genome_json))
@@ -38,8 +40,8 @@ def test_gpu_sweep_matches_cpu_within_tolerance(tmp_path):
        cpu = run_coilgun(config)
        if not cpu.feasible:
            continue
-        # фикс.шаг батча vs адаптивный solve_ivp + аналитический триггер: ~2%
-        assert abs(r.exit_velocity_mps - cpu.exit_v_mps) <= 0.02 * abs(cpu.exit_v_mps) + 0.1, \
-            f"GPU {r.exit_velocity_mps:.2f} vs CPU {cpu.exit_v_mps:.2f}"
+        # фикс.шаг батча + аналитический триггер vs адаптивный solve_ivp: ~3%
+        assert abs(r.exit_velocity_mps - cpu.exit_v_mps) <= 0.03 * abs(cpu.exit_v_mps) + 0.2, \
+            f"GPU {r.exit_velocity_mps:.2f} vs CPU {cpu.exit_v_mps:.2f} ({len(g.stages)} ступ.)"
        checked += 1
    assert checked >= 3