--- /dev/null +++ b/tests/validation/diagnose_j3_lags.py @@ -0,0 +1,276 @@ +"""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()