In [ ]:
%pdb
import numpy as np
import util
import multiprocessing
import matplotlib
matplotlib.rcParams['savefig.dpi'] = 150
import matplotlib.pyplot as plt
import yt
import logging
logging.getLogger('yt').setLevel(logging.WARNING)
from yt.utilities.logger import ytLogger as mylog
print yt.__version__
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
def xcenter_slice(pfilter, data):
    filter = abs(data[("io", 'particle_posx')]) < 3.08567758E21
    return filter
add_particle_filter("jetp", function=JetP, filtered_type='io', requires=["particle_shok"])
add_particle_filter("xcenter", function=xcenter_slice, filtered_type='io', requires=["particle_posx"])

In [ ]:
fname = '/home/ychen/d9/FLASH4/2015_production_runs/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 _particle_pressure(field, data):
#    required = data['gas', 'pressure']
#    mylog.debug("Mapping grid data to particles for %s in %s" % (field, data))
#    return data.ds.find_field_values_at_points(('gas', 'pressure'), data[ptype, 'particle_position'])
#logging.getLogger('yt').setLevel(logging.INFO)
#ds.add_field((ptype, 'particle_pressure'), function=_particle_pressure, particle_type=True, units='g/(cm*s**2)')

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')
#nu = 1.5E8 # 150MHz

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*B**1.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)
nu = yt.YTQuantity(1.4, 'GHz') # 1.4GHz
ds.add_field(('io', 'particle_emissivity'), function=_synchrotron_emissivity, particle_type=True, 
             display_name="%s Emissivity" % nu, 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 [ ]:
logging.getLogger('yt').setLevel(logging.INFO)
ad['jetp', 'particle_emissivity']
#ds.find_field_values_at_points('pres', [[1E23,0.,0.],[1E22,0.,0.]])

In [ ]:
logging.getLogger('yt').setLevel(logging.DEBUG)
p = None
pemis = None
def _deposit_average_filling(field, grid, method=1):
    #data._debug()
    ptype = 'jetp'
    deposit_field = 'particle_emissivity'
    uq = grid['gas', 'pressure'].uq
    if isinstance(grid, FieldDetector): return grid[ptype, deposit_field]*uq
    if len(grid[ptype, deposit_field]) > 0: return grid['gas', 'pressure']*grid[ptype, deposit_field].mean()
    #else: return data['zeros']*uq
    elif grid.Level == 0: return grid['zeros']*uq
    else:
        if method == 1:
            pfield = np.zeros(0)
            grids_to_go = grid.Parent.Children
            #for i, gr in enumerate(grids_to_go):
            #    if len(gr.Children) > 0:
            #        del grids_to_go[i]
            #        grids_to_go.extend(gr.Children)
            for gr in grids_to_go:
                pfield = np.concatenate([pfield, gr[ptype, deposit_field]])
            if len(pfield) == 0: 
                return grid['zeros']*uq
            else:
                return grid['gas', 'pressure']*pfield.mean()        
        if method == 2:
            global p, pemis
            if p is None or pemis is None:
                import time
                t0 = time.time()
                p = ds.all_data()['jetp', 'particle_position']
                mylog.debug('Load particle position time:', time.time() - t0)
                t0 = time.time()
                pemis = ds.all_data()['jetp', 'particle_emissivity']
                mylog.debug('Calculate particle emissivity time:', time.time() - t0)
            l = grid.LeftEdge
            r = grid.RightEdge
            c = (l+r)/2.
            i = 1
            filter = p[:,0]>l[0] 
            for i in range(3):
                filter *= (p[:,i]>l[i])*(p[:,i]<r[i])
            #reg = grid.ds.region(c, c+(l-c)*fac, c+(r-c)*fac)
            while len(pemis[filter]) == 0 and i < grid.Level:
                i += 1
                l = c+(l-c)*2.
                r = c+(r-c)*2.
                filter = p[:,0]>l[0] 
                for i in range(3):
                    filter *= (p[:,i]>l[i])*(p[:,i]<r[i])
                #reg = grid.ds.region(c, c+(l-c)*fac, c+(r-c)*fac)
            return grid['gas', 'pressure']*pemis[filter].mean()



    #pdata = data['all', 'particle_gamc']
    #grid_id = np.unique(data['index', 'grid_indices'])[0]
    #if len(np.unique(data['index', 'grid_indices'])) > 1:
    #    raise RuntimError
    #grid_id = data.id
    #data._debug()
    #l = data.Level
    #while grid_id not in GridParticles and l > data.ds.index.max_level:
        #id_offset = data._id_offset
        #grid_id = data.Parent.id
        #l-=1
    #if grid_id not in GridParticles: return data['zeros']
    
    #return data['zeros'] + pdata[GridParticles[grid_id]].mean()

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)

