Skip to content

Commit

Permalink
Fixed distribute and post processing bugs
Browse files Browse the repository at this point in the history
  • Loading branch information
@mehdiataei
mehdiataei committed Aug 2, 2024
1 parent d7b35b3 commit 1810e9c
Show file tree
Hide file tree
Showing 6 changed files with 147 additions and 66 deletions.
55 changes: 35 additions & 20 deletions examples/cfd/flow_past_sphere_3d.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,44 +15,51 @@
import numpy as np
import jax.numpy as jnp


class FlowOverSphere:
def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):

# initialize backend
xlb.init(
velocity_set=velocity_set,
default_backend=backend,
default_precision_policy=precision_policy,
)

self.grid_shape = grid_shape
self.velocity_set = velocity_set
self.backend = backend
self.precision_policy = precision_policy
self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = create_nse_fields(grid_shape)
self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = (
create_nse_fields(grid_shape)
)
self.stepper = None
self.boundary_conditions = []

# Setup the simulation BC, its initial conditions, and the stepper
self._setup(omega)

def _setup(self, omega):
self.setup_boundary_conditions()
self.setup_boundary_masks()
self.initialize_fields()
self.setup_stepper(omega)

def define_boundary_indices(self):
inlet = self.grid.boundingBoxIndices['left']
outlet = self.grid.boundingBoxIndices['right']
walls = [self.grid.boundingBoxIndices['bottom'][i] + self.grid.boundingBoxIndices['top'][i] +
self.grid.boundingBoxIndices['front'][i] + self.grid.boundingBoxIndices['back'][i] for i in range(self.velocity_set.d)]

inlet = self.grid.boundingBoxIndices["left"]
outlet = self.grid.boundingBoxIndices["right"]
walls = [
self.grid.boundingBoxIndices["bottom"][i]
+ self.grid.boundingBoxIndices["top"][i]
+ self.grid.boundingBoxIndices["front"][i]
+ self.grid.boundingBoxIndices["back"][i]
for i in range(self.velocity_set.d)
]

sphere_radius = self.grid_shape[1] // 12
x = np.arange(self.grid_shape[0])
y = np.arange(self.grid_shape[1])
z = np.arange(self.grid_shape[2])
X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
X, Y, Z = np.meshgrid(x, y, z, indexing="ij")
indices = np.where(
(X - self.grid_shape[0] // 6) ** 2
+ (Y - self.grid_shape[1] // 2) ** 2
Expand All @@ -62,7 +69,7 @@ def define_boundary_indices(self):
sphere = [tuple(indices[i]) for i in range(self.velocity_set.d)]

return inlet, outlet, walls, sphere

def setup_boundary_conditions(self):
inlet, outlet, walls, sphere = self.define_boundary_indices()
bc_left = EquilibriumBC(rho=1.0, u=(0.02, 0.0, 0.0), indices=inlet)
Expand All @@ -83,24 +90,31 @@ def setup_boundary_masks(self):

def initialize_fields(self):
self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.backend)

def setup_stepper(self, omega):
self.stepper = IncompressibleNavierStokesStepper(
omega, boundary_conditions=self.boundary_conditions
)

def run(self, num_steps):
def run(self, num_steps, post_process_interval=100):
for i in range(num_steps):
self.f_1 = self.stepper(self.f_0, self.f_1, self.boundary_mask, self.missing_mask, i)
self.f_1 = self.stepper(
self.f_0, self.f_1, self.boundary_mask, self.missing_mask, i
)
self.f_0, self.f_1 = self.f_1, self.f_0

if i % post_process_interval == 0 or i == num_steps - 1:
self.post_process(i)

def post_process(self, i):
# Write the results. We'll use JAX backend for the post-processing
if not isinstance(self.f_0, jnp.ndarray):
self.f_0 = wp.to_jax(self.f_0)
f_0 = wp.to_jax(self.f_0)
else:
f_0 = self.f_0

macro = Macroscopic(compute_backend=ComputeBackend.JAX)
rho, u = macro(self.f_0)
rho, u = macro(f_0)

# remove boundary cells
u = u[:, 1:-1, 1:-1, 1:-1]
Expand All @@ -120,6 +134,7 @@ def post_process(self, i):
precision_policy = PrecisionPolicy.FP32FP32
omega = 1.6

simulation = FlowOverSphere(omega, grid_shape, velocity_set, backend, precision_policy)
simulation.run(num_steps=10000)
simulation.post_process(i=10000)
simulation = FlowOverSphere(
omega, grid_shape, velocity_set, backend, precision_policy
)
simulation.run(num_steps=10000, post_process_interval=1000)
47 changes: 30 additions & 17 deletions examples/cfd/lid_driven_cavity_2d.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,22 +13,23 @@

class LidDrivenCavity2D:
def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):

# initialize backend
xlb.init(
velocity_set=velocity_set,
default_backend=backend,
default_precision_policy=precision_policy,
)

self.grid_shape = grid_shape
self.velocity_set = velocity_set
self.backend = backend
self.precision_policy = precision_policy
self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = create_nse_fields(grid_shape)
self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = (
create_nse_fields(grid_shape)
)
self.stepper = None
self.boundary_conditions = []

# Setup the simulation BC, its initial conditions, and the stepper
self._setup(omega)

Expand All @@ -39,11 +40,15 @@ def _setup(self, omega):
self.setup_stepper(omega)

def define_boundary_indices(self):
lid = self.grid.boundingBoxIndices['top']
walls = [self.grid.boundingBoxIndices['bottom'][i] + self.grid.boundingBoxIndices['left'][i] +
self.grid.boundingBoxIndices['right'][i] for i in range(self.velocity_set.d)]
lid = self.grid.boundingBoxIndices["top"]
walls = [
self.grid.boundingBoxIndices["bottom"][i]
+ self.grid.boundingBoxIndices["left"][i]
+ self.grid.boundingBoxIndices["right"][i]
for i in range(self.velocity_set.d)
]
return lid, walls

def setup_boundary_conditions(self):
lid, walls = self.define_boundary_indices()
bc_top = EquilibriumBC(rho=1.0, u=(0.02, 0.0), indices=lid)
Expand All @@ -62,25 +67,32 @@ def setup_boundary_masks(self):

def initialize_fields(self):
self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.backend)

def setup_stepper(self, omega):
self.stepper = IncompressibleNavierStokesStepper(
omega, boundary_conditions=self.boundary_conditions
)

def run(self, num_steps):
def run(self, num_steps, post_process_interval=100):
for i in range(num_steps):
self.f_1 = self.stepper(self.f_0, self.f_1, self.boundary_mask, self.missing_mask, i)
self.f_1 = self.stepper(
self.f_0, self.f_1, self.boundary_mask, self.missing_mask, i
)
self.f_0, self.f_1 = self.f_1, self.f_0

if i % post_process_interval == 0 or i == num_steps - 1:
self.post_process(i)

def post_process(self, i):
# Write the results. We'll use JAX backend for the post-processing
if not isinstance(self.f_0, jnp.ndarray):
self.f_0 = wp.to_jax(self.f_0)
f_0 = wp.to_jax(self.f_0)
else:
f_0 = self.f_0

macro = Macroscopic(compute_backend=ComputeBackend.JAX)

rho, u = macro(self.f_0)
rho, u = macro(f_0)

# remove boundary cells
rho = rho[:, 1:-1, 1:-1]
Expand All @@ -95,13 +107,14 @@ def post_process(self, i):

if __name__ == "__main__":
# Running the simulation
grid_size = 128
grid_size = 500
grid_shape = (grid_size, grid_size)
backend = ComputeBackend.JAX
velocity_set = xlb.velocity_set.D2Q9()
precision_policy = PrecisionPolicy.FP32FP32
omega = 1.6

simulation = LidDrivenCavity2D(omega, grid_shape, velocity_set, backend, precision_policy)
simulation.run(num_steps=500)
simulation.post_process(i=500)
simulation = LidDrivenCavity2D(
omega, grid_shape, velocity_set, backend, precision_policy
)
simulation.run(num_steps=5000, post_process_interval=1000)
3 changes: 1 addition & 2 deletions examples/cfd/lid_driven_cavity_2d_distributed.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,5 +31,4 @@ def setup_stepper(self, omega):
omega=1.6

simulation = LidDrivenCavity2D_distributed(omega, grid_shape, velocity_set, backend, precision_policy)
simulation.run(num_steps=5000)
simulation.post_process(i=5000)
simulation.run(num_steps=5000, post_process_interval=1000)
60 changes: 40 additions & 20 deletions examples/cfd/windtunnel_3d.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,23 +19,26 @@


class WindTunnel3D:
def __init__(self, omega, wind_speed, grid_shape, velocity_set, backend, precision_policy):

def __init__(
self, omega, wind_speed, grid_shape, velocity_set, backend, precision_policy
):
# initialize backend
xlb.init(
velocity_set=velocity_set,
default_backend=backend,
default_precision_policy=precision_policy,
)

self.grid_shape = grid_shape
self.velocity_set = velocity_set
self.backend = backend
self.precision_policy = precision_policy
self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = create_nse_fields(grid_shape)
self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = (
create_nse_fields(grid_shape)
)
self.stepper = None
self.boundary_conditions = []

# Setup the simulation BC, its initial conditions, and the stepper
self._setup(omega, wind_speed)

Expand All @@ -54,31 +57,38 @@ def voxelize_stl(self, stl_filename, length_lbm_unit):
return mesh_matrix, pitch

def define_boundary_indices(self):
inlet = self.grid.boundingBoxIndices['left']
outlet = self.grid.boundingBoxIndices['right']
walls = [self.grid.boundingBoxIndices['bottom'][i] + self.grid.boundingBoxIndices['top'][i] +
self.grid.boundingBoxIndices['front'][i] + self.grid.boundingBoxIndices['back'][i] for i in range(self.velocity_set.d)]

inlet = self.grid.boundingBoxIndices["left"]
outlet = self.grid.boundingBoxIndices["right"]
walls = [
self.grid.boundingBoxIndices["bottom"][i]
+ self.grid.boundingBoxIndices["top"][i]
+ self.grid.boundingBoxIndices["front"][i]
+ self.grid.boundingBoxIndices["back"][i]
for i in range(self.velocity_set.d)
]

stl_filename = "examples/cfd/stl-files/DrivAer-Notchback.stl"
grid_size_x = self.grid_shape[0]
car_length_lbm_unit = grid_size_x / 4
car_voxelized, pitch = self.voxelize_stl(stl_filename, car_length_lbm_unit)

car_area = np.prod(car_voxelized.shape[1:])
tx, ty, tz = np.array([grid_size_x, grid_size_y, grid_size_z]) - car_voxelized.shape
tx, ty, tz = (
np.array([grid_size_x, grid_size_y, grid_size_z]) - car_voxelized.shape
)
shift = [tx // 4, ty // 2, 0]
car = np.argwhere(car_voxelized) + shift
car = np.array(car).T
car = [tuple(car[i]) for i in range(self.velocity_set.d)]

return inlet, outlet, walls, car

def setup_boundary_conditions(self, wind_speed):
inlet, outlet, walls, car = self.define_boundary_indices()
bc_left = EquilibriumBC(rho=1.0, u=(wind_speed, 0.0, 0.0), indices=inlet)
bc_walls = FullwayBounceBackBC(indices=walls)
bc_do_nothing = DoNothingBC(indices=outlet)
bc_car= FullwayBounceBackBC(indices=car)
bc_car = FullwayBounceBackBC(indices=car)
self.boundary_conditions = [bc_left, bc_walls, bc_do_nothing, bc_car]

def setup_boundary_masks(self):
Expand All @@ -93,26 +103,35 @@ def setup_boundary_masks(self):

def initialize_fields(self):
self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.backend)

def setup_stepper(self, omega):
self.stepper = IncompressibleNavierStokesStepper(
omega, boundary_conditions=self.boundary_conditions, collision_type="KBC"
)

def run(self, num_steps, print_interval):
def run(self, num_steps, print_interval, post_process_interval=100):
start_time = time.time()
for i in range(num_steps):
self.f_1 = self.stepper(self.f_0, self.f_1, self.boundary_mask, self.missing_mask, i)
self.f_1 = self.stepper(
self.f_0, self.f_1, self.boundary_mask, self.missing_mask, i
)
self.f_0, self.f_1 = self.f_1, self.f_0

if (i + 1) % print_interval == 0:
elapsed_time = time.time() - start_time
print(f"Iteration: {i+1}/{num_steps} | Time elapsed: {elapsed_time:.2f}s")
print(
f"Iteration: {i+1}/{num_steps} | Time elapsed: {elapsed_time:.2f}s"
)

if i % post_process_interval == 0 or i == num_steps - 1:
self.post_process(i)

def post_process(self, i):
# Write the results. We'll use JAX backend for the post-processing
if not isinstance(self.f_0, jnp.ndarray):
f_0 = wp.to_jax(self.f_0)
else:
f_0 = self.f_0

macro = Macroscopic(compute_backend=ComputeBackend.JAX)

Expand Down Expand Up @@ -159,6 +178,7 @@ def post_process(self, i):
print(f"Max iterations: {num_steps}")
print("\n" + "=" * 50 + "\n")

simulation = WindTunnel3D(omega, wind_speed, grid_shape, velocity_set, backend, precision_policy)
simulation.run(num_steps, print_interval)
simulation.post_process(i=num_steps)
simulation = WindTunnel3D(
omega, wind_speed, grid_shape, velocity_set, backend, precision_policy
)
simulation.run(num_steps, print_interval, post_process_interval=1000)
Loading

0 comments on commit 1810e9c

Please sign in to comment.