commit 7153315dfdeea57ab24c5bb1abcfe44622f24295 Author: Michael Krayer Date: Sat May 15 00:57:03 2021 +0200 init: needs testing/cleanup. Included ucftools here because of some overlap diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3cb42ab --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ +from . import particle +from . import visu +from . import field +from . import ucf \ No newline at end of file diff --git a/field.py b/field.py new file mode 100644 index 0000000..80c72ba --- /dev/null +++ b/field.py @@ -0,0 +1,195 @@ +import numpy +class Field3d: + def __init__(self,data,origin,spacing,period=(None,None,None)): + assert len(origin)==3, "'origin' must be of length 3" + assert len(spacing)==3, "'spacing' must be of length 3" + assert len(period)==3, "'period' must be of length 3" + self.data = numpy.array(data) + self.data.setflags(write=False) + self.numpoints = self.data.shape + self.origin = tuple([float(x) for x in origin]) + self.spacing = tuple([float(x) for x in spacing]) + self.period = tuple(period) + return + + def __str__(self): + str = 'Field3d with\n' + str+= ' dimension {}, {}, {}\n'.format(*self.numpoints) + str+= ' origin {:G}, {:G}, {:G}\n'.format(*self.origin) + str+= ' spacing {:G}, {:G}, {:G}\n'.format(*self.spacing) + str+= ' period {}, {}, {}'.format(*self.period) + return str + + @classmethod + def from_chunk(cls,chunk,gridg,period=(None,None,None)): + '''Initialize Field3d from chunk data and global grid.''' + xg,yg,zg = gridg + ib,jb,kb = chunk['ibeg']-1, chunk['jbeg']-1, chunk['kbeg']-1 + dx,dy,dz = xg[1]-xg[0], yg[1]-yg[0], zg[1]-zg[0] + xo,yo,zo = xg[ib]-chunk['ighost']*dx, yg[jb]-chunk['ighost']*dy, zg[kb]-chunk['ighost']*dz + nx,ny,nz = chunk['data'].shape + assert (chunk['nxl']+2*chunk['ighost'])==nx, "Invalid chunk data: nxl != chunk['data'].shape[0]" + assert (chunk['nyl']+2*chunk['ighost'])==ny, "Invalid chunk data: nyl != chunk['data'].shape[1]" + assert (chunk['nzl']+2*chunk['ighost'])==nz, "Invalid chunk data: nzl != chunk['data'].shape[2]" + return cls(chunk['data'],origin=(xo,yo,zo),spacing=(dx,dy,dz),period=period) + + def dim(self,axis=None): + if axis is None: + return tuple(self.dim(axis=ii) for ii in range(0,3)) + assert axis<3, "'axis' must be one of 0,1,2." + return self.numpoints[axis] + + def endpoint(self,axis): + assert axis<3, "'axis' must be one of 0,1,2." + return self.origin[axis]+(self.numpoints[axis]-1)*self.spacing[axis] + + def derivative(self,axis,keep_origin=False): + if not keep_origin: + # 2nd order FD with h = dx/2 + origin = list(self.origin) + if self.period[axis] is not None: + slhi = [slice(None),slice(None),slice(None)] + sllo = [slice(None),slice(None),slice(None)] + slmi = [slice(None),slice(None),slice(None)] + if numpy.isclose(self.origin[axis],-self.spacing[axis]): + slhi[axis] = slice(1,-1) + sllo[axis] = slice(0,-2) + origin[axis] += 0.5*self.spacing[axis] + else: + slhi[axis] = slice(2,None) + sllo[axis] = slice(1,-1) + origin[axis] -= 0.5*self.spacing[axis] + slmi[axis] = slice(1,-1) + slhi = tuple(slhi) + sllo = tuple(sllo) + slmi = tuple(slmi) + data = numpy.zeros_like(self.data) + data[slmi] = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo]) + Field3d.__update_ghostcells_periodic(data,axis) + else: + slhi = [slice(None),slice(None),slice(None)] + sllo = [slice(None),slice(None),slice(None)] + slhi[axis] = slice(1,None) + sllo[axis] = slice(0,-1) + slhi = tuple(slhi) + sllo = tuple(sllo) + data = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo]) + origin[axis] += 0.5*self.spacing[axis] + else: + # 2nd order FD with h = dx, one-sided at boundaries + # https://numpy.org/doc/stable/reference/generated/numpy.gradient.html + origin = self.origin + if self.period[axis] is not None: + data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=1) + Field3d.__update_ghostcells_periodic(data,axis) + else: + data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=2) + return Field3d(data,origin,self.spacing,self.period) + + def gradient(self,axis,keep_origin=False): + return [self.derivative(axis,keep_origin=keep_origin) for axis in range(0,3)] + + def change_grid(self,origin,spacing,numpoints,EPS_COLLAPSE=1e-12): + # TBD: periodicity + assert all([origin[ii]>=self.origin[ii] for ii in range(0,3)]), "New origin is out of bounds." + endpoint = [origin[ii]+(numpoints[ii]-1)*spacing[ii] for ii in range(0,3)] + assert all([endpoint[ii]<=self.endpoint(ii) for ii in range(0,3)]), "New end point is out of bounds." + data = numpy.zeros(numpoints) + if numpy.allclose(spacing,self.spacing): + # spacing is the same: we can construct universal weights for the stencil + i0,j0,k0 = [numpy.floor((origin[ii]-self.origin[ii])/self.spacing[ii]).astype('int') for ii in range(0,3)] + cx,cy,cz = [numpy.remainder(origin[ii]-self.origin[ii],self.spacing[ii]) for ii in range(0,3)] + c = weights_trilinear((cx,cy,cz)) + for ii in range(0,2): + for jj in range(0,2): + for kk in range(0,2): + if c[ii,jj,kk]>EPS_COLLAPSE: + data += c[ii,jj,kk]*self.data[ + i0+ii:i0+ii+numpoints[0], + j0+jj:j0+jj+numpoints[1], + k0+kk:k0+kk+numpoints[2]] + else: + for ii in range(0,numpoints[0]): + for jj in range(0,numpoints[1]): + for kk in range(0,numpoints[2]): + coordinate = ( + origin[0]+ii*spacing[0], + origin[1]+jj*spacing[1], + origin[2]+kk*spacing[2]) + data[ii,jj,kk] = self.interpolate(coordinate) + return Field3d(data,origin,spacing,self.period) + + def interpolate(self,coordinate,EPS_COLLAPSE=1e-12): + assert all([coordinate[ii]>=self.origin[ii] for ii in range(0,3)]), "'coordinate' is out of bounds." + assert all([coordinate[ii]<=self.endpoint(ii) for ii in range(0,3)]), "'coordinate' is out of bounds." + i0,j0,k0 = [numpy.floor((coordinate[ii]-self.origin[ii])/self.spacing[ii]).astype('int') for ii in range(0,3)] + cx,cy,cz = [numpy.remainder(coordinate[ii]-self.origin[ii],self.spacing[ii]) for ii in range(0,3)] + c = weights_trilinear((cx,cy,cz)) + val = 0.0 + for ii in range(0,2): + for jj in range(0,2): + for kk in range(0,2): + if c[ii,jj,kk]>EPS_COLLAPSE: + val += c[ii,jj,kk]*self.data[i0+ii,j0+jj,k0+kk] + return val + + def move_origin_periodic(self,shift,axis): + '''Moves origin of grid in periodic direction by shift*spacing.''' + assert self.period[axis], "Origin can only be moved in periodic directions." + self.data = numpy.roll(self.data,shift,axis=axis) + self.data.setflags(write=False) + origin = list(self.origin) + origin[axis] += shift*self.spacing[axis] + self.origin = tuple(origin) + return + @staticmethod + def __update_ghostcells_periodic(data,axis): + slsrc = [slice(None),slice(None),slice(None)] + sldst = [slice(None),slice(None),slice(None)] + slsrc[axis] = slice(1,2) + sldst[axis] = slice(-1,None) + data[sldst] = data[slsrc] + slsrc = [slice(None),slice(None),slice(None)] + sldst = [slice(None),slice(None),slice(None)] + slsrc[axis] = slice(-2,-1) + sldst[axis] = slice(0,1) + data[sldst] = data[slsrc] + return + +class ChunkIterator: + '''Iterates through all chunks. 'snapshot' must be an instance + of a class which returns a Field3d from the method call + snapshot.field_chunk(rank,key,keep_ghost=keep_ghost). + One example implementation is UCFSnapshot from suspendtools.ucf.''' + def __init__(self,snapshot,key,keep_ghost=True): + self.snapshot = snapshot + self.key = key + self.keep_ghost = keep_ghost + self.iter_rank = 0 + def __iter__(self): + self.iter_rank = 0 + return self + def __next__(self): + if self.iter_rank=0.0 for ii in range(0,3)]), "'rel_dist' must be >=0" + assert all([rel_dist[ii]<=1.0 for ii in range(0,3)]), "'rel_dist' must be <=1" + cx,cy,cz = rel_dist + c = numpy.zeros((2,2,2)) + c[0,0,0] = 1.0-(cx+cy+cz)+(cx*cy+cx*cz+cy*cz)-(cx*cy*cz) + c[1,0,0] = cx-(cx*cy+cx*cz)+(cx*cy*cz) + c[0,1,0] = cy-(cx*cy+cy*cz)+(cx*cy*cz) + c[0,0,1] = cz-(cx*cz+cy*cz)+(cx*cy*cz) + c[1,1,0] = (cx*cy)-(cx*cy*cz) + c[1,0,1] = (cx*cz)-(cx*cy*cz) + c[0,1,1] = (cy*cz)-(cx*cy*cz) + c[1,1,1] = (cx*cy*cz) + return c diff --git a/particle.py b/particle.py new file mode 100644 index 0000000..62ec46c --- /dev/null +++ b/particle.py @@ -0,0 +1,26 @@ +def sort(pp,col): + ncol,npart,ntime = pp.shape + assert('id' in col) + for itime in range(0,ntime): + idx = pp[col['id'],:,itime].squeeze().argsort() + pp[:,:,itime] = pp[:,idx,itime] + return pp + +def slice_columns(pp,col,keys): + idx_col = [] + col_new = {} + ii = 0 + for key in keys: + idx_col.append(col[key]) + col_new[key] = ii + ii+=1 + return pp[idx_col,:,:], col_new + +def translate_circular(pp,col,translation,bounds,axis=0): + '''Translates particles while taking into account + the bounding box''' + assert(axis<3) + L = bounds[2*axis+1] + keys = ('x','y','z') + pp[col[keys[axis]],:,:] = (pp[col[keys[axis]],:,:]+translation)%L + return pp \ No newline at end of file diff --git a/ucf.py b/ucf.py new file mode 100644 index 0000000..81c09b2 --- /dev/null +++ b/ucf.py @@ -0,0 +1,859 @@ +import warnings +def __warning_format(message, category, filename, lineno, file=None, line=None): + return '%s:%s: %s:%s\n' % (filename, lineno, category.__name__, message) +warnings.formatwarning = __warning_format + +class UCF: + """UCF low-level access class""" + def __init__(self,file=None,verbosity=False,debug=False): + self.__initializeConstants() + self.__resetPublicProperties() + self.__resetPrivateProperties() + self.__resetCurrentStep() + self.__resetCurrentSet() + if file is not None: + self.open(file) + self.Debug = debug + self.Verbosity = verbosity + + def open(self,file): + """Opens an input stream for reading access. The variable 'file' can be of the following types: + str opens a file on disk whose path is specified by 'file' + bytes data which is already located in memory as a bytes or bytearray object (can be used for streams) + """ + from io import BytesIO + import numpy + # Check what class 'file' belongs to and treat it accordingly + if isinstance(file,str): + self.File = file + self.__fileobj = open(self.File,'rb') + self.__inputAvailable = True + elif isinstance(file,bytes): + self.File = 'byte-stream' + self.__fileobj = BytesIO(file) + self.__inputAvailable = True + # Determine file size + self.__fileobj.seek(0,2) + self.FileSize = self.__fileobj.tell() + self.__fileobj.seek(self.__fileBeg,0) + # Read the header of the file + self.__readHeaderFile() + # Scan through file to get the basic structure (steps/sets) + self.__timeStep = numpy.zeros(self.__scanBuffSize,dtype=numpy.float64) + self.__posStep = numpy.zeros(self.__scanBuffSize,dtype=numpy.int64) + self.__numSetPerStep = numpy.zeros(self.__scanBuffSize,dtype=numpy.int32) + istep = 0; + while self.__fileobj.tell()self.__scanBuffSize: + warnings.warn('Buffer overflow detected: increase scanBuffSize.') + self.__timeStep = self.__timeStep[0:nstep] + self.__posStep = self.__posStep[0:nstep] + self.__numSetPerStep = self.__numSetPerStep[0:nstep] + # Set some internal variables + self.NumDataset = numpy.max(self.__numSetPerStep) + self.NumTimestep = nstep + self.__isFileHeaderWritten = True + self.__isStepHeaderWritten = True + + def close(self): + """Closes input file object""" + self.__fileobj.close() + self.__init__() + + def addFileHeaderToBuffer(self,rank=0,rankijk=(0,0,0),ftype=1999): + """Initialize a buffer to generate a new UCF file.""" + self.__bufNumSteps = 0 + self.__bufStep = [] + self.__bufParams = [] + self.__bufData = [] + self.__bufRank = rank + self.__bufRankijk = rankijk + self.__bufFileType = ftype + self.__bufAvailable = True + + def addStepToBuffer(self,step=1,time=0.0): + """Add a new step to buffer.""" + if not self.__bufAvailable: + raise BufferError('Buffer has not been initialized.') + if step>self.__bufNumSteps: + self.__bufStep.extend([None] for ii in range(self.__bufNumSteps,step)) + self.__bufParams.extend([] for ii in range(self.__bufNumSteps,step)) + self.__bufData.extend([] for ii in range(self.__bufNumSteps,step)) + self.__bufNumSteps = step + self.__bufStep[step-1] = time + + def addDatasetToBuffer(self,data,params=None,step=1,dset=1): + import numpy + """Add a new dataset to specified step of buffer.""" + if not self.__bufAvailable: + raise BufferError('Buffer has not been initialized.') + if step>self.__bufNumSteps: + raise ValueError('Requested step does not exist.') + if not hasattr(data,'dtype'): + raise TypeError('Cannot determine datatype of provided data') + if not hasattr(data,'nbytes'): + raise TypeError('Cannot determine number of bytes of provided data') + if not hasattr(data,'tobytes'): + raise TypeError('Cannot convert provided data to bytes') + if not hasattr(data,'shape'): + raise TypeError('Cannot determine shape of provided data') + if params is not None and not all(numpy.issubdtype(type(ii),numpy.integer) for ii in params): + raise TypeError('Parameters must be provided as integer') + nset = len(self.__bufData[step-1]) + if dset>nset: + self.__bufParams[step-1].extend(None for ii in range(nset,dset)) + self.__bufData[step-1].extend(None for ii in range(nset,dset)) + self.__bufParams[step-1][dset-1] = params + self.__bufData[step-1][dset-1] = data + + def copyFileHeaderToBuffer(self): + if not self.__inputAvailable: + raise IOError('No input file available') + self.addFileHeaderToBuffer(rank=self.IORank[0],rankijk=self.IORank[1:],ftype=self.__typeID) + + def copyStepToBuffer(self,step_in,step_out=1,recursive=False,singlePrecision=False): + """Copy a step from an input file to output buffer. If recursive copying is activated, all datasets + within the step will be copied, otherwise only the step header is copied without datasets. + If datasets are copied, the precision can be reduced using the 'singlePrecision' flag.""" + if not self.__inputAvailable: + raise IOError('No input file available') + if not self.__bufAvailable: + raise BufferError('Buffer has not been initialized.') + self.addStepToBuffer(step=step_out,time=self.__timeStep[step_in-1]) + if recursive: + for dset in range(0,self.__numSetPerStep[step_in-1]): + self.copyDatasetToBuffer(step_in,dset+1,step_out=step_out,dset_out=dset+1,singlePrecision=singlePrecision) + + def copyDatasetToBuffer(self,step_in,dset_in,step_out=1,dset_out=1,singlePrecision=False): + """Copy a dataset from an input file to output buffer at specified step. The precision of the + dataset can be reduced using the 'singlePrecision' flag.""" + import numpy + if not self.__inputAvailable: + raise IOError('No input file available') + if not self.__bufAvailable: + raise BufferError('Buffer has not been initialized.') + (data,params) = self.readSet(step_in,dset_in) + if singlePrecision: + if data.dtype==numpy.dtype('float64'): + data = numpy.float32(data) + elif data.dtype==numpy.dtype('int64'): + data = numpy.int32(data) + self.addDatasetToBuffer(data,params=params,step=step_out,dset=dset_out) + + def flushBuffer(self): + """Returns the buffer as a bytes object, which can be written to a file using a file object.""" + from time import time + from sys import stderr + from struct import pack + import numpy as np + # Sanity check and size gathering + sizeStep = [] + sizeSet = [[]] + for step in range(0,self.__bufNumSteps): + nset = len(self.__bufData[step]) + tmpSizeStep = 0 + if nset==0: + warnings.warn('Step #{} in buffer does not contain any dataset.'.format(step+1),RuntimeWarning) + for dset in range(0,nset): + if self.__bufData[step][dset] is None: + raise ValueError('Step #{}, dataset #{} does not contain any data.'.format(step+1,dset+1)) + if self.__bufParams[step][dset] is None: + warnings.warn('No parameters were provided for step #{}, dataset #{}.'.format(step+1,dset+1),RuntimeWarning) + nparam==0 + else: + nparam = len(self.__bufParams[step][dset]) + sizeSet[step].append(self.__bufData[step][dset].nbytes) + tmpSizeStep += (self.__nHeaderSet+nparam)*self.__nByteHeaderSet + tmpSizeStep += self.__bufData[step][dset].nbytes + sizeStep.append(tmpSizeStep) + # Create output buffer + obuff = b'' + # Build file header + magicFile = self.__magicFile + fileVersion = 2 + unixTime = int(time()) + fileType = self.__bufFileType + rank = self.__bufRank + (iproc,jproc,kproc) = self.__bufRankijk + if self.Debug: + print('Write the following file header at {} bytes'.format(len(obuff)),file=stderr) + print((magicFile,fileVersion,unixTime,fileType,rank,iproc,jproc,kproc),file=stderr) + obuff += pack('qqqqqqqq',magicFile,fileVersion,unixTime,fileType,rank,iproc,jproc,kproc) + # Build step header + for step in range(0,self.__bufNumSteps): + if self.Verbosity: + print('Adding step #{} to output buffer'.format(step+1),file=stderr) + magicStep = self.__magicStep + stepBytes = sizeStep[step] + stepTime = self.__bufStep[step] + nset = len(self.__bufData[step]) + if self.Debug: + print('Write the following step header at {} bytes'.format(len(obuff)),file=stderr) + print((magicStep,stepBytes,stepTime,nset),file=stderr) + obuff += pack('qqdq',magicStep,stepBytes,stepTime,nset) + # Build dataset headers + attach data + for dset in range(0,nset): + if self.Verbosity: + print(' dataset #{}'.format(dset+1),file=stderr) + magicSet = self.__magicSet + setSize = sizeSet[step][dset] + nptype = self.__bufData[step][dset].dtype + if nptype==numpy.dtype('int32'): + dtEncoded = 11 + elif nptype==numpy.dtype('int64'): + dtEncoded = 12 + elif nptype==numpy.dtype('float32'): + dtEncoded = 21 + elif nptype==numpy.dtype('float64'): + dtEncoded = 22 + else: + raise TypeError('Data at step #{}, dataset #{} has an invalid datatype.'.format(step+1,dset+1)) + if self.__bufParams[step][dset] is None: + nparam = 0 + else: + nparam = len(self.__bufParams[step][dset]) + if self.Debug: + print('Write the following set header at {} bytes'.format(len(obuff)),file=stderr) + print((magicSet,setSize,dtEncoded,nparam),file=stderr) + print('with parameters:',file=stderr) + print(self.__bufParams[step][dset],file=stderr) + + obuff += pack('qqqq',magicSet,setSize,dtEncoded,nparam) + if nparam!=0: + obuff += pack(nparam*'q',*self.__bufParams[step][dset]) + obuff += self.__bufData[step][dset].tobytes('F') + # Return bytes + return obuff + + def getTime(self,step=1): + assert(step-1self.NumTimestep: + raise ValueError("Out of bounds: timestep. %d, %d" %(tstep,self.NumTimestep)) + if dset>self.__numSetPerStep[tstep-1]: + raise ValueError("Out of bounds: dataset. %d, %d" % (dset,self.NumDataset)) + # Navigate to correct set + self.__fileobj.seek(self.__posStep[tstep-1],0) + for iset in range(0,dset-1): + self.__readHeaderSet() + self.__fileobj.seek(self.__currentSetSize,1) + posHeader = self.__fileobj.tell() + if self.Debug: + print("Found step #%d, set #%d at position %d" % (tstep,dset,posHeader),file=stderr) + return posHeader + + def __initializeConstants(self): + self.__magicFile = 81985529216486895; + self.__magicStep = 11944304052957; + self.__magicSet = 240217520921210; + self.__nHeaderFile = 8; + self.__nHeaderStep = 4; + self.__nHeaderSet = 4; + self.__nByteHeaderFile = 8; + self.__nByteHeaderStep = 8; + self.__nByteHeaderSet = 8; + self.__nSetParamsField = 10; + self.__nSetParamsParticle = 16; + self.__factorMajor = 1000000000; + self.__factorMinor = 1000000; + self.__factorPatch = 1000; + self.__factorTypeIDClass = 1000; + self.__factorTypeIDKind = 10; + self.__typeIDmatlabField = 1999; + self.__typeIDmatlabParticle = 2999; + self.__scanBuffSize = 131072; + def __resetPublicProperties(self): + self.File = '' # file name + self.Type = '' # file type + self.Class = '' # file class + self.Endian = '' # endianess + self.CodeVersion = '' # version of the simulation code + self.UCFVersion = '' # version of the data format ("unified container format") + self.NumDataset = 0 # maximum number of datasets in this file (over all time steps) + self.NumTimestep = 0 # number of time steps in this file + self.FileSize = 0 # file size + self.CreationTime = 0 # time of creation + self.IOMode = '' # file opened in read-only or read-write mode? + self.IORank = 0 # rank of processor + col,row,pln + self.Verbosity = 0 # verbose output? + self.Debug = 0 # debug information? + def __resetPrivateProperties(self): + self.__fileobj = None + self.__fileBeg = 0 + self.__typeID = 0 + self.__creationTimeUnix = '' + self.__versionMajor = 0 + self.__versionMinor = 0 + self.__versionPatch = 0 + self.__versionFile = 0 + self.__posStep = 0 + self.__numSetPerStep = 0 + self.__timeStep = 0 + self.__inputAvailable = False + self.__bufAvailable = False + def __resetCurrentStep(self): + self.__currentStep = 0 + self.__currentStepPosHeader = 0 + self.__currentStepPosData = 0 + self.__currentStepSize = 0 + self.__currentStepTime = 0 + self.__currentStepNumSet = 0 + def __resetCurrentSet(self): + self.__currentSet = 0 + self.__currentSetPosHeader = 0 + self.__currentSetPosData = 0 + self.__currentSetSize = 0 + self.__currentSetDatatype = 0 + self.__currentSetDatatypeNumeric = 0 + self.__currentSetSizeof = 0 + self.__currentSetNumParams = 0 + self.__currentSetParams = 0 + self.__currentSetNumElements = 0 + +class UCFSnapshot: + '''Handles a snapshot.ucf.tar file.''' + def __init__(self,file_tar,file_index=None): + self.handler = Ustar(file_tar,file_index) + self.verbose = False + self.debug = False + file_name_string = '\t'.join(self.handler.file_name) + if 'uvwp' in file_name_string: + self.type = 'uvwp' + elif 'scal' in file_name_string: + self.type = 'scal' + else: + raise ValueError("Archive does not contain 'uvwp' nor 'scal' files.") + self.__nproc = sum(self.type in s for s in self.handler.file_name) + self.__nprocijk = None + self.__buffer_grid = None + self.__buffer_proc = None + self.__buffer_params = None + def read_particles(self): + data = self.handler.read('particles.bin') + return read_particles(data,step=1,verbosity=self.verbose,debug=self.debug) + def read_chunk(self,rank,dset=None,keep_ghost=True): + file_target = self.type+'.{:05d}'.format(rank) + data = self.handler.read(file_target) + return read_chunk(data,step=1,dset=dset,keep_ghost=keep_ghost, + verbosity=self.verbose,debug=self.debug) + def read_grid(self,key=None): + self.__initialize_buffer_grid() + if key is not None: + gidx = self.gidx_from_key(key) + return self.__buffer_grid[gidx] + return self.__buffer_grid + def read_procgrid(self,key=None): + self.__initialize_buffer_proc() + if key is not None: + gidx = self.gidx_from_key(key) + return self.__buffer_proc[gidx] + return self.__buffer_proc + def read_parameters(self): + self.__initialize_buffer_params() + return self.__buffer_params + def field_chunk(self,rank,key,keep_ghost=True): + '''Returns chunk data as Field3d object.''' + from .field import Field3d + # Support a list/tuple of keys: this is not the most efficient + # implementation but convenient + if isinstance(key,(list,tuple)): + return [self.field_chunk(rank,k,keep_ghost=keep_ghost) + for k in key] + dset = self.dset_from_key(key) + chunk = self.read_chunk(rank,dset=dset,keep_ghost=keep_ghost) + grid = self.read_grid(key=key) + dx,dy,dz = (grid[0][1]-grid[0][0], + grid[1][1]-grid[1][0], + grid[2][1]-grid[2][0]) + xo,yo,zo = (grid[0][chunk['ibeg']]-chunk['ighost']*dx, + grid[1][chunk['jbeg']]-chunk['ighost']*dy, + grid[2][chunk['kbeg']]-chunk['ighost']*dz) + # Per convention periodicity requires ghost cells for Field3d + if keep_ghost: + period = tuple(self.period(ii) if self.nproc(axis=ii)==1 + else None for ii in range(0,3)) + else: + period = (None,None,None) + return Field3d(chunk['data'],(xo,yo,zo),(dx,dy,dz),period) + def field_complete(self): + return + def period(self,axis=None): + if axis is None: + return tuple(self.period(axis=ii) for ii in range(0,3)) + else: + assert axis<3, "'axis' must be smaller than 3" + params = self.read_parameters() + key_bound = ('b','d','f') + key_period = ('xperiodic','yperiodic','zperiodic') + period = params.getfloat('geometry',key_bound[axis]) + isPeriodic = params.getboolean('geometry',key_period[axis]) + return period if isPeriodic else None + def nproc(self,axis=None): + if axis is None: + return self.__nproc + else: + assert axis<3, "'axis' must be smaller than 3" + self.__initialize_nprocijk() + return self.__nprocijk[axis] + def ijk_from_rank(self,rank): + assert rank0: + col['rank'] = ioffset; ioffset+=1 + if ihybrid>0: + col['id'] = ioffset; ioffset+=1 + col['x'] = ioffset; ioffset+=1 + col['y'] = ioffset; ioffset+=1 + col['z'] = ioffset; ioffset+=1 + col['r'] = ioffset; ioffset+=1 + col['rho']= ioffset; ioffset+=1 + col['ax'] = ioffset; ioffset+=1 + col['ay'] = ioffset; ioffset+=1 + col['az'] = ioffset; ioffset+=1 + col['u'] = ioffset; ioffset+=1 + col['v'] = ioffset; ioffset+=1 + col['w'] = ioffset; ioffset+=1 + col['ox'] = ioffset; ioffset+=1 + col['oy'] = ioffset; ioffset+=1 + col['oz'] = ioffset; ioffset+=1 + col['fx'] = ioffset; ioffset+=1 + col['fy'] = ioffset; ioffset+=1 + col['fz'] = ioffset; ioffset+=1 + col['tx'] = ioffset; ioffset+=1 + col['ty'] = ioffset; ioffset+=1 + col['tz'] = ioffset; ioffset+=1 + if idem>0: + col['fxc'] = ioffset; ioffset+=1 + col['fyc'] = ioffset; ioffset+=1 + col['fzc'] = ioffset; ioffset+=1 + col['txc'] = ioffset; ioffset+=1 + col['tyc'] = ioffset; ioffset+=1 + col['tzc'] = ioffset; ioffset+=1 + if iscal>0: + for ii in range(0,iscal): + col['s'+str(ii)] = ioffset; ioffset+=1 + col['q'+str(ii)] = ioffset; ioffset+=1 + return col + +def grid_chunk(chunk,gridg): + '''Returns the grid vectors for chunk (including ghost cells if they exist)''' + import numpy + xg,yg,zg = gridg + # Shift indices so that they start from zero + ib = chunk['ibeg']-1 + jb = chunk['jbeg']-1 + kb = chunk['kbeg']-1 + nxl = chunk['nxl'] + nyl = chunk['nyl'] + nzl = chunk['nzl'] + ighost = chunk['ighost'] + if ighost: + nxg = len(xg) + nyg = len(yg) + nzg = len(zg) + dx = xg[1]-xg[0] + dy = yg[1]-yg[0] + dz = zg[1]-zg[0] + xl = numpy.zeros(nxl+2) + yl = numpy.zeros(nyl+2) + zl = numpy.zeros(nzl+2) + xl[0] = xg[ib]-dx + xl[1:-1] = xg[ib:ib+nxl] + xl[-1] = xg[ib+nxl-1]+dx + yl[0] = yg[jb]-dy + yl[1:-1] = yg[jb:jb+nyl] + yl[-1] = yg[jb+nyl-1]+dy + zl[0] = zg[kb]-dz + zl[1:-1] = zg[kb:kb+nzl] + zl[-1] = zg[kb+nzl-1]+dz + else: + xl = xg[ibeg:ib+nxl-1] + yl = yg[jbeg:jb+nyl-1] + zl = zg[kbeg:kb+nzl-1] + return (xl,yl,zl) + +def grid_chunk_origin_spacing(chunk,gridg): + '''Returns origin and spacing of regular grid of chunk data''' + xg,yg,zg = gridg + # Shift indices so that they start from zero + ib = chunk['ibeg']-1 + jb = chunk['jbeg']-1 + kb = chunk['kbeg']-1 + dx = xg[1]-xg[0] + dy = yg[1]-yg[0] + dz = zg[1]-zg[0] + xo = xg[ib]-chunk['ighost']*dx + yo = yg[jb]-chunk['ighost']*dy + zo = zg[kb]-chunk['ighost']*dz + return (xo,yo,zo),(dx,dy,dz) \ No newline at end of file diff --git a/visu.py b/visu.py new file mode 100644 index 0000000..fa429ed --- /dev/null +++ b/visu.py @@ -0,0 +1,141 @@ +def add_domain_bounds(plotter,bounds,color='black',line_width=2): + import pyvista + domain = pyvista.Box(bounds=bounds) + plotter.add_mesh(domain,color=color,style='wireframe',line_width=line_width) + return + +def add_particles(plotter,pp,col,itime=0, + name=None, + color='white', + scalars=None,cmap=None,clim=None, + opacity=1.0, + theta_resolution=10,phi_resolution=10): + ''' + name (str, optional) – The name for the added mesh/actor so that it can be easily updated. If an actor of this name already exists in the rendering window, it will be replaced by the new actor. + + color (string or 3 item list, optional, defaults to white) – + Use to make the entire mesh have a single solid color. Either a string, RGB list, or hex color string. For example: color='white', color='w', color=[1, 1, 1], or color='#FFFFFF'. Color will be overridden if scalars are specified. + + scalars (str, optional) – + Scalars used to “color” the mesh. Accepts a key of the col dictionary. If both color and scalars are None, then the active scalars are used. + + cmap (str, list, optional) – + Name of the Matplotlib colormap to use when mapping the scalars. See available Matplotlib colormaps. Only applicable for when displaying scalars. Requires Matplotlib to be installed. colormap is also an accepted alias for this. If colorcet or cmocean are installed, their colormaps can be specified by name. + You can also specify a list of colors to override an existing colormap with a custom one. For example, to create a three color colormap you might specify ['green', 'red', 'blue'] + + clim (2 item list, optional) – + Color bar range for scalars. Defaults to minimum and maximum of scalars array. Example: [-1, 2]. rng is also an accepted alias for this. + + opacity (float, str, array-like) – + Opacity of the mesh. If a single float value is given, it will be the global opacity of the mesh and uniformly applied everywhere - should be between 0 and 1. A string can also be specified to map the scalars range to a predefined opacity transfer function (options include: ‘linear’, ‘linear_r’, ‘geom’, ‘geom_r’). A string could also be used to map a scalars array from the mesh to the opacity (must have same number of elements as the scalars argument). Or you can pass a custom made transfer function that is an array either n_colors in length or shorter. + ''' + import pyvista + ncol,npart,ntime = pp.shape + assert(itime