def speedcurve(result, t=0):
velocities = utils.get_velocities(result)[t]
positions = utils.get_positions(result)[t]
plt.plot(
np.linalg.norm(positions, axis=1)[1:],
np.linalg.norm(velocities, axis=1)[1:], '.')
plt.xlim(0, sc.Rmw * 1.2)
# plt.ylim(0, 0.1e6)
plt.xlabel("r (m)")
plt.ylabel("v (m/s)")
plt.show()
def energy(result, t=0):
masses = utils.get_masses(result)[t]
r = utils.get_positions(result)[t]
pmag = linMom(result, t)[3]
Ekin = np.sum(pmag**2 / (2 * masses))
Epot = 0
for i in range(len(masses)):
for j in range(i + 1, len(masses)):
dist = np.linalg.norm(r[i] - r[j])
Epot -= sc.G * masses[i] * masses[j] / dist
return Ekin + Epot, Ekin, Epot
def load_data(filename, mass_scale, show_dm, no_show_m, darkmatter_intensity):
print("Start loading data...")
preloaded = cs.Result.load(filename).numpy()
print("Selecting data...")
func = np.vectorize(lambda b: (show_dm and b.dark_matter) or
(not no_show_m and not b.dark_matter))
preloaded = preloaded[:, func(preloaded[0])]
print(preloaded.shape)
print("Extracting positions and adding color...")
positions = get_positions(preloaded).astype(np.float32)
c = bodies_to_color(preloaded[0], mass_scale)
if show_dm:
dm_index = [i for i, b in enumerate(preloaded[0]) if b.dark_matter]
c[dm_index] = [0, darkmatter_intensity, 0]
colors = gloo.VertexBuffer(c)
print("Data loaded and ready for GPU")
return positions, colors
bodies = [cs.Body3(pos=np.zeros(3), vel=np.zeros(3), mass=M)]
for i in range(1,n):
bodies.append(cs.Body3(pos=posarray[i], vel=velarray[i], mass=massarray[i]))
return cs.BodyList3(np.array(bodies))
thetamax = 0.5
n_steps = 500 # int(30/1e-4)
begin = time.time()
result = cs.LeapFrogSaveC(genEGalaxy(2,sc.Msgra,space=True,elliptical=True,spiralarms=2), 1e12, n_steps, thetamax, sc.G)
end = time.time()
s = utils.get_positions(result)
print(s)
print(s[100])
plt.scatter(s[0][:,0],s[0][:,1])
plt.show()
large_xyz = 1e20
medium_xyz = 1e19
large_limits = {"xlim": (-large_xyz, large_xyz), "ylim": (-large_xyz, large_xyz), "zlim": (-large_xyz, large_xyz)}
medium_limits = {"xlim": (-medium_xyz, medium_xyz), "ylim": (-medium_xyz, medium_xyz), "zlim": (-medium_xyz, medium_xyz)}
# plotting.movie3d(s, np.arange(2), until_timestep=1000, skip_steps=10, mode="line", **medium_limits)
until_timestep = int(n_steps)
starting_angle = 225 # default 270, 0 for sun zoom
rotation_speed = 0 # default 40
def angMom(result, t=0):
r = utils.get_positions(result)[t]
p = linMom(result, t)[0]
L = np.cross(r, p)
sumL = np.sum(L, axis=0)
return L, sumL, np.linalg.norm(sumL)
def movie3d(dataset, particle_indices, **CONFIG):
"""Makes a 3D plot movie of particles over time.
-Dataset is a numpy array containing an array of 3D coordinates for each time step
-particle_indices is a list of the indices of the particles you want to plot
-CONFIG allows you to pass a bunch of plotting parameters. The parameters present in VIDEO_CONFIG and FRAME_CONFIG
(see source code) can be used to customize your plot. All other keywords are passed on to the matplotlib
ax.plot3D function, so matplotlib plotting settings such as color, linestyle or marker style can be set this
way
The parameter particle_config can be used to stylize each individual particle. This config will override any
settings set by the "mode" parameter or by other parameters set in CONFIG. Usage: list of dicts.
Explanation of some CONFIG parameters:
- mode (str): sets some default parameters for the matplotlib ax.plot3D call. Current options are "line" and "point"
- preview (bool): show a preview still before rendering. By default renders 1st frame, change with ⏎
- preview_frame (int): choose which frame to preview
- skip_steps (int): how many time steps to skip between each frame. By default no steps are skipped,
but if your time step is small, consider increasing this value for faster rendering
- rotation speed (float): how many arcseconds the scene will rotate per frame (default 30)
- init_azimuth (float): viewing angle around z-axis in the first frame
- elevation (float): viewing angle relative to x-y plane in degrees
- axes (bool): true by default
- grid (bool): false by default
General tips: if rendering takes too long, increase skip_steps. If video plays too fast, decrease fps.
"""
VIDEO_CONFIG = {
"filename": 'test.mp4',
"preview": True,
"preview_frame": 0,
"until_timestep": len(dataset) - 1,
"skip_steps": 1,
"fps": 30,
"artist": "",
"bitrate": 1800,
}
FRAME_CONFIG = {
"mode": "line",
"xlim": (-100, 100),
"ylim": (-100, 100),
"zlim": (-100, 100),
"rotation_speed": 30, # arcseconds per frame
"init_azimuth": 270,
"elevation": 10,
"axes": True,
"grid": False,
"particle_config": None
}
dataset = utils.get_positions(dataset)
plotting_config = {}
frame_config = {}
video_config = {}
for key, value in CONFIG.items():
if key in FRAME_CONFIG:
frame_config[key] = value
elif key in VIDEO_CONFIG:
video_config[key] = value
else: # other kwargs are passed on to mpl ax.plot3D
plotting_config[key] = value
frame_config = {**FRAME_CONFIG, **frame_config}
video_config = {**VIDEO_CONFIG, **video_config}
frame_config["skip_steps"] = video_config[
"skip_steps"] # pass along to data gen to normalize rotation speed
if video_config['preview']:
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d')
data_gen(ax, dataset, video_config["preview_frame"], particle_indices,
frame_config, plotting_config)
plt.show()
fig = plt.figure()
# ax = plt.axes(projection='3d')
ax = fig.add_subplot(1, 1, 1, projection='3d')
data_gen_formatted = lambda index: data_gen(
ax, dataset, index, particle_indices, frame_config, plotting_config)
grav_ani = animation.FuncAnimation(fig,
data_gen_formatted,
frames=np.arange(
0, video_config["until_timestep"],
video_config["skip_steps"]),
interval=30,
blit=False)
Writer = animation.writers['ffmpeg']
writer = Writer(fps=video_config["fps"],
metadata=dict(artist=video_config["artist"]),
bitrate=video_config["bitrate"])
filename = video_config["filename"]
print("Started writing video to {}".format(filename))
grav_ani.save(filename, writer=writer)
print("Video saved at {}".format(filename))
allbodies = np.concatenate((bodies,bodiesDM))
print(len(allbodies))
return cs.BodyList3(np.array(allbodies))
thetamax = 0.5
n_steps = 5000 # int(30/1e-4)
begin = time.time()
result = cs.LeapFrogSaveC(genDMG(100,M=sc.Msgra,spherical=True,nDM=1000), 1e12, n_steps, thetamax, sc.G)
end = time.time()
result.save("testDMG.binv")
r = utils.get_positions(result.numpy())
rnorm = np.sqrt(r[-1][:,0]**2+r[-1][:,1]**2+r[-1][:,2]**2)
rnormM = rnorm[0:100]
rnormDM = rnorm[101:1100]
v = utils.get_velocities(result.numpy())
vnorm = np.sqrt(v[-1][:,0]**2+v[-1][:,1]**2+v[-1][:,2]**2)
vnormM = vnorm[0:100]
vnormDM = vnorm[101:1100]
plt.scatter(rnormM,vnormM)
plt.scatter(rnormDM,vnormDM)
plt.show()
"""
large_xyz = 1e20
medium_xyz = 2*1e19