1
"""Methods for converting volumetric data objects"""
2

3 26
__all__ = [
4
    'get_volume_shape',
5
    'volume_grid_geom_to_vtk',
6
    'volume_to_vtk',
7
]
8

9 26
__displayname__ = 'Volume'
10

11 26
import numpy as np
12 26
import pyvista
13

14 26
from omfvista.utilities import check_orientation, check_orthogonal
15

16

17 26
def get_volume_shape(vol):
18
    """Returns the shape of a gridded volume"""
19 26
    return ( len(vol.tensor_u), len(vol.tensor_v), len(vol.tensor_w))
20

21

22 26
def volume_grid_geom_to_vtk(volgridgeom, origin=(0.0, 0.0, 0.0)):
23
    """Convert the 3D gridded volume to a :class:`pyvista.StructuredGrid`
24
    (or a :class:`pyvista.RectilinearGrid` when apprropriate) object contatining
25
    the 2D surface.
26

27
    Args:
28
        volgridgeom (:class:`omf.volume.VolumeGridGeometry`): the grid geometry
29
            to convert
30
    """
31 26
    volgridgeom._validate_mesh()
32

33 26
    ox, oy, oz = volgridgeom.origin
34

35
    # Make coordinates along each axis
36 26
    x = ox + np.cumsum(volgridgeom.tensor_u)
37 26
    x = np.insert(x, 0, ox)
38 26
    y = oy + np.cumsum(volgridgeom.tensor_v)
39 26
    y = np.insert(y, 0, oy)
40 26
    z = oz + np.cumsum(volgridgeom.tensor_w)
41 26
    z = np.insert(z, 0, oz)
42

43
    # If axis orientations are standard then use a vtkRectilinearGrid
44 26
    if check_orientation(volgridgeom.axis_u, volgridgeom.axis_v, volgridgeom.axis_w):
45 26
        return pyvista.RectilinearGrid(x + origin[0], y + origin[1], z + origin[2])
46

47
    # Otherwise use a vtkStructuredGrid
48
    # Build out all nodes in the mesh
49 26
    xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')
50 26
    points = np.c_[xx.ravel('F'), yy.ravel('F'), zz.ravel('F')]
51

52
    # Rotate the points based on the axis orientations
53 26
    rotation_mtx = np.array([volgridgeom.axis_u, volgridgeom.axis_v, volgridgeom.axis_w])
54 26
    points = points.dot(rotation_mtx)
55

56 26
    output = pyvista.StructuredGrid()
57 26
    output.points = points
58 26
    output.dimensions = len(x), len(y), len(z)
59 26
    output.points += np.array(origin)
60 26
    return output
61

62

63 26
def volume_to_vtk(volelement, origin=(0.0, 0.0, 0.0)):
64
    """Convert the volume element to a VTK data object.
65

66
    Args:
67
        volelement (:class:`omf.volume.VolumeElement`): The volume element to
68
            convert
69

70
    """
71 26
    output = volume_grid_geom_to_vtk(volelement.geometry, origin=origin)
72 26
    shp = get_volume_shape(volelement.geometry)
73
    # Add data to output
74 26
    for data in volelement.data:
75 26
        arr = data.array.array
76 26
        arr = np.reshape(arr, shp).flatten(order='F')
77 26
        output[data.name] = arr
78 26
    return output
79

80

81
# Now set up the display names for the docs
82 26
volume_to_vtk.__displayname__ = 'Volume to VTK'
83 26
volume_grid_geom_to_vtk.__displayname__ = 'Volume Grid Geometry to VTK'
84 26
get_volume_shape.__displayname__ = 'Volume Shape'

Read our documentation on viewing source code .

Loading