Use direct momentum in opt-in shifted blocked equilibrium

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.