def COM(snap, quantity='x', type="default", unit="default"):
''' Computes the centre-of-mass value of a given vector component'''
xunitinfo, x, xscaling_factor, xlabel = UserQuantity(quantity).fetch(type, snap, unit)
m = UserQuantity('m').fetch(type, snap)[1]
com = (x*m).sum()/m.sum()
return xunitinfo, com, xscaling_factor, xlabel+'_COM'
def lagrangian_radii(snap, mfrac=0.5, type="default", unit="default"):
'''Computes the Lagrangian radii from all particles in simulation'''
runitinfo, r, rscaling_factor, rlabel = UserQuantity('r').fetch(type, snap, unit)
m = UserQuantity('m').fetch(type, snap, unit)[1]
# Find particle ids in order of increasing radial distance
porder = np.argsort(r)
m_ordered = m[porder]
mcumulative = np.cumsum(m_ordered)
mtotal = mcumulative[-1]
mlag = mfrac*mtotal
index = np.searchsorted(mcumulative,mlag)
lag_radius = 0.5*(r[porder[index-1]] + r[porder[index]])
return runitinfo, lag_radius, rscaling_factor, 'lag_radius_' + str(mfrac)
def energy_error(snap, etot0, type="default", unit="default"):
'''Computes the energy error of all particles in the simulation'''
vx = UserQuantity('vx').fetch(type, snap, unit)[1]
vy = UserQuantity('vy').fetch(type, snap, unit)[1]
vz = UserQuantity('vz').fetch(type, snap, unit)[1]
m = UserQuantity('m').fetch(type, snap, unit)[1]
gpot = UserQuantity('gpot').fetch(type, snap, unit)[1]
N = m.size
# Loop over all particles and compute
energy = 0.5*m*(vx*vx + vy*vy + vz*vz) - 0.5*m*gpot
etot = np.sum(energy)
error = abs((etot - etot0)/etot0)
return error
def evaluateStack(s, type, snap):
op = s.pop()
if debug_flag:
print op
if op in "+-*/^":
op2 = evaluateStack(s, type, snap)
op1 = evaluateStack(s, type, snap)
return opn[op](op1, op2)
elif op == "pi":
return np.pi
elif op == "e":
return np.e
elif op in functions1arg:
operand = evaluateStack(s, type, snap)
return getattr(np, op)(operand)
elif op in functions2arg:
operand2 = evaluateStack(s, type, snap)
operand1 = evaluateStack(s, type, snap)
return getattr(np, op)(operand1, operand2)
elif op in _KnownQuantities():
return UserQuantity(op).fetch(type, snap, "default")[1]
elif op in snap.sim.GetIntAndFloatParameterKeys():
return float(snap.sim.GetParam(op))
elif re.search('^[-+]?[0-9]+$', op):
return int(op)
elif re.search('[-+]?[0-9]+[.][0-9]*', op):
return float(op)
else:
raise ParseException("Unable to parse string: " + op)
def density_pdf(snap, type="default", nbin=32, rhomin="auto", rhomax="auto"):
'''Calculate the probability density function of the density field'''
rho = UserQuantity("rho").fetch(type, snap)[1]
if rhomin == "auto": rhomin = np.amin(rho)
if rhomax == "auto": rhomax = np.amax(rho)
n = rho.size
rho = np.log10(rho)
bins = np.linspace(log10(rhomin), log10(rhomax), nbin + 1)
#rhopdf = np.zeros(nbin+2)
rhopdf = np.histogram(rho, bins=bins)[0]
#binpos = np.digitize(rho,bins,right=True)
#for i in range(n):
# rhopdf[binpos[i]] += 1
rhopdf = np.float32(rhopdf)
# Normalise logarithmic bins
for j in range(nbin):
rhopdf[j] = rhopdf[j] / (10**bins[j + 1] - 10**bins[j])
print bins
print rhopdf
return bins[0:nbin], rhopdf
def get_array(self, axis, snap):
'''Helper function to get the array of the quantity on the x, y or
rendered axis.
Inputs:
axis
String with the desidered axis ('x', 'y' or 'render')
snap
Snapshot object
Output:
data
Array containing data, in code units
scaling_factor
Scaling factor for unit conversion (multiply data by this
number to get the requested unit)
unitinfo
UnitInfo object with the label and the name of the unit
label
Label of plotted quantity (string before the unit)
'''
quantity = getattr(self, axis + 'quantity')
unit = getattr(self, axis + 'unit')
unitinfo, data, scaling_factor, label = UserQuantity(quantity).fetch(
self._type, snap, unit)
return unitinfo, data, scaling_factor, label
def mass(self, unit='default'):
'''Returns total mass of gas inside blob'''
mass_info = UserQuantity('m').fetch(self.type, self.snap, unit=unit)
mass_type = mass_info[1]
scaling_factor = mass_info[2]
return mass_type[self.ids].sum() * scaling_factor
def time_derivative(snap, quantity, type="default", unit="default", id=None):
'''Return for a given snapshot, the time derivative of a given quantity
of a given type. If possible, use central difference. Otherwise, use
forward/backward difference. If id is not specified, return an array with
the quantity for each particle. Otherwise, return a scalar quantity only
for the given particle.
'''
from data_fetcher import get_time_snapshot
if unit != "default":
raise NotImplementedError("""time_derivative implemented only with default units
at the moment!""")
# Find previous and next snapshots. If first or last snapshot, then
# return the current snapshot to compute a value
try:
snap1 = SimBuffer.get_previous_snapshot_from_object(snap)
except BufferException:
snap1 = snap
try:
snap2 = SimBuffer.get_next_snapshot_from_object(snap)
except BufferException:
snap2 = snap
# Return array of values of quantity. If either are empty, return nan
quantityunitinfo, values1, quantityscaling_factor, label = UserQuantity(quantity).fetch(type, snap1)
values2 = UserQuantity(quantity).fetch(type, snap2)[1]
timeunitinfo, time, timescaling_factor, tlabel = get_time_snapshot(snap)
scaling_factor = quantityscaling_factor/timescaling_factor
unitinfo = UnitInfo()
unitinfo.name= quantityunitinfo.name + "_" + timeunitinfo.name
unitinfo.label= quantityunitinfo.label + "\\ " + timeunitinfo.label + "^{-1}"
# Calculate the time derivative with central difference and return value
tdiff = snap2.t - snap1.t
if values1.size == 0 or values2.size == 0:
timederiv = np.nan
else:
if id == None:
timederiv = (values2 - values1)/tdiff
else:
timederiv = (values2[id] - values1[id])/tdiff
return unitinfo, timederiv, scaling_factor, label+"_t"
def SPH_velocities_sim_frame(self, unit='default'):
'''Returns array of velocities of particles comprising gas blob'''
ndim = self.snap.ndim
velocities = np.zeros((self.n_particles(), ndim))
first_coordinate_info = UserQuantity('vx').fetch(self.type,
self.snap,
unit=unit)
scaling_factor = first_coordinate_info[2]
velocities[:, 0] = first_coordinate_info[1][self.ids]
if ndim > 1:
velocities[:, 1] = UserQuantity('vy').fetch(self.type,
self.snap,
unit=unit)[1][self.ids]
if ndim > 2:
velocities[:, 2] = UserQuantity('vz').fetch(self.type,
self.snap,
unit=unit)[1][self.ids]
return velocities * scaling_factor
def SPH_positions_sim_frame(self, unit='default'):
'''Returns array of positions of particles comprising gas blob'''
ndim = self.snap.ndim
positions = np.zeros((self.n_particles(), ndim))
first_coordinate_info = UserQuantity('x').fetch(self.type,
self.snap,
unit=unit)
scaling_factor = first_coordinate_info[2]
positions[:, 0] = first_coordinate_info[1][self.ids]
if ndim > 1:
positions[:, 1] = UserQuantity('y').fetch(self.type,
self.snap,
unit=unit)[1][self.ids]
if ndim > 2:
positions[:, 2] = UserQuantity('z').fetch(self.type,
self.snap,
unit=unit)[1][self.ids]
return positions * scaling_factor
def get_data(quantity, snap="current",type="default",sim="current",unit="default" ):
'''Returns the array with the data for the given quantity.
The data is returned scaled to the specified unit
Required argument:
quantity :The quantity required. Must be a string
Optional arguments:
type :The type of the particles (e.g. 'star')
snap :Number of the snapshot. Defaults to 'current'
sim :Number of the simulation. Defaults to 'current'
unit :Specifies the unit to use to return the data
'''
simno = get_sim_no(sim)
sim = SimBuffer.get_sim_no(simno)
snapobject = SimBuffer.get_snapshot_extended(sim, snap)
nspecies = snapobject.GetNTypes()
if type=="all":
raise Exception("You requested all particle types to get_data, but we can return only one array!")
fetcher=UserQuantity(quantity)
unitinfo,data,scaling,label=fetcher.fetch(type=type,snap=snapobject,unit=unit)
return data*scaling
def particle_data(snap, quantity, type="default", unit="default", id=None):
'''Return for a given snapshot, a given quantity of a given type. If id is
not specified, return an array with the quantity for each particle. Otherwise,
return a scalar quantity only for the given particle.
'''
unitinfo, values, scaling_factor, label = UserQuantity(quantity).fetch(type, snap, unit=unit)
if values.size == 0:
values_to_return = np.nan
else:
if id == None:
values_to_return = values
else:
values_to_return = values[id]
return unitinfo, values_to_return, scaling_factor, label
def structure_function(snap,
type="default",
nbin=8,
npoints=1000,
rmin=0.001,
rmax=10.0):
'''Calculate the structure function for a given snapshot'''
# Return all relevant particle data (positions and velocity)
x = UserQuantity("x").fetch(type, snap)[1]
y = UserQuantity("y").fetch(type, snap)[1]
z = UserQuantity("z").fetch(type, snap)[1]
vx = UserQuantity("vx").fetch(type, snap)[1]
vy = UserQuantity("vy").fetch(type, snap)[1]
vz = UserQuantity("vz").fetch(type, snap)[1]
n = x.size
r = np.zeros(n)
vsqd = np.zeros(n)
# Create logarithmic bins based on inputted range
bins = np.linspace(log10(rmin), log10(rmax), nbin + 1)
vsqdmean = np.zeros(nbin + 2)
npart = np.zeros(nbin + 2)
# Loop through a random selection of points
for j in range(npoints):
i = random.randrange(0, n - 1)
x0 = x[i]
y0 = y[i]
z0 = z[i]
vx0 = vx[i]
vy0 = vy[i]
vz0 = vz[i]
r[:] = (x[:] - x0)**2
r[:] += (y[:] - y0)**2
r[:] += (z[:] - z0)**2
r = np.log10(np.sqrt(r))
vsqd[:] = (vx[:] - vx0)**2
vsqd[:] += (vy[:] - vy0)**2
vsqd[:] += (vz[:] - vz0)**2
# Now discretise values into bins and sum up values
binpos = np.digitize(r, bins, right=True)
for jj in range(npoints):
ii = random.randrange(0, n - 1)
vsqdmean[binpos[ii]] += vsqd[ii]
npart[binpos[ii]] += 1
# Normalise velocity bins to arithmetic mean values
for j in range(nbin):
if npart[j] > 0: vsqdmean[j] = vsqdmean[j] / npart[j]
vsqdmean = np.log10(vsqdmean)
# Print out results when finished (for now)
print bins
print vsqdmean
print npart
return bins[0:nbin], vsqdmean[0:nbin]
def extract_discs(snapno,
sim,
type='default',
eccenlimit=0.9,
distancelimit=1.,
limiteigenvalues=0.2):
'''This function takes a snapshot and a simulation number ("current" is
also fine) and looks for which particles are bound to the stars. It
returns a tuple, consisting of an Ambient_gas object (representing the
gas that is not bound to any star) and of a list of Disc objects, one for
each star.
'''
snap = SimBuffer.get_snapshot_extended(sim, snapno)
parameters = dict(
eccenlimit=eccenlimit,
distancelimit=distancelimit,
limiteigenvalues=limiteigenvalues,
)
# Query the number of dimensions
ndim = snap.ndim
# Fetch the data - we use code units
# First extract coordinates, velocities and masses for the given type
x_type = UserQuantity('x').fetch(type, snap)[1]
vx_type = UserQuantity('vx').fetch(type, snap)[1]
x_star = UserQuantity('x').fetch('star', snap)[1]
vx_star = UserQuantity('vx').fetch('star', snap)[1]
if ndim > 1:
y_type = UserQuantity('y').fetch(type, snap)[1]
vy_type = UserQuantity('vy').fetch(type, snap)[1]
y_star = UserQuantity('y').fetch('star', snap)[1]
vy_star = UserQuantity('vy').fetch('star', snap)[1]
if ndim > 2:
z_type = UserQuantity('z').fetch(type, snap)[1]
vz_type = UserQuantity('vz').fetch(type, snap)[1]
z_star = UserQuantity('z').fetch('star', snap)[1]
vz_star = UserQuantity('vz').fetch('star', snap)[1]
m_type = UserQuantity('m').fetch(type, snap)[1]
m_star = UserQuantity('m').fetch('star', snap)[1]
n_star = snap.GetNparticlesType('star')
if ndim == 2:
owner = extract_disc_cython.flag_owner2d(x_type, y_type, vx_type,
vy_type, m_type, x_star,
y_star, vx_star, vy_star,
m_star, parameters)
elif ndim == 3:
owner = extract_disc_cython.flag_owner3d(x_type, y_type, z_type,
vx_type, vy_type, vz_type,
m_type, x_star, y_star,
z_star, vx_star, vy_star,
vz_star, m_star, parameters)
# Loops over all stars and creates the disc objects
disclist = []
for istar in range(n_star):
ids = (owner == istar)
disc = Disc(istar, ids, snap, type)
disclist.append(disc)
print 'mass disc number', istar, ':', disc.mass()
# Create the ambient gas object
ids = (owner == -1)
ambient = Ambient_gas(ids, snap, type)
print 'ambient gas mass:', ambient.mass()
# Return a tuple with the ambient gas and a list of discs
return (ambient, disclist)