fluid dynamics - code

import torch import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D import numpy as np from matplotlib.animation import FuncAnimation class FluidSimulator3D: def __init__(self, width, height, depth, viscosity=0.1, dt=0.1): self.width = width self.height = height self.depth = depth self.viscosity = viscosity self.dt = dt # Initialize velocity fields (u: x, v: y, w: z) self.u = torch.zeros((depth, height, width)) self.v = torch.zeros((depth, height, width)) self.w = torch.zeros((depth, height, width)) # Initialize pressure field self.p = torch.zeros((depth, height, width)) # Initialize density/dye field for visualization self.dye = torch.zeros((depth, height, width)) def set_initial_conditions(self, center_x, center_y, center_z, radius, strength): """Set initial toroidal vortex conditions (smoke ring)""" z, y, x = torch.meshgrid( torch.arange(self.depth), torch.arange(self.height), torch.arange(self.width) ) # Calculate distances from center dx = x - center_x dy = y - center_y dz = z - center_z # Radial distance from vertical axis r = torch.sqrt(dx**2 + dy**2) # Distance from ring centerline ring_radius = radius * 1.5 d = torch.sqrt((r - ring_radius)**2 + dz**2) # Create toroidal vortex mask = d < radius # Set initial velocities for smoke ring # Circular motion around the ring's core theta = torch.atan2(dy, dx) self.u[mask] = -strength * torch.sin(theta[mask]) * dz[mask] / radius self.v[mask] = strength * torch.cos(theta[mask]) * dz[mask] / radius self.w[mask] = strength * (r[mask] - ring_radius) / radius # Add dye for visualization self.dye[mask] = 1.0 def diffuse(self, field): """Apply viscous diffusion in 3D""" laplacian = ( field[:-2, 1:-1, 1:-1] + field[2:, 1:-1, 1:-1] + field[1:-1, :-2, 1:-1] + field[1:-1, 2:, 1:-1] + field[1:-1, 1:-1, :-2] + field[1:-1, 1:-1, 2:] - 6 * field[1:-1, 1:-1, 1:-1] ) field[1:-1, 1:-1, 1:-1] += self.viscosity * self.dt * laplacian def advect(self): """Advect velocity and dye fields in 3D""" z, y, x = torch.meshgrid( torch.arange(self.depth), torch.arange(self.height), torch.arange(self.width) ) # Calculate previous positions x_prev = x - self.dt * self.u y_prev = y - self.dt * self.v z_prev = z - self.dt * self.w # Clip to bounds x_prev = torch.clamp(x_prev, 0, self.width-1) y_prev = torch.clamp(y_prev, 0, self.height-1) z_prev = torch.clamp(z_prev, 0, self.depth-1) # Convert to indices x0 = torch.floor(x_prev).long() x1 = x0 + 1 y0 = torch.floor(y_prev).long() y1 = y0 + 1 z0 = torch.floor(z_prev).long() z1 = z0 + 1 # Clip again x0 = torch.clamp(x0, 0, self.width-1) x1 = torch.clamp(x1, 0, self.width-1) y0 = torch.clamp(y0, 0, self.height-1) y1 = torch.clamp(y1, 0, self.height-1) z0 = torch.clamp(z0, 0, self.depth-1) z1 = torch.clamp(z1, 0, self.depth-1) # Get weights for trilinear interpolation sx = x_prev - x0.float() sy = y_prev - y0.float() sz = z_prev - z0.float() # Interpolate def interpolate(field): c000 = field[z0, y0, x0] c001 = field[z0, y0, x1] c010 = field[z0, y1, x0] c011 = field[z0, y1, x1] c100 = field[z1, y0, x0] c101 = field[z1, y0, x1] c110 = field[z1, y1, x0] c111 = field[z1, y1, x1] return ( (1-sz) * (1-sy) * (1-sx) * c000 + (1-sz) * (1-sy) * sx * c001 + (1-sz) * sy * (1-sx) * c010 + (1-sz) * sy * sx * c011 + sz * (1-sy) * (1-sx) * c100 + sz * (1-sy) * sx * c101 + sz * sy * (1-sx) * c110 + sz * sy * sx * c111 ) # Update fields self.u = interpolate(self.u) self.v = interpolate(self.v) self.w = interpolate(self.w) self.dye = interpolate(self.dye) def project(self): """Pressure projection to enforce incompressibility in 3D""" # Calculate divergence div = torch.zeros((self.depth, self.height, self.width)) div[1:-1, 1:-1, 1:-1] = ( (self.u[1:-1, 1:-1, 2:] - self.u[1:-1, 1:-1, :-2]) / 2 + (self.v[1:-1, 2:, 1:-1] - self.v[1:-1, :-2, 1:-1]) / 2 + (self.w[2:, 1:-1, 1:-1] - self.w[:-2, 1:-1, 1:-1]) / 2 ) # Solve pressure Poisson equation for _ in range(20): self.p[1:-1, 1:-1, 1:-1] = ( (self.p[:-2, 1:-1, 1:-1] + self.p[2:, 1:-1, 1:-1] + self.p[1:-1, :-2, 1:-1] + self.p[1:-1, 2:, 1:-1] + self.p[1:-1, 1:-1, :-2] + self.p[1:-1, 1:-1, 2:] - div[1:-1, 1:-1, 1:-1]) / 6 ) # Enforce Neumann boundary conditions self.p[0, :, :] = self.p[1, :, :] self.p[-1, :, :] = self.p[-2, :, :] self.p[:, 0, :] = self.p[:, 1, :] self.p[:, -1, :] = self.p[:, -2, :] self.p[:, :, 0] = self.p[:, :, 1] self.p[:, :, -1] = self.p[:, :, -2] # Subtract pressure gradient self.u[1:-1, 1:-1, 1:-1] -= (self.p[1:-1, 1:-1, 2:] - self.p[1:-1, 1:-1, :-2]) / 2 self.v[1:-1, 1:-1, 1:-1] -= (self.p[1:-1, 2:, 1:-1] - self.p[1:-1, :-2, 1:-1]) / 2 self.w[1:-1, 1:-1, 1:-1] -= (self.p[2:, 1:-1, 1:-1] - self.p[:-2, 1:-1, 1:-1]) / 2 def step(self): """Perform one simulation step""" self.diffuse(self.u) self.diffuse(self.v) self.diffuse(self.w) self.project() self.advect() self.project() def run_simulation(self, num_steps=100): """Run simulation and return animation frames""" frames = [] for _ in range(num_steps): self.step() # Store a slice of the 3D volume for visualization frames.append(self.dye[:, :, self.width//2].clone()) return frames def create_and_run_simulation(): # Create simulator size = 50 # smaller size for 3D sim = FluidSimulator3D(size, size, size, viscosity=0.1, dt=0.1) # Set initial conditions (smoke ring) sim.set_initial_conditions( center_x=size//2, center_y=size//2, center_z=size//4, # start lower radius=8, strength=2.0 ) # Run simulation frames = sim.run_simulation(num_steps=200) # Create animation fig, ax = plt.subplots() def animate(frame): ax.clear() ax.imshow(frame, cmap='viridis') ax.set_title('Y-Z slice at middle of X axis') ax.set_xticks([]) ax.set_yticks([]) anim = FuncAnimation( fig, animate, frames=frames, interval=50, blit=False ) plt.show() return anim if __name__ == "__main__": anim = create_and_run_simulation()