1
"""Methods to convert line set objects to VTK data objects"""
2

3

4 26
__all__ = [
5
    'line_set_to_vtk',
6
]
7

8 26
__displayname__ = 'Line Set'
9

10 26
import numpy as np
11 26
import pyvista
12

13 26
from omfvista.utilities import add_data
14

15

16 26
def line_set_to_vtk(lse, origin=(0.0, 0.0, 0.0)):
17
    """Convert the line set to a :class:`pyvista.PolyData` data object.
18

19
    Args:
20
        lse (:class:`omf.lineset.LineSetElement`): The line set to convert
21

22
    Return:
23
        :class:`pyvista.PolyData`
24
    """
25 26
    ids = np.array(lse.geometry.segments).reshape(-1, 2).astype(np.int_)
26 26
    lines = np.c_[np.full(len(ids), 2, dtype=np.int_), ids]
27

28 26
    output = pyvista.PolyData()
29 26
    output.points = np.array(lse.geometry.vertices)
30 26
    output.lines = lines
31

32 26
    indices = output.connectivity().cell_arrays['RegionId']
33 26
    output['Line Index'] = indices
34

35
    # Now add data to lines:
36 26
    add_data(output, lse.data)
37

38
    # TODO: if subtype is borehole make a tube
39

40 26
    output.points += np.array(origin)
41 26
    return output
42

43

44 26
line_set_to_vtk.__displayname__ = 'Line Set to VTK'

Read our documentation on viewing source code .

Loading