1
"""This module provides a wrapper that will work for any OMF data object or
2
project files.
3

4
Example Use
5
-----------
6

7
Use the wrapper provided in ``omfvista`` to wrap any ``omf`` data object:
8

9
.. code-block:: python
10

11
    import omfvista
12

13
    omfvista.wrap(data)
14

15

16
Here's an example using the sample data hosted in the `OMF repository`_.
17

18
.. _OMF repository: https://github.com/gmggroup/omf/tree/master/assets
19

20
.. code-block:: python
21

22
    import omf
23
    import omfvista
24

25
    # Read all elements
26
    reader = omf.OMFReader('test_file.omf')
27
    project = reader.get_project()
28

29
    # Iterate over the elements and add converted VTK objects to dictionary:
30
    data = dict()
31
    for e in project.elements:
32
        d = omfvista.wrap(e)
33
        data[e.name] = d
34

35
Or better yet, just use the project loader:
36

37
.. code-block:: python
38

39
    import omfvista
40
    data = omfvista.load_project('test_file.omf')
41

42
"""
43

44

45 26
__all__ = [
46
    'wrap',
47
    'project_to_vtk',
48
    'load_project',
49
]
50

51 26
__displayname__ = 'Wrapper'
52

53 26
import numpy as np
54 26
import omf
55 26
import pyvista
56

57 26
import omfvista
58

59

60 26
def wrap(data, origin=(0.0, 0.0, 0.0)):
61
    """Wraps the OMF data object/project as a VTK data object. This is the
62
    primary function that an end user will harness.
63

64
    Args:
65
        data: any OMF data object
66

67
    Example:
68
        >>> import omf
69
        >>> import omfvista
70

71
        >>> # Read all elements
72
        >>> reader = omf.OMFReader('test_file.omf')
73
        >>> project = reader.get_project()
74

75
        >>> # Iterate over the elements and add converted VTK objects to dictionary:
76
        >>> data = dict()
77
        >>> for e in project.elements:
78
        >>>     d = omfvista.wrap(e)
79
        >>>     data[e.name] = d
80

81
    """
82
    # Allow recursion
83 26
    if isinstance(data, (list, tuple)):
84 26
        multi = pyvista.MultiBlock()
85 26
        for i, item in enumerate(data):
86 26
            multi.append(wrap(item))
87 26
            multi.set_block_name(i, item.name)
88 26
        return multi
89
    # Define wrappers
90 26
    wrappers = {
91
        'LineSetElement': omfvista.line_set_to_vtk,
92
        'PointSetElement': omfvista.point_set_to_vtk,
93
        # Surfaces
94
        'SurfaceGeometry': omfvista.surface_geom_to_vtk,
95
        'SurfaceGridGeometry': omfvista.surface_grid_geom_to_vtk,
96
        'SurfaceElement': omfvista.surface_to_vtk,
97
        # Volumes
98
        'VolumeGridGeometry': omfvista.volume_grid_geom_to_vtk,
99
        'VolumeElement': omfvista.volume_to_vtk,
100
        'Project': omfvista.project_to_vtk,
101
    }
102
    # get the class name
103 26
    key = data.__class__.__name__
104 26
    try:
105 26
        if key != 'Project':
106 26
            return wrappers[key](data, origin=origin)
107
        else:
108 26
            return wrappers[key](data)
109 0
    except KeyError:
110 0
        raise RuntimeError('Data of type ({}) is not supported currently.'.format(key))
111

112

113 26
def project_to_vtk(project):
114
    """Converts an OMF project (:class:`omf.base.Project`) to a
115
    :class:`pyvista.MultiBlock` data boject
116
    """
117
    # Iterate over the elements and add converted VTK objects a MultiBlock
118 26
    data = pyvista.MultiBlock()
119 26
    origin = np.array(project.origin)
120 26
    for i, e in enumerate(project.elements):
121 26
        d = omfvista.wrap(e, origin=origin)
122 26
        data[i, e.name] = d
123 26
    return data
124

125

126 26
def load_project(filename):
127
    """Loads an OMF project file into a :class:`pyvista.MultiBlock` dataset"""
128 26
    reader = omf.OMFReader(filename)
129 26
    project = reader.get_project()
130 26
    return project_to_vtk(project)
131

132

133
# Now set up the display names for the docs
134 26
load_project.__displayname__ = 'Load Project File'
135 26
project_to_vtk.__displayname__ = 'Project to VTK'
136 26
wrap.__displayname__ = 'The Wrapper'

Read our documentation on viewing source code .

Loading