Use direct momentum in opt-in shifted blocked equilibrium

docs/shifted_blocked_candidate.md

@@ -41,3 +41,21 @@ this is not full-solver validation. Host algebra/storage/remap tests and staged
 GPU continuation, three-grid cylinder physics, raw budgets, Couette helper
 stress and measured runtime are separate gates. Original failed experiments
 remain evidence; no acceptance threshold is relaxed.
+
+## Direct-input-momentum equilibrium candidate
+
+The shifted BlockStep now constructs the linear equilibrium term from input
+momentum j, without reconstructing rho*u. For opposite axis pairs the odd term
+is j/3; for diagonal pairs it is (jx +/- jy)/12, evaluated with explicit
+round-to-nearest division. The even term is
+w*[delta + 4.5*(c dot j)*(c dot u) - 1.5*(j dot u)].
+This is the same second-order D2Q9 equilibrium in exact arithmetic. Original SRT,
+storage, boundaries, and transfers are unchanged. No collision residual or force
+is corrected afterward. Existing four/five-argument equilibrium APIs are retained;
+only shifted blocked collision calls the new six-argument momentum helper.
+
+The all-fluid read-only CUDA replay reduced signed x-momentum error by 120–2285x
+at five sampled states, with lower mass/momentum and second-moment RMS errors.
+Signed mass bias was larger (~1e-12/cell/hop), so full evolving mass budgets remain
+mandatory. This local evidence does not establish physical convergence or speed.
+The unshifted default and its arithmetic path remain unchanged.

src/CelerisLab/lbm/kernels/blocks.cu

@@ -146,7 +146,9 @@ extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int hal
 #if USE_DDF_SHIFTING
             // Preserve the density perturbation before adding the unit reference.
             const float drho = f[0]+f[1]+f[2]+f[3]+f[4]+f[5]+f[6]+f[7]+f[8];
-            compute_feq(rho, ux, uy, feq, drho);
+            const float jx = f[1]-f[2]+f[5]-f[6]+f[7]-f[8];
+            const float jy = f[3]-f[4]+f[5]-f[6]-f[7]+f[8];
+            compute_feq_from_momentum(drho, jx, jy, ux, uy, feq);
 #else
             compute_feq(rho, ux, uy, feq);
 #endif

src/CelerisLab/lbm/kernels/operators/macro.cuh

@@ -113,6 +113,29 @@ __device__ __forceinline__ void compute_feq(float rho, float ux, float uy,
     #undef FEQ_CORE
 }
 
+#if USE_DDF_SHIFTING
+// Shifted D2Q9 equilibrium using input momentum in the linear term.
+// Algebraically identical to the usual second-order equilibrium; explicit
+// rational divisions avoid 3*w rounding and the rho*(j/rho) round trip.
+__device__ __forceinline__ void compute_feq_from_momentum(
+    float delta, float jx, float jy, float ux, float uy,
+    float* __restrict__ feq)
+{
+    const float common = fmaf(-1.5f, fmaf(jx, ux, jy * uy), delta);
+    feq[0] = W0_VAL * common;
+    const float jp[4] = {jx, jy, jx + jy, jx - jy};
+    const float up[4] = {ux, uy, ux + uy, ux - uy};
+    #pragma unroll
+    for (int p = 0; p < 4; ++p) {
+        const float weight = p < 2 ? WS_VAL : WE_VAL;
+        const float odd = __fdiv_rn(jp[p], p < 2 ? 3.0f : 12.0f);
+        const float even = weight * fmaf(4.5f * jp[p], up[p], common);
+        feq[1 + 2*p] = even + odd;
+        feq[2 + 2*p] = even - odd;
+    }
+}
+#endif
+
 #if USE_DDF_SHIFTING
 // Existing callers recover the offset from rho; blocked shifted collision
 // supplies sum(g) directly to the overload above, retaining small perturbations.

