tests.make_tri_pyramid()   B
last analyzed

Complexity

Conditions 5

Size

Total Lines 15

Duplication

Lines 0
Ratio 0 %
Metric Value
cc 5
dl 0
loc 15
rs 8.5454
1
import numpy as np
2
import math
3
4
5
def make_triangle(sx, sy, dx, dy, delta):
6
    '''Create a right rectangle, alinged with X and Y axes, with x-side size sX
7
    and y-side size sY. Triangle is offset by dX and dY.'''
8
    x1 = np.arange(0, sx, delta)
9
    y1 = np.zeros(x1.shape)
10
11
    y2 = np.arange(0, sy, delta)
12
    x2 = np.zeros(y2.shape)
13
14
    x3 = np.arange(0, sx, delta)
15
    y3 = sy - x3 * sy / sx
16
17
    xs = np.hstack([x1, x2, x3]) - dx
18
    ys = np.hstack([y1, y2, y3]) - dy
19
20
    return xs, ys
21
22
23
def make_tri_pyramid(sx, sy, sz, dx, dy, dz, delta):
24
    '''Create a right rectangle triangular pyramid, alinged with X and Y axes,
25
    with x-side at the base size of sX and y-side size at the base of sY.
26
    Pyramid has high sZ. It is offset by dX, dY and dZ.'''
27
    points = []
28
    for z in np.arange(0, sz, delta):
29
        ai = sx - z * sx / sz
30
        bi = sy - z * sy / sz
31
        xs, ys = make_triangle(ai, bi, dx, dy, delta)
32
        points.append((xs, ys, z * np.ones(xs.shape)))
33
    xs = np.hstack([x for x, y, z in points])
34
    ys = np.hstack([y for x, y, z in points])
35
    zs = np.hstack([z for x, y, z in points]) - dz
36
    points = np.vstack([xs, ys, zs]).T
37
    return points
38
39
40
def make_tri_pyramid_footprint(sx, sy, sz, dx, dy, dz):
41
    '''Create the footprint of a pyramid created by make_tri_pyramid'''
42
    footprint = np.array([
43
        [0, 0, 0],
44
        [0, sy, 0],
45
        [sx, 0, 0],
46
        [0, 0, 0],
47
    ])
48
    footprint[:, 0] -= dx
49
    footprint[:, 1] -= dy
50
    footprint[:, 2] -= dz
51
    return footprint
52
53
54
def make_tri_pyramid_with_base(side, delta, offset):
55
    '''Create a pyramid as per make_tri_pyramid, suroundeded by a triangular
56
    flat base.'''
57
    rng = np.random.RandomState(0)
58
    sx = side / 2
59
    sy = side
60
    sz = side / 4
61
62
    dx = offset[0] + side / 2
63
    dy = offset[1] + side / 2
64
    dz = offset[2]
65
66
    points = make_tri_pyramid(sx, sy, sz, dx, dy, dz, delta)
67
    _add_noise(points, 0.1, rng)
68
69
    for s in np.arange(0, side * 0.05, delta):
70
        xs, ys = make_triangle(sx * (1 + s), sy * (1 + s), dx + s, dy + s,
71
                               delta)
72
        zs = np.zeros(xs.shape) - dz
73
        tmp = np.vstack([xs, ys, zs]).T
74
        points = np.vstack([points, tmp])
75
76
    footprint = make_tri_pyramid_footprint(sx, sy, sz, dx, dy, dz)
77
    return points, footprint
78
79
80
def _add_noise(points, size, rng):
81
    '''Add noise to an array of 2D points'''
82
    points += (rng.rand(points.shape[0], points.shape[1]) - 0.5) * size
83
84
85
def perpendicular_2d(a):
86
    '''Create a vector perpendicular to the original'''
87
    b = np.zeros(a.shape)
88
    b[0] = -a[1]
89
    b[1] = a[0]
90
    return b
91
92
93
def rotation_around_axis(axis, theta):
94
    '''
95
    Return the rotation matrix associated with counterclockwise rotation about
96
    the given axis by theta radians.
97
    '''
98
    axis = np.asarray(axis)
99
    theta = np.asarray(theta)
100
    axis = axis / math.sqrt(np.dot(axis, axis))
101
    a = math.cos(theta / 2)
102
    b, c, d = -axis * math.sin(theta / 2)
103
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
104
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
105
    return np.array([[aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
106
                     [2 * (bc - ad), aa+cc-bb-dd, 2 * (cd + ab)],
107
                     [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]])
108
109
110
def make_half_red_stick(point_from, point_to, width=0.04, num_pts_per_line=50,
111
                        num_lines_per_stick=25):
112
    ''' Make a hollow red-white-red-white stick '''
113
    point_from = np.asarray(point_from, dtype=float)[:3]
114
    point_to = np.asarray(point_to, dtype=float)[:3]
115
    length = np.linalg.norm(point_to - point_from)
116
    width = length * 0.04
117
    axis = (point_to - point_from) * 1.0 / length
118
    origin = perpendicular_2d(axis) * width
119
120
    points = np.zeros((num_pts_per_line * num_lines_per_stick, 6))
121
    points[:, 3:6] = 255
122
123
    idx = 0
124
    unitline = np.linspace(0, 1, num_pts_per_line)
125
    for theta in np.linspace(0, math.pi, num_lines_per_stick):
126
        src = np.dot(rotation_around_axis(axis, theta), origin)
127
128
        # straight slope [0, 1]
129
        line = np.array((unitline, unitline, unitline)).T
130
        # linear function with slope
131
        line = point_from + src + (point_to - point_from) * line
132
        points[idx:idx + num_pts_per_line, :3] = line
133
134
        points[idx:idx + num_pts_per_line/2, 3:6] = [210, 25, 30]
135
        idx += num_pts_per_line
136
137
    return points
138
139
140
def make_red_stick(point_from, point_to, **kwargs):
141
    ''' Make a hollow red-white-red-white stick '''
142
    point_from = np.asarray(point_from, dtype=float)[:3]
143
    point_to = np.asarray(point_to, dtype=float)[:3]
144
    halfway = (point_to + point_from)/2
145
    half1 = make_half_red_stick(point_from, halfway, **kwargs)
146
    half2 = make_half_red_stick(halfway, point_to, **kwargs)
147
    return np.concatenate((half1, half2), axis=0)
148