In [ ]:
%pdb
%matplotlib inline  
import numpy as np
import matplotlib
matplotlib.rcParams['savefig.dpi'] = 150
import matplotlib.pyplot as plt
import yt
from yt.data_objects.particle_filters import add_particle_filter
from yt.fields.derived_field import ValidateGridType
from yt.fields.field_detector import FieldDetector

def JetP(pfilter, data):
    #filter = data[("all", "particle_shok")] == 0
    filter = np.logical_and((data[("io", "particle_shok")] == 0), 
                            (data[("io", "particle_gamc")] > 0.0))
    return filter

add_particle_filter("jetp", function=JetP, filtered_type='io', requires=["particle_shok"])

In [ ]:
fname = '/home/ychen/d9/FLASH4/stampede/0529_L45_M10_b1_h1/MHD_Jet_hdf5_plt_cnt_0620'
ds = yt.load(fname, particle_filename=fname.replace('plt_cnt', 'part'))
ad = ds.all_data()

In [ ]:
def _jet_volume_fraction(field, data):
    g = yt.YTQuantity(1.0, 'g')
    cm = yt.YTQuantity(1.0, 'cm')
    Kelvin = yt.YTQuantity(1.0, 'K')
    rhoCore = data.ds.parameters['sim_rhocore']*g/cm**3
    rCore   = data.ds.parameters['sim_rcore']*cm
    densitybeta = data.ds.parameters['sim_densitybeta']
    Tout    = data.ds.parameters['sim_tout']*Kelvin
    Tcore   = data.ds.parameters['sim_tcore']*Kelvin
    rCoreT  = data.ds.parameters['sim_rcoret']*cm
    gammaICM= data.ds.parameters['sim_gammaicm']
    mu      = data.ds.parameters['sim_mu']
    mp = yt.utilities.physical_constants.mass_hydrogen
    k  = yt.utilities.physical_constants.boltzmann_constant

    r = data['index', 'spherical_radius']
    
    density0 = rhoCore*(1.0 + (r/rCore)**2)**(-1.5*densitybeta)
    T0 = Tout*(1.0+(r/rCoreT)**3)/(Tout/Tcore+(r/rCoreT)**3)
    P0 = density0/mu/mp*k*T0
    icm_mass_fraction = 1.0 - data['flash', 'jet ']
    P = data['gas', 'pressure']
    density = data['gas', 'density']
    
    icm_volume_fraction = (P0/P)**(1/gammaICM)*icm_mass_fraction*density/density0
    
    icm_volume_fraction = np.where(icm_volume_fraction < 1.0, icm_volume_fraction, 1.0)
    return 1.0 - icm_volume_fraction


ds.add_field(('gas', 'jet_volume_fraction'), function=_jet_volume_fraction,
             display_name="Jet Volume Fraction", force_override=True)

In [ ]:
me = yt.utilities.physical_constants.mass_electron #9.109E-28 g
c  = yt.utilities.physical_constants.speed_of_light #2.998E10 cm/s
e  = yt.utilities.physical_constants.elementary_charge #4.803E-10 esu
nu = yt.YTQuantity(1.4, 'GHz')
gamma_min = yt.YTQuantity(10, 'dimensionless')

def _synchrotron_emissivity(field, data):
    ptype = 'io'
    # To convert from FLASH "none" unit to cgs unit, times the B field from FLASH by sqrt(4*pi)
    B = np.sqrt(data[(ptype, 'particle_magx')]**2+data[(ptype, 'particle_magy')]**2+data[(ptype, 'particle_magz')]**2)\
        *np.sqrt(4.0*np.pi)
    B = data.apply_units(B, 'G')
    nuc = 3.0*data[(ptype, 'particle_gamc')]**2*e*B/(4.0*np.pi*me*c)
    nu = data.get_field_parameter("frequency", default=yt.YTQuantity(1.4, 'GHz'))
    #P = data[(ptype, 'particle_pressure')]
    #P = data[(ptype, 'particle_dens')]*yt.YTQuantity(1, 'dyne/cm**2')
    fit_const = 5.8
    norm = 0.5*e**3.5/(c**2.5*me**1.5*(4.*np.pi)**0.5)
    N0 = 3.0/me/c/c/(np.log(np.abs(data[(ptype, 'particle_gamc')]/gamma_min)))

    return N0*norm*fit_const*nu**(-0.5)*np.exp(-nu/nuc)
#logging.getLogger('yt').setLevel(logging.DEBUG)

ds.add_field(('io', 'particle_emissivity'), function=_synchrotron_emissivity, particle_type=True, 
             display_name="%s Emissivity" % nu, units="auto", force_override=True)

ds.add_particle_filter('jetp')
#ds.add_field(('jetp', 'particle_emissivity'), function=_synchrotron_emissivity, particle_type=True,
              #units='erg/s/cm**3/Hz')

In [ ]:
ad = ds.all_data()
ptype = 'jetp'
deposit_field = 'particle_emissivity'
yt.YTArray([0], input_units=ad[ptype, deposit_field].units)

In [ ]:
def _deposit_average_filling(field, grid):
    #data._debug()
    ptype = 'jetp'
    deposit_field = 'particle_emissivity'
    P = grid['gas', 'pressure']
    B = grid['gas', 'magnetic_field_strength']
    uq = P.uq*(B.uq)**1.5*grid[ptype, deposit_field].uq
    if isinstance(grid, FieldDetector): return grid[ptype, deposit_field]*P*B**1.5
    if len(grid[ptype, deposit_field]) > 0: return P*B**1.5*grid[ptype, deposit_field].mean()
    elif grid.Level == 0: return grid['zeros']*uq
    else:
        pfield = yt.YTArray([0], input_units=ad[ptype, deposit_field].units)
        for gr in grid.Parent.Children:
            pfield = np.concatenate([pfield, gr[ptype, deposit_field]])
        if len(pfield) == 0: 
            return grid['zeros']*uq
        else:
            return P*B**1.5*pfield.mean()*grid[ptype, deposit_field].uq

