Add bounded paired Karman pinball VIV research driver

research/karman_viv/run.py

@@ -0,0 +1,292 @@
+"""One paired zero/PPO Karman pinball VIV experiment; see docs/karman_viv.md.
+
+Research-only adapter: static precursor -> two fresh, identical moving plants.
+No moving checkpoint/rollback, policy retraining, or force filtering.
+"""
+from __future__ import annotations
+
+import argparse
+import hashlib
+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.steady_pinball.run import Structure, config, scaled_points, CENTRES, PROBES, COLUMNS
+
+BIAS = np.array([0., -4., 4.])
+MODEL_SHA = '148240c1dcb8b11d8e5a0a9e991f6874c4e799c9744707bf0fb32a8efeb0468d'
+NORM_SHA = '534fda818c030a56960e95d0f3ac1e350d887d03db7f80f1d6453195a502b7f0'
+STEP_COLUMNS = (*COLUMNS, 'omega_F', 'omega_U', 'omega_L')
+
+
+def sha(path):
+    h = hashlib.sha256()
+    with open(path, 'rb') as f:
+        for block in iter(lambda: f.read(1024*1024), b''):
+            h.update(block)
+    return h.hexdigest()
+
+
+def array_sha(a):
+    return hashlib.sha256(memoryview(np.ascontiguousarray(a)).cast('B')).hexdigest()
+
+
+def observation(sensor_mean, force_mean, scale, norm):
+    """Fine-grid area means/CME -> trained D20 sensor sums/forces -> frozen norm."""
+    sensors = np.asarray(sensor_mean, np.float64)*scale*78.
+    forces = np.asarray(force_mean, np.float64)*scale
+    raw = np.asarray([*sensors, *forces], np.float32)
+    preclip = np.hstack((raw[6:]/norm['force_norm_fact'],
+                        (raw[:6]-np.asarray(norm['sens_deviation'], np.float32)) /
+                        np.asarray(norm['sens_norm_fact'], np.float32)))
+    if not np.isfinite(preclip).all():
+        raise RuntimeError('nonfinite policy observation')
+    return raw, np.clip(preclip, -1, 1).astype(np.float32), int((np.abs(preclip)>1).sum())
+
+
+def wall_target(action, scale):
+    return (8.*np.asarray(action)+BIAS)*(.01/scale)
+
+
+def ema(eff, target, scale):
+    alpha = .9**(1./scale**2)
+    return alpha*eff + (1.-alpha)*target
+
+
+def load_policy(model):
+    # Same inference spaces as legacy_test.core.model_loader, without importing it.
+    import gymnasium as gym
+    from stable_baselines3 import PPO
+    class InferenceEnv(gym.Env):
+        observation_space = gym.spaces.Box(-1., 1., (12,), np.float32)
+        action_space = gym.spaces.Box(-1., 1., (3,), np.float32)
+        def reset(self, *, seed=None, options=None):
+            super().reset(seed=seed)
+            return np.zeros(12, np.float32), {}
+        def step(self, action):
+            return np.zeros(12, np.float32), 0., False, False, {}
+    return PPO.load(str(model), env=InferenceEnv(), device='cpu')
+
+
+def add_scene(sim, scale, moving=False, pinball=True):
+    from CelerisLab.body.motion import Translation
+    sim.add_body('circle', scaled_points(((200., 255.5),), scale)[0], 20.*scale)
+    probes = [sim.add_body('sensor', c, 5.*scale) for c in scaled_points(PROBES, scale)]
+    bodies = []
+    if pinball:
+        envelope = Translation((0., 0., -20.*scale, 20.*scale), .05/scale, shared=True) if moving else None
+        bodies = [sim.add_body('circle', c, 10.*scale, translation=envelope)
+                  for c in scaled_points(CENTRES, scale)]
+    return probes, bodies
+
+
+def capture_seed(sim):
+    """Capture both FP32 buffers at fixed geometry; never a moving-state checkpoint."""
+    import pycuda.driver as cuda
+    with sim.ctx.activate():
+        sim.stream.synchronize()
+        buffers = [np.empty_like(sim.field.ddf) for _ in range(2)]
+        for dst, src in zip(buffers, (sim.field.ddf_gpu, sim.field.temp_gpu)):
+            cuda.memcpy_dtoh(dst, src)
+    return (*buffers, sim.get_flags(), sim.stepper.step_count)
+
+
+def install_seed(sim, seed):
+    """Only a new, zero-offset instance may receive the static release state."""
+    import pycuda.driver as cuda
+    current, spare, flags, clock = seed
+    if sim.stepper.step_count != 0 or any(a.shape != sim.field.ddf.shape or a.dtype != sim.field.ddf.dtype
+                                         for a in (current, spare)):
+        raise ValueError('release transfer requires a fresh matching field')
+    if not np.array_equal(sim.get_flags(), flags):
+        raise ValueError('static/moving release flags differ')
+    for obj in sim.bodies.objects:
+        if obj.motion is not None and (obj.motion.applied != (0., 0.) or obj.motion.pending is not None):
+            raise ValueError('release transfer requires zero geometry and no target')
+    with sim.ctx.activate():
+        sim.stream.synchronize()
+        cuda.memcpy_htod(sim.field.ddf_gpu, current)
+        cuda.memcpy_htod(sim.field.temp_gpu, spare)
+    sim.stepper._step_count = clock  # one-time research initialization, not restore()
+    sim.field.invalidate_host_ddf_cache()
+    copied = capture_seed(sim)
+    if copied[3] != clock or any(not np.array_equal(a, b, equal_nan=True) for a, b in zip(copied[:3], seed[:3])):
+        raise RuntimeError('release transfer did not preserve both buffers/flags')
+    return dict(current=array_sha(copied[0]), spare=array_sha(copied[1]), flags=array_sha(copied[2]), clock=copied[3])
+
+
+def check_time(start, cap):
+    if time.monotonic()-start > cap:
+        raise TimeoutError('paired experiment wall cap')
+
+
+def warmup(sim, count, start, cap):
+    while count:
+        check_time(start, cap)
+        n = min(count, 10000)
+        sim.run(n)
+        count -= n
+
+
+def precursor(cfgpath, scale, warm, si, targets, out, start, cap):
+    from CelerisLab import Simulation
+    with Simulation(str(cfgpath)) as sim:
+        probes, _ = add_scene(sim, scale, pinball=False)
+        sim.initialize()
+        warmup(sim, warm, start, cap)
+        target = []
+        for _ in range(targets):
+            check_time(start, cap)
+            sim.run(si)
+            target.append([v*scale*78. for p in probes for v in sim.read_sensor(p)[:2]])
+        np.save(out/'target.npy', np.asarray(target, np.float32))
+        first = sim.bodies.count
+        for c in scaled_points(CENTRES, scale):
+            sim.add_body('circle', c, 10.*scale)
+        sim.sync_bodies()
+        bodies = list(range(first, first+3))
+        warmup(sim, warm-1, start, cap)
+        sim.run(1)
+        force = sum(float(sim.read_force(p)[1]) for p in bodies)
+        seed = capture_seed(sim)
+        np.savez(out/'release.npz', current=seed[0], spare=seed[1], flags=seed[2], clock=seed[3], force=force)
+        print(json.dumps(dict(event='precursor_complete', clock=seed[3], force=force,
+                              elapsed_s=time.monotonic()-start)), flush=True)
+        return seed, force
+
+
+def run_arm(name, cfgpath, scale, si, intervals, seed, initial_force, model, norm, out, start, cap):
+    from CelerisLab import Simulation
+    from CelerisLab.body.motion import MovingTables
+    cols = {k:i for i, k in enumerate(MovingTables.DIAG)}
+    diameter = 20.*scale
+    structure = Structure(2., 1.e-4/scale**2, 0., diameter)
+    arm = out/name
+    arm.mkdir()
+    rows = np.lib.format.open_memmap(arm/'steps.npy', mode='w+', dtype=np.float64,
+                                    shape=(si*intervals, len(STEP_COLUMNS)))
+    rows[:] = np.nan
+    signals = {k:np.full((intervals, width), np.nan, np.float32) for k, width in
+               [('sensors',6), ('forces',6), ('actions',3), ('observations',12)]}
+    completed = completed_intervals = clipped = 0
+    status, transfer = 'running', None
+    force = initial_force
+    obs = np.zeros(12, np.float32)
+    arm_start = time.monotonic()
+    try:
+        with Simulation(str(cfgpath)) as sim:
+            probes, bodies = add_scene(sim, scale, moving=True)
+            sim.initialize()
+            transfer = install_seed(sim, seed)
+            for i in range(intervals):
+                action = -BIAS/8 if name == 'zero' else 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 policy action')
+                target, eff = wall_target(action, scale), np.zeros(3)
+                sensor_sum, force_sum = np.zeros(6), np.zeros(6)
+                for k in range(si):
+                    if k % 1000 == 0:
+                        check_time(start, cap)
+                    y1, v1 = structure.candidate(force)
+                    if not np.isfinite([y1,v1]).all() or abs(y1)>diameter:
+                        raise RuntimeError('nonfinite structure or |y| > D; no retuning')
+                    eff = ema(eff, target, scale)
+                    omega = eff/(diameter/2)
+                    for p, w in zip(bodies, omega):
+                        sim.set_body(p, 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')
+                    # Policy sees packed CME as in fixed-body training; structure sees CME+Wen.
+                    force_sum += f
+                    sensor_sum += s
+                    force = float(cme.sum(dtype=np.float64)+wen.sum(dtype=np.float64))
+                    rows[completed] = row
+                    completed += 1
+                    structure.y, structure.v = y1,v1
+                raw, obs, nclip = observation(sensor_sum/si, force_sum/si, scale, norm)
+                clipped += nclip
+                for key, val in [('sensors',raw[:6]),('forces',raw[6:]),('actions',action),('observations',obs)]:
+                    signals[key][i] = val
+                completed_intervals = i+1
+                if (i+1)%10 == 0 or i+1 == intervals:
+                    rows.flush()
+                    np.savez(arm/'signals.npz', **signals)
+                    print(json.dumps(dict(event='progress', arm=name, interval=i+1, step=completed,
+                                          tU_D=completed*.01/scale/diameter, y_D=y1/diameter,
+                                          elapsed_s=time.monotonic()-arm_start)), flush=True)
+            status = 'duration_completed_not_qualified'
+    except BaseException as exc:
+        status = f'error: {type(exc).__name__}: {exc}'
+        raise
+    finally:
+        rows.flush()
+        np.savez(arm/'signals.npz', **signals)
+        report = dict(status=status, completed_steps=completed, completed_intervals=completed_intervals,
+                      elapsed_s=time.monotonic()-arm_start, release_transfer=transfer, clipped_observation_channels=clipped,
+                      columns=STEP_COLUMNS, M=structure.mass,K=structure.stiffness,fn=structure.frequency,damping=0.,
+                      max_abs_y_D=float(np.max(np.abs(rows[:completed,1]))/diameter) if completed else None)
+        (arm/'report.json').write_text(json.dumps(report,indent=2)+'\n')
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument('--model', type=Path, required=True)
+    ap.add_argument('--norm', type=Path, required=True)
+    ap.add_argument('--out', type=Path, required=True)
+    ap.add_argument('--smoke', action='store_true', help='short plumbing only; no physical conclusion')
+    ap.add_argument('--wall-seconds', type=float, default=14400.)
+    args = ap.parse_args()
+    if not np.isfinite(args.wall_seconds) or args.wall_seconds<=0:
+        ap.error('positive finite wall cap required')
+    if sha(args.model)!=MODEL_SHA or sha(args.norm)!=NORM_SHA:
+        ap.error('model/normalizer must match frozen provenance')
+    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)
+    norm = json.loads(args.norm.read_text())
+    model = load_policy(args.model)
+    scale = 2.
+    warm, si, targets, intervals = (32,8,2,4) if args.smoke else (2048000,3200,150,400)
+    cfgpath = out/'config.json'
+    cfgpath.write_text(json.dumps(config(40),indent=2)+'\n')
+    receipt = dict(source=subprocess.check_output(['git','-C',str(ROOT),'rev-parse','HEAD'],text=True).strip(),
+                   model_sha256=sha(args.model),norm_sha256=sha(args.norm),smoke=args.smoke,D=40,U=.005,nu=.004,
+                   warmup_each=warm,target_intervals=targets,SI=si,response_intervals=intervals,
+                   response_steps=si*intervals,response_tU_D=si*intervals*.005/40,
+                   rhs='previous native step sum CME_y+Wen_y; policy uses interval packed CME',
+                   probe_frame='lab-fixed; fine area mean * 2 * trained 78 cells',
+                   ema_native_alpha=.9**.25, structural_damping=0., columns=STEP_COLUMNS)
+    (out/'case.json').write_text(json.dumps(receipt,indent=2)+'\n')
+    start = time.monotonic()
+    status = 'running'
+    try:
+        seed, force = precursor(cfgpath,scale,warm,si,targets,out,start,args.wall_seconds)
+        for name in ('zero','controlled'):
+            run_arm(name,cfgpath,scale,si,intervals,seed,force,model,norm,out,start,args.wall_seconds)
+        status = 'smoke_completed' if args.smoke else 'pair_completed_not_qualified'
+    except BaseException as exc:
+        status = f'error: {type(exc).__name__}: {exc}'
+        raise
+    finally:
+        (out/'terminal.json').write_text(json.dumps(dict(status=status,elapsed_s=time.monotonic()-start),indent=2)+'\n')
+
+
+if __name__ == '__main__':
+    main()