Merge pull request #114 from simpeg/feat/vtk

feat/vtk: Add VTK object interface

Author: lheagy

.bumpversion.cfg

@@ -1,4 +1,4 @@
 [bumpversion]
-current_version = 0.3.4
+current_version = 0.3.5
 files = setup.py discretize/__init__.py docs/conf.py
 

discretize/BaseMesh.py

@@ -4,8 +4,12 @@
 import json
 from .utils.matutils import mkvc
 
+try:
+    from .mixins import vtkInterface
+except ImportError as err:
+    vtkInterface = object
 
-class BaseMesh(properties.HasProperties):
+class BaseMesh(properties.HasProperties, vtkInterface):
     """
     BaseMesh does all the counting you don't want to do.
     BaseMesh should be inherited by meshes with a regular structure.
@@ -29,14 +33,16 @@ class BaseMesh(properties.HasProperties):
     )
 
     # Instantiate the class
-    def __init__(self, n, x0=None):
+    def __init__(self, n, x0=None, **kwargs):
         self._n = n  # number of dimensions
 
         if x0 is None:
             self.x0 = np.zeros(len(self._n))
         else:
             self.x0 = x0
 
+        super(BaseMesh, self).__init__(**kwargs)
+
     # Validators
     @properties.validator('_n')
     def check_n_shape(self, change):
@@ -52,7 +58,7 @@ def check_n_shape(self, change):
                 "supported".format(change['value'])
             )
 
-        if change['previous'] != properties.undefined:
+        if np.any(change['previous'] != properties.undefined):
             # can't change dimension of the mesh
             assert len(change['previous']) == len(change['value']), (
                 "Cannot change dimensionality of the mesh. Expected {} "
@@ -373,11 +379,92 @@ def copy(self):
         """
         return properties.copy(self)
 
+    axis_u = properties.Vector3(
+        'Vector orientation of u-direction. For more details see the docs for the :attr:`~discretize.BaseMesh.BaseMesh.rotation_matrix` property.',
+        default='X',
+        length=1
+    )
+    axis_v = properties.Vector3(
+        'Vector orientation of v-direction. For more details see the docs for the :attr:`~discretize.BaseMesh.BaseMesh.rotation_matrix` property.',
+        default='Y',
+        length=1
+    )
+    axis_w = properties.Vector3(
+        'Vector orientation of w-direction. For more details see the docs for the :attr:`~discretize.BaseMesh.BaseMesh.rotation_matrix` property.',
+        default='Z',
+        length=1
+    )
+
+    @properties.validator
+    def _validate_orientation(self):
+        """Check if axes are orthogonal"""
+        if not (np.abs(self.axis_u.dot(self.axis_v) < 1e-6) and
+                np.abs(self.axis_v.dot(self.axis_w) < 1e-6) and
+                np.abs(self.axis_w.dot(self.axis_u) < 1e-6)):
+            raise ValueError('axis_u, axis_v, and axis_w must be orthogonal')
+        return True
+
+    @property
+    def reference_is_rotated(self):
+        """True if the axes are rotated from the traditional <X,Y,Z> system
+        with vectors of :math:`(1,0,0)`, :math:`(0,1,0)`, and :math:`(0,0,1)`
+        """
+        if (    np.allclose(self.axis_u, (1, 0, 0)) and
+                np.allclose(self.axis_v, (0, 1, 0)) and
+                np.allclose(self.axis_w, (0, 0, 1)) ):
+            return False
+        return True
+
+    @property
+    def rotation_matrix(self):
+        """Builds a rotation matrix to transform coordinates from their coordinate
+        system into a conventional cartesian system. This is built off of the
+        three `axis_u`, `axis_v`, and `axis_w` properties; these mapping
+        coordinates use the letters U, V, and W (the three letters preceding X,
+        Y, and Z in the alphabet) to define the projection of the X, Y, and Z
+        durections. These UVW vectors describe the placement and transformation
+        of the mesh's coordinate sytem assuming at most 3 directions.
+
+        Why would you want to use these UVW mapping vectors the this
+        `rotation_matrix` property? They allow us to define the realationship
+        between local and global coordinate systems and provide a tool for
+        switching between the two while still maintaing the connectivity of the
+        mesh's cells. For a visual example of this, please see the figure in the
+        docs for the :class:`~discretize.mixins.vtkModule.vtkInterface`.
+        """
+        return np.array([self.axis_u, self.axis_v, self.axis_w])
+
+
+    reference_system = properties.String(
+        'The type of coordinate reference frame. Can take on the values ' +
+        'cartesian, cylindrical, or spherical. Abbreviations of these are allowed.',
+        default='cartesian',
+        change_case='lower',
+    )
+
+    @properties.validator
+    def _validate_reference_system(self):
+        """Check if the reference system is of a known type."""
+        choices = ['cartesian', 'cylindrical', 'spherical']
+        # Here are a few abbreviations that users can harnes
+        abrevs = {
+            'car': choices[0],
+            'cart': choices[0],
+            'cy': choices[1],
+            'cyl': choices[1],
+            'sph': choices[2],
+        }
+        # Get the name and fix it if it is abbreviated
+        self.reference_system = abrevs.get(self.reference_system, self.reference_system)
+        if self.reference_system not in choices:
+            raise ValueError('Coordinate system ({}) unknown.'.format(self.reference_system))
+        return True
+
 
 class BaseRectangularMesh(BaseMesh):
     """BaseRectangularMesh"""
-    def __init__(self, n, x0=None):
-        BaseMesh.__init__(self, n, x0=x0)
+    def __init__(self, n, x0=None, **kwargs):
+        BaseMesh.__init__(self, n, x0=x0, **kwargs)
 
     @property
     def nCx(self):

discretize/CylMesh.py

@@ -43,6 +43,7 @@ class CylMesh(
 
     def __init__(self, h=None, x0=None, **kwargs):
         super(CylMesh, self).__init__(h=h, x0=x0, **kwargs)
+        self.reference_system = 'cylindrical'
 
         if not np.abs(self.hy.sum() - 2*np.pi) < 1e-10:
             raise AssertionError("The 2nd dimension must sum to 2*pi")

discretize/MeshIO.py

@@ -7,6 +7,11 @@
 from . import utils
 from .BaseMesh import BaseMesh
 
+try:
+    from .mixins import vtkTensorRead
+except ImportError as err:
+    vtkTensorRead = object
+
 
 def load_mesh(filename):
     """
@@ -24,7 +29,7 @@ def load_mesh(filename):
     return data
 
 
-class TensorMeshIO(object):
+class TensorMeshIO(vtkTensorRead):
 
     @classmethod
     def _readUBC_3DMesh(TensorMesh, fileName):
@@ -154,164 +159,6 @@ def readUBC(TensorMesh, fileName, directory=''):
             raise Exception('File format not recognized')
         return Tnsmsh
 
-    @classmethod
-    def readVTK(TensorMesh, fileName, directory=''):
-        """Read VTK Rectilinear (vtr xml file) and return Tensor mesh and model
-
-        Input:
-        :param str fileName: path to the vtr model file to read or just its name if directory is specified
-        :param str directory: directory where the UBC GIF file lives
-
-        Output:
-        :rtype: tuple
-        :return: (TensorMesh, modelDictionary)
-        """
-        from vtk import vtkXMLRectilinearGridReader as vtrFileReader
-        from vtk.util.numpy_support import vtk_to_numpy
-
-        fname = os.path.join(directory, fileName)
-        # Read the file
-        vtrReader = vtrFileReader()
-        vtrReader.SetFileName(fname)
-        vtrReader.Update()
-        vtrGrid = vtrReader.GetOutput()
-        # Sort information
-        hx = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetXCoordinates())))
-        xR = vtk_to_numpy(vtrGrid.GetXCoordinates())[0]
-        hy = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetYCoordinates())))
-        yR = vtk_to_numpy(vtrGrid.GetYCoordinates())[0]
-        zD = np.diff(vtk_to_numpy(vtrGrid.GetZCoordinates()))
-        # Check the direction of hz
-        if np.all(zD < 0):
-            hz = np.abs(zD[::-1])
-            zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[-1]
-        else:
-            hz = np.abs(zD)
-            zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[0]
-        x0 = np.array([xR, yR, zR])
-
-        # Make the object
-        tensMsh = TensorMesh([hx, hy, hz], x0=x0)
-
-        # Grap the models
-        models = {}
-        for i in np.arange(vtrGrid.GetCellData().GetNumberOfArrays()):
-            modelName = vtrGrid.GetCellData().GetArrayName(i)
-            if np.all(zD < 0):
-                modFlip = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
-                tM = tensMsh.r(modFlip, 'CC', 'CC', 'M')
-                modArr = tensMsh.r(tM[:, :, ::-1], 'CC', 'CC', 'V')
-            else:
-                modArr = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
-            models[modelName] = modArr
-
-        # Return the data
-        return tensMsh, models
-
-    def writeVTK(mesh, fileName, models=None, directory=''):
-        """Makes and saves a VTK rectilinear file (vtr)
-        for a Tensor mesh and model.
-
-        Input:
-        :param str fileName:  path to the output vtk file or just its name if directory is specified
-        :param str directory: directory where the UBC GIF file lives
-        :param dict models: dictionary of numpy.array - Name('s) and array('s).
-        Match number of cells
-        """
-        from vtk import vtkRectilinearGrid as rectGrid, vtkXMLRectilinearGridWriter as rectWriter, VTK_VERSION
-        from vtk.util.numpy_support import numpy_to_vtk
-
-        fname = os.path.join(directory, fileName)
-        # Deal with dimensionalities
-        if mesh.dim >= 1:
-            vX = mesh.vectorNx
-            xD = mesh.nNx
-            yD, zD = 1, 1
-            vY, vZ = np.array([0, 0])
-        if mesh.dim >= 2:
-            vY = mesh.vectorNy
-            yD = mesh.nNy
-        if mesh.dim == 3:
-            vZ = mesh.vectorNz
-            zD = mesh.nNz
-        # Use rectilinear VTK grid.
-        # Assign the spatial information.
-        vtkObj = rectGrid()
-        vtkObj.SetDimensions(xD, yD, zD)
-        vtkObj.SetXCoordinates(numpy_to_vtk(vX, deep=1))
-        vtkObj.SetYCoordinates(numpy_to_vtk(vY, deep=1))
-        vtkObj.SetZCoordinates(numpy_to_vtk(vZ, deep=1))
-
-        # Assign the model('s) to the object
-        if models is not None:
-            for item in six.iteritems(models):
-                # Convert numpy array
-                vtkDoubleArr = numpy_to_vtk(item[1], deep=1)
-                vtkDoubleArr.SetName(item[0])
-                vtkObj.GetCellData().AddArray(vtkDoubleArr)
-            # Set the active scalar
-            vtkObj.GetCellData().SetActiveScalars(list(models.keys())[0])
-
-        # Check the extension of the fileName
-        ext = os.path.splitext(fname)[1]
-        if ext is '':
-            fname = fname + '.vtr'
-        elif ext not in '.vtr':
-            raise IOError('{:s} is an incorrect extension, has to be .vtr')
-        # Write the file.
-        vtrWriteFilter = rectWriter()
-        if float(VTK_VERSION.split('.')[0]) >= 6:
-            vtrWriteFilter.SetInputData(vtkObj)
-        else:
-            vtuWriteFilter.SetInput(vtuObj)
-        vtrWriteFilter.SetFileName(fname)
-        vtrWriteFilter.Update()
-
-    def _toVTRObj(mesh, models=None):
-        """
-        Makes and saves a VTK rectilinear file (vtr) for a
-        Tensor mesh and model.
-
-        Input:
-        :param str, path to the output vtk file
-        :param mesh, TensorMesh object - mesh to be transfer to VTK
-        :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells
-
-        """
-        # Import
-        from vtk import vtkRectilinearGrid as rectGrid, VTK_VERSION
-        from vtk.util.numpy_support import numpy_to_vtk
-
-        # Deal with dimensionalities
-        if mesh.dim >= 1:
-            vX = mesh.vectorNx
-            xD = mesh.nNx
-            yD, zD = 1, 1
-            vY, vZ = np.array([0, 0])
-        if mesh.dim >= 2:
-            vY = mesh.vectorNy
-            yD = mesh.nNy
-        if mesh.dim == 3:
-            vZ = mesh.vectorNz
-            zD = mesh.nNz
-        # Use rectilinear VTK grid.
-        # Assign the spatial information.
-        vtkObj = rectGrid()
-        vtkObj.SetDimensions(xD, yD, zD)
-        vtkObj.SetXCoordinates(numpy_to_vtk(vX, deep=1))
-        vtkObj.SetYCoordinates(numpy_to_vtk(vY, deep=1))
-        vtkObj.SetZCoordinates(numpy_to_vtk(vZ, deep=1))
-
-        # Assign the model('s) to the object
-        if models is not None:
-            for item in models.iteritems():
-                # Convert numpy array
-                vtkDoubleArr = numpy_to_vtk(item[1], deep=1)
-                vtkDoubleArr.SetName(item[0])
-                vtkObj.GetCellData().AddArray(vtkDoubleArr)
-            # Set the active scalar
-            vtkObj.GetCellData().SetActiveScalars(list(models.keys())[0])
-        return vtkObj
 
     def _readModelUBC_2D(mesh, fileName):
         """
@@ -655,52 +502,3 @@ def writeUBC(mesh, fileName, models=None):
             for item in six.iteritems(models):
                 # Save the data
                 np.savetxt(item[0], item[1][ubc_order], fmt='%3.5e')
-
-
-    def writeVTK(self, fileName, models=None):
-        """Function to write a VTU file from a TreeMesh and model."""
-        import vtk
-        from vtk import vtkXMLUnstructuredGridWriter as Writer, VTK_VERSION
-        from vtk.util.numpy_support import numpy_to_vtk
-
-        # Make the data parts for the vtu object
-        # Points
-        ptsMat = np.vstack((self.gridN, self.gridhN))
-
-        vtkPts = vtk.vtkPoints()
-        vtkPts.SetData(numpy_to_vtk(ptsMat, deep=True))
-        # Cells
-        cellArray = [c for c in self]
-        cellConn = np.array([cell.nodes for cell in cellArray])
-
-        cellsMat = np.concatenate((np.ones((cellConn.shape[0], 1))*cellConn.shape[1], cellConn), axis=1).ravel()
-        cellsArr = vtk.vtkCellArray()
-        cellsArr.SetNumberOfCells(cellConn.shape[0])
-        cellsArr.SetCells(cellConn.shape[0], numpy_to_vtk(cellsMat, deep=True, array_type=vtk.VTK_ID_TYPE))
-
-        # Make the object
-        vtuObj = vtk.vtkUnstructuredGrid()
-        vtuObj.SetPoints(vtkPts)
-        vtuObj.SetCells(vtk.VTK_VOXEL, cellsArr)
-        # Add the level of refinement as a cell array
-        cell_levels = np.array([cell._level for cell in cellArray])
-        refineLevelArr = numpy_to_vtk(cell_levels, deep=1)
-        refineLevelArr.SetName('octreeLevel')
-        vtuObj.GetCellData().AddArray(refineLevelArr)
-        # Assign the model('s) to the object
-        if models is not None:
-            for item in six.iteritems(models):
-                # Convert numpy array
-                vtkDoubleArr = numpy_to_vtk(item[1], deep=1)
-                vtkDoubleArr.SetName(item[0])
-                vtuObj.GetCellData().AddArray(vtkDoubleArr)
-
-        # Make the writer
-        vtuWriteFilter = Writer()
-        if float(VTK_VERSION.split('.')[0]) >= 6:
-            vtuWriteFilter.SetInputData(vtuObj)
-        else:
-            vtuWriteFilter.SetInput(vtuObj)
-        vtuWriteFilter.SetFileName(fileName)
-        # Write the file
-        vtuWriteFilter.Update()

discretize/TensorMesh.py

@@ -88,7 +88,7 @@ def __init__(self, h=None, x0=None, **kwargs):
             n = np.array([x.size for x in h])
 
         super(BaseTensorMesh, self).__init__(
-            n, x0=x0
+            n, x0=x0, **kwargs
         )
 
         # Ensure h contains 1D vectors