def diff(x, n=1, axis=-1):
dx_2 = 0.5 * x.axes[axis].increment
@metasl
def __diff(x, n=1, ax=-1):
return np.diff(x, n=n, axis=ax)
r = __diff(x, n=n, ax=axis)
r.axes[axis] = osh5def.DataAxis(axis_min=x.axes[axis].min+n*dx_2, axis_max=x.axes[axis].max-n*dx_2,
axis_npoints=x.axes[axis].size-n, attrs=copy.deepcopy(x.axes[axis].attrs))
return r
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 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)
# h5_output = hdf_data()
# h5_output.shape = [total_time, nx]
print('nx=' + repr(nx))
print('ny=' + repr(ny))
print('time_step=' + repr(time_step))
print('total_time=' + repr(total_time))
# calculate the z-range
dz=h5_data.axis[1][1]-h5_data.axis[1][0]
nx_begin=int((zmin-h5_data.axis[1][0])/dz)
nx_end=int((zmax-h5_data.axis[1][0])/dz)
xaxis=h5_data.axes[1][nx_begin:nx_end]
#
taxis=osh5def.DataAxis(0, time_step * (total_time -1), total_time,
attrs={'NAME':'t', 'LONG_NAME':'time', 'UNITS':'1 / \omega_p'})
delta_k = 2*np.pi/(h5_data.axis[0].max-h5_data.axis[0].min)
nk = int(kmax/delta_k)
nz=nx_end-nx_begin+1
# here we allocate the history array
total_es = np.zeros((total_time,nk,nz))
total_em =np.zeros((total_time,nk,nz))
total_es = 0
total_em = 0
kaxis=osh5def.DataAxis(0, nk* delta_k, nk, attrs{'NAME':'k_{\perp}', 'LONG_NAME':'k_{\perp}', 'UNITS':osh5def.OSUnits('\omega_{0}/c')})
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])
