import numpy as np
from scipy.interpolate import griddata
[docs]class BaseGrid(object):
"""
Base class for grids. Implements a simple data and grid cell interface.
"""
dim = None #: dimensionality of the grid
_cellCenters = None #: coordinates of the grid cell centers
_cellVariables = () #: names of the cell variables
_cellValues = () #: values of the cell variables
_dataDict = None #: dictionary holding variable data as key-value pairs
[docs] def merge(self, other):
"""
Merge two grids into a new grid instance
"""
if self.dim != other.dim:
raise RuntimeError("Grid dimensions do not match!")
self._calcCellCenters()
other._calcCellCenters()
new = self._merge_derived(other)
new._calcCellCenters()
for var in self._cellVariables:
data = np.append(self.data()[var],other.data()[var],axis=0)
new.add_data(var,data)
return new
[docs] def cellCenters(self):
"""
:returns: The cell centers of the grid.
:rtype: array_like
"""
if self._cellCenters is None:
self._calcCellCenters()
return self._cellCenters #np.array((self._cellCenters[...,0],self._cellCenters[...,1]))
[docs] def add_data(self,name,data):
"""
Used for attaching data to the grid.
:param str name: Name of the variable (will be used to reference it)
:param array_like data: Variable data. First dimensions must match the cell dimensions.
"""
if name in self._cellVariables:
raise ValueError("Cell variable {} already added".format(name))
data = np.squeeze(np.array(data))
if any([a != b for a,b in zip(data.shape, self.cellCenters().shape[:-1])]):
raise TypeError("First dimensions of data array must match cell dimensions")
self._cellVariables = self._cellVariables + (name,)
self._cellValues = self._cellValues + (data,)
[docs] def data(self):
"""
:returns: The cell variables as key-value pairs
:rtype: dict
"""
if self._dataDict is None:
self._calcDataDict()
if not all([x in self._dataDict.keys() for x in self._cellVariables]):
self._calcDataDict()
return self._dataDict
def _calcDataDict(self):
self._dataDict = {cvar: cval for cvar, cval in zip(self._cellVariables,self._cellValues)}
def _calcCellCenters(self):
return NotImplementedError
def _merge_derived(self,other):
return NotImplementedError
[docs]class RegularGrid(BaseGrid):
"""
Class implementing a grid interface for given cell centers on a regular grid.
"""
def __init__(self,cellCenters):
super(RegularGrid,self).__init__()
cellCenters = np.squeeze(np.array(cellCenters))
if not len(cellCenters.shape) in [1,2,3]:
raise ValueError("cellCenters must either be a 1-d, 2-d, or 3-d array")
if len(cellCenters.shape) != 1:
if not cellCenters.shape[-1] in [1,2,3]:
raise ValueError("Last dimension of cellCenters must be of size 1, 2, or 3.")
self._dim = cellCenters.shape[-1]
else:
self._dim = 1
self._cellCenters = cellCenters
def _merge_derived(self,other):
"""
Merge this grid with another into a new regular grid instance
"""
cellCenters = np.append(self.cellCenters(),other.cellCenters(),axis=0)
return RegularGrid(cellCenters)
[docs]class UnstructuredVTKGrid(BaseGrid):
"""
Class for storing an unstructured grid and cell variables defined on it.
:param pointCoords: Coordinates of the grid points
:type pointCoords: array_like of shape (npoints,dim)
:param array_like connectivity: For each cell, the indices of the points that \
make up the cell
:param ppc: Points per grid cell
:type ppc: 3, 4, 8
:param int dim: Real dimension of the grid
"""
def __init__(self,pointCoords,connectivity,ppc):
super(UnstructuredVTKGrid,self).__init__()
self.points = np.array(pointCoords)
self.connectivity = np.array(connectivity,dtype=int)
self.ppc = ppc
if len(self.points.shape) != 2:
raise TypeError("pointCoords must be a 2D array")
self._numPoints = self.points.shape[0]
self._dim = self.points.shape[1]
# filter trivial dimensions
self._trivial = []
for i in range(self._dim):
if np.allclose(self.points[0,i],self.points[:,i]):
self._trivial.append(i)
self.points = self._remove_trivial(self.points)
self.dim = self.points.shape[1]
if self.dim > 3:
raise TypeError("Only 2- and 3-dimensional grids are supported")
if not isinstance(self.ppc,int):
raise TypeError("ppc (points per cell) must be an integer")
if len(self.connectivity) % self.ppc != 0:
raise TypeError("Non-integer number of grid cells")
self._numCells = int(len(self.connectivity) / self.ppc)
def _merge_derived(self, other):
"""
Merge this grid with another into a new unstructured grid instance
"""
if self.ppc != other.ppc:
raise RuntimeError("Using different numbers of points per grid cell")
points_n = np.append(self.points,other.points,axis=0)
connectivity_n = np.append(self.connectivity,other.connectivity + len(self.points),axis=0)
return UnstructuredVTKGrid(points_n,connectivity_n,self.ppc)
[docs] def triangulation(self):
"""
:returns: Cell point coordinates x, y, z and connectivity c as (x, y, z), c
:rtype: tuple
"""
return self.points.T, self.connectivity.reshape(self._numCells,self.ppc)
[docs] def integrate(self,v):
"""
Integrates an array v defined on each cell center across the grid.
:param v: Input array of shape (ncells, dim) or scalar
"""
if np.asarray(v).size > 1 and self._numCells != np.asarray(v).shape[0]:
raise ValueError("Shape mismatch (got: {}, {})".format(self._numCells, np.asarray(v).shape[0]))
if self.dim == 2:
if self.ppc == 3:
areafun = lambda t: .5 * abs(t[:,0,0] * t[:,1,1] - t[:,0,1] * t[:,1,0])
elif self.ppc == 4:
areafun = lambda t: abs(t[:,0,0] * t[:,1,1])
else:
raise ValueError("Only triangular and rectangular geometry supported")
elif self.dim == 3:
if self.ppc == 4:
areafun = lambda t: 1./6. * abs(np.sum(t[:,0,:] * np.cross(t[:,1,:],t[:,2,:]),axis=-1))
elif self.ppc == 6:
areafun = lambda t: abs(t[:,0,0] * t[:,1,1] * t[:,2,2])
else:
raise ValueError("Only tetrahedral and cubic geometry supported")
else:
raise ValueError("Only two and three-dimensional grids supported")
corners = self.points[self.connectivity.reshape(self._numCells,self.ppc)]
t = corners[:,1:,:] - corners[:,0,:][:,np.newaxis,:]
cell_area = areafun(t)
broadcast_shape = (self._numCells,) + (1,) * (np.asarray(v).ndim-1)
integral = np.sum(cell_area.reshape(broadcast_shape) * v, axis=0)
return np.squeeze(integral)
def _calcCellCenters(self):
"""
:returns: Coordinates of cell centers
:rtype: array_like
"""
cellCoords = self.points[self.connectivity,:]
cells = cellCoords.reshape(self._numCells,self.ppc,self.dim)
cellCenters = np.mean(cells,axis=1) # average over cell coordinates to get cell center
self._cellCenters = np.squeeze(cellCenters)
def _remove_trivial(self,x):
"""
Removes all trivial dimensions from the input, depending on the grid cells.
:param array_like x: Data containing the trivial dimension
:returns: Data without the trivial dimension
"""
if self._trivial:
return np.delete(x,self._trivial,axis=1)
else:
return x
[docs]class StructuredVTKGrid(BaseGrid):
def __init__(self):
raise NotImplementedError