Add research measured-strain fluid refinement example

research/shear_refinement/example.py

@@ -0,0 +1,140 @@
+"""One measured L1 strain observation, one growth-only L2 remap, no scheduler.
+
+Run from the repository root: PYTHONPATH=src:. python -m research.shear_refinement.example
+Requires CUDA; CPU helpers and tests never import CUDA.
+"""
+from contextlib import ExitStack
+from numbers import Integral
+import numpy as np
+
+from CelerisLab.config import LBMConfig
+from CelerisLab.body.geometry.circle import CircleGeometry
+from CelerisLab.body.objects import SimObject
+from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
+from CelerisLab.lbm.refinement import FixedPatchSet2D
+from CelerisLab.lbm.synchronized_remap import RemapTopology
+from .boundary_shear_tagging import tag_cells
+
+CFG = LBMConfig(nx=48, ny=32, viscosity=.1)
+PARENTS = ((4, 4, 20, 28), (24, 4, 44, 28))
+CHILDREN = ((16, 16, 32, 40), (56, 16, 72, 40))
+CORRIDOR = (72, 20, 76, 36)  # L1 allocated indices; bounded to right of fluid L2
+
+
+def grow_rectangle(current, request, *, alignment=4, area_cap=480):
+    """Bounding-box growth only, in request-array cell-edge coordinates.
+
+    No body or solver knowledge. Reject invalid/budget-exceeding proposals with
+    ValueError; never clip a request, shrink old coverage or mutate an input.
+    Alignment is edge divisibility, not a substitute for full topology guards.
+    """
+    request = np.asarray(request)
+    if request.dtype != bool or request.ndim != 2:
+        raise ValueError('request must be a 2D boolean mask')
+    values = (*current, alignment, area_cap)
+    if len(current) != 4 or any(isinstance(v, (bool, np.bool_)) or not isinstance(v, Integral) for v in values):
+        raise ValueError('rectangle, alignment and area cap must be integers')
+    x0, y0, x1, y1 = current
+    ny, nx = request.shape
+    if alignment <= 0 or area_cap <= 0 or not (0 <= x0 < x1 <= nx and 0 <= y0 < y1 <= ny):
+        raise ValueError('invalid rectangle, alignment or area cap')
+    if any(v % alignment for v in current):
+        raise ValueError('current rectangle must be aligned')
+    y, x = np.nonzero(request)
+    if x.size:
+        x0, y0 = min(x0, int(x.min())), min(y0, int(y.min()))
+        x1, y1 = max(x1, int(x.max())+1), max(y1, int(y.max())+1)
+    a = alignment
+    result = (x0//a*a, y0//a*a, (x1+a-1)//a*a, (y1+a-1)//a*a)
+    if result[2] > nx or result[3] > ny:
+        raise ValueError('aligned growth exceeds observation frame')
+    if (result[2]-result[0])*(result[3]-result[1]) > area_cap:
+        raise ValueError('growth exceeds area cap')
+    return result
+
+
+def objects():
+    return (SimObject(71, geometry=CircleGeometry(12., 14., .8),
+                      center=(12., 14.), radius=.8),)
+
+
+def topology(rectangles=CHILDREN):
+    return RemapTopology.from_config(CFG, FixedPatchSet2D.from_config(CFG, PARENTS),
+                                     rectangles, child_objects=objects())
+
+
+def observe(g, mesh, *, shear_on=.001):
+    """Moments of completed-time postcollision G; only real L1 fluid leaves.
+
+    No reconstructed covered values, ghosts or holes are derivative donors.
+    solid is independent geometry (empty on L1: the body is pinned at L2).
+    The array is cropped/sparse, so neither derivative axis wraps.
+    """
+    valid = mesh.leaf_masks[1]
+    f = g[:, valid].astype(np.float64)
+    rho = f.sum(axis=0)
+    if not np.isfinite(f).all() or np.any(rho <= 0):
+        raise ValueError('invalid completed-time L1 fluid populations')
+    ux, uy = (np.full(valid.shape, np.nan) for _ in range(2))
+    ux[valid] = (np.asarray(D2Q9_EX)[:, None]*f).sum(axis=0)/rho
+    uy[valid] = (np.asarray(D2Q9_EY)[:, None]*f).sum(axis=0)/rho
+    tags = tag_cells(ux, uy, np.zeros(valid.shape, bool), valid=valid, owned=valid,
+                     dx_level=.5, origin=mesh.parent.fine_origin, periodic=(False, False),
+                     U_ref=.05, L_ref=1., shear_on=shear_on, shear_off=0.,
+                     band_on=0., band_off=0., n_buffer=0)
+    corridor = np.zeros(valid.shape, bool)
+    x0, y0, x1, y1 = CORRIDOR
+    corridor[y0:y1, x0:x1] = True
+    request = tags.tagged & corridor
+    proposed = (CHILDREN[0], grow_rectangle(CHILDREN[1], request))
+    # Host guards cover nesting, donor footprints, complete body and separation.
+    candidate = topology(proposed)
+    candidate.child.fine_layout((4, 4))
+    report = dict(requested=int(request.sum()), unknown=int(tags.unevaluable.sum()),
+                  corridor_unknown=int((tags.unevaluable & corridor).sum()),
+                  corridor_unavailable=int((corridor & ~valid).sum()),
+                  out_of_scope=int((tags.tagged & ~corridor).sum()), rectangles=proposed)
+    return proposed, report
+
+
+def initial_fields(mesh):
+    """Smooth periodic shear wave, sampled in each level's physical frame."""
+    from CelerisLab.lbm.equilibrium import compute_feq_d2q9
+    frames = ((np.zeros(2), 1.), (np.asarray(mesh.parent.fine_origin), .5),
+              (np.asarray(mesh.parent.fine_origin)+.5*np.asarray(mesh.child.fine_origin), .25))
+    arrays = []
+    for leaf, (origin, dx) in zip(mesh.leaf_masks, frames):
+        y = origin[1]+dx*(np.arange(leaf.shape[0])+.5)
+        ux = .01*np.sin(2*np.pi*y/CFG.ny)
+        rows = np.stack([compute_feq_d2q9(1., u, 0.) for u in ux], axis=1)
+        arrays.append(np.broadcast_to(rows[:, :, None], (9, *leaf.shape)).copy())
+    return tuple(arrays)
+
+
+def initialize(run):
+    run.initialize(*initial_fields(topology()))
+
+
+def make(context):
+    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
+    return FixedPatchStepper(CFG, PARENTS, context=context, base_block=(4, 4),
+                             fine_block=(4, 4), child_block=(4, 4),
+                             child_rectangles_in_parent=CHILDREN, child_objects=objects())
+
+
+def main():
+    from CelerisLab.cuda.context import CudaContext
+    with CudaContext(0) as ctx, ExitStack() as owned:
+        run = owned.enter_context(make(ctx))
+        initialize(run)
+        run.step(2)  # Returns only at completed synchronized base time.
+        proposed, report = observe(run.download()[1], topology())
+        print(report)
+        if proposed != CHILDREN:
+            run = owned.enter_context(run.remap_child(proposed, child_objects=objects()))
+        run.step(2)
+        print(f'completed base step {run.step_count}; body rectangle {run.child_patch.bounds[0]}')
+
+
+if __name__ == '__main__':
+    main()