Add continuous OFF ON OFF VIV video research driver

research/continuous_viv/run.py

@@ -0,0 +1,222 @@
+"""Continuous OFF -> ON -> OFF VIV illustration; docs/continuous_viv_video.md.
+
+Each invocation owns ONE Simulation and ONE oscillator. Rendering never advances
+the solver or changes the policy clock. No topology changes or moving restart.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+import subprocess
+import sys
+import time
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT))
+from research.karman_viv import run as base
+from research.continuous_viv.render import Video, frame_schedule
+
+D, U, SCALE, SI = 40., .005, 2., 3200
+PHASES = ('off_before', 'on', 'off_after')
+STEADY_WALL = np.array([0., -5.1, 5.1]) * U
+
+
+def timing(scene, smoke):
+    if smoke:
+        return 1,32,32
+    # Prior video's steady interval was 1000 at D20; diffusive D40 uses 4000.
+    return (320,4000,2048000) if scene=='steady' else (400,SI,2048000)
+
+
+def command(scene, phase, obs, model):
+    """Physical surface speeds; OFF is not a zero normalized PPO action."""
+    if phase != 'on':
+        return np.zeros(3), np.array([0., .5, -.5])
+    if scene == 'steady':
+        return STEADY_WALL.copy(), np.full(3, np.nan)
+    action = np.asarray(model.predict(obs, deterministic=True)[0], np.float32)
+    if action.shape != (3,) or not np.isfinite(action).all() or (np.abs(action)>1).any():
+        raise RuntimeError('invalid PPO action')
+    return base.wall_target(action, SCALE), action
+
+
+def add_scene(sim, scene):
+    if scene == 'karman':
+        return base.add_scene(sim, SCALE, moving=True)
+    from CelerisLab.body.motion import Translation
+    probes = [sim.add_body('sensor', c, 10.) for c in base.scaled_points(base.PROBES, SCALE)]
+    envelope = Translation((0., 0., -D, D), .025, shared=True)
+    bodies = [sim.add_body('circle', c, D/2, translation=envelope)
+              for c in base.scaled_points(base.CENTRES, SCALE)]
+    return probes, bodies
+
+
+def evolve(sim, probes, bodies, structure, force, *, scene, model, norm,
+           intervals, si, out, render, start, wall_cap):
+    """One native-step loop across all phases; completed SI feeds the next SI."""
+    from CelerisLab.body.motion import MovingTables
+    cols = {k:i for i,k in enumerate(MovingTables.DIAG)}
+    count = len(PHASES)*intervals*si
+    rows = np.lib.format.open_memmap(out/'steps.npy', mode='w+', dtype=np.float64,
+                                   shape=(count, len(base.STEP_COLUMNS)))
+    # The completed count is authoritative after any interrupted run.
+    completed = 0
+    obs = np.zeros(12, np.float32)
+    signals, switches = [], []
+    origin_clock = sim.stepper.step_count
+    status = 'running'
+    try:
+        for phase in PHASES:
+            switches.append(dict(phase=phase, step=completed, absolute_clock=sim.stepper.step_count,
+                                 y=structure.y, v=structure.v, previous_rhs=force,
+                                 input_observation=obs.tolist()))
+            for i in range(intervals):
+                target, action = command(scene, phase, obs, model)
+                # Match the frozen inference contract: reset the EMA every SI.
+                eff = np.zeros(3)
+                sensor_sum, force_sum = np.zeros(6), np.zeros(6)
+                for k in range(si):
+                    if k % 1000 == 0:
+                        base.check_time(start, wall_cap)
+                    y1, v1 = structure.candidate(force)
+                    if not np.isfinite([y1,v1]).all() or abs(y1)>D:
+                        raise RuntimeError('nonfinite structure or |y| > D; no retuning')
+                    eff = base.ema(eff, target, SCALE)
+                    omega = eff/(D/2)
+                    for bid,w in zip(bodies,omega):
+                        sim.set_body(bid, omega=float(w), offset=(0.,y1), vy=y1-structure.y)
+                    sim.run(1, zero_obs=True, sync_obs=True)
+                    d = sim.moving_diagnostics()
+                    cme,wen = d[:,cols['cme_y']], d[:,cols['wen_y']]
+                    node = d[:,cols['cover_jy']]-d[:,cols['uncover_jy']]
+                    row = (completed+1,y1,v1,*cme,*wen,*node,
+                           d[:,cols['eq_only']].sum(),d[:,cols['no_donor']].sum(),
+                           d[:,cols['crossings']].sum(),*omega)
+                    if not np.isfinite(row).all() or row[13] != 0:
+                        raise RuntimeError('nonfinite moving diagnostic or no_donor')
+                    f=np.array([v for p in bodies for v in sim.read_force(p)[:2]])
+                    s=np.array([v for p in probes for v in sim.read_sensor(p)[:2]])
+                    if not np.isfinite([f,s]).all():
+                        raise RuntimeError('nonfinite force/sensor observation')
+                    sensor_sum += s
+                    force_sum += f
+                    force=float(cme.sum(dtype=np.float64)+wen.sum(dtype=np.float64))
+                    rows[completed]=row
+                    completed += 1
+                    structure.y,structure.v=y1,v1
+                    if sim.stepper.step_count != origin_clock+completed:
+                        raise RuntimeError('noncontinuous solver clock')
+                    render(sim, completed, phase, eff, rows[:completed])
+                if scene == 'karman':
+                    raw, obs, clipped = base.observation(sensor_sum/si,force_sum/si,SCALE,norm)
+                else:
+                    # Steady uses prescribed rotations, not a model or borrowed norm.
+                    raw=np.r_[sensor_sum/si,force_sum/si]
+                    clipped=0
+                signals.append(dict(phase=phase, step=completed, raw=raw.tolist(),
+                                    observation=obs.tolist(), clipped=clipped,
+                                    action=action.tolist() if np.isfinite(action).all() else None,
+                                    target_surface_speed=target.tolist()))
+                if (i+1)%10==0 or i+1==intervals:
+                    rows.flush()
+                    progress=dict(phase=phase, interval=i+1, completed_steps=completed,
+                                  tU_D=completed*U/D, y_D=structure.y/D,
+                                  elapsed_s=time.monotonic()-start)
+                    (out/'progress.json').write_text(json.dumps(progress)+'\n')
+                    print(json.dumps(progress),flush=True)
+        status='duration_completed_not_physics_qualified'
+    except BaseException as exc:
+        status=f'error: {type(exc).__name__}: {exc}'
+        raise
+    finally:
+        rows.flush()
+        (out/'signals.json').write_text(json.dumps(signals,allow_nan=False)+'\n')
+        report=dict(status=status,completed_steps=completed,origin_clock=origin_clock,
+                    final_clock=sim.stepper.step_count,switches=switches,
+                    elapsed_s=time.monotonic()-start,columns=base.STEP_COLUMNS,
+                    M=structure.mass,K=structure.stiffness,fn=structure.frequency,damping=0.,
+                    max_abs_y_D=float(np.abs(rows[:completed,1]).max()/D) if completed else None)
+        (out/'report.json').write_text(json.dumps(report,indent=2,allow_nan=False)+'\n')
+
+
+def main():
+    ap=argparse.ArgumentParser(description=__doc__)
+    ap.add_argument('--scene',choices=('steady','karman'),required=True)
+    ap.add_argument('--out',type=Path,required=True)
+    ap.add_argument('--model',type=Path)
+    ap.add_argument('--norm',type=Path)
+    ap.add_argument('--release-dir',type=Path)
+    ap.add_argument('--smoke',action='store_true',help='3 x 32 steps, all switches, no physical result')
+    ap.add_argument('--wall-seconds',type=float,default=10800.)
+    args=ap.parse_args()
+    if not np.isfinite(args.wall_seconds) or args.wall_seconds<=0:
+        ap.error('finite positive wall cap required')
+    if args.scene=='karman':
+        if args.model is None or args.norm is None:
+            ap.error('Karman requires the frozen model and norm')
+        if base.sha(args.model)!=base.MODEL_SHA or base.sha(args.norm)!=base.NORM_SHA:
+            ap.error('frozen Karman model/normalizer hash mismatch')
+        if not args.smoke and args.release_dir is None:
+            ap.error('physical Karman requires its verified static release')
+    elif any(p is not None for p in (args.model,args.norm,args.release_dir)):
+        ap.error('steady uses constant control and its own warmup, no borrowed policy or release')
+    out=args.out.resolve()
+    if out==ROOT or ROOT in out.parents:
+        ap.error('output must be outside checkout')
+    out.mkdir(parents=True,exist_ok=False)
+    intervals,si,warm=timing(args.scene,args.smoke)
+    cfgpath=out/'config.json'
+    cfgpath.write_text(json.dumps(base.config(40),indent=2)+'\n')
+    source=subprocess.check_output(['git','-C',str(ROOT),'rev-parse','HEAD'],text=True).strip()
+    case=dict(source=source,scene=args.scene,smoke=args.smoke,D=D,U=U,nu=.004,
+              phases=PHASES,SI=si,intervals_per_phase=intervals,
+              phase_duration_tU_D=intervals*si*U/D, mass_ratio=4.,fn_D_U=.2,zeta=0.,
+              first_ppo_observation='previous complete OFF interval; no reset at switches',
+              rhs='previous native step CME_y+Wen_y; no filtering; policy uses SI-mean CME',
+              control='fixed (0,-5.1,+5.1)U' if args.scene=='steady' else 'frozen karman_re100 PPO; actual Re_D50',
+              model_sha256=base.sha(args.model) if args.model else None,
+              norm_sha256=base.sha(args.norm) if args.norm else None,
+              config_sha256=base.sha(cfgpath),render='sparse post-step snapshots only; no step skipping')
+    model=base.load_policy(args.model) if args.model else None
+    norm=json.loads(args.norm.read_text()) if args.norm else None
+    seed=None
+    if args.scene=='karman' and not args.smoke:
+        seed,force=base.saved_release(args.release_dir)
+        case['release_sha256']=base.sha(args.release_dir/'release.npz')
+    (out/'case.json').write_text(json.dumps(case,indent=2)+'\n')
+    from CelerisLab import Simulation
+    from CelerisLab.body.motion import MovingTables
+    started=time.monotonic()
+    structure=base.Structure(4.,2.5e-5,0.,D)
+    video=None
+    try:
+        with Simulation(str(cfgpath)) as sim:
+            probes,bodies=add_scene(sim,args.scene)
+            sim.initialize()
+            if seed is not None:
+                case['release_transfer']=base.install_seed(sim,seed)
+                del seed
+            else:
+                base.warmup(sim,warm-1,started,args.wall_seconds)
+                sim.run(1)
+                d=sim.moving_diagnostics()
+                force=float(d[:,MovingTables.DIAG.index('cme_y')].sum(dtype=np.float64)
+                            +d[:,MovingTables.DIAG.index('wen_y')].sum(dtype=np.float64))
+            case.update(warmup_executed=0 if 'release_transfer' in case else warm,
+                        origin_clock=sim.stepper.step_count)
+            (out/'case.json').write_text(json.dumps(case,indent=2)+'\n')
+            phase_steps=intervals*si
+            video=Video(out,args.scene,bodies,phase_steps,smoke=args.smoke,source=source)
+            evolve(sim,probes,bodies,structure,force,scene=args.scene,model=model,norm=norm,
+                   intervals=intervals,si=si,out=out,render=video.capture,
+                   start=started,wall_cap=args.wall_seconds)
+    finally:
+        if video is not None:
+            video.close()
+
+
+if __name__=='__main__':
+    main()