Source code for PVGeo.model_build.grids

__all__ = [
    'CreateEvenRectilinearGrid',
    'CreateUniformGrid',
    'CreateTensorMesh',
]

__displayname__ = 'Grids'

import numpy as np
import vtk

from .. import interface
from ..base import AlgorithmBase


def _makeSpatialCellData(nx, ny, nz):
    """Used for testing"""
    arr = np.fromfunction(lambda k, j, i: k * j * i, (nz, ny, nx))
    return arr.flatten()


[docs] class CreateUniformGrid(AlgorithmBase): """Create uniform grid (``vtkImageData``)""" __displayname__ = 'Create Uniform Grid' __category__ = 'source' def __init__( self, extent=(10, 10, 10), spacing=(1.0, 1.0, 1.0), origin=(0.0, 0.0, 0.0) ): AlgorithmBase.__init__( self, nInputPorts=0, nOutputPorts=1, outputType='vtkImageData' ) self.__extent = extent self.__spacing = spacing self.__origin = origin
[docs] def RequestData(self, request, inInfo, outInfo): """Used by pipeline to generate the output""" pdo = self.GetOutputData(outInfo, 0) nx, ny, nz = self.__extent[0], self.__extent[1], self.__extent[2] sx, sy, sz = self.__spacing[0], self.__spacing[1], self.__spacing[2] ox, oy, oz = self.__origin[0], self.__origin[1], self.__origin[2] # Setup the ImageData pdo.SetDimensions(nx, ny, nz) pdo.SetOrigin(ox, oy, oz) pdo.SetSpacing(sx, sy, sz) # pdo.SetExtent(0,nx-1, 0,ny-1, 0,nz-1) # Add CELL data data = _makeSpatialCellData( nx - 1, ny - 1, nz - 1 ) # minus 1 b/c cell data not point data data = interface.convert_array(data, name='Spatial Cell Data', deep=True) # THIS IS CELL DATA! Add the model data to CELL data: pdo.GetCellData().AddArray(data) # Add Point data data = _makeSpatialCellData(nx, ny, nz) data = interface.convert_array(data, name='Spatial Point Data', deep=True) # THIS IS CELL DATA! Add the model data to CELL data: pdo.GetPointData().AddArray(data) return 1
[docs] def RequestInformation(self, request, inInfo, outInfo): """Used by pipeline to handle output extents""" # Now set whole output extent ext = [ 0, self.__extent[0] - 1, 0, self.__extent[1] - 1, 0, self.__extent[2] - 1, ] info = outInfo.GetInformationObject(0) # Set WHOLE_EXTENT: This is absolutely necessary info.Set(vtk.vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT(), ext, 6) return 1
#### Setters / Getters ####
[docs] def set_extent(self, nx, ny, nz): """Set the extent of the output grid.""" if self.__extent != [nx, ny, nz]: self.__extent = [nx, ny, nz] self.Modified()
[docs] def set_spacing(self, dx, dy, dz): """Set the spacing for the points along each axial direction.""" if self.__spacing != [dx, dy, dz]: self.__spacing = [dx, dy, dz] self.Modified()
[docs] def set_origin(self, x0, y0, z0): """Set the origin of the output grid.""" if self.__origin != [x0, y0, z0]: self.__origin = [x0, y0, z0] self.Modified()
[docs] class CreateEvenRectilinearGrid(AlgorithmBase): """This creates a vtkRectilinearGrid where the discretization along a given axis is uniformly distributed. """ __displayname__ = 'Create Even Rectilinear Grid' __category__ = 'source' def __init__( self, extent=[10, 10, 10], xrng=[-1.0, 1.0], yrng=[-1.0, 1.0], zrng=[-1.0, 1.0] ): AlgorithmBase.__init__( self, nInputPorts=0, nOutputPorts=1, outputType='vtkRectilinearGrid' ) self.__extent = extent self.__xrange = xrng self.__yrange = yrng self.__zrange = zrng
[docs] def RequestData(self, request, inInfo, outInfo): """Used by pipeline to generate the output""" # Get output of Proxy pdo = self.GetOutputData(outInfo, 0) # Perform task nx, ny, nz = self.__extent[0] + 1, self.__extent[1] + 1, self.__extent[2] + 1 xcoords = np.linspace(self.__xrange[0], self.__xrange[1], num=nx) ycoords = np.linspace(self.__yrange[0], self.__yrange[1], num=ny) zcoords = np.linspace(self.__zrange[0], self.__zrange[1], num=nz) # CONVERT TO VTK # xcoords = interface.convert_array(xcoords, deep=True) ycoords = interface.convert_array(ycoords, deep=True) zcoords = interface.convert_array(zcoords, deep=True) pdo.SetDimensions(nx, ny, nz) pdo.SetXCoordinates(xcoords) pdo.SetYCoordinates(ycoords) pdo.SetZCoordinates(zcoords) data = _makeSpatialCellData(nx - 1, ny - 1, nz - 1) data = interface.convert_array(data, name='Spatial Data', deep=True) # THIS IS CELL DATA! Add the model data to CELL data: pdo.GetCellData().AddArray(data) return 1
[docs] def RequestInformation(self, request, inInfo, outInfo): """Used by pipeline to handle output extents""" # Now set whole output extent ext = [0, self.__extent[0], 0, self.__extent[1], 0, self.__extent[2]] info = outInfo.GetInformationObject(0) # Set WHOLE_EXTENT: This is absolutely necessary info.Set(vtk.vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT(), ext, 6) return 1
#### Setters / Getters ####
[docs] def set_extent(self, nx, ny, nz): """Set the extent of the output grid.""" if self.__extent != [nx, ny, nz]: self.__extent = [nx, ny, nz] self.Modified()
[docs] def set_x_range(self, start, stop): """Set range (min, max) for the grid in the X-direction.""" if self.__xrange != [start, stop]: self.__xrange = [start, stop] self.Modified()
[docs] def set_y_range(self, start, stop): """Set range (min, max) for the grid in the Y-direction""" if self.__yrange != [start, stop]: self.__yrange = [start, stop] self.Modified()
[docs] def set_z_range(self, start, stop): """Set range (min, max) for the grid in the Z-direction""" if self.__zrange != [start, stop]: self.__zrange = [start, stop] self.Modified()
[docs] class CreateTensorMesh(AlgorithmBase): """This creates a vtkRectilinearGrid where the discretization along a given axis is not uniform. Cell spacings along each axis can be set via strings with repeating patterns or explicitly using the ``Set*Cells`` methods. """ __displayname__ = 'Create Tensor Mesh' __category__ = 'source' def __init__( self, origin=[-350.0, -400.0, 0.0], data_name='Data', xcellstr='200 100 50 20*50.0 50 100 200', ycellstr='200 100 50 21*50.0 50 100 200', zcellstr='20*25.0 50 100 200', ): AlgorithmBase.__init__( self, nInputPorts=0, nOutputPorts=1, outputType='vtkRectilinearGrid' ) self.__origin = origin self.__xcells = CreateTensorMesh._read_cell_line(xcellstr) self.__ycells = CreateTensorMesh._read_cell_line(ycellstr) self.__zcells = CreateTensorMesh._read_cell_line(zcellstr) self.__data_name = data_name
[docs] @staticmethod def _read_cell_line(line): """Read cell sizes for each line in the UBC mesh line strings""" # OPTIMIZE: work in progress # TODO: when optimized, make sure to combine with UBC reader line_list = [] for seg in line.split(): if '*' in seg: sp = seg.split('*') seg_arr = np.ones((int(sp[0]),), dtype=float) * float(sp[1]) else: seg_arr = np.array([float(seg)], dtype=float) line_list.append(seg_arr) return np.concatenate(line_list)
[docs] def get_extent(self): """Get the extent of the created mesh""" ne, nn, nz = len(self.__xcells), len(self.__ycells), len(self.__zcells) return (0, ne, 0, nn, 0, nz)
[docs] def _make_model(self, pdo): """Generates the output data object""" ox, oy, oz = self.__origin[0], self.__origin[1], self.__origin[2] # Read the cell sizes cx = self.__xcells cy = self.__ycells cz = self.__zcells # Invert the indexing of the vector to start from the bottom. cz = cz[::-1] # Adjust the reference point to the bottom south west corner oz = oz - np.sum(cz) # Now generate the coordinates for from cell width and origin cox = ox + np.cumsum(cx) cox = np.insert(cox, 0, ox) coy = oy + np.cumsum(cy) coy = np.insert(coy, 0, oy) coz = oz + np.cumsum(cz) coz = np.insert(coz, 0, oz) # Set the dims and coordinates for the output ext = self.get_extent() nx, ny, nz = ext[1] + 1, ext[3] + 1, ext[5] + 1 pdo.SetDimensions(nx, ny, nz) # Convert to VTK array for setting coordinates pdo.SetXCoordinates(interface.convert_array(cox, deep=True)) pdo.SetYCoordinates(interface.convert_array(coy, deep=True)) pdo.SetZCoordinates(interface.convert_array(coz, deep=True)) return pdo
[docs] def _add_model_data(self, pdo, data): """Add an array to the output data object. If data is None, random values will be generated. """ nx, ny, nz = pdo.GetDimensions() nx, ny, nz = nx - 1, ny - 1, nz - 1 # ADD DATA to cells if data is None: data = np.random.rand(nx * ny * nz) data = interface.convert_array(data, name='Random Data', deep=True) else: data = interface.convert_array(data, name=self.__data_name, deep=True) pdo.GetCellData().AddArray(data) return pdo
[docs] def RequestData(self, request, inInfo, outInfo): """Used by pipeline to generate output data object""" # Get input/output of Proxy pdo = self.GetOutputData(outInfo, 0) # Perform the task self._make_model(pdo) self._add_model_data(pdo, None) # TODO: add ability to set input data return 1
[docs] def RequestInformation(self, request, inInfo, outInfo): """Used by pipeline to set output whole extent""" # Now set whole output extent ext = self.get_extent() info = outInfo.GetInformationObject(0) # Set WHOLE_EXTENT: This is absolutely necessary info.Set(vtk.vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT(), ext, 6) return 1
#### Getters / Setters ####
[docs] def set_origin(self, x0, y0, z0): """Set the origin of the output""" if self.__origin != [x0, y0, z0]: self.__origin = [x0, y0, z0] self.Modified()
[docs] def set_x_cells(self, xcells): """Set the spacings for the cells in the X direction Args: xcells (list or np.array(floats)) : the spacings along the X-axis""" if len(xcells) != len(self.__xcells) or not np.allclose(self.__xcells, xcells): self.__xcells = xcells self.Modified()
[docs] def set_y_cells(self, ycells): """Set the spacings for the cells in the Y direction Args: ycells (list or np.array(floats)) : the spacings along the Y-axis""" if len(ycells) != len(self.__ycells) or not np.allclose(self.__ycells, ycells): self.__ycells = ycells self.Modified()
[docs] def set_z_cells(self, zcells): """Set the spacings for the cells in the Z direction Args: zcells (list or np.array(floats)): the spacings along the Z-axis""" if len(zcells) != len(self.__zcells) or not np.allclose(self.__zcells, zcells): self.__zcells = zcells self.Modified()
[docs] def set_x_cells_str(self, xcellstr): """Set the spacings for the cells in the X direction Args: xcellstr (str) : the spacings along the X-axis in the UBC style""" xcells = CreateTensorMesh._read_cell_line(xcellstr) self.set_x_cells(xcells)
[docs] def set_y_cells_str(self, ycellstr): """Set the spacings for the cells in the Y direction Args: ycellstr (str) : the spacings along the Y-axis in the UBC style""" ycells = CreateTensorMesh._read_cell_line(ycellstr) self.set_y_cells(ycells)
[docs] def set_z_cells_str(self, zcellstr): """Set the spacings for the cells in the Z direction Args: zcellstr (str) : the spacings along the Z-axis in the UBC style""" zcells = CreateTensorMesh._read_cell_line(zcellstr) self.set_z_cells(zcells)