In [ ]:
logging.getLogger('yt').setLevel(logging.INFO)
slice = yt.SlicePlot(ds, 'x', ('deposit', 'avgfill_emissivity'), center=(0,0,0))
slice.set_zlim(('deposit', 'avgfill_emissivity'), 1E-36, 1E-32)
slice.annotate_grids()
slice.zoom(10)

In [ ]:
logging.getLogger('yt').setLevel(logging.INFO)
proj = yt.ProjectionPlot(ds, 'x', ('deposit', 'avgfill_emissivity'), center=(0,0,0))
proj.set_zlim(('deposit', 'avgfill_emissivity'), 1E-36, 1E-32)
#proj.annotate_grids()
proj.zoom(10)
proj.show()

In [ ]:
#logging.getLogger('yt').setLevel(logging.INFO)
#proj = yt.ProjectionPlot(ds, 'x', ('deposit', 'emissivity_nn'), center=(0,0,0))
#proj.set_zlim(('deposit', 'avgfill_emissivity'), 1E-36, 1E-32)
proj.set_zlim(('deposit', 'emissivity_nn'), 1E-14, 1E-10)
#proj.annotate_grids()
#proj.zoom(8)
proj.show()

In [ ]:
pos = yt.YTArray([0.0, -1.2, 9.0])*3.0857E21
grids = ds.index.grids
for grid in grids:
    if grid['jetp', 'particle_emissivity'].mean() == 0:
        print grid, grid.LeftEdge.in_units('kpc')

In [ ]:
ds.add_deposited_particle_field(('io', 'particle_emissivity'), 'nearest')
def _synchrotron_emissivity_pres(field, data):
    return data['deposit', 'io_nn_emissivity']*data['gas', 'pressure']

fn = ds.add_field(('deposit', 'emissivity_nn'), function=_synchrotron_emissivity_pres,
             display_name='Emissivity', units='erg/s/cm**3/Hz', take_log=True, force_override=True)

In [ ]:
fn = 'emissivity_nn'
proj = yt.ProjectionPlot(ds, 'x', fn, center=(0,0,0))
proj.set_zlim(fn, 1E-14, 1E-10)
#proj.annotate_grids()
proj.zoom(10)

In [ ]:
proj.save('/home/ychen/d9/FLASH4/stampede/0529_L45_M10_b1_h1/0620_nn_deposited_emissivity.png')

In [ ]:
def map_particles_to_grid(ds):
    ad = ds.all_data()
    pos = ad['particle_position']
    grids, indices = ds.index._find_points(pos[:,0], pos[:,1], pos[:,2])
    GridParticles = {}
    for ind, grid in enumerate(grids):
        if grid.id not in GridParticles:
            GridParticles[grid.id] = []
        GridParticles[grid.id].append(ind)
        
    # Fill particles for parent blocks
    for level in reversed(range(ds.index.max_level)):
        for grid.id in ds.index.select_grids(level):
            if grid.id not in GridParticles and len(grid.Children) > 0:
                #print grid.Children
                GridParticles[grid.id] = []
                for cgr in grid.Children:
                    if cgr in GridParticles:
                        #print 'ChildrenParticles:', GridParticles[cgr]
                        GridParticles[grid].extend(GridParticles[cgr.id])
                #print 'AllParticles:', GridParticles[grid]
                if len(GridParticles[grid.id]) == 0:
                    del GridParticles[grid.id]
    return GridParticles

In [ ]:
import time
t0 = time.time()
ad = ds.all_data()
pos = ad['particle_position']
print 'Load particle position time:', time.time() - t0
t0 = time.time()
GridParticles = map_particles_to_grid(ds)
print 'Mapping time:', time.time() - t0
print len(GridParticles)
print 'number of particles:', len(pos)

In [ ]:
#@yt.derived_field(name=('deposit', 'average_filling'), validators=[ValidateGridType()])
def _deposit_average_filling(field, data):
    if isinstance(data, FieldDetector): return data['zeros']
    if len(data['all', 'particle_gamc']) > 0: return data['zeros'] + data['all', 'particle_gamc'].mean()
    elif data.Level == 0: return data['zeros']
    else:
        mean = np.concatenate([ngrid['all', 'particle_gamc'] for ngrid in data.Parent.Children]).mean()
        return data['zeros'] + mean
    #pdata = data['all', 'particle_gamc']
    #grid_id = np.unique(data['index', 'grid_indices'])[0]
    #if len(np.unique(data['index', 'grid_indices'])) > 1:
    #    raise RuntimError
    #grid_id = data.id
    #data._debug()
    #l = data.Level
    #while grid_id not in GridParticles and l > data.ds.index.max_level:
        #id_offset = data._id_offset
        #grid_id = data.Parent.id
        #l-=1
    #if grid_id not in GridParticles: return data['zeros']
    
    #return data['zeros'] + pdata[GridParticles[grid_id]].mean()

ds.add_field(('deposit', 'average_filling'), function=_deposit_average_filling, validators=[ValidateGridType()])

In [ ]:
ad['deposit', 'average_filling']

In [ ]:


In [ ]:
slice.annotate_grids(alpha=0.1, draw_ids=True)

slice.show()

In [ ]:
# Fill particles for parent blocks
for level in reversed(range(ds.index.max_level)):
    nPart = 0
    for grid in ds.index.select_grids(level):
        if grid not in GridParticles and len(grid.Children) > 0:
            #print grid.Children
            GridParticles[grid] = []
            for cgr in grid.Children:
                if cgr in GridParticles:
                    #print 'ChildrenParticles:', GridParticles[cgr]
                    GridParticles[grid].extend(GridParticles[cgr])
                    nPart += len(GridParticles[cgr])
            #print 'AllParticles:', GridParticles[grid]
            if len(GridParticles[grid]) == 0:
                del GridParticles[grid]
    print 'level ', level, ', number of particles:', nPart
print len(GridParticles)

In [ ]:
nPart = 0
for grid in ds.index.select_grids(0):
    if grid in GridParticles:
        #print len(GridParticles[grid])
        nPart += len(GridParticles[grid])
print nPart

In [ ]:


In [ ]:
max = 0
maxg = None
maxp = None
for grid, particles in GridParticles.iteritems():
    if len(particles) > max:
        max = len(particles)
        maxg = grid
        maxp = particles
print max, maxg, maxg.Level, maxg.LeftEdge

In [ ]:
for grid in ds.index.grids:
    if grid in GridParticles:
        grid[('gas','test')].fill(np.average(ad['particle_gamc'][GridParticles[grid]]))
    
    #elif grid.Parent is not None and grid.Parent in GridParticles:
    #    grid[('gas','test')].fill(np.average(ad['particle_gamc'][GridParticles[grid.Parent]]))
    #    print np.average(ad['particle_gamc'][GridParticles[grid]])
    #    print len(ad['particle_gamc'][GridParticles[grid]])

In [ ]:
for grid in ds.index.select_grids(0):
    print grid['gas','test']
    break
    if sum(grid['gas','test'])>0:
        print grid['gas','test']

In [ ]:
ad['gas', 'test'].max()

In [ ]:
for grid, particles in GridParticles.iteritems():
    if len(particles) != grid.NumberOfParticles:
        print grid, len(particles), grid.NumberOfParticles

In [ ]:
plt.hist(ad['jetp', 'particle_dens'])

In [ ]: