def __init__(self, varfile='', datadir='data/', proc=-1, ivar=-1, quiet=False, trimall=False, format='native', param=None, dim=None, index=None, run2D=False, magic=None, setup=None): """ Description: ----------- Read VAR files from pencil code. if proc < 0, then load all data and assemble. otherwise, load VAR file from specified processor. format -- one of (['native', 'n'], ['ieee-le', 'l'], ['ieee-be', 'B']) for byte-ordering Params: ------ varfile='' datadir='data/' proc=-1 ivar=-1 quiet=False trimall=False format='native' param=None dim=None index=None run2D=False Example of usage ------ ff=pc.read_var(trimall=True,ivar=100,magic=['tt','vort']) """ if (setup is not None): datadir = os.path.expanduser(setup.datadir) dim = setup.dim param = setup.param index = setup.index run2D = setup.run2D else: datadir = os.path.expanduser(datadir) if dim is None: ldownsampled = 'VARd' in varfile dim = read_dim(datadir, proc, down=ldownsampled) if param is None: param = read_param(datadir=datadir, quiet=quiet) if index is None: index = read_index(datadir=datadir, down=ldownsampled) if dim.precision == 'D': precision = 'd' else: precision = 'f' if param.lwrite_aux: totalvars = dim.mvar + dim.maux else: totalvars = dim.mvar if 'VARd' in varfile: if param.mvar_down > 0: totalvars = param.mvar_down # Read index.pro to get positions and "names" # of variables in f(mx,my,mz,nvar). # Thomas: seems useless now ? #exec(index) # this loads the indicies. if (not varfile): if ivar < 0: varfile = 'var.dat' else: varfile = 'VAR' + str(ivar) if proc < 0: procdirs = natural_sort( filter(lambda s: s.startswith('proc'), os.listdir(datadir))) else: procdirs = ['proc' + str(proc)] #global array if (not run2D): f = np.zeros((totalvars, dim.mz, dim.my, dim.mx), dtype=precision) else: if dim.ny == 1: f = np.zeros((totalvars, dim.mz, dim.mx), dtype=precision) else: f = np.zeros((totalvars, 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 procdirs: proc = int(directory[4:]) procdim = read_dim(datadir, proc, down=ldownsampled) if (not quiet): #print "reading data from processor %i of %i ..." \ # Python 2 #% (proc, len(procdirs)) # Python 2 print("reading data from processor {0} of {1} ...".format( proc, len(procdirs))) mxloc = procdim.mx myloc = procdim.my mzloc = procdim.mz #read data filename = os.path.join(datadir, directory, varfile) infile = npfile(filename, endian=format) if (not run2D): f_loc = infile.fort_read(precision, shape=(-1, mzloc, myloc, mxloc)) else: if dim.ny == 1: f_loc = infile.fort_read(precision, shape=(-1, mzloc, mxloc)) else: f_loc = infile.fort_read(precision, shape=(-1, myloc, mxloc)) raw_etc = infile.fort_read(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(procdirs) > 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 #endif MPI run #endfor directories loop if (magic is not None): if ('bb' in magic): # Compute the magnetic field before doing trimall. aa = f[index['ax'] - 1:index['az'], ...] self.bb = curl(aa, dx, dy, dz, x, y, z, run2D=param.lwrite_2d) if (trimall): 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 trimall. aa = f[index['ax'] - 1:index['az'], ...] self.jj = curl2(aa, dx, dy, dz, x, y, z) if (trimall): 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 trimall. uu = f[index['ux'] - 1:index['uz'], ...] self.vort = curl(uu, dx, dy, dz, x, y, z, run2D=param.lwrite_2d) if (trimall): if (param.lwrite_2d): 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 trimall: 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, value in index.items(): # print key,value. if key != 'global_gg': setattr(self, key, self.f[value - 1, ...]) # special treatment for vector quantities if 'uu' in index.keys(): self.uu = self.f[index['ux'] - 1:index['uz'], ...] if 'aa' in index.keys(): self.aa = self.f[index['ax'] - 1:index['az'], ...] # Also treat Fcr (from cosmicrayflux) as a vector. if 'fcr' in index.keys(): self.fcr = self.f[index['fcr'] - 1:index['fcr'] + 2, ...] self.fcrx = self.fcr[0] self.fcry = self.fcr[1] self.fcrz = self.fcr[2] 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 trimall (i.e. no additional curl...). self.magic = magic if self.magic is not None: self.__magicAttributes(param)
def __init__(self, varfile='', datadir='data/', proc=-1, ivar=-1, quiet=False, trimall=False, format='native', param=None, dim=None, index=None, run2D=False, magic=None, setup=None): """ Description: ----------- Read VAR files from pencil code. if proc < 0, then load all data and assemble. otherwise, load VAR file from specified processor. format -- one of (['native', 'n'], ['ieee-le', 'l'], ['ieee-be', 'B']) for byte-ordering Params: ------ varfile='' datadir='data/' proc=-1 ivar=-1 quiet=False trimall=False format='native' param=None dim=None index=None run2D=False """ if (setup is not None): datadir = os.path.expanduser(setup.datadir) dim = setup.dim param = setup.param index = setup.index run2D = setup.run2D 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) if dim.precision == 'D': precision = 'd' else: precision = 'f' if param.lwrite_aux: totalvars = dim.mvar+dim.maux else: totalvars = dim.mvar # Read index.pro to get positions and "names" # of variables in f(mx,my,mz,nvar). # Thomas: seems useless now ? #exec(index) # this loads the indicies. if (not varfile): if ivar < 0: varfile = 'var.dat' else: varfile = 'VAR'+str(ivar) if proc < 0: procdirs = natural_sort(filter(lambda s:s.startswith('proc'), os.listdir(datadir))) else: procdirs = ['proc'+str(proc)] #global array if (not run2D): f = np.zeros((totalvars, dim.mz, dim.my, dim.mx), dtype=precision) else: if dim.ny == 1: f = np.zeros((totalvars, dim.mz, dim.mx), dtype=precision) else: f = np.zeros((totalvars, 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 procdirs: proc = int(directory[4:]) procdim = read_dim(datadir, proc) if (not quiet): #print "reading data from processor %i of %i ..." \ # Python 2 #% (proc, len(procdirs)) # Python 2 print("reading data from processor {0} of {1} ...".format(proc, len(procdirs))) mxloc = procdim.mx myloc = procdim.my mzloc = procdim.mz #read data filename = os.path.join(datadir,directory,varfile) infile = npfile(filename, endian=format) if (not run2D): f_loc = infile.fort_read(precision, shape=(-1, mzloc, myloc, mxloc)) else: if dim.ny == 1: f_loc = infile.fort_read(precision, shape=(-1, mzloc, mxloc)) else: f_loc = infile.fort_read(precision, shape=(-1, myloc, mxloc)) raw_etc = infile.fort_read(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(procdirs) > 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 #endif MPI run #endfor directories loop if (magic is not None): if ('bb' in magic): # Compute the magnetic field before doing trimall. aa = f[index['ax']-1:index['az'],...] self.bb = curl(aa,dx,dy,dz,run2D=param.lwrite_2d) if (trimall): 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 trimall. aa = f[index['ax']-1:index['az'],...] self.jj = curl2(aa,dx,dy,dz) if (trimall): 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 trimall. uu = f[index['ux']-1:index['uz'],...] self.vort = curl(uu,dx,dy,dz,run2D=param.lwrite_2d) if (trimall): if (param.lwrite_2d): 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 trimall: 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,value in index.items(): # print key,value. if key != 'global_gg': setattr(self,key,self.f[value-1,...]) # special treatment for vector quantities if 'uu' in index.keys(): self.uu = self.f[index['ux']-1:index['uz'],...] if 'aa' in index.keys(): self.aa = self.f[index['ax']-1:index['az'],...] # Also treat Fcr (from cosmicrayflux) as a vector. if 'fcr' in index.keys(): self.fcr = self.f[index['fcr']-1:index['fcr']+2,...] self.fcrx = self.fcr[0] self.fcry = self.fcr[1] self.fcrz = self.fcr[2] 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 trimall (i.e. no additional curl...). self.magic = magic if self.magic is not None: self.__magicAttributes(param)