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.

Sign up for GitHub

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

Jump to bottom

Fixed distribute and post processing bugs #54

Merged
merged 1 commit into from
Aug 2, 2024
Merged

Fixed distribute and post processing bugs #54

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
55 changes: 35 additions & 20 deletions examples/cfd/flow_past_sphere_3d.py
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading