def add_N2(axis, n0, p):
""" Makes a field that when projected is a correction to the column density used
in calculating the polarization fraction. """
field_horizontal = {'x': 'By', 'y': 'Bz', 'z': 'Bx'}[axis]
field_vertical = {'x': 'Bz', 'y': 'Bx', 'z': 'By'}[axis]
def _N2_local(field, data):
""" Calculate n2 = n * (cos^2(gamma)/2 - 1/3) where gamma is the
inclination angle between the B field and the plane of the sky.
cos(gamma) = B_sky dot B / |B_sky|*|B|
e.g. Line of sight along z axis. B = (Bx,By,Bz); B_sky = (Bx,By,0)
cos(gamma) = Bx^2 + By^2 / (sqrt(Bx^2 + By^2) * sqrt(Bx^2 + By^2 + Bz^2))
cos(gamma)^2 = Bx^2 + By^2 / (Bx^2 + By^2 + Bz^2)
This function returns n2 as a dimensionless value since epsilon was calculated as dimensionless"""
B_sq = data['Bx']**2.0 + data['By']**2.0 + data['Bz']**2.0
cos_gamma_sq = (data[field_horizontal]**2.0 +
data[field_vertical]**2.0) / B_sq
epsilon = np.ones(data['density'].shape)
n = data['density']
epsilon[n <= n0] = (n.v)[n <= n0]
epsilon[n > n0] = (n0**(1 - p) * n.v**p)[n > n0]
return epsilon * (0.5 * cos_gamma_sq - 1 / 3)
yt.add_field('N2%s_n0-%04d_p-%d' % (axis, n0, p),
units='dimensionless',
function=_N2_local)
def add_epsilon(axis, factor):
"""
makes an epsilon field for yt.
epsilon is a scaling factor that scales the stokes parameters
for high density regions where there is depolarization.
Above a visual extinction of 3 mag.
"""
n0, n0_str = get_n0(factor)
NH_Av = 1.8e21 # [atoms/cm^2/mag]
N_avg = 1000 * 4.6 * 3.09e18 # [atoms/cm^2 * code_length^2)]
def _epsilon_local(field, data):
"""
Calculate epsilon.
epsilon = code_density/(1.8e21 atoms/cm^2/mag) * 1000 atoms/cm^3/code_density * 4.6pc/code_length
= code_density * N_avg/NH_Av
[epsilon] = mag/code_length
-> Extinction, Av = Integral of epsilon along line of sight.
"""
p = data.get_field_parameter('p')
if p is None: p = 0
n = data['density']
epsilon = np.ones(data['density'].shape)
epsilon[n <= n0] = (n.v)[n <= n0] * N_avg / NH_Av
epsilon[n > n0] = (n0**(1 - p) * n.v**p)[n > n0] * N_avg / NH_Av
return epsilon
yt.add_field('epsilon_n0-%s_%s' % (n0_str, axis), function=_epsilon_local)
return n0_str
def add_unweighted_stokes(axis):
"""makes a stokes field for yt.
axis should be x,y,z.
These should test the projection and the rest of the pipeline.
Second it's an interesting demonstration of what Stokes traces."""
field_horizontal = {'x': 'By', 'y': 'Bz', 'z': 'Bx'}[axis]
field_vertical = {'x': 'Bz', 'y': 'Bx', 'z': 'By'}[axis]
def _unweighted_Q_local(field, data):
"""This function calculates the Stokes Parameter "Q"."""
B_sq = data['Bx']**2.0 + data['By']**2.0 + data['Bz']**2.0
return (data[field_horizontal]**2.0 - data[field_vertical]**2.0) / B_sq
yt.add_field('unweighted_Q%s' % axis,
units='dimensionless',
function=_unweighted_Q_local)
def _unweighted_U_local(field, data):
"""Makes stokes U."""
B_sq = data['Bx']**2.0 + data['By']**2.0 + data['Bz']**2.0
return 2 * (data[field_horizontal]) * (data[field_vertical]) / B_sq
yt.add_field('unweighted_U%s' % axis,
units='dimensionless',
function=_unweighted_U_local)
def add_fields():
"""Add three Eulerian fields Sunrise uses"""
def _MetalMass(field, data):
return data["Metallicity"] * data["CellMassMsun"]
def _convMetalMass(data):
return 1.0
add_field("MetalMass",
function=_MetalMass,
convert_function=_convMetalMass)
def _initial_mass_cen_ostriker(field, data):
# SFR in a cell. This assumes stars were created by the Cen & Ostriker algorithm
# Check Grid_AddToDiskProfile.C and star_maker7.src
star_mass_ejection_fraction = data.ds.get_parameter(
"StarMassEjectionFraction", float)
xv1 = ((data.ds["InitialTime"] - data["creation_time"]) /
data["dynamical_time"])
denom = (1.0 - star_mass_ejection_fraction *
(1.0 - (1.0 + xv1) * np.exp(-xv1)))
minitial = data["ParticleMassMsun"] / denom
return minitial
add_field("InitialMassCenOstriker", function=_initial_mass_cen_ostriker)
def getPointData(file, posx, posy):
"""get all unk data (plus speed, soundspeed and fermiDeg) from a
single cell in a file to track.
otpn: pointTrack
# get cell masses
# log.warning('Assuming cylindrical slice')
# dx = data['path_element_x']
# dy = data['path_element_y']
# r = data['x']
# cylvol = 2.0*np.pi*dy*dx*r
# cell_masses = cylvol*data['dens']
# print('cell mass (calc <> data)', cell_masses, data['cell_mass'])
# print('cell vol (calc <> data)', cylvol, data['cell_volume'])
Args:
file(str): filepath.
posx(float): x-axis position.
posy(float): y-axis position.
"""
ds = yt.load(file)
# read in custom fields
cfl = ['speed', 'sound_speed', 'fermiDeg']
for f in cfl:
if f in dir(ytf):
meta = getattr(ytf, '_' + f)
yt.add_field(("flash", f), function=getattr(ytf, f), **meta)
ad = ds.all_data()
cellsize = np.min(ad['dx'].v)
sp = ds.sphere([posx, posy, 0.0], cellsize)
flds, sps = getFields(ds.field_list)
sps = [s.ljust(4) for s in sps]
qu = sp.quantities
ffl = [
'x', 'y', 'cell_volume', 'cell_mass', 'sound_speed',
'courant_time_step', 'dynamical_time', 'path_element_x',
'path_element_y'
]
allf = ffl + flds + sps
# there's only one cell so min/max is irrelevant, nor the field queried
vals = qu.sample_at_max_field_values('dens', sample_fields=allf)
vals = [x.v for x in vals[1:]] # drop the first value (repeated)
# write all data from the cell to file
# use radius as an alias for time so that datamatrix works
# change x and y to posx posy again so that dm wont take them as n or p.
allf[0] = 'posx'
allf[1] = 'posy'
metatxt = ['# radius {}'.format(' '.join(allf))]
parsedvals = ['{:.8E}'.format(float(ds.current_time))]
parsedvals += ['{:.8E}'.format(x) for x in vals]
metatxt.append(' '.join(parsedvals))
dest, num, name = setFolders(file, 'pointTrack')
with open(name + '.meta', 'w') as f:
f.write('\n'.join(metatxt))
print("Wrote: {}".format(name + '.meta'))
def getLineout(fname,
fields=['density', 'temperature', 'pressure'],
species=True,
radius=5e11,
geom='cartesian',
direction=[],
origin=[0.0, 0.0, 0.0],
srcnames=True):
"""Returns a np.array with radius, dens, temp, pres and species specified.
Args:
fname(str): filename to extract data from.
fields(str list): specify fields to extract (species are added).
species(bool): toggle nuclide data.
radius(float): reach of lineout.
geom(str): geometry (['cartesian'], 'spherical').
direction(list of float): angles of lineout.
takes x/(x,z) as normals for 2D/3D.
also sets dimensionality.
origin(float list): override origin of ray to build profile.
srcnames(bool): passthrough to hdf5yt.getFields.
Returns:
dblock (numpy array): matrix with fields as columns.
species (list of str): names of species in the checkpoint.
"""
for f in fields:
if f in dir(ytf):
meta = getattr(ytf, '_' + f)
yt.add_field(("flash", f), function=getattr(ytf, f), **meta)
ds = yt.load(fname)
# spherical (r, theta (polar), phi(azimuth))
# cartesian (x, y, z)
cname, bearing = getBearing(direction, geom=geom)
# print("calculated bearing: ",bearing)
# ray = ds.ray([0.0, 0.0, 0.0], radius*bearing)
ray = ds.ray(origin, radius * bearing)
rs = np.argsort(ray['t'])
dblock = ray[cname][rs].value
for f in fields:
dblock = np.vstack((dblock, ray[f][rs].value))
_, sps = getFields(ds.field_list)
if species:
for s in sps:
# force ('flash', s) to ensure it retireves the checkpoint data
# instead of a yt variable.
try:
dblock = np.vstack(
(dblock, ray[('flash', "{}".format(s))][rs].value))
except YTFieldNotFound:
dblock = np.vstack(
(dblock, ray[('flash', "{:<4}".format(s))][rs].value))
_, sps = getFields(ds.field_list, srcnames=srcnames)
return dblock, sps
def getDerFieldInfo(Param_Dict, ComboBox_Dict, Misc_Dict, dialog):
"""If the user has filled everything in correctly and presses apply, this
function is called to store the function"""
# get all of the information from the dialog:
fieldName = dialog.fieldName
displayName = dialog.displayName
fieldFunction = dialog.fieldFunction
unit = dialog.unit
dim = dialog.dim
override = dialog.override
if dim == 1:
dim = "dimensionless"
text = dialog.functionText
# Add the field to the known fields in Param_Dict and comboboxes:
for axis in ["X", "Y", "Z"]:
Param_Dict[axis + "Fields"].insert(3, fieldName)
ComboBox_Dict[axis + "Axis"].insertItem(3, fieldName)
if axis != "X":
ComboBox_Dict[axis + "Weight"].insertItem(3, fieldName)
Param_Dict["WeightFields"].insert(3, fieldName)
yt.add_field(("gas", fieldName),
function=fieldFunction,
units=unit,
dimensions=dim,
force_override=override,
display_name=displayName,
take_log=False)
reloadFiles(Param_Dict, Misc_Dict)
# Save the parameters in Param_Dict for the ScriptWriting function
Param_Dict["NewDerFieldDict"][fieldName] = {}
Param_Dict["NewDerFieldDict"][fieldName]["Override"] = override
Param_Dict["NewDerFieldDict"][fieldName]["DisplayName"] = displayName
Param_Dict["NewDerFieldDict"][fieldName]["FunctionText"] = text
if unit == "auto":
Param_Dict["NewDerFieldDict"][fieldName]["Unit"] = dialog.baseUnit
else:
Param_Dict["NewDerFieldDict"][fieldName]["Unit"] = unit
Param_Dict["NewDerFieldDict"][fieldName]["Dimensions"] = dim
if unit == "auto":
unit = fieldFunction(
"Hello, I'm a fancy placeholder, why did you \
call me?", dialog.sp).units
ds = Param_Dict["CurrentDataSet"]
Param_Dict["FieldUnits"][fieldName] = yt.YTQuantity(ds.arr(1, unit)).units
GUILogger.log(
29, f"The new field <b>{fieldName}</b> with {dim}-dimension \
has been added successfully.")
GUILogger.info("You can find it in the comboBoxes for the field selection,"
" and it will be added to every new dataset.")
def run_stuff():
# different data - FLASH
filename = '~/data/GasSloshing/sloshing_nomag2_hdf5_plt_cnt_0100'
sphere_rad = 15.0 # in kpc
#Enzo
filename = '~/data/IsolatedGalaxy/galaxy0030/galaxy0030'
sphere_rad = 200.0 # in kpc
outfile = 'galsurfaces'
rho = 1e-27 # for each surface
trans = 1.0 # for transparency of each surface
color_field = 'temperature' # color your surface by this
pf = yt.load(filename)
# emissivity of the material
# this needs to be a combination of the color_field and surface field
#def _Emissivity(field, data):
# return (eval("data['gas','density']*data['density']*np.sqrt(data['gas','temperature'])"))
def _Emissivity(field, data):
return (data['gas', 'density'] * data['density'] *
np.sqrt(data['gas', 'temperature']))
#yt.add_field(("gas","emissivity"), units="g**2*K**0.5/cm**6", function=_Emissivity)
yt.add_field("emissivity",
units="g**2*K**0.5/cm**6",
function=_Emissivity,
force_override=True)
# for testing
#yt.SlicePlot(pf, 'z', "emissivity", width = (200.0, 'kpc')).save('~/Desktop/mytest.png')
dd = pf.sphere("max", (sphere_rad, "kpc"))
surf = pf.surface(dd, 'density', rho)
#surf.export_obj(outfile, transparency = trans,
# color_field=color_field, emit_field = 'emissivity')
emit_field_name = ('gas', 'emissivity')
emit_field_max = None
emit_field_min = None
color_field_max = None
color_field_min = None
color_map = "algae"
vertices, colors, alpha, emisses, colorindex = surf.export_blender(
transparency=trans,
color_field=color_field,
emit_field=emit_field_name,
color_map=color_map,
plot_index=0,
color_field_max=color_field_max,
color_field_min=color_field_min,
emit_field_max=emit_field_max,
emit_field_min=emit_field_min)
return vertices
def add_stokes(axis, n0, p):
"""makes a stokes field for yt.
axis should be x,y,z."""
field_horizontal = {'x': 'By', 'y': 'Bz', 'z': 'Bx'}[axis]
field_vertical = {'x': 'Bz', 'y': 'Bx', 'z': 'By'}[axis]
def _Q_local(field, data):
"""This function calculates the Stokes Parameter "Q" along an axis x, y, or z.
Makes use of the depolarization factor "epsilon" using a power exponent.
"""
epsilon = np.ones(data['density'].shape)
n = data['density']
B_sq = data['Bx']**2.0 + data['By']**2.0 + data['Bz']**2.0
epsilon[n <= n0] = (n.v)[n <= n0]
epsilon[n > n0] = (n0**(1 - p) * n.v**p)[n > n0]
return (epsilon * (((data[field_horizontal])**2.0) -
((data[field_vertical])**2.0)) / B_sq)
print 'adding yt field Q%s_n0-%04d_p-%d' % (axis, n0, p)
yt.add_field('Q%s_n0-%04d_p-%d' % (axis, n0, p),
units='dimensionless',
function=_Q_local,
force_override=True)
def _U_local(field, data):
"""Makes stokes U."""
epsilon = np.ones(data['density'].shape)
n = data['density']
B_sq = data['Bx']**2.0 + data['By']**2.0 + data['Bz']**2.0
epsilon[n <= n0] = (n.v)[n <= n0]
epsilon[n > n0] = (n0**(1 - p) * n.v**p)[n > n0]
return (2.0 * epsilon * ((data[field_horizontal]) *
(data[field_vertical])) / B_sq)
print 'adding yt field U%s_n0-%04d_p-%d' % (axis, n0, p)
yt.add_field('U%s_n0-%04d_p-%d' % (axis, n0, p),
units='dimensionless',
function=_U_local,
force_override=True)
def add(field):
fDict = {'HI_Column_Density': _HI_Column_Density,
'Warm_Gas_Density': _Warm_Gas_Density,
'Hot_Gas_Density': _Hot_Gas_Density,
'Cold_Gas_Density': _Cold_Gas_Density,
'All_Gas_Density': _All_Gas_Density,
'HI_Analytical_HM96': _HI_Analytical_HM96,
'HI_Analytical_HM12': _HI_Analytical_HM12,
'HI_Analytical_HM01': _HI_Analytical_HM01,
'HI_NumberDensity_2': _HI_NumberDensity_2,
'CellVolume' : _CellVolume}
display_name_dict = {'HI_Column_Density': r"N$_{\rm{HI}}$",
'Warm_Gas_Density': r"\rm{g} \rm{cm}^{-3}",
'Cold_Gas_Density': r"\rm{g} \rm{cm}^{-3}",
'Hot_Gas_Density': r"\rm{g} \rm{cm}^{-3}",
'All_Gas_Density': r"\rm{g} \rm{cm}^{-3}",
'HI_Analytical_HM96': r"\rm{cm}^{-3}",
'HI_Analytical_HM12': r"\rm{cm}^{-3}",
'HI_Analytical_HM01': r"\rm{cm}^{-3}",
'HI_NumberDensity_2': r"\rm{cm}^{-3}",
'CellVolume': 'Volume'}
units_dict = {'HI_Column_Density': "cm**(-2)",
"Warm_Gas_Density" : "g/cm**3",
"Cold_Gas_Density" : "g/cm**3",
"Hot_Gas_Density" : "g/cm**3",
"All_Gas_Density" : "g/cm**3",
"HI_Analytical_HM96" : "cm**(-3)",
"HI_Analytical_HM12" : "cm**(-3)",
"HI_Analytical_HM01" : "cm**(-3)",
"HI_NumberDensity_2" : "cm**(-3)",
"CellVolume" : "cm**(3)"}
yt.add_field(field, function=fDict[field], display_name = display_name_dict[field], units=units_dict[field])
def add_field(field):
"""
Wrapper to yt.add_field function. Supply field name from
list of possible fields given by derived_fields.fields
Parameters
----------
field : string
Name of field to add. Must already have function
definitions contained within the derived_fields.
See derived_fields.fields for list of field names.
"""
if field == 'gas_column_density':
print " ----- WARNING ----- "
print "Column calculated using horrible assumptions on"
print "composition and mean molecular weight"
print " ----- ------ ------ "
yt.add_field(field, function = function_dict[field],
display_name = display_name_dict[field],
units = units_dict[field])
# Define some constant parameters to be used.
mp = 1.6726e-24 * gram # g
mu = 1.2924
kb = 1.3806e-16 *erg / Kelvin # erg K-1
GNewton = 6.6743e-8 * cm**3 / (gram * second**2 )# cm3 g-1 s-2
Msun = 1.9884e33 * gram
mm = mu*mp
ppc = 3.0856776e18
# Create a derived field.
@derived_field(name="numdens", units="1/cm**3", force_override=True)
def numdens(field, data):
return data["dens"].in_cgs()/mm
yt.add_field('numdens', function=numdens, units="1/cm**3", force_override=True)
# Create all the derived fields calculating the force, per unit volume.
# Generalize for all the other clouds.
if Cloud_name == "M4e3":
px = 180
py = -25
pz = 10
cloud_alpha = 0.4
rad = 50
real_rad = 30
exec(open('/home/lawsmith/jYT/my_settings.py').read())
clusterdir = '/data/groups/ramirez-ruiz/lawsmith/FLASH4.3/'
savepath = '/data/groups/ramirez-ruiz/lawsmith/FLASH4.3/lux_slices/'
def _unbound_mass(field, data):
e = 0.5 * (data['velocity_magnitude'].d)**2 + data['gpot'].d
factor = np.sign(e)
return factor * data['cell_mass'].d / M_SUN_CGS
yt.add_field(("gas", "unbound_mass"),
display_name=r'$M_{\rm ej}$',
function=_unbound_mass,
take_log=False,
units=None,
force_override=True)
if args.chkplt == 'chk':
LOAD_FILES = clusterdir + args.run + '/multitidal_hdf5_chk_' + args.files
elif args.chkplt == 'plt':
LOAD_FILES = clusterdir + args.run + '/multitidal_hdf5_plt_cnt_' + args.files
yt.enable_parallelism()
ts = yt.DatasetSeries(LOAD_FILES)
t_array = []
unbound_m_array = []
bound_m_array = []
from radial_data_nozeros import *
from astropy.table import Table
from holoviews.operation.datashader import aggregate, datashade, dynspread, shade
from holoviews.operation import decimate
from holoviews.operation import histogram
from get_halo_center import get_halo_center
track_name = '/Users/dalek/data/Jason/symmetric_box_tracking/complete_track_symmetric_50kpc'
#track_name = '/astro/simulations/FOGGIE/halo_008508/complete_track_symmetric_50kpc'
def _cooling_criteria(field,data):
return data['cooling_time'] / ((data['dx']/data['sound_speed']).in_units('s'))
yt.add_field(("gas","cooling_criteria"),function=_cooling_criteria,units=None)
def sym_refine_box(ds,halo_center):
dx = ds.arr(20.,'kpc').in_units('code_length').value
dy = ds.arr(50.,'kpc').in_units('code_length').value
box_left = [halo_center[0]-dx, halo_center[1]-dy, halo_center[2]-dx]
box_right = [halo_center[0]+dx, halo_center[1]+dy, halo_center[2]+dx]
refine_box = ds.r[box_left[0]:box_right[0],
box_left[1]:box_right[1],
box_left[2]:box_right[2]]
return refine_box
def initial_center_guess(ds,track_name):
track = Table.read(track_name, format='ascii')
track.sort('col1')
zsnap = ds.get_parameter('CosmologyCurrentRedshift')
def get_N_mod(runfile,
d=(1, 'kpc'),
bmaj=(1.132, 'arcsec'),
width=None,
center='max',
cmap='magma',
resolution=512,
fontsize=13,
title=None,
vmin=None,
vmax=None,
savefig='',
show=True,
mask=True,
verbose=False):
"""
Obtain the column density of oH2Dp from the model file via averaging
over the volume defined by region and radius.
[Parameters]
runfile : (string) Name of the hdf5 file that contains the model snapshot.
d : (tuple) Artificial distance to the source used to convert angular scales into a physical width. Must be a tuple (val, "unit").
bmaj : (tuple) Value of an angular scale to be converted into a map width. Must be a tuple (val, "unit").
width : (tuple) Width of half of the map in physical scales. If width is provided, it overrides bmaj and d parameters. Must be a tuple (val, "unit").
cmap : (string) Colorscheme to be used in the colorbar.
resolution : (int) Number of pixels per side in the resulting map.
fontsize: (float) Size of the plot labels and title.
title : (string) Optional title for the plot.
vmin : (float) Minimun value for the plot colorbar.
vmin : (float) Maximum value for the plot colorbar.
savefig : (string) Optional filename to save the plot figure.
show : (boolean) Whether to show the plot or not.
mask : (boolean) Whether to apply a circular mask or not.
"""
import yt
# Suppress INFO messages setting the log level to "ERROR"
yt.funcs.mylog.setLevel(50)
fname = sys._getframe().f_code.co_name
start_time = time.time()
def circle_mask(h, w, center=None, radius=None):
"""
applies a circular mask to the input data
aiming to mimic the telescope beam area.
"""
if center is None:
center = [int(w / 2), int(h / 2)]
if radius is None:
radius = min(center[0], center[1], w - center[0], h - center[1])
Y, X = np.ogrid[:h, :w]
r = np.sqrt((X - center[0])**2 + (Y - center[1])**2)
mask = r > radius
return mask
# Create a new field containing the oH2D+ number density
def n_oh2dp(field, data):
import yt.units as u
n = data['dens'] * data['h2do'] / (m_H2D * u.g)
return n
def n_tot(field, data):
import yt.units as u
return data['dens'] / (2.4 * m_H * u.g)
# Validate bmaj parameter type
if not isinstance(bmaj, (tuple)):
raise ValueError(
'Not a valid type for bmaj. Must be a tuple (val, "unit").')
else:
if str(bmaj[1]).lower() == 'arcsec':
bmaj = bmaj[0] * u.arcsec.to(u.rad)
elif str(bmaj[1]).lower() == 'deg':
bmaj = bmaj[0] * u.deg.to(u.rad)
elif str(bmaj[1]).lower() == 'rad':
bmaj = bmaj[0]
# Validate d parameter type
if not isinstance(d, (tuple)):
raise ValueError(
'Not a valid type for d. Must be a tuple (val, "unit").')
else:
if str(d[1]).lower() == 'kpc':
d = d[0] * u.kpc.to(u.cm)
elif str(d[1]).lower() == 'pc':
d = d[0] * u.pc.to(u.cm)
elif str(d[1]).lower() == 'au':
d = d[0] * u.au.to(u.cm)
elif str(d[1]).lower() == 'm':
d = d[0] * u.m.to(u.cm)
elif str(d[1]).lower() == 'cm':
d = d[0]
# Validate width parameter type
if width == None:
# In case width is not provided, calculate it based on bmaj and d
width = (bmaj * d, 'cm')
elif isinstance(width, (tuple)) and \
type(width[0]) in (float, int) and \
type(width[1]) == str:
width = width
else:
raise ValueError(
'Not a valid type for width. Must be a tuple (val, "unit").')
# Load the data from the snapshot file
ds = yt.load(runfile)
d = ds.all_data()
d_trimmed = d.cut_region(
["(obj['density'] > 4.008e-20) & (obj['density'] < 4.008e-18)"])
yt.add_field(('gas', 'n_oh2dp'),
function=n_oh2dp,
units='cm**-3',
force_override=True)
yt.add_field(('gas', 'n_tot'),
function=n_tot,
units='cm**-3',
force_override=True)
# Projection of data
axis = 'z'
field = 'n_oh2dp'
weight = None
data_trimmed = None
#data_trimmed = d_trimmed
# info about method and weight_field parameters
# at https://yt-project.org/doc/visualizing/plots.html
proj = ds.proj(field,
axis,
weight_field=weight,
data_source=data_trimmed,
method='integrate')
if center in ['max', 'm']:
# center on the density peak
center = ds.find_max(field)[1]
elif center in ['center', 'c']:
# center on the box center
center = [0, 0, 0]
# Resolution of the numpy grid
resolution = np.array([resolution, resolution])
# Convert grid data into a 2D grid for easy plotting
column = proj.to_frb(width, resolution, center)
# Compute the column density
cdensity = np.array(column[field])
# Evolutionary time of the snapshot
time = np.round(ds.current_time.in_units('kyr'), 1)
cdensity = np.flipud(cdensity)
if mask:
# create a circular mask to mimic the telescope beam area
Mask = circle_mask(cdensity.shape[0], cdensity.shape[1])
cdensity[Mask] = 0
# Print the column density averaged over the map area
N_mod = np.nanmean(cdensity)
print_(f"N_mod [max, mean]:\t{cdensity.max()}\t{N_mod}",
fname,
verbose=verbose)
if show:
import matplotlib.pyplot as plt
# Plotting commands
plt.imshow(np.log10(cdensity), cmap=cmap, vmin=vmin, vmax=vmax)
plt.colorbar().set_label(r'$N\,$[o-H$_2$D$^+$] / cm$^{-2}$',
size=fontsize)
#plt.colorbar().set_label(r'$N_{\rm total}$ / cm$^{-2}$', size=fontsize)
nticks = 3
width = (np.round(width[0] * u.cm.to(u.au), 1), 'AU')
plt.xticks(np.linspace(0, resolution.max(), nticks),
np.linspace(-width[0] / 2., width[0] / 2., nticks))
plt.yticks(np.linspace(0, resolution.max(), nticks),
np.linspace(-width[0] / 2., width[0] / 2., nticks))
plt.xlabel(r'x [' + width[1].upper() + ']', size=fontsize)
plt.ylabel(r'y [' + width[1].upper() + ']', size=fontsize)
if title in ['', None]:
plt.title(f't = {time} kyr', size=fontsize)
else:
plt.title(title, size=fontsize)
if savefig != '':
plt.savefig(savefig)
plt.show(block=True)
# Print the time taken by the function
elapsed_time(time.time() - start_time, fname, verbose)
return cdensity
nC = data["dens"] * (data["icp "] + data["ico "]) / (1.4 * mC)
xCp = data["dens"] * data["icp "] / (1.4 * mC) / nC
return xCp
# electron fraction
def xe(field, data):
nH = data["dens"] * (data["ihp "] + data["iha "] + data["ih2 "]) / (1.4 *
mH)
nC = data["dens"] * (data["icp "] + data["ico "]) / (1.4 * mC)
ne = data["dens"] * (data["ihp "] / (1.4 * mH) + data["icp "] / (1.4 * mC))
xe = ne / (nH + nC)
return xe
yt.add_field('nH', function=numdensH, units="1/cm**3", force_override=True)
yt.add_field('nH2', function=numdensH2, units="1/cm**3", force_override=True)
yt.add_field('nC', function=numdensC, units="1/cm**3", force_override=True)
yt.add_field('ne', function=numdense, units="1/cm**3", force_override=True)
yt.add_field('xHp', function=xHp, units="dimensionless", force_override=True)
yt.add_field('xH2', function=xH2, units="dimensionless", force_override=True)
yt.add_field('xCp', function=xCp, units="dimensionless", force_override=True)
yt.add_field('xe', function=xe, units="dimensionless", force_override=True)
#yt.add_field('G', function=GG, units="dimensionless", force_override=True)
# Input variables.
# DustBox
data_dir = "/home/jcibanezm/codes/run/Silcc/CF_Prabesh"
#Daikaiju
#data_dir = "/data/gamera/jcibanezm/DustAnalysis/CF_Data"
def prepare_octree(ds, ile, start_level=0, debug=True, dd=None, center=None):
if dd is None:
#we keep passing dd around to not regenerate the data all the time
dd = ds.all_data()
try:
dd['MetalMass']
except KeyError:
add_fields() #add the metal mass field that sunrise wants
def _temp_times_mass(field, data):
return data["Temperature"] * data["CellMassMsun"]
add_field("TemperatureTimesCellMassMsun", function=_temp_times_mass)
fields = [
"CellMassMsun", "TemperatureTimesCellMassMsun", "MetalMass",
"CellVolumeCode"
]
#gather the field data from octs
pbar = get_pbar("Retrieving field data", len(fields))
field_data = []
for fi, f in enumerate(fields):
field_data += dd[f],
pbar.update(fi)
pbar.finish()
del field_data
#first we cast every cell as an oct
#ngrids = np.max([g.id for g in ds._grids])
grids = {}
levels_all = {}
levels_finest = {}
for l in range(100):
levels_finest[l] = 0
levels_all[l] = 0
pbar = get_pbar("Initializing octs ", len(ds.index.grids))
for gi, g in enumerate(ds.index.grids):
ff = np.array([g[f] for f in fields])
og = amr_utils.OctreeGrid(g.child_index_mask.astype('int32'),
ff.astype("float64"),
g.LeftEdge.astype("float64"),
g.ActiveDimensions.astype("int32"),
np.ones(1, dtype="float64") * g.dds[0],
g.Level, g.id)
grids[g.id] = og
#how many refinement cells will we have?
#measure the 'volume' of each mesh, but many
#cells do not exist. an overestimate
levels_all[g.Level] += g.ActiveDimensions.prod()
#how many leaves do we have?
#this overestimates. a child of -1 means no child,
#but that cell may still be expanded on a submesh because
#(at least in ART) the meshes are inefficient.
g.clear_data()
pbar.update(gi)
pbar.finish()
#create the octree grid list
#oct_list = amr_utils.OctreeGridList(grids)
#initialize arrays to be passed to the recursion algo
o_length = np.sum(levels_all.values())
r_length = np.sum(levels_all.values())
output = np.zeros((o_length, len(fields)), dtype='float64')
refined = np.zeros(r_length, dtype='int32')
levels = np.zeros(r_length, dtype='int32')
ids = np.zeros(r_length, dtype='int32')
pos = position()
hs = hilbert_state()
start_time = time.time()
if debug:
printing = lambda x: print_oct(x)
else:
printing = None
pbar = get_pbar("Building Hilbert DFO octree", len(refined))
RecurseOctreeDepthFirstHilbert(
ile,
pos,
grids[0], #we always start on the root grid
hs,
output,
refined,
levels,
grids,
start_level,
ids,
debug=printing,
tracker=pbar)
pbar.finish()
#by time we get it here the 'current' position is actually
#for the next spot, so we're off by 1
print('took %1.2e seconds' % (time.time() - start_time))
print('refinement tree # of cells %i, # of leaves %i' %
(pos.refined_pos, pos.output_pos))
print('first few entries :', refined[:12])
output = output[:pos.output_pos]
refined = refined[:pos.refined_pos]
levels = levels[:pos.refined_pos]
return output, refined, dd, pos.refined_pos
parser.add_argument("-hdf5",action="store",help='hdf5 data dump to load')
parser.add_argument("-cs",action="store",help='adiabatic sound speed in cm/sec')
args = parser.parse_args()
if args.hdf5 and args.cs:
hdf5= str(args.hdf5)
cs= float(args.cs)
else: raise ValueError
import matplotlib
matplotlib.use('Agg')
import yt
import numpy as np
def _velocity_x(field, data):
return data['X-momentum']/data['density']
yt.add_field(("gas","velocity_x"), function=_velocity_x) #, units="cm/s")
def _velocity_y(field, data):
return data['Y-momentum']/data['density']
yt.add_field(("gas","velocity_y"), function=_velocity_y) #, units="cm/s")
def _velocity_z(field, data):
return data['Z-momentum']/data['density']
yt.add_field(("gas","velocity_z"), function=_velocity_z) #, units="cm/s")
ds= load(hdf5)
#yt method
box=ds.region(ds.domain_center,ds.domain_left_edge,ds.domain_right_edge)
rho= np.array(box[('chombo', 'density')])
vx= np.array(box[("gas","velocity_x")])
vy= np.array(box[("gas","velocity_y")])
y = data['y'] - center[1]
z = data['z'] - center[2]
top = accx * x + accy * y + accz * z
bottom = (x**2 + y**2 + z**2)**0.5
#
return top/bottom
# return np.abs(top/bottom)
def _blob_center(field, data):
return np.array([8.00E21,3.0857E21,3.0857E21]) * yt.units.cm
#yt.add_field('v_r', function=_v_r, units="cm/s", display_name = r'v$_{\rm{r}}$', take_log=False)
#yt.add_field('g_r', function=_g_r, units="cm/s**2", display_name = r'g$_{\rm{r}}$')
yt.add_field('blob_center', function=_blob_center, units="cm")
mydf.add_field('gas_column_density')
mydf.add_field('gas_number_density')
#name = 'dwarf_fullp_hdf5_chk_'
#name = 'blobtest_hdf5_chk_'
#fields_list = ['velx','g_r']
#fields_list = ['mach_number','sound_speed']
#fields_list = ['velx']
#fields_list = ['density','temperature','pressure','density','velx','vely','velx']#,'v_r','velx']
#fields_list = ['density','temperature','pressure','density']
#fields_list = ['gas_number_density']
fields_list = ['density','temperature']#,'pressure','velx','vely','velz']
weight_field = 'density'
#blob_center = np.array([8.00E21,3.0857E21,3.0857E21])
#blob_center = np.array([3.0857E21,3.0857E21,3.0857E21])
exec(open('/groups/dark/lawsmith/jYT/my_settings.py').read())
clusterdir = '/storage/dark/lawsmith/FLASH4.3/runs/'
savepath = '/groups/dark/lawsmith/fend_slices/'
elif args.cluster == 'lux':
exec(open('/home/lawsmith/jYT/my_settings.py').read())
clusterdir = '/data/groups/ramirez-ruiz/lawsmith/FLASH4.3/'
savepath = '/data/groups/ramirez-ruiz/lawsmith/FLASH4.3/lux_slices/'
def _ekin_plus_gpot(field, data):
return 0.5 * (data['velocity_magnitude'].d)**2 + data['gpot'].d
yt.add_field(("gas", "ekin_plus_gpot"),
display_name='v^2/2 + gpot',
function=_ekin_plus_gpot,
take_log=False,
units=None,
force_override=True)
if args.chkplt == 'chk':
LOAD_FILES = clusterdir + args.run + '/multitidal_hdf5_chk_' + args.files
elif args.chkplt == 'plt':
LOAD_FILES = clusterdir + args.run + '/multitidal_hdf5_plt_cnt_' + args.files
yt.enable_parallelism()
ts = yt.DatasetSeries(LOAD_FILES)
t_array = []
e_array = []
for ds in ts.piter():
# 3. Neutral fraction and temperature slices at the last output
first = 1
last = 15
Myr = 3.1557e13
kpc = 3.086e21
########################################################################
def _MyNeutralFrac(field, data):
return data['H_p0_fraction'] / 0.75908798
yt.add_field("Neutral_Fraction",
function=_MyNeutralFrac,
take_log=True,
units='')
########################################################################
center = [0.75757575, 0.5, 0.5]
radius = 0.121212
time = []
xe = []
temp = []
for i in range(first, last + 1):
amrfile = "DD%4.4d/data%4.4d" % (i, i)
pf = yt.load(amrfile)
sphere = pf.h.sphere(center, radius)
from yt.utilities.physical_constants import planck_constant_cgs
import halo_analysis as ha
# Add J21_LW and J_Lyman fields to yt
sigma_H2I = yt.YTQuantity(3.71e-18, 'cm**2')
sigma_HI = yt.YTQuantity(6.3e-18, 'cm**2')
J21_norm = yt.YTQuantity(1e21, 'cm**2/erg')
E_LW = yt.YTQuantity(12.4, 'eV')
E_HI = yt.YTQuantity(13.6, 'eV')
nu_H = yt.YTQuantity(3.2881e15, 'Hz')
def _J21_LW(field, data):
return J21_norm*data['enzo', 'H2I_kdiss']*E_LW/(4.0*np.pi*np.pi*sigma_H2I \
*nu_H)
yt.add_field(name=('gas', 'J21_LW'), function=_J21_LW, units='dimensionless', \
display_name='J$_{21}$')
def _J_Lyman(field, data):
return data['enzo', 'HI_kph']*planck_constant_cgs/(4.0*np.pi*np.pi*sigma_HI)
yt.add_field(name=('gas', 'J_Lyman'), function=_J_Lyman, units='erg/cm**2', \
display_name='J(Lyman)')
# Load ytree arbor to update.
a = ytree.load('wise_tree_0_0_0.dat')
# Load full timeseries of datasets.
ts = yt.load('wise_data/DD00??/output_00??')
# Add all analysis fields to the arbor.
a.add_analysis_field('cosmological_time', units='yr')
bv = data.get_field_parameter("bulk_velocity").in_units("cm/s")
else:
bv = data.ds.arr(np.zeros(3), "cm/s")
xv = data["gas","velocity_x"] - bv[0]
yv = data["gas","velocity_y"] - bv[1]
center = data.get_field_parameter('center')
x_hat = data["x"] - center[0]
y_hat = data["y"] - center[1]
r = np.sqrt(x_hat*x_hat+y_hat*y_hat)
x_hat /= r
y_hat /= r
return (yv*x_hat-xv*y_hat)
yt.add_field("vc", function=_vc,take_log=False, units=r"km/s",validators=[ValidateParameter('bulk_velocity')])
def _radial_velocity(field,data):
center = data.get_field_parameter('center')
if data.has_field_parameter("bulk_velocity"):
bv = data.get_field_parameter("bulk_velocity").in_units("cm/s")
else:
bv = data.ds.arr(np.zeros(3), "cm/s")
x = data["x"] - center[0]
y = data["y"] - center[1]
r = np.sqrt(x**2+y**2)
vx = data["gas","velocity_x"] - bv[0]
vy = data["gas","velocity_y"] - bv[1]
gmpt = G * edata[6]
#for chk in p[2]:
chk = p[2]
f = '/pfs/lawsmith/FLASH4.3/runs/' + p[
0] + '/multitidal_hdf5_chk_' + chk
pf = yt.load(f)
tindex = abs(time - pf.current_time.v).argmin()
myvars = var('sb_mass')
yt.add_field(("gas", "sb_mass"),
function=myvars.mesh,
take_log=False,
force_override=True,
units="g")
# todo might want to add density cut later & see if affects results.
#ad = pf.all_data()
#dense_ad = ad.cut_region(['obj["dens"] > 1e-11'])
sp = pf.sphere("center", (1.0, "Mpc"))
dm = sp.quantities.total_quantity("sb_mass").v
b_array.append(p[1])
dm_array.append(dm)
ax.plot(b_array,
np.log10(np.array(dm_array) / M_sun),
def Brho(field, data):
return np.sqrt(data['magx']**2 + data['magy']**2 +
data['magz']**2) / data['dens']
def CRPres(field, data):
return 0.33333 * (data['cray'] * data['density'])
def GasPres(field, data):
return data['pressure'] - 0.33333 * (data['cray'] * data['density'] *
pUnit)
yt.add_field(("gas", "GasPres"), function=GasPres, units="g/cm/s**2")
yt.add_field(("gas", "CRPres"), function=CRPres, units="g/cm**3")
yt.add_field(("gas", "Brho"), function=Brho, units="G*cm**3/g")
massISM = []
for ds in ts:
time = ds.current_time.in_units("Myr")
times.append(time)
dd = ds.all_data()
print(ds.derived_field_list)
print(dd.quantities.extrema("dx"))
dsk1 = ds.disk("center", [0, 1, 0], (100.0, "kpc"), (1.0, "kpc"))
dsk2 = ds.disk("center", [0, 1, 0], (100.0, "kpc"), (2.0, "kpc"))
dsk3 = ds.disk("center", [0, 1, 0], (100.0, "kpc"), (3.0, "kpc"))
dsk4 = ds.disk("center", [0, 1, 0], (100.0, "kpc"), (4.0, "kpc"))
NBINS = 16
Myr = 3.1557e13
kpc = 3.086e21
########################################################################
def _MyRadius(field, data):
center = data.get_field_parameter("center")
dx = data["x"] - center[0]
dy = data["y"] - center[1]
dz = data["z"] - center[2]
return np.sqrt(dx * dx + dy * dy + dz * dz)
yt.add_field("Radius", function=_MyRadius, take_log=False, units='')
def _MyNeutralFrac(field, data):
return data['H_p0_fraction'] / 0.75908798
yt.add_field("Neutral_Fraction",
function=_MyNeutralFrac,
take_log=False,
units='')
########################################################################
center = [1e-3] * 3
time = []
radius = []
from yt.data_objects.particle_filters import add_particle_filter
from yt.fields.derived_field import \
ValidateGridType, \
ValidateParameter, \
ValidateSpatial, \
NeedsParameter
from astropy.table import Table , Column
def _Disk_H(field, data):
center = data.get_field_parameter('center')
z = data["z"] - center[2]
return np.abs(z)
yt.add_field("Disk_H",
function=_Disk_H,
units="pc",
take_log=False,
validators=[ValidateParameter('center')])
def _radial_velocity(field,data):
if data.has_field_parameter("bulk_velocity"):
bv = data.get_field_parameter("bulk_velocity").in_units("cm/s")
else:
bv = data.ds.arr(np.zeros(3), "cm/s")
xv = data["gas","velocity_x"] - bv[0]
yv = data["gas","velocity_y"] - bv[1]
center = data.get_field_parameter('center')
x_hat = data["x"] - center[0]
y_hat = data["y"] - center[1]
r = np.sqrt(x_hat*x_hat+y_hat*y_hat)
x_hat /= r
if __name__ == '__main__':
parser = ArgumentParser(description="test")
parser.add_argument("--hdf5",action="store",default="data/data.0006.3d.hdf5",help='hdf5 file to load',required=True)
parser.add_argument("--model",choices=['bInf','b1','b1e-2'],action="store",help='hdf5 file to load',required=True)
parser.add_argument("--init_sink_file",action="store",default="modules/hdf5_params/init.sink",help='init.sink file',required=True)
parser.add_argument("--sink_id",type=int,action="store",help='[0-63] are possible',required=True)
parser.add_argument("--cores",type=int,default=1,action="store",help='if > 1 do multiprocessing',required=False)
args = parser.parse_args()
# Outdir
outdir= os.path.join(args.model, os.path.basename(args.hdf5).replace('.','_'))
if not os.path.exists(outdir): os.makedirs(outdir)
# Autotmatically add fields to all subsequent data loads
yt.add_field("my_vx",function=_my_vx,units='cm/s')
yt.add_field("my_vy",function=_my_vy,units='cm/s')
yt.add_field("my_vz",function=_my_vz,units='cm/s')
yt.add_field("my_radius",
function=_my_radius,
units="cm",
take_log=False,
validators=[ValidateParameter('center')])
yt.add_field("my_radial_velocity",
function=_my_radial_velocity,
units="cm/s",
take_log=False,
validators=[ValidateParameter('center'),
ValidateParameter('bulk_velocity')])
if args.model != 'bInf':
# # FIX ME, units?
# Generate specific volume field
def _spVol(field, data):
dens = data['gas', 'density']
return yt.numpy.power(dens, -1)
fname = "cylindrical_rmi_2d_hdf5_chk_" # Establish base name and path for the simulation output files
fieldName = ['density', 'specific_volume']
path = os.getcwd()
path = path + '/'
yt.add_field(('gas', 'specific_volume'),
function=_spVol,
units="cm**3/g",
sampling_type="cell",
take_log=False)
r_max = (3.5, "cm")
p_res = 512
r_tar = 2.5
#mproc.cpu_count() <= Nproc
# Initialize MPI Communicator
comm = MPI.COMM_WORLD
Nproc = int(comm.Get_size())
Pid = int(comm.Get_rank())
initfile = 0
finalfile = 100
nfiles = finalfile - initfile
r = np.arctan2(y, x)
return r
def _vertical_velocity(field, data):
if data.has_field_parameter("bulk_velocity"):
bv = data.get_field_parameter("bulk_velocity").in_units("cm/s")
else:
bv = data.ds.arr(np.zeros(3), "cm/s")
v = data["gas", "velocity_z"] - bv[2]
return v
yt.add_field("Disk_Radius",
function=_Disk_Radius,
units="cm",
take_log=False,
validators=[ValidateParameter('center')])
yt.add_field("Disk_H",
function=_Disk_H,
units="pc",
take_log=False,
validators=[ValidateParameter('center')])
yt.add_field("vc",
function=_vc,
take_log=False,
units=r"km/s",
validators=[ValidateParameter('bulk_velocity')])
yt.add_field("Disk_Angle",
function=_Disk_Angle,
units="dimensionless",
size = comm.Get_size()
#yt.enable_parallelism()
def cooling_function(field,data):
LAMBDA = np.zeros_like(kb*data['temperature'] * cm**3/s)
kbT = (kb*data['temperature']).in_units('keV').value
LAMBDA[np.where(kbT > 5.0)[0]] = 6.10540967286092 * (kbT[np.where(kbT > 5.0)[0]]/100.)**0.4
LAMBDA[np.where((kbT<= 5.0)&(kbT > 0.7))[0]] = 1.842055928644565 * (kbT[np.where((kbT<= 5.0)&(kbT > 0.7))[0]]/5.)**0.22
LAMBDA[np.where((kbT<= 0.7)&(kbT > 0.13))[0]] = 1.1952286093343938 * (kbT[np.where((kbT<= 0.7)&(kbT > 0.13))[0]]/0.7)**-0.5
LAMBDA[np.where((kbT<= 0.13)&(kbT > 0.02))[0]] = 2.773500981126145 * (kbT[np.where((kbT<= 0.13)&(kbT > 0.02))[0]]/0.13)**-1.7
LAMBDA[np.where((kbT<= 0.02)&(kbT > 0.0017235))[0]] = 67.2 * (kbT[np.where((kbT<= 0.02)&(kbT > 0.0017235))[0]]/0.02)**0.6
LAMBDA[np.where((kbT<= 0.0017235)&(kbT > kbTfloor))[0]] = 15.44 * (kbT[np.where((kbT<= 0.0017235)&(kbT > kbTfloor))[0]]/0.0017235)**6.0
return LAMBDA*1e-23
yt.add_field(("gas", "cooling_function"), function=cooling_function, units='erg*cm**3/s', display_name=r"$\Lambda$")
def tcool(field,data):
return 1.5*kb*data['temperature']/((data['density']/mp/0.62) * data['cooling_function'])
yt.add_field(("gas","tcool"),function=tcool,units="Myr", display_name=r"$\t_{cool}$")
def _my_radial_velocity(field, data):
if data.has_field_parameter("bulk_velocity"):
bv = data.get_field_parameter("bulk_velocity").in_units("cm/s")
else:
bv = data.ds.arr(np.zeros(3), "cm/s")
xv = data["gas","velocity_x"] - bv[0]
yv = data["gas","velocity_y"] - bv[1]
zv = data["gas","velocity_z"] - bv[2]
bl_HA = interpolate.LinearNDInterpolator(pts, sr_HA)
def _Emission_HAlpha(field, data):
H_N = np.log10(np.array(data['H_nuclei_density']))
Temperature = np.log10(np.array(data['Temperature']))
dia1 = bl_HA(H_N, Temperature)
idx = np.isnan(dia1)
dia1[idx] = -200.
emission_line = (10.**dia1) * ((10.**H_N)**2.0)
emission_line = emission_line / (4. * np.pi * 3.03e-12)
return emission_line
yt.add_field('Emission_HAlpha',
units=emission_units,
function=_Emission_HAlpha)
hden1, T1, table_SiIV_1 = make_Cloudy_table(13)
hden1, T1, table_SiIV_2 = make_Cloudy_table(14)
sr_SiIV_1 = table_SiIV_1.T.ravel()
sr_SiIV_2 = table_SiIV_2.T.ravel()
bl_SiIV_1 = interpolate.LinearNDInterpolator(pts, sr_SiIV_1)
bl_SiIV_2 = interpolate.LinearNDInterpolator(pts, sr_SiIV_2)
def _Emission_SiIV(field, data):
H_N = np.log10(np.array(data["H_NumberDensity"]))
Temperature = np.log10(np.array(data["Temperature"]))
dia1 = bl_SiIV_1(H_N, Temperature)
dia2 = bl_SiIV_2(H_N, Temperature)
def ge_dve(field, data):
try:
return data[('enzo', 'GasEnergy')].v
except:
return data[('enzo', 'TotalEnergy')].v - data['kinetic_energy'].v
def pe_dave(field, data):
p = data[('enzo', 'TotalEnergy')].v - 0.5 * (data[
('enzo', 'x-velocity')].v**2 + data[
('enzo', 'y-velocity')].v**2 + data[('enzo', 'z-velocity')].v**2)
p *= data[('enzo', 'Density')] / (data.ds['Gamma'] - 1)
return p
yt.add_field('pe_dave', pe_dave)
yt.add_field('ge', ge_dve)
def temp_vs_ge(field, data):
return data[('enzo', 'Temperature')].v / data['ge']
yt.add_field('tvg', temp_vs_ge)
default_frames = [266]
framelist = {'f03_ded_nodef': [125]}
temps = {}
dirnames = {}
plt.clf()
def _bern(field, data) :
PE = data[('Gas','Phi')]/cl
posCM, velCM = getCM(data.ds, IE=useIE)
v = np.linalg.norm( data[('Gas','Velocities')]/cv - velCM, axis=1 )
KE = 0.5*np.multiply(v,v)
if useIE:
enthalpy = data[('Gas','ie')]
else:
enthalpy = gamma / (gamma-1.0) * R * data[('Gas','Temperature')] / K
bern = PE + KE + enthalpy
if userho :
rho = data[('Gas','rho')]
bern = np.multiply( bern, rho )
return bern
yt.add_field(('Gas','bernoulli'), function = _bern, particle_type = True )
plt.clf()
fig = plt.figure(figsize=(10,8))
ad = ds.all_data()
time = getTime(ds)
pos = ad[('Gas','Coordinates')]/cl
x = pos[:,0]
y = pos[:,1]
z = pos[:,2]
bern = ad[('Gas','bernoulli')]
boolArray = np.absolute(z) < dPeriod * 0.5 * bernslice
x = x[boolArray]
def get_use_gas():
"""
returns whether to use gas when calculate center velocity and position
"""
global use_gas
return use_gas
def _Neg_z(field, data):
"""
returns the negative of the z-positions
"""
return -1 * data['z']
yt.add_field("Neg_z", function=_Neg_z, units=r"cm")
def _Neg_dz(field, data):
"""
returns the negative of the dz
"""
return -1 * data['dz']
yt.add_field("Neg_dz", function=_Neg_z, units=r"cm")
def _Center_Position(field, data):
"""
Returns the center position for the current set center.
import yt
from yt.utilities.physical_constants import kboltz, mh
import numpy as np
def _hspec(field, data):
return data['h ']
yt.add_field("hspec", function=_hspec)
def _hpspec(field, data):
return data['hp ']
yt.add_field("hpspec", function=_hpspec)
def _hespec(field, data):
return data['he ']
yt.add_field("hespec", function=_hespec)
def _hepspec(field, data):
return data['hep ']
yt.add_field("hepspec", function=_hepspec)
data_shape = np.shape(M)
M = M.flatten()
Z = Z.flatten()
L = np.zeros(np.size(M))
for i in range(np.size(M)):
L[i] = isp.ZAMSData.interpolate(M[i], Z[i], return_only = 'luminosity')
L = L.reshape(data_shape)
L = L * isp.const.Lsun
L = L * yt.units.erg / yt.units.s
return L
yt.add_field('star_luminosity', function=_stellar_luminosity, units='erg*s**(-1)')
def _stellar_radius(field, data):
# need to know star mass
M = data['particle_mass'].convert_to_units('Msun')
M = M.value
Z = data[('io','metallicity_fraction')]
Z = Z.value
data_shape = np.shape(M)
# reshape here
M = M.flatten()
Z = Z.flatten()
#
