def __read_2d_aver(self, plane, datadir, aver_file_name, n_vars): """ Read the yaverages.dat, xzaverages.dat, yzaverages.dat Return the raw data and the time array. """ import os import numpy as np from pencilnew import read # Determine the structure of the xy/xz/yz averages. if plane == 'xy': nw = getattr(read.dim(), 'nz') if plane == 'xz': nw = getattr(read.dim(), 'ny') if plane == 'yz': nw = getattr(read.dim(), 'nx') file_id = open(os.path.join(datadir, aver_file_name)) aver_lines = file_id.readlines() file_id.close() entry_length = int(np.ceil(nw * n_vars / 8.)) n_times = int(len(aver_lines) / (1. + entry_length)) # Prepare the data arrays. t = np.zeros(n_times, dtype=np.float32) raw_data = np.zeros([n_times, n_vars * nw]) # Read the data line_idx = 0 t_idx = -1 for current_line in aver_lines: if line_idx % (entry_length + 1) == 0: t_idx += 1 t[t_idx] = np.float32(current_line) raw_idx = 0 else: raw_data[t_idx, raw_idx*8:(raw_idx*8+8)] = \ list(map(np.float32, current_line.split())) raw_idx += 1 line_idx += 1 # Restructure the raw data and add it to the Averages object. raw_data = np.reshape(raw_data, [n_times, n_vars, nw]) return t, raw_data
def read(self, data_dir='data', param=None, dim=None): """ Read Pencil Code index data from index.pro. call signature: read(self, data_dir='data/', param=None, dim=None) Keyword arguments: *data_dir*: Directory where the data is stored. *param* Parameter object. *dim* Dimension object. """ import os import pencilnew.read as read if param is None: param = read.param(data_dir=data_dir, quiet=True) if dim is None: dim = read.dim(data_dir=data_dir) if param.lwrite_aux: totalvars = dim.mvar + dim.maux else: totalvars = dim.mvar index_file = open(os.path.join(data_dir, 'index.pro')) for line in index_file.readlines(): clean = line.strip() name = clean.split('=')[0].strip().replace('[', '').replace(']', '') if (clean.split('=')[1].strip().startswith('intarr(370)')): continue val = int(clean.split('=')[1].strip()) if (val != 0 and val <= totalvars \ and not name.startswith('i_') and name.startswith('i')): name = name.lstrip('i') if (name == 'lnTT' and param.ltemperature_nolog): name = 'tt' setattr(self, name, val)
def update(self, hard=False, quiet=True): """Update simulation object: if not read in: - read param.nml - read grid and ghost grid Set hard=True to force update. """ from os.path import exists from os.path import join from pencilnew.read import param, grid, dim REEXPORT = False if hard == True: self.param = False self.grid = False self.ghost_grid = False self.dim = False REEXPORT = True if self.param == False: try: if exists(join(self.datadir, 'param.nml')): print('~ Reading param.nml.. ') param = param(quiet=quiet, datadir=self.datadir) self.param = {} # read params into Simulation object for key in dir(param): if key.startswith('_') or key == 'read': continue self.param[key] = getattr(param, key) REEXPORT = True else: if not quiet: print( '? WARNING: for ' + self.path + '\n? Simulation has not run yet! Meaning: No param.nml found!' ) REEXPORT = True except: print('! ERROR: while reading param.nml for ' + self.path) self.param = False REEXPORT = True if self.param != False and (self.grid == False or self.ghost_grid == False): try: # read grid only if param is not False #import pencilnew as pcn; pcn.io.debug_breakpoint() print('~ Reading grid.. ') self.grid = grid(datadir=self.datadir, trim=True, quiet=True) print('~ Reading ghost_grid.. ') self.ghost_grid = grid(datadir=self.datadir, trim=False, quiet=True) print('~ Reading dim.. ') self.dim = dim(datadir=self.datadir) if not quiet: print('# Updating grid and ghost_grid succesfull') REEXPORT = True # adding lx, dx etc to params self.param['Lx'] = self.grid.Lx self.param['Ly'] = self.grid.Ly self.param['Lz'] = self.grid.Lz self.param['lx'] = self.grid.Lx self.param['ly'] = self.grid.Ly self.param['lz'] = self.grid.Lz self.param['dx'] = self.grid.dx self.param['dy'] = self.grid.dy self.param['dz'] = self.grid.dz except: if not quiet: print( '? WARNING: Updating grid and ghost_grid was not succesfull, since reading grid had an error' ) if self.started() or (not quiet): print('? WARNING: Couldnt load grid for ' + self.path) self.grid = False self.ghost_grid = False self.dim = False REEXPORT = True elif self.param == False: if not quiet: print( '? WARNING: Updating grid and ghost_grid was not succesfull, since run did is not started yet.' ) self.grid = False self.ghost_grid = False self.dim = False REEXPORT = True if REEXPORT == True: self.export() #import pencilnew as pcn; pcn.io.debug_breakpoint() return self
def read(self, field='', extension='', datadir='data', proc=-1, old_file=False, precision='f'): """ Read Pencil Code slice data. call signature: read(self. field='', extension='', datadir='data', proc=-1, old_file=False, precision='f') Keyword arguments: *field*: Name of the field(s) to be read. *extension* Specifies the slice(s). *datadir*: Directory where the data is stored. *proc*: Processor to be read. If -1 read all and assemble to one array. *old_file* Flag for reading old file format. *precision*: Precision of the data. Either float 'f' or double 'd'. """ import os import numpy as np from scipy.io import FortranFile from pencilnew import read # Define the directory that contains the slice files. if proc < 0: slice_dir = datadir else: slice_dir = os.path.join(datadir, 'proc{0}'.format(proc)) # Initialize the fields list. if field: if isinstance(field, list): field_list = field else: field_list = [field] else: # Find the existing fields. field_list = [] for file_name in os.listdir(slice_dir): if (file_name[:6] == 'slice_'): field_list.append(file_name.split('.')[0][6:]) # Remove duplicates. field_list = list(set(field_list)) # Initialize the extensions list. if extension: if isinstance(extension, list): extension_list = extension else: extension_list = [extension] else: # Find the existing extensions. extension_list = [] for file_name in os.listdir(slice_dir): if (file_name[:6] == 'slice_'): extension_list.append(file_name.split('.')[1]) # Remove duplicates. extension_list = list(set(extension_list)) class Foo(): pass for extension in extension_list: # This one will store the data. ext_object = Foo() for field in field_list: # Compose the file name according to field and extension. datadir = os.path.expanduser(datadir) if proc < 0: file_name = os.path.join( datadir, 'slice_' + field + '.' + extension) else: file_name = os.path.join( datadir, 'proc{0}'.format(proc), 'slice_' + field + '.' + extension) dim = read.dim(datadir, proc) if dim.precision == 'D': precision = 'd' else: precision = 'f' # Set up slice plane. if (extension == 'xy' or extension == 'Xy'): hsize = dim.nx vsize = dim.ny if (extension == 'xz'): hsize = dim.nx vsize = dim.nz if (extension == 'yz'): hsize = dim.ny vsize = dim.nz infile = FortranFile(file_name) islice = 0 self.t = np.zeros(1, dtype=precision) slice_series = np.zeros(1, dtype=precision) while True: try: raw_data = infile.read_record(dtype=precision) except ValueError: break except TypeError: break if old_file: self.t = np.concatenate((self.t, raw_data[-1:])) slice_series = np.concatenate( (slice_series, raw_data[:-1])) else: self.t = np.concatenate((self.t, raw_data[-2:-1])) slice_series = np.concatenate( (slice_series, raw_data[:-2])) islice += 1 # Reshape and remove first entry. self.t = self.t[1:] slice_series = slice_series[1:].reshape(islice, vsize, hsize) setattr(ext_object, field, slice_series) setattr(self, extension, ext_object)
def read(self, datadir='data', proc=-1, quiet=False, trim=False): """ Read the grid data from the pencil code simulation. If proc < 0, then load all data and assemble. Otherwise, load grid from specified processor. call signature: read(self, file_name='time_series.dat', datadir='data', double=0, quiet=0, comment_char='#') Keyword arguments: *datadir*: Directory where the data is stored. *proc* Processor to be read. If proc is -1, then read the 'global' grid. If proc is >=0, then read the grid.dat in the corresponding processor directory. *quiet* Flag for switching of output. *trim* Cuts off the ghost points. """ import numpy as np import os from scipy.io import FortranFile import pencilnew.read as read datadir = os.path.expanduser(datadir) dim = read.dim(datadir, proc) if dim.precision == 'D': precision = 'd' else: precision = 'f' if proc < 0: import pencilnew proc_dirs = list(filter(lambda s: s.startswith('proc'), os.listdir(datadir))) else: proc_dirs = ['proc'+str(proc)] # Define the global arrays. x = np.zeros(dim.mx, dtype=precision) y = np.zeros(dim.my, dtype=precision) z = np.zeros(dim.mz, dtype=precision) dx_1 = np.zeros(dim.mx, dtype=precision) dy_1 = np.zeros(dim.my, dtype=precision) dz_1 = np.zeros(dim.mz, dtype=precision) dx_tilde = np.zeros(dim.mx, dtype=precision) dy_tilde = np.zeros(dim.my, dtype=precision) dz_tilde = np.zeros(dim.mz, dtype=precision) for directory in proc_dirs: proc = int(directory[4:]) procdim = read.dim(datadir, proc) if not quiet: print("reading grid data from processor {0} of {1} ...".format(proc, len(proc_dirs))) mxloc = procdim.mx myloc = procdim.my mzloc = procdim.mz # Read the grid data. file_name = os.path.join(datadir, directory, 'grid.dat') infile = FortranFile(file_name, 'r') grid_raw = infile.read_record(dtype=precision) dx, dy, dz = tuple(infile.read_record(dtype=precision)) Lx, Ly, Lz = tuple(infile.read_record(dtype=precision)) dx_1_raw = infile.read_record(dtype=precision) dx_tilde_raw = infile.read_record(dtype=precision) infile.close() # Reshape the arrays. t = grid_raw[0] x_loc = grid_raw[1:mxloc+1] y_loc = grid_raw[mxloc+1:mxloc+myloc+1] z_loc = grid_raw[mxloc+myloc+1:mxloc+myloc+mzloc+1] dx_1_loc = dx_1_raw[0:mxloc] dy_1_loc = dx_1_raw[mxloc:mxloc+myloc] dz_1_loc = dx_1_raw[mxloc+myloc:mxloc+myloc+mzloc] dx_tilde_loc = dx_tilde_raw[0:mxloc] dy_tilde_loc = dx_tilde_raw[mxloc:mxloc+myloc] dz_tilde_loc = dx_tilde_raw[mxloc+myloc:mxloc+myloc+mzloc] if len(proc_dirs) > 1: if procdim.ipx == 0: i0x = 0 i1x = i0x + procdim.mx i0x_loc = 0 i1x_loc = procdim.mx else: i0x = procdim.ipx*procdim.nx + procdim.nghostx i1x = i0x + procdim.mx - procdim.nghostx i0x_loc = procdim.nghostx i1x_loc = procdim.mx if procdim.ipy == 0: i0y = 0 i1y = i0y + procdim.my i0y_loc = 0 i1y_loc = procdim.my else: i0y = procdim.ipy*procdim.ny + procdim.nghosty i1y = i0y + procdim.my - procdim.nghosty i0y_loc = procdim.nghosty i1y_loc = procdim.my if procdim.ipz == 0: i0z = 0 i1z = i0z + procdim.mz i0z_loc = 0 i1z_loc = procdim.mz else: i0z = procdim.ipz*procdim.nz + procdim.nghostz i1z = i0z + procdim.mz - procdim.nghostz i0z_loc = procdim.nghostz i1z_loc = procdim.mz x[i0x:i1x] = x_loc[i0x_loc:i1x_loc] y[i0y:i1y] = y_loc[i0y_loc:i1y_loc] z[i0z:i1z] = z_loc[i0z_loc:i1z_loc] dx_1[i0x:i1x] = dx_1_loc[i0x_loc:i1x_loc] dy_1[i0y:i1y] = dy_1_loc[i0y_loc:i1y_loc] dz_1[i0z:i1z] = dz_1_loc[i0z_loc:i1z_loc] dx_tilde[i0x:i1x] = dx_tilde_loc[i0x_loc:i1x_loc] dy_tilde[i0y:i1y] = dy_tilde_loc[i0y_loc:i1y_loc] dz_tilde[i0z:i1z] = dz_tilde_loc[i0z_loc:i1z_loc] else: x = x_loc y = y_loc z = z_loc dx_1 = dx_1_loc dy_1 = dy_1_loc dz_1 = dz_1_loc dx_tilde = dx_tilde_loc dy_tilde = dy_tilde_loc dz_tilde = dz_tilde_loc if trim: self.x = x[dim.l1:dim.l2+1] self.y = y[dim.m1:dim.m2+1] self.z = z[dim.n1:dim.n2+1] self.dx_1 = dx_1[dim.l1:dim.l2+1] self.dy_1 = dy_1[dim.m1:dim.m2+1] self.dz_1 = dz_1[dim.n1:dim.n2+1] self.dx_tilde = dx_tilde[dim.l1:dim.l2+1] self.dy_tilde = dy_tilde[dim.m1:dim.m2+1] self.dz_tilde = dz_tilde[dim.n1:dim.n2+1] else: self.x = x self.y = y self.z = z self.dx_1 = dx_1 self.dy_1 = dy_1 self.dz_1 = dz_1 self.dx_tilde = dx_tilde self.dy_tilde = dy_tilde self.dz_tilde = dz_tilde self.t = t self.dx = dx self.dy = dy self.dz = dz self.Lx = Lx self.Ly = Ly self.Lz = Lz
def update(self, quiet=True): """Update simulation object: - read param.nml - read grid and ghost grid Someday this might also change the real simulation parameter. """ from os.path import exists from os.path import join from pencilnew.read import param, grid, dim if self.param == False or self.param != param(quiet=quiet, datadir=self.datadir): try: if exists(join(self.datadir, 'param.nml')): param = param(quiet=quiet, datadir=self.datadir) self.param = {} # read params into Simulation object for key in dir(param): if key.startswith('_') or key == 'read': continue self.param[key] = getattr(param, key) else: if not quiet: print( '? WARNING: for ' + self.path + '\n? Simulation has not run yet! Meaning: No param.nml found!' ) except: print('! ERROR: while reading param.nml for ' + self.path) self.param = False if self.param != False and (self.grid == False or self.ghost_grid == False): try: # read grid only if param is not False #import pencilnew as pcn; pcn.io.debug_breakpoint() self.grid = grid(datadir=self.datadir, trim=True, quiet=True) self.ghost_grid = grid(datadir=self.datadir, trim=False, quiet=True) self.dim = dim(datadir=self.datadir) if not quiet: print('# Updating grid and ghost_grid succesfull') except: if not quiet: print( '? WARNING: Updating grid and ghost_grid was not succesfull, since reading grid had an error' ) if self.started() or (not quiet): print('? WARNING: Couldnt load grid for ' + self.path) self.grid = False self.ghost_grid = False self.dim = False elif self.param == False: if not quiet: print( '? WARNING: Updating grid and ghost_grid was not succesfull, since run did is not started yet.' ) self.grid = False self.ghost_grid = False self.dim = False self.export() #import pencilnew as pcn; pcn.io.debug_breakpoint() return self
def __read_1d_aver(self, plane, datadir, aver_file_name, n_vars, proc): """ Read the yaverages.dat, zaverages.dat. Return the raw data and the time array. """ import os import numpy as np from scipy.io import FortranFile from pencilnew import read if proc < 0: proc_dirs = self.__natural_sort( filter(lambda s: s.startswith('proc'), os.listdir(datadir))) else: proc_dirs = ['proc' + str(proc)] dim = read.dim(datadir, proc) if dim.precision == 'S': dtype = np.float32 if dim.precision == 'D': dtype = np.float64 # Prepare the raw data. # This will be reformatted at the end. raw_data = [] for directory in proc_dirs: proc = int(directory[4:]) proc_dim = read.dim(datadir, proc) # Read the data. t = [] proc_data = [] file_id = FortranFile( os.path.join(datadir, directory, aver_file_name)) while True: try: t.append(file_id.read_record(dtype=dtype)[0]) proc_data.append(file_id.read_record(dtype=dtype)) except: # Finished reading. break file_id.close() # Reshape the proc data into [len(t), pnu, pnv]. if plane == 'y': pnu = proc_dim.nx pnv = proc_dim.nz if plane == 'z': pnu = proc_dim.nx pnv = proc_dim.ny proc_data = np.array(proc_data) print(proc_data.shape, len(t), n_vars, pnv, pnu) proc_data = proc_data.reshape([len(t), n_vars, pnv, pnu]) # Add the proc_data (one proc) to the raw_data (all procs) if plane == 'y': nu = dim.nx nv = dim.nz idx_u = proc_dim.ipx * proc_dim.nx idx_v = proc_dim.ipz * proc_dim.nz if plane == 'z': nu = dim.nx nv = dim.ny idx_u = proc_dim.ipx * proc_dim.nx idx_v = proc_dim.ipy * proc_dim.ny if not isinstance(raw_data, np.ndarray): # Initialize the raw_data array with the right dimensions. raw_data = np.zeros([len(t), n_vars, nv, nu]) raw_data[:, :, idx_v:idx_v + pnv, idx_u:idx_u + pnu] = proc_data.copy() t = np.array(t) raw_data = np.swapaxes(raw_data, 2, 3) return t, raw_data
def read(self, datadir='data', proc=-1, quiet=False, trim=False): """ Read the grid data from the pencil code simulation. If proc < 0, then load all data and assemble. Otherwise, load grid from specified processor. call signature: read(self, file_name='time_series.dat', datadir='data', double=0, quiet=0, comment_char='#') Keyword arguments: *datadir*: Directory where the data is stored. *proc* Processor to be read. If proc is -1, then read the 'global' grid. If proc is >=0, then read the grid.dat in the corresponding processor directory. *quiet* Flag for switching of output. *trim* Cuts off the ghost points. """ import numpy as np import os from scipy.io import FortranFile import pencilnew.read as read datadir = os.path.expanduser(datadir) dim = read.dim(datadir, proc) if dim.precision == 'D': precision = 'd' else: precision = 'f' if proc < 0: proc_dirs = list(filter(lambda string: string.startswith('proc'), os.listdir(datadir))) else: proc_dirs = ['proc' + str(proc)] # Define the global arrays. x = np.zeros(dim.mx, dtype=precision) y = np.zeros(dim.my, dtype=precision) z = np.zeros(dim.mz, dtype=precision) dx_1 = np.zeros(dim.mx, dtype=precision) dy_1 = np.zeros(dim.my, dtype=precision) dz_1 = np.zeros(dim.mz, dtype=precision) dx_tilde = np.zeros(dim.mx, dtype=precision) dy_tilde = np.zeros(dim.my, dtype=precision) dz_tilde = np.zeros(dim.mz, dtype=precision) for directory in proc_dirs: proc = int(directory[4:]) procdim = read.dim(datadir, proc) if not quiet: print("reading grid data from processor {0} of {1} ...".format(proc, len(proc_dirs))) mxloc = procdim.mx myloc = procdim.my mzloc = procdim.mz # Read the grid data. file_name = os.path.join(datadir, directory, 'grid.dat') infile = FortranFile(file_name, 'r') grid_raw = infile.read_record(dtype=precision) dx, dy, dz = tuple(infile.read_record(dtype=precision)) Lx, Ly, Lz = tuple(infile.read_record(dtype=precision)) dx_1_raw = infile.read_record(dtype=precision) dx_tilde_raw = infile.read_record(dtype=precision) infile.close() # Reshape the arrays. t = grid_raw[0] x_loc = grid_raw[1:mxloc+1] y_loc = grid_raw[mxloc+1:mxloc+myloc+1] z_loc = grid_raw[mxloc+myloc+1:mxloc+myloc+mzloc+1] dx_1_loc = dx_1_raw[0:mxloc] dy_1_loc = dx_1_raw[mxloc:mxloc+myloc] dz_1_loc = dx_1_raw[mxloc+myloc:mxloc+myloc+mzloc] dx_tilde_loc = dx_tilde_raw[0:mxloc] dy_tilde_loc = dx_tilde_raw[mxloc:mxloc+myloc] dz_tilde_loc = dx_tilde_raw[mxloc+myloc:mxloc+myloc+mzloc] if len(proc_dirs) > 1: if procdim.ipx == 0: i0x = 0 i1x = i0x + procdim.mx i0x_loc = 0 i1x_loc = procdim.mx else: i0x = procdim.ipx*procdim.nx + procdim.nghostx i1x = i0x + procdim.mx - procdim.nghostx i0x_loc = procdim.nghostx i1x_loc = procdim.mx if procdim.ipy == 0: i0y = 0 i1y = i0y + procdim.my i0y_loc = 0 i1y_loc = procdim.my else: i0y = procdim.ipy*procdim.ny + procdim.nghosty i1y = i0y + procdim.my - procdim.nghosty i0y_loc = procdim.nghosty i1y_loc = procdim.my if procdim.ipz == 0: i0z = 0 i1z = i0z + procdim.mz i0z_loc = 0 i1z_loc = procdim.mz else: i0z = procdim.ipz*procdim.nz + procdim.nghostz i1z = i0z + procdim.mz - procdim.nghostz i0z_loc = procdim.nghostz i1z_loc = procdim.mz x[i0x:i1x] = x_loc[i0x_loc:i1x_loc] y[i0y:i1y] = y_loc[i0y_loc:i1y_loc] z[i0z:i1z] = z_loc[i0z_loc:i1z_loc] dx_1[i0x:i1x] = dx_1_loc[i0x_loc:i1x_loc] dy_1[i0y:i1y] = dy_1_loc[i0y_loc:i1y_loc] dz_1[i0z:i1z] = dz_1_loc[i0z_loc:i1z_loc] dx_tilde[i0x:i1x] = dx_tilde_loc[i0x_loc:i1x_loc] dy_tilde[i0y:i1y] = dy_tilde_loc[i0y_loc:i1y_loc] dz_tilde[i0z:i1z] = dz_tilde_loc[i0z_loc:i1z_loc] else: x = x_loc y = y_loc z = z_loc dx_1 = dx_1_loc dy_1 = dy_1_loc dz_1 = dz_1_loc dx_tilde = dx_tilde_loc dy_tilde = dy_tilde_loc dz_tilde = dz_tilde_loc if trim: self.x = x[dim.l1:dim.l2+1] self.y = y[dim.m1:dim.m2+1] self.z = z[dim.n1:dim.n2+1] self.dx_1 = dx_1[dim.l1:dim.l2+1] self.dy_1 = dy_1[dim.m1:dim.m2+1] self.dz_1 = dz_1[dim.n1:dim.n2+1] self.dx_tilde = dx_tilde[dim.l1:dim.l2+1] self.dy_tilde = dy_tilde[dim.m1:dim.m2+1] self.dz_tilde = dz_tilde[dim.n1:dim.n2+1] else: self.x = x self.y = y self.z = z self.dx_1 = dx_1 self.dy_1 = dy_1 self.dz_1 = dz_1 self.dx_tilde = dx_tilde self.dy_tilde = dy_tilde self.dz_tilde = dz_tilde self.t = t self.dx = dx self.dy = dy self.dz = dz self.Lx = Lx self.Ly = Ly self.Lz = Lz
def read(self, var_file='', sim=None, datadir='data', proc=-1, ivar=-1, quiet=True, trim_all=False, trimall=False, dim=None, magic=None): """ Read VAR files from pencil code. If proc < 0, then load all data and assemble. otherwise, load VAR file from specified processor. The file format written by output() (and used, e.g. in var.dat) consists of the followinig Fortran records: 1. data(mx, my, mz, nvar) 2. t(1), x(mx), y(my), z(mz), dx(1), dy(1), dz(1), deltay(1) Here nvar denotes the number of slots, i.e. 1 for one scalar field, 3 for one vector field, 8 for var.dat in the case of MHD with entropy. but, deltay(1) is only there if lshear is on! need to know parameters. call signature: read(var_file='', datadir='data/', proc=-1, ivar=-1, quiet=True, trimall=False, magic=None, sim=None) Keyword arguments: var_file: Name of the VAR file. sim: Simulation sim object. magic: Values to be computed from the data, e.g. B = curl(A). trimall: Trim the data cube to exclude ghost zones. quiet: Flag for switching off output. datadir: Directory where the data is stored. proc: Processor to be read. If -1 read all and assemble to one array. ivar: Index of the VAR file, if var_file is not specified. """ import numpy as np import os from scipy.io import FortranFile from pencilnew.math.derivatives import curl, curl2 from pencilnew import read if sim is not None: datadir = os.path.expanduser(sim.datadir) dim = sim.dim param = read.param(datadir=sim.datadir, quiet=True) index = read.index(datadir=sim.datadir) else: datadir = os.path.expanduser(datadir) if dim is None: dim = read.dim(datadir, proc) if param is None: param = read.param(datadir=datadir, quiet=quiet) if index is None: index = read.index(datadir=datadir) run2D = param.lwrite_2d if dim.precision == 'D': precision = 'd' else: precision = 'f' if param.lwrite_aux: total_vars = dim.mvar + dim.maux else: total_vars = dim.mvar if not var_file: if ivar < 0: var_file = 'var.dat' else: var_file = 'VAR' + str(ivar) if proc < 0: proc_dirs = self.__natural_sort( filter(lambda s: s.startswith('proc'), os.listdir(datadir))) else: proc_dirs = ['proc' + str(proc)] if trimall != False: trim_all = trimall # Set up the global array. if not run2D: f = np.zeros((total_vars, dim.mz, dim.my, dim.mx), dtype=precision) else: if dim.ny == 1: f = np.zeros((total_vars, dim.mz, dim.mx), dtype=precision) else: f = np.zeros((total_vars, dim.my, dim.mx), dtype=precision) x = np.zeros(dim.mx, dtype=precision) y = np.zeros(dim.my, dtype=precision) z = np.zeros(dim.mz, dtype=precision) for directory in proc_dirs: proc = int(directory[4:]) procdim = read.dim(datadir, proc) if not quiet: print("Reading data from processor {0} of {1} ...".format( \ proc, len(proc_dirs))) mxloc = procdim.mx myloc = procdim.my mzloc = procdim.mz # Read the data. file_name = os.path.join(datadir, directory, var_file) infile = FortranFile(file_name) if not run2D: f_loc = infile.read_record(dtype=precision) f_loc = f_loc.reshape((-1, mzloc, myloc, mxloc)) else: if dim.ny == 1: f_loc = infile.read_record(dtype=precision) f_loc = f_loc.reshape((-1, mzloc, mxloc)) else: f_loc = infile.read_record(dtype=precision) f_loc = f_loc.reshape((-1, myloc, mxloc)) raw_etc = infile.read_record(precision) infile.close() t = raw_etc[0] x_loc = raw_etc[1:mxloc + 1] y_loc = raw_etc[mxloc + 1:mxloc + myloc + 1] z_loc = raw_etc[mxloc + myloc + 1:mxloc + myloc + mzloc + 1] if param.lshear: shear_offset = 1 deltay = raw_etc[-1] else: shear_offset = 0 dx = raw_etc[-3 - shear_offset] dy = raw_etc[-2 - shear_offset] dz = raw_etc[-1 - shear_offset] if len(proc_dirs) > 1: # Calculate where the local processor will go in # the global array. # # Don't overwrite ghost zones of processor to the left (and # accordingly in y and z direction -- makes a difference on the # diagonals) # # Recall that in NumPy, slicing is NON-INCLUSIVE on the right end # ie, x[0:4] will slice all of a 4-digit array, not produce # an error like in idl. if procdim.ipx == 0: i0x = 0 i1x = i0x + procdim.mx i0xloc = 0 i1xloc = procdim.mx else: i0x = procdim.ipx * procdim.nx + procdim.nghostx i1x = i0x + procdim.mx - procdim.nghostx i0xloc = procdim.nghostx i1xloc = procdim.mx if procdim.ipy == 0: i0y = 0 i1y = i0y + procdim.my i0yloc = 0 i1yloc = procdim.my else: i0y = procdim.ipy * procdim.ny + procdim.nghosty i1y = i0y + procdim.my - procdim.nghosty i0yloc = procdim.nghosty i1yloc = procdim.my if procdim.ipz == 0: i0z = 0 i1z = i0z + procdim.mz i0zloc = 0 i1zloc = procdim.mz else: i0z = procdim.ipz * procdim.nz + procdim.nghostz i1z = i0z + procdim.mz - procdim.nghostz i0zloc = procdim.nghostz i1zloc = procdim.mz x[i0x:i1x] = x_loc[i0xloc:i1xloc] y[i0y:i1y] = y_loc[i0yloc:i1yloc] z[i0z:i1z] = z_loc[i0zloc:i1zloc] if not run2D: f[:, i0z:i1z, i0y:i1y, i0x:i1x] = \ f_loc[:, i0zloc:i1zloc, i0yloc:i1yloc, i0xloc:i1xloc] else: if dim.ny == 1: f[:, i0z:i1z, i0x:i1x] = \ f_loc[:, i0zloc:i1zloc, i0xloc:i1xloc] else: f[:, i0y:i1y, i0x:i1x] = \ f_loc[:, i0yloc:i1yloc, i0xloc:i1xloc] else: f = f_loc x = x_loc y = y_loc z = z_loc if magic is not None: if 'bb' in magic: # Compute the magnetic field before doing trim_all. aa = f[index.ax - 1:index.az, ...] self.bb = curl(aa, dx, dy, dz, run2D=run2D) if trim_all: self.bb = self.bb[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1, dim.l1:dim.l2 + 1] if 'jj' in magic: # Compute the electric current field before doing trim_all. aa = f[index.ax - 1:index.az, ...] self.jj = curl2(aa, dx, dy, dz) if trim_all: self.jj = self.jj[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1, dim.l1:dim.l2 + 1] if 'vort' in magic: # Compute the vorticity field before doing trim_all. uu = f[index.ux - 1:index.uz, ...] self.vort = curl(uu, dx, dy, dz, run2D=run2D) if trim_all: if run2D: if dim.nz == 1: self.vort = self.vort[:, dim.m1:dim.m2 + 1, dim.l1:dim.l2 + 1] else: self.vort = self.vort[:, dim.n1:dim.n2 + 1, dim.l1:dim.l2 + 1] else: self.vort = self.vort[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1, dim.l1:dim.l2 + 1] # Trim the ghost zones of the global f-array if asked. if trim_all: self.x = x[dim.l1:dim.l2 + 1] self.y = y[dim.m1:dim.m2 + 1] self.z = z[dim.n1:dim.n2 + 1] if not run2D: self.f = f[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1, dim.l1:dim.l2 + 1] else: if dim.ny == 1: self.f = f[:, dim.n1:dim.n2 + 1, dim.l1:dim.l2 + 1] else: self.f = f[:, dim.m1:dim.m2 + 1, dim.l1:dim.l2 + 1] else: self.x = x self.y = y self.z = z self.f = f self.l1 = dim.l1 self.l2 = dim.l2 + 1 self.m1 = dim.m1 self.m2 = dim.m2 + 1 self.n1 = dim.n1 self.n2 = dim.n2 + 1 # Assign an attribute to self for each variable defined in # 'data/index.pro' so that e.g. self.ux is the x-velocity for key in index.__dict__.keys(): if key != 'global_gg' and key != 'keys': value = index.__dict__[key] setattr(self, key, self.f[value - 1, ...]) # Special treatment for vector quantities. if hasattr(index, 'uu'): self.uu = self.f[index.ux - 1:index.uz, ...] if hasattr(index, 'aa'): self.aa = self.f[index.ax - 1:index.az, ...] self.t = t self.dx = dx self.dy = dy self.dz = dz if param.lshear: self.deltay = deltay # Do the rest of magic after the trim_all (i.e. no additional curl...). self.magic = magic if self.magic is not None: self.magic_attributes(param)
def update(self, hard=False, quiet=True): """Update simulation object: if not read in: - read param.nml - read grid and ghost grid Set hard=True to force update. """ from os.path import exists from os.path import join from pencilnew.read import param, grid, dim REEXPORT = False if hard == True: self.param = False self.grid = False self.ghost_grid = False self.dim = False REEXPORT = True if self.param == False: try: if exists(join(self.datadir,'param.nml')): print('~ Reading param.nml.. ') param = param(quiet=quiet, datadir=self.datadir) self.param = {} # read params into Simulation object for key in dir(param): if key.startswith('_') or key == 'read': continue self.param[key] = getattr(param, key) REEXPORT = True else: if not quiet: print('? WARNING: for '+self.path+'\n? Simulation has not run yet! Meaning: No param.nml found!') REEXPORT = True except: print('! ERROR: while reading param.nml for '+self.path) self.param = False REEXPORT = True if self.param != False and (self.grid == False or self.ghost_grid == False): try: # read grid only if param is not False #import pencilnew as pcn; pcn.io.debug_breakpoint() print('~ Reading grid.. ') self.grid = grid(datadir=self.datadir, trim=True, quiet=True) print('~ Reading ghost_grid.. ') self.ghost_grid = grid(datadir=self.datadir, trim=False, quiet=True) print('~ Reading dim.. ') self.dim = dim(datadir=self.datadir) if not quiet: print('# Updating grid and ghost_grid succesfull') REEXPORT = True # adding lx, dx etc to params self.param['Lx'] = self.grid.Lx; self.param['Ly'] = self.grid.Ly; self.param['Lz'] = self.grid.Lz self.param['lx'] = self.grid.Lx; self.param['ly'] = self.grid.Ly; self.param['lz'] = self.grid.Lz self.param['dx'] = self.grid.dx; self.param['dy'] = self.grid.dy; self.param['dz'] = self.grid.dz except: if not quiet: print('? WARNING: Updating grid and ghost_grid was not succesfull, since reading grid had an error') if self.started() or (not quiet): print('? WARNING: Couldnt load grid for '+self.path) self.grid = False self.ghost_grid = False self.dim = False REEXPORT = True elif self.param == False: if not quiet: print('? WARNING: Updating grid and ghost_grid was not succesfull, since run did is not started yet.') self.grid = False self.ghost_grid = False self.dim = False REEXPORT = True if REEXPORT == True: self.export() #import pencilnew as pcn; pcn.io.debug_breakpoint() return self
def animate(self, field='uu1', extension='xz', datadir='data', proc=-1, interval=50., amin=0., amax=1., transform='', old_file=False, precision='f', verbose=False): """ Read Pencil Code slice data. call signature: animate(self. field='', extension='', datadir='data', proc=-1, old_file=False, precision='f') Keyword arguments: *field*: Name of the field(s) to be read. *extension* Specifies the slice(s). *datadir*: Directory where the data is stored. *proc*: Processor to be read. If -1 read all and assemble to one array. *old_file* Flag for reading old file format. *precision*: Precision of the data. Either float 'f' or double 'd'. *verbose*: Print progress *tmin*: Start time of the animation *tmax*: End time of the animation *amin*: Minimum value for image scaling *amax*: Maximum value for image scaling *transform*: Insert arbitrary numerical code to modify the slice *wait*: Pause in seconds between animation slices """ import os import numpy as np from pencilnew import read from scipy.io import FortranFile import matplotlib.pyplot as plt import matplotlib.animation as animation # Define the directory that contains the slice files. if proc < 0: slice_dir = datadir else: slice_dir = os.path.join(datadir, 'proc{0}'.format(proc)) # Initialize the fields list. if field: if isinstance(field, list): field_list = field else: field_list = [field] else: # Find the existing fields. field_list = [] for file_name in os.listdir(slice_dir): if file_name[:6] == 'slice_': field_list.append(file_name.split('.')[0][6:]) # Remove duplicates. field_list = list(set(field_list)) # Initialize the extensions list. if extension: if isinstance(extension, list): extension_list = extension else: extension_list = [extension] else: # Find the existing extensions. extension_list = [] for file_name in os.listdir(slice_dir): if file_name[:6] == 'slice_': extension_list.append(file_name.split('.')[1]) # Remove duplicates. extension_list = list(set(extension_list)) class Foo(object): pass for extension in extension_list: if verbose: print('Extension: ' + str(extension)) # This one will store the data. ext_object = Foo() for field in field_list: if verbose: print(' -> Field: ' + str(field)) # Compose the file name according to field and extension. datadir = os.path.expanduser(datadir) if proc < 0: file_name = os.path.join( datadir, 'slice_' + field + '.' + extension) else: file_name = os.path.join( datadir, 'proc{0}'.format(proc), 'slice_' + field + '.' + extension) dim = read.dim(datadir, proc) if dim.precision == 'D': precision = 'd' else: precision = 'f' # Set up slice plane. if extension == 'xy' or extension == 'Xy' or extension == 'xy2': hsize = dim.nx vsize = dim.ny if extension == 'xz': hsize = dim.nx vsize = dim.nz if extension == 'yz': hsize = dim.ny vsize = dim.nz infile = FortranFile(file_name) islice = 0 self.t = np.zeros(1, dtype=precision) self.t = [0] slice_series = [0] plane = np.zeros((vsize, hsize), dtype=precision) fig = plt.figure() # ims is a list of lists, each row is a list of artists to draw in the # current frame; here we are just animating one artist, the image, in # each frame ax = plt.axes() ax.set_xlabel('x') ax.set_ylabel('y') ims = [] # image series ifirst = True while True: try: if verbose: print(' -> Reading... ') raw_data = infile.read_record(dtype=precision) except ValueError: break except TypeError: break if old_file: self.t.append(list(raw_data[-1])) t_now = raw_data[-1] slice_series.extend(list(raw_data[:-1])) plane = raw_data[:-1] else: self.t.append(list(raw_data[-2:-1])) t_now = raw_data[-2:-1] slice_series.extend(list(raw_data[:-2])) plane = raw_data[:-2] if verbose: print(' -> Done') if ifirst: print( "----islice----------t---------min-------max-------delta" ) print("%10i %10.3e %10.3e %10.3e %10.3e" % (islice, t_now, plane.min(), plane.max(), plane.max() - plane.min())) ifirst = False title = 't = %11.3e' % t_now ax.set_title(title) im = plt.imshow(plane, vmin=amin, vmax=amax, animated=True) ims.append([im]) # 每张图片都用单独list的形式加入到图片序列中 islice += 1 # Reshape and remove first entry. ani = animation.ArtistAnimation(fig, ims, interval=interval, blit=True, repeat_delay=1000) ani.save('animate_' + extension + field + '.mp4') plt.show() if verbose: print('Reshaping array') self.t = np.array(self.t[1:], dtype=precision) slice_series = np.array(slice_series, dtype=precision) slice_series = slice_series[1:].reshape(islice, vsize, hsize) setattr(ext_object, field, slice_series) setattr(self, extension, ext_object)
def calc(self, aver=[], datatopdir='.', lskip_zeros=False, proc=0, rank=0, rmfzeros=1, rmbzeros=1, iy=None, l_correction=False, t_correction=0., dim=None, timereducer=None, trargs=[], tindex=(0,None), imask=None ): """object returns time dependent meridional tensors from Averages object aver.z. u, acoef and bcoef and aver.t For long DNS runs the 'zaverages.dat' file can be very large so MPI may be required and the data is loaded by processor as default. lskip_zeros=True identifies the resetting of the testfield and rmbzeros and rmfzeros number to exclude before and following By default none are removed. iy is the index array that is computed in this MPI process, which may be a subset of the array on this processor l_correction=True permits the pencil coefficients computed prior to the Pencil Code correction implemented after time=t_correction to be rescaled accordingly to match the new formulation. trargs contain optional arguments for the time treatments: mean, smoothing, etc. tindex is set to limit the range of the iterations loaded from Averages in zaverages.dat The index imask, excluding the resets, can be specified to ensure all processes use the same mask """ import numpy as np import os from pencilnew import read os.chdir(datatopdir) # return to working directory grid = read.grid(proc=proc,trim=True, quiet=True) # if iy None or scalar create numpy array try: iy.size>0 except: print('exception') if iy==None: print('exception None') iy=np.arange(grid.y.size) else: print('exception int') iy=np.array(iy) if rank==0: print('iy size is {0}'.format(iy.shape)) r, theta = np.meshgrid(grid.x,grid.y[iy],indexing='ij') del(grid,theta) #conserve memory print('rank {0} calculating tensors for proc {1}'.format(rank,proc)) # string containers for zaverages.z keys uformat = 'u{0}mxy' alpformat = 'alp{0}{1}xy' etaformat = 'eta{0}{1}{2}xy' # imask calculated once for MPI/processor consistency if rank==0: print("Removing zeros") old_size = aver.t.shape # if imask is not provided either exclude the zeros or use the full time series try: imask.size>0 print('imask shape is {}'.format(imask.shape)) except: if lskip_zeros: index = alpformat.format(1,1) izero=np.array(np.where(aver.z.__getattribute__(index)[:, aver.z.__getattribute__(index).shape[-2]/2, aver.z.__getattribute__(index).shape[-1]/2]==0))[0] rmfrange = np.arange(0,rmfzeros-1) rmbrange = np.arange(0,rmbzeros-1) rmpoints = np.array([],dtype=int) for zero in izero: rmpoints = np.append(rmpoints, rmfrange + zero) rmpoints = np.append(rmpoints, zero - rmbrange) if izero.size>0: imask=np.delete(np.where(aver.t),rmpoints) if rank==0: print("Removed {0} zeros from {1} resets".format(len(rmpoints), len(izero))) print("Resets occured at save points {0}".format(izero)) else: imask=np.where(aver.t)[0] del(rmpoints,rmbrange,rmfrange) else: imask=np.arange(aver.t.size) if rank==0: print("Skipped zero removals.") # update the time of the snapshots included self.t=aver.t[imask] # Correction to Pencil Code error may be required on old data if l_correction: if dim==None: dim=read.dim(quiet=True) itcorr = np.where(aver.t[imask]<t_correction)[0] index = alpformat.format(1,3) aver.z.__getattribute__(index)[itcorr] *=\ -dim.nprocz/(dim.nprocz-2.) for j in range(0,3): index = alpformat.format(3,j+1) aver.z.__getattribute__(index)[itcorr] *=\ -dim.nprocz/(dim.nprocz-2.) index = etaformat.format(1,1,1) aver.z.__getattribute__(index)[itcorr] *=\ -dim.nprocz/(dim.nprocz-2.) for j in range(0,3): index = etaformat.format(j+1,2,1) aver.z.__getattribute__(index)[itcorr] *=\ -dim.nprocz/(dim.nprocz-2.) index = etaformat.format(1,1,2) aver.z.__getattribute__(index)[itcorr] *=\ -dim.nprocz/(dim.nprocz-2.) for j in range(0,3): index = etaformat.format(j+1,2,2) aver.z.__getattribute__(index)[itcorr] *=\ -dim.nprocz/(dim.nprocz-2.) # set up place holders for the Pencil Code tensor coefficients index = alpformat.format(1,1) u =np.zeros([3, len(imask),aver.z.__getattribute__(index).shape[-2],iy.size]) alp=np.zeros([3,3, len(imask),aver.z.__getattribute__(index).shape[-2],iy.size]) eta=np.zeros([3,3,3,len(imask),aver.z.__getattribute__(index).shape[-2],iy.size]) if rank==0: print(u.shape,aver.z.__getattribute__(index)[imask,:,:].shape) # store the individual components in the z-averages as tensors for i,coord in zip(range(0,3),('x','y','z')): try: index = uformat.format(coord) if iy.size>1: tmp=aver.z.__getattribute__(index)[:,:,iy] u[i,:,:,:] = tmp[imask] else: u[i,:,:,0] = aver.z.__getattribute__(index)[imask,:,iy] except KeyError: pass for i in range(0,3): for j in range(0,3): index = alpformat.format(i+1,j+1) if iy.size>1: tmp=aver.z.__getattribute__(index)[:,:,iy] alp[j,i,:,:,:] = tmp[imask] else: alp[j,i,:,:,0] = aver.z.__getattribute__(index)[imask,:,iy] for i in range(0,3): for j in range(0,3): index1 = etaformat.format(i+1,j+1,1) index2 = etaformat.format(i+1,j+1,2) # Sign difference with Schrinner + r correction if iy.size>1: tmp=aver.z.__getattribute__(index1)[:,:,iy] eta[0,j,i,:,:,:] = -tmp[imask] tmp=aver.z.__getattribute__(index2)[:,:,iy] eta[1,j,i,:,:,:] = -tmp[imask]*r del(tmp) else: eta[0,j,i,:,:,0] = -aver.z.__getattribute__(index1)[imask,:,iy] eta[1,j,i,:,:,0] = -aver.z.__getattribute__(index2)[imask,:,iy]*r[:,0] # apply the specified averaging or smoothing: 'None' returns unprocessed arrays if callable(timereducer): u=timereducer(u,trargs) alp=timereducer(alp,trargs) eta=timereducer(eta,trargs) if rank==0: print("Old time dimension has length: {0}".format(old_size)) print("New time dimension has length: {0}".format(alp.shape[-3])) # Create output tensors datatype = alp.dtype datashape = [alp.shape[-3], alp.shape[-2], alp.shape[-1], 1] setattr(self,'utensor', np.zeros([3]+datashape,dtype=datatype)) setattr(self,'alpha' , np.zeros([3,3]+datashape,dtype=datatype)) setattr(self,'beta' , np.zeros([3,3]+datashape,dtype=datatype)) setattr(self,'gamma' , np.zeros([3]+datashape,dtype=datatype)) setattr(self,'delta' , np.zeros([3]+datashape,dtype=datatype)) setattr(self,'kappa' , np.zeros([3,3,3]+datashape,dtype=datatype)) setattr(self,'acoef' , np.zeros([3,3]+datashape,dtype=datatype)) setattr(self,'bcoef' , np.zeros([3,3,3]+datashape,dtype=datatype)) """ All tensors need to be reordered nz,ny,nx,nt for efficient writing to disk """ # Calculating a and b matrices self.acoef[:,:,:,:,:,0] = np.copy(alp) self.acoef=np.swapaxes(self.acoef,-4,-1) self.acoef=np.swapaxes(self.acoef,-3,-2) self.bcoef[:,:,:,:,:,:,0] = np.copy(eta) self.bcoef=np.swapaxes(self.bcoef,-4,-1) self.bcoef=np.swapaxes(self.bcoef,-3,-2) irr, ith, iph = 0,1,2 # u-tensor print("Calculating utensor on rank {}".format(rank)) #utensor[:,:,:,:,0] = u[:,:,:,:] - np.mean(u[:,:,:,:],axis=1,keepdims=True) self.utensor[:,:,:,:,0] = u[:,:,:,:] self.utensor=np.swapaxes(self.utensor,-4,-1) self.utensor=np.swapaxes(self.utensor,-3,-2) # Alpha tensor print("Calculating alpha on rank {}".format(rank)) self.alpha[irr,irr,:,:,:,0] = (alp[irr,irr,:,:,:]-eta[ith,ith,irr,:,:,:]/r) self.alpha[irr,ith,:,:,:,0] = 0.5*(alp[ith,irr,:,:,:]+eta[ith,irr,irr,:,:,:]/r+alp[irr,ith,:,:,:]-eta[ith,ith,ith,:,:,:]/r) self.alpha[irr,iph,:,:,:,0] = 0.5*(alp[iph,irr,:,:,:]+alp[irr,iph,:,:,:] - eta[ith,ith,iph,:,:,:]/r) self.alpha[ith,irr,:,:,:,0] = self.alpha[irr,ith,:,:,:,0] self.alpha[ith,ith,:,:,:,0] = (alp[ith,ith,:,:,:]+eta[ith,irr,ith,:,:,:]/r) self.alpha[ith,iph,:,:,:,0] = 0.5*(alp[iph,ith,:,:,:]+alp[ith,iph,:,:,:]+eta[ith,irr,iph,:,:,:]/r) self.alpha[iph,irr,:,:,:,0] = self.alpha[irr,iph,:,:,:,0] self.alpha[iph,ith,:,:,:,0] = self.alpha[ith,iph,:,:,:,0] self.alpha[iph,iph,:,:,:,0] = alp[iph,iph,:,:,:] self.alpha=np.swapaxes(self.alpha,-4,-1) self.alpha=np.swapaxes(self.alpha,-3,-2) # Gamma vector print("Calculating gamma on rank {}".format(rank)) self.gamma[irr,:,:,:,0] = -0.5*(alp[iph,ith,:,:,:]-alp[ith,iph,:,:,:]-eta[ith,irr,iph,:,:,:]/r) self.gamma[ith,:,:,:,0] = -0.5*(alp[irr,iph,:,:,:]-alp[iph,irr,:,:,:]-eta[ith,ith,iph,:,:,:]/r) self.gamma[iph,:,:,:,0] = -0.5*(alp[ith,irr,:,:,:]-alp[irr,ith,:,:,:]+eta[ith,irr,irr,:,:,:]/r +eta[ith,ith,ith,:,:,:]/r) self.gamma=np.swapaxes(self.gamma,-4,-1) self.gamma=np.swapaxes(self.gamma,-3,-2) # Beta tensor print("Calculating beta on rank {}".format(rank)) self.beta[irr,irr,:,:,:,0] = -0.5* eta[ith,iph,irr,:,:,:] self.beta[irr,ith,:,:,:,0] = 0.25*(eta[irr,iph,irr,:,:,:] - eta[ith,iph,ith,:,:,:]) self.beta[irr,iph,:,:,:,0] = 0.25*(eta[ith,irr,irr,:,:,:] - eta[ith,iph,iph,:,:,:] - eta[irr,ith,irr,:,:,:]) self.beta[ith,irr,:,:,:,0] = self.beta[irr,ith,:,:,:,0] self.beta[ith,ith,:,:,:,0] = 0.5*eta[irr,iph,ith,:,:,:] self.beta[ith,iph,:,:,:,0] = 0.25*(eta[ith,irr,ith,:,:,:] + eta[irr,iph,iph,:,:,:] - eta[irr,ith,ith,:,:,:]) self.beta[iph,irr,:,:,:,0] = self.beta[irr,iph,:,:,:,0] self.beta[iph,ith,:,:,:,0] = self.beta[ith,iph,:,:,:,0] self.beta[iph,iph,:,:,:,0] = 0.5*(eta[ith,irr,iph,:,:,:] - eta[irr,ith,iph,:,:,:]) # Sign convention to match with meanfield_e_tensor self.beta = -self.beta self.beta=np.swapaxes(self.beta,-4,-1) self.beta=np.swapaxes(self.beta,-3,-2) # Delta vector print("Calculating delta on rank {}".format(rank)) self.delta[irr,:,:,:,0] = 0.25*(eta[irr,ith,ith,:,:,:] - eta[ith,irr,ith,:,:,:] + eta[irr,iph,iph,:,:,:]) self.delta[ith,:,:,:,0] = 0.25*(eta[ith,irr,irr,:,:,:] - eta[irr,ith,irr,:,:,:] + eta[ith,iph,iph,:,:,:]) self.delta[iph,:,:,:,0] = -0.25*(eta[irr,iph,irr,:,:,:] + eta[ith,iph,ith,:,:,:]) # Sign convention to match with meanfield_e_tensor self.delta = -self.delta self.delta=np.swapaxes(self.delta,-4,-1) self.delta=np.swapaxes(self.delta,-3,-2) # Kappa tensor print("Calculating kappa on rank {}".format(rank)) for i in range(0,3): self.kappa[irr,irr,i,:,:,:,0]= -eta[irr,irr,i,:,:,:] self.kappa[ith,irr,i,:,:,:,0]= -0.5*(eta[ith,irr,i,:,:,:]+eta[irr,ith,i,:,:,:]) self.kappa[iph,irr,i,:,:,:,0]= -0.5* eta[irr,iph,i,:,:,:] self.kappa[irr,ith,i,:,:,:,0]= self.kappa[ith,irr,i,:,:,:,0] self.kappa[ith,ith,i,:,:,:,0]= - eta[ith,ith,i,:,:,:] self.kappa[iph,ith,i,:,:,:,0]= -0.5* eta[ith,iph,i,:,:,:] self.kappa[irr,iph,i,:,:,:,0]= self.kappa[iph,irr,i,:,:,:,0] self.kappa[ith,iph,i,:,:,:,0]= self.kappa[iph,ith,i,:,:,:,0] self.kappa[iph,iph,i,:,:,:,0]= 1e-91 # Sign convention to match with meanfield_e_tensor self.kappa = -self.kappa self.kappa=np.swapaxes(self.kappa,-4,-1) self.kappa=np.swapaxes(self.kappa,-3,-2) setattr(self, 'imask', imask)