ds.add_field(('deposit', 'avgfill_emissivity'), function=_deposit_average_filling, validators=[ValidateGridType()],
             display_name="%s Emissivity" % nu, units='erg/s/cm**3/Hz', take_log=True, force_override=True)

def _intensity(field, data):
    return data['deposit', 'avgfill_emissivity']/yt.YTQuantity(4.*np.pi, 'sr')

ds.add_field(('deposit', 'avgfill_intensity'), function=_intensity, display_name="%s Intensity" % nu,
             units='Jy/cm/arcsec**2', take_log=True, force_override=True)

In [ ]:
#negative_ad = ad.cut_region(["obj['jet_volume_fraction'] < 0.0"])
field = ('gas', 'jet_volume_fraction')
#field = ('flash', 'jet ')
slice = yt.SlicePlot(ds, 'x', field, center=(0,0,0))
#                    field_parameters={'frequency': nu})
#slice.set_zlim(('deposit', field), 1E-36/nu.in_units('GHz').value**0.5,
#               1E-32/nu.in_units('GHz').value**0.5)
slice.annotate_grids()
#slice.set_log(field, False)
#slice.set_cmap(field, 'gist_heat')
slice.zoom(10)

In [ ]:
field = 'avgfill_intensity'
proj = yt.ProjectionPlot(ds, 'x', ('deposit', field), center=(0,0,0),
                        field_parameters={'frequency': nu})
proj.set_zlim(('deposit', field), 1E-3/nu.in_units('GHz').value**0.5, 
              1E1/nu.in_units('GHz').value**0.5)
proj.zoom(10)
proj.show()

In [ ]:
proj.set_zlim(('deposit', 'avgfill_intensity'), 1E-3, 1E1)
proj.annotate_timestamp(corner='upper_left', time_format="{time:6.3f} {units}", time_unit='Myr', draw_inset_box=True)
proj.annotate_text((0.9, 0.95), 'hinf', coord_system='axis', text_args={'color':'k'})
proj.show()
#proj.save('/home/ychen/d9/FLASH4/stampede/0529_L45_M10_b1_h1/0620_avgfill_deposited_intensity.png')

In [ ]:
ds.add_deposited_particle_field(('jetp', 'particle_emissivity'), 'nearest')
print ad['jetp', 'particle_emissivity'].units
print ad['deposit', 'jetp_nn_emissivity'].units
def _nn_intensity(field, data):
    '''
    Intensity per length using nearest neighbot. Integrate over line of sight to get intensity.
    '''
    P = data['gas', 'pressure']
    B = data['gas', 'magnetic_field_strength']
    print data['jetp', 'particle_emissivity'].units
    print data['deposit', 'jetp_nn_emissivity'].units
    return P*B**1.5*data['deposit', 'jetp_nn_emissivity']/yt.YTQuantity(4.*np.pi, 'sr')

f5 = ds.add_field(('deposit', 'nn_intensity'), function=_nn_intensity, display_name='%s NN Intensity' % nu,
                  units='Jy/cm/arcsec**2', take_log=True, force_override=True)

In [ ]:
fn = 'emissivity_nn'
slice = yt.SlicePlot(ds, 'x', fn, center=(0,0,0))
slice.set_zlim(fn, 1E-36, 1E-32)
slice.annotate_grids()
slice.zoom(10)

In [ ]:
den_ratio = ad['jetp', 'particle_dens']/ad['jetp', 'particle_den0']
print den_ratio

In [ ]:
print den_ratio.max()
print den_ratio.min()
print den_ratio.mean()
print len(den_ratio)

In [ ]:
plt.hist(den_ratio, bins=200, range=(0,2))
plt.show()

In [ ]:
ndens = ad['deposit', 'jetp_count']/ad['index', 'cell_volume']

In [ ]:
densjet = ad['gas', 'density']*ad['flash', 'jet ']
color = {8:'r', 7:'orange', 6:'magenta', 5:'g', 4:'b', 3:'cyan'}
for ilevel in [8,7,6,5,4,3]:
    levelfilter = ad['index', 'grid_level'] == ilevel
    filter = ad['deposit', 'jetp_count'][levelfilter] > 0
    #print len(ndens)
    print len(ndens[levelfilter][filter])
    #print len(densjet[filter])
    plt.scatter(np.log10(ndens[levelfilter][filter]), 
                np.log10(densjet[levelfilter][filter]), s=1, lw=0, c=color[ilevel])

In [ ]:
plt.hist([ad['deposit', 'io_count'], ad['jetp', 'particle_dens']], bins=50,
        range=(0, 3E-26))

plt.show()

In [ ]:
L = [1,1,0] # vector normal to cutting plane
north_vector = [-1,1,0]
fn = 'emissivity_nn'
proj = yt.ProjectionPlot(ds, 'x', ('deposit', field), center=(0,0,0),
                        field_parameters={'frequency': nu})
proj.set_zlim(('deposit', field), 1E-3/nu.in_units('GHz').value**0.5, 
              1E1/nu.in_units('GHz').value**0.5)
proj.zoom(10)
proj.show()

In [ ]:
ad.set_field_parameter?

In [ ]: