Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fixed a few 2D bugs #85

Merged
merged 2 commits into from
Oct 31, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 12 additions & 5 deletions examples/cfd/lid_driven_cavity_2d.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@


class LidDrivenCavity2D:
def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
def __init__(self, omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy):
# initialize backend
xlb.init(
velocity_set=velocity_set,
Expand All @@ -29,6 +29,7 @@ def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
self.grid, self.f_0, self.f_1, self.missing_mask, self.bc_mask = create_nse_fields(grid_shape)
self.stepper = None
self.boundary_conditions = []
self.prescribed_vel = prescribed_vel

# Setup the simulation BC, its initial conditions, and the stepper
self._setup(omega)
Expand All @@ -49,7 +50,7 @@ def define_boundary_indices(self):

def setup_boundary_conditions(self):
lid, walls = self.define_boundary_indices()
bc_top = EquilibriumBC(rho=1.0, u=(0.02, 0.0), indices=lid)
bc_top = EquilibriumBC(rho=1.0, u=(self.prescribed_vel, 0.0), indices=lid)
bc_walls = HalfwayBounceBackBC(indices=walls)
self.boundary_conditions = [bc_walls, bc_top]

Expand Down Expand Up @@ -112,7 +113,13 @@ def post_process(self, i):
precision_policy = PrecisionPolicy.FP32FP32

velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend)
omega = 1.6

simulation = LidDrivenCavity2D(omega, grid_shape, velocity_set, backend, precision_policy)
simulation.run(num_steps=5000, post_process_interval=1000)
# Setting fluid viscosity and relaxation parameter.
Re = 200.0
prescribed_vel = 0.05
clength = grid_shape[0] - 1
visc = prescribed_vel * clength / Re
omega = 1.0 / (3.0 * visc + 0.5)

simulation = LidDrivenCavity2D(omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy)
simulation.run(num_steps=50000, post_process_interval=1000)
mehdiataei marked this conversation as resolved.
Show resolved Hide resolved
16 changes: 11 additions & 5 deletions examples/cfd/lid_driven_cavity_2d_distributed.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@


class LidDrivenCavity2D_distributed(LidDrivenCavity2D):
def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
super().__init__(omega, grid_shape, velocity_set, backend, precision_policy)
def __init__(self, omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy):
super().__init__(omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy)

def setup_stepper(self, omega):
stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions)
Expand All @@ -29,7 +29,13 @@ def setup_stepper(self, omega):
precision_policy = PrecisionPolicy.FP32FP32

velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend)
omega = 1.6

simulation = LidDrivenCavity2D_distributed(omega, grid_shape, velocity_set, backend, precision_policy)
simulation.run(num_steps=5000, post_process_interval=1000)
# Setting fluid viscosity and relaxation parameter.
Re = 200.0
prescribed_vel = 0.05
clength = grid_shape[0] - 1
visc = prescribed_vel * clength / Re
omega = 1.0 / (3.0 * visc + 0.5)

simulation = LidDrivenCavity2D_distributed(omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy)
simulation.run(num_steps=50000, post_process_interval=1000)
10 changes: 6 additions & 4 deletions tests/boundary_conditions/mask/test_bc_indices_masker_warp.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,6 @@ def test_indices_masker_warp(dim, velocity_set, grid_shape):
[test_bc],
bc_mask,
missing_mask,
start_index=(0, 0, 0) if dim == 3 else (0, 0),
)
assert missing_mask.dtype == xlb.Precision.BOOL.wp_dtype

Expand All @@ -69,9 +68,12 @@ def test_indices_masker_warp(dim, velocity_set, grid_shape):
bc_mask = bc_mask.numpy()
missing_mask = missing_mask.numpy()

assert bc_mask.shape == (1,) + grid_shape

assert missing_mask.shape == (velocity_set.q,) + grid_shape
if len(grid_shape) == 2:
assert bc_mask.shape == (1,) + grid_shape + (1,), "bc_mask shape is incorrect got {}".format(bc_mask.shape)
assert missing_mask.shape == (velocity_set.q,) + grid_shape + (1,), "missing_mask shape is incorrect got {}".format(missing_mask.shape)
else:
assert bc_mask.shape == (1,) + grid_shape, "bc_mask shape is incorrect got {}".format(bc_mask.shape)
assert missing_mask.shape == (velocity_set.q,) + grid_shape, "missing_mask shape is incorrect got {}".format(missing_mask.shape)

