Python code examples

Set up a 3d plot

import matplotlib as mpl
mpl.use('TkAgg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d

boxsize = ···some value···

#--- Set up the figure and axis ranges. ---#
figure = plt.figure()
axes = figure.add_subplot(111, projection='3d')
axes.set_xlim3d([0,boxsize])
axes.set_ylim3d([0,boxsize])
axes.set_zlim3d([0,boxsize])
axes.set_box_aspect([1,1,1])

···plot something···

#--- Show the plot. ---#
# plt.savefig('plot.pdf')
plt.show()

Read data from a snapshot file

import g3read

#--- Read the data from the snapshot. ---#
snap = g3read.GadgetFile('snap_144')
print(snap.header.__dict__.keys())
boxsize = snap.header.BoxSize
#--- We are interested in particle type 0 (gas particles). ---#
pos  = snap.read_new('POS ', 0)
vel  = snap.read_new('VEL ', 0)
mass = snap.read_new('MASS', 0)
npart = len(pos) # We have this many particles.

Create an average-velocity grid from particle positions and velocities

import math
import numpy as np

#--- Set up the velocity grid. ---#
ngrid = 20
vel_grid  = np.zeros((ngrid,ngrid,ngrid,3))
mass_grid = np.zeros((ngrid,ngrid,ngrid))

#--- Find grid index for each particle and update grid cell values. ---#
for i in range(npart):
  ix = math.floor(pos[i,0]/boxsize*ngrid)
  iy = math.floor(pos[i,1]/boxsize*ngrid)
  iz = math.floor(pos[i,2]/boxsize*ngrid)
  vel_grid [ix,iy,iz,:] += mass[i]*vel[i,:] # Weighted sum of velocities.
  mass_grid[ix,iy,iz]   += mass[i]          # Total mass in grid cell.

#--- Weighted average velocity: ---#
for ix in range(ngrid):
  for iy in range(ngrid):
    for iz in range(ngrid):
      if mass_grid[ix,iy,iz] != 0.:
        vel_grid[ix,iy,iz,:] /= mass_grid[ix,iy,iz]

Create a single streamline

nsteps = 100
dt = 1.0

#--- Starting position. ---#
x0 = ···some starting value x···
y0 = ···some starting value y···
z0 = ···some starting value z···

#--- Array to hold the streamline. ---#
stream = np.zeros((nsteps,3))
stream[0,:] = [x0, y0, z0]

#--- Do steps to create the streamline. ---#
for i in range(1,nsteps):
#--- Find grid cell we're in (to get the current velocity). ---#
  ix0 = math.floor(x0/boxsize*ngrid)
  iy0 = math.floor(y0/boxsize*ngrid)
  iz0 = math.floor(z0/boxsize*ngrid)
#--- Update position by following the velocity for a small distance. ---#
  x1 = x0 + vel_grid[ix0,iy0,iz0,0]*dt
  y1 = y0 + vel_grid[ix0,iy0,iz0,1]*dt
  z1 = z0 + vel_grid[ix0,iy0,iz0,2]*dt
#--- If we have left the box, wrap around to the other side. ---#
  x0 = x1 - boxsize*math.floor(x1/boxsize)
  y0 = y1 - boxsize*math.floor(y1/boxsize)
  z0 = z1 - boxsize*math.floor(z1/boxsize)
#--- Save position to streamline. ---#
  stream[i,:] = [x0, y0, z0]
#--- If we have wrapped, don't connect the lines. ---#
  if(x0!=x1 or y0!=y1 or z0!=z1):
    stream[i,:] = np.nan

axes.plot3D(stream[:,0],stream[:,1],stream[:,2])