Add synchronized CPU leaf population remap

src/CelerisLab/lbm/synchronized_remap.py

@@ -0,0 +1,147 @@
+"""CPU-only whole-volume G transfer, not a runtime adaptive solver.
+
+Only body-free periodic three-level grids with a fixed L0/L1 topology are
+supported. The caller certifies that current owned leaves are quiescent,
+unshifted D2Q9 postcollision G at the supplied completed base interval.
+No device buffers, boundary state, ghosts, collision or time advance occur.
+"""
+from dataclasses import dataclass
+from numbers import Integral
+
+import numpy as np
+
+from .refinement import ACTIVE, FixedPatch2D, FixedPatchSet2D, nest_patch_set
+
+
+@dataclass(frozen=True, eq=False)
+class RemapTopology:
+    """Validated fixed parent and separated L2 set; use ``from_config``.
+
+    Child rectangles are in the parent's allocated L1 coordinates, including
+    its window offset/halo when the parent is a scalar patch. Dense masks are
+    rebuilt from geometry, never inferred from finite values in input arrays.
+    """
+    parent: object
+    child: FixedPatchSet2D
+    leaf_masks: tuple
+
+    @classmethod
+    def from_config(cls, cfg, parent, child_rectangles, *, coarse_boundary="periodic", objects=()):
+        if coarse_boundary != "periodic" or tuple(objects) or cfg.nz != 1:
+            raise ValueError("remap supports only body-free periodic 2D state")
+        if type(parent) not in (FixedPatch2D, FixedPatchSet2D):
+            raise ValueError("parent must be a fixed patch or fixed patch set")
+        parent = type(parent).from_config(cfg, parent.bounds)
+        child, middle_kind = nest_patch_set(cfg, parent, child_rectangles)
+        masks = (~parent.covered, middle_kind == ACTIVE, child.fine_kind == ACTIVE)
+        for mask in masks:
+            mask.setflags(write=False)
+        return cls(parent, child, masks)
+
+
+@dataclass(frozen=True, eq=False)
+class RemapCandidate:
+    """Separate dense arrays plus explicit synchronization/frame metadata.
+
+    Passing these arrays to the existing stepper.initialize resets its count;
+    this candidate is deliberately NOT an installed runtime continuation.
+    """
+    topology: RemapTopology
+    g: tuple
+    base_count: int
+    level_steps: tuple
+    phase: str
+    origin: tuple
+    spacing: float
+
+
+def _same_parent(a, b):
+    return (type(a) is type(b) and a.bounds == b.bounds
+            and a.coarse_shape == b.coarse_shape
+            and a.tau_coarse == b.tau_coarse and a.tau_fine == b.tau_fine)
+
+
+def _read_leaves(topology, arrays, layouts):
+    """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")
+    result = []
+    for a, layout, leaf in zip(arrays, layouts, topology.leaf_masks):
+        if not isinstance(a, np.ndarray) or a.dtype not in (np.dtype('float32'), np.dtype('float64')):
+            raise ValueError("G arrays must be FP32 or FP64 NumPy arrays")
+        if layout is not None:
+            if layout.shape != leaf.shape or not layout.owned[leaf].all():
+                raise ValueError("layout must own every required leaf donor")
+            a = layout.to_dense(a)
+        if a.shape != (9, *leaf.shape):
+            raise ValueError("G must have Q-major (9, ny, nx) shape")
+        values = a[:, leaf]
+        if not np.isfinite(values).all() or (values.sum(axis=0, dtype=np.float64) <= 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:
+        raise ValueError("all levels must use the same population dtype")
+    return result
+
+
+def remap_synchronized(old, new, old_g, *, base_count, level_steps, phase,
+                       origin=(0., 0.), spacing=1., layouts=None):
+    """Return a non-mutating candidate at the SAME physical time and frame.
+
+    ``old``/``new`` are RemapTopology.from_config results. L0/L1 geometry and
+    relaxation parameters must match; only the L2 rectangle set may change.
+    ``old_g`` contains current owned-leaf G, dense or (with ``layouts``) blocked.
+    ``level_steps`` must equal (base_count, 2*base_count, 4*base_count), and
+    ``phase`` must be 'postcollision_g'. These stamps are caller assertions,
+    not proof of device quiescence or active-buffer freshness.
+
+    Preserve same-level leaf overlap bitwise. Copy all nine parent G values
+    to newly refined children; average ALL four old children (including q0)
+    for newly exposed L1. Restriction accumulates in FP64, then rounds once to
+    the source dtype. No equilibrium reconstruction or nonequilibrium scaling.
+    All destination nonleaves are NaN. Inputs and their write flags are intact.
+    """
+    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):
+        raise ValueError("remap requires unchanged L0/L1 geometry and relaxation")
+    if not isinstance(level_steps, (tuple, list)):
+        raise ValueError("level_steps must be three completed integer counts")
+    counts = (base_count, *level_steps)
+    if (len(counts) != 4 or any(isinstance(n, (bool, np.bool_))
+                              or not isinstance(n, Integral) or n < 0 for n in counts)
+            or tuple(level_steps) != (base_count, 2*base_count, 4*base_count)
+            or phase != 'postcollision_g'):
+        raise ValueError("requires synchronized completed counts and postcollision_g")
+    try:
+        frame = np.asarray(origin, dtype=np.float64)
+        spacing = float(spacing)
+    except (TypeError, ValueError):
+        raise ValueError("frame needs a finite (x,y) origin and positive spacing") from None
+    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)
+    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)):
+        overlap = before & after
+        candidate[level][:, overlap] = source[level][:, overlap]
+
+    # A child set uses a level-wide grid in its parent's allocated coordinates.
+    # Thus child (y,x) belongs to parent (y//2,x//2), even for a shifted L1 window.
+    y, x = np.nonzero(new.leaf_masks[2] & ~old.leaf_masks[2])
+    if not old.leaf_masks[1][y//2, x//2].all():
+        raise ValueError("new L2 requires an old owned L1 leaf")
+    candidate[2][:, y, x] = source[1][:, y//2, x//2]
+
+    y, x = np.nonzero(new.leaf_masks[1] & ~old.leaf_masks[1])
+    children = [(2*y+dy, 2*x+dx) for dy in (0, 1) for dx in (0, 1)]
+    if any(not old.leaf_masks[2][cy, cx].all() for cy, cx in children):
+        raise ValueError("reactivated L1 requires all four old owned L2 leaves")
+    # 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)
+    return RemapCandidate(new, tuple(candidate), int(base_count), tuple(map(int, level_steps)),
+                          phase, tuple(map(float, frame)), float(spacing))