implemented padding, laplacian etc
This commit is contained in:
parent
4a33b5f00a
commit
8dd8161537
143
field.py
143
field.py
|
|
@ -123,50 +123,58 @@ class Field3d:
|
|||
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,preserve_grid=False,padding=None):
|
||||
if not preserve_grid:
|
||||
origin = list(self.origin)
|
||||
if padding is None:
|
||||
data = numpy.zeros(self.data.shape,dtype=self.data.dtype)
|
||||
slout = [slice(None),slice(None),slice(None)]
|
||||
origin[axis] += 0.5*self.spacing[axis]
|
||||
elif padding=='before':
|
||||
shape = numpy.array(self.data.shape)
|
||||
shape[axis] += 1
|
||||
data = numpy.zeros(tuple(shape),dtype=self.data.dtype)
|
||||
slout = [slice(None),slice(None),slice(None)]
|
||||
slout[axis] = slice(1,None)
|
||||
origin[axis] -= 0.5*self.spacing[axis]
|
||||
elif padding=='after':
|
||||
shape = numpy.array(self.data.shape)
|
||||
shape[axis] += 1
|
||||
data = numpy.zeros(tuple(shape),dtype=self.data.dtype)
|
||||
slout = [slice(None),slice(None),slice(None)]
|
||||
slout[axis] = slice(0,-1)
|
||||
origin[axis] += 0.5*self.spacing[axis]
|
||||
elif padding=='both':
|
||||
shape = numpy.array(self.data.shape)
|
||||
shape[axis] += 2
|
||||
data = numpy.zeros(tuple(shape),dtype=self.data.dtype)
|
||||
slout = [slice(None),slice(None),slice(None)]
|
||||
slout[axis] = slice(1,-1)
|
||||
origin[axis] -= 0.5*self.spacing[axis]
|
||||
def derivative(self,axis,only_keep_interior=False,add_border='before',preserve_origin=False):
|
||||
if axis is None:
|
||||
return tuple(self.derivative(axis,preserve_grid=preserve_grid,padding=padding,numpad=numpad)
|
||||
for axis in range(0,3))
|
||||
assert axis<3, "'axis' must be one of 0,1,2."
|
||||
if not preserve_origin:
|
||||
origin = list(self.origin)
|
||||
numpoints = list(self.numpoints)
|
||||
origin[axis] += 0.5*self.spacing[axis]
|
||||
numpoints[axis] -= 1
|
||||
if not only_keep_interior:
|
||||
assert add_border in ('before','after'), "'add_border' must be one of {'before','after'}."
|
||||
origin,numpoints,sl_out = padding(origin,self.spacing,numpoints,add_border,1)
|
||||
else:
|
||||
raise ValueError("'padding' must be one of 'none','before','after'.")
|
||||
slout = tuple(slout)
|
||||
sl_out = tuple(slice(None),slice(None),slice(None))
|
||||
data = numpy.zeros(numpoints,dtype=self.data.dtype)
|
||||
# 2nd order FD with h = dx/2
|
||||
slhi = [slice(None),slice(None),slice(None)]
|
||||
sllo = [slice(None),slice(None),slice(None)]
|
||||
slhi = 3*[slice(None)]
|
||||
sllo = 3*[slice(None)]
|
||||
slhi[axis] = slice(1,None)
|
||||
sllo[axis] = slice(0,-1)
|
||||
slhi = tuple(slhi)
|
||||
sllo = tuple(sllo)
|
||||
data[slout] = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo])
|
||||
data[sl_out] = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo])
|
||||
else:
|
||||
origin,numpoints,sl_out = padding(self.origin,self.spacing,self.numpoints,padding,numpad)
|
||||
data = numpy.zeros(numpoints,dtype=self.data.dtype)
|
||||
# 2nd order FD with h = dx, one-sided at boundaries
|
||||
# https://numpy.org/doc/stable/reference/generated/numpy.gradient.html
|
||||
origin = self.origin
|
||||
data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=2)
|
||||
data[sl_out] = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=2)
|
||||
if only_keep_interior:
|
||||
sl_out = 3*[slice(None)]
|
||||
sl_out[axis] = slice(1,-1)
|
||||
data = data[tuple(sl_out)]
|
||||
origin[axis] += self.spacing[axis]
|
||||
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,[1.,-2.,1.]/self.spacing[0],axis=0,mode='constant',cval=numpy.nan,origin=0)
|
||||
data += ndimage.correlate1d(self.data,[1.,-2.,1.]/self.spacing[1],axis=1,mode='constant',cval=numpy.nan,origin=0)
|
||||
data += ndimage.correlate1d(self.data,[1.,-2.,1.]/self.spacing[2],axis=2,mode='constant',cval=numpy.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):
|
||||
|
|
@ -204,22 +212,22 @@ class Field3d:
|
|||
return (out,weight)
|
||||
else:
|
||||
return out/weight
|
||||
|
||||
def gradient(self,axis,preserve_origin=False,padding=None):
|
||||
return [self.derivative(axis,preserve_origin=preserve_origin,padding=padding) for axis in range(0,3)]
|
||||
|
||||
def gaussian_filter(self,sigma,truncate=4.0,border='constant',const_val=numpy.nan):
|
||||
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,numpy.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=border,cval=const_val)
|
||||
print(data.shape,self.data.shape)
|
||||
print(numpy.linalg.norm(data-self.data))
|
||||
print(numpy.max(data))
|
||||
return Field3d(data,self.origin,self.spacing)
|
||||
data = ndimage.gaussian_filter(self.data,sigma_img,truncate=truncate,mode='constant',cval=numpy.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.'''
|
||||
|
|
@ -233,6 +241,7 @@ class Field3d:
|
|||
return tuple(radius)
|
||||
|
||||
def shift_origin(self,rel_shift):
|
||||
raise NotImplementedError("Routine has not been verified yet.")
|
||||
#TBD: verify this routine
|
||||
assert isinstance(rel_shift,(tuple,list,numpy.ndarray)) and len(rel_shift)==3,\
|
||||
"'shift' must be tuple/list with length 3."
|
||||
|
|
@ -267,11 +276,14 @@ class Field3d:
|
|||
data[sl_clr] = numpy.nan
|
||||
return Field3d(data,origin,self.spacing)
|
||||
|
||||
def change_grid(self,origin,spacing,numpoints):
|
||||
def change_grid(self,origin,spacing,numpoints,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]+(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)
|
||||
# Allocate (possibly padded array)
|
||||
origin_pad,numpoints_pad,sl_out = padding(origin,spacing,numpoints,padding,numpad):
|
||||
data = numpy.zeros(numpoints_pad,dtype=self.data.dtype)
|
||||
# Trilinear interpolation
|
||||
if numpy.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)
|
||||
|
|
@ -282,11 +294,12 @@ class Field3d:
|
|||
for jj in range(0,2):
|
||||
for kk in range(0,2):
|
||||
if c[ii,jj,kk]>self.eps_collapse:
|
||||
data += c[ii,jj,kk]*self.data[
|
||||
data[sl_out] += 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:
|
||||
data_ref = data[sl_out]
|
||||
for ii in range(0,numpoints[0]):
|
||||
for jj in range(0,numpoints[1]):
|
||||
for kk in range(0,numpoints[2]):
|
||||
|
|
@ -294,8 +307,8 @@ class Field3d:
|
|||
origin[0]+ii*spacing[0],
|
||||
origin[1]+jj*spacing[1],
|
||||
origin[2]+kk*spacing[2])
|
||||
data[ii,jj,kk] = self.interpolate(coord)
|
||||
return Field3d(data,origin,spacing)
|
||||
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."
|
||||
|
|
@ -312,6 +325,38 @@ class Field3d:
|
|||
val += c[ii,jj,kk]*self.data[i0+ii,j0+jj,k0+kk]
|
||||
return val
|
||||
|
||||
@staticmethod
|
||||
def padding(origin,spacing,numpoints,padding,numpad):
|
||||
if isinstance(numpad,int):
|
||||
numpad = numpy.fill(3,numpad,dtype=int)
|
||||
else:
|
||||
numpad = numpy.array(numpad,dtype=int)
|
||||
assert len(numpad)==3, "'numpad' must be either an integer or tuple/list of length 3."
|
||||
origin_pad = numpy.array(origin)
|
||||
numpoints_pad = numpy.array(numpoints)
|
||||
sl_out = [slice(None),slice(None),slice(None)]
|
||||
if padding is not None:
|
||||
if padding=='before':
|
||||
numpoints_pad += numpad
|
||||
origin_pad -= numpad*spacing
|
||||
for axis in range(3):
|
||||
sl_out[axis] = slice(numpad[axis],None)
|
||||
elif padding=='after':
|
||||
numpoints_pad += numpad
|
||||
for axis in range(3):
|
||||
sl_out[axis] = slice(0,-numpad[axis])
|
||||
elif padding=='both':
|
||||
numpoints_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)
|
||||
numpoints_pad = tuple(numpoints_pad)
|
||||
return (origin_pad,numpoints_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
|
||||
|
|
|
|||
Loading…
Reference in New Issue