Support tagged shifted storage in blocked host transfers

src/CelerisLab/lbm/block_stepper.py

@@ -21,6 +21,7 @@ import pycuda.driver as cuda
 from ..cuda import compiler_v2 as compiler
 from .blocks import NQ, BlockLayout2D, blocked_link_tables
 from .descriptors import D2Q9_EX, D2Q9_EY, FLUID, SOLID
+from .population_storage import _StoredPopulations, _encode_physical, _decode_stored
 
 MAX_GRID_DIM = 65535  # CUDA gridDim.y/z limit (rows, blocks)
 
@@ -85,10 +86,10 @@ class PeriodicBlockStepper:
         checked.validate()
         if (not checked.is_d2q9 or checked.collision != "SRT"
                 or checked.streaming != "double_buffer" or checked.store_precision != "FP32"
-                or checked.ddf_shifting or checked.les_enabled
+                or checked.les_enabled
                 or (checked.ny, checked.nx) != layout.shape):
             raise ValueError("PeriodicBlockStepper requires matching D2Q9/SRT/FP32/"
-                             "unshifted/double_buffer config without LES")
+                             "double_buffer config without LES")
         ly, lx = layout.local_shape
         pairs = len(layout.halo_dst)
         tpb = checked.threads_per_block
@@ -99,6 +100,7 @@ class PeriodicBlockStepper:
         links = 0 if self.boundary is None else len(self.boundary["cell"])
         if links >= 2**31:
             raise ValueError("cut links exceed launch limits")
+        self._ddf_shifting = bool(checked.ddf_shifting)
         self.layout = layout
         self.ctx = context
         self._links = links
@@ -177,14 +179,20 @@ class PeriodicBlockStepper:
         """Upload dense G_0; the spare buffer starts NaN so stale halos cannot hide."""
         with self._lock, self.ctx.activate():
             self._require("new", "ready")
-            dense = np.array(g, dtype=np.float32, order="C", copy=True)
+            dense = np.array(g, dtype=np.float64 if self._ddf_shifting else np.float32,
+                             order="C", copy=True)
             if dense.shape != (NQ, *self.layout.shape):
                 raise ValueError("Expected a Q-major (9, ny, nx) array")
             fluid = np.ones(self.layout.shape, bool) if self.boundary is None else ~self.boundary["solid"]
             if not np.isfinite(dense[:, fluid]).all() or (dense[:, fluid].sum(axis=0) <= 0).any():
                 raise ValueError("Fluid populations must be finite with positive density")
             dense[:, ~fluid] = np.nan  # Solid cores are deliberately invalid.
-            blocked = self.layout.to_blocked(dense)
+            encoded = _encode_physical(dense, self._ddf_shifting)
+            values = encoded[:, fluid]
+            if (not np.isfinite(values).all()
+                    or (values.sum(axis=0, dtype=np.float64) + int(self._ddf_shifting) <= 0).any()):
+                raise ValueError("Stored fluid populations must be finite with positive density")
+            blocked = self.layout.to_blocked(encoded)
             spare = np.full_like(blocked, np.nan)
             try:
                 cuda.memcpy_htod(self._current, blocked)
@@ -225,7 +233,12 @@ class PeriodicBlockStepper:
                 raise
 
     def download(self):
-        """Return dense (9, ny, nx) FP32 core populations; no host cache."""
+        """Return physical core populations (FP64 when shifted); no host cache."""
+        snapshot = self._download_raw()
+        return _decode_stored(snapshot.values[0], snapshot.shifted)
+
+    def _download_raw(self):
+        """Return one tagged dense FP32 population array without conversion."""
         with self._lock, self.ctx.activate():
             self._require("ready")
             blocked = np.empty(self.layout.blocked_shape, np.float32)
@@ -234,7 +247,7 @@ class PeriodicBlockStepper:
             except BaseException:
                 self._status = "failed"
                 raise
-            return self.layout.to_dense(blocked)
+            return _StoredPopulations((self.layout.to_dense(blocked),), self._ddf_shifting)
 
     def close(self):
         """Release owned resources; retain them for retry if context activation fails."""

src/CelerisLab/lbm/population_storage.py