tests/unit/cpp/collision_regression.cpp

@@ -13,6 +13,8 @@
 #define OMEGA_COLLISION_MAX 1.99f
 struct Params {float omega, fx, fy, fz;} d_params;
 #include "core/descriptors.cuh"
+// Host round-to-nearest shim; CUDA intrinsic behavior is tested on-device.
+inline float __fdiv_rn(float a, float b) { return a / b; }
 #include "operators/macro.cuh"
 #include "operators/collision_srt.cuh"
 #include "operators/collision_trt.cuh"

tests/unit/cpp/semantic_regression.cpp

@@ -34,6 +34,8 @@ inline int obs_torque_index(int,int){return 0;}
 #include "core/descriptors.cuh"
 #include "core/precision.cuh"
 #include "core/layout.cuh"
+// Host round-to-nearest shim; CUDA intrinsic behavior is tested on-device.
+inline float __fdiv_rn(float a, float b) { return a / b; }
 #include "operators/macro.cuh"
 #include "operators/collision_srt.cuh"
 #include "operators/collision_trt.cuh"

tests/unit/cpp/shifted_block_operators.cpp

@@ -13,6 +13,8 @@
 #define WALL_SLOT_WIDTH 0
 #define NX 16
 #include "core/descriptors.cuh"
+// Host round-to-nearest shim; CUDA intrinsic behavior is tested on-device.
+inline float __fdiv_rn(float a, float b) { return a / b; }
 #include "operators/macro.cuh"
 #include "operators/collision_srt.cuh"
 #include "boundary/planar_velocity.cuh"
@@ -23,6 +25,13 @@ int main() {
         float eq[9]; compute_feq(1.f+delta,u,-.2f*u,eq,delta);
         printf("E %.9g %.9g %.9g\n",delta,u,-.2f*u);emit(eq);
     }
+    for(float delta : {0.f, 0x1p-30f, -0x1p-30f, .04f, -.01f})
+    for(float jx : {0.f,.005f,-.02f}) {
+        const float jy = -.2f*jx, rho = 1.f+delta;
+        const float ux = jx/rho, uy = jy/rho;
+        float eq[9];compute_feq_from_momentum(delta,jx,jy,ux,uy,eq);
+        printf("J %.9g %.9g %.9g %.9g %.9g\n",delta,jx,jy,ux,uy);emit(eq);
+    }
     for(int mode : {0,1,2}) for(float rho : {1.f,1.04f}) {
         float f[9];const float u=.007f,v=.003f;
         compute_feq(rho,u,v,f,rho-1.f);

tests/unit/test_shifted_block_operators.py

@@ -19,6 +19,15 @@ def test_retained_density_and_shifted_planar_headers(tmp_path):
             if u==0:
                 np.testing.assert_allclose(a.sum(),delta,rtol=5e-8,atol=0)
                 if delta: assert np.any(a!=0)
+        elif kind=='J':
+            delta,jx,jy,ux,uy=np.array(numbers,dtype=np.float32).astype(float)
+            jd=cx*jx+cy*jy;ud=cx*ux+cy*uy
+            expected=w*(delta+3*jd+4.5*jd*ud-1.5*(jx*ux+jy*uy))
+            np.testing.assert_allclose(a,expected,rtol=3e-6,atol=2e-9)
+            actual=np.array([a.sum(),a@cx,a@cy,a@(cx*cx),a@(cx*cy),a@(cy*cy)])
+            target=np.array([delta,jx,jy,delta/3+jx*ux,jx*uy,delta/3+jy*uy])
+            np.testing.assert_allclose(actual,target,atol=4e-9,rtol=3e-6)
+            if jx==0 and delta: assert np.any(a!=0)
         else:
             mode=int(numbers[0]);rho,u,v=np.array(numbers[1:],dtype=np.float32).astype(float)
             b=np.fromstring(next(it),sep=' ',dtype=np.float32).astype(float)