In [ ]:
%matplotlib inline
import h5py
import numpy as np
import matplotlib.pyplot as plt
import pickle
path = '/home/ychen/d9/FLASH4/2015_production_runs/0529_L45_M10_b1_h1/'
read = pickle.load(open('%s_particles_leave_dict.pickle' % path.split('/')[-2], 'r'))

In [ ]:
#tadd, tag = 218857882571403.5, 127225.0
read.viewkeys()

In [ ]:
fname = 'MHD_Jet_hdf5_part_0600'
f = h5py.File(path+fname, 'r')
fr = h5py.File(path+fname+'_updated', 'w')

In [ ]:
for dataset in f.values():
    print dataset

In [ ]:
#print read
colname = [item[0].strip() for item in f['particle names']]
itag  = colname.index('tag')
itadd = colname.index('tadd')
iden0 = colname.index('den0')
ivelz = colname.index('velz')
itau  = colname.index('tau')
tp = f['tracer particles'].value

#tau1 = np.zeros(tp.shape[0])
dtau = np.zeros(tp.shape[0])
den1 = tp[:,iden0]

for i, (tag, tadd)  in enumerate(zip(tp[:,itag], tp[:,itadd])):
    #print tadd, tag, read[(tag, tadd)]
    try:
        den1[i] = read[(tag, tadd)][2]
        #if den1[i] < 0: print den1[i]
        tau1 = read[(tag, tadd)][3]
        dtau[i] = max(tp[i,itau]-tau1, 1E-100)
    except KeyError:
        print tag, tadd, '%e' % tp[i,ivelz], 'not in pickled data [dict]'
        

newcols = np.column_stack((den1,dtau))
print newcols.shape

In [ ]:
for v in f.values():
    if 'particle names' in v.name:
        shape = (v.shape[0]+2, 1)
        newfields = [['den1                    '], ['dtau                    ']]
        data = np.concatenate((v.value, newfields), axis=0)
        print shape
        print data
        fr.create_dataset(v.name, shape, v.dtype, data)
    elif 'tracer particles' in v.name:
        print v
        shape = (v.shape[0], v.shape[1]+2)
        print shape
        #newcol = np.expand_dims(v.value[:,0], axis=1)
        #newcols = v.value[:,0:2]
        data = np.concatenate((v.value, newcols), axis=1)
        fr.create_dataset(v.name, shape, v.dtype, data)
    else:
        fr.create_dataset(v.name, v.shape, v.dtype, v.value)
f.close()
fr.close()

In [ ]:
fr = h5py.File(path+fname+'_updated', 'r')
for v, vr in zip( f.values(), fr.values() ):
    print v
    print vr

In [ ]:
f = h5py.File(path+fname, 'r')
f['tracer particles'][:,itau]
plt.figure(figsize=(8,6))

xx = np.linspace(-8, -2, 10)
tau = f['tracer particles'][:,itau]
idens = colname.index('dens')
iden0 = colname.index('den0')
dens = f['tracer particles'][:,idens]
den0 = f['tracer particles'][:,iden0]
tadd = f['tracer particles'][:,itadd]
plt.scatter(np.log10(tau), np.log10(newcols[:,1]), c=tadd, vmin=2E14, vmax=3E14, lw=0, s=2)
plt.plot(xx,xx, c='r')
plt.colorbar()
plt.xlabel('tau')
plt.ylabel('dtau')
plt.xlim(-8,-2)
plt.ylim(-101,-2)

In [ ]:
read_array = np.array([[key[0] + key[1]] + value for key, value in read.iteritems()])

In [ ]:
import matplotlib.pyplot as plt
den1 = read_array[:,3]
tadd = read_array[:,1]
mask = np.where(den1<0)
print len(den1[mask])
#null = plt.hist(den1, bins=100, range=(-1E-26, 0))
plt.scatter(den1[mask], tadd[mask])
plt.xlim(-1E-26, 0)

In [ ]:
import pickle
dir = '/home/ychen/d9/FLASH4/2015_production_runs/0529_L45_M10_b1_h1/'
read = pickle.load(open('%s_particles_leave_dict.pickle' % dir.split('/')[-2], 'rb'))
print type(read)

In [ ]:


In [ ]: