From 06430408cd443e711b6d6b4c65af6cd9a16b9a33 Mon Sep 17 00:00:00 2001 From: tgautam03 <30533136+tgautam03@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:20:46 -0700 Subject: [PATCH 1/3] Rename 'prescribed_value' to 'prescribed_values' for consistency in ZouHeBC class --- xlb/operator/boundary_condition/bc_zouhe.py | 40 ++++++++++----------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/xlb/operator/boundary_condition/bc_zouhe.py b/xlb/operator/boundary_condition/bc_zouhe.py index 5cad5048..98636e7b 100644 --- a/xlb/operator/boundary_condition/bc_zouhe.py +++ b/xlb/operator/boundary_condition/bc_zouhe.py @@ -38,7 +38,7 @@ def __init__( self, bc_type, profile=None, - prescribed_value: Union[float, Tuple[float, ...], np.ndarray] = None, + prescribed_values: Union[float, Tuple[float, ...], np.ndarray] = None, velocity_set: VelocitySet = None, precision_policy: PrecisionPolicy = None, compute_backend: ComputeBackend = None, @@ -63,28 +63,28 @@ def __init__( ) # Handle prescribed value if provided - if prescribed_value is not None: + if prescribed_values is not None: if profile is not None: - raise ValueError("Cannot specify both profile and prescribed_value") + raise ValueError("Cannot specify both profile and prescribed_values") # Convert input to numpy array for validation - if isinstance(prescribed_value, (tuple, list)): - prescribed_value = np.array(prescribed_value, dtype=np.float64) - elif isinstance(prescribed_value, (int, float)): + if isinstance(prescribed_values, (tuple, list)): + prescribed_values = np.array(prescribed_values, dtype=np.float64) + elif isinstance(prescribed_values, (int, float)): if bc_type == "pressure": - prescribed_value = float(prescribed_value) + prescribed_values = float(prescribed_values) else: - raise ValueError("Velocity prescribed_value must be a tuple or array") - elif isinstance(prescribed_value, np.ndarray): - prescribed_value = prescribed_value.astype(np.float64) + raise ValueError("Velocity prescribed_values must be a tuple or array") + elif isinstance(prescribed_values, np.ndarray): + prescribed_values = prescribed_values.astype(np.float64) # Validate prescribed value if bc_type == "velocity": - if not isinstance(prescribed_value, np.ndarray): - raise ValueError("Velocity prescribed_value must be an array-like") + if not isinstance(prescribed_values, np.ndarray): + raise ValueError("Velocity prescribed_values must be an array-like") # Check for non-zero elements - only one element should be non-zero - non_zero_count = np.count_nonzero(prescribed_value) + non_zero_count = np.count_nonzero(prescribed_values) if non_zero_count > 1: raise ValueError("This BC only supports normal prescribed values (only one non-zero element allowed)") @@ -94,10 +94,10 @@ def __init__( # a vector of zeros associated with no-slip BC. # Accounting for all scenarios here. if self.compute_backend is ComputeBackend.WARP: - idx = np.nonzero(prescribed_value)[0] - prescribed_value = prescribed_value[idx][0] if idx.size else 0.0 - prescribed_value = self.precision_policy.store_precision.wp_dtype(prescribed_value) - self.prescribed_value = prescribed_value + idx = np.nonzero(prescribed_values)[0] + prescribed_values = prescribed_values[idx][0] if idx.size else 0.0 + prescribed_values = self.precision_policy.store_precision.wp_dtype(prescribed_values) + self.prescribed_values = prescribed_values self.profile = self._create_constant_prescribed_profile() # This BC needs auxilary data initialization before streaming @@ -113,7 +113,7 @@ def __init__( self.needs_padding = True def _create_constant_prescribed_profile(self): - _prescribed_value = self.prescribed_value + _prescribed_value = self.prescribed_values @wp.func def prescribed_profile_warp(index: wp.vec3i): @@ -299,8 +299,8 @@ def functional_velocity( # Find the value of u from the missing directions # Since we are only considering normal velocity, we only need to find one value (stored at the center of f_1) # Create velocity vector by multiplying the prescribed value with the normal vector - prescribed_value = f_1[0, index[0], index[1], index[2]] - _u = -prescribed_value * normals + prescribed_values = f_1[0, index[0], index[1], index[2]] + _u = -prescribed_values * normals for d in range(_d): unormal += _u[d] * normals[d] From 908e785248217c5fdd446e60da5ac1a87dc3a861 Mon Sep 17 00:00:00 2001 From: tgautam03 <30533136+tgautam03@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:26:17 -0700 Subject: [PATCH 2/3] Add 2D cylinder flow simulation and visualization scripts --- examples/cfd/2D_cylinder_flow.py | 210 ++++++++++++++++++++++ examples/cfd/2D_cylinder_flow_plot_viz.py | 163 +++++++++++++++++ 2 files changed, 373 insertions(+) create mode 100644 examples/cfd/2D_cylinder_flow.py create mode 100644 examples/cfd/2D_cylinder_flow_plot_viz.py diff --git a/examples/cfd/2D_cylinder_flow.py b/examples/cfd/2D_cylinder_flow.py new file mode 100644 index 00000000..228f8e05 --- /dev/null +++ b/examples/cfd/2D_cylinder_flow.py @@ -0,0 +1,210 @@ +import jax +import jax.numpy as jnp +import numpy as np +import xlb +from xlb import ComputeBackend, PrecisionPolicy +from xlb.velocity_set.d2q9 import D2Q9 +from xlb.grid import grid_factory +from xlb.operator.collision.bgk import BGK +from xlb.operator.equilibrium.quadratic_equilibrium import QuadraticEquilibrium +from xlb.operator.boundary_condition import ( + ZouHeBC, + FullwayBounceBackBC, +) +from xlb.operator.stream import Stream +from xlb.operator.macroscopic import Macroscopic +from xlb.helper.nse_solver import create_nse_fields + +# --- Configuration --- +# Physical parameters +Re = 100.0 +u_max = 0.1 +cylinder_radius = 10 +L_y = 100 +nu = u_max * (2 * cylinder_radius) / Re +tau = 3.0 * nu + 0.5 +omega = 1.0 / tau +print(f"Re: {Re}, u_max: {u_max}, nu: {nu}, tau: {tau}, omega: {omega}") + +# Grid parameters +nx, ny = 800, 200 + +# Backend and Precision +compute_backend = ComputeBackend.JAX +precision_policy = PrecisionPolicy.FP32FP32 + +# Initialize Velocity Set +velocity_set = D2Q9(precision_policy, compute_backend) + +# Initialize XLB +xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, +) + +# --- Setup --- +# Create Grid +grid = grid_factory((nx, ny)) + +# Create Fields +grid, f_0, f_1, missing_mask, bc_mask = create_nse_fields(grid=grid) + +# Define Geometry (Cylinder) +# Coordinate arrays +x_coords = jnp.arange(nx) +y_coords = jnp.arange(ny) +X, Y = jnp.meshgrid(x_coords, y_coords, indexing="ij") # Shape (nx, ny) + +cylinder_mask = (X - ny//2)**2 + (Y - ny//2)**2 <= cylinder_radius**2 +# Add cylinder to missing mask (solid obstacle) +missing_mask = missing_mask.at[:, cylinder_mask].set(True) + + +# --- Boundary Conditions --- +# Instantiate BCs first to get IDs +# Use NumPy array for prescribed values as ZouHeBC expects np.ndarray +u_inlet = np.array([u_max, 0.0]) # ux, uy + +zouhe_inlet = ZouHeBC( + bc_type="velocity", + prescribed_values=u_inlet, +) + +zouhe_outlet = ZouHeBC( + bc_type="pressure", + prescribed_values=1.0, # Scalar for pressure +) + +wall_bc = FullwayBounceBackBC() + +cylinder_bc = FullwayBounceBackBC() + +# Get IDs +BC_INLET = zouhe_inlet.id +BC_OUTLET = zouhe_outlet.id +BC_WALL = wall_bc.id +BC_CYLINDER = cylinder_bc.id + +print(f"BC IDs: Inlet={BC_INLET}, Outlet={BC_OUTLET}, Wall={BC_WALL}, Cylinder={BC_CYLINDER}") + +# Set BC Mask +# Inlet (Left, x=0) - all y positions at x=0 +bc_mask = bc_mask.at[0, 0, :].set(BC_INLET) +# Outlet (Right, x=nx-1) - all y positions at x=nx-1 +bc_mask = bc_mask.at[0, -1, :].set(BC_OUTLET) +# Top Wall (y=ny-1) - all x positions at y=ny-1 +bc_mask = bc_mask.at[0, :, -1].set(BC_WALL) +# Bottom Wall (y=0) - all x positions at y=0 +bc_mask = bc_mask.at[0, :, 0].set(BC_WALL) +# Cylinder Surface +bc_mask = bc_mask.at[0, cylinder_mask].set(BC_CYLINDER) + +# Note: ZouHe BC computes missing_mask internally from bc_mask +# Only the cylinder needs missing_mask set (solid obstacle) + +# List of BCs for iteration +bcs = [zouhe_inlet, zouhe_outlet, wall_bc, cylinder_bc] + + +# --- Operators --- +eq_op = QuadraticEquilibrium() +collision_op = BGK() +stream_op = Stream() +macroscopic_op = Macroscopic() + + +# --- Simulation Loop --- +# Initialize fields +rho_init = jnp.ones((1, nx, ny)) +u_init = jnp.zeros((2, nx, ny)) +# Add slight perturbation or initial flow to speed up? +u_init = u_init.at[0, ...].set(u_max) + +# Equilibrium initialization +f_0 = eq_op(rho_init, u_init) +f_1 = jnp.zeros_like(f_0) + +# JIT compile the step +@jax.jit +def step(f_pre, f_post, bc_mask, missing_mask): + # 1. Macroscopic + rho, u = macroscopic_op(f_pre) + + # 2. Equilibrium + feq = eq_op(rho, u) + + # 3. Collision + f_out = collision_op(f_pre, feq, rho, u, omega) + + # 4. Stream + f_streamed = stream_op(f_out) + + # 5. Boundary Conditions + # Apply BCs sequentially + f_curr = f_streamed + for bc in bcs: + f_curr = bc(f_streamed, f_curr, bc_mask, missing_mask) + + return f_curr + +# Run +num_steps = 20000 +save_interval = 10 # Save every 10 steps +print("Starting simulation...") +import time +start_time = time.time() + +# Storage for visualization +saved_steps = [] +saved_rho = [] +saved_u = [] +saved_vorticity = [] + +# We need to swap f_0 and f_1 buffers +current_f = f_0 +next_f = f_1 + +for i in range(num_steps): + next_f = step(current_f, next_f, bc_mask, missing_mask) + + # Swap + current_f, next_f = next_f, current_f + + # Save data for visualization + if i % save_interval == 0: + rho, u = macroscopic_op(current_f) + # Compute vorticity (curl of velocity in 2D) + # vorticity = du_y/dx - du_x/dy + du_y_dx = jnp.gradient(u[1], axis=0) + du_x_dy = jnp.gradient(u[0], axis=1) + vorticity = du_y_dx - du_x_dy + + saved_steps.append(i) + saved_rho.append(np.array(rho[0])) + saved_u.append(np.array(u)) + saved_vorticity.append(np.array(vorticity)) + + if i % 100 == 0: + rho, u = macroscopic_op(current_f) + u_mag = jnp.sqrt(u[0]**2 + u[1]**2) + print(f"Step {i}: Max U = {jnp.max(u_mag):.4f}, Min Rho = {jnp.min(rho):.4f}") + +end_time = time.time() +print(f"Simulation finished in {end_time - start_time:.2f} seconds.") +print(f"MLUPS: {num_steps * nx * ny / (end_time - start_time) / 1e6:.2f}") + +# Convert to numpy arrays for easier handling +saved_rho = np.array(saved_rho) +saved_u = np.array(saved_u) +saved_vorticity = np.array(saved_vorticity) + +print(f"\nSaved {len(saved_steps)} frames for visualization") + +# Save data to files for visualization +np.save('saved_steps.npy', np.array(saved_steps)) +np.save('saved_rho.npy', saved_rho) +np.save('saved_u.npy', saved_u) +np.save('saved_vorticity.npy', saved_vorticity) +print("Data saved to .npy files") +print("\nRun 'python plot_flow.py' to visualize the results") diff --git a/examples/cfd/2D_cylinder_flow_plot_viz.py b/examples/cfd/2D_cylinder_flow_plot_viz.py new file mode 100644 index 00000000..2cb7dc4e --- /dev/null +++ b/examples/cfd/2D_cylinder_flow_plot_viz.py @@ -0,0 +1,163 @@ +""" +Interactive visualization for 2D cylinder flow simulation. +Optimized for smooth, real-time slider interaction. + +Usage: python plot_flow.py +""" + +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.patches import Circle +from matplotlib.widgets import Slider + +# Load the saved data +print("Loading simulation data...") +try: + saved_steps = np.load('saved_steps.npy') + saved_rho = np.load('saved_rho.npy') + saved_u = np.load('saved_u.npy') + saved_vorticity = np.load('saved_vorticity.npy') + print(f"Loaded {len(saved_steps)} frames") +except FileNotFoundError: + print("Error: Data files not found. Please run 2D_cylinder_flow.py first.") + exit(1) + +# Simulation parameters +nx, ny = 800, 200 +cylinder_radius = 10 + +# Pre-compute velocity magnitude for all frames +print("Pre-computing velocity magnitudes...") +saved_u_mag = np.sqrt(saved_u[:, 0]**2 + saved_u[:, 1]**2) + +# Pre-compute global min/max for consistent colormaps (except vorticity) +u_mag_max = np.max(saved_u_mag) +rho_min, rho_max = np.min(saved_rho), np.max(saved_rho) +ux_min, ux_max = np.min(saved_u[:, 0]), np.max(saved_u[:, 0]) + +print("Setting up visualization...") + +# Create figure with subplots +fig = plt.figure(figsize=(16, 9)) +gs = fig.add_gridspec(3, 2, height_ratios=[1, 1, 0.05], hspace=0.3, wspace=0.3) + +ax1 = fig.add_subplot(gs[0, 0]) +ax2 = fig.add_subplot(gs[0, 1]) +ax3 = fig.add_subplot(gs[1, 0]) +ax4 = fig.add_subplot(gs[1, 1]) +ax_slider = fig.add_subplot(gs[2, :]) + +# Initialize with first frame +frame_idx = 0 + +# Velocity magnitude +im1 = ax1.imshow( + saved_u_mag[frame_idx].T, + origin='lower', + cmap='jet', + extent=[0, nx, 0, ny], + vmin=0, + vmax=u_mag_max, + aspect='equal' +) +circle1 = Circle((ny//2, ny//2), cylinder_radius, color='black', zorder=10) +ax1.add_patch(circle1) +ax1.set_title('Velocity Magnitude') +ax1.set_xlabel('X') +ax1.set_ylabel('Y') +cbar1 = plt.colorbar(im1, ax=ax1, label='|u|') + +# Vorticity – no fixed vmin/vmax here +im2 = ax2.imshow( + saved_vorticity[frame_idx].T, + origin='lower', + cmap='RdBu_r', + extent=[0, nx, 0, ny], + aspect='equal' +) +circle2 = Circle((ny//2, ny//2), cylinder_radius, color='black', zorder=10) +ax2.add_patch(circle2) +ax2.set_title('Vorticity (ω = ∂v/∂x - ∂u/∂y)') +ax2.set_xlabel('X') +ax2.set_ylabel('Y') +cbar2 = plt.colorbar(im2, ax=ax2, label='ω') + +# Density +im3 = ax3.imshow( + saved_rho[frame_idx].T, + origin='lower', + cmap='viridis', + extent=[0, nx, 0, ny], + vmin=rho_min, + vmax=rho_max, + aspect='equal' +) +circle3 = Circle((ny//2, ny//2), cylinder_radius, color='black', zorder=10) +ax3.add_patch(circle3) +ax3.set_title('Density (ρ)') +ax3.set_xlabel('X') +ax3.set_ylabel('Y') +cbar3 = plt.colorbar(im3, ax=ax3, label='ρ') + +# X-velocity +im4 = ax4.imshow( + saved_u[frame_idx, 0].T, + origin='lower', + cmap='coolwarm', + extent=[0, nx, 0, ny], + vmin=ux_min, + vmax=ux_max, + aspect='equal' +) +circle4 = Circle((ny//2, ny//2), cylinder_radius, color='black', zorder=10) +ax4.add_patch(circle4) +ax4.set_title('X-Velocity Component (u_x)') +ax4.set_xlabel('X') +ax4.set_ylabel('Y') +cbar4 = plt.colorbar(im4, ax=ax4, label='u_x') + +title = fig.suptitle( + f'2D Cylinder Flow - Step {saved_steps[frame_idx]}', + fontsize=16, + fontweight='bold' +) + +# Create slider +slider = Slider( + ax=ax_slider, + label='Frame', + valmin=0, + valmax=len(saved_steps) - 1, + valinit=0, + valstep=1 +) + +def update(val): + """Update plots with dynamic vorticity scaling.""" + frame_idx = int(slider.val) + + # Velocity magnitude, density, x-velocity + im1.set_data(saved_u_mag[frame_idx].T) + im3.set_data(saved_rho[frame_idx].T) + im4.set_data(saved_u[frame_idx, 0].T) + + # Vorticity with per-frame scaling (use percentiles to ignore outliers) + vort = saved_vorticity[frame_idx] + vmin = np.percentile(vort, 5) + vmax = np.percentile(vort, 95) + im2.set_data(vort.T) + im2.set_clim(vmin, vmax) # update color limits + cbar2.update_normal(im2) # keep colorbar in sync + + # Update title + title.set_text(f'2D Cylinder Flow - Step {saved_steps[frame_idx]}') + + fig.canvas.draw_idle() + +slider.on_changed(update) + +print("\nVisualization ready!") +print("Use the slider to scroll through timesteps.") +print("Close the window to exit.") + +plt.show() From fd48688e3f733910ff9eb91aaa174aac5391006f Mon Sep 17 00:00:00 2001 From: tgautam03 <30533136+tgautam03@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:44:39 -0700 Subject: [PATCH 3/3] Refactor step function to improve readability by removing unnecessary blank lines --- examples/cfd/2D_cylinder_flow.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/examples/cfd/2D_cylinder_flow.py b/examples/cfd/2D_cylinder_flow.py index 228f8e05..5ff118f7 100644 --- a/examples/cfd/2D_cylinder_flow.py +++ b/examples/cfd/2D_cylinder_flow.py @@ -130,22 +130,22 @@ def step(f_pre, f_post, bc_mask, missing_mask): # 1. Macroscopic rho, u = macroscopic_op(f_pre) - + # 2. Equilibrium feq = eq_op(rho, u) - + # 3. Collision f_out = collision_op(f_pre, feq, rho, u, omega) - + # 4. Stream f_streamed = stream_op(f_out) - + # 5. Boundary Conditions # Apply BCs sequentially f_curr = f_streamed for bc in bcs: f_curr = bc(f_streamed, f_curr, bc_mask, missing_mask) - + return f_curr # Run @@ -167,10 +167,10 @@ def step(f_pre, f_post, bc_mask, missing_mask): for i in range(num_steps): next_f = step(current_f, next_f, bc_mask, missing_mask) - + # Swap current_f, next_f = next_f, current_f - + # Save data for visualization if i % save_interval == 0: rho, u = macroscopic_op(current_f) @@ -179,12 +179,12 @@ def step(f_pre, f_post, bc_mask, missing_mask): du_y_dx = jnp.gradient(u[1], axis=0) du_x_dy = jnp.gradient(u[0], axis=1) vorticity = du_y_dx - du_x_dy - + saved_steps.append(i) saved_rho.append(np.array(rho[0])) saved_u.append(np.array(u)) saved_vorticity.append(np.array(vorticity)) - + if i % 100 == 0: rho, u = macroscopic_op(current_f) u_mag = jnp.sqrt(u[0]**2 + u[1]**2)