@@ -0,0 +1,24 @@
+"""Private D2Q9 FP32 storage boundary for the opt-in blocked steppers."""
+from dataclasses import dataclass
+import numpy as np
+
+# Unit-density reference weights in FP64; sum(w)=1 matches rho=1+sum(g).
+_WEIGHTS = np.array([4/9, 1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36],
+                    dtype=np.float64)[:, None, None]
+
+
+@dataclass(frozen=True, eq=False)
+class _StoredPopulations:
+    values: tuple
+    shifted: bool
+
+
+def _encode_physical(values, shifted):
+    if not shifted:
+        return np.array(values, dtype=np.float32, order="C", copy=True)
+    return np.ascontiguousarray(np.asarray(values, dtype=np.float64) - _WEIGHTS,
+                                dtype=np.float32)
+
+
+def _decode_stored(values, shifted):
+    return values.astype(np.float64) + _WEIGHTS if shifted else values

src/CelerisLab/lbm/refinement_stepper.py

@@ -25,6 +25,7 @@ from .refinement_tables import (blocked_base_tables, blocked_fine_tables,  # noq
                                 BLOCKED_BOUNDARY_DEVICE_COLUMNS)
 from .descriptors import SOLID
 from .raw_exchange import summarize_exchange
+from .population_storage import _StoredPopulations, _encode_physical, _decode_stored
 
 
 @dataclass
@@ -238,6 +239,9 @@ class FixedPatchStepper:
         level_blocks = [None if b is None else tuple(b) for b in level_blocks]
         _check_levels(refinements, level_blocks, level_objects, level_jets)
         depth = len(refinements) + 1
+        self._ddf_shifting = bool(cfg.ddf_shifting)
+        if self._ddf_shifting and any(b is None for b in level_blocks):
+            raise ValueError("shifted populations require all-blocked levels; dense/mixed is unsupported")
         self._wall_aware = coarse_boundary == "slot_reservoir"
         if self._wall_aware and (cfg.y_wall_bc != "on_site_velocity" or cfg.top_wall_bc != "reservoir"
                                  or cfg.velocity != 0 or not 2 <= cfg.wall_slot_width < cfg.nx):
@@ -469,23 +473,54 @@ class FixedPatchStepper:
         for array, leaf in zip(arrays, self._leaf):
             array[:, ~leaf] = np.nan
 
-    def initialize(self, *level_g, coarse_g=None, fine_g=None, child_g=None):
-        """Upload an explicit G_0, one Q-major array per level; non-leaves are discarded.
+    def _check_storage(self):
+        if any(bool(c.ddf_shifting) != self._ddf_shifting for c in self._level_cfgs):
+            raise ValueError("mixed population representations are unsupported")
+        if self._ddf_shifting and any(layout is None for layout in self._layouts):
+            raise ValueError("shifted populations require all-blocked levels")
 
-        Pass the level arrays positionally (L0 first), or use the original
-        two/three-level names coarse_g, fine_g[, child_g]; mixing is rejected.
+    def initialize(self, *level_g, coarse_g=None, fine_g=None, child_g=None):
+        """Upload physical postcollision G, encoding before any FP32 rounding.
+
+        Shifted input should be physical FP64 to retain small perturbations.
+        Public initialization resets the step count; non-leaves are discarded.
         """
         given = _initial_levels(level_g, coarse_g, fine_g, child_g)
         with self._lock, self.ctx.activate():
             self._require("new", "ready")
+            self._check_storage()
             if len(given) != len(self._levels):
                 raise ValueError("Expected one Q-major array per level")
-            arrays = [np.array(a, dtype=np.float32, order="C", copy=True) for a in given]
+            dtype = np.float64 if self._ddf_shifting else np.float32
+            arrays = [np.array(a, dtype=dtype, order="C", copy=True) for a in given]
             if any(a.shape != level.shape for a, level in zip(arrays, self._levels)):
                 raise ValueError("Expected Q-major (9,ny,nx) arrays for every level")
             for array, mask in zip(arrays, self._leaf):
                 if not np.isfinite(array[:, mask]).all() or (array[:, mask].sum(axis=0) <= 0).any():
                     raise ValueError("Leaf populations must be finite with positive density")
+            encoded = tuple(_encode_physical(a, self._ddf_shifting) for a in arrays)
+            self._initialize_raw(_StoredPopulations(encoded, self._ddf_shifting))
+
+    def _initialize_raw(self, snapshot):
+        """Upload an internally tagged same-representation state without a codec."""
+        with self._lock, self.ctx.activate():
+            self._require("new", "ready")
+            self._check_storage()
+            if (not isinstance(snapshot, _StoredPopulations)
+                    or type(snapshot.shifted) is not bool or snapshot.shifted != self._ddf_shifting):
+                raise ValueError("stored population representation mismatch")
+            if len(snapshot.values) != len(self._levels):
+                raise ValueError("Expected one stored array per level")
+            if any(not isinstance(a, np.ndarray) or a.dtype != np.float32 for a in snapshot.values):
+                raise ValueError("stored populations must be FP32")
+            arrays = [np.array(a, order="C", copy=True) for a in snapshot.values]
+            if any(a.shape != level.shape for a, level in zip(arrays, self._levels)):
+                raise ValueError("Expected Q-major (9,ny,nx) arrays for every level")
+            for array, mask in zip(arrays, self._leaf):
+                values = array[:, mask]
+                if (not np.isfinite(values).all()
+                        or (values.sum(axis=0, dtype=np.float64) + int(self._ddf_shifting) <= 0).any()):
+                    raise ValueError("Leaf populations must be finite with positive density")
             self._mask_nonleaves(*arrays)
             arrays = [a if level.layout is None else level.layout.to_blocked(a)
                       for a, level in zip(arrays, self._levels)]
@@ -518,7 +553,7 @@ class FixedPatchStepper:
         construction/upload errors never invalidate the old instance; failure
         to synchronize/read its device state does, as with step/download.
         """
-        from .synchronized_remap import (RemapTopology, remap_synchronized,
+        from .synchronized_remap import (RemapTopology, remap_synchronized, _remap_stored,
                                          same_physical_boundary, boundary_velocity)
 
         with self._lock:
@@ -558,14 +593,16 @@ class FixedPatchStepper:
             with self.ctx.activate():
                 try:
                     self._stream.synchronize()
-                    source = self.download()
+                    source = self._download_raw() if self._ddf_shifting else self.download()
                 except BaseException:
                     self._status = "failed"
                     raise
             count = self._step_count
-            transferred = remap_synchronized(
+            transfer = _remap_stored if self._ddf_shifting else remap_synchronized
+            options = dict(shifted=self._ddf_shifting) if self._ddf_shifting else {}
+            transferred = transfer(
                 old, new, source, base_count=count, level_steps=(count, 2*count, 4*count),
-                phase="postcollision_g")
+                phase="postcollision_g", **options)
             candidate = type(self)(
                 cfg, pairs[0].bounds, context=self.ctx, base_block=base_block,
                 fine_block=fine_block, child_block=child_block,
@@ -573,7 +610,10 @@ class FixedPatchStepper:
                 child_objects=child_objects, child_jets=jets,
                 observe_raw_exchange=self._observe_raw_exchange)
             try:
-                candidate.initialize(*transferred.g)
+                if self._ddf_shifting:
+                    candidate._initialize_raw(_StoredPopulations(transferred.g, transferred.shifted))
+                else:
+                    candidate.initialize(*transferred.g)
                 candidate._step_count = count
             except BaseException:
                 try:
@@ -772,9 +812,15 @@ class FixedPatchStepper:
                                       reference_center=self._exchange_centers[level][body_id])
 
     def download(self):
-        """Return synchronized G arrays (one per level), NaN on all non-leaves; no host cache."""
+        """Return physical G (FP64 for shifted storage), NaN on non-leaves."""
+        snapshot = self._download_raw()
+        return tuple(_decode_stored(a, snapshot.shifted) for a in snapshot.values)
+
+    def _download_raw(self):
+        """Current owned populations with a representation tag; no add/subtract w."""
         with self._lock, self.ctx.activate():
             self._require("ready")
+            self._check_storage()
             arrays = [np.empty(level.shape if level.layout is None else level.layout.blocked_shape,
                                np.float32) for level in self._levels]
             try:
@@ -786,7 +832,7 @@ class FixedPatchStepper:
             arrays = [a if level.layout is None else level.layout.to_dense(a)
                       for a, level in zip(arrays, self._levels)]
             self._mask_nonleaves(*arrays)
-            return tuple(arrays)
+            return _StoredPopulations(tuple(arrays), self._ddf_shifting)
 
     def close(self):
         """Release owned resources; retain them for retry if context activation fails."""

src/CelerisLab/lbm/synchronized_remap.py

@@ -16,6 +16,7 @@ from .refinement import (ACTIVE, FixedPatch2D, FixedPatchSet2D, FineBoundary2D,
                          nest_patch_set, prepare_fine_boundary)
 from .descriptors import SOLID
 from ..body.jets import surface_link_velocity
