From fa164353bd245e0cec243ead72252fc48282eab9 Mon Sep 17 00:00:00 2001 From: Jonas Erbesdobler Date: Mon, 10 Jun 2024 01:08:03 +0200 Subject: [PATCH] clean up and fixes --- cases/db.yaml | 2 +- jax_sph/solver.py | 44 ++++++++++++++++++++++++-------------------- tests/test_pf2d.py | 2 -- 3 files changed, 25 insertions(+), 23 deletions(-) diff --git a/cases/db.yaml b/cases/db.yaml index 4ad6e13..d0748c4 100644 --- a/cases/db.yaml +++ b/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 diff --git a/jax_sph/solver.py b/jax_sph/solver.py index 9902f7b..93254fb 100644 --- a/jax_sph/solver.py +++ b/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 ) diff --git a/tests/test_pf2d.py b/tests/test_pf2d.py index 1a8a58d..6c8b7cb 100644 --- a/tests/test_pf2d.py +++ b/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."