if dim == 2:
assert np.all(bc_mask[0, indices[0], indices[1]] == test_bc.id)
Expand Down
12 changes: 8 additions & 4 deletions xlb/operator/boundary_condition/bc_extrapolation_outflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,6 @@ def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask):
def _construct_warp(self):
# Set local constants
sound_speed = self.compute_dtype(1.0 / wp.sqrt(3.0))
_f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
_c = self.velocity_set.c
_q = self.velocity_set.q
_opp_indices = self.velocity_set.opp_indices
Expand All @@ -143,9 +142,14 @@ def _construct_warp(self):
def get_normal_vectors(
missing_mask: Any,
):
for l in range(_q):
if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
return -wp.vec3i(_c[0, l], _c[1, l], _c[2, l])
if wp.static(self.velocity_set.d == 3):
for l in range(_q):
if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
return -wp.vec3i(_c[0, l], _c[1, l], _c[2, l])
else:
for l in range(_q):
if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) == 1:
return -wp.vec2i(_c[0, l], _c[1, l])

# Construct the functionals for this BC
@wp.func
Expand Down
12 changes: 8 additions & 4 deletions xlb/operator/boundary_condition/bc_regularized.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,6 @@ def _construct_warp(self):
# Set local constants TODO: This is a hack and should be fixed with warp update
# _u_vec = wp.vec(_d, dtype=self.compute_dtype)
# compute Qi tensor and store it in self
_f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
_u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
_rho = self.compute_dtype(rho)
_u = _u_vec(u[0], u[1], u[2]) if _d == 3 else _u_vec(u[0], u[1])
Expand Down Expand Up @@ -162,9 +161,14 @@ def _get_fsum(
def get_normal_vectors(
missing_mask: Any,
):
for l in range(_q):
if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l])
if wp.static(_d == 3):
for l in range(_q):
if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This does not make sense. As we only have 3D vectors in Warp.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's not true. The lattice velocity "c" is still represents a 2D vector per each lattice dimenion for a 2D lattice and a 3D vector per each lattice dimension for a 3D lattice.

return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l])
else:
for l in range(_q):
if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) == 1:
return -_u_vec(_c_float[0, l], _c_float[1, l])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We're not working with 2D vectors anymore in the Warp backend and their use should be avoided. The vectors shouldn't depend on the grid dimension. I guess this is being done because this is a normal. But the fix should be done more fundamentally to avoid downstream issues.

For example in 2D, you may have only (1,0,0) and (0,1,0),...Why would this not be working? I think the change should be down in the code that utilizes this normals later.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please see the above comment.


@wp.func
def bounceback_nonequilibrium(
Expand Down
11 changes: 8 additions & 3 deletions xlb/operator/boundary_condition/bc_zouhe.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,9 +207,14 @@ def _get_fsum(
def get_normal_vectors(
missing_mask: Any,
):
for l in range(_q):
if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l])
if wp.static(_d == 3):
for l in range(_q):
if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l])
else:
for l in range(_q):
if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) == 1:
return -_u_vec(_c_float[0, l], _c_float[1, l])

@wp.func
def bounceback_nonequilibrium(
Expand Down
20 changes: 12 additions & 8 deletions xlb/operator/collision/kbc.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ def jax_implementation(
fneq = f - feq
if isinstance(self.velocity_set, D2Q9):
shear = self.decompose_shear_d2q9_jax(fneq)
delta_s = shear * rho / 4.0 # TODO: Check this
delta_s = shear * rho / 4.0
elif isinstance(self.velocity_set, D3Q27):
shear = self.decompose_shear_d3q27_jax(fneq)
delta_s = shear * rho
Expand Down Expand Up @@ -191,16 +191,16 @@ def _construct_warp(self):
@wp.func
def decompose_shear_d2q9(fneq: Any):
pi = self.momentum_flux.warp_functional(fneq)
N = pi[0] - pi[1]
N = pi[0] - pi[2]
s = _f_vec()
s[3] = N
s[6] = N
s[2] = -N
s[1] = -N
s[8] = pi[2]
s[4] = -pi[2]
s[5] = -pi[2]
s[7] = pi[2]
s[8] = pi[1]
s[4] = -pi[1]
s[5] = -pi[1]
s[7] = pi[1]
return s

# Construct functional for decomposing shear
Expand Down Expand Up @@ -271,8 +271,12 @@ def functional(
):
# Compute shear and delta_s
fneq = f - feq
shear = decompose_shear_d3q27(fneq)
delta_s = shear * rho # TODO: Check this
if wp.static(self.velocity_set.d == 3):
shear = decompose_shear_d3q27(fneq)
delta_s = shear * rho
else:
shear = decompose_shear_d2q9(fneq)
delta_s = shear * rho / self.compute_dtype(4.0)

# Perform collision
delta_h = fneq - delta_s
Expand Down
1 change: 0 additions & 1 deletion xlb/velocity_set/velocity_set.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,6 @@ def _init_jax_properties(self):
self.w = jnp.array(self._w, dtype=dtype)
self.opp_indices = jnp.array(self._opp_indices, dtype=jnp.int32)
self.cc = jnp.array(self._cc, dtype=dtype)
self.c_float = jnp.array(self._c_float, dtype=dtype)
self.qi = jnp.array(self._qi, dtype=dtype)

def _init_backend_constants(self):
Expand Down
Loading