def combine(dir_or_filelist,
prefix=None,
file_slice=slice(None, ),
preprocess=None,
axesdata=None,
save=None):
"""
stack a directory of grid data and optionally save the result to a file
:param dir_or_filelist: name of the directory
:param prefix: string, match file names with specific prefix
:param file_slice: a slice that applies to the file name list, useful for skipping every other files or start
in the middle of the list for example
:param preprocess: a list of callable (and their args and kwargs if any) that will act on the data before stacking
happens. It has to accept one H5Data objects as argument and return one H5Data objects.
Note that it has to be a list, not tuple or any other type, aka if isinstance(preprocess, list)
returns False then this parameter will be ignored entirely.
:param axesdata: user difined axes, see stack for more detail
:param save: name of the save file. user can also set it to true value and the output will use write_h5 defaults
:return: combined grid data, one dimension more than the preprocessed original data
Usage of preprocess:
The functino list should look like:
[(func1, arg11, arg21, ..., argn1, {'kwarg11': val11, 'kwarg21': val21, ..., 'kwargn1', valn1}),
(func2, arg12, arg22, ..., argn2, {'kwarg12': val12, 'kwarg22': val22, ..., 'kwargn2', valn2}),
...,
(func2, arg1n, arg2n, ..., argnn, {'kwarg1n': val1n, 'kwarg2n': val2n, ..., 'kwargnn', valnn})] where
any of the *args and/or **args can be omitted. see also __parse_func_param for limitations.
if preprocess=[(numpy.power, 2), (numpy.average, {axis=0}), numpy.sqrt], then the data to be stacked is
numpy.sqrt( numpy.average( numpy.power( read_h5(file_name), 2 ), axis=0 ) )
"""
prfx = str(prefix).strip() if prefix else ''
if isinstance(dir_or_filelist, str):
flist = sorted(glob.glob(dir_or_filelist + '/' + prfx +
'*.h5'))[file_slice]
else: # dir_or_filelist is a list of file names
flist = dir_or_filelist[file_slice]
if isinstance(preprocess, list) and preprocess:
func_list = [__parse_func_param(item) for item in preprocess]
tmp = [
reduce(lambda x, y: y[0](x, *y[1], **y[2]), func_list,
osh5io.read_h5(fn)) for fn in flist
]
else:
tmp = [osh5io.read_h5(f) for f in flist]
res = stack(tmp, axesdata=axesdata)
if save:
if not isinstance(save, str):
save = dir_or_filelist if isinstance(
dir_or_filelist, str) else './' + res.name + '.h5'
osh5io.write_h5(res, save)
return res
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)
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()
# run_attrs = h5_output.run_attrs
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)
h5_data = osh5io.read_h5(h5_filename)
for i in range(list_length):
temp = np.average((h5_data[i]), axis=0)
h5_output[i, file_number, 1:ny] = temp[1:ny]
# file_number+=1
print('before write')
for i in range(list_length):
list_filename = "%s-%02d.h5" % (outfilename, i)
print(list_filename)
b = osh5def.H5Data(h5_output[i, :, :],
timestamp='x',
data_attrs=data_attrs,
run_attrs=run_attrs,
axes=[taxis, xaxis])
osh5io.write_h5(b, filename=list_filename)
print('after write')
# print('Before barrier'+repr(rank))
