some physics with python and cuda
some physics with python and cudaunknown
python
a year ago
6.2 kB
25
Indexable
import numpy as np import pycuda.autoinit import pycuda.driver as cuda from pycuda.compiler import SourceModule import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D class GravitationalModel: def __init__(self, num_particles): self.num_particles = num_particles self.block_dim = (256, 1, 1) self.grid_dim = ((num_particles + self.block_dim[0] - 1) // self.block_dim[0], 1) self.module = SourceModule(""" __global__ void gravity(float *pos, float *force, int num_particles, float G, float dt, float *pendulum_attractor, float *fold_curve) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < num_particles) { float fx = 0.0f; float fy = 0.0f; float mass = 1.0f; // assuming all particles have same mass for simplicity for (int i = 0; i < num_particles; ++i) { if (i != idx) { float dx = pos[i * 2] - pos[idx * 2]; float dy = pos[i * 2 + 1] - pos[idx * 2 + 1]; float dist = sqrt(dx * dx + dy * dy); float inv_dist = 1.0f / (dist + 0.01f); // avoid division by zero fx += G * mass * mass * (dx * inv_dist) * dt; fy += G * mass * mass * (dy * inv_dist) * dt; } } fx *= pendulum_attractor[idx]; fy *= pendulum_attractor[idx]; fx *= fold_curve[idx]; fy *= fold_curve[idx]; force[idx * 2] = fx; force[idx * 2 + 1] = fy; } } __global__ void riemann_zeta(int n, double s, double *result) { int idx = threadIdx.x + blockIdx.x * blockDim.x; double sum = 0.0; for (int k = 1; k <= n; k++) { sum += 1.0 / pow((double)k, s); } result[idx] = sum; } __global__ void compute_sine_wave(float *x, float *y) { int idx = threadIdx.x + blockIdx.x * blockDim.x; y[idx] = sinf(x[idx]); } """) self.gravity_func = self.module.get_function("gravity") self.riemann_zeta_func = self.module.get_function("riemann_zeta") self.sine_wave_func = self.module.get_function("compute_sine_wave") def compute_gravity(self, pos, G, dt, pendulum_attractor, fold_curve): force = np.zeros_like(pos) pos_gpu = cuda.to_device(pos.astype(np.float32)) force_gpu = cuda.to_device(force.astype(np.float32)) pendulum_attractor_gpu = cuda.to_device(pendulum_attractor.astype(np.float32)) fold_curve_gpu = cuda.to_device(fold_curve.astype(np.float32)) self.gravity_func(pos_gpu, force_gpu, np.int32(self.num_particles), np.float32(G), np.float32(dt), pendulum_attractor_gpu, fold_curve_gpu, block=self.block_dim, grid=self.grid_dim) cuda.memcpy_dtoh(force, force_gpu) return force def compute_riemann_zeta(self, n, s): result = np.zeros(n, dtype=np.double) result_gpu = cuda.to_device(result) self.riemann_zeta_func(np.int32(n), np.double(s), result_gpu, block=self.block_dim, grid=self.grid_dim) cuda.memcpy_dtoh(result, result_gpu) return result def compute_sine_wave(self, x): y = np.zeros_like(x) x_gpu = cuda.to_device(x.astype(np.float32)) y_gpu = cuda.to_device(y.astype(np.float32)) self.sine_wave_func(x_gpu, y_gpu, block=self.block_dim, grid=self.grid_dim) cuda.memcpy_dtoh(y, y_gpu) return y class Menu: def __init__(self): self.model = None self.pos = None self.G = 6.67430e-11 self.dt = 1.0 self.pendulum_attractor = None self.fold_curve = None def create_model(self, num_particles): self.model = GravitationalModel(num_particles) def initialize_particles(self, num_particles): self.pos = np.random.uniform(-100, 100, (num_particles, 2)) self.pendulum_attractor = np.random.uniform(0.5, 2, (num_particles,)) self.fold_curve = np.random.uniform(0.5, 2, (num_particles,)) def run_simulation(self): force = self.model.compute_gravity(self.pos, self.G, self.dt, self.pendulum_attractor, self.fold_curve) return force def plot_simulation(self, force): fig = plt.figure() ax = fig.add_subplot(111, projection='3d') ax.quiver(self.pos[:, 0], self.pos[:, 1], np.zeros_like(self.pos[:, 0]), force[:, 0], force[:, 1], np.zeros_like(force[:, 0])) ax.set_xlabel('X') ax.set_ylabel('Y') ax.set_zlabel('Z') plt.show() def main_menu(self): while True: print("1. Create Model") print("2. Initialize Particles") print("3. Run Simulation") print("4. Plot Simulation") print("5. Exit") choice = int(input("Enter your choice: ")) if choice == 1: num_particles = int(input("Enter number of particles: ")) self.create_model(num_particles) elif choice == 2: num_particles = int(input("Enter number of particles: ")) self.initialize_particles(num_particles) elif choice == 3: if self.model is None or self.pos is None: print("Please create model and initialize particles first.") continue force = self.run_simulation() print("Simulation completed.") elif choice == 4: if self.model is None or self.pos is None: print("Please create model and initialize particles first.") continue self.plot_simulation(force) elif choice == 5: break else: print("Invalid choice. Please try again.") if __name__ == "__main__": menu = Menu() menu.main_menu()
Editor is loading...
Leave a Comment