Skip to content

Instantly share code, notes, and snippets.

@jwise77
Created April 9, 2024 13:08
Show Gist options
  • Save jwise77/40e6d2fa96a37fadcba4de1f1c72ee7e to your computer and use it in GitHub Desktop.
Save jwise77/40e6d2fa96a37fadcba4de1f1c72ee7e to your computer and use it in GitHub Desktop.
import time
import numpy as np
from mpi4py import MPI
import matplotlib.pyplot as plt
Mtotal = 1000 # system mass [kg]
Npart = 1000 # number of particles
Mpart = Mtotal / Npart
radius = 1 # sphere radius [m]
Vcf = 0.5 # fraction of circular velocity
G = 6.673e-11 # gravitational constant
tstop = 5000 # stop time (seconds)
dt = 10.0
np.random.seed(12345)
#
#
def initialize(N):
# Initialize a uniform distribution in a cube centered on the origin and then restrict to
# sphere. A cube that bounds a sphere has about twice the volume. So to be on the safe
# side (because randomness), I generate `3*Npart` particle positions and then select the
# first `Npart` particles that lie within the sphere.
ninside = 0
while ninside < N:
cubepos = 2*np.random.random([3*N,3]) - 1
r = (cubepos**2).sum(1)**0.5
inside = (r<1)
ninside = inside.sum()
pos = cubepos[inside][:N,:]
r = (pos**2).sum(1)**0.5
# Circular radius of a cross-section (xy-plane) of the sphere at coordinate z of the
# particle. This is the particle's initial orbit for a solid body rotator. We need this to
# calculate the x- and y-components of the initial velocity.
rcirc = (pos[:,:2]**2).sum(1)**0.5
vc = np.sqrt(G * Mtotal * r**2 / radius**3)
vel = np.empty([N,3])
vel[:,0] = -Vcf*vc * (pos[:,1] / rcirc) # ratio is (y/r), which is sin(theta)
vel[:,1] = +Vcf*vc * (pos[:,0] / rcirc) # ratio is (x/r), which is cos(theta)
vel[:,2] = 0.0
return pos, vel
def subrange(comm):
### MODIFY HERE to determine the starting and ending point that an MPI process
### will compute in the force calculation.
start = 0
end = Npart
return start, end, end-start
def calc_accel(comm, pos):
start, end, Nsub = subrange(comm)
# include forces from all particles
idx = np.arange(Npart)
Nsub = end - start
# acceleration for particles being calculated on this MPI process (core)
accel = np.empty([Nsub,3])
# only calculate acceleration for start -> end in the particle list
for i in range(start,end):
rij = pos[idx != i, :] - pos[i,:] # exclude self (idx != i)
rij_mag = (rij**2).sum(1)**0.5
accel[i-start,:] = (rij / rij_mag[:,np.newaxis]**1.5).sum(0)
# multiply by G*M at the very end
accel *= G*Mpart
### MODIFY HERE to share the accelerations from each MPI process to
### all other processes
all_accel = accel # serial only (comment out after parallelized)
return all_accel
def leapfrog_initialstep(comm, pos, vel, dt):
accel = calc_accel(comm, pos)
vhalf = vel + 0.5*dt * accel
return vhalf
def leapfrog(comm, pos, vel, vhalf, dt):
pos = pos + dt * vhalf
anext = calc_accel(comm, pos)
k = dt * anext
# Only need v(t+dt) if we need to calculate associated quantities (e.g. energy) or plot it
vel = vhalf + 0.5*k
vhalf = vhalf + k
return pos, vel, vhalf
#######################
#######################
#######################
comm = MPI.COMM_WORLD
if comm.rank == 0:
pos, vel = initialize(Npart)
else:
pos, vel = None, None
# MODIFY HERE to share the initial conditions that were only generated on the root (rank-0)
# process to the other MPI processes
vhalf = leapfrog_initialstep(comm, pos, vel, dt)
istep = 0
_time = 0.0
nstep = tstop/dt
time0 = time.perf_counter()
while _time < tstop:
pos, vel, vhalf = leapfrog(comm, pos, vel, vhalf, dt)
istep += 1
_time += dt
if (comm.rank == 0) and (istep % (nstep/10) == 0 or istep == 1):
plt.clf()
plt.scatter(pos[:,1], pos[:,2], s=10, c=np.linspace(0,1,Npart), cmap='coolwarm')
plt.title(f"Time = {_time:.0f} s")
plt.xlim(-1,1)
plt.ylim(-1,1)
plt.savefig(f"positions-{istep:04d}.png")
print(f"time = {_time:.0f}/{tstop:.0f} seconds")
time1 = time.perf_counter()
telapsed = time1-time0
if comm.rank == 0:
print(f"Compute time elapsed = {telapsed:.3f} seconds ({nstep/telapsed:.1f} steps/sec // {Npart*nstep/telapsed:.3g} particle-evolve/sec)")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment