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