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