some physics with python and cuda

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()