def makeColDist(directory, runName, f_list, savefolder, ion, fieldname):
for i in f_list:
data = yt.load(directory + runName + '/KH_hdf5_chk_' + i)
data.add_field(('gas', 'metallicity'),
function=_metallicity,
display_name="Metallicity",
units='Zsun')
#add an ion field to the dataset
tri.add_ion_fields(data, ions=[ion])
#projection of the number density
p = data.proj(fieldname, 1)
#make fixed resolution buffer (same as when a figure is made)
#to have an array with every value in each pixel
frb = yt.FixedResolutionBuffer(
p, (-2.464e21, 2.464e21, -2.464e21, 2.464e21), (800, 800))
#flatten the frb to a 1d array
flattened_coldens = frb[fieldname].flatten()
if min(flattened_coldens) == 0:
mincol = 8
else:
mincol = np.log10(min(flattened_coldens))
maxcol = np.log10(max(flattened_coldens))
#make histogram!
plt.hist(flattened_coldens,
bins=np.logspace(mincol - 1, maxcol + 1, 50),
fill=False,
color='blue',
log=True)
plt.xscale('log')
#plt.yscale('log')
plt.ylabel('Number of pixels')
plt.xlabel(ion + ' Column Density per pixel')
plt.ylim(1, 1e6)
fig = plt.gcf()
fig.savefig(savefolder + '/' + runName + '_' + i + '_' + fieldname +
'.png')
plt.close()
def make_metrics(ds):
""" Generates an frb from a dataset, and returns the corresponding Metrics object. """
axis = "theta"
origin = "native"
fields = ["enuc", "r"]
xlim = 0.0, 6e4
ylim = 0.0, 1.5e4
width = get_width(ds, xlim, ylim)
center = get_center(ds, xlim, ylim)
dim = 256, 1024
axis = ds.coordinates.axis_id.get(axis, axis)
bounds, center, display_center = get_window_parameters(
ds, axis, width, center)
slc = ds.slice(axis, center[axis], center=center)
slc.get_data(fields)
frb = yt.FixedResolutionBuffer(slc, bounds, dim)
return Metrics(frb)
def run(self):
# self.ds already exists
sl = self.ds.slice(2, 0.5)
frb = yt.FixedResolutionBuffer(sl, (0.0, 1.0, 0.0, 1.0), (400, 400), antialias=False)
dens = frb["Density"]
return np.array([dens.mean(), dens.std(), dens.min(), dens.max()])
def calcRankTau(directory, runName, runNumber, ions, velocityBin, frameVel):
#load and add ions to the dataset
velocityBin = velocityBin * 1.0e5 #convert to cm/s
frameVel = frameVel * 1.0e5 #convert to cm/s
data = yt.load(directory + runName + '/KH_hdf5_chk_' + runNumber)
data.add_field(('gas', 'metallicity'),
function=_metallicity,
display_name="Metallicity",
units='Zsun')
absorbFracs = []
#add ion fields to the dataset
for ion in ions:
tri.add_ion_fields(data, ions=[ion['ion']])
#define the function for b(x) for this particular ion
### function definition in a for loop... Ew... ###
def _soundSpeed(field, data):
topFrac = 2.0 * kb * data['temperature']
botFrac = ion['massNum'] * mp
b = np.sqrt(topFrac / botFrac)
return b
#add b_soundSpeed to the dataset
data.add_field(('gas', 'b_soundSpeed'),
function=_soundSpeed,
display_name="B Sound Speed",
units='cm/s',
force_override=True)
#define the function for dTau for this particular ion/velocity
### function definition in a for loop... Ew... ###
def _dTau(field, data):
term1 = data[ion['fieldname']] / data['b_soundSpeed']
#assumes velocity is now cm/s
#must add the frame velocity to the y velocity!
expTerm = (((data['velocity_y'] + frameVel *
(centimeter / second)) - velocityBin *
(centimeter / second)) / data['b_soundSpeed'])**2.0
dtau = term1 * np.exp(-1.0 * expTerm)
return dtau
#add dTau to the dataset
data.add_field(('gas', 'tau'),
function=_dTau,
display_name="dTau",
units='s/cm**4',
force_override=True)
#make cylinder, project through the y axis
#reg = data.disk([0.0, 0.2*kpc, 0.0], [0,1,0], (0.25, 'kpc'), (1.6, 'kpc'))
#select entire region
reg = data.all_data()
plot = yt.ProjectionPlot(data, 'y', 'density')
plot.set_zlim('density', 1e-5, 5e-5)
plot.save()
#projection of the number density
p = data.proj('tau', 1)
#make fixed resolution buffer (same as when a figure is made)
#to have an array with every value in each pixel
frb = yt.FixedResolutionBuffer(
p, (-2.464e21, 2.464e21, -2.464e21, 2.464e21), (800, 800))
#flatten the frb to a 1d array
flattened_tau = frb['tau'].flatten()
tau_coef = c_speed * ion['sigma'] / np.sqrt(np.pi)
projectedTau = tau_coef * flattened_tau
#rank sort the projected taus and then make plot
sorted_tau = np.sort(projectedTau, kind='quicksort')
totalpixel = len(sorted_tau)
pixelNum = np.arange(0.0, totalpixel, 1.0)
pixelFrac = pixelNum / totalpixel
print(pixelFrac)
writeFile = open(
'../rankTau' + ion['ionfolder'] + 'rankTau' + runName + '_v' +
str(int(round(velocityBin / 1e5, 0))) + '.txt', 'w')
for i in range(len(pixelFrac)):
writeString = str(sorted_tau[i].value) + '\n'
writeFile.write(writeString)
writeFile.close()
#clip = 620800 #0.97
#plt.plot(pixelFrac[clip:], sorted_tau[clip:])
#fig = plt.gcf()
#fig.savefig('rankTau/test_rankTau'+runName+'_v'+str(int(round(velocityBin/1e5, 0)))+'_'+ion['ion']+'.png')
#fig.clear()
return absorbFracs
def calcRankTau(directory, runName, runNumber, ions, velocityBin, frameVel):
#load and add ions to the dataset
frameVel = frameVel*1.0e5*(centimeter/second) #convert to cm/s
data = yt.load(directory+runName+'/KH_hdf5_chk_'+runNumber)
data.add_field(('gas', 'metallicity'), function=_metallicity, display_name="Metallicity", units='Zsun')
bList = []
vList = []
TList = []
#add ion fields to the dataset
for ion in ions:
tri.add_ion_fields(data, ions=[ion['ion']])
#select entire region
reg = data.all_data()
#projection of the number density (= fieldname)
p = data.proj(ion['fieldname'], 1)
#make fixed resolution buffer (same as when a figure is made)
#to have an array with every value in each pixel
frb = yt.FixedResolutionBuffer(p, (-2.464e21, 2.464e21, -2.464e21, 2.464e21), (800, 800))
#flatten the frb to a 1d array
flattened_num = frb[ion['fieldname']].flatten()
projectedNum = flattened_num
#rank sort the projected taus and then make plot
sorted_num = np.sort(projectedNum, kind='quicksort')
#find the average velocity for the ion
average_vely = reg.quantities.weighted_average_quantity('vely', ion['fieldname'])+frameVel
average_temp = reg.quantities.weighted_average_quantity('temperature', ion['fieldname'])
#write the ranked column densities to a file
writeFile = open('../rankNum'+ion['ionfolder']+'rankNum'+runName+'_v'+str(velocityBin)+'.txt', 'w')
for i in range(len(sorted_num)):
writeString = str(sorted_num[i].value)+'\n'
writeFile.write(writeString)
writeFile.close()
#add b^2 thermal as a field
def _btherm(field, data):
topFrac = 2.0*kb*data['temperature']
botFrac = ion['massNum']*mp
b = topFrac/botFrac
return b
data.add_field(('gas', 'b_therm'), function=_btherm, display_name="B thermal squared", units = 'cm**2/s**2', force_override=True)
#project b thermal
btherm = data.proj('b_therm', 1, weight_field=ion['fieldname'])
frb_therm = yt.FixedResolutionBuffer(btherm, (-2.464e21, 2.464e21, -2.464e21, 2.464e21), (800, 800))
flattened_btherm = frb_therm['b_therm'].flatten()
#add b^2 doppler as a field
def _bdop(field, data):
b = (data['vely'] - average_vely)**2
return b #multiply by number density to weight the b value
data.add_field(('gas', 'b_doppler'), function=_bdop, display_name="B doppler squared", units = 'cm**2/s**2', force_override=True)
#project b doppler
bdop = data.proj('b_doppler', 1, weight_field=ion['fieldname'])
frb_dop = yt.FixedResolutionBuffer(bdop, (-2.464e21, 2.464e21, -2.464e21, 2.464e21), (800, 800))
flattened_bdop = frb_dop['b_doppler'].flatten()
good_btherm = flattened_btherm[~np.isnan(flattened_btherm)]
good_bdop = flattened_bdop[~np.isnan(flattened_bdop)]
#add the two b's and sqrt
b_tot_Sq = good_btherm+good_bdop
b_tot = np.sqrt(b_tot_Sq)
#compute the average of the summation of b projections
b_avg = np.average(b_tot)
bList.append(b_avg)
vList.append(average_vely)
TList.append(average_temp)
print(b_avg)
print(average_vely)
return bList, vList, TList
def run(self):
sl = self.ds.slice(2, 0.5)
frb = yt.FixedResolutionBuffer(sl, (0.0, 1.0, 0.0, 1.0), (200, 200))
dd = frb[self.field]
return np.array([dd.mean(), dd.std(), dd.min(), dd.max()])
def getDensities(runName, time, ionList, velBin, frameVel):
frameVel = frameVel * 1.0e5 * (centimeter / second)
data = yt.load('../' + runName + '/chkfiles/CT_hdf5_chk_' + time)
data.add_field(('gas', 'metallicity'),
function=_metallicity,
display_name="Metallicity",
units='Zsun',
sampling_type='cell')
bList = []
vList = []
TList = []
for ion in ionList:
print(ion['name'], end=' ')
for source in ['chem']: #here you can define chem or trident
#add the ion to the dataset
_ = addfield(data, ion['atom'], ion['ion'], 'dens', source)
#select the region
reg = data.all_data()
p = data.proj(ion[source], 1)
frb = yt.FixedResolutionBuffer(
p, (-2.464e21, 2.464e21, -2.464e21, 2.464e21), (800, 800))
#flatten the frb to a 1d array
flattened_num = frb[ion[source]].flatten()
projectedNum = flattened_num
#rank sort the projected taus and then make plot
sorted_num = np.sort(projectedNum, kind='quicksort')
#find the average velocity for the ion
average_vely = reg.quantities.weighted_average_quantity(
'vely', ion[source]) + frameVel
average_temp = reg.quantities.weighted_average_quantity(
'temperature', ion[source])
#write out the ranked Column densities to be used to make a fit
writeFile = open(
'../rankNum' + ion['ionfolder'] + 'rankNum' + runName + '_v' +
str(velBin) + '.txt', 'w')
for i in range(len(sorted_num)):
writeString = str(sorted_num[i].value) + '\n'
writeFile.write(writeString)
writeFile.close()
#add b^2 thermal as a field
def _btherm(field, data):
topFrac = 2.0 * kb * data['temperature']
botFrac = ion['massNum'] * mp
b = topFrac / botFrac
return b
data.add_field(('gas', 'b_therm'),
function=_btherm,
display_name="B thermal squared",
units='cm**2/s**2',
force_override=True,
sampling_type='cell')
#project b thermal
btherm = data.proj('b_therm', 1, weight_field=ion[source])
frb_therm = yt.FixedResolutionBuffer(
btherm, (-2.464e21, 2.464e21, -2.464e21, 2.464e21), (800, 800))
flattened_btherm = frb_therm['b_therm'].flatten()
#add b^2 doppler as a field
def _bdop(field, data):
b = (data['vely'] - average_vely)**2
return b #multiply by number density to weight the b value
data.add_field(('gas', 'b_doppler'),
function=_bdop,
display_name="B doppler squared",
units='cm**2/s**2',
force_override=True,
sampling_type='cell')
#project b doppler
bdop = data.proj('b_doppler', 1, weight_field=ion[source])
frb_dop = yt.FixedResolutionBuffer(
bdop, (-2.464e21, 2.464e21, -2.464e21, 2.464e21), (800, 800))
flattened_bdop = frb_dop['b_doppler'].flatten()
good_btherm = flattened_btherm[~np.isnan(flattened_btherm)]
good_bdop = flattened_bdop[~np.isnan(flattened_bdop)]
#add the two b's and sqrt
b_tot_Sq = good_btherm + good_bdop
b_tot = np.sqrt(b_tot_Sq)
#compute the average of the summation of b projections
b_avg = np.average(b_tot)
bList.append(b_avg)
vList.append(average_vely)
TList.append(average_temp)
#print(b_avg)
#print(average_vely)
return bList, vList, TList
def plot_a_region(
ram,
out,
ds,
center,
fields='den',
kind='slc',
axis='z',
width=None,
center_vel=None,
L=None,
direcs='face',
zlims={},
l_max=None,
is_id=False,
bar_length=2e3,
sketch=False,
time_offset='tRelax',
is_set_size=True,
is_time=True,
streamplot_kwargs={},
more_kwargs={},
):
"""Do projection or slice plots in a region.
Args:
ram: a ramtools.Ramses instance
out (int): output frame
ds (yt.ds): yt.load instance
center (list_like): the center of the region like in yt.SlicePlot
fields (str or tuple): the field to plot. The avaialble options are: 'den' or 'density' - density. 'logden' - log of density. 'T' or 'temperature' - temperature. 'T_vel_rela' - temperature overplot with velocity field. 'pressure' - thermal pressure. 'magstream' - stream lines of magnetic fields on top of density slice. The magnetic strength is be indicated by the color of the stream lines. 'beta' - plasma beta parameter. 'AlfMach' - Alfvenic Mach number. 'xHII' - hydrogen ionization fraction. 'mach' - thermal Mach number. 'mag' - magnetic field strength. 'vel' - velocity field on top of density slice. 'vel_rela' - relative velocity field, i.e. the velocity field in the frame with a velocity defined by center_vel. 'mach2d' - thermal Mach number in the plane. 'p_mag' - magnetic pressure.
kind (str): 'slc' or 'prj'. Default: 'slc'
axis (str or int): One of (0, 1, 2, 'x', 'y', 'z').
width (float or tuple): width of the field of view. e.g. 0.01,
(1000, 'AU').
center_vel (list_like or None): the velocity of the center. (Default
None)
L (list_like): the line-of-sight vector. Will overwrite axis.
Default: None
direcs (str): 'face' or 'edge'. Will only be used if L is not None.
zlims (dict): The limits of the fields. e.g. {'density': [1e4, 1e8]},
{'B': [1e-5, 1e-2]} (the default of B field).
l_max (int):
is_id (bool): toggle marking sink ID (the order of creation)
bar_length (float): not using
sketch (bool): If True, will do a faster plot of a simple imshow for
test purpose. Default: False.
is_time (bool): toggle overplotting time tag
time_offset (str or int or float): One of the following cases:
str: 'rRelax'
int: the frame of which the time is used
float: time in Myr
is_set_size (bool): toggle setting figure size to 6 to make texts
bigger. Default: True.
streamplot_kwargs (dict): More kwargs for plt.streamplot that is used when fields='magstream'. Default: {}
more_kwargs (dict): more field-specific kwargs. Default: {}
Returns:
yt figure or plt.figure
"""
from matplotlib import colors
r = ram
ytfast.set_data_dir(os.path.join(r.ramses_dir, 'h5_data'))
assert width is not None
width = to_boxlen(width, r.ds1)
if isinstance(time_offset, int):
time_offset = r.get_time(time_offset)
elif isinstance(time_offset, str):
if time_offset is 'tRelax':
time_offset = r.tRelax
else:
time_offset = time_offset
if l_max is None: l_max = ds.max_level
use_h5 = False # TODO: enable use_h5
# calculated north and right, if L is not None
if L is None:
if axis in [0, 1, 2]:
axis = ['x', 'y', 'z'][axis]
is_onaxis = True
los = axis
axismap = {'x': [1, 0, 0], 'y': [0, 1, 0], 'z': [0, 0, 1]}
los_v = axismap[axis]
north = axismap[{'x': 'z', 'y': 'x', 'z': 'y'}[axis]]
logger.debug('is_onaxis = {}'.format(is_onaxis))
else:
is_onaxis = False
L /= norm(L)
tmpright = np.cross([0, 0, 1], L)
tmpright = tmpright / norm(tmpright)
if direcs == 'face':
los = L
north = -1 * tmpright
elif direcs == 'edge':
los = tmpright
north = L
los_v = los
right = np.cross(north, los_v)
logger.debug('is_onaxis = {}'.format(is_onaxis))
logger.info(f'doing {r.jobPath}, los={los}, north={north}, right={right}')
# ------------------ doing the plot --------------------
if 'density' in zlims:
den_zlim = zlims['density']
elif 'den' in zlims:
den_zlim = zlims['den']
else:
den_zlim = [None, None]
gas_field = fields
is_log = True
zlim = None
cb_label = None
cmap = 'viridis'
if fields in ['T', 'temperature', 'T_vel_rela']:
gas_field = 'temperature'
if fields in ['pressure', 'p_mag']:
cmap = 'gist_heat'
if fields in ['beta', 'magbeta', 'AlfMach']:
cmap = 'seismic'
add_beta_corrected(ds)
zlim = [-3, 3]
is_log = False
if fields == 'magbeta':
gas_field = "logPlasmaBeta"
def _logBeta(field, data):
return np.log10(data['beta'])
ds.add_field(("gas", gas_field), function=_logBeta)
cb_label = "log (Plasma Beta)"
if fields == 'AlfMach':
gas_field = ("gas", "LogMachA")
def _LogMachA(field, data):
""" MachA = v_rms / v_A = gamma / 2 * Mach^2 * beta """
gamma = 5 / 3
return np.log10(gamma / 2 * data["mach"]**2 * data["beta"])
ds.add_field(gas_field, function=_LogMachA)
cb_label = "log $M_A$"
if fields in ['xHII']:
zlim = [1e-6, 1]
if fields in ['mach', 'AlfMach']:
add_mach(ds, center_vel)
if gas_field == 'mach':
gas_field = 'logmach'
add_logmach(ds)
cb_label = r'log ($\mathcal{M}_s$)'
is_log = False
cmap = 'seismic'
zlim = [-2, 2]
if fields in ['den', 'density', 'mag', 'vel', 'vel_rela']:
gas_field = 'density'
if fields in ['logden', ('gas', 'logden')]:
def _logden(_field, data):
mu = 1.4
return data[('gas', 'density')] / (
mu * yt.physical_constants.mass_hydrogen)
gas_field = ('gas', 'logden')
ds.add_field(gas_field,
function=_logden,
take_log=True,
sampling_type='cell',
units="cm**(-3)")
if fields == 'magstream':
def _rel_B_1(_field, data):
Bx = (data["x-Bfield-left"] + data["x-Bfield-right"]) / 2.
By = (data["y-Bfield-left"] + data["y-Bfield-right"]) / 2.
Bz = (data["z-Bfield-left"] + data["z-Bfield-right"]) / 2.
B = yt.YTArray([Bx, By, Bz])
return np.tensordot(right, B, 1)
def _rel_B_2(_field, data):
Bx = (data["x-Bfield-left"] + data["x-Bfield-right"]) / 2.
By = (data["y-Bfield-left"] + data["y-Bfield-right"]) / 2.
Bz = (data["z-Bfield-left"] + data["z-Bfield-right"]) / 2.
B = yt.YTArray([Bx, By, Bz])
return np.tensordot(north, B, 1)
ds.add_field(('gas', 'rel_B_1'), function=_rel_B_1)
ds.add_field(('gas', 'rel_B_2'), function=_rel_B_2)
if fields in ['vel_rela', 'mach2d']:
assert north is not None
def _rel_vel_1(_field, data):
rel_vel = yt.YTArray([
data['velocity_x'] - yt.YTArray(center_vel[0], 'cm/s'),
data['velocity_y'] - yt.YTArray(center_vel[1], 'cm/s'),
data['velocity_z'] - yt.YTArray(center_vel[2], 'cm/s')
])
return yt.YTArray(np.tensordot(right, rel_vel, 1), 'cm/s')
def _rel_vel_2(_field, data):
rel_vel = yt.YTArray([
data['velocity_x'] - yt.YTArray(center_vel[0], 'cm/s'),
data['velocity_y'] - yt.YTArray(center_vel[1], 'cm/s'),
data['velocity_z'] - yt.YTArray(center_vel[2], 'cm/s')
])
return yt.YTArray(np.tensordot(north, rel_vel, 1), 'cm/s')
def _mach2d(_field, data):
rel_speed_sq = data['rel_vel_1'].in_cgs().value ** 2 + \
data['rel_vel_2'].in_cgs().value ** 2 # cm/s
T = data['temperature'].value
# km/s, assuming gamma=5/3. mu is inside T, therefore this
# is precise.
cs = 0.11729 * np.sqrt(T)
return 1e-5 * np.sqrt(rel_speed_sq) / cs # dimensionless
ds.add_field(('gas', 'rel_vel_1'), function=_rel_vel_1, units="cm/s")
ds.add_field(('gas', 'rel_vel_2'), function=_rel_vel_2, units="cm/s")
ds.add_field(
('gas', 'mach2d'),
function=_mach2d,
)
if fields == 'mach2d':
gas_field = 'logmach2d'
def _logmach2d(_field, data):
return np.log10(data["mach2d"])
ds.add_field(
('gas', gas_field),
function=_logmach2d,
)
cb_label = r'log ($\mathcal{M}_{s, 2D}$)'
is_log = False
cmap = "seismic"
zlim = [-2, 2]
if fields != 'magstream':
if kind == 'slc':
if is_onaxis:
p = yt.SlicePlot(ds,
los,
gas_field,
width=width,
center=center)
else:
p = yt.OffAxisSlicePlot(ds,
los,
gas_field,
width=width,
center=center,
north_vector=north)
elif kind == 'prj':
if is_onaxis:
p = ytfast.ProjectionPlot(ds,
los,
gas_field,
width=width,
center=center,
weight_field='density')
else:
print("Doing OffAxis projection plot...")
p = yt.OffAxisProjectionPlot(ds,
los,
gas_field,
width=width,
center=center,
weight_field='density',
max_level=l_max,
north_vector=north)
print("Done")
elif kind == 'colden':
if is_onaxis:
p = ytfast.ProjectionPlot(ds,
los,
gas_field,
width=width,
center=center,
weight_field=None)
else:
p = yt.OffAxisProjectionPlot(ds,
los,
gas_field,
width=width,
center=center,
weight_field=None,
max_level=l_max,
north_vector=north)
# p.set_colorbar_label(field, _label)
if is_set_size:
p.set_figure_size(6)
p.set_axes_unit('AU')
if gas_field in zlims:
if zlims[gas_field] is not None:
p.set_zlim(gas_field, *zlims[gas_field])
if is_time:
p.annotate_timestamp(
time_format='{time:.2f} {units}',
time_offset=time_offset,
text_args={
'fontsize': 8,
'color': 'w'
},
corner='upper_left',
)
# Overplot sink particles
args = {}
if 'sink_colors' in more_kwargs.keys():
args['colors'] = more_kwargs['sink_colors']
if 'mass_lims' in more_kwargs.keys():
args['lims'] = more_kwargs['mass_lims']
if 'colors' in more_kwargs.keys():
args['colors'] = more_kwargs['colors']
r.overplot_sink_with_id(p,
out,
center,
width / 2,
is_id=is_id,
zorder='mass',
withedge=1,
**args)
# set cmap, zlim, and other styles
p.set_log(gas_field, is_log)
p.set_cmap(gas_field, cmap)
if zlim is not None:
p.set_zlim(gas_field, zlim)
p.set_colorbar_label(gas_field, cb_label)
if gas_field == 'density':
den_setup(p, den_zlim, time_offset=time_offset)
if gas_field == 'temperature':
T_setup(p, [3, 1e4], time_offset=time_offset)
if fields == 'vel':
p.annotate_velocity(factor=30,
scale=scale,
scale_units='x',
plot_args={"color": "cyan"})
p.set_colorbar_label(gas_field, r'log ($\mathcal{M}_s$)')
if fields == ['vel_rela', 'T_vel_rela']:
p.annotate_quiver('rel_vel_1',
'rel_vel_2',
factor=30,
scale=scale,
scale_units='x',
plot_args={"color": "cyan"})
if fields in ['vel', 'vel_rela', 'T_vel_rela']:
# scale ruler for velocities
ruler_v_kms = 4 # km/s
width_au = width * r.boxlen * 2e5
bar_length = 2e3 * width_au / 2e4
coeff = bar_length # length of ruler in axis units [AU]
scale = 1e5 * ruler_v_kms / coeff # number of cm/s per axis units
# add a scale ruler
p.annotate_scale(coeff=coeff,
unit='AU',
scale_text_format="{} km/s".format(ruler_v_kms))
return p
else: # streamplot
if not sketch:
sl = ds.cutting(los_v, center, north_vector=north)
hw = width / 2.
bounds = [-hw, hw, -hw, hw]
size = 2**10
frb = yt.FixedResolutionBuffer(sl, bounds, [size, size])
m_h = 1.6735575e-24
mu = 1.4
den = frb['density'].value / (mu * m_h)
unitB = utilities.get_unit_B(ds)
logger.info(f"unitB.value = {unitB.value}")
Bx = frb['rel_B_1'].value * unitB.value # Gauss
By = frb['rel_B_2'].value * unitB.value # Gauss
Bmag = np.sqrt(Bx**2 + By**2)
logger.info(f"Bmag min = {Bmag.min()}, max = {Bmag.max()}")
hw_au = hw * r.unit_l / AU
grid_base = np.linspace(-hw_au, hw_au, size)
bounds_au = [-hw_au, hw_au, -hw_au, hw_au]
else:
den = np.random.random([3, 3])
hw_au = 10
den_zlim = [None, None]
bounds_au = [-hw_au, hw_au, -hw_au, hw_au]
# fig, ax = plt.subplots()
if 'figax' in more_kwargs.keys():
fig, ax = more_kwargs['figax']
else:
fig, ax = plt.subplots()
im = ax.imshow(np.log10(den),
cmap='inferno',
extent=bounds_au,
vmin=den_zlim[0],
vmax=den_zlim[1],
origin='lower')
colornorm_den = colors.Normalize(vmin=den_zlim[0], vmax=den_zlim[1])
ax.set(xlabel="Image x (AU)",
ylabel="Image y (AU)",
xlim=[-hw_au, hw_au],
ylim=[-hw_au, hw_au])
cmap = mpl.cm.get_cmap(
"Greens",
6,
)
Bvmin, Bvmax = -5, -2
if "Blims_log" in more_kwargs.keys():
Bvmin, Bvmax = more_kwargs["Blims_log"]
if "B" in zlims.keys():
Bvmin = np.log10(zlims['B'][0])
Bvmax = np.log10(zlims['B'][1])
colornorm_stream = colors.Normalize(vmin=Bvmin, vmax=Bvmax)
# streamplot_kwargs = {}
# if "streamplot" in more_kwargs.keys():
# streamplot_kwargs = more_kwargs["streamplot"]
# linewidth = 0.2
if "stream_linewidth" in more_kwargs.keys():
linewidth = more_kwargs["stream_linewidth"]
if not sketch:
logmag = np.log10(Bmag)
# scaled_mag = (logmag - Bvmin) / (Bvmax - Bvmin)
strm = ax.streamplot(
grid_base,
grid_base,
Bx,
By,
color=np.log10(Bmag),
density=1.2,
# linewidth=linewidth,
# linewidth=1.5 * scaled_mag,
# cmap='Greens',
cmap=cmap,
norm=colornorm_stream,
arrowsize=0.5,
**streamplot_kwargs,
)
# plt.subplots_adjust(right=0.8)
plot_cb = True
plot_cb2 = True
if 'plot_cb' in more_kwargs.keys():
plot_cb = more_kwargs['plot_cb']
if 'plot_cb2' in more_kwargs.keys():
plot_cb2 = more_kwargs['plot_cb2']
if plot_cb:
if 'cb_axis' in more_kwargs.keys():
ax2 = more_kwargs['cb_axis']
# pos0 = ax.get_position()
# cbaxis = fig.add_axes([pos0.x1, pos0.y0, 0.06, pos0.height])
# cb = plt.colorbar(im, cax=cbaxis)
cb2 = mpl.colorbar.ColorbarBase(
ax2,
orientation='vertical',
cmap=mpl.cm.get_cmap("inferno"),
norm=colornorm_den,
)
else:
cb2 = fig.colorbar(im, ax=ax, pad=0)
cb2.set_label("log $n$ (cm$^{-3}$)")
if plot_cb2:
if 'cb2_axis' in more_kwargs.keys():
ax2 = more_kwargs['cb2_axis']
elif plot_cb:
pos1 = ax.get_position()
ax2 = fig.add_axes([pos1.x1 + .15, pos1.y0, 0.02, pos1.height])
else:
pos1 = ax.get_position()
ax2 = fig.add_axes([pos1.x0, pos1.y0, 0.02, pos1.height])
colornorm2 = colors.Normalize(vmin=den_zlim[0], vmax=den_zlim[1])
cb2 = mpl.colorbar.ColorbarBase(
ax2,
orientation='vertical',
cmap=cmap,
norm=colornorm_stream,
# ticks=[-4, -3, -2]
)
cb2.set_label("log B (Gauss)")
# plt.savefig(fs['magstream'], dpi=300)
return fig
z_tick_labels = ['$' + '{0:0.1f}'.format(zt) + '$' for zt in z_tick_locations]
#### For setting the resolution and figure size in inches
aspect_z_div_r = (z_up - z_down) / (r_up - r_down)
dpi = args.dpi # Figure resolution desired (in both dimensions)
inches_r = 2.5 # Width of the resulting desired image
inches_z = inches_r * aspect_z_div_r
npix_r = int(dpi *
inches_r) # Number of FRB pixels to use in the radial direction
npix_z = int(dpi * inches_z)
res = (npix_z, npix_r)
## Create a fixed-resolution buffer object with the above parameters
#### Given the geometry, center the plot so the z-axis is the left edge
#### r_down, r_up, z_down, z_up multiplied by 1e5 to convert (km) to data units (cm)
frb = yt.FixedResolutionBuffer(
slc, (r_down * 1e5, r_up * 1e5, z_down * 1e5, z_up * 1e5), res)
if args.field == 'rgbcomp':
## Get a numpy array of shape res for each progress variable
phfa = np.array(frb['phfa'])
phaq = np.array(frb['phaq'])
phqn = np.array(frb['phqn'])
## Create an array of shape (npix_z, npix_r, 3) with the RGB values
## WHITE: Fuel
## RED: Ash
## GREEN: QNSE
## BLACK: NSE
rgb = np.empty((npix_z, npix_r, 3), dtype=float)
rgb[:, :, 0] = 1.0 - phaq[:, :] # Red
rgb[:, :, 1] = 1.0 - phfa[:, :] + phaq[:, :] - phqn[:, :] # Green