In [ ]:

    

%matplotlib inline
import h5py
import numpy as np
import matplotlib.pyplot as plt
import os
import yt
import logging
logging.getLogger('yt').setLevel(logging.ERROR)

from synchrotron.yt_synchrotron_emissivity import *
from yt.utilities.file_handler import HDF5FileHandler
from yt.funcs import mylog

def setup_part_file(ds):
    filename = os.path.join(ds.directory,ds.basename)
    ds._particle_handle = HDF5FileHandler(filename.replace('plt_cnt', 'part')+'_updated')
    ds.particle_filename = filename.replace('plt_cnt', 'part')+'_updated'
    mylog.info('Changed particle files to:' + ds.particle_filename)
    
dir = '/home/ychen/d9/FLASH4/2015_production_runs/0529_L45_M10_b1_h1/'
fname = dir + '/MHD_Jet_hdf5_plt_cnt_0630'

f = h5py.File(fname, 'r')

ds = yt.load(fname)
setup_part_file(ds)


    

In [ ]:

    

def prep_field_data(ds, field):
    #print(field)
    offset = 1
    data = np.zeros([ds.index.num_grids, *ds.index.grid_dimensions[0]], dtype='float32')
    for g in ds.index.grids:
        #print(g)
        if np.nan in g[field].v or np.inf in g[field].v:
            print(g[field].v)
        data[g.id-offset] = np.nan_to_num(g[field].v.transpose())
    #print(data[g.id-offset])
    return data


    

In [ ]:

    

os.path.exists
for k, v in f.items():
    print('{0:30s},{2},{1}'.format(k, v.shape, v.dtype))


    

In [ ]:

    

for name, value in f.items():
    print('%26s' % name, value)


    

In [ ]:

    

field_list = [field.decode() for field in f['unknown names'].value[:,0]]
print(field_list)


    

In [ ]:

    

'dens' in field_list


    

In [ ]:

    

for name, v in f.items():
    #print(name)
    if name not in field_list:
        print(v)


    

In [ ]:

    

from numpy.testing import assert_array_equal

def check_grids():
    for grid in ds.index.grids:
        h5id = grid.id-1
        #print(f['gid']).value[h5id]
        assert_array_equal(f['dens'].value[h5id], grid['dens'].v.transpose())

%timeit check_grids()


    

In [ ]:

    

ptype = 'lobe'
proj_axis = [1,0,1]
nu = (150, 'MHz')
#pars = add_synchrotron_emissivity(ds, ptype=ptype, nu=nu)
pars = add_synchrotron_dtau_emissivity(ds, ptype=ptype, nu=nu, proj_axis=proj_axis)
fields = []
for pol in ['i', 'q', 'u']:
    field = ('deposit', ('nn_emissivity_%s_%s_%%.1f%%s' % (pol, ptype)) % nu)
    fields.append(field)


    

In [ ]:

    

g = ds.index.grids[0]
print(field)
g[field]


    

In [ ]:

    

orig_field_list = [field.decode() for field in f['unknown names'].value[:,0]]
print(orig_field_list)

write_fields = np.array([f for ptype, f in fields])
#write_fields = np.array(['density'])
print(write_fields.shape)
print(write_fields)


    

In [ ]:

    

if type(proj_axis) is str:
    postfix = ('_synchrotron_%%.1f%%s_%s' % proj_axis) % nu
elif type(proj_axis) is list:
    postfix = ('_synchrotron_%%.1f%%s_%.1f_%.1f_%.1f' % tuple(proj_axis)) % nu

h5fw = h5py.File(f.filename+postfix, 'w')
print('Writing to ', f.filename+postfix)
for name, v in f.items():
    #print(name)
    if name in orig_field_list:
        #print('Skip %s' % name)
        pass
    elif name == 'unknown names':
        bnames = [f.encode('utf8') for f in write_fields]
        h5fw.create_dataset('unknown names', data=bnames)
    else:
        h5fw.create_dataset(v.name, v.shape, v.dtype, v.value)
for field in write_fields:
    data = prep_field_data(ds, field)
    fieldname = field[1] if type(field) is tuple else field
    print('writing field: %s' % fieldname)
    h5fw.create_dataset(fieldname, data.shape, data.dtype, data)
h5fw.close()


    

In [ ]:

    

h5fw.close()


    

In [ ]:

    

g = ds.index.grids[0]


    

In [ ]:

    

h5fr = h5py.File(f.filename+postfix, 'r')
for name, value in h5fr.items():
    print('%26s' % name, value)
h5fr.close()


    

In [ ]:

    

ds2 = yt.load(f.filename+postfix)
ds2.field_list


    

In [ ]:

    

from numpy.testing import assert_array_equal

assert_array_equal(ds2.all_data()['density'].v, ds.all_data()['dens'].v)


    

In [ ]:

    

yt.ProjectionPlot(ds2, 'x', ('flash','nn_emissivity_i_lobe_150.0MHz')).zoom(16)


    

In [ ]: