Riemann tvf BC fix

tumaer · Jun 13, 2024 · ac1f31b · ac1f31b
1 parent 32ce4fb
commit ac1f31b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/jax_sph/solver.py b/jax_sph/solver.py
@@ -257,7 +257,11 @@ def acceleration_fn_riemann(
         _kernel_grad = kernel_fn.grad_w(d_ij)
         _c = _weighted_volume * _kernel_grad / (d_ij + EPS)
 
-        _A = (tvf_stress_fn(rho_i, u_i, u_i) + tvf_stress_fn(rho_j, u_j, u_j)) / 2
+        _A = jnp.where(
+            jnp.isin(wall_mask_j, wall_tags),
+            (tvf_stress_fn(rho_i, u_i, u_i) + tvf_stress_fn(rho_j, u_d, u_d)) / 2,
+            (tvf_stress_fn(rho_i, u_i, u_i) + tvf_stress_fn(rho_j, u_j, u_j)) / 2,
+        )
         a_eq_8 = _c * jnp.dot(_A, r_ij)
 
         return eq_9 + eq_6 + a_eq_8