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])