Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import random
from cymem.cymem cimport Pool
from libc.math cimport sqrt
cimport cython
ctypedef double[2] double_pair
cdef class World:
cdef Pool mem
cdef int N
cdef double* m
cdef double_pair* r
cdef double_pair* v
cdef double_pair* F
cdef readonly double dt
def __init__(self, N, threads=1, m_min=1, m_max=30.0, r_max=50.0, v_max=4.0, dt=1e-3):
self.mem = Pool()
self.N = N
#m = np.random.uniform(m_min, m_max, N)
#r = np.random.uniform(-r_max, r_max, (N, 2))
# v = np.random.uniform(-v_max, v_max, (N, 2))
#F = np.zeros_like(self.r)
self.m = <double*>self.mem.alloc(N, sizeof(double))
self.r = <double_pair*>self.mem.alloc(N, sizeof(double_pair))
self.v = <double_pair*>self.mem.alloc(N, sizeof(double_pair))
self.F = <double_pair*>self.mem.alloc(N, sizeof(double_pair))
for i in range(N):
self.m[i] = random.uniform(m_min, m_max)
for j in range(2):
self.r[i][j] = random.uniform(-r_max, r_max)
self.v[i][j] = random.uniform(-v_max, v_max)
self.F[i][j] = 0
self.dt = dt
@cython.cdivision(True)
def compute_F(World w):
"""Compute the force on each body in the world, w."""
cdef int i, j
# Set all forces to zero.
for i in range(w.N):
w.F[i][0] = 0
w.F[i][1] = 0
cdef double sx, sy, s3, tmp, Fx, Fy
for i in range(w.N):
for j in range(i+1, w.N):
sx = w.r[j][0] - w.r[i][0];
sy = w.r[j][1] - w.r[i][1];
s3 = sqrt(sx*sx + sy*sy);
s3 *= s3 * s3;
tmp = w.m[i] * w.m[j] / s3;
Fx = tmp * sx;
Fy = tmp * sy;
w.F[i][0] += Fx;
w.F[i][1] += Fy;
w.F[j][0] -= Fx;
w.F[j][1] -= Fy;
@cython.cdivision(True)
def evolve(World w, int steps):
"""Evolve the world, w, through the given number of steps."""
cdef int _, i
for _ in range(steps):
compute_F(w)
for i in range(w.N):
w.v[i][0] += w.F[i][0] * w.dt / w.m[i]
w.v[i][1] += w.F[i][1] * w.dt / w.m[i]
w.r[i][0] += w.v[i][0] * w.dt
w.r[i][1] += w.v[i][1] * w.dt
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment