qudarap/tests/spline.py

0001 #!/usr/bin/env python
0002 """
0003
0004 https://docs.pyvista.org/examples/00-load/create-spline.html
0005
0006 """
0007 import numpy as np
0008 import pyvista as pv
0009 from opticks.ana.pvplt import *
0010
0011 def make_points():
0012     """
0013     In [1]: points
0014     Out[1]:
0015     array([[ 0.   ,  5.   , -2.   ],
0016            [ 1.216,  4.685, -1.96 ],
0017            [ 2.277,  4.092, -1.919],
0018            [ 3.126,  3.278, -1.879],
0019            [ 3.722,  2.309, -1.838],
0020            [ 4.042,  1.257, -1.798],
0021            [ 4.084,  0.195, -1.758],
0022            [ 3.865, -0.809, -1.717],
0023            [ 3.415, -1.693, -1.677],
0024
0025     """
0026     theta = np.linspace(-4 * np.pi, 4 * np.pi, 100)
0027     z = np.linspace(-2, 2, 100)
0028     r = z**2 + 1
0029     x = r * np.sin(theta)
0030     y = r * np.cos(theta)
0031     return np.column_stack((x, y, z))
0032
0033 def make_line_cells(num_points):
0034     """
0035     Cells of 2 indices like this yield visible divisions between the line segments
0036
0037     In [2]: cells[:5]
0038     Out[2]:
0039     array([[2, 0, 1],
0040            [2, 1, 2],
0041            [2, 2, 3],
0042            [2, 3, 4],
0043            [2, 4, 5]])
0044
0045     In [3]: cells[-5:]
0046     Out[3]:
0047     array([[ 2, 94, 95],
0048            [ 2, 95, 96],
0049            [ 2, 96, 97],
0050            [ 2, 97, 98],
0051            [ 2, 98, 99]])
0052
0053
0054     """
0055     dtype = np.int_
0056     shape = (num_points - 1, 3)
0057     cells = np.zeros( shape, dtype=dtype)
0058     cells[:, 0] = 2
0059     cells[:, 1] = np.arange(0, num_points - 1, dtype=dtype )
0060     cells[:, 2] = cells[:,1] + 1
0061     return cells
0062
0063 def make_manual_cells():
0064     """
0065     disjoint spline
0066     """
0067     dtype = np.int_
0068     cells = np.zeros( (2,11), dtype=dtype )
0069     cells[0] = [10,  0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ]
0070     cells[1] = [10, 20,21,22,23,24,25,26,27,28,29 ]
0071     return cells
0072
0073 def make_single_cells(n):
0074     """
0075     All the indices in a single cell yields a tube without visible segments
0076
0077     :param n: int
0078     :return: for n=5  np.array( [5, 0,1,2,3,4], dtype=np.int_ )
0079     """
0080     cells = np.zeros( n+1, dtype=np.int_ )
0081     cells[0] = n
0082     cells[1:] = np.arange(n)
0083     return cells
0084
0085 def make_points_lines_polydata(points, lines):
0086     """
0087     :param points: float positions (n,3)
0088     :param lines: int
0089
0090     Given an array of points, make a line set
0091     """
0092     poly = pv.PolyData()
0093     poly.points = points
0094     poly.lines = lines
0095     return poly
0096
0097 if __name__ == '__main__':
0098     points = make_points()
0099
0100     num_points = len(points)
0101
0102     #all_cells = make_line_cells(num_points)
0103     #cells = all_cells[:20] + all_cells[-20:]    ## segments
0104     #cells = all_cells[:20]
0105
0106     cells = make_manual_cells()
0107
0108     #cells = make_single_cells(num_points)
0109
0110     line = make_points_lines_polydata(points, cells)
0111     line["scalars"] = np.arange(line.n_points)
0112
0113     pl = pvplt_plotter()
0114     tube = line.tube(radius=0.1)
0115
0116     pl.add_mesh( tube, smooth_shading=True)
0117     #pl.add_mesh( line, smooth_shading=True, line_width=5 )
0118
0119     pl.show()
0120
0121
0122