--- a/src/CelerisLab/common/checkpoint.py +++ b/src/CelerisLab/common/checkpoint.py @@ -21,8 +21,10 @@ import numpy as np import pycuda.driver as cuda +from ..lbm.descriptors import FLUID, SOLID, OBSTACLE, SENSOR_FLAG, FRC_REGION + # File format version — bump on incompatible layout changes. -CKPT_VERSION = 2 +CKPT_VERSION = 3 _SUPPORTED_STORE_PRECISIONS = ("FP32", "FP16S") @@ -35,10 +37,25 @@ ) -def _geometry_records(field, bodies): +def _domain_flags(field, bodies, flags): + """Keep actual domain boundaries, removing only registered object overlays. + + A solid object's hidden underlay is unknowable from merged flags: use FLUID + there canonically. This base is for fixed-geometry restart, not body removal. + """ + base = flags & np.uint16(0xffff ^ (SENSOR_FLAG | FRC_REGION)) + for obj in bodies.objects: + mask = obj.get_flag_mask(field.nx, field.ny, field.nz) + base[(mask & SOLID) != 0] = FLUID + if np.any(base & OBSTACLE): + raise ValueError("Checkpoint flags contain an unregistered obstacle") + return base + + +def _geometry_records(field, bodies, base_flags=None): """Describe actual geometry/projection, not just the number of objects.""" records = [] - flags = bodies.build_flags(field.build_channel_flags()) + flags = bodies.build_flags(field.build_channel_flags() if base_flags is None else base_flags) for obj in bodies.objects: digest = hashlib.sha256() digest.update(obj.get_flag_mask(field.nx, field.ny, field.nz).tobytes()) @@ -108,6 +125,11 @@ cuda.memcpy_dtoh(tmp_buf, field.temp_gpu) flag_buf = np.empty(field.n, dtype=np.uint16) cuda.memcpy_dtoh(flag_buf, field.flag_gpu) + domain_flags = _domain_flags(field, bodies, flag_buf) + geometry, expected_flags = _geometry_records(field, bodies, domain_flags) + if not np.array_equal(flag_buf, expected_flags): + raise ValueError("Checkpoint flags do not match current object/domain geometry") + bodies.build_compact_lists_for(bodies.objects, domain_flags=expected_flags) with h5py.File(path, "w") as hf: # Version and metadata @@ -130,8 +152,9 @@ compression_opts=4) hf.create_dataset("flag", data=flag_buf, compression="gzip", compression_opts=4) - - geometry, _ = _geometry_records(field, bodies) + hf.create_dataset("domain_flag", data=domain_flags, compression="gzip", + compression_opts=4) + hf.attrs["geometry_json"] = json.dumps(geometry) hf.attrs["runtime_params_json"] = json.dumps(field._runtime_params) hf.attrs["obs_accum_steps"] = bodies._obs_accum_steps @@ -167,6 +190,8 @@ The current config's lattice_model, nx, ny, nz, and store_precision must match the checkpoint. Other parameters (collision, LES, …) may differ for parameter-study restarts. + v3 restores the saved domain boundaries; v1/v2 retain their default-channel + validation. Object geometry and roles must match in every version. Args: path: HDF5 checkpoint file path. @@ -181,7 +206,7 @@ with h5py.File(path, "r") as hf: # Version check ver = hf.attrs.get("ckpt_version", 0) - if ver not in (1, CKPT_VERSION): + if ver not in (1, 2, CKPT_VERSION): raise ValueError( f"Checkpoint version mismatch: file={ver}, " f"expected={CKPT_VERSION}. " @@ -218,7 +243,14 @@ if (len(obj_data) != bodies.count or [od.get("obj_id") for od in obj_data] != list(range(bodies.count))): raise ValueError("Checkpoint object IDs must match the dense current registry") - geometry, expected_flags = _geometry_records(field, bodies) + domain_flags = None + if ver >= 3: + if "domain_flag" not in hf: + raise ValueError("Checkpoint is missing domain_flag") + domain_flags = hf["domain_flag"][:] + if domain_flags.shape != (field.n,) or domain_flags.dtype != np.uint16: + raise ValueError("Invalid checkpoint domain_flag shape/dtype") + geometry, expected_flags = _geometry_records(field, bodies, domain_flags) if ver == 1: _validate_legacy_objects(obj_data, bodies) elif json.loads(hf.attrs["geometry_json"]) != geometry: @@ -236,7 +268,9 @@ if array.shape != shape or array.dtype != expected_dtype: raise ValueError(f"Invalid checkpoint {name} shape/dtype") if not np.array_equal(flag_buf, expected_flags): - raise ValueError("Checkpoint flags do not match current object/channel geometry") + raise ValueError("Checkpoint flags do not match current object/domain geometry") + if ver >= 3 and not np.array_equal(domain_flags, _domain_flags(field, bodies, flag_buf)): + raise ValueError("Checkpoint domain_flag is not a canonical object-free domain") restored_step = int(hf.attrs["step_count"]) if restored_step < 0: raise ValueError("Checkpoint step_count must be nonnegative") --- a/tests/unit/test_state_regressions.py +++ b/tests/unit/test_state_regressions.py @@ -18,7 +18,7 @@ from CelerisLab.body.ddf_patch import patch_ddf_for_body_sync, upload_patched_ddf from CelerisLab.common.checkpoint import save_checkpoint, load_checkpoint from CelerisLab.lbm.field import LBMField -from CelerisLab.lbm.descriptors import FLUID, SOLID, SENSOR_FLAG, FRC_REGION +from CelerisLab.lbm.descriptors import FLUID, SOLID, OBSTACLE, SENSOR_FLAG, FRC_REGION from CelerisLab.lbm.equilibrium import compute_feq_d2q9 W = np.array([4/9] + [1/9]*4 + [1/36]*4, dtype=np.float32) @@ -306,11 +306,42 @@ assert field._runtime_params['fx'] == .0002 +@pytest.mark.parametrize("periodic_x", [False, True]) +@pytest.mark.parametrize("with_objects", [False, True]) +def test_checkpoint_restores_actual_domain(device, tmp_path, periodic_x, with_objects): + source = make_field(device) + objects = [obj(x=2, r=1), obj(x=10, r=2, is_sensor=True)] if with_objects else [] + bodies = manager(source, objects) + base = source.build_channel_flags() + if periodic_x: + base.reshape(source.ny, source.nx)[1:-1, [0, -1]] = FLUID + source.flag = bodies.build_flags(base) + source.upload_flags() + bodies.sync_to_gpu(source, rebuild_flags=False) + saved_flags = source.flag.copy() + # The device, not a possibly stale host mirror, is authoritative at save. + source.flag = source.build_channel_flags() + path = save_checkpoint(source, SimpleNamespace(step_count=3), source.cfg, + bodies, str(tmp_path / 'domain.h5')) + target = make_field(device) + target_bodies = manager(target, objects) + stepper = SimpleNamespace(_step_count=0) + load_checkpoint(path, target, stepper, target.cfg, target_bodies) + assert stepper._step_count == 3 + np.testing.assert_array_equal(target.flag, saved_flags) + np.testing.assert_array_equal(np.frombuffer(target.flag_gpu.data, np.uint16), saved_flags) + assert target.ddf_gpu.data == source.ddf_gpu.data + + +@pytest.mark.parametrize("periodic_x", [False, True]) @pytest.mark.parametrize("change", ["geometry", "role", "ddf_shape"]) -def test_checkpoint_rejects_before_device_writes(device,tmp_path,change): +def test_checkpoint_rejects_before_device_writes(device,tmp_path,change,periodic_x): field = make_field(device) bodies = manager(field,[obj()]) bodies.sync_to_gpu(field) + if periodic_x: + field.flag.reshape(field.ny, field.nx)[1:-1, [0, -1]] = FLUID + field.upload_flags() path = save_checkpoint(field,SimpleNamespace(step_count=3),field.cfg,bodies,str(tmp_path/'state.h5')) if change == 'geometry': bodies.get(0).geometry = CircleGeometry(6,5,2) @@ -335,6 +366,89 @@ field.cfg.validate() load_checkpoint(path,field,SimpleNamespace(_step_count=0),field.cfg,bodies) assert field._runtime_params['omega'] == field.cfg.omega + + +@pytest.mark.parametrize('version', [1, 2]) +def test_legacy_checkpoint_default_domain(device, tmp_path, version): + field = make_field(device) + bodies = manager(field, [obj()]) + bodies.sync_to_gpu(field) + path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg, + bodies, str(tmp_path / 'legacy-domain.h5')) + with h5py.File(path, 'a') as hf: + hf.attrs['ckpt_version'] = version + del hf['domain_flag'] + load_checkpoint(path, field, SimpleNamespace(_step_count=0), field.cfg, bodies) + # Legacy files have no authoritative base-domain record. Keep their check. + with h5py.File(path, 'a') as hf: + hf['flag'][field.nx] = FLUID + writes = device.writes + with pytest.raises(ValueError, match='geometry'): + load_checkpoint(path, field, SimpleNamespace(_step_count=0), field.cfg, bodies) + assert device.writes == writes + + +@pytest.mark.parametrize('damage', ['missing', 'shape', 'dtype', 'boundary', 'hidden', 'annotation']) +def test_checkpoint_rejects_invalid_domain_before_writes(device, tmp_path, damage): + field = make_field(device) + bodies = manager(field, [obj()]) + bodies.sync_to_gpu(field) + path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg, + bodies, str(tmp_path / 'invalid-domain.h5')) + with h5py.File(path, 'a') as hf: + if damage in ('missing', 'shape', 'dtype'): + del hf['domain_flag'] + if damage != 'missing': + hf.create_dataset('domain_flag', data=np.ones( + 2 if damage == 'shape' else field.n, + dtype=np.uint16 if damage == 'shape' else np.float32)) + elif damage == 'boundary': + hf['domain_flag'][field.nx] = FLUID + elif damage == 'hidden': + hf['domain_flag'][5 + 5 * field.nx] = SOLID + else: + # Mutate both records: overlay idempotence alone would accept this. + for name in ('domain_flag', 'flag'): + hf[name][field.nx + 1] = FLUID | SENSOR_FLAG + before = bytes(field.ddf_gpu.data), field.flag.copy() + writes = device.writes + with pytest.raises(ValueError): + load_checkpoint(path, field, SimpleNamespace(_step_count=0), field.cfg, bodies) + assert device.writes == writes + assert bytes(field.ddf_gpu.data) == before[0] + np.testing.assert_array_equal(field.flag, before[1]) + + +@pytest.mark.parametrize('extra', [OBSTACLE | SOLID, SENSOR_FLAG, FRC_REGION]) +def test_checkpoint_save_rejects_unregistered_overlay(device, tmp_path, extra): + field = make_field(device) + bodies = manager(field) + field.flag[field.nx + 1] |= extra + field.upload_flags() + path = tmp_path / 'unchanged.h5' + path.write_bytes(b'existing checkpoint') + with pytest.raises(ValueError): + save_checkpoint(field, SimpleNamespace(step_count=0), field.cfg, bodies, str(path)) + assert path.read_bytes() == b'existing checkpoint' + + +def test_checkpoint_domain_handles_plain_solid_mask(device, tmp_path): + class PlainSolidCircle(CircleGeometry): + def build_flag_mask(self, nx, ny): + mask = super().build_flag_mask(nx, ny) + mask[(mask & SOLID) != 0] = SOLID + return mask + + field = make_field(device) + body = obj() + body.geometry = PlainSolidCircle(5, 5, 2) + bodies = manager(field, [body]) + bodies.sync_to_gpu(field) + path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg, + bodies, str(tmp_path / 'plain-solid.h5')) + with h5py.File(path) as hf: + assert hf['domain_flag'][5 + 5 * field.nx] == FLUID + load_checkpoint(path, field, SimpleNamespace(_step_count=0), field.cfg, bodies) @pytest.mark.parametrize("kind",['circle','sensor','force_region']) --- a/tests/integration/test_state_gpu.py +++ b/tests/integration/test_state_gpu.py @@ -91,10 +91,48 @@ np.testing.assert_array_equal(sim.get_ddf(),expected) -def test_checkpoint_geometry_mismatch_leaves_device_unchanged(tmp_path): +@pytest.mark.parametrize('streaming', ['double_buffer', 'esopull']) +@pytest.mark.parametrize('periodic_x', [False, True]) +def test_checkpoint_domain_exact_continuation(streaming, periodic_x, tmp_path): + # Same low-Mach forced channel as the Poiseuille runner, bounded to 7 steps. + kwargs = dict(nx=16, ny=18, streaming=streaming, collision='MRT', + shifted=True, viscosity=.1, velocity=0.) + with make_sim(**kwargs) as source: + source.initialize() + if periodic_x: + source.field.flag.reshape(18, 16)[1:-1, [0, -1]] = FLUID + source.field.upload_flags() + source.update_runtime_params(fx=8 * .1 * .03 / 16**2) + source.run(3) + flags = source.get_flags().copy() + saved = source.get_ddf().copy() + path = source.save_checkpoint(str(tmp_path / 'domain.h5')) + source.run(4) + expected = source.get_ddf().copy() + # Restore into this instance and also a fresh default-domain instance. + source.load_checkpoint(path) + np.testing.assert_array_equal(source.get_flags(), flags) + np.testing.assert_array_equal(source.get_ddf(), saved) + source.run(4) + np.testing.assert_array_equal(source.get_ddf(), expected) + with make_sim(**kwargs) as target: + target.initialize() + target.load_checkpoint(path) + assert target.stepper.step_count == 3 + np.testing.assert_array_equal(target.get_flags(), flags) + np.testing.assert_array_equal(target.get_ddf(), saved) + target.run(4) + np.testing.assert_array_equal(target.get_ddf(), expected) + + +@pytest.mark.parametrize('periodic_x', [False, True]) +def test_checkpoint_geometry_mismatch_leaves_device_unchanged(tmp_path, periodic_x): with make_sim(nx=48,ny=24) as source: source.add_body('circle',center=(14,12),radius=3) source.initialize() + if periodic_x: + source.field.flag.reshape(24, 48)[1:-1, [0, -1]] = FLUID + source.field.upload_flags() source.run(2) path = source.save_checkpoint(str(tmp_path/'geometry.h5')) with make_sim(nx=48,ny=24) as target: