def argmaxloc(data, axis=None, out=None):
d = data.meta2dict()
ax = d['axes'].pop(axis)
out = data.argmax(axis=axis, out=out)
out = ax[out]
out = osh5def.H5Data(out, **d)
return out
def sl(*args, **kwargs):
# search for all specified input arguments
saved, kl = [], []
if isinstance(mapping[0], tuple): # multiple save/load
kl = sorted(list(mapping), key=lambda a: a[0]) # sort by input param. pos.
else:
kl = [mapping]
if len(args) < kl[-1][0]: # TODO(1) they can be in the kwargs
raise Exception('Cannot find the ' + str(kl[-1][0]) + '-th argument for meta data saving')
# save meta data into a list
for tp in kl:
if hasattr(args[tp[0]], 'meta2dict'):
saved.insert(0, (args[tp[0]].meta2dict(), tp[1]))
if not iter(saved):
raise ValueError('Illegal mapping parameters')
# execute user function
out = func(*args, **kwargs)
# load saved meta data into specified positions
ol, tl = [], out if isinstance(out, tuple) else (out, )
try:
for tp in saved:
if isinstance(tl[tp[1]], osh5def.H5Data):
ol.append(osh5def.H5Data.__dict__.update(tp[0]))
else:
aaa = tl[tp[1]]
ol.append(osh5def.H5Data(aaa, **tp[0]))
except IndexError:
raise IndexError('Output does not have ' + str(tp[1]) + ' elements')
except:
raise TypeError('Output[' + str(tp[1]) + '] is not/cannot convert to H5Data')
tmp = tuple(ol) if len(ol) > 1 else ol[0] if ol else out
return tmp
def spectrogram(h5data, axis=-1, **kwargs):
"""
A wrapper function of scipy.signal.spectrogram
:param h5data:
:param kwargs:
:param axis:
:return: h5data with one more dimension appended
"""
meta = h5data.meta2dict()
k, x, Sxx = signal.spectrogram(h5data.values, fs=1/h5data.axes[axis].increment, axis=axis, **kwargs)
meta['axes'][axis].ax = x - x[0] + h5data.axes[axis].min
meta['axes'].insert(axis, osh5def.DataAxis(attrs=copy.deepcopy(meta['axes'][axis].attrs), data=2*np.pi*k))
__update_axes_label(meta['axes'], axis - 1 if axis < 0 else axis)
return osh5def.H5Data(Sxx, **meta)
def sl(*args, **kwargs):
# save meta data into a list
saved = args[0].meta2dict() if hasattr(args[0], 'meta2dict') else None
# execute user function
out = func(args[0].view(np.ndarray), *args[1:], **kwargs)
# load saved meta data into specified positions
try:
if isinstance(out, osh5def.H5Data):
out.__dict__.update(saved)
else:
out = osh5def.H5Data(out, **saved)
except:
raise TypeError('Output is not/cannot convert to H5Data')
# Update unit if necessary
if unit is not None:
out.data_attrs['UNITS'] = osh5def.OSUnits(unit)
if axes is not None:
out.axes = axes
return out
def os_rwigner(real_array,xaxis):
nx=real_array.shape
nxh=(nx[0]+1)/2
dx=(xaxis.max-xaxis.min)/nx
#
# the wigner transform is not FFT because it is roughly the FFT of the correlation
# so the k-range is a little different than those of the FFT,
# hence the 0.5
#
kmax=0.5*(np.pi)/dx
dk=(0.5*np.pi)/xaxis.max
w_data= rwigner(real_array)
kaxis=osh5def.DataAxis(0,kmax,nxh,attrs={'NAME':'k','LONG_NAME':'wave number', 'UNITS': '\omega_0/c' })
data_attrs={'UNITS': '[a.u.]','NAME': '$W_{\phi}$', 'LONG_NAME': 'Wigner Transform' }
run_attrs = {'XMAX' : np.array( [xaxis.max, kmax] ) , 'XMIN' : np.array( [ xaxis.min, 0 ] ) }
os5data_wigner=osh5def.H5Data(w_data,timestamp='x', data_attrs=data_attrs, axes=[xaxis,kaxis])
return os5data_wigner
def stack(arr, axis=0, axesdata=None):
"""Similar to numpy.stack. Arr is the list of H5Data to be stacked,
or a filename filter (such as './FLD/s1-line/s1-line-x2-01*.h5') can be read into H5Data.
By default the newly created dimension will be labeled as time axis.
Other meta data will be copied from the last element of arr
"""
try:
if not isinstance(arr[-1], osh5def.H5Data):
raise TypeError(
'Input be list of H5Data or filenames of Osiris data (such as "./FLD/e1/*.h5")'
)
except (
TypeError, IndexError
): # not an array or an empty array, just return what ever passed in
return arr
md = arr[-1]
ax = copy.deepcopy(md.axes)
if axesdata:
if axesdata.size != len(arr):
raise ValueError(
'Number of points in axesdata is different from the new dimension to be created'
)
ax.insert(axis, axesdata)
else: # we assume the new dimension is time
taxis_attrs = {
'UNITS': "\omega_p^{-1}",
'LONG_NAME': "time",
'NAME': "t"
}
ax.insert(
axis,
osh5def.DataAxis(arr[0].run_attrs['TIME'],
arr[-1].run_attrs['TIME'],
len(arr),
attrs=taxis_attrs))
r = np.stack(arr, axis=axis)
return osh5def.H5Data(r,
md.timestamp,
md.data_attrs,
md.run_attrs,
axes=ax)
el_inv = np.abs(np.fft.fft(el,axis=0))
et_inv = np.abs(np.fft.fft(et,axis=0))
# temp = np.abs(np.fft.fft(e2_data,axis=0))
total_es[file_number,:,:]=temp[0:nk,nx_begin:nx_end]
total_em[file_number,:,:]=temp[0:nk,nx_begin:nx_end]
# temp = np.sum(s1_data, axis=0) / nx
# sum up the results to node 0 and output, 2 datasets, one for +
# component abd one for the - component
# first let's do the + root
comm.Reduce(e2_plus_output, total, op=MPI.SUM, root=0)
if rank == 0:
b=osh5def.H5Data(total, timestamp='x', data_attrs=data_attrs_eplus,
run_attrs=run_attrs, axes=[taxis,xaxis])
outFilename=outDir+'/'+'hfhi-es-data.h5'
osh5io.write_h5(b,filename=outFilename)
# now let's do the - component
comm.Reduce(total2, total, op=MPI.SUM, root=0)
if rank == 0:
b=osh5def.H5Data(total2, timestamp='x', data_attrs=data_attrs_eminus,
run_attrs=run_attrs, axes=[taxis,xaxis])
outFilename=outDir+'/'+'e2-minus.h5'
osh5io.write_h5(b,filename=outFilename)
# write_hdf(h5_output, outFilename)
print('Before barrier'+repr(rank))
comm.barrier()
s1_data = e2_data * b3_data - e3_data * b2_data
print(s1_data.shape)
h5_output[file_number, 1:nx] = s1_data[1:nx]
# temp = np.sum(s1_data, axis=0) / nx
# h5_output.data[file_number, 1:ny] = temp[1:ny]
# temp = np.sum(s1_data[0:n_avg, :], axis=0) / n_avg
# income[file_number, 1:ny] = temp[1:ny]
# temp = np.sum(s1_data, axis=1) / ny
# h5_output2.data[file_number, 1:nx] = temp[1:nx]
# temp = np.sum(s1_data[:, 0:n_avg], axis=1) / n_avg
# income2[file_number, 1:nx] = temp[1:nx]
# file_number+=1
comm.Reduce(h5_output, total, op=MPI.SUM, root=0)
if rank == 0:
b=osh5def.H5Data(total, timestamp='x', data_attrs=data_attrs,
run_attrs=run_attrs, axes=[xaxis,taxis])
osh5io.write_h5(b,filename=outFilename)
# write_hdf(h5_output, outFilename)
print('Before barrier'+repr(rank))
comm.barrier()
# comm.Reduce(income, total, op=MPI.SUM, root=0)
# if rank == 0:
# h5_output.data = total
# newName = outFilename.rsplit('.', 1)[0] + '-x-n' + str(n_avg) + '.h5'
# write_hdf(h5_output, newName)
# comm.barrier()
# comm.Reduce(h5_output2.data, total2, op=MPI.SUM, root=0)
# if rank == 0:
# h5_output2.data = total2
# write_hdf(h5_output2, outFilename)
# comm.barrier()
run_attrs = {
'XMAX': np.array([time_step * (total_time - 1), xaxis.max]),
'XMIN': np.array([0, xaxis.min, 0])
}
file_number = 0
for file_number in range(i_begin, i_end):
h5_filename = filelist[file_number]
if (file_number % 10 == 0):
print(h5_filename)
try:
h5_data_old = h5_data
h5_data = osh5io.read_h5(h5_filename)
except OSError:
h5_data = h5_data_old
temp = np.average((h5_data.data), axis=0)
h5_output[file_number, 1:ny] = temp[1:ny]
# file_number+=1
print('before write')
print(outfilename)
b = osh5def.H5Data(h5_output,
timestamp='x',
data_attrs=data_attrs,
run_attrs=run_attrs,
axes=[taxis, xaxis])
osh5io.write_h5(b, filename=outfilename)
print('after write')
print('Before barrier' + repr(rank))
def q3d_to_3d(rundir, plasma_field, fileno, mode_max, x1_min, x1_max, nx1,
x2_min, x2_max, nx2, x3_min, x3_max, nx3):
from scipy import interpolate
dx1 = (x1_max - x1_min) / (nx1 - 1)
dx2 = (x2_max - x2_min) / (nx2 - 1)
dx3 = (x3_max - x3_min) / (nx3 - 1)
x1_axis = np.arange(x1_min, x1_max + dx1, dx1)
x2_axis = np.arange(x2_min, x2_max + dx2, dx2)
x3_axis = np.arange(x3_min, x3_max + dx3, dx3)
a = np.zeros((nx1, nx2, nx3), dtype=float)
filename_out = filename_3d(rundir, plasma_field, fileno)
x1 = osh5def.DataAxis(x1_min,
x1_max,
nx1,
attrs={
'NAME': 'x1',
'LONG_NAME': 'x_1',
'UNITS': 'c / \omega_0'
})
x2 = osh5def.DataAxis(x2_min,
x2_max,
nx2,
attrs={
'NAME': 'x2',
'LONG_NAME': 'x_2',
'UNITS': 'c / \omega_0'
})
x3 = osh5def.DataAxis(x3_min,
x3_max,
nx3,
attrs={
'NAME': 'x3',
'LONG_NAME': 'x_3',
'UNITS': 'c / \omega_0'
})
# More attributes associated with the data/simulation. Again no need to worry about the details.
data_attrs = {
'UNITS': osh5def.OSUnits('m_e c \omega_0 / e'),
'NAME': plasma_field,
'LONG_NAME': plasma_field
}
run_attrs = {
'NOTE': 'parameters about this simulation are stored here',
'TIME UNITS': '1/\omega_0',
'XMAX': np.array([1., 15.]),
'XMIN': np.array([0., 10.])
}
# Now "wrap" the numpy array into osh5def.H5Data. Note that the data and the axes are consistent and are in fortran ordering
b = osh5def.H5Data(a,
timestamp='123456',
data_attrs=data_attrs,
run_attrs=run_attrs,
axes=[x1, x2, x3])
# I am doing mode 0 outside of the loop
fname_re = filename_re(rundir, plasma_field, 0, fileno)
# DEBUG
print(fname_re)
# DEBUG
data_re = osh5io.read_h5(fname_re)
print(data_re.shape)
print(data_re.axes[1].ax.shape)
print(data_re.axes[0].ax.shape)
func_re = interpolate.interp2d(data_re.axes[1].ax,
data_re.axes[0].ax,
data_re,
kind='cubic')
for i1 in range(0, nx1):
for i2 in range(0, nx2):
for i3 in range(0, nx3):
z = x1_axis[i1]
x = x3_axis[i3]
y = x2_axis[i2]
r = np.sqrt(x * x + y * y)
# if r != 0:
# cos_th = x/r
# sin_th = y/r
# else:
# cos_th = 1
# sin_th = 0
a[i1, i2, i3] = a[i1, i2, i3] + func_re(z, r)
for i_mode in range(1, mode_max + 1):
fname_re = filename_re(rundir, plasma_field, i_mode, fileno)
fname_im = filename_im(rundir, plasma_field, i_mode, fileno)
# DEBUG
print(fname_re)
# DEBUG
if (plasma_field == 'e2' or plasma_field == 'e3'):
if (plasma_field == 'e2'):
field_comp = 'e3'
else:
field_comp = 'e2'
data_re_self = osh5io.read_h5(
filename_re(rundir, plasma_field, i_mode, fileno))
data_im_self = osh5io.read_h5(
filename_im(rundir, plasma_field, i_mode, fileno))
data_re_comp = osh5io.read_h5(
filename_re(rundir, field_comp, i_mode, fileno))
data_im_comp = osh5io.read_h5(
filename_im(rundir, field_comp, i_mode, fileno))
else:
data_re = osh5io.read_h5(
filename_re(rundir, plasma_field, i_mode, fileno))
data_im = osh5io.read_h5(
filename_im(rundir, plasma_field, i_mode, fileno))
func_re = interpolate.interp2d(data_re.axes[1].ax,
data_re.axes[0].ax,
data_re,
kind='cubic')
func_im = interpolate.interp2d(data_im.axes[1].ax,
data_im.axes[0].ax,
data_im,
kind='cubic')
for i1 in range(0, nx1):
for i2 in range(0, nx2):
for i3 in range(0, nx3):
z = x1_axis[i1]
x = x3_axis[i3]
y = x2_axis[i2]
r = np.sqrt(x * x + y * y)
if r > 0.000001:
cos_th = x / r
sin_th = y / r
else:
cos_th = 1
sin_th = 0
# start the recursion relation to evaluate cos(n*theta) and sin(n_theta)
sin_n = sin_th
cos_n = cos_th
for int_mode in range(2, i_mode + 1):
temp_s = sin_n
temp_c = cos_n
cos_n = temp_c * cos_th - temp_s * sin_th
sin_n = temp_s * cos_th + temp_c * sin_th
#
# here we perform the addition of the N-th mode
# to the data in 3D
#
a[i1, i2,
i3] = a[i1, i2, i3] + func_re(z, r) * cos_n - func_im(
z, r) * sin_n
osh5io.write_h5(b, filename=filename_out)
# ******************************************************************
from h5_utilities import *
import matplotlib.pyplot as plt
import sys
# Han Wen's pyVisOS
sys.path.append('/Volumes/Lacie-5TB/codes/pyVisOS/')
def os_rwigner(real_array,xaxis):
nx=real_array.shape
nxh=(nx+1)/2
dx=(xaxis.max-xaxis.min)/nx
kmax=(np.pi)/dx
dk=(2.0*np.pi)/xaxis.max
w_data= rwigner(real_array)
kaxis=osh5def.DataAxis(0,kmax,nx_h,
attrs={'NAME':'k','LONG_NAME':'wave number', 'UNITS': '\omega_0/c'}
data_attrs = data_attrs = { 'UNITS': osh5def.OSUnits('[a.u.]'),
'NAME': 'W_{\phi}', 'LONG_NAME': 'Wigner Transform' }
run_attrs = {'XMAX' : np.array( [kmax, xaxis.max] ) ,
'XMIN' : np.array( [0, xaxis.min ] ) }
os5data_wigner=osh5def.H5Data(w_data,timestamp='x', data_attrs=data_attrs, axes=[kaxis,xaxis])
