def _add_mesh_sampling_particle_field(self, deposit_field, ftype, ptype):
units = self.ds.field_info[ftype, deposit_field].units
take_log = self.ds.field_info[ftype, deposit_field].take_log
field_name = f"cell_{ftype}_{deposit_field}"
def _mesh_sampling_particle_field(field, data):
pos = data[ptype, "particle_position"]
field_values = data[ftype, deposit_field]
if isinstance(data, FieldDetector):
return np.zeros(pos.shape[0])
i, j, k = np.floor(
(pos - data.LeftEdge) / data.dds).astype("int64").T
# Make sure all particles are within the current grid, otherwise return nan
maxi, maxj, maxk = field_values.shape
mask = (i < maxi) & (j < maxj) & (k < maxk)
mask &= (i >= 0) & (j >= 0) & (k >= 0)
result = np.full(len(pos), np.nan, dtype="float64")
if result.shape[0] > 0:
result[mask] = field_values[i[mask], j[mask], k[mask]]
return data.ds.arr(result, field_values.units)
self.ds.add_field(
(ptype, field_name),
function=_mesh_sampling_particle_field,
sampling_type="particle",
units=units,
take_log=take_log,
validators=[ValidateSpatial()],
)
def add_particle_average(registry,
ptype,
field_name,
weight="particle_mass",
density=True):
field_units = registry[ptype, field_name].units
def _pfunc_avg(field, data):
pos = data[ptype, "particle_position"]
f = data[ptype, field_name]
wf = data[ptype, weight]
f *= wf
v = data.deposit(pos, [f], method="sum")
w = data.deposit(pos, [wf], method="sum")
v /= w
if density: v /= data["index", "cell_volume"]
v[np.isnan(v)] = 0.0
return v
fn = ("deposit", "%s_avg_%s" % (ptype, field_name))
registry.add_field(fn,
function=_pfunc_avg,
validators=[ValidateSpatial(0)],
particle_type=False,
units=field_units)
return fn
def setup_counts_fields(ds, ebounds, ftype="gas"):
r"""
Create deposited image fields from X-ray count data in energy bands.
Parameters
----------
ds : Dataset
The FITS events file dataset to add the counts fields to.
ebounds : list of tuples
A list of tuples, one for each field, with (emin, emax) as the
energy bounds for the image.
ftype : string, optional
The field type of the resulting field. Defaults to "gas".
Examples
--------
>>> ds = yt.load("evt.fits")
>>> ebounds = [(0.1,2.0),(2.0,3.0)]
>>> setup_counts_fields(ds, ebounds)
"""
for (emin, emax) in ebounds:
cfunc = _make_counts(emin, emax)
fname = "counts_%s-%s" % (emin, emax)
mylog.info("Creating counts field %s." % fname)
ds.add_field((ftype, fname),
function=cfunc,
units="counts/pixel",
validators=[ValidateSpatial()],
display_name="Counts (%s-%s keV)" % (emin, emax))
def add_particle_average(registry,
ptype,
field_name,
weight=None,
density=True):
if weight is None:
weight = (ptype, "particle_mass")
field_units = registry[ptype, field_name].units
def _pfunc_avg(field, data):
pos = data[ptype, "particle_position"]
f = data[ptype, field_name]
wf = data[ptype, weight]
f *= wf
v = data.deposit(pos, [f], method="sum")
w = data.deposit(pos, [wf], method="sum")
v /= w
if density:
v /= data["index", "cell_volume"]
v[np.isnan(v)] = 0.0
return v
fn = ("deposit", f"{ptype}_avg_{field_name}")
registry.add_field(
fn,
sampling_type="cell",
function=_pfunc_avg,
validators=[ValidateSpatial(0)],
units=field_units,
)
return fn
def add_nearest_neighbor_value_field(ptype, coord_name, sampled_field, registry):
"""
This adds a nearest-neighbor field, where values on the mesh are assigned
based on the nearest particle value found. This is useful, for instance,
with voronoi-tesselations.
"""
field_name = ("deposit", f"{ptype}_nearest_{sampled_field}")
field_units = registry[ptype, sampled_field].units
unit_system = registry.ds.unit_system
def _nearest_value(field, data):
pos = data[ptype, coord_name]
pos = pos.convert_to_units("code_length")
value = data[ptype, sampled_field].in_base(unit_system.name)
rv = data.smooth(
pos, [value], method="nearest", create_octree=True, nneighbors=1
)
rv = data.apply_units(rv, field_units)
return rv
registry.add_field(
field_name,
sampling_type="cell",
function=_nearest_value,
validators=[ValidateSpatial(0)],
units=field_units,
)
return [field_name]
def add_volume_weighted_smoothed_field(ptype, coord_name, mass_name,
smoothing_length_name, density_name, smoothed_field, registry,
nneighbors = None):
field_name = ("deposit", "%s_smoothed_%s" % (ptype, smoothed_field))
field_units = registry[ptype, smoothed_field].units
def _vol_weight(field, data):
pos = data[ptype, coord_name].in_units("code_length")
mass = data[ptype, mass_name].in_cgs()
dens = data[ptype, density_name].in_cgs()
quan = data[ptype, smoothed_field].in_units(field_units)
if smoothing_length_name is None:
hsml = np.zeros(quan.shape, dtype='float64') - 1
hsml = data.apply_units(hsml, "code_length")
else:
hsml = data[ptype, smoothing_length_name].in_units("code_length")
kwargs = {}
if nneighbors:
kwargs['nneighbors'] = nneighbors
rv = data.smooth(pos, [mass, hsml, dens, quan],
method="volume_weighted",
create_octree = True)[0]
rv[np.isnan(rv)] = 0.0
# Now some quick unit conversions.
rv = data.apply_units(rv, field_units)
return rv
registry.add_field(field_name, function = _vol_weight,
validators = [ValidateSpatial(0)],
units = field_units)
return [field_name]
def _add_mesh_sampling_particle_field(self, deposit_field, ftype, ptype):
units = self.ds.field_info[ftype, deposit_field].units
take_log = self.ds.field_info[ftype, deposit_field].take_log
field_name = f"cell_{ftype}_{deposit_field}"
def _mesh_sampling_particle_field(field, data):
pos = data[ptype, "particle_position"]
field_values = data[ftype, deposit_field]
if isinstance(data, FieldDetector):
return np.zeros(pos.shape[0])
i, j, k = np.floor(
(pos - data.LeftEdge) / data.dds).astype("int64").T
return field_values[i, j, k]
self.ds.add_field(
(ptype, field_name),
function=_mesh_sampling_particle_field,
sampling_type="particle",
units=units,
take_log=take_log,
validators=[ValidateSpatial()],
)
def add_density_kernel(ptype, coord_name, mass_name, registry, nneighbors=64):
field_name = (ptype, "smoothed_density")
field_units = registry[ptype, mass_name].units
def _nth_neighbor(field, data):
pos = data[ptype, coord_name]
pos.convert_to_units("code_length")
mass = data[ptype, mass_name]
mass.convert_to_units("g")
densities = mass * 0.0
data.particle_operation(pos, [mass, densities],
method="density",
nneighbors=nneighbors)
ones = pos.prod(axis=1) # Get us in code_length**3
ones[:] = 1.0
densities /= ones
# Now some quick unit conversions.
return densities
registry.add_field(field_name,
function=_nth_neighbor,
validators=[ValidateSpatial(0)],
particle_type=True,
units="g/cm**3")
return [field_name]
def add_deposited_particle_field(self, deposit_field, method):
"""Add a new deposited particle field
Creates a new deposited field based on the particle *deposit_field*.
Parameters
----------
deposit_field : tuple
The field name tuple of the particle field the deposited field will
be created from. This must be a field name tuple so yt can
appropriately infer the correct particle type.
method : one of 'count', 'sum', or 'cic'
The particle deposition method to use.
Returns
-------
The field name tuple for the newly created field.
"""
self.index
if isinstance(deposit_field, tuple):
ptype, deposit_field = deposit_field[0], deposit_field[1]
else:
raise RuntimeError
units = self.field_info[ptype, deposit_field].units
def _deposit_field(field, data):
"""
Create a grid field for particle wuantities weighted by particle
mass, using cloud-in-cell deposition.
"""
pos = data[ptype, "particle_position"]
# get back into density
if method != 'count':
pden = data[ptype, "particle_mass"]
top = data.deposit(pos, [data[(ptype, deposit_field)] * pden],
method=method)
bottom = data.deposit(pos, [pden], method=method)
top[bottom == 0] = 0.0
bnz = bottom.nonzero()
top[bnz] /= bottom[bnz]
d = data.ds.arr(top, input_units=units)
else:
d = data.ds.arr(
data.deposit(pos, [data[ptype, deposit_field]],
method=method))
return d
name_map = {"cic": "cic", "sum": "nn", "count": "count"}
field_name = "%s_" + name_map[method] + "_%s"
field_name = field_name % (ptype, deposit_field.replace(
'particle_', ''))
self.add_field(("deposit", field_name),
function=_deposit_field,
units=units,
take_log=False,
validators=[ValidateSpatial()])
return ("deposit", field_name)
def add_volume_weighted_smoothed_field(ptype,
coord_name,
mass_name,
smoothing_length_name,
density_name,
smoothed_field,
registry,
nneighbors=64,
kernel_name='cubic'):
unit_system = registry.ds.unit_system
if kernel_name == 'cubic':
field_name = ("deposit", "%s_smoothed_%s" % (ptype, smoothed_field))
else:
field_name = ("deposit", "%s_%s_smoothed_%s" %
(ptype, kernel_name, smoothed_field))
field_units = registry[ptype, smoothed_field].units
def _vol_weight(field, data):
pos = data[ptype, coord_name]
pos = pos.convert_to_units("code_length")
mass = data[ptype, mass_name].in_base(unit_system.name)
dens = data[ptype, density_name].in_base(unit_system.name)
quan = data[ptype, smoothed_field]
if hasattr(quan, "units"):
quan = quan.convert_to_units(field_units)
if smoothing_length_name is None:
hsml = np.zeros(quan.shape, dtype='float64') - 1
hsml = data.apply_units(hsml, "code_length")
else:
hsml = data[ptype, smoothing_length_name]
hsml.convert_to_units("code_length")
# This is for applying cutoffs, similar to in the SPLASH paper.
smooth_cutoff = data["index", "cell_volume"]**(1. / 3)
smooth_cutoff.convert_to_units("code_length")
# volume_weighted smooth operations return lists of length 1.
rv = data.smooth(pos, [mass, hsml, dens, quan],
index_fields=[smooth_cutoff],
method="volume_weighted",
create_octree=True,
nneighbors=nneighbors,
kernel_name=kernel_name)[0]
rv[np.isnan(rv)] = 0.0
# Now some quick unit conversions.
# This should be used when seeking a non-normalized value:
rv /= hsml.uq**3 / hsml.uq.in_base(unit_system.name).uq**3
rv = data.apply_units(rv, field_units)
return rv
registry.add_field(field_name,
sampling_type="cell",
function=_vol_weight,
validators=[ValidateSpatial(0)],
units=field_units)
registry.find_dependencies((field_name, ))
return [field_name]
def add_nearest_neighbor_field(ptype, coord_name, registry, nneighbors = 64):
field_name = (ptype, "nearest_neighbor_distance_%s" % (nneighbors))
def _nth_neighbor(field, data):
pos = data[ptype, coord_name]
pos.convert_to_units("code_length")
distances = 0.0 * pos[:,0]
data.particle_operation(pos, [distances],
method="nth_neighbor",
nneighbors = nneighbors)
# Now some quick unit conversions.
return distances
registry.add_field(field_name, sampling_type="particle", function = _nth_neighbor,
validators = [ValidateSpatial(0)],
units = "code_length")
return [field_name]
def add_contour_field(ds, contour_key):
def _contours(field, data):
fd = data.get_field_parameter("contour_slices_%s" % contour_key)
vals = data["index", "ones"] * -1
if fd is None or fd == 0.0:
return vals
for sl, v in fd.get(data.id, []):
vals[sl] = v
return vals
ds.add_field(("index", "contours_%s" % contour_key),
function=_contours,
validators=[ValidateSpatial(0)],
take_log=False,
display_field=False)
ValidateSpatial, \
NeedsParameter
import h5py
from yt.funcs import \
just_one
from common_functions import *
sl_left = slice(None, -2, None)
sl_right = slice(2, None, None)
div_fac = 2.0
sl_center = slice(1, -1, None)
ftype='gas'
vort_validators = [ValidateSpatial(1,
[(ftype, "velocity_x"),
(ftype, "velocity_y"),
(ftype, "velocity_z")])]
def _Disk_H(field, data):
center = data.get_field_parameter('center')
z = data["z"] - center[2]
return np.abs(z)
def _vturb(field, data):
fx = data[ftype, "velocity_x"][sl_right,sl_center,sl_center]/6.0
fx += data[ftype, "velocity_x"][sl_left,sl_center,sl_center]/6.0
fx += data[ftype, "velocity_x"][sl_center,sl_right,sl_center]/6.0
fx += data[ftype, "velocity_x"][sl_center,sl_left,sl_center]/6.0
fx += data[ftype, "velocity_x"][sl_center,sl_center,sl_right]/6.0
def func(field, data):
return -data['spitzer_conduction_coefficient']*yt.physical_constants.kb*data['H_nuclei_density']*data['temperature_gradient_%s' % ax]
return func
for ax in 'xyz':
f = _spitzer_heat_flux(ax)
yt.add_field('%s_%s' % (basename, ax), function=f, sampling_type='cell', units='erg/s/cm**2')
sl_left = slice(None, -2, None)
sl_right = slice(2, None, None)
div_fac = 2
@yt.derived_field(('gas', "%s_divergence" % basename), sampling_type="cell", units='erg/s/cm**3', validators=[ValidateSpatial(2)])
def _divergence(field, data):
ds = div_fac * just_one(data["index", "dx"])
f = data[xn][sl_right,1:-1,1:-1]/ds
f -= data[xn][sl_left ,1:-1,1:-1]/ds
ds = div_fac * just_one(data["index", "dy"])
f += data[yn][1:-1,sl_right,1:-1]/ds
f -= data[yn][1:-1,sl_left ,1:-1]/ds
ds = div_fac * just_one(data["index", "dz"])
f += data[zn][1:-1,1:-1,sl_right]/ds
f -= data[zn][1:-1,1:-1,sl_left ]/ds
new_field = data.ds.arr(np.zeros(data[xn].shape, dtype=np.float64),
f.units)
new_field[1:-1,1:-1,1:-1] = f
return new_field
def add_synchrotron_dtau_emissivity(ds,
ptype='lobe',
nu=(1.4, 'GHz'),
method='nearest_weighted',
proj_axis='x',
extend_cells=32):
me = yt.utilities.physical_constants.mass_electron #9.109E-28
c = yt.utilities.physical_constants.speed_of_light #2.998E10
e = yt.utilities.physical_constants.elementary_charge #4.803E-10 esu
gamma_min = yt.YTQuantity(10, 'dimensionless')
# Index for electron power law distribution
p = 2.0
pol_ratio = (p + 1.) / (p + 7. / 3.)
# Fitted constants for the approximated power-law + exponential spectra
# Integral of 2*F(x) -> tot_const*(nu**-2)*exp(-nu/nuc)
# 2*F(x) for the total intensity (parallel + perpendicular)
tot_const = 4.1648
stokes = StokesFieldName(ptype, nu, proj_axis)
nu = yt.YTQuantity(*nu)
if proj_axis == 'x':
los = [1., 0., 0.]
xvec = [0., 1., 0.]
yvec = [0., 0., 1.]
elif proj_axis == 'y':
los = [0., 1., 0.]
xvec = [0., 0., 1.]
yvec = [1., 0., 0.]
elif proj_axis == 'z':
los = [0., 0., 1.]
xvec = [1., 0., 0.]
yvec = [0., 1., 0.]
elif type(proj_axis) is list:
los = proj_axis
if los[0] != 0.: # not perpendicular to z-axis
xvec = [0., 1., 0.]
yvec = [0., 0., 1.]
# TODO: xvec and yvec for arbitrary proj_axis
else:
raise NotImplementedError
else:
raise NotImplementedError
los = np.array(los)
xvec = np.array(xvec)
yvec = np.array(yvec)
los = los / np.sqrt(np.sum(los * los))
# Determine the version of the simulation
if ('io', 'particle_tau1') in ds.field_list:
version = 2018
mylog.info(
'Field particle_tau1 detected. This is a version 2018 simulation.')
elif ('io', 'particle_type') in ds.field_list:
version = 2016
mylog.info(
'Field particle_type detected. This is a version 2016 simulation.')
else:
version = 2015
# Update the particle file handler in yt; raise exception if not successful
success = setup_part_file(ds)
mylog.info('Assuming this is a version 2015 simulation.')
if not success:
raise IOError
if ('gas', 'jet_volume_fraction') not in ds.derived_field_list:
ds.add_field(('gas', 'jet_volume_fraction'),
function=_jet_volume_fraction,
display_name="Jet Volume Fraction",
sampling_type='cell')
#def _gamc(field, data):
# # The new cutoff gamma
# # Note that particle_dens could be negative due to quadratic interpolation!
# gamc = (np.abs(data['particle_dens'] / data['particle_den0']))**(1./3.) \
# / (data['particle_dtau'] + np.finfo(np.float64).tiny)
# ind = np.where(gamc < 0.0)[0]
# if ind.shape[0] > 0:
# print(ind)
# print(gamc)
# return gamc
#pfname = 'particle_gamc_dtau'
#ds.add_field(pfname, function=_gamc, sampling_type='particle',
# units='', force_override=True)
#deposit_field = 'particle_gamc_dtau'
def _synchrotron_spec(field, data):
# To convert from FLASH "none" unit to cgs unit, times the B field from FLASH by sqrt(4*pi)
Bvec = np.array([data['particle_magx'],\
data['particle_magy'],\
data['particle_magz']])*np.sqrt(4.0*np.pi)
Bvec = data.apply_units(Bvec, 'gauss')
# Calculate sin(a), in which a is the pitch angle of the electrons relative to B field.
# See _nn_emissivity_i for more comments
cross = np.cross(los, Bvec, axisb=0)
Bsina = np.sqrt(np.sum(cross * cross, axis=-1))
Bsina = data.apply_units(Bsina, 'gauss')
#B = np.sqrt(np.sum(Bvec*Bvec, axis=0))
if version == 2018:
# Return for the FieldDetector; do nothing
if isinstance(data, FieldDetector):
return data['particle_dens']/data['particle_den0']**(1./3.)/ \
(data['particle_tau1']+data['particle_cmb1'])
gamc = (np.abs(data['particle_dens'] / data['particle_den0']))**(1./3.) \
/ (data['particle_tau1'] + data['particle_cmb1'] + np.finfo(np.float64).tiny)
else:
# Return for the FieldDetector; do nothing
if isinstance(data, FieldDetector):
return data['particle_dens']/data['particle_den0']**(1./3.)/ \
(data['particle_dtau'])
if np.any(data['particle_dtau'] < 0.0):
print('negative tau!')
print(data)
print(data['particle_tau'])
print(data['particle_dtau'])
# The new cutoff gamma
# Note that particle_dens could be negative due to quadratic interpolation!
gamc = (np.abs(data['particle_dens'] / data['particle_den0']))**(1./3.) \
/ (data['particle_dtau'] + np.finfo(np.float64).tiny)
ind = np.where(gamc <= 0.0)[0]
if ind.shape[0] > 0:
print(ind)
print(gamc[ind])
#gamc = data[(ptype, 'particle_gamc')]
# Cutoff frequency
nuc = 3.0 * gamc**2 * e * Bsina / (4.0 * np.pi * me * c)
#nu = data.get_field_parameter("frequency", default=yt.YTQuantity(1.4, 'GHz'))
# B**1.5 is taken from the grid data
norm = 3.0 / 8.0 * e**3.5 / (c**2.5 * me**1.5 * (np.pi)**0.5)
# P is taken from the grid data
N0 = 3.0 / me / c / c / (np.log(gamc / gamma_min)) / yt.YTQuantity(
4. * np.pi, 'sr')
# Fix where the cutoff gamma < 0
N0[ind] = 0.0
return np.clip(N0 * norm * nu**(-0.5) * np.exp(-nu / nuc),
np.finfo(np.float64).tiny, None)
# particle field name
pfname = 'particle_sync_spec_%s' % stokes.nu_str
ds.add_field(pfname,
function=_synchrotron_spec,
sampling_type='particle',
units='cm**(3/4)*s**(3/2)/g**(3/4)/sr',
force_override=True)
deposit_field = pfname
#try:
ds.add_particle_filter(ptype)
#except:
# raise NotImplementedError
###########################################################################
## Nearest Neighbor method
###########################################################################
#fname_nn = ds.add_deposited_particle_field(
# (ptype, 'particle_sync_spec_%s' % stokes.nu_str), 'nearest', extend_cells=extend_cells)
sync_unit = ds.field_info[deposit_field].units
if method == "nearest":
field_name = "%s_nn_%s"
elif method == "nearest_weighted":
field_name = "%s_nnw_%s"
else:
raise NotImplementedError
field_name = field_name % (ptype, deposit_field.replace('particle_', ''))
ad = ds.all_data()
#print(ad[ptype, "particle_position"])
def _nnw_deposit_field(field, data):
if isinstance(data, FieldDetector):
jetfluid = data['velocity_magnitude'] > lobe_v
d = data.deposit(data[ptype, "particle_position"],
(ptype, deposit_field),
method=method)
d[jetfluid] = 0.0
return d
#pos = ad[ptype, "particle_position"]
if ptype == 'lobe':
# Calling other fields must preceed the more intensive
# deposit function to prevent repeated iterations
jetfluid = data['velocity_magnitude'] > lobe_v
alldata = True
if alldata:
pos = ad[ptype, "particle_position"]
# Deposit using the distance weighted log field
#fields = [np.log(ad[ptype, deposit_field])]
fields = [ad[ptype, deposit_field]]
fields = [np.ascontiguousarray(f) for f in fields]
else:
left_edge = np.maximum(data.LeftEdge-data.dds*extend_cells,\
data.ds.domain_left_edge)
right_edge = np.minimum(data.RightEdge+data.dds*extend_cells,\
data.ds.domain_right_edge)
box = data.ds.box(left_edge, right_edge)
pos = box[ptype, "particle_position"]
fields = [box[ptype, deposit_field]]
d = data.deposit(pos, fields, method=method, extend_cells=extend_cells)
if np.all(d == 0.0):
d = data.ds.arr(d, input_units=sync_unit)
else:
# Conver the log back to real value
#d = data.ds.arr(np.exp(d), input_units=sync_unit)
d = data.ds.arr(d, input_units=sync_unit)
if ptype == 'lobe':
d[jetfluid] = 0.0
return d
fname_nn = ("deposit", field_name)
ds.add_field(fname_nn,
function=_nnw_deposit_field,
sampling_type="cell",
units=sync_unit,
take_log=True,
force_override=True,
validators=[ValidateSpatial()])
def _nn_emissivity_i(field, data):
'''
Emissivity using nearest neighbor. Integrate over line of sight to get intensity.
'''
Bvec = np.array([data[('gas', 'magnetic_field_x')],\
data[('gas', 'magnetic_field_y')],\
data[('gas', 'magnetic_field_z')]])
# Calculate sin(a), in which a is the pitch angle of the electrons relative to B field.
# We only see the radiation from electrons with pitch angles pointing to line of sight.
cross = np.cross(los, Bvec, axisb=0)
# B * sin(alpha) = (B * |(los x Bvec)|/|los|/|Bvec|)
# = |(los x Bvec)|
Bsina = np.sqrt(np.sum(cross * cross, axis=-1))
Bsina = data.apply_units(Bsina, 'gauss')
# P * (B*sina)^1.5
PBsina = data['gas', 'pressure'] * Bsina**1.5
frac = data['gas', 'jet_volume_fraction']
# Use gamc from deposit field
#############################
# fname_nn = 'ptype_nnw_gamc_dtau'
#if ('flash', fname_nn[1]) in data.ds.field_list:
# gamc = data[('flash', fname_nn[1])]
#else:
# gamc = data[fname_nn]
#bad_mask = gamc <= gamma_min
## Cutoff frequency
#nuc = 3.0*gamc**2*e*Bsina/(4.0*np.pi*me*c)
##nu = data.get_field_parameter("frequency", default=yt.YTQuantity(1.4, 'GHz'))
#norm = 3.0*Bsina**1.5/8.0*e**3.5/(c**2.5*me**1.5*(np.pi)**0.5)
#N0 = 3.0*data['gas', 'pressure']/me/c/c/(np.log(gamc/gamma_min))/yt.YTQuantity(4.*np.pi, 'sr')
#N0[bad_mask] = 0.0
return PBsina * frac * tot_const * data[fname_nn]
ds.add_field(stokes.I,
function=_nn_emissivity_i,
sampling_type='cell',
display_name=stokes.display_name('I'),
units='Jy/cm/arcsec**2',
take_log=True,
force_override=True,
validators=[ValidateSpatial()])
def _cos_theta(field, data):
Bvec = np.stack([data[('gas', 'magnetic_field_x')],\
data[('gas', 'magnetic_field_y')],\
data[('gas', 'magnetic_field_z')]], axis=-1)
Bproj = Bvec - np.expand_dims(np.inner(Bvec, los), -1) * los
# cos = cos(theta), theta is the angle between projected B and xvec
# Ignore invalid 0/0 warnings
with np.errstate(invalid='ignore'):
cos = np.inner(Bproj, xvec) / np.sqrt(
np.sum(Bproj * Bproj, axis=-1))
return cos
ds.add_field('cos',
function=_cos_theta,
sampling_type='cell',
display_name='cos theta',
units='',
take_log=False,
force_override=True)
def _nn_emissivity_q(field, data):
# pol_ratio = (perp - para) / (perp + para)
# The minus accounts for the perpendicular polarization
rtn = -data[stokes.I] * pol_ratio * (2 * data['cos']**2 - 1.0)
rtn[~np.isfinite(data['cos'])] = 0.0
return rtn
ds.add_field(stokes.Q,
function=_nn_emissivity_q,
sampling_type='cell',
display_name=stokes.display_name('Q'),
units='Jy/cm/arcsec**2',
take_log=False,
force_override=True)
def _nn_emissivity_u(field, data):
sin = np.sqrt(1.0 - data['cos']**2)
rtn = -data[stokes.I] * pol_ratio * 2 * sin * data['cos']
rtn[~np.isfinite(data['cos'])] = 0.0
return rtn
ds.add_field(stokes.U,
function=_nn_emissivity_u,
sampling_type='cell',
display_name=stokes.display_name('U'),
units='Jy/cm/arcsec**2',
take_log=False,
force_override=True)
return pfname, fname_nn, stokes.I, stokes.nu_str
def create_fields(ds):
units_override = {
"length_unit": (1, "Rsun"),
"time_unit": (6.955e+05, "s"),
"mass_unit": (3.36427433875e+17, "g"),
"magnetic_unit": (1.121e-02, "G")
}
def _radialvelocity(field, data):
return data['velocity_x']*data['x']/data['radius'] + \
data['velocity_y']*data['y']/data['radius'] + \
data['velocity_z']*data['z']/data['radius']
ds.add_field(('gas', "radialvelocity"),
function=_radialvelocity,
units="cm/s",
take_log=False)
def _sound_speed(field, data):
gamma = 1.05
ftype = field.name[0]
tr = gamma * data[ftype, "pressure"] / data[ftype, "density"]
return np.sqrt(tr)
ds.add_field(('gas', "sound_speed"),
function=_sound_speed,
units="cm/s",
take_log=False)
def _mach_number(field, data):
""" M{|v|/c_sound} """
ftype = field.name[0]
return data[ftype, "velocity_magnitude"] / data[ftype, "sound_speed"]
ds.add_field(('gas', "mach_number"),
function=_mach_number,
units="",
take_log=False)
def _temperature(field, data):
return (data["gas", "pressure"] * mp) / (2.0 * data["gas", "density"] *
kb)
ds.add_field(('gas', "temperature"),
function=_temperature,
units="K",
take_log=True)
def _radius_planet(field, data):
a = 0.047 * 1.496e+13 / 6.955e+10
shift = data.ds.arr(np.ones_like(data['x'])) * a
x_planet = data['x'] - shift
return np.sqrt(x_planet*x_planet \
+ data['y']*data['y'] \
+ data['z']*data['z'])
ds.add_field(('index', "radius_planet"),
function=_radius_planet,
units="cm",
take_log=False)
def _ni(field, data):
return data["density"] / (1.09 * mp)
ds.add_field(("gas", "ni"), function=_ni, units="cm**-3")
def _BGx1(field, data):
B0s = YTQuantity(2.0, "G")
B0p = YTQuantity(1.0, "G")
Rs = YTQuantity(6.955e+10, "cm")
Rp = YTQuantity(1.5 * 0.10045 * Rs, "cm")
a = YTQuantity(0.047, "au").in_units("cm")
center = data.get_field_parameter('center')
x1 = data["x"].in_units('cm')
x2 = data["y"].in_units('cm')
x3 = data["z"].in_units('cm')
rs = np.sqrt(x1 * x1 + x2 * x2 + x3 * x3)
rp = np.sqrt((x1 - a) * (x1 - a) + x2 * x2 + x3 * x3)
BGx1 = data.ds.arr(np.zeros_like(data["magnetic_field_x"]), "G")
BGx1 = 3.0 * x1 * x3 * B0s * Rs**3 * rs**(-5) + 3.0 * (
x1 - a) * x3 * B0p * Rp**3 * rp**(-5)
BGx1[rs <= Rs] = 3.0 * x1[rs <= Rs] * x3[rs <= Rs] * B0s * Rs**3 * rs[
rs <= Rs]**(-5)
BGx1[rs <= 0.5 * Rs] = 0.0
BGx1[rp <= Rp] = 3.0*(x1[rp <= Rp] - a)*x3[rp <= Rp]\
*B0p*Rp**3*rp[rp <= Rp]**(-5)
BGx1[rp <= 0.5 * Rp] = 0.0
return BGx1
ds.add_field(("gas", "BGx1"), function=_BGx1, units="G", take_log=False)
def _BGx2(field, data):
B0s = YTQuantity(2.0, "G")
B0p = YTQuantity(1.0, "G")
Rs = YTQuantity(6.955e+10, "cm")
Rp = YTQuantity(1.5 * 0.10045 * Rs, "cm")
a = YTQuantity(0.047, "au").in_units("cm")
center = data.get_field_parameter('center')
x1 = data["x"].in_units('cm')
x2 = data["y"].in_units('cm')
x3 = data["z"].in_units('cm')
rs = np.sqrt(x1 * x1 + x2 * x2 + x3 * x3)
rp = np.sqrt((x1 - a) * (x1 - a) + x2 * x2 + x3 * x3)
BGx2 = data.ds.arr(np.zeros_like(data["magnetic_field_y"]), "G")
BGx2 = 3.0 * x3 * x2 * B0s * Rs**3 * rs**(
-5) + 3.0 * x3 * x2 * B0p * Rp**3 * rp**(-5)
BGx2[rs <= Rs] = 3.0*x3[rs <= Rs]*x2[rs <= Rs]\
*B0s*Rs**3*rs[rs <= Rs]**(-5)
BGx2[rs <= 0.5 * Rs] = 0.0
BGx2[rp <= Rp] = 3.0*x3[rp <= Rp]*x2[rp <= Rp]\
*B0p*Rp**3*rp[rp <= Rp]**(-5)
BGx2[rp <= 0.5 * Rp] = 0.0
return BGx2
ds.add_field(("gas", "BGx2"), function=_BGx2, units="G", take_log=False)
def _BGx3(field, data):
B0s = YTQuantity(2.0, "G")
B0p = YTQuantity(1.0, "G")
Rs = YTQuantity(6.955e+10, "cm")
Rp = YTQuantity(1.5 * 0.10045 * Rs, "cm")
a = YTQuantity(0.047, "au").in_units("cm")
x1 = data["x"].in_units('cm')
x2 = data["y"].in_units('cm')
x3 = data["z"].in_units('cm')
rs = np.sqrt(x1 * x1 + x2 * x2 + x3 * x3)
rp = np.sqrt((x1 - a) * (x1 - a) + x2 * x2 + x3 * x3)
BG_z = data.ds.arr(np.zeros_like(data["magnetic_field_z"]), "G")
BGx3 = (3.0*x3*x3 - rs*rs)*B0s*Rs**3*rs**(-5) \
+ (3.0*x3*x3 - rp*rp)*B0p*Rp**3*rp**(-5)
BGx3[rs <= Rs] = (3.0*x3[rs <= Rs]*x3[rs <= Rs] - \
rs[rs <= Rs]*rs[rs <= Rs])*B0s*Rs**3*rs[rs <= Rs]**(-5)
BGx3[rs <= 0.5 * Rs] = 16.0 * B0s
BGx3[rp <= Rp] = (3.0*x3[rp <= Rp]*x3[rp <= Rp] - \
rp[rp <= Rp]*rp[rp <= Rp])*B0p*Rp**3*rp[rp <= Rp]**(-5)
BGx3[rp <= 0.5 * Rp] = 16.0 * B0p
return BGx3
ds.add_field(("gas", "BGx3"), function=_BGx3, units="G", take_log=False)
def _Bx1(field, data):
return data["gas", "magnetic_field_x"] + data["gas", "BGx1"]
ds.add_field(("gas", "Bx1"), function=_Bx1, units="G", take_log=False)
def _Bx2(field, data):
return data["gas", "magnetic_field_y"] + data["gas", "BGx2"]
ds.add_field(("gas", "Bx2"), function=_Bx2, units="G", take_log=False)
def _Bx3(field, data):
return data["gas", "magnetic_field_z"] + data["gas", "BGx3"]
ds.add_field(("gas", "Bx3"), function=_Bx3, units="G", take_log=False)
def _mag_energy(field, data):
return (data["Bx1"]**2 + data["Bx2"]**2 + data["Bx3"]**2) / (8 * np.pi)
ds.add_field(("gas", "mag_energy"),
function=_mag_energy,
units="g*cm**-1*s**-2",
take_log=True)
def _mag_field_strength(field, data):
return np.sqrt(8. * np.pi * data["mag_energy"])
ds.add_field(("gas", "mag_field_strength"),
function=_mag_field_strength,
units="G",
take_log=True)
def _mag_field_magnitude(field, data):
return np.sqrt(data["Bx1"]**2 + data["Bx2"]**2 + data["Bx3"]**2)
ds.add_field(("gas", "mag_field_magnitude"),
function=_mag_field_magnitude,
units="G",
take_log=True)
def _plasma_b(field, data):
return data['pressure'] / data['mag_energy']
ds.add_field(("gas", "plasma_b"), function=_plasma_b, units="")
def _B_divergence(field, data):
sl_right = slice(None, -2, None)
sl_left = slice(2, None, None)
div_fac = 2.0
ds = div_fac * just_one(data["index", "dx"])
f = data["Bx1"][sl_right, 1:-1, 1:-1] / ds
f -= data["Bx1"][sl_left, 1:-1, 1:-1] / ds
ds = div_fac * just_one(data["index", "dy"])
f += data["Bx2"][1:-1, sl_right, 1:-1] / ds
f -= data["Bx2"][1:-1, sl_left, 1:-1] / ds
ds = div_fac * just_one(data["index", "dz"])
f += data["Bx3"][1:-1, 1:-1, sl_right] / ds
f -= data["Bx3"][1:-1, 1:-1, sl_left] / ds
new_field = data.ds.arr(np.zeros(data["Bx1"].shape, dtype=np.float64),
f.units)
new_field[1:-1, 1:-1, 1:-1] = f
return np.abs(new_field)
ds.add_field(("gas", "B_divergence"),
function=_B_divergence,
units="G/code_length",
validators=[ValidateSpatial(1)],
take_log=True)
def _divB_measure(field, data):
return data["index", "dx"]*np.abs(data['B_divergence'])\
/data['mag_field_magnitude']
ds.add_field(("gas", "divB_measure"),
function=_divB_measure,
units="dimensionless",
take_log=True)
def _mag_alfven_speed(field, data):
ftype = field.name[0]
B = data[ftype, 'mag_field_strength']
return B / np.sqrt(4.0 * np.pi * data[ftype, 'density'])
ds.add_field(("gas", "mag_alfven_speed"),
function=_mag_alfven_speed,
units="cm/s",
take_log=True)
def _mag_mach_alfven(field, data):
ftype = field.name[0]
return data[ftype, 'velocity_magnitude'] / data[ftype,
'mag_alfven_speed']
ds.add_field(("gas", "mag_mach_alfven"),
function=_mag_mach_alfven,
units="",
take_log=True)
def _fc(field, data):
return YTQuantity(2.8, 'G**-1*MHz') * data["gas", "mag_field_strength"]
ds.add_field(('gas', "fc"), function=_fc, units="MHz", take_log=True)
def _fp(field, data):
return YTQuantity(8.98e-3, 'cm**(3/2)*MHz')*\
np.sqrt(1.01*data["gas", "ni"])
ds.add_field(('gas', "fp"), function=_fp, units="MHz", take_log=True)
def _f_ratio(field, data):
return data['gas', 'fc'] / data['gas', 'fp']
ds.add_field(('gas', "fc/fp"), function=_f_ratio, units="", take_log=True)
def _specific_angular_momentum_density_x(field, data):
return data["gas", "specific_angular_momentum_x"] * data["gas",
"density"]
ds.add_field(('index', "specific_angular_momentum_density_x"),
function=_specific_angular_momentum_density_x,
units="cm**-1*g*s**-1",
take_log=False,
force_override=True)
def _specific_angular_momentum_density_y(field, data):
return data["gas", "specific_angular_momentum_y"] * data["gas",
"density"]
ds.add_field(('index', "specific_angular_momentum_density_y"),
function=_specific_angular_momentum_density_y,
units="cm**-1*g*s**-1",
take_log=False,
force_override=True)
def _specific_angular_momentum_density_z(field, data):
return data["gas", "specific_angular_momentum_z"] * data["gas",
"density"]
ds.add_field(('index', "specific_angular_momentum_density_z"),
function=_specific_angular_momentum_density_z,
units="cm**-1*g*s**-1",
take_log=False,
force_override=True)
ds.periodicity = (True, True, True)
def setup_fluid_vector_fields(registry, ftype="gas", slice_info=None):
# Current implementation for gradient is not valid for curvilinear geometries
if is_curvilinear(registry.ds.geometry): return
unit_system = registry.ds.unit_system
# slice_info would be the left, the right, and the factor.
# For example, with the old Enzo-ZEUS fields, this would be:
# slice(None, -2, None)
# slice(1, -1, None)
# 1.0
# Otherwise, we default to a centered difference.
if slice_info is None:
sl_left = slice(None, -2, None)
sl_right = slice(2, None, None)
div_fac = 2.0
else:
sl_left, sl_right, div_fac = slice_info
sl_center = slice(1, -1, None)
def _baroclinic_vorticity_x(field, data):
rho2 = data[ftype, "density"].astype(np.float64)**2
return (data[ftype, "pressure_gradient_y"] *
data[ftype, "density_gradient_z"] -
data[ftype, "pressure_gradient_z"] *
data[ftype, "density_gradient_z"]) / rho2
def _baroclinic_vorticity_y(field, data):
rho2 = data[ftype, "density"].astype(np.float64)**2
return (data[ftype, "pressure_gradient_z"] *
data[ftype, "density_gradient_x"] -
data[ftype, "pressure_gradient_x"] *
data[ftype, "density_gradient_z"]) / rho2
def _baroclinic_vorticity_z(field, data):
rho2 = data[ftype, "density"].astype(np.float64)**2
return (data[ftype, "pressure_gradient_x"] *
data[ftype, "density_gradient_y"] -
data[ftype, "pressure_gradient_y"] *
data[ftype, "density_gradient_x"]) / rho2
bv_validators = [
ValidateSpatial(1, [(ftype, "density"), (ftype, "pressure")])
]
for ax in 'xyz':
n = "baroclinic_vorticity_%s" % ax
registry.add_field((ftype, n),
sampling_type="cell",
function=eval("_%s" % n),
validators=bv_validators,
units=unit_system["frequency"]**2)
create_magnitude_field(registry,
"baroclinic_vorticity",
unit_system["frequency"]**2,
ftype=ftype,
slice_info=slice_info,
validators=bv_validators)
def _vorticity_x(field, data):
vz = data[ftype, "relative_velocity_z"]
vy = data[ftype, "relative_velocity_y"]
f = ((vz[sl_center, sl_right, sl_center] -
vz[sl_center, sl_left, sl_center]) /
(div_fac * just_one(data["index", "dy"])))
f -= ((vy[sl_center, sl_center, sl_right] -
vy[sl_center, sl_center, sl_left]) /
(div_fac * just_one(data["index", "dz"])))
new_field = data.ds.arr(np.zeros_like(vz, dtype=np.float64), f.units)
new_field[sl_center, sl_center, sl_center] = f
return new_field
def _vorticity_y(field, data):
vx = data[ftype, "relative_velocity_x"]
vz = data[ftype, "relative_velocity_z"]
f = ((vx[sl_center, sl_center, sl_right] -
vx[sl_center, sl_center, sl_left]) /
(div_fac * just_one(data["index", "dz"])))
f -= ((vz[sl_right, sl_center, sl_center] -
vz[sl_left, sl_center, sl_center]) /
(div_fac * just_one(data["index", "dx"])))
new_field = data.ds.arr(np.zeros_like(vz, dtype=np.float64), f.units)
new_field[sl_center, sl_center, sl_center] = f
return new_field
def _vorticity_z(field, data):
vx = data[ftype, "relative_velocity_x"]
vy = data[ftype, "relative_velocity_y"]
f = ((vy[sl_right, sl_center, sl_center] -
vx[sl_left, sl_center, sl_center]) /
(div_fac * just_one(data["index", "dx"])))
f -= ((vx[sl_center, sl_right, sl_center] -
vx[sl_center, sl_left, sl_center]) /
(div_fac * just_one(data["index", "dy"])))
new_field = data.ds.arr(np.zeros_like(vy, dtype=np.float64), f.units)
new_field[sl_center, sl_center, sl_center] = f
return new_field
vort_validators = [
ValidateSpatial(1, [(ftype, "velocity_%s" % d) for d in 'xyz']),
ValidateParameter('bulk_velocity')
]
for ax in 'xyz':
n = "vorticity_%s" % ax
registry.add_field((ftype, n),
sampling_type="cell",
function=eval("_%s" % n),
units=unit_system["frequency"],
validators=vort_validators)
create_magnitude_field(registry,
"vorticity",
unit_system["frequency"],
ftype=ftype,
slice_info=slice_info,
validators=vort_validators)
create_squared_field(registry,
"vorticity",
unit_system["frequency"]**2,
ftype=ftype,
slice_info=slice_info,
validators=vort_validators)
def _vorticity_stretching_x(field, data):
return data[ftype, "velocity_divergence"] * data[ftype, "vorticity_x"]
def _vorticity_stretching_y(field, data):
return data[ftype, "velocity_divergence"] * data[ftype, "vorticity_y"]
def _vorticity_stretching_z(field, data):
return data[ftype, "velocity_divergence"] * data[ftype, "vorticity_z"]
for ax in 'xyz':
n = "vorticity_stretching_%s" % ax
registry.add_field((ftype, n),
sampling_type="cell",
function=eval("_%s" % n),
units=unit_system["frequency"]**2,
validators=vort_validators)
create_magnitude_field(registry,
"vorticity_stretching",
unit_system["frequency"]**2,
ftype=ftype,
slice_info=slice_info,
validators=vort_validators)
def _vorticity_growth_x(field, data):
return -data[ftype, "vorticity_stretching_x"] - \
data[ftype, "baroclinic_vorticity_x"]
def _vorticity_growth_y(field, data):
return -data[ftype, "vorticity_stretching_y"] - \
data[ftype, "baroclinic_vorticity_y"]
def _vorticity_growth_z(field, data):
return -data[ftype, "vorticity_stretching_z"] - \
data[ftype, "baroclinic_vorticity_z"]
for ax in 'xyz':
n = "vorticity_growth_%s" % ax
registry.add_field((ftype, n),
sampling_type="cell",
function=eval("_%s" % n),
units=unit_system["frequency"]**2,
validators=vort_validators)
def _vorticity_growth_magnitude(field, data):
result = np.sqrt(data[ftype, "vorticity_growth_x"]**2 +
data[ftype, "vorticity_growth_y"]**2 +
data[ftype, "vorticity_growth_z"]**2)
dot = data.ds.arr(np.zeros(result.shape), "")
for ax in "xyz":
dot += (data[ftype, "vorticity_%s" % ax] *
data[ftype, "vorticity_growth_%s" % ax]).to_ndarray()
result = np.sign(dot) * result
return result
registry.add_field((ftype, "vorticity_growth_magnitude"),
sampling_type="cell",
function=_vorticity_growth_magnitude,
units=unit_system["frequency"]**2,
validators=vort_validators,
take_log=False)
def _vorticity_growth_magnitude_absolute(field, data):
return np.sqrt(data[ftype, "vorticity_growth_x"]**2 +
data[ftype, "vorticity_growth_y"]**2 +
data[ftype, "vorticity_growth_z"]**2)
registry.add_field((ftype, "vorticity_growth_magnitude_absolute"),
sampling_type="cell",
function=_vorticity_growth_magnitude_absolute,
units=unit_system["frequency"]**2,
validators=vort_validators)
def _vorticity_growth_timescale(field, data):
domegax_dt = data[ftype, "vorticity_x"] / data[ftype,
"vorticity_growth_x"]
domegay_dt = data[ftype, "vorticity_y"] / data[ftype,
"vorticity_growth_y"]
domegaz_dt = data[ftype, "vorticity_z"] / data[ftype,
"vorticity_growth_z"]
return np.sqrt(domegax_dt**2 + domegay_dt**2 + domegaz_dt**2)
registry.add_field((ftype, "vorticity_growth_timescale"),
sampling_type="cell",
function=_vorticity_growth_timescale,
units=unit_system["time"],
validators=vort_validators)
########################################################################
# With radiation pressure
########################################################################
def _vorticity_radiation_pressure_x(field, data):
rho = data[ftype, "density"].astype(np.float64)
return (data[ftype, "radiation_acceleration_y"] *
data[ftype, "density_gradient_z"] -
data[ftype, "radiation_acceleration_z"] *
data[ftype, "density_gradient_y"]) / rho
def _vorticity_radiation_pressure_y(field, data):
rho = data[ftype, "density"].astype(np.float64)
return (data[ftype, "radiation_acceleration_z"] *
data[ftype, "density_gradient_x"] -
data[ftype, "radiation_acceleration_x"] *
data[ftype, "density_gradient_z"]) / rho
def _vorticity_radiation_pressure_z(field, data):
rho = data[ftype, "density"].astype(np.float64)
return (data[ftype, "radiation_acceleration_x"] *
data[ftype, "density_gradient_y"] -
data[ftype, "radiation_acceleration_y"] *
data[ftype, "density_gradient_x"]) / rho
vrp_validators = [
ValidateSpatial(1, [(ftype, "density"),
(ftype, "radiation_acceleration_x"),
(ftype, "radiation_acceleration_y"),
(ftype, "radiation_acceleration_z")])
]
for ax in 'xyz':
n = "vorticity_radiation_pressure_%s" % ax
registry.add_field((ftype, n),
sampling_type="cell",
function=eval("_%s" % n),
units=unit_system["frequency"]**2,
validators=vrp_validators)
create_magnitude_field(registry,
"vorticity_radiation_pressure",
unit_system["frequency"]**2,
ftype=ftype,
slice_info=slice_info,
validators=vrp_validators)
def _vorticity_radiation_pressure_growth_x(field, data):
return -data[ftype, "vorticity_stretching_x"] - \
data[ftype, "baroclinic_vorticity_x"] \
-data[ftype, "vorticity_radiation_pressure_x"]
def _vorticity_radiation_pressure_growth_y(field, data):
return -data[ftype, "vorticity_stretching_y"] - \
data[ftype, "baroclinic_vorticity_y"] \
-data[ftype, "vorticity_radiation_pressure_y"]
def _vorticity_radiation_pressure_growth_z(field, data):
return -data[ftype, "vorticity_stretching_z"] - \
data[ftype, "baroclinic_vorticity_z"] \
-data[ftype, "vorticity_radiation_pressure_z"]
for ax in 'xyz':
n = "vorticity_radiation_pressure_growth_%s" % ax
registry.add_field((ftype, n),
sampling_type="cell",
function=eval("_%s" % n),
units=unit_system["frequency"]**2,
validators=vrp_validators)
def _vorticity_radiation_pressure_growth_magnitude(field, data):
result = np.sqrt(
data[ftype, "vorticity_radiation_pressure_growth_x"]**2 +
data[ftype, "vorticity_radiation_pressure_growth_y"]**2 +
data[ftype, "vorticity_radiation_pressure_growth_z"]**2)
dot = data.ds.arr(np.zeros(result.shape), "")
for ax in "xyz":
dot += (data[ftype, "vorticity_%s" % ax] *
data[ftype, "vorticity_growth_%s" % ax]).to_ndarray()
result = np.sign(dot) * result
return result
registry.add_field(
(ftype, "vorticity_radiation_pressure_growth_magnitude"),
sampling_type="cell",
function=_vorticity_radiation_pressure_growth_magnitude,
units=unit_system["frequency"]**2,
validators=vrp_validators,
take_log=False)
def _vorticity_radiation_pressure_growth_magnitude_absolute(field, data):
return np.sqrt(
data[ftype, "vorticity_radiation_pressure_growth_x"]**2 +
data[ftype, "vorticity_radiation_pressure_growth_y"]**2 +
data[ftype, "vorticity_radiation_pressure_growth_z"]**2)
registry.add_field(
(ftype, "vorticity_radiation_pressure_growth_magnitude_absolute"),
sampling_type="cell",
function=_vorticity_radiation_pressure_growth_magnitude_absolute,
units="s**(-2)",
validators=vrp_validators)
def _vorticity_radiation_pressure_growth_timescale(field, data):
domegax_dt = data[ftype, "vorticity_x"] / \
data[ftype, "vorticity_radiation_pressure_growth_x"]
domegay_dt = data[ftype, "vorticity_y"] / \
data[ftype, "vorticity_radiation_pressure_growth_y"]
domegaz_dt = data[ftype, "vorticity_z"] / \
data[ftype, "vorticity_radiation_pressure_growth_z"]
return np.sqrt(domegax_dt**2 + domegay_dt**2 + domegaz_dt**2)
registry.add_field(
(ftype, "vorticity_radiation_pressure_growth_timescale"),
sampling_type="cell",
function=_vorticity_radiation_pressure_growth_timescale,
units=unit_system["time"],
validators=vrp_validators)
def _shear(field, data):
"""
Shear is defined as [(dvx/dy + dvy/dx)^2 + (dvz/dy + dvy/dz)^2 +
(dvx/dz + dvz/dx)^2 ]^(0.5)
where dvx/dy = [vx(j-1) - vx(j+1)]/[2dy]
and is in units of s^(-1)
(it's just like vorticity except add the derivative pairs instead
of subtracting them)
"""
if data.ds.dimensionality > 1:
vx = data[ftype, "relative_velocity_x"]
vy = data[ftype, "relative_velocity_y"]
dvydx = ((vy[sl_right, sl_center, sl_center] -
vy[sl_left, sl_center, sl_center]) /
(div_fac * just_one(data["index", "dx"])))
dvxdy = ((vx[sl_center, sl_right, sl_center] -
vx[sl_center, sl_left, sl_center]) /
(div_fac * just_one(data["index", "dy"])))
f = (dvydx + dvxdy)**2.0
del dvydx, dvxdy
if data.ds.dimensionality > 2:
vz = data[ftype, "relative_velocity_z"]
dvzdy = ((vz[sl_center, sl_right, sl_center] -
vz[sl_center, sl_left, sl_center]) /
(div_fac * just_one(data["index", "dy"])))
dvydz = ((vy[sl_center, sl_center, sl_right] -
vy[sl_center, sl_center, sl_left]) /
(div_fac * just_one(data["index", "dz"])))
f += (dvzdy + dvydz)**2.0
del dvzdy, dvydz
dvxdz = ((vx[sl_center, sl_center, sl_right] -
vx[sl_center, sl_center, sl_left]) /
(div_fac * just_one(data["index", "dz"])))
dvzdx = ((vz[sl_right, sl_center, sl_center] -
vz[sl_left, sl_center, sl_center]) /
(div_fac * just_one(data["index", "dx"])))
f += (dvxdz + dvzdx)**2.0
del dvxdz, dvzdx
np.sqrt(f, out=f)
new_field = data.ds.arr(np.zeros_like(data[ftype, "velocity_x"]),
f.units)
new_field[sl_center, sl_center, sl_center] = f
return new_field
registry.add_field((ftype, "shear"),
sampling_type="cell",
function=_shear,
validators=[
ValidateSpatial(1, [(ftype, "velocity_x"),
(ftype, "velocity_y"),
(ftype, "velocity_z")]),
ValidateParameter('bulk_velocity')
],
units=unit_system["frequency"])
def _shear_criterion(field, data):
"""
Divide by c_s to leave shear in units of length**-1, which
can be compared against the inverse of the local cell size (1/dx)
to determine if refinement should occur.
"""
return data[ftype, "shear"] / data[ftype, "sound_speed"]
registry.add_field((ftype, "shear_criterion"),
sampling_type="cell",
function=_shear_criterion,
units=unit_system["length"]**-1,
validators=[
ValidateSpatial(1, [(ftype, "sound_speed"),
(ftype, "velocity_x"),
(ftype, "velocity_y"),
(ftype, "velocity_z")])
])
def _shear_mach(field, data):
"""
Dimensionless shear (shear_mach) is defined nearly the same as shear,
except that it is scaled by the local dx/dy/dz and the local sound speed.
So it results in a unitless quantity that is effectively measuring
shear in mach number.
In order to avoid discontinuities created by multiplying by dx/dy/dz at
grid refinement boundaries, we also multiply by 2**GridLevel.
Shear (Mach) = [(dvx + dvy)^2 + (dvz + dvy)^2 +
(dvx + dvz)^2 ]^(0.5) / c_sound
"""
if data.ds.dimensionality > 1:
vx = data[ftype, "relative_velocity_x"]
vy = data[ftype, "relative_velocity_y"]
dvydx = (vy[sl_right, sl_center, sl_center] -
vy[sl_left, sl_center, sl_center]) / div_fac
dvxdy = (vx[sl_center, sl_right, sl_center] -
vx[sl_center, sl_left, sl_center]) / div_fac
f = (dvydx + dvxdy)**2.0
del dvydx, dvxdy
if data.ds.dimensionality > 2:
vz = data[ftype, "relative_velocity_z"]
dvzdy = (vz[sl_center, sl_right, sl_center] -
vz[sl_center, sl_left, sl_center]) / div_fac
dvydz = (vy[sl_center, sl_center, sl_right] -
vy[sl_center, sl_center, sl_left]) / div_fac
f += (dvzdy + dvydz)**2.0
del dvzdy, dvydz
dvxdz = (vx[sl_center, sl_center, sl_right] -
vx[sl_center, sl_center, sl_left]) / div_fac
dvzdx = (vz[sl_right, sl_center, sl_center] -
vz[sl_left, sl_center, sl_center]) / div_fac
f += (dvxdz + dvzdx)**2.0
del dvxdz, dvzdx
f *= (
2.0**data["index", "grid_level"][sl_center, sl_center, sl_center] /
data[ftype, "sound_speed"][sl_center, sl_center, sl_center])**2.0
np.sqrt(f, out=f)
new_field = data.ds.arr(np.zeros_like(vx), f.units)
new_field[sl_center, sl_center, sl_center] = f
return new_field
vs_fields = [(ftype, "sound_speed"), (ftype, "velocity_x"),
(ftype, "velocity_y"), (ftype, "velocity_z")]
registry.add_field((ftype, "shear_mach"),
sampling_type="cell",
function=_shear_mach,
units="",
validators=[
ValidateSpatial(1, vs_fields),
ValidateParameter('bulk_velocity')
])
def particle_deposition_functions(ptype, coord_name, mass_name, registry):
unit_system = registry.ds.unit_system
orig = set(registry.keys())
ptype_dn = ptype.replace("_", " ").title()
def particle_count(field, data):
pos = data[ptype, coord_name]
d = data.deposit(pos, method="count")
d = data.ds.arr(d, input_units="cm**-3")
return data.apply_units(d, field.units)
registry.add_field(("deposit", "%s_count" % ptype),
function=particle_count,
validators=[ValidateSpatial()],
units='',
display_name=r"\mathrm{%s Count}" % ptype_dn)
def particle_mass(field, data):
pos = data[ptype, coord_name]
pmass = data[ptype, mass_name]
pmass.convert_to_units(field.units)
d = data.deposit(pos, [pmass], method="sum")
return data.apply_units(d, field.units)
registry.add_field(("deposit", "%s_mass" % ptype),
function=particle_mass,
validators=[ValidateSpatial()],
display_name=r"\mathrm{%s Mass}" % ptype_dn,
units=unit_system["mass"])
def particle_density(field, data):
pos = data[ptype, coord_name].convert_to_units("code_length")
mass = data[ptype, mass_name].convert_to_units("code_mass")
d = data.deposit(pos, [mass], method="sum")
d = data.ds.arr(d, "code_mass")
d /= data["index", "cell_volume"]
return d
registry.add_field(("deposit", "%s_density" % ptype),
function=particle_density,
validators=[ValidateSpatial()],
display_name=r"\mathrm{%s Density}" % ptype_dn,
units=unit_system["density"])
def particle_cic(field, data):
pos = data[ptype, coord_name]
d = data.deposit(pos, [data[ptype, mass_name]], method="cic")
d = data.apply_units(d, data[ptype, mass_name].units)
d /= data["index", "cell_volume"]
return d
registry.add_field(("deposit", "%s_cic" % ptype),
function=particle_cic,
validators=[ValidateSpatial()],
display_name=r"\mathrm{%s CIC Density}" % ptype_dn,
units=unit_system["density"])
def _get_density_weighted_deposit_field(fname, units, method):
def _deposit_field(field, data):
"""
Create a grid field for particle quantities weighted by particle
mass, using cloud-in-cell deposit.
"""
pos = data[ptype, "particle_position"]
# Get back into density
pden = data[ptype, 'particle_mass']
top = data.deposit(pos, [data[(ptype, fname)] * pden],
method=method)
bottom = data.deposit(pos, [pden], method=method)
top[bottom == 0] = 0.0
bnz = bottom.nonzero()
top[bnz] /= bottom[bnz]
d = data.ds.arr(top, input_units=units)
return d
return _deposit_field
for ax in 'xyz':
for method, name in zip(("cic", "sum"), ("cic", "nn")):
function = _get_density_weighted_deposit_field(
"particle_velocity_%s" % ax, "cm/s", method)
registry.add_field(
("deposit", ("%s_" + name + "_velocity_%s") % (ptype, ax)),
function=function,
units=unit_system["velocity"],
take_log=False,
validators=[ValidateSpatial(0)])
for method, name in zip(("cic", "sum"), ("cic", "nn")):
function = _get_density_weighted_deposit_field("age", "s", method)
registry.add_field(("deposit", ("%s_" + name + "_age") % (ptype)),
function=function,
units=unit_system["time"],
take_log=False,
validators=[ValidateSpatial(0)])
# Now some translation functions.
def particle_ones(field, data):
v = np.ones(data[ptype, mass_name].shape, dtype="float64")
return data.apply_units(v, field.units)
registry.add_field((ptype, "particle_ones"),
function=particle_ones,
particle_type=True,
units="",
display_name=r"Particle Count")
def particle_mesh_ids(field, data):
pos = data[ptype, coord_name]
ids = np.zeros(pos.shape[0], dtype="float64") - 1
# This is float64 in name only. It will be properly cast inside the
# deposit operation.
#_ids = ids.view("float64")
data.deposit(pos, [ids], method="mesh_id")
return data.apply_units(ids, "")
registry.add_field((ptype, "mesh_id"),
function=particle_mesh_ids,
validators=[ValidateSpatial()],
units='',
particle_type=True)
return list(set(registry.keys()).difference(orig))
import fastcluster
import scipy.cluster.hierarchy as sch
from mpi4py import MPI
import os, sys
from astropy.table import Table, Column, vstack, hstack
from common_functions import *
sl_left = slice(None, -2, None)
sl_right = slice(2, None, None)
div_fac = 2.0
sl_center = slice(1, -1, None)
ftype = 'gas'
vort_validators = [
ValidateSpatial(1, [(ftype, "velocity_x"), (ftype, "velocity_y"),
(ftype, "velocity_z")])
]
def mad_based_outlier(points, thresh=5):
if len(points.shape) == 1:
points = points[:, None]
median = np.median(points, axis=0)
diff = np.sum((points - median)**2, axis=-1)
diff = np.sqrt(diff)
med_abs_deviation = np.median(diff)
modified_z_score = 0.6745 * diff / med_abs_deviation
return modified_z_score > thresh