def save_to_fits(ds, mhd=True):
units = pa.set_units(muH=1.4271)
coolftn = pa.coolftn()
if not os.path.isdir(ds.dir + 'fits'): os.mkdir(ds.dir + 'fits')
fields = ['density', 'temperature', 'velocity']
if mhd: fields.append('magnetic_field')
for field in fields:
if field is 'temperature':
data = coolftn.get_temp(ds.read_all_data('T1'))
else:
data = ds.read_all_data(field)
fitsname = '%s/fits/%s.%s.%s.fits' % (ds.dir, ds.id, ds.step,
field[:4])
hdr = fits.Header()
hdr['field'] = field
hdr['time'] = ds.domain['time']
hdr['tMyr'] = (ds.domain['time'] * units['time'].to('Myr').value,
'Myr')
hdr['xmin'] = (ds.domain['left_edge'][0], 'pc')
hdr['xmax'] = (ds.domain['right_edge'][0], 'pc')
hdr['ymin'] = (ds.domain['left_edge'][1], 'pc')
hdr['ymax'] = (ds.domain['right_edge'][1], 'pc')
hdr['zmin'] = (ds.domain['left_edge'][2], 'pc')
hdr['zmax'] = (ds.domain['right_edge'][2], 'pc')
hdr['dx'] = (ds.domain['dx'][0], 'pc')
hdr['dy'] = (ds.domain['dx'][1], 'pc')
hdr['dz'] = (ds.domain['dx'][2], 'pc')
hdr['unit'] = (units[field].value, units[field].unit)
if ds.domain.has_key('qshear'):
hdr['qshear'] = ds.domain['qshear']
if ds.domain.has_key('Omega'):
hdr['Omega'] = (ds.domain['Omega'], 'km/s/pc')
hdu = fits.PrimaryHDU(data, header=hdr)
hdu.writeto(fitsname, overwrite=True)
def sp_plot(ax, sp, plot_runaway=True):
#This is comment for the function.
#This will be printed when you call your function in help.
#You may want to write usage or examples.
import pyathena as pa
unit = pa.set_units(muH=1.4271)
Msun = unit['mass'].to('Msun').value
Myr = unit['time'].to('Myr').value
young_sp = sp[sp['age'] * Myr < 40.]
runaway = young_sp[young_sp['mass'] == 0]
cluster = young_sp[young_sp['mass'] != 0]
mass = cluster['mass'] * Msun
age = cluster['age'] * Myr
cl = ax.scatter(cluster['x1'],
cluster['x2'],
marker='o',
s=mass_norm(mass),
c=age,
vmax=40,
vmin=0,
cmap=plt.cm.cool_r)
if (plot_runaway):
ax.scatter(runaway['x1'], runaway['x2'], marker='.', color='k')
return cl
def get_proj_ims(ds,slab_range=None,test=False):
if slab_range is None:
ns_start=ds.NGrids[2]/4+1
ns_end=ds.NGrids[2]/4*3+1
else:
ns_start,ns_end = slab_range
Lx=ds.domain['Lx'][1]
x2min=ds.domain['left_edge'][1]-Lx
x2max=ds.domain['right_edge'][1]+Lx
x3min=ds.domain['right_edge'][2]
x3max=ds.domain['left_edge'][2]
for ns in range(ns_start,ns_end):
g=ds._get_slab_grid(slab=ns,verbose=test)
x3min=min(x3min,g[0]['left_edge'][2])
x3max=max(x3max,g[0]['right_edge'][2])
extent=(x2min,x2max,x3min,x3max)
dlist=[]
for ns in range(ns_start,ns_end):
dlist.append(ds.read_all_data('density',slab=ns))
den=np.concatenate(dlist)
dlist=[]
for ns in range(ns_start,ns_end):
dlist.append(ds.read_all_data('T1',slab=ns))
T1=np.concatenate(dlist)
dlist=[]
for ns in range(ns_start,ns_end):
dlist.append(ds.read_all_data('magnetic_field',slab=ns))
B=np.concatenate(dlist)
Bmag=np.sqrt((B**2).sum(axis=-1))
coolftn=pa.coolftn()
temp=coolftn.get_temp(T1)
Nz,Ny,Nx=den.shape
unit=pa.set_units(muH=1.4271)
unit['magnetic_field']
dproj=den.mean(axis=2)
Tproj=(den*temp).mean(axis=2)/dproj
Bproj=(den*Bmag).mean(axis=2)*unit['magnetic_field'].value/dproj
#dproj=np.roll(dproj,Nx/2,axis=1)
#Tproj=np.roll(Tproj,Nx/2,axis=1)
#Bproj=np.roll(Bproj,Nx/2,axis=1)
#Tproj=(temp).mean(axis=2)
dproj=np.concatenate([dproj,dproj,dproj],axis=1)
Tproj=np.concatenate([Tproj,Tproj,Tproj],axis=1)
Bproj=np.concatenate([Bproj,Bproj,Bproj],axis=1)
return dproj,Tproj,Bproj,extent
def set_units(self):
"""set units and a few conversion constants
"""
units = pa.set_units(muH=1.4271)
self.to_pok = (units['pressure'] / ac.k_B).cgs.value
self.to_Myr = units['time'].to('Myr').value
self.to_flux = (units['density'] *
units['velocity']).to('Msun/(kpc^2*yr)').value
self.to_Eflux = (units['density'] *
units['velocity']**3).to('erg/(kpc^2*yr)').value
self.to_gacc = ((units['pressure'] / units['length']) /
units['density']).cgs.value
self.to_Msun = units['mass'].to('Msun').value
self.to_erg = (units['pressure'] * units['length']**3).to('erg').value
def scatter_sp(sp,ax,axis=0,thickness=10,norm_factor=4., \
type='slice',kpc=True,runaway=True,active=True):
unit = set_units(muH=1.4271)
Msun = unit['mass'].to('Msun').value
Myr = unit['time'].to('Myr').value
#print len(sp)
if 'flag' in sp and active:
sp = sp[sp['flag'] > -2]
if len(sp) > 0:
runaways = (sp['mass'] == 0.0)
sp_runaway = sp[runaways]
sp_normal = sp[-runaways]
#print len(sp_runaway)
if len(sp_runaway) > 0 and runaway:
spx, spy, spz = projection(sp_runaway, axis)
spvx, spvy, spvz = projection_v(sp_runaway, axis)
if kpc:
spx = spx / 1.e3
spy = spy / 1.e3
if type == 'slice':
islab = np.where(abs(spz) < thickness)
#ax.scatter(spx.iloc[islab],spy.iloc[islab],marker='.',c='k',alpha=1.0)
#ax.quiver(spx.iloc[islab],spy.iloc[islab],
# spvx.iloc[islab],spvy.iloc[islab],color='k',alpha=1.0)
ax.scatter(spx,
spy,
marker='o',
color='k',
alpha=1.0,
s=10.0 / norm_factor)
#ax.quiver(spx,spy,spvx,spvy,color='w',alpha=0.5)
if len(sp_normal) > 0:
spx, spy, spz = projection(sp_normal, axis)
if kpc:
spx = spx / 1.e3
spy = spy / 1.e3
if type == 'slice': xbool = abs(spz) < thickness
spm = np.sqrt(sp_normal['mass'] * Msun) / norm_factor
spa = sp_normal['age'] * Myr
iyoung = np.where(spa < 40.)
#print len(iyoung[0])
if type == 'slice':
islab = np.where(xbool * (spa < 40))
ax.scatter(spx.iloc[islab],spy.iloc[islab],marker='o',\
s=spm.iloc[islab],c=spa.iloc[islab],\
vmax=40,vmin=0,cmap=plt.cm.cool_r,alpha=1.0)
ax.scatter(spx.iloc[iyoung],spy.iloc[iyoung],marker='o',\
s=spm.iloc[iyoung],c=spa.iloc[iyoung],\
vmax=40,vmin=0,cmap=plt.cm.cool_r,alpha=0.7)
import sys import pyathena as pa from scipy.interpolate import interp1d from matplotlib.ticker import MultipleLocator import copy from mpl_toolkits import axes_grid1 from matplotlib.colors import LogNorm from six.moves import cPickle as pickle # In[7]: import matplotlib as mpl # In[8]: unit = pa.set_units(muH=1.4271) print(unit) # print (unit['density'].cgs / 1.4271 / c.m_p.cgs, unit['velocity'], unit['length']) # print unit['density'] kb = 1.3806504 * 1e-16 # boltzmann constant erg/K / erg = g cm2/s2K vpc = 7168. * 1024 * 1024 / (128 * 128 * 896) # volume per cell z = np.arange(-3588, 3588, 8) g = 4.5181 * 1e-30 # gravitational constant : pc3/solmass*s2 sig_star = 42 # solmass/pc2 z_s = 245 # pc r_0 = 8000 # pc rho_dm = 0.0064 # solmass/pc3 km = 3.24078 * 1e-14 # 1km in parsec cm = 3.24078 * 1e-19 # 1cm in parsec gram = 5.02785 * 1e-34 # 1g in solar mass
import glob
import os
import matplotlib.colorbar as colorbar
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.colors import LogNorm, SymLogNorm, NoNorm, Normalize
from pyathena import read_starvtk, texteffect, set_units
import numpy as np
import string
from .scatter_sp import scatter_sp
import cPickle as pickle
unit = set_units(muH=1.4271)
to_Myr = unit['time'].to('Myr').value
def slice(slcfname,starfname,fields_to_draw,zoom=1.,aux={},\
writefile=True,tstamp=True,stars=True,field_label=True,norm_factor=2):
plt.rc('font', size=14)
plt.rc('xtick', labelsize=14)
plt.rc('ytick', labelsize=14)
slc_data = pickle.load(open(slcfname, 'rb'))
x0 = slc_data['yextent'][0]
y0 = slc_data['yextent'][3]
Lx = slc_data['yextent'][1] - slc_data['yextent'][0]
Lz = slc_data['yextent'][3] - slc_data['yextent'][2]
Lz = Lz / zoom
ix = 2
iz = ix * Lz / Lx
def get_zprof(data,domain,par,hst):
Omega=par['problem']['Omega']
qshear=par['problem']['qshear']
shearing_box,cooling,MHD,nscalar,Gamma = get_configure(par)
import pyathena as pa
x,y,z,=pa.cc_arr(domain)
Nx=domain['Nx']
x3d=np.tile(x.reshape(1,1,Nx[0]),(Nx[2],Nx[1],1))
z3d=np.tile(z.reshape(Nx[2],1,1),(1,Nx[1],Nx[0]))
if shearing_box:
vy0=-qshear*Omega*x3d
Phit=-qshear*Omega**2*x3d**2
else:
vy0=np.zeros(x3d.shape)
if cooling:
heat_ratio=np.interp(domain['time'],hst['time'],hst['heat_ratio'])
unit=pa.set_units(muH=1.4271)
unitC=unit['density']*unit['velocity']**3/unit['length']
unitC=unitC.cgs.value
coolftn=pa.coolftn()
dA=domain['dx'][0]*domain['dx'][1]
dz=domain['dx'][2]
d=np.copy(data['density'].data)
v1=np.copy(data['velocity1'].data)
v2=np.copy(data['velocity2'].data)
v3=np.copy(data['velocity3'].data)
dv2=v2-vy0
if MHD:
B1=np.copy(data['magnetic_field1'].data)
B2=np.copy(data['magnetic_field2'].data)
B3=np.copy(data['magnetic_field3'].data)
P=np.copy(data['pressure'].data)
Phi=np.copy(data['gravitational_potential'].data)
T1=np.copy(data['T1'].data)
if cooling:
lam=coolftn.get_cool(T1)
gam=coolftn.get_heat(T1)*heat_ratio
zprof={}
zprof['d']=d
zprof['v1']=v1
zprof['v2']=v2
zprof['v3']=v3
zprof['P']=P
if MHD:
zprof['B1']=B1
zprof['B2']=B2
zprof['B3']=B3
dsqrt=np.sqrt(d)
vdotB=B1*v1+B2*v2+B3*v3
Emag=0.5*(B1**2+B2**2+B3**2)
meanB={'1':B1.mean(axis=2).mean(axis=1),'2':B2.mean(axis=2).mean(axis=1),'3':B3.mean(axis=2).mean(axis=1)}
for vec in ['1','2','3']:
zprof['M%s' % vec] = zprof['d']*zprof['v%s' % vec]
zprof['Ek%s' % vec] = 0.5*zprof['d']*zprof['v%s' % vec]**2
if MHD:
zprof['PB%s' % vec] = 0.5*zprof['B%s' % vec]**2
zprof['vA%s' % vec] = zprof['B%s' % vec]/dsqrt
Bbar=np.tile(meanB[vec].reshape(Nx[2],1,1),(1,Nx[1],Nx[0])).astype('double')
zprof['dB%s' % vec] = zprof['B%s' % vec] - Bbar
zprof['dPB%s' % vec] = 0.5*zprof['dB%s' % vec]**2
zprof['dvA%s' % vec] = zprof['dB%s' % vec]/dsqrt
zprof['S%s' % vec] = 2.0*Emag*zprof['v%s' % vec] - zprof['B%s' % vec]*vdotB
if shearing_box:
zprof['v2']=dv2
zprof['dM2'] = zprof['d']*zprof['v2']
zprof['dEk2'] = 0.5*zprof['d']*zprof['v2']**2
# zprof['Phit'] = Phit
zprof['T']=coolftn.get_temp(T1)
if cooling:
zprof['cool']=d**2*lam/unitC
zprof['heat']=d*gam/unitC
for ns in range(nscalar):
zprof['s%s' % (ns +1)]=data['specific_scalar%s' % ns].data*d
param={}
param['R0']=par['problem']['R0']*c.pc
param['rho_dm']=par['problem']['rhodm']*c.M_sun/c.pc**3
param['surf_s']=par['problem']['SurfS']*c.M_sun/c.pc**2
param['z0']=par['problem']['zstar']*c.pc
phiext=tigress_extpot(param).phiext
Phie=phiext(z3d*c.pc).to('km**2/s**2').value
gext=phiext((z3d+dz/2)*c.pc)-phiext((z3d-dz/2)*c.pc)
zprof['Phie']=Phie
zprof['gext']=gext.to('km**2/s**2').value/dz
zprof['dWext']=d*zprof['gext']
zprof['Phisg']=Phi
dPhi=np.copy(Phi)
dPhi[1:-1,:,:]=(Phi[2:,:,:]-np.roll(Phi,2,axis=0)[2:,:,:])/2.0/dz
dPhi[0,:,:]=(Phi[1,:,:]-Phi[0,:,:])/dz
dPhi[-1,:,:]=(Phi[-1,:,:]-Phi[-2,:,:])/dz
zprof['gsg']=dPhi
zprof['dWsg']=d*zprof['gsg']
zprof['Ber'] = 0.5*(v1**2+v2**2+v3**2) + Gamma/(Gamma-1)*P/d+Phie+Phi+Phit
for pm in ['p','m']:
zprof[pm+'A'] = np.ones(d.shape)*dA
zprof[pm+'d'] = np.copy(d)
zprof[pm+'vz'] = np.copy(v3)
for f in ['d','M1','M2','M3']:
zprof['%sFz%s' % (pm,f)] = zprof[f]*v3
for f in ['E1','E2','E3','Ege','Egsg','Etidal']:
if f in ['E1','E2','E3']:
zf='%sk%s' %(f[0],f[1])
tmp = zprof[zf]*v3
elif f == 'Ege': tmp = d*Phie*v3
elif f == 'Egsg': tmp = d*Phi*v3
elif f == 'Etidal': tmp = d*Phit*v3
zprof['%sFz%s' % (pm,f)] = tmp
zprof['%sFzP' % pm] = Gamma/(Gamma-1)*zprof['P']*v3
if MHD:
zprof['%sSzEm1' % pm] = 2.0*zprof['PB1']*v3
zprof['%sSzEm2' % pm] = 2.0*zprof['PB2']*v3
zprof['%sSzvB1' % pm] = -B3*B1*v1
zprof['%sSzvB2' % pm] = -B3*B2*v2
for ns in range(nscalar):
zprof['%sFzs%s' % (pm,ns+1)]=data['specific_scalar%s' % ns].data*d*v3
if shearing_box:
zprof['Rxy']=d*v1*dv2
zprof['Mxy']=-B1*B2
nv3=zprof['v3']<0
pv3=~nv3
for k in zprof:
if k.startswith('p'): zprof[k][nv3]=0
if k.startswith('m'): zprof[k][pv3]=0
zprof['A']=np.ones(d.shape)*dA
return z,zprof
