"""One exact J3 replay and one 520-step audit trajectory; no production edits.""" import argparse from dataclasses import asdict import hashlib import json import os from pathlib import Path import sys import h5py import numpy as np import pycuda.driver as cuda import CelerisLab from CelerisLab.cuda import compiler_v2 as compiler from run_j1 import ROOT, AUDIT, make_sim, populations, observables from run_j2_slot import targets START, LENGTH, SCALE = 20760, 520, .002 PHASE = 'postcollision g(t); D2Q9 paired ordering; unshifted FP32 double_buffer; q,y,x' def sha(path): return hashlib.sha256(Path(path).read_bytes()).hexdigest() def buffers(sim): arrays = dict(ddf=np.empty(sim.field.n*sim.field.nq, dtype=np.float32), temp=np.empty(sim.field.n*sim.field.nq, dtype=np.float32), flag=np.empty(sim.field.n, dtype=np.uint16)) with sim.ctx.activate(): for name, array in arrays.items(): cuda.memcpy_dtoh(array, getattr(sim.field, name+'_gpu')) return arrays def save_verified(sim, path): """Persist both real buffers, then restore into a new case without advancing.""" assert not path.exists(), f'Refusing to overwrite {path}' before, config = buffers(sim), asdict(sim.lbm_cfg) params, count = sim.field._runtime_params.copy(), sim.stepper.step_count sim.save_checkpoint(str(path)) with h5py.File(path, 'a') as hf: hf.attrs['distribution_phase'] = PHASE hf.attrs['diagnostic_script_sha256'] = sha(__file__) for name, array in before.items(): assert hf[name][:].tobytes() == array.tobytes(), name assert json.loads(hf.attrs['config_json']) == config assert json.loads(hf.attrs['runtime_params_json']) == params assert int(hf.attrs['step_count']) == count with make_sim(.002, nx=48, top_u=0., slot_width=16, top_mode='reservoir') as restored: restored.load_checkpoint(str(path)) after = buffers(restored) assert all(before[k].tobytes() == after[k].tobytes() for k in before) assert asdict(restored.lbm_cfg) == config assert restored.field._runtime_params == params assert restored.stepper.step_count == count return dict(path=str(path), sha256=sha(path), step=count, phase=PHASE, no_advance_restore='passed', buffers=['ddf', 'temp', 'flag']) def replay(sim): sim.initialize() rho, old_u, old_v = observables(populations(sim)) initial_mass = old_mass = float(rho.sum()) history = [dict(step=0, mass=initial_mass)] for _ in range(40): sim.run(519) rho, u, v = observables(populations(sim)) mass = float(rho.sum()) history.append(dict(step=sim.stepper.step_count, mass=mass, u_change=float(np.sqrt(np.mean((u-old_u)**2))/SCALE), v_change=float(np.sqrt(np.mean((v-old_v)**2))/SCALE), mass_change=abs(mass-old_mass)/initial_mass)) old_u, old_v, old_mass = u, v, mass assert sim.stepper.step_count == START return history, (rho, u, v) def record_trajectory(sim, path, reference): """Exactly LENGTH advances. Reuses J3's already validated audit source.""" build = compiler.generate_config(sim.lbm_cfg) original_step = sim.stepper.step_fn try: source = Path(build)/'step/one_step_double.cu' code = source.read_text() call, marker = 'collide_dispatch(f, rho_n, ux, uy);', '__global__ __launch_bounds__(NT, 2)' assert code.count(call) == 1 and code.count(marker) == 1 source.write_text(code.replace(call, 'j1_collide_audit(f, rho_n, ux, uy);') .replace(marker, AUDIT+'\n'+marker)) audit_hash = sha(source) assert audit_hash == reference['audit_source_sha256'], 'Audit source changed; stop' ptx = compiler.compile_kernel(build_dir=build) with sim.ctx.activate(), h5py.File(path, 'x') as hf: module = compiler.load_module(ptx) params, size = module.get_global('d_params') cuda.memcpy_dtod(params, sim.field.module.get_global('d_params')[0], size) incoming_ptr, _ = module.get_global('j1_incoming') collision_ptr, size = module.get_global('j1_collision') cuda.memset_d8(collision_ptr, 0, size) incoming = np.empty((2, 48), dtype=np.float64) collision = np.empty((35, 48), dtype=np.float64) hf.attrs.update(phase=PHASE, first_step=START, last_step=START+LENGTH, first_fluid_y=1, velocity_scale=SCALE, rho_scale=1., recorded_samples=0, audit_source_sha256=audit_hash) hf.create_dataset('step', data=np.arange(START, START+LENGTH+1)) for name in ('rho', 'u', 'v'): hf.create_dataset(name, (LENGTH+1, 33, 48), dtype='f8', chunks=(1, 33, 48), compression='gzip') for name in ('mass_flux_by_y', 'volume_flux_by_y'): hf.create_dataset(name, (LENGTH+1, 33), dtype='f8') hf.create_dataset('mass', (LENGTH+1,), dtype='f8') for name in ('incoming', 'outgoing'): hf.create_dataset(name, (LENGTH, 2, 48), dtype='f8', compression='gzip') hf.create_dataset('collision', (LENGTH, 35, 48), dtype='f8', compression='gzip') hf.attrs['boundary_order'] = 'bottom,top; all x including zero-target flanks' bottom, top_guess, slot = targets(.002) hf.create_dataset('bottom_slot_mask', data=slot) hf.create_dataset('bottom_requested_velocity', data=bottom) hf.attrs.update(requested_bottom_volume=float(bottom.sum()), top_initial_guess=float(top_guess)) def store_frame(n, f): rho, u, v = observables(f) for name, array in zip(('rho', 'u', 'v'), (rho, u, v)): hf[name][n] = array hf['mass'][n] = rho.sum() hf['mass_flux_by_y'][n] = (rho*v).sum(axis=1) hf['volume_flux_by_y'][n] = v.sum(axis=1) hf.attrs['recorded_samples'] = n+1 before = populations(sim) store_frame(0, before) sim.stepper.step_fn = module.get_function('OneStep') for n in range(LENGTH): outgoing = np.stack((before[[4,6,7], 1].sum(axis=0), before[[3,5,8], -2].sum(axis=0))) mass = before[:, 1:-1].sum() sim.run(1) after = populations(sim) assert np.isfinite(after).all() cuda.memcpy_dtoh(incoming, incoming_ptr) cuda.memcpy_dtoh(collision, collision_ptr) hf['incoming'][n], hf['outgoing'][n], hf['collision'][n] = incoming, outgoing, collision store_frame(n+1, after) if n < 16: row = reference['steps'][n] measured = dict(mass=mass, delta_mass=after[:, 1:-1].sum()-mass, bottom_net_in=incoming[0].sum()-outgoing[0].sum(), top_net_in=incoming[1].sum()-outgoing[1].sum(), collision_roundoff=collision.sum()) assert all(measured[k] == row[k] for k in measured), 'Prior J3 audit not reproduced' before = after assert sim.stepper.step_count == START+LENGTH hf.attrs['first_16_match_j3'] = True return audit_hash finally: sim.stepper.step_fn = original_step compiler.cleanup_build_dir(build) def variation(u, v): """Per-time/per-origin spatial metrics; no assumption of independent samples.""" delta = np.stack((u, v), axis=1) return dict(rms=np.sqrt(np.mean(delta**2, axis=(2, 3))), max_abs=np.max(np.abs(delta), axis=(2, 3)), signed_mean=np.mean(delta, axis=(2, 3)), vector_rms=np.sqrt(np.mean(u*u+v*v, axis=(1, 2))), vector_max=np.sqrt(np.max(u*u+v*v, axis=(1, 2)))) def summarize(path): summary = dict(lags={}) with h5py.File(path, 'a') as hf: assert hf.attrs['recorded_samples'] == LENGTH+1 u, v, mass = hf['u'][:], hf['v'][:], hf['mass'][:] for lag in (1, 16, 519): group = hf.create_group(f'lag/{lag}') group.create_dataset('origin_step', data=np.arange(START, START+LENGTH-lag+1)) values = variation(u[lag:]-u[:-lag], v[lag:]-v[:-lag]) for name, array in values.items(): group.create_dataset(name, data=array) summary['lags'][str(lag)] = dict(origins=LENGTH+1-lag, rms_mean=values['rms'].mean(axis=0).tolist(), rms_min=values['rms'].min(axis=0).tolist(), rms_max=values['rms'].max(axis=0).tolist(), max_abs=values['max_abs'].max(axis=0).tolist(), rms_mean_scaled=(values['rms'].mean(axis=0)/SCALE).tolist(), signed_mean_change=values['signed_mean'].mean(axis=0).tolist(), signed_mean_change_per_step=(values['signed_mean'].mean(axis=0)/lag).tolist(), vector_rms_mean=float(values['vector_rms'].mean()), vector_max=float(values['vector_max'].max())) summary['lags'][str(lag)]['metrics'] = { name: {label: func(array, axis=0).tolist() for label, func in (('mean', np.mean), ('minimum', np.min), ('maximum', np.max))} for name, array in values.items()} summary['lags'][str(lag)]['metrics_scaled'] = { name: {label: (np.asarray(value)/SCALE).tolist() for label, value in metrics.items()} for name, metrics in summary['lags'][str(lag)]['metrics'].items()} means = hf.create_group('time_mean') for name in ('rho', 'u', 'v'): means.create_dataset(name, data=hf[name][:].mean(axis=0)) deviation = variation(u-u.mean(axis=0), v-v.mean(axis=0)) group = hf.create_group('deviation') for name, array in deviation.items(): group.create_dataset(name, data=array) summary['time_mean_deviation'] = dict(rms=np.sqrt(np.mean(deviation['rms']**2, axis=0)).tolist(), max_abs=deviation['max_abs'].max(axis=0).tolist()) end_delta = np.stack((u[-1]-u[0], v[-1]-v[0])) hf.create_dataset('end_minus_start_uv', data=end_delta) hf.create_dataset('mean_drift_per_step_uv', data=end_delta/LENGTH) summary['signed_end_change'] = dict(mean=end_delta.mean(axis=(1, 2)).tolist(), minimum=end_delta.min(axis=(1, 2)).tolist(), maximum=end_delta.max(axis=(1, 2)).tolist(), mean_per_step=(end_delta.mean(axis=(1, 2))/LENGTH).tolist()) incoming, outgoing, collision = hf['incoming'][:], hf['outgoing'][:], hf['collision'][:] slot = hf['bottom_slot_mask'][:] groups = hf.create_group('budget') net = incoming.sum(axis=(1, 2))-outgoing.sum(axis=(1, 2)) roundoff, delta = collision.sum(axis=(1, 2)), np.diff(mass) residual = delta-net-roundoff for name, array in dict(delta_mass=delta, net_boundary=net, collision_roundoff=roundoff, physical_residual=delta-net, residual=residual).items(): groups.create_dataset(name, data=array) summary['budget'] = dict(mean_delta_mass=float(delta.mean()), net_boundary=float(net.mean()), collision_roundoff=float(roundoff.mean()), max_residual=float(np.abs(residual).max()), max_relative_residual=float(np.max(np.abs(residual)/mass[:-1])), max_physical_relative_residual=float(np.max(np.abs(delta-net)/mass[:-1])), groups={name:dict(incoming=float(i.mean()), outgoing=float(o.mean()), net=float((i-o).mean())) for name, i, o in ( ('slot', incoming[:, 0, slot].sum(axis=1), outgoing[:, 0, slot].sum(axis=1)), ('flanks', incoming[:, 0, ~slot].sum(axis=1), outgoing[:, 0, ~slot].sum(axis=1)), ('top', incoming[:, 1].sum(axis=1), outgoing[:, 1].sum(axis=1)))}) summary['mass'] = dict(start=float(mass[0]), end=float(mass[-1]), delta=float(mass[-1]-mass[0]), relative_span=float(np.ptp(mass)/mass[0]), relative_delta=float((mass[-1]-mass[0])/mass[0])) return summary def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--output-dir', type=Path, required=True) args = parser.parse_args() assert Path.cwd() == ROOT == Path('/home/frank14f/CelerisLab') assert sys.executable == '/home/frank14f/anaconda3/envs/pycuda_3_10/bin/python' assert Path(CelerisLab.__file__).resolve() == ROOT/'src/CelerisLab/__init__.py' assert Path(compiler.kernel_path('kernel_v2.cu')).resolve() == ROOT/'src/CelerisLab/lbm/kernels/kernel_v2.cu' assert os.environ.get('CUDA_VISIBLE_DEVICES') in ('0', '1') out = args.output_dir.resolve() assert out == ROOT/'research/j3-lag-2026-10-12' baseline = json.loads((out/'baseline.json').read_text()) assert all(sha(ROOT/name) == digest for name, digest in baseline['hashes'].items()) for name in ('results.json', 'trajectory.h5', 'start-20760.h5', 'end-21280.h5'): assert not (out/name).exists(), f'Refusing to overwrite {name}' original = json.loads((ROOT/'research/j3-reservoir-2026-10-12/results.json').read_text())['cases'][0] report = dict(python=sys.executable, package=CelerisLab.__file__, cuda_source=compiler.kernel_path('kernel_v2.cu'), physical_gpu=os.environ['CUDA_VISIBLE_DEVICES'], script_sha256=sha(__file__), phase=PHASE, velocity_scale=SCALE, replay_steps=START, planned_diagnostic_steps=LENGTH, diagnostic_steps=0, j3_acceptance='failed_unchanged') def save_report(): (out/'results.json').write_text(json.dumps(report, indent=2)+'\n') with make_sim(.002, nx=48, top_u=0., slot_width=16, top_mode='reservoir') as sim: history, fields = replay(sim) report['start_checkpoint'] = save_verified(sim, out/'start-20760.h5') report['replay_history'] = history report['replay_matches_j3'] = bool(history == original['history'] and all( np.array_equal(array, original[name]) for array, name in zip(fields, ('rho', 'u', 'v')))) save_report() assert report['replay_matches_j3'], 'Replay differs from J3; saved real checkpoint, stopping' print('20760 replay and no-advance checkpoint restore passed; matches J3.', flush=True) try: report['audit_source_sha256'] = record_trajectory(sim, out/'trajectory.h5', original['population_budget']) finally: report['diagnostic_steps'] = sim.stepper.step_count-START save_report() report['end_checkpoint'] = save_verified(sim, out/'end-21280.h5') report.update(summarize(out/'trajectory.h5')) report['trajectory_sha256'] = sha(out/'trajectory.h5') save_report() print(json.dumps({k:report[k] for k in ('lags','signed_end_change','mass','budget')}), flush=True) if __name__ == '__main__': main()