+from .population_storage import _StoredPopulations
 
 
 @dataclass(frozen=True, eq=False)
@@ -118,6 +119,7 @@ class RemapCandidate:
     phase: str
     origin: tuple
     spacing: float
+    shifted: bool = False
 
 
 def _same_parent(a, b):
@@ -126,7 +128,7 @@ def _same_parent(a, b):
             and a.tau_coarse == b.tau_coarse and a.tau_fine == b.tau_fine)
 
 
-def _read_leaves(topology, arrays, layouts):
+def _read_leaves(topology, arrays, layouts, *, shifted=False):
     """Unpack only owned cores; nonleaf values are not evidence of validity."""
     if len(arrays) != 3 or len(layouts) != 3:
         raise ValueError("expected three G arrays and three layouts")
@@ -140,8 +142,10 @@ def _read_leaves(topology, arrays, layouts):
             a = layout.to_dense(a)
         if a.shape != (9, *leaf.shape):
             raise ValueError("G must have Q-major (9, ny, nx) shape")
+        if shifted and a.dtype != np.float32:
+            raise ValueError("shifted stored populations must be FP32")
         values = a[:, leaf]
-        if not np.isfinite(values).all() or (values.sum(axis=0, dtype=np.float64) <= 0).any():
+        if not np.isfinite(values).all() or (values.sum(axis=0, dtype=np.float64) + int(shifted) <= 0).any():
             raise ValueError("required leaf donors must be finite with positive density")
         result.append(a)
     if len({a.dtype for a in result}) != 1:
@@ -166,6 +170,20 @@ def remap_synchronized(old, new, old_g, *, base_count, level_steps, phase,
     the source dtype. No equilibrium reconstruction or nonequilibrium scaling.
     All destination nonleaves are NaN. Inputs and their write flags are intact.
     """
+    return _remap(old, new, old_g, base_count=base_count, level_steps=level_steps,
+                  phase=phase, origin=origin, spacing=spacing, layouts=layouts)
+
+
+def _remap_stored(old, new, snapshot, *, shifted, **kwargs):
+    """Private same-representation transfer; public remap remains physical."""
+    if (not isinstance(snapshot, _StoredPopulations) or type(shifted) is not bool
+            or type(snapshot.shifted) is not bool or snapshot.shifted != shifted):
+        raise ValueError("stored population representation mismatch")
+    return _remap(old, new, snapshot.values, shifted=shifted, **kwargs)
+
+
+def _remap(old, new, old_g, *, base_count, level_steps, phase,
+           origin=(0., 0.), spacing=1., layouts=None, shifted=False):
     if (not isinstance(old, RemapTopology) or not isinstance(new, RemapTopology)
             or not _same_parent(old.parent, new.parent)
             or old.child.tau_fine != new.child.tau_fine):
@@ -189,7 +207,7 @@ def remap_synchronized(old, new, old_g, *, base_count, level_steps, phase,
     if (frame.shape != (2,) or not np.isfinite(frame).all()
             or not np.isscalar(spacing) or not np.isfinite(spacing) or spacing <= 0):
         raise ValueError("frame needs a finite (x,y) origin and positive spacing")
-    source = _read_leaves(old, old_g, (None,)*3 if layouts is None else layouts)
+    source = _read_leaves(old, old_g, (None,)*3 if layouts is None else layouts, shifted=shifted)
     candidate = [np.full((9, *mask.shape), np.nan, dtype=source[0].dtype)
                  for mask in new.leaf_masks]
     for level, (before, after) in enumerate(zip(old.leaf_masks, new.leaf_masks)):
@@ -210,6 +228,6 @@ def remap_synchronized(old, new, old_g, *, base_count, level_steps, phase,
     # Scale first to avoid an overflowing intermediate sum for finite FP64 G.
     mean = sum(source[2][:, cy, cx].astype(np.float64) * .25 for cy, cx in children)
     candidate[1][:, y, x] = mean
-    _read_leaves(new, candidate, (None,)*3)
+    _read_leaves(new, candidate, (None,)*3, shifted=shifted)
     return RemapCandidate(new, tuple(candidate), int(base_count), tuple(map(int, level_steps)),
-                          phase, tuple(map(float, frame)), float(spacing))
+                          phase, tuple(map(float, frame)), float(spacing), shifted)