309 lines
14 KiB
Diff
309 lines
14 KiB
Diff
--- 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:
|