clean up and fixes

cases/db.yaml

@@ -9,7 +9,7 @@ case:
   viscosity: 0.00005
   special:
     L_wall: 5.366
-    H_wall: 5.366 #2.0
+    H_wall: 2.0
     L: 2.0  # water column length
     H: 1.0  # water column height
     W: 0.2  # width in 3D case

jax_sph/solver.py

@@ -57,18 +57,22 @@ def rho_evol_riemann_fn(
         kernel_grad = kernel_fn.grad_w(d_ij) * (e_ij)
 
         # Compute average states eq. (6)/(12)/(13), Zhang (2017)
-        u_L = jnp.where(wall_mask_j == 1, jnp.dot(u_i, -n_w_j), jnp.dot(u_i, -e_ij))
+        u_L = jnp.where(
+            jnp.isin(wall_mask_j, wall_tags), jnp.dot(u_i, -n_w_j), jnp.dot(u_i, -e_ij)
+        )
         p_L = p_i
         rho_L = rho_i
 
         #  u_w from eq. (15), Yang (2020)
         u_R = jnp.where(
-            wall_mask_j == 1,
-            -u_L + 2 * jnp.dot(u_j, -n_w_j),
+            jnp.isin(wall_mask_j, wall_tags),
+            -u_L + 2 * jnp.dot(u_j, n_w_j),
             jnp.dot(u_j, -e_ij),
         )
-        p_R = jnp.where(wall_mask_j == 1, p_L + rho_L * jnp.dot(g_ext_i, -r_ij), p_j)
-        rho_R = jnp.where(wall_mask_j == 1, eos.rho_fn(p_R), rho_j)
+        p_R = jnp.where(
+            jnp.isin(wall_mask_j, wall_tags), p_L + rho_L * jnp.dot(g_ext_i, -r_ij), p_j
+        )
+        rho_R = jnp.where(jnp.isin(wall_mask_j, wall_tags), eos.rho_fn(p_R), rho_j)
 
         U_avg = (u_L + u_R) / 2
         v_avg = (u_i + u_j) / 2
@@ -208,18 +212,22 @@ def acceleration_fn_riemann(
         kernel_grad = kernel_part_diff * (e_ij)
 
         # Compute average states eq. (6)/(12)/(13), Zhang (2017)
-        u_L = jnp.where(wall_mask_j == 1, jnp.dot(u_i, -n_w_j), jnp.dot(u_i, -e_ij))
+        u_L = jnp.where(
+            jnp.isin(wall_mask_j, wall_tags), jnp.dot(u_i, -n_w_j), jnp.dot(u_i, -e_ij)
+        )
         p_L = p_i
         rho_L = rho_i
 
         # u_w from eq. (15), Yang (2020)
         u_R = jnp.where(
-            wall_mask_j == 1,
-            -u_L + 2 * jnp.dot(u_j, -n_w_j),
+            jnp.isin(wall_mask_j, wall_tags),
+            -u_L + 2 * jnp.dot(u_j, n_w_j),
             jnp.dot(u_j, -e_ij),
         )
-        p_R = jnp.where(wall_mask_j == 1, p_L + rho_L * jnp.dot(g_ext_i, -r_ij), p_j)
-        rho_R = jnp.where(wall_mask_j == 1, eos.rho_fn(p_R), rho_j)
+        p_R = jnp.where(
+            jnp.isin(wall_mask_j, wall_tags), p_L + rho_L * jnp.dot(g_ext_i, -r_ij), p_j
+        )
+        rho_R = jnp.where(jnp.isin(wall_mask_j, wall_tags), eos.rho_fn(p_R), rho_j)
 
         P_avg = (p_L + p_R) / 2
         rho_avg = (rho_L + rho_R) / 2
@@ -229,9 +237,6 @@ def acceleration_fn_riemann(
         eta_ij = 2 * eta_i * eta_j / (eta_i + eta_j + EPS)
 
         # Compute Riemann states eq. (7) and (10), Zhang (2017)
-        # u_R = jnp.where(
-        # wall_mask_j == 1, -u_L - 2 * jnp.dot(v_j, -n_w_j), jnp.dot(v_j, -e_ij)
-        # )
         P_star = P_avg + 0.5 * rho_avg * (u_L - u_R) * beta_fn(u_L, u_R, eta_limiter)
 
         # pressure term with linear Riemann solver eq. (9), Zhang (2017)
@@ -240,7 +245,7 @@ def acceleration_fn_riemann(
         # viscosity term eq. (6), Zhang (2019)
         u_d = 2 * u_j - u_tilde_j
         v_ij = jnp.where(
-            wall_mask_j == 1,
+            jnp.isin(wall_mask_j, wall_tags),
             u_i - u_d,
             u_i - u_j,
         )
@@ -396,14 +401,14 @@ def gwbc_fn_riemann_wrapper(is_free_slip, is_heat_conduction):
         def free_weight(fluid_mask_i, tag_i):
             return fluid_mask_i
 
-        def Riemann_velocities(u, w_dist, fluid_mask, i_s, j_s, N):
+        def riemann_velocities(u, w_dist, fluid_mask, i_s, j_s, N):
             return u
     else:
 
         def free_weight(fluid_mask_i, tag_i):
             return jnp.ones_like(tag_i)
 
-        def Riemann_velocities(u, w_dist, fluid_mask, i_s, j_s, N):
+        def riemann_velocities(u, w_dist, fluid_mask, i_s, j_s, N):
             w_dist_fluid = w_dist * fluid_mask[j_s]
             u_wall_nom = ops.segment_sum(w_dist_fluid[:, None] * u[j_s], i_s, N)
             u_wall_denom = ops.segment_sum(w_dist_fluid, i_s, N)
@@ -427,7 +432,7 @@ def heat_bc(mask_j_s_fluid, w_dist, temperature, i_s, j_s, tag, N):
         def heat_bc(mask_j_s_fluid, w_dist, temperature, i_s, j_s, tag, N):
             return temperature
 
-    return free_weight, Riemann_velocities, heat_bc
+    return free_weight, riemann_velocities, heat_bc
 
 
 def limiter_fn_wrapper(eta_limiter, c_ref):
@@ -522,7 +527,7 @@ def __init__(
         self._gwbc_fn = gwbc_fn_wrapper(is_free_slip, is_heat_conduction, eos)
         (
             self._free_weight,
-            self._Riemann_velocities,
+            self._riemann_velocities,
             self._heat_bc,
         ) = gwbc_fn_riemann_wrapper(is_free_slip, is_heat_conduction)
         self._acceleration_tvf_fn = acceleration_tvf_fn_wrapper(self._kernel_fn)
@@ -593,7 +598,7 @@ def forward(state, neighbors):
 
             ##### Riemann velocity BCs
             if self.is_bc_trick and (self.solver == "RIE"):
-                u_tilde = self._Riemann_velocities(u, w_dist, fluid_mask, i_s, j_s, N)
+                u_tilde = self._riemann_velocities(u, w_dist, fluid_mask, i_s, j_s, N)
 
             ##### Density summation or evolution
 
@@ -645,7 +650,6 @@ def forward(state, neighbors):
                 )
             elif self.is_bc_trick and (self.solver == "RIE"):
                 mask = self._free_weight(fluid_mask[i_s], tag[i_s])
-                # u_tilde = self._Riemann_velocities(u, w_dist, fluid_mask, i_s, j_s, N)
                 temperature = self._heat_bc(
                     fluid_mask[j_s], w_dist, temperature, i_s, j_s, tag, N
                 )

tests/test_pf2d.py

@@ -108,8 +108,6 @@ def test_pf2d(tvf, solver, tmp_path, setup_simulation):
     """Test whether the poiseuille flow simulation matches the analytical solution"""
     y_axis, t_dimless, ref_solutions = setup_simulation
     data_path = run_simulation(tmp_path, tvf, solver)
-    # print(f"tmp_path = {tmp_path}, subdirs = {subdirs}")
     solutions = get_solution(data_path, t_dimless, y_axis)
-    # print(f"solution: {solutions[-1]} \nref_solution: {ref_solutions[-1]}")
     for sol, ref_sol in zip(solutions, ref_solutions):
         assert np.allclose(sol, ref_sol, atol=1e-2), "Velocity profile does not match."