1441 lines
68 KiB
Python
1441 lines
68 KiB
Python
import numpy as np
|
||
class Field3d:
|
||
def __init__(self,data,origin,spacing,deep=False):
|
||
assert len(origin)==3, "'origin' must be of length 3"
|
||
assert len(spacing)==3, "'spacing' must be of length 3"
|
||
assert isinstance(data,np.ndarray), "'data' must be numpy.ndarray."
|
||
if deep:
|
||
self.data = data.copy()
|
||
else:
|
||
self.data = data
|
||
self.origin = tuple([float(x) for x in origin])
|
||
self.spacing = tuple([float(x) for x in spacing])
|
||
self.eps_collapse = 1e-4*np.min(spacing)
|
||
self._dim = None
|
||
return
|
||
|
||
def __str__(self):
|
||
str = 'Field3d with\n'
|
||
str+= ' dimension: {}, {}, {}\n'.format(*self.dim())
|
||
str+= ' origin: {:G}, {:G}, {:G}\n'.format(*self.origin)
|
||
str+= ' spacing: {:G}, {:G}, {:G}\n'.format(*self.spacing)
|
||
str+= ' datatype: {}'.format(self.dtype())
|
||
return str
|
||
|
||
def __add__(self,other):
|
||
if isinstance(other,Field3d):
|
||
assert self.has_same_grid(other), "Grid mismatch."
|
||
return Field3d(self.data+other.data,self.origin,self.spacing)
|
||
else:
|
||
return Field3d(self.data+other,self.origin,self.spacing)
|
||
|
||
def __sub__(self,other):
|
||
if isinstance(other,Field3d):
|
||
assert self.has_same_grid(other), "Grid mismatch."
|
||
return Field3d(self.data-other.data,self.origin,self.spacing)
|
||
else:
|
||
return Field3d(self.data-other,self.origin,self.spacing)
|
||
|
||
def __mul__(self,other):
|
||
if isinstance(other,Field3d):
|
||
assert self.has_same_grid(other), "Grid mismatch."
|
||
return Field3d(self.data*other.data,self.origin,self.spacing)
|
||
else:
|
||
return Field3d(self.data*other,self.origin,self.spacing)
|
||
|
||
def __truediv__(self,other):
|
||
if isinstance(other,Field3d):
|
||
assert self.has_same_grid(other), "Grid mismatch."
|
||
return Field3d(self.data/other.data,self.origin,self.spacing)
|
||
else:
|
||
return Field3d(self.data/other,self.origin,self.spacing)
|
||
|
||
def __radd__(self,other):
|
||
return Field3d(other+self.data,self.origin,self.spacing)
|
||
|
||
def __rmul__(self,other):
|
||
return Field3d(other*self.data,self.origin,self.spacing)
|
||
|
||
def __pow__(self,other):
|
||
return Field3d(self.data**other,self.origin,self.spacing)
|
||
|
||
def __iadd__(self,other):
|
||
if isinstance(other,Field3d):
|
||
assert self.has_same_grid(other), "Grid mismatch."
|
||
self.data += other.data
|
||
else:
|
||
self.data += other
|
||
return self
|
||
|
||
def __isub__(self,other):
|
||
if isinstance(other,Field3d):
|
||
assert self.has_same_grid(other), "Grid mismatch."
|
||
self.data -= other.data
|
||
else:
|
||
self.data -= other
|
||
return self
|
||
|
||
def __imul__(self,other):
|
||
if isinstance(other,Field3d):
|
||
assert self.has_same_grid(other), "Grid mismatch."
|
||
self.data *= other.data
|
||
else:
|
||
self.data *= other
|
||
return self
|
||
|
||
def __itruediv__(self,other):
|
||
if isinstance(other,Field3d):
|
||
assert self.has_same_grid(other), "Grid mismatch."
|
||
self.data /= other.data
|
||
else:
|
||
self.data /= other
|
||
return self
|
||
|
||
# TBD: this should return another Field3d object
|
||
# def __getitem__(self,val):
|
||
# assert isinstance(val,tuple) and len(val)==3, "Field3d must be indexed by [ii,jj,kk]."
|
||
# sl = []
|
||
# for x in val:
|
||
# if isinstance(x,int):
|
||
# lo,hi = x,x+1
|
||
# hi = hi if hi!=0 else None
|
||
# sl.append(slice(lo,hi))
|
||
# elif isinstance(x,slice):
|
||
# sl.append(x)
|
||
# else:
|
||
# raise TypeError("Trajectories can only be sliced by slice objects or integers.")
|
||
# return self.data[sl[0],sl[1],sl[2]]
|
||
|
||
@classmethod
|
||
def from_chunk(cls,chunk,gridg):
|
||
'''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))
|
||
|
||
@classmethod
|
||
def from_file(cls,file,name='Field3d'):
|
||
import h5py
|
||
is_open = isinstance(file,(h5py.File,h5py.Group))
|
||
f = file if is_open else h5py.File(file,'r')
|
||
g = f[name]
|
||
origin = tuple(g['origin'])
|
||
spacing = tuple(g['spacing'])
|
||
data = g['data'][:]
|
||
if not is_open: f.close()
|
||
return cls(data,origin,spacing)
|
||
|
||
@classmethod
|
||
def from_snapshot(cls,snap,key,procslice=None,dtype=np.float64):
|
||
'''
|
||
Construct a Field3d from a field saved in an UCF snapshot.
|
||
Use procslice = (ib,ie,jb,je,kb,je) to load only a block
|
||
of the full data.
|
||
'''
|
||
from .helper import eval_procslice, key_grid, rank_from_position
|
||
proc = snap.procgrid()
|
||
origin = snap.origin()
|
||
spacing = snap.spacing()
|
||
periodicity = snap.periodicity()
|
||
bounds = snap.bounds()
|
||
keyg = key_grid(key)
|
||
nxp,nyp,nzp = [len(proc[keyg][2*ii]) for ii in range(3)]
|
||
# Get field dimensions
|
||
(ipb,ipe,jpb,jpe,kpb,kpe),ori,(nxl,nyl,nzl) = eval_procslice(
|
||
proc,origin,spacing,keyg,procslice)
|
||
# Create a Field3d: ghost cells will be included!
|
||
data = np.empty((nxl+2,nyl+2,nzl+2),dtype=dtype)
|
||
field_ = cls(data,[ori[ii]-spacing[ii] for ii in range(3)],spacing)
|
||
# Go through each chunk to be read and construct the field
|
||
for ip_ in range(ipb,ipe+1):
|
||
for jp_ in range(jpb,jpe+1):
|
||
for kp_ in range(kpb,kpe+1):
|
||
# Determine rank of the chunk to be read
|
||
rank_ = rank_from_position(ip_,jp_,kp_,nyp,nzp)
|
||
# Compute bounds of this chunk
|
||
ib_ = proc[keyg][0][ip_]
|
||
ie_ = proc[keyg][1][ip_]
|
||
jb_ = proc[keyg][2][jp_]
|
||
je_ = proc[keyg][3][jp_]
|
||
kb_ = proc[keyg][4][kp_]
|
||
ke_ = proc[keyg][5][kp_]
|
||
nxl_ = ie_-ib_+1
|
||
nyl_ = je_-jb_+1
|
||
nzl_ = ke_-kb_+1
|
||
# Read data and insert it
|
||
subfield = snap.field_chunk(rank_,key)
|
||
field_.insert_subfield(subfield)
|
||
return field_
|
||
|
||
@classmethod
|
||
def allocate(cls,dim,origin,spacing,fill=None,dtype=np.float64):
|
||
'''Allocates an empty field, or a field filled with 'fill'.'''
|
||
assert isinstance(dim,(tuple,list,np.ndarray)) and len(dim)==3,\
|
||
"'dim' must be a tuple/list of length 3."
|
||
assert isinstance(origin,(tuple,list,np.ndarray)) and len(origin)==3,\
|
||
"'origin' must be a tuple/list of length 3."
|
||
assert isinstance(spacing,(tuple,list,np.ndarray)) and len(spacing)==3,\
|
||
"'spacing' must be a tuple/list of length 3."
|
||
if fill is None:
|
||
data = np.empty(dim,dtype=dtype)
|
||
else:
|
||
data = np.full(dim,fill,dtype=dtype)
|
||
return cls(data,origin,spacing)
|
||
|
||
@classmethod
|
||
def pseudo_field(cls,dim,origin,spacing):
|
||
'''Creates a Field3d instance without allocating any memory.'''
|
||
assert isinstance(dim,(tuple,list,np.ndarray)) and len(dim)==3,\
|
||
"'dim' must be a tuple/list of length 3."
|
||
assert isinstance(origin,(tuple,list,np.ndarray)) and len(origin)==3,\
|
||
"'origin' must be a tuple/list of length 3."
|
||
assert isinstance(spacing,(tuple,list,np.ndarray)) and len(spacing)==3,\
|
||
"'spacing' must be a tuple/list of length 3."
|
||
data = np.empty((0,0,0))
|
||
r = cls(data,origin,spacing)
|
||
r._dim = dim
|
||
return r
|
||
|
||
def save(self,file,name='Field3d',truncate=False,single_precision=False):
|
||
import h5py
|
||
is_open = isinstance(file,(h5py.File,h5py.Group))
|
||
flag = 'w' if truncate else 'a'
|
||
f = file if is_open else h5py.File(file,flag)
|
||
g = f.create_group(name)
|
||
g.create_dataset('origin',data=self.origin)
|
||
g.create_dataset('spacing',data=self.spacing)
|
||
if single_precision:
|
||
g.create_dataset('data',data=self.data,dtype='f4')
|
||
else:
|
||
g.create_dataset('data',data=self.data)
|
||
if not is_open: f.close()
|
||
return
|
||
|
||
def copy(self):
|
||
return Field3d(self.data.copy(),self.origin,self.spacing)
|
||
|
||
def pseudo_copy(self):
|
||
return self.pseudo_field(self.dim(),self.origin,self.spacing)
|
||
|
||
def insert_subfield(self,subfield):
|
||
assert all([abs(subfield.spacing[ii]-self.spacing[ii])<self.eps_collapse
|
||
for ii in range(3)]), "spacing differs. Got {}, have {}".format(subfield.spacing,self.spacing)
|
||
assert all([self.distance_to_nearest_gridpoint(subfield.origin[ii],axis=ii)<self.eps_collapse
|
||
for ii in range(3)]), "subfield has shifted origin."
|
||
assert all(self.is_within_bounds(subfield.origin,axis=None)), "subfield origin is out-of-bounds."
|
||
assert all(self.is_within_bounds(subfield.endpoint(),axis=None)), "subfield endpoint is out-of-bounds."
|
||
#ib,jb,kb = [int(round((subfield.origin[ii]-self.origin[ii])/self.spacing[ii])) for ii in range(3)]
|
||
ib,jb,kb = self.nearest_gridpoint(subfield.origin,axis=None)
|
||
nx,ny,nz = subfield.dim()
|
||
self.data[ib:ib+nx,jb:jb+ny,kb:kb+nz] = subfield.data[:,:,:]
|
||
return
|
||
|
||
def extract_subfield(self,idx_origin,dim,stride=(1,1,1),deep=False,strict_bounds=True):
|
||
if strict_bounds:
|
||
assert all(idx_origin[ii]>=0 and idx_origin[ii]<self.dim(axis=ii) for ii in range(3)),\
|
||
"'origin' is out-of-bounds."
|
||
assert all(idx_origin[ii]+stride[ii]*(dim[ii]-1)<self.dim(axis=ii) for ii in range(3)),\
|
||
"endpoint is out-of-bounds."
|
||
else:
|
||
for ii in range(3):
|
||
if idx_origin[ii]<0:
|
||
dim[ii] += idx_origin[ii]
|
||
idx_origin[ii] = 0
|
||
if idx_origin[ii]+stride[ii]*(dim[ii]-1)>=self.dim(axis=ii):
|
||
dim[ii] = (self.dim(axis=ii)-idx_origin[ii])//stride[ii]
|
||
if dim[ii]<=0:
|
||
return None
|
||
sl = tuple(slice(idx_origin[ii],
|
||
idx_origin[ii]+stride[ii]*dim[ii],
|
||
stride[ii]) for ii in range(3))
|
||
origin = self.coordinate(idx_origin)
|
||
spacing = tuple(self.spacing[ii]*stride[ii] for ii in range(3))
|
||
data = self.data[sl]
|
||
return Field3d(data,origin,spacing,deep=deep)
|
||
|
||
def clip(self,position,axis,invert=False,deep=False,is_relative=False):
|
||
'''Extracts a subfield by clipping with a plane with normal pointing
|
||
direction specified by 'axis'.'''
|
||
if is_relative:
|
||
coord = self.origin[axis] + coord*self.dim(axis=axis)*self.spacing[axis]
|
||
idx_clip = self.nearest_gridpoint(coord,axis=axis,lower=True)
|
||
sl = 3*[slice(None)]
|
||
origin_ = self.origin
|
||
spacing_ = self.spacing
|
||
if invert:
|
||
sl[axis] = slice(idx_clip+1,None)
|
||
origin_[axis] = self.origin[axis]+idx_clip*self.spacing[axis]
|
||
else:
|
||
sl[axis] = slice(0,idx_clip+1)
|
||
data_ = self.data[tuple(sl)]
|
||
return Field3d(data_,origin_,spacing_,deep=deep)
|
||
|
||
def clip_box(self,bounds,deep=False,is_relative=False):
|
||
'''Extracts a subfield by clipping with a box.'''
|
||
if is_relative:
|
||
bounds = tuple(self.origin[ii//2] + bounds[ii]*self.dim(ii//2)*self.spacing[ii//2] for ii in range(6))
|
||
idx_lo = self.nearest_gridpoint((bounds[0],bounds[2],bounds[4]),lower=True)
|
||
idx_hi = self.nearest_gridpoint((bounds[1],bounds[3],bounds[5]),lower=True)
|
||
idx_lo = tuple(0 if idx_lo[axis]<0 else idx_lo[axis] for axis in range(3))
|
||
idx_hi = tuple(0 if idx_hi[axis]<0 else idx_hi[axis] for axis in range(3))
|
||
sl_ = tuple([slice(idx_lo[0],idx_hi[0]+1),
|
||
slice(idx_lo[1],idx_hi[1]+1),
|
||
slice(idx_lo[2],idx_hi[2]+1)])
|
||
origin_ = tuple(self.origin[axis]+idx_lo[axis]*self.spacing[axis] for axis in range(3))
|
||
spacing_ = self.spacing
|
||
data_ = self.data[sl_]
|
||
return Field3d(data_,origin_,spacing_,deep=deep)
|
||
|
||
def threshold(self,val,invert=False):
|
||
'''Returns a binary array indicating which grid points are above,
|
||
or in case of invert=True below, a given threshold.'''
|
||
if invert:
|
||
return self.data<val
|
||
else:
|
||
return self.data>=val
|
||
|
||
def coordinate(self,idx,axis=None):
|
||
if axis is None:
|
||
assert len(idx)==3, "If 'axis' is None, 'idx' must be a tuple/list of length 3."
|
||
return tuple(self.coordinate(idx[ii],axis=ii) for ii in range(3))
|
||
assert axis<3, "'axis' must be one of 0,1,2."
|
||
assert idx<self.dim(axis=axis), "'idx' is out-of-bounds."
|
||
return self.origin[axis]+idx*self.spacing[axis]
|
||
|
||
def grid(self,axis=None):
|
||
if axis is None:
|
||
return tuple(self.grid(axis=ii) for ii in range(3))
|
||
assert axis<3, "'axis' must be one of 0,1,2."
|
||
return self.origin[axis]+np.arange(0,self.dim(axis=axis))*self.spacing[axis]
|
||
def x(self): return self.grid(axis=0)
|
||
def y(self): return self.grid(axis=1)
|
||
def z(self): return self.grid(axis=2)
|
||
|
||
def nearest_gridpoint(self,coord,axis=None,lower=False):
|
||
if axis is None:
|
||
assert len(coord)==3, "If 'axis' is None, 'coord' must be a tuple/list of length 3."
|
||
return tuple(self.nearest_gridpoint(coord[ii],axis=ii,lower=lower) for ii in range(3))
|
||
assert axis<3, "'axis' must be one of 0,1,2."
|
||
if lower:
|
||
return np.floor((coord+self.eps_collapse-self.origin[axis])/self.spacing[axis]).astype('int')
|
||
else:
|
||
return np.round((coord-self.origin[axis])/self.spacing[axis]).astype('int')
|
||
|
||
def distance_to_nearest_gridpoint(self,coord,axis=None,lower=False):
|
||
if axis is None:
|
||
assert len(coord)==3, "If 'axis' is None, 'coord' must be a tuple/list of length 3."
|
||
return tuple(self.distance_to_nearest_gridpoint(coord[ii],axis=ii,lower=lower) for ii in range(3))
|
||
assert axis<3, "'axis' must be one of 0,1,2."
|
||
val = np.remainder(coord+self.eps_collapse-self.origin[axis],self.spacing[axis])-self.eps_collapse
|
||
if not lower and val>0.5*self.spacing[axis]:
|
||
val = self.spacing[axis]-val
|
||
return val
|
||
|
||
def is_within_bounds(self,coord,axis=None):
|
||
if axis is None:
|
||
assert len(coord)==3, "If 'axis' is None, 'coord' must be a tuple/list of length 3."
|
||
return tuple(self.is_within_bounds(coord[ii],axis=ii) for ii in range(3))
|
||
assert axis<3, "'axis' must be one of 0,1,2."
|
||
idx_nearest = self.nearest_gridpoint(coord,axis=axis)
|
||
if idx_nearest>0 and idx_nearest<self.dim(axis=axis)-1:
|
||
return True
|
||
dist_nearest = self.distance_to_nearest_gridpoint(coord,axis=axis)
|
||
if (idx_nearest==0 or idx_nearest==self.dim(axis=axis)-1) and abs(dist_nearest)<self.eps_collapse:
|
||
return True
|
||
else:
|
||
return False
|
||
|
||
def has_same_grid(self,other):
|
||
if not self.has_same_origin(other): return False
|
||
if not self.has_same_spacing(other): return False
|
||
if not self.has_same_origin(other): return False
|
||
return True
|
||
|
||
def has_same_origin(self,other,axis=None):
|
||
if axis is None:
|
||
return all([self.has_same_origin(other,axis=ii) for ii in range(3)])
|
||
origin = other.origin if isinstance(other,Field3d) else other
|
||
return abs(origin[axis]-self.origin[axis])<self.eps_collapse
|
||
|
||
def has_same_spacing(self,other,axis=None):
|
||
if axis is None:
|
||
return all([self.has_same_spacing(other,axis=ii) for ii in range(3)])
|
||
spacing = other.spacing if isinstance(other,Field3d) else other
|
||
return abs(spacing[axis]-self.spacing[axis])<self.eps_collapse
|
||
|
||
def has_same_dim(self,other,axis=None):
|
||
if axis is None:
|
||
return all([self.has_same_dim(other,axis=ii) for ii in range(3)])
|
||
dim = other.dim(axis=axis) if isinstance(other,Field3d) else other
|
||
return dim==self.dim(axis=axis)
|
||
|
||
def dim(self,axis=None):
|
||
if axis is None:
|
||
return tuple(self.dim(axis=ii) for ii in range(3))
|
||
assert axis<3, "'axis' must be one of 0,1,2."
|
||
if self._dim is not None:
|
||
return self._dim[axis]
|
||
else:
|
||
return self.data.shape[axis]
|
||
|
||
def endpoint(self,axis=None):
|
||
if axis is None:
|
||
return tuple(self.endpoint(axis=ii) for ii in range(3))
|
||
assert axis<3, "'axis' must be one of 0,1,2."
|
||
return self.origin[axis]+(self.dim(axis=axis)-1)*self.spacing[axis]
|
||
|
||
def dtype(self):
|
||
return self.data.dtype
|
||
|
||
def convert_dtype(self,dtype):
|
||
self.data = self.data.astype(dtype,copy=False)
|
||
return
|
||
|
||
def derivative(self,axis,only_keep_interior=False,shift_origin='before'):
|
||
'''Computes derivative wrt to direction 'axis' with 2nd order finite differences
|
||
centered between the origin grid points'''
|
||
from scipy import ndimage
|
||
assert axis<3, "'axis' must be one of 0,1,2."
|
||
origin = list(self.origin)
|
||
assert shift_origin in ('before','after'), "'shift_origin' must be one of {'before','after'}."
|
||
if shift_origin=='before':
|
||
corr_orig = 0
|
||
origin[axis] -= 0.5*self.spacing[axis]
|
||
else:
|
||
corr_orig = -1
|
||
origin[axis] += 0.5*self.spacing[axis]
|
||
data = ndimage.correlate1d(self.data,np.array([-1.,1.])/self.spacing[axis],axis=axis,
|
||
mode='constant',cval=np.nan,origin=corr_orig)
|
||
if only_keep_interior:
|
||
sl = 3*[slice(None)]
|
||
if shift_origin=='before':
|
||
sl[axis] = slice(-1,None)
|
||
origin[axis] += self.spacing[axis]
|
||
else:
|
||
sl[axis] = slice(0,-1)
|
||
data = data[tuple(sl)]
|
||
return Field3d(data,origin,self.spacing)
|
||
|
||
def gradient(self,axis,preserve_origin=False,only_keep_interior=False,add_border='before'):
|
||
return [self.derivative(axis,preserve_origin=preserve_origin,
|
||
only_keep_interior=only_keep_interior,add_border=add_border) for axis in range(0,3)]
|
||
|
||
def laplacian(self,only_keep_interior=False):
|
||
'''Computes the Laplacian of a field.'''
|
||
from scipy import ndimage
|
||
data = ndimage.correlate1d(self.data,np.array([1.,-2.,1.])/self.spacing[0]**2,axis=0,mode='constant',cval=np.nan,origin=0)
|
||
data += ndimage.correlate1d(self.data,np.array([1.,-2.,1.])/self.spacing[1]**2,axis=1,mode='constant',cval=np.nan,origin=0)
|
||
data += ndimage.correlate1d(self.data,np.array([1.,-2.,1.])/self.spacing[2]**2,axis=2,mode='constant',cval=np.nan,origin=0)
|
||
origin = list(self.origin)
|
||
if only_keep_interior:
|
||
data = data[1:-1,1:-1,1:-1]
|
||
for axis in range(3):
|
||
origin[axis] = origin[axis]+self.spacing[axis]
|
||
return Field3d(data,origin,self.spacing)
|
||
|
||
def integral(self,integrate_axis,average=False,ignore_nan=False,return_weights=False,ufunc=None):
|
||
'''Computes the integral or average along a given axis applying the
|
||
function 'ufunc' to each node.'''
|
||
assert isinstance(integrate_axis,(list,tuple,np.ndarray)) and len(integrate_axis)==3,\
|
||
"'integrate_axis' must be a tuple/list of length 3."
|
||
assert all([isinstance(integrate_axis[ii],(bool,int)) for ii in range(3)]),\
|
||
"'integrate_axis' requires bool values."
|
||
assert any(integrate_axis), "'integrate_axis' must contain at least one True."
|
||
axes = []
|
||
weight = 1.0
|
||
for axis in range(3):
|
||
if integrate_axis[axis]:
|
||
axes.append(axis)
|
||
if average:
|
||
weight *= self.dim(axis=axis)
|
||
else:
|
||
weight *= self.spacing[axis]
|
||
axes = tuple(axes)
|
||
if ignore_nan:
|
||
if average:
|
||
weight = np.sum(~np.isnan(self.data),axis=axes,keepdims=True)
|
||
func_sum = np.nansum
|
||
else:
|
||
func_sum = np.sum
|
||
if ufunc is None:
|
||
out = func_sum(self.data,axis=axes,keepdims=True)
|
||
else:
|
||
assert isinstance(ufunc,np.ufunc), "'ufunc' needs to be a numpy ufunc. "\
|
||
"Check out https://numpy.org/doc/stable/reference/ufuncs.html for reference."
|
||
assert ufunc.nin==1, "Only ufunc with single input argument are supported for now."
|
||
out = func_sum(ufunc(self.data),axis=axes,keepdims=True)
|
||
if return_weights:
|
||
return (out,weight)
|
||
else:
|
||
return out/weight
|
||
|
||
def gaussian_filter(self,sigma,truncate=4.0,only_keep_interior=False):
|
||
'''Applies a gaussian filter: sigma is standard deviation for Gaussian kernel for each axis.'''
|
||
from scipy import ndimage
|
||
assert isinstance(sigma,(tuple,list,np.ndarray)) and len(sigma)==3,\
|
||
"'sigma' must be a tuple/list of length 3"
|
||
# Convert sigma from simulation length scales to grid points as required by ndimage
|
||
sigma_img = tuple(sigma[ii]/self.spacing[ii] for ii in range(3))
|
||
data = ndimage.gaussian_filter(self.data,sigma_img,truncate=truncate,mode='constant',cval=np.nan)
|
||
origin = list(self.origin)
|
||
if only_keep_interior:
|
||
r = self.gaussian_filter_radius(sigma,truncate=truncate)
|
||
data = data[r[0]:-r[0],r[1]:-r[1],r[2]:-r[2]]
|
||
for axis in range(3):
|
||
origin[axis] = origin[axis]+r[axis]*self.spacing[axis]
|
||
return Field3d(data,origin,self.spacing)
|
||
|
||
def gaussian_filter_radius(self,sigma,truncate=4.0):
|
||
'''Radius of Gaussian filter. Stencil width is 2*radius+1.'''
|
||
assert isinstance(sigma,(tuple,list,np.ndarray)) and len(sigma)==3,\
|
||
"'sigma' must be a tuple/list of length 3"
|
||
# Convert sigma from simulation length scales to grid points as required by ndimage
|
||
sigma_img = tuple(sigma[ii]/self.spacing[ii] for ii in range(3))
|
||
radius = []
|
||
for ii in range(3):
|
||
radius.append(int(truncate*sigma_img[ii]+0.5))
|
||
return tuple(radius)
|
||
|
||
def shift_origin(self,rel_shift,only_keep_interior=False):
|
||
'''Shifts the origin of a field by multiple of spacing.'''
|
||
from scipy import ndimage
|
||
assert isinstance(rel_shift,(tuple,list,np.ndarray)) and len(rel_shift)==3,\
|
||
"'shift' must be tuple/list with length 3."
|
||
assert all([rel_shift[ii]>=-1.0-self.eps_collapse and rel_shift[ii]<=1.0+self.eps_collapse for ii in range(3)]),\
|
||
"'shift' must be in (-1.0,1.0). {}".format(rel_shift)
|
||
data = self.data.copy()
|
||
origin = list(self.origin)
|
||
sl = 3*[slice(None)]
|
||
for axis in range(3):
|
||
if abs(rel_shift[axis])<self.eps_collapse:
|
||
continue
|
||
elif rel_shift[axis]>0:
|
||
w = rel_shift[axis] if rel_shift[axis]<=1.0 else 1.0
|
||
weights = (1.0-w,w)
|
||
data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=np.nan,origin=-1)
|
||
origin[axis] += w*self.spacing[axis]
|
||
if only_keep_interior:
|
||
sl[axis] = slice(0,-1)
|
||
else:
|
||
w = rel_shift[axis] if rel_shift[axis]>=-1.0 else -1.0
|
||
weights = (-w,1.0+w)
|
||
data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=np.nan,origin=0)
|
||
origin[axis] += w*self.spacing[axis]
|
||
if only_keep_interior:
|
||
sl[axis] = slice(1,None)
|
||
origin[axis] += self.spacing[axis]
|
||
if only_keep_interior:
|
||
data = data[sl]
|
||
return Field3d(data,origin,self.spacing)
|
||
|
||
def relative_shift(self,field):
|
||
'''Compute the relative shift (in terms of spacing) to shift self onto field.'''
|
||
assert self.has_same_spacing(field), "spacing differs."
|
||
rel_shift = [0.0,0.0,0.0]
|
||
for axis in range(3):
|
||
dist = field.origin[axis]-self.origin[axis]
|
||
if abs(dist)>self.eps_collapse:
|
||
rel_shift[axis] = dist/self.spacing[axis]
|
||
return tuple(rel_shift)
|
||
|
||
def change_grid(self,origin,spacing,dim,padding=None,numpad=1):
|
||
assert all([origin[ii]>=self.origin[ii] for ii in range(0,3)]), "New origin is out of bounds."
|
||
endpoint = [origin[ii]+(dim[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."
|
||
# Allocate (possibly padded array)
|
||
origin_pad,dim_pad,sl_out = padding(origin,spacing,dim,padding,numpad)
|
||
data = np.zeros(dim_pad,dtype=self.data.dtype)
|
||
# Trilinear interpolation
|
||
if np.allclose(spacing,self.spacing):
|
||
# spacing is the same: we can construct universal weights for the stencil
|
||
i0,j0,k0 = self.nearest_gridpoint(origin,axis=None,lower=True)
|
||
cx,cy,cz = [self.distance_to_nearest_gridpoint(origin[ii],axis=ii,lower=True)/self.spacing[ii]
|
||
for ii in range(3)]
|
||
c = self.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]>self.eps_collapse:
|
||
data[sl_out] += c[ii,jj,kk]*self.data[
|
||
i0+ii:i0+ii+dim[0],
|
||
j0+jj:j0+jj+dim[1],
|
||
k0+kk:k0+kk+dim[2]]
|
||
else:
|
||
data_ref = data[sl_out]
|
||
for ii in range(0,dim[0]):
|
||
for jj in range(0,dim[1]):
|
||
for kk in range(0,dim[2]):
|
||
coord = (
|
||
origin[0]+ii*spacing[0],
|
||
origin[1]+jj*spacing[1],
|
||
origin[2]+kk*spacing[2])
|
||
data_ref[ii,jj,kk] = self.interpolate(coord)
|
||
return Field3d(data,origin_pad,spacing)
|
||
|
||
def interpolate(self,coord):
|
||
assert all([coord[ii]>=self.origin[ii] for ii in range(0,3)]), "'coord' is out of bounds."
|
||
assert all([coord[ii]<=self.endpoint(ii) for ii in range(0,3)]), "'coord' is out of bounds."
|
||
i0,j0,k0 = self.nearest_gridpoint(coord,axis=None,lower=True)
|
||
cx,cy,cz = [self.distance_to_nearest_gridpoint(coord[ii],axis=ii,lower=True)/self.spacing[ii]
|
||
for ii in range(3)]
|
||
c = self.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]>self.eps_collapse:
|
||
val += c[ii,jj,kk]*self.data[i0+ii,j0+jj,k0+kk]
|
||
return val
|
||
|
||
def scipy_interpolate(self,pts,oob_fill=None):
|
||
from scipy.interpolate import RegularGridInterpolator
|
||
if oob_fill is None:
|
||
bounds_error = True
|
||
fill_value = None
|
||
else:
|
||
bounds_error = False
|
||
fill_value = oob_fill
|
||
rgi = RegularGridInterpolator(self.grid(),self.data,method='linear',
|
||
bounds_error=bounds_error,fill_value=fill_value)
|
||
return rgi(pts)
|
||
|
||
@staticmethod
|
||
def padding(origin,spacing,dim,padding,numpad):
|
||
if isinstance(numpad,int):
|
||
numpad = np.fill(3,numpad,dtype=int)
|
||
else:
|
||
numpad = np.array(numpad,dtype=int)
|
||
assert len(numpad)==3, "'numpad' must be either an integer or tuple/list of length 3."
|
||
origin_pad = np.array(origin)
|
||
dim_pad = np.array(dim)
|
||
sl_out = [slice(None),slice(None),slice(None)]
|
||
if padding is not None:
|
||
if padding=='before':
|
||
dim_pad += numpad
|
||
origin_pad -= numpad*spacing
|
||
for axis in range(3):
|
||
sl_out[axis] = slice(numpad[axis],None)
|
||
elif padding=='after':
|
||
dim_pad += numpad
|
||
for axis in range(3):
|
||
sl_out[axis] = slice(0,-numpad[axis])
|
||
elif padding=='both':
|
||
dim_pad += 2*numpad
|
||
origin_pad -= numpad*spacing
|
||
for axis in range(3):
|
||
sl_out[axis] = slice(numpad[axis],-numpad[axis])
|
||
else:
|
||
raise ValueError("'padding' must either be None or one of {'before','after','both'}.")
|
||
sl_out = tuple(sl_out)
|
||
origin_pad = tuple(origin_pad)
|
||
dim_pad = tuple(dim_pad)
|
||
return (origin_pad,dim_pad,sl_out)
|
||
|
||
def weights_trilinear(self,rel_dist):
|
||
assert len(rel_dist)==3, "len(rel_dist) must be 3."
|
||
cx,cy,cz = rel_dist
|
||
if cx<0.0 and cx>-self.eps_collapse: cx=0.0
|
||
if cy<0.0 and cy>-self.eps_collapse: cy=0.0
|
||
if cz<0.0 and cz>-self.eps_collapse: cz=0.0
|
||
if cx>1.0 and cx<1.0+self.eps_collapse: cx=1.0
|
||
if cy>1.0 and cy<1.0+self.eps_collapse: cy=1.0
|
||
if cz>1.0 and cz<1.0+self.eps_collapse: cz=1.0
|
||
assert cx>=0.0 and cy>=0.0 and cz>=0.0, "'rel_dist' must be >=0"
|
||
assert cx<=1.0 and cy<=1.0 and cz<=1.0, "'rel_dist' must be <=1"
|
||
c = np.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
|
||
|
||
def set_writable(self,flag):
|
||
self.data.setflags(write=flag)
|
||
return
|
||
|
||
def to_vtk(self,deep=False):
|
||
import pyvista as pv
|
||
mesh = pv.UniformGrid()
|
||
mesh.dimensions = self.dim(axis=None)
|
||
mesh.origin = self.origin
|
||
mesh.spacing = self.spacing
|
||
# order needs to be F no matter how array is stored in memory
|
||
if deep:
|
||
mesh.point_arrays['data'] = self.data.flatten(order='F')
|
||
else:
|
||
if not self.data.flags['F_CONTIGUOUS']:
|
||
self.data = np.asfortranarray(self.data)
|
||
mesh.point_arrays['data'] = self.data.ravel(order='F')
|
||
return mesh
|
||
|
||
def vtk_contour(self,val,deep=False,method='contour'):
|
||
if not isinstance(val,(tuple,list)):
|
||
val = [val]
|
||
return self.to_vtk(deep=deep).contour(val,method=method)
|
||
|
||
def vtk_slice(self,normal,origin,deep=False):
|
||
assert (normal in ('x','y','z') or (isinstance(normal,(tuple,list))
|
||
and len(normal)==3)), "'normal' must be 'x','y','z' or tuple of length 3."
|
||
assert isinstance(origin,(tuple,list)) and len(origin)==3,\
|
||
"'origin' must be tuple of length 3."
|
||
return self.to_vtk(deep=deep).slice(normal=normal,origin=origin)
|
||
|
||
def gaussian_filter_umean_channel(array,spacing,sigma,truncate=4.0):
|
||
'''Applies a Gaussian filter to a numpy array of dimension (1,ny,1) which
|
||
contains the mean streamwise velocity of the channel (or possibly sth else).
|
||
Since yu[0] = c and yu[-1] = d, we can use scipy's mirror settings and don't
|
||
need ghost cells.'''
|
||
from scipy import ndimage
|
||
assert array.ndim==3, "Expected an array with three dimensions/axes."
|
||
assert array.shape[0]==1 and array.shape[1]>1 and array.shape[2]==1,\
|
||
"Expected an array with shape (1,ny,1)."
|
||
sigma_img = sigma/spacing
|
||
array = ndimage.gaussian_filter1d(array,sigma_img,axis=1,truncate=truncate,mode='mirror')
|
||
return array
|
||
|
||
|
||
class Features3d:
|
||
def __init__(self,input,threshold,origin,spacing,periodicity,
|
||
invert=False,has_ghost=False,keep_input=False,
|
||
contour_method='flying_edges',
|
||
report=False):
|
||
assert len(origin)==3, "'origin' must be of length 3"
|
||
assert len(spacing)==3, "'spacing' must be of length 3"
|
||
assert len(periodicity)==3, "'periodicity' must be of length 3"
|
||
assert isinstance(input,np.ndarray), "'input' must be numpy.ndarray."
|
||
# Assign basic properties to class variables
|
||
self.spacing = np.array(spacing,dtype=np.float)
|
||
self.origin = np.array(origin,dtype=np.float)-self.spacing
|
||
self.periodicity = tuple(bool(x) for x in periodicity)
|
||
# If regions are supposed to be inverted, i.e. the interior consists of values
|
||
# smaller than the threshold instead of larger, change the sign of the array.
|
||
sign_invert = -1 if invert else +1
|
||
self._threshold = sign_invert*threshold
|
||
# '_input' is the array which is to be triangulated. Here one trailing ghost cell is
|
||
# required to get closed surfaces. The data will always be copied since it probably
|
||
# needs to be copied for vtk anyway (needs FORTRAN contiguous array) and this way
|
||
# there will never be side effects on the input data. The array is padded with
|
||
# a very high value (cannot use NaN or Inf for vtk) in order to close the contour at
|
||
# the boundaries.
|
||
if has_ghost:
|
||
self.dimensions = input.shape+np.array((2,2,2))
|
||
self._input = np.full(self.dimensions,-1e30,dtype=input.dtype,order='F')
|
||
self._input[1:-1,1:-1,1:-1] = sign_invert*input
|
||
else:
|
||
pw = tuple((0,1) if x else (0,0) for x in periodicity)
|
||
self.dimensions = input.shape+np.array((2,2,2))+np.array([sum(x) for x in pw])
|
||
self._input = np.full(self.dimensions,-1e30,dtype=input.dtype,order='F')
|
||
self._input[1:-1,1:-1,1:-1] = np.pad(sign_invert*input,pw,mode='wrap')
|
||
# Triangulate
|
||
self.triangulate(contour_method=contour_method,
|
||
report=report)
|
||
# Set some state variable
|
||
self._kdtree = None
|
||
self._kdaxis = None
|
||
self._kdradius = None
|
||
# Drop input if requested: this may free memory in case it was copied/temporary.
|
||
if not keep_input: self._input = None
|
||
return
|
||
|
||
@classmethod
|
||
def from_field(cls,fld,threshold,periodicity,invert=False,has_ghost=False,
|
||
keep_input=False,contour_method='flying_edges',
|
||
correct_normals=True,report=False):
|
||
return cls(fld.data,threshold,fld.origin,fld.spacing,periodicity,
|
||
invert=invert,has_ghost=has_ghost,keep_input=keep_input,
|
||
contour_method=contour_method,
|
||
report=report)
|
||
|
||
def triangulate(self,contour_method='flying_edges',report=False):
|
||
import pyvista as pv
|
||
import vtk
|
||
from scipy import ndimage, spatial
|
||
from scipy.spatial import KDTree
|
||
from numba import jit
|
||
from time import time
|
||
# Check if '_input' is available: might have been dropped after initialization
|
||
assert self._input is not None, "'_input' not available. Initialize object with keep_input=True flag."
|
||
# Wrap data for VTK using pyvista
|
||
datavtk = pv.UniformGrid()
|
||
datavtk.dimensions = self.dimensions
|
||
datavtk.origin = self.origin
|
||
datavtk.spacing = self.spacing
|
||
datavtk.point_arrays['data'] = self._input.ravel('F')
|
||
# Compute full contour: due to the padding, the contour will result in closed object.
|
||
# This is necessary for proper inside/outside checks and volume computation. However,
|
||
# for surface area computation, the boundary faces have to be removed. Therefore, we
|
||
# will isolate them and store them in an extra array.
|
||
if report: print('[Features3d.triangulate] computing isocontour using {}...'.format(contour_method))
|
||
contour = datavtk.contour([self._threshold],method=contour_method,compute_scalars=False,compute_gradients=True)
|
||
assert contour.is_all_triangles(), "Contouring produced non-triangle cells."
|
||
points = contour.points
|
||
faces = contour.faces.reshape(-1,4)
|
||
gradients = contour.point_arrays['Gradients'][faces[:,1],:]
|
||
# Python for loops are horribly slow, but for loops are the easiest way to perform
|
||
# some of the necessary check. Therefore let's define some functions to be jit compiled
|
||
# by numba to speed things up.
|
||
@jit(nopython=True)
|
||
def __get_boundary_faces(faces,points,bounds,tolerance):
|
||
assert faces.shape[1]==4
|
||
assert points.shape[1]==3
|
||
assert len(bounds)==6
|
||
assert len(tolerance)==3
|
||
nfaces = faces.shape[0]
|
||
output = np.full((nfaces,),-1,dtype=np.int8)
|
||
for ii in range(nfaces):
|
||
for jj in range(3):
|
||
if (points[faces[ii,1],jj]<bounds[2*jj]+tolerance[jj] and
|
||
points[faces[ii,2],jj]<bounds[2*jj]+tolerance[jj] and
|
||
points[faces[ii,3],jj]<bounds[2*jj]+tolerance[jj]):
|
||
output[ii] = 2*jj
|
||
elif (points[faces[ii,1],jj]>bounds[2*jj+1]-tolerance[jj] and
|
||
points[faces[ii,2],jj]>bounds[2*jj+1]-tolerance[jj] and
|
||
points[faces[ii,3],jj]>bounds[2*jj+1]-tolerance[jj]):
|
||
output[ii] = 2*jj+1
|
||
return output
|
||
@jit(nopython=True)
|
||
def __connect_faces_periodic(face_uni_lo,face_uni_hi,points,dist,idx,search_dist):
|
||
N = len(face_uni_hi)
|
||
output = np.zeros((N,4),dtype=face_uni_lo.dtype)
|
||
jj = 0
|
||
for ii in range(N):
|
||
if dist[ii]<search_dist:
|
||
output[jj,:] = (3,face_uni_hi[ii],face_uni_lo[idx[ii]],face_uni_hi[ii])
|
||
jj += 1
|
||
return output[:jj,:]
|
||
# In order to treat periodicity when labeling the various regions, we connect the
|
||
# overlapping regions with degenerate triangles. Then vtkPolyDataConnectivityFilter
|
||
# returns the correct result right away, and we will just ignore the trailing
|
||
# virtual faces afterwards.
|
||
tol_overlap = 1e-5*self.spacing
|
||
bd_face_tag = __get_boundary_faces(faces,points,tuple(contour.bounds),tol_overlap)
|
||
faces_connect = [faces]
|
||
for axis in range(3):
|
||
if not self.periodicity[axis]: continue
|
||
sl_pln = [0,1,2]
|
||
del sl_pln[axis]
|
||
face_uni_lo = np.unique(faces[bd_face_tag==2*axis,1:].ravel())
|
||
face_uni_hi = np.unique(faces[bd_face_tag==2*axis+1,1:].ravel())
|
||
kdt = KDTree(points[np.ix_(face_uni_lo,sl_pln)],leafsize=10,compact_nodes=False,
|
||
copy_data=False,balanced_tree=False)
|
||
dist,idx = kdt.query(points[np.ix_(face_uni_hi,sl_pln)],k=1,
|
||
distance_upper_bound=tol_overlap[axis])
|
||
faces_connect += [__connect_faces_periodic(face_uni_lo,face_uni_hi,points,
|
||
dist,idx,tol_overlap[axis])]
|
||
# Save original number of faces and add the virtual connecting faces to the polydata.
|
||
nfaces = faces.shape[0]
|
||
contour.faces = np.concatenate(faces_connect,axis=0)
|
||
# Compute the connectivity using pure VTK (pyvista only supports vtkConnectivityFilter
|
||
# which takes around twice as long).
|
||
alg = vtk.vtkPolyDataConnectivityFilter()
|
||
alg.SetInputData(contour)
|
||
alg.SetExtractionModeToAllRegions()
|
||
alg.SetColorRegions(True)
|
||
alg.Update()
|
||
contour = pv.filters._get_output(alg)
|
||
# RegionIds are now stored as point data. To efficiently convert it to cell data, the first
|
||
# point of each cell determines the value for this cell.
|
||
regionid = contour.point_arrays['RegionId'][contour.faces.reshape(-1,4)[:nfaces,1]].ravel()
|
||
nfeatures = np.amax(regionid)+1
|
||
# Store the face/vertex data in class arrays. Sort the faces by RegionId for quick access
|
||
# later on.
|
||
if report: print('[Features3d.triangulate] sorting faces by labels...')
|
||
self._points = np.array(points)
|
||
self._nfeatures = nfeatures
|
||
idx = np.flatnonzero(bd_face_tag<0)
|
||
idxbd = np.flatnonzero(bd_face_tag>=0)
|
||
self._offset = np.zeros((2*nfeatures+1,),dtype=np.int64)
|
||
self._offset[1:nfeatures+1] = np.cumsum(np.bincount(regionid[idx],minlength=nfeatures))
|
||
self._offset[nfeatures+1:] = np.cumsum(np.bincount(regionid[idxbd],minlength=nfeatures))
|
||
self._offset[nfeatures+1:] += self._offset[nfeatures]
|
||
self._faces = np.concatenate((
|
||
faces[idx,:][np.argsort(regionid[idx]),:],
|
||
faces[idxbd,:][np.argsort(regionid[idxbd]),:]),axis=0)
|
||
gradients = np.concatenate((
|
||
gradients[idx,:][np.argsort(regionid[idx]),:],
|
||
gradients[idxbd,:][np.argsort(regionid[idxbd]),:]),axis=0)
|
||
# Compute the volume and area per cell.
|
||
if report: print('[Features3d.triangulate] calculating surface area and volume...')
|
||
A = self._points[self._faces[:,1],:]
|
||
B = self._points[self._faces[:,2],:]
|
||
C = self._points[self._faces[:,3],:]
|
||
cn = np.cross(B-A,C-A)
|
||
# Check if cell normal points in direction of gradient. If not, switch vertex order.
|
||
idx = (gradients*cn).sum(axis=-1)>0
|
||
self._faces[np.ix_(idx,[2,3])] = self._faces[np.ix_(idx,[3,2])]
|
||
cn[idx,:] = -cn[idx,:]
|
||
# Compute area and signed volume per cell
|
||
cc = (A+B+C)/3
|
||
cell_areas = 0.5*np.sqrt(np.square(cn).sum(axis=1))
|
||
cell_volumes = 0.5*np.mean(cn*cc,axis=1)
|
||
@jit(nopython=True)
|
||
def __sum_up_cell_data(cell_data,offset,incl_boundary):
|
||
nfeat = (len(offset)-1)//2
|
||
niter = 2*nfeat if incl_boundary else nfeat
|
||
output = np.zeros(nfeat,dtype=cell_data.dtype)
|
||
for ifeat in range(niter):
|
||
for ii in range(offset[ifeat],offset[ifeat+1]):
|
||
output[ifeat%nfeat] += cell_data[ii]
|
||
return output
|
||
self._areas = __sum_up_cell_data(cell_areas, self._offset,False)
|
||
self._volumes = __sum_up_cell_data(cell_volumes,self._offset,True )
|
||
return
|
||
|
||
@property
|
||
def threshold(self): return self._threshold
|
||
|
||
@property
|
||
def nfeatures(self): return self._nfeatures
|
||
|
||
@property
|
||
def areas(self): return self._areas
|
||
|
||
@property
|
||
def volumes(self): return self._volumes
|
||
|
||
def reduce_noise(self,threshold=1e-5,is_relative=True,report=False):
|
||
'''Discards all objects with smaller volume than a threshold.'''
|
||
if is_relative:
|
||
vol_domain = np.prod(self.spacing*self.dimensions)
|
||
threshold = threshold*vol_domain
|
||
features = np.flatnonzero(np.abs(self._volumes)<threshold)
|
||
if len(features)>0:
|
||
self.discard_features(features,report=report)
|
||
return
|
||
|
||
def discard_features(self,features,clean_points=False,report=False):
|
||
'''Removes features from triangulated data.'''
|
||
features = self.list_of_features(features)
|
||
# Get index ranges which are to be deleted and also create an array
|
||
# which determines the size of the block to be deleted
|
||
print('Discarding',features)
|
||
idx = []
|
||
gapsize = np.zeros((self._nfeatures,),dtype=np.int)
|
||
for feature in features:
|
||
rng_ = np.arange(self._offset[feature],self._offset[feature+1])
|
||
idx.append(rng_)
|
||
gapsize[feature] = len(rng_)
|
||
idx = np.concatenate(idx,axis=0)
|
||
# Save former number of faces for report
|
||
nfaces = self._faces.shape[0]
|
||
# Delete indexed elements from arrays
|
||
self._faces = np.delete(self._faces,idx,axis=0)
|
||
self._areas = np.delete(self._areas,features,axis=0)
|
||
self._volumes = np.delete(self._volumes,features,axis=0)
|
||
# Correct offset
|
||
self._offset[1:] = self._offset[1:]-np.cumsum(gapsize)
|
||
self._offset = np.delete(self._offset,features,axis=0)
|
||
self._nfeatures = self._offset.shape[0]-1
|
||
if report:
|
||
print('[Features3d.discard_features]',end=' ')
|
||
print('discarded {:} features with {:} faces.'.format(
|
||
len(features),nfaces-self._faces.shape[0]))
|
||
if clean_points:
|
||
self.clean_points(report=report)
|
||
return
|
||
|
||
def clean_points(self,report=False):
|
||
'''Removes points which are not referenced by any face.'''
|
||
nfaces = self._faces.shape[0]
|
||
ind,inv = np.unique(self._faces[:,1:],return_inverse=True)
|
||
if report:
|
||
print('[Features3d.clean_points]',end=' ')
|
||
print('removed {:} orphan points.'.format(self._points.shape[0]-len(ind)))
|
||
self._faces[:,1:4] = inv.reshape(nfaces,3)
|
||
self._points = self._points[ind,:]
|
||
return
|
||
|
||
def area(self,feature):
|
||
'''Returns the surface area of feature. If feature is None, total surface
|
||
area of all features is returned.'''
|
||
if feature is None:
|
||
return np.sum(self.areas)
|
||
else:
|
||
return self.areas[feature]
|
||
|
||
def volume(self,feature):
|
||
'''Returns volume enclosed by feature. If feature isNone, total volume of
|
||
all features is returned.'''
|
||
if feature is None:
|
||
return np.sum(self.volumes)
|
||
else:
|
||
return self.volumes[feature]
|
||
|
||
def build_kdtree(self,kdaxis=0,leafsize=100,compact_nodes=False,balanced_tree=False,report=False):
|
||
'''Builds a KD-tree for feature search.'''
|
||
from scipy.spatial import KDTree
|
||
from numba import jit
|
||
@jit(nopython=True,cache=True)
|
||
def __get_center_and_search_radius(points,faces,kdaxis):
|
||
nfaces = faces.shape[0]
|
||
center = np.zeros((nfaces,2),dtype=points.dtype)
|
||
radius = 0.0
|
||
if kdaxis==0: i1,i2 = 1,2
|
||
elif kdaxis==1: i1,i2 = 0,2
|
||
elif kdaxis==2: i1,i2 = 0,1
|
||
for iface in range(nfaces):
|
||
min1 = min(min(points[faces[iface,1],i1],points[faces[iface,2],i1]),points[faces[iface,3],i1])
|
||
max1 = max(max(points[faces[iface,1],i1],points[faces[iface,2],i1]),points[faces[iface,3],i1])
|
||
min2 = min(min(points[faces[iface,1],i2],points[faces[iface,2],i2]),points[faces[iface,3],i2])
|
||
max2 = max(max(points[faces[iface,1],i2],points[faces[iface,2],i2]),points[faces[iface,3],i2])
|
||
center[iface,0] = 0.5*(max1+min1)
|
||
center[iface,1] = 0.5*(max2+min2)
|
||
radius = max(radius,(max1-min1)**2+(max2-min2)**2)
|
||
radius = 0.5*np.sqrt(radius)
|
||
return center,radius
|
||
center,radius = __get_center_and_search_radius(self._points,self._faces,kdaxis)
|
||
self._kdtree = KDTree(center,leafsize=leafsize,compact_nodes=compact_nodes,
|
||
copy_data=False,balanced_tree=balanced_tree)
|
||
self._kdaxis = kdaxis
|
||
self._kdradius = radius
|
||
if report:
|
||
print('[Features3d.build_kdtree]',end=' ')
|
||
print('KD-Tree built for axis {}.'.format(kdaxis))
|
||
|
||
def inside_feature(self,coords,report=True):
|
||
'''Find nearest triangle from point R in direction ±dR and its orientation w.r.t dR.
|
||
The algorithm determines the closest intersection of a ray cast from the origin of the
|
||
point in the specified direction (with positiv and negative sign). The ray is supposed to
|
||
be sent in a direction perpendicular to any boundary. Periodicity does not matter here.
|
||
The ray-triangle intersection is implemented by the Möller–Trumbore algorithm (see method
|
||
'ray_triangle_intersection()' for details). Whether a hit occured from the interior
|
||
or the exterior is determined by the direction of the surface normal of the triangle
|
||
and the direction of the ray.
|
||
This method is very similar to 'ray_triangle_intersection()', but takes advantage of
|
||
the fact that only the nearest intersection is to be returned.
|
||
Input:
|
||
R, dR: the point to be checked and the direction of the ray as (3,) numpy arrays
|
||
A,B,C: vertices of N triangles as (N,3) numpy arrays
|
||
Returns:
|
||
is_inside: is point inside? [bool]
|
||
t: parameter to determine intersection point (x = R+t*dR) [float]
|
||
hit_dir: direction from which the triangle was hit, from inward/outward = +1,-1 [int]
|
||
'''
|
||
from numba import jit
|
||
#
|
||
coords = np.array(coords)
|
||
assert coords.ndim==2 and coords.shape[1]==3, "'coords' need to be provided as Nx3 array."
|
||
#
|
||
if not self._kdtree:
|
||
self.build_kdtree()
|
||
if self._kdaxis==0: query_axis = [1,2]
|
||
elif self._kdaxis==1: query_axis = [0,2]
|
||
elif self._kdaxis==2: query_axis = [0,1]
|
||
#
|
||
candidates = self._kdtree.query_ball_point(coords[:,query_axis],self._kdradius).tolist()
|
||
cand_num = np.asarray(tuple(map(len,candidates)))
|
||
cand_arr = np.concatenate(candidates).astype(np.int)
|
||
#
|
||
raydir = np.zeros((3,),dtype=self._points.dtype)
|
||
raydir[self._kdaxis] = 1.0
|
||
#
|
||
@jit(nopython=True,error_model='numpy',cache=True)
|
||
def __ray_triangle_intersect(R,dR,A,B,C):
|
||
'''Implements the Möller–Trumbore ray-triangle intersection algorithm. Taken from
|
||
https://stackoverflow.com/questions/42740765/intersection-between-line-and-triangle-in-3d'''
|
||
E1 = B-A
|
||
E2 = C-A
|
||
N = np.cross(E1,E2)
|
||
norm = np.sqrt(np.sum(N*N))
|
||
det = -np.sum(dR*N)
|
||
invdet = 1.0/det
|
||
AO = R-A
|
||
DAO = np.cross(AO,dR)
|
||
u = np.sum(E2*DAO)*invdet
|
||
v = -np.sum(E1*DAO)*invdet
|
||
t = np.sum(AO*N) *invdet
|
||
return (np.abs(det)>=1e-7*norm and t>=0.0 and u>=0 and v>=0 and (u+v)-1.0<=0), t
|
||
@jit(nopython=True,cache=True)
|
||
def __get_nearest_face(coords,raydir,cand_arr,cand_num,faces,points):
|
||
N = coords.shape[0]
|
||
nearface = np.zeros(N,dtype=np.int64)
|
||
isinside = np.zeros(N,dtype=np.bool_)
|
||
jj = 0
|
||
for ii in range(N):
|
||
t = np.inf
|
||
f = -1
|
||
n = 0
|
||
for kk in range(cand_num[ii]):
|
||
idx = cand_arr[jj]
|
||
hit_,t_ = __ray_triangle_intersect(coords[ii],raydir,
|
||
points[faces[idx,1],:],
|
||
points[faces[idx,2],:],
|
||
points[faces[idx,3],:])
|
||
if hit_:
|
||
n += 1
|
||
if t_<t:
|
||
t = t_
|
||
f = idx
|
||
jj += 1
|
||
nearface[ii] = f
|
||
isinside[ii] = (n%2)!=0
|
||
return nearface,isinside
|
||
nf,isinside = __get_nearest_face(coords,raydir,cand_arr,cand_num,self._faces,self._points)
|
||
output = self.feature_from_face(nf)
|
||
output[np.logical_not(isinside)] = -1
|
||
if report:
|
||
print('[Features3d.inside_feature]',end=' ')
|
||
print('{} of {} points are located inside of features.'.format(sum(isinside),coords.shape[0]))
|
||
return output
|
||
|
||
def feature_from_face(self,idx_cell):
|
||
'''Gets feature ID for a given cell.'''
|
||
# This is an optimized solution for monotonically increasing arrays:
|
||
# as long as offset is smaller than the cell index, the boolean operation
|
||
# returns False. As soon as the first True is returned, i.e.
|
||
# as soon as idx_cell is equal or larger than offset, the argmax function
|
||
# short circuits and returns the index of the first occurence without
|
||
# checking any value afterwards.
|
||
return (np.searchsorted(self._offset,idx_cell,side='right')-1)%self._nfeatures
|
||
|
||
def faces_from_feature(self,features,incl_boundary=False):
|
||
'''Returns indices of cells which belong to given features.'''
|
||
features = self.list_of_features(features)
|
||
faces = []
|
||
for feature in features:
|
||
idx = np.arange(self._offset[feature],self._offset[feature+1])
|
||
faces.append(self._faces[idx,:])
|
||
if incl_boundary:
|
||
for feature in features:
|
||
idx = np.arange(self._offset[self._nfeatures+feature],
|
||
self._offset[self._nfeatures+feature+1])
|
||
faces.append(self._faces[idx,:])
|
||
return np.concatenate(faces,axis=0)
|
||
|
||
def regionid_from_feature(self,features,incl_boundary=False):
|
||
'''Returns regionid cell array to be added to PolyData.'''
|
||
features = self.list_of_features(features)
|
||
regionid = []
|
||
for feature in features:
|
||
nfaces = self._offset[feature+1]-self._offset[feature]
|
||
regionid.append(np.full(nfaces,feature,dtype=np.int64))
|
||
if incl_boundary:
|
||
for feature in features:
|
||
nfaces = (self._offset[self._nfeatures+feature+1]-
|
||
self._offset[self._nfeatures+feature])
|
||
regionid.append(np.full(nfaces,feature,dtype=np.int64))
|
||
return np.concatenate(regionid,axis=0)
|
||
|
||
def nfaces_of_feature(self,features,incl_boundary=False):
|
||
'''Returns number of cells which belong to given features. Can be used
|
||
to construct new offsets.'''
|
||
features = self.list_of_features(features)
|
||
nfaces = []
|
||
for feature in features:
|
||
nfaces.append(self._offset[feature+1]-self._offset[feature])
|
||
if incl_boundary:
|
||
for feature in features:
|
||
nfaces.append(self._offset[self._nfeatures+feature+1]-
|
||
self._offset[self._nfeatures+feature])
|
||
return np.concatenate(nfaces,axis=0)
|
||
|
||
def list_of_features(self,features):
|
||
'''Ensures that 'features' is a list.'''
|
||
from collections.abc import Iterable
|
||
if features is None:
|
||
return np.arange(0,self._nfeatures)
|
||
if not isinstance(features,Iterable):
|
||
return [features]
|
||
else:
|
||
return list(features)
|
||
|
||
@staticmethod
|
||
def cmap_features(nfeatures,name='tab10'):
|
||
from matplotlib.colors import ListedColormap
|
||
if name=='tab10': # seaborn/matplotlib tab10
|
||
cmap = [(0.12156862745098039, 0.4666666666666667, 0.7058823529411765),
|
||
(1.0, 0.4980392156862745, 0.054901960784313725),
|
||
(0.17254901960784313, 0.6274509803921569, 0.17254901960784313),
|
||
(0.8392156862745098, 0.15294117647058825, 0.1568627450980392),
|
||
(0.5803921568627451, 0.403921568627451, 0.7411764705882353),
|
||
(0.5490196078431373, 0.33725490196078434, 0.29411764705882354),
|
||
(0.8901960784313725, 0.4666666666666667, 0.7607843137254902),
|
||
(0.4980392156862745, 0.4980392156862745, 0.4980392156862745),
|
||
(0.7372549019607844, 0.7411764705882353, 0.13333333333333333),
|
||
(0.09019607843137255, 0.7450980392156863, 0.8117647058823529)]
|
||
else:
|
||
raise NotImplementedError('Colormap not implemented.')
|
||
ncolor = len(cmap)
|
||
icolor = 0
|
||
output = np.empty((nfeatures,3))
|
||
for ii in range(nfeatures):
|
||
output[ii] = cmap[icolor]
|
||
icolor = (icolor+1)%ncolor
|
||
return ListedColormap(output)
|
||
|
||
def to_vtk(self,features,incl_boundary=False,incl_regionid=False,remap_regionid=True):
|
||
import pyvista as pv
|
||
faces = self.faces_from_feature(features,incl_boundary=incl_boundary)
|
||
output = pv.PolyData(self._points,faces)
|
||
if incl_regionid:
|
||
output.cell_arrays['RegionId'] = self.regionid_from_feature(features,incl_boundary=incl_boundary)
|
||
if remap_regionid:
|
||
output.cell_arrays['RegionId'] = np.unique(output.cell_arrays['RegionId'],return_inverse=True)[1]
|
||
return output
|
||
|
||
# def to_vtk2(self,features):
|
||
# import pyvista as pv
|
||
# idx = self.faces_from_feature(features)
|
||
# faces_ = self._faces[idx,:]
|
||
# A = self._points[faces_[:,1],:]
|
||
# B = self._points[faces_[:,2],:]
|
||
# C = self._points[faces_[:,3],:]
|
||
# cn = np.cross(B-A,C-A)
|
||
# cn = cn/np.sqrt(np.square(cn).sum(axis=-1,keepdims=True))
|
||
# cc = (A+B+C)/3
|
||
# # print(A[0,:],B[0,:],C[0,:],cn[0,:])
|
||
# # print(A.shape,B.shape,C.shape,cn.shape,faces_.shape,self._faces.shape)
|
||
# output = pv.PolyData(cc)
|
||
# output['cellnormal'] = cn
|
||
# return output
|
||
|
||
# def vtk_faces(self,face_idx):
|
||
# import pyvista as pv
|
||
# return pv.PolyData(self._points,self._faces[face_idx,:])
|
||
|
||
class BinaryFieldNd:
|
||
def __init__(self,input,periodicity,connect_diagonals=False,deep=False,has_ghost=False):
|
||
assert isinstance(input,np.ndarray) and input.dtype==bool,\
|
||
"'input' must be a numpy array of type bool."
|
||
assert all([isinstance(x,(bool,int)) for x in periodicity]),\
|
||
"'periodicity' requires bool values."
|
||
assert len(periodicity)==input.ndim,\
|
||
"Number of entries in 'periodicity' must match dimension of binary field."
|
||
if has_ghost and deep:
|
||
self._data = input.copy()
|
||
elif has_ghost:
|
||
self._data = input
|
||
else:
|
||
pw = tuple((0,1) if x else (0,0) for x in periodicity)
|
||
self._data = np.pad(input,pw,mode='wrap')
|
||
self._sldata = tuple(slice(0,-1) if x else None for x in periodicity)
|
||
self._dim = self._data.shape
|
||
self._ndim = self._data.ndim
|
||
self._labels = None
|
||
self.nlabels = None
|
||
self.wrap = tuple(self._ndim*[None])
|
||
self.periodicity = tuple(bool(x) for x in periodicity)
|
||
self.connect_diagonals = connect_diagonals
|
||
|
||
@property
|
||
def data(self):
|
||
return self._data[self._sldata]
|
||
|
||
@property
|
||
def labels(self):
|
||
if self._labels is None:
|
||
return None
|
||
return self._labels[self._sldata]
|
||
|
||
def label(self,use_cc3d=False):
|
||
'''Labels connected regions in binary fields.'''
|
||
if use_cc3d:
|
||
import cc3d
|
||
if self._ndim==2: connectivity = 8 if self.connect_diagonals else 4
|
||
elif self._ndim==3: connectivity = 18 if self.connect_diagonals else 6
|
||
else: raise RuntimeError("'use_cc3d' can only be used with 2D or 3D data.")
|
||
else:
|
||
from scipy import ndimage
|
||
if self.connect_diagonals: structure = ndimage.generate_binary_structure(self._ndim,self._ndim)
|
||
else: structure = ndimage.generate_binary_structure(self._ndim,1)
|
||
#
|
||
if use_cc3d:
|
||
self._labels,self.nlabels = cc3d.connected_components(self._data,connectivity=connectivity,return_N=True)
|
||
else:
|
||
self._labels,self.nlabels = ndimage.label(self._data,structure=structure)
|
||
if not any(self.periodicity):
|
||
return
|
||
# Get a mapping of labels which differ at periodic overlap
|
||
map_ = np.array(range(self.nlabels+1),dtype=self._labels.dtype)
|
||
wrap_ = self._ndim*[None]
|
||
for axis in range(self._ndim):
|
||
if not self.periodicity[axis]: continue
|
||
sl_lo = tuple(slice(0,1) if ii==axis else slice(None) for ii in range(self._ndim))
|
||
sl_hi = tuple(slice(-1,None) if ii==axis else slice(None) for ii in range(self._ndim))
|
||
sl_pre = tuple(slice(-2,-1) if ii==axis else slice(None) for ii in range(self._ndim))
|
||
lab_lo = self._labels[sl_lo]
|
||
lab_hi = self._labels[sl_hi]
|
||
lab_pre = np.unique(self._labels[sl_pre]) # all labels in last (unwrapped) slice
|
||
# Initialize array to keep track of wrapping
|
||
wrap_[axis] = np.zeros(self.nlabels+1,dtype=bool)
|
||
# Determine new label and map
|
||
lab_new = np.minimum(lab_lo,lab_hi)
|
||
for lab_ in [lab_lo,lab_hi]:
|
||
li = (lab_!=lab_new) #if not map_to_zero else (lab_!=0)
|
||
lab_li = lab_[li]
|
||
lab_new_li = lab_new[li]
|
||
for idx_ in np.unique(lab_li,return_index=True)[1]:
|
||
source_ = lab_li[idx_] # the label to be changed
|
||
target_ = lab_new_li[idx_] # the label which will be newly assigned
|
||
while target_ != map_[target_]: # map it recursively
|
||
target_ = map_[target_]
|
||
map_[source_] = target_
|
||
if source_ in lab_pre: # check if source is not a ghost
|
||
wrap_[axis][target_] = True
|
||
# Remove gaps from target mapping
|
||
idx_,map_ = np.unique(map_,return_index=True,return_inverse=True)[1:3]
|
||
# Relabel
|
||
self._labels = map_[self._labels]
|
||
self.nlabels = np.max(map_)
|
||
self.wrap = tuple(None if x is None else x[idx_] for x in wrap_)
|
||
return
|
||
|
||
def fill_holes(self,keep_wall_attached=True,use_cc3d=False,return_mask=False):
|
||
'''Fill the holes in binary objects while taking into account periodicity.
|
||
In the non-periodic sense, a hole is a region of zeros which does not connect
|
||
to a boundary. When keep_wall_attached==False, only regions are kept which
|
||
fully connect from top to bottom wall.
|
||
In the periodic sense, a hole is a region of zeros which is not
|
||
connected to itself accross the opposite periodic boundary.'''
|
||
if return_mask: mask = self._data.copy()
|
||
np.logical_not(self._data,self._data)
|
||
if use_cc3d:
|
||
import cc3d
|
||
if self._ndim==2: connectivity = 8 if self.connect_diagonals else 4
|
||
elif self._ndim==3: connectivity = 18 if self.connect_diagonals else 6
|
||
else: raise RuntimeError("'use_cc3d' can only be used with 2D or 3D data.")
|
||
labels_,nlabels_ = cc3d.connected_components(self._data,connectivity=connectivity,return_N=True)
|
||
else:
|
||
from scipy import ndimage
|
||
if self.connect_diagonals: structure = ndimage.generate_binary_structure(self._ndim,self._ndim)
|
||
else: structure = ndimage.generate_binary_structure(self._ndim,1)
|
||
labels_,nlabels_ = ndimage.label(self._data,structure=structure)
|
||
labels_keep = set()
|
||
for axis in range(3):
|
||
sl_lo = tuple(slice(None) if ii!=axis else 0 for ii in range(3))
|
||
sl_hi = tuple(slice(None) if ii!=axis else -1 for ii in range(3))
|
||
if self.periodicity[axis]:
|
||
labels_keep |= set(np.unique(labels_[sl_lo])) & set(np.unique(labels_[sl_hi]))
|
||
elif keep_wall_attached:
|
||
labels_keep |= set(np.unique(labels_[sl_lo])) | set(np.unique(labels_[sl_hi]))
|
||
labels_keep.discard(0)
|
||
self._data = np.isin(labels_,tuple(labels_keep),invert=True)
|
||
# If labels have been computed already, set them to None because they are now invalid
|
||
if self._labels is not None:
|
||
self._labels = None
|
||
self.nlabels = None
|
||
self.wrap = tuple(self._ndim*[None])
|
||
# Compute mask if it is to be returned, otherwise we are done
|
||
if return_mask:
|
||
np.logical_xor(mask,self._data,out=mask)
|
||
return mask
|
||
return
|
||
|
||
def probe(self,idx,probe_label=False):
|
||
'''Returns whether or not a point at idx is True or False.'''
|
||
if probe_label:
|
||
return self._labels[tuple(idx)]
|
||
else:
|
||
return self._data[tuple(idx)]
|
||
|
||
def volume(self,label):
|
||
'''Returns the sum of True values.'''
|
||
if label is None:
|
||
return np.sum(self.data)
|
||
else:
|
||
if self._labels is None: self.label()
|
||
return np.sum(self.labels==label)
|
||
|
||
def volumes(self):
|
||
'''Returns volume of all features, i.e. connected regions which have been
|
||
labeled using the label() method including the volume of the background region.
|
||
The array is sorted by labels, i.e. vol[0] is the volume of the background
|
||
region, vol[1] the volume of label 1, etc. If 'label' is an integer value,
|
||
only the volume of the corresponding region is returned.
|
||
Note: it is more efficient to retrieve all volumes at once than querying
|
||
single labels.'''
|
||
if self._labels is None: self.label()
|
||
return np.bincount(self.labels.ravel())
|
||
|
||
def volume_domain(self):
|
||
'''Returns volume of entire domain. Should be equal to sum(volume_feature()).'''
|
||
dim = tuple(self._dim[axis]-1 if self.periodicity[axis] else self._dim[axis] for axis in range(self._ndim))
|
||
return np.prod(dim)
|
||
|
||
def feature_labels_by_volume(self,descending=True,return_volumes=False):
|
||
'''Returns labels of connected regions sorted by volume.'''
|
||
vol = self.volumes()[1:]
|
||
labels = np.argsort(vol)+1
|
||
if descending:
|
||
labels = labels[::-1]
|
||
vol = vol[::-1]
|
||
if return_volumes:
|
||
return labels,vol
|
||
else:
|
||
return labels
|
||
|
||
def discard_features(self,selection,return_mask=False):
|
||
'''Removes features from data. Optionally returns a boolean array 'mask'
|
||
which marks position of features which have been removed.'''
|
||
if self._labels is None:
|
||
self.label()
|
||
selection = self._select_features(selection)
|
||
# Map tagged regions to zero in order to discard them
|
||
map_ = np.array(range(self.nlabels+1),dtype=self._labels.dtype)
|
||
map_[selection] = 0
|
||
# Remove gaps from target mapping
|
||
idx_,map_ = np.unique(map_,return_index=True,return_inverse=True)[1:3]
|
||
# Discard regions
|
||
self._labels = map_[self._labels]
|
||
self.nlabels = np.max(map_)
|
||
print(self.nlabels)
|
||
self.wrap = tuple(None if x is None else x[idx_] for x in self.wrap)
|
||
# for axis in range(self._ndim):
|
||
# if self.wrap[axis] is not None:
|
||
# self.wrap[axis] = self.wrap[axis][idx_]
|
||
if return_mask:
|
||
mask = np.logical_and(self._labels==0,self._data)
|
||
self._data = self._labels>0
|
||
if return_mask: return mask
|
||
|
||
def _select_features(self,selection):
|
||
dtype = self._labels.dtype
|
||
if selection is None:
|
||
return np.array(range(1,self.nlabels+1))
|
||
elif np.issubdtype(type(selection),np.integer):
|
||
return np.array(selection,dtype=dtype,ndmin=1)
|
||
elif isinstance(selection,(list,tuple,np.ndarray)):
|
||
selection = np.array(selection)
|
||
if selection.dtype==np.dtype('bool'):
|
||
assert selection.ndim==1 and selection.shape[0]==self.nlabels+1,\
|
||
"Boolean indexing must provide count+1 values."
|
||
else:
|
||
selection = selection.astype(dtype)
|
||
assert np.max(selection)<=self.nlabels and np.min(selection)>=0,\
|
||
"Entry in selection is out-of-bounds."
|
||
return selection
|
||
else:
|
||
raise ValueError('Invalid input. Accepting int,list,tuple,ndarray.')
|
||
|
||
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<self.snapshot.nproc():
|
||
field = self.snapshot.field_chunk(
|
||
self.iter_rank,self.key,keep_ghost=self.keep_ghost)
|
||
self.iter_rank += 1
|
||
return field
|
||
else:
|
||
raise StopIteration
|