1"""
2sample.py
3------------
4
5Randomly sample surface and volume of meshes.
6"""
7import numpy as np
8
9from . import util
10from . import transformations
11
12
13def sample_surface(mesh, count):
14    """
15    Sample the surface of a mesh, returning the specified
16    number of points
17
18    For individual triangle sampling uses this method:
19    http://mathworld.wolfram.com/TrianglePointPicking.html
20
21    Parameters
22    ---------
23    mesh : trimesh.Trimesh
24      Geometry to sample the surface of
25    count : int
26      Number of points to return
27
28    Returns
29    ---------
30    samples : (count, 3) float
31      Points in space on the surface of mesh
32    face_index : (count,) int
33      Indices of faces for each sampled point
34    """
35
36    # len(mesh.faces) float, array of the areas
37    # of each face of the mesh
38    area = mesh.area_faces
39    # total area (float)
40    area_sum = np.sum(area)
41    # cumulative area (len(mesh.faces))
42    area_cum = np.cumsum(area)
43    face_pick = np.random.random(count) * area_sum
44    face_index = np.searchsorted(area_cum, face_pick)
45
46    # pull triangles into the form of an origin + 2 vectors
47    tri_origins = mesh.triangles[:, 0]
48    tri_vectors = mesh.triangles[:, 1:].copy()
49    tri_vectors -= np.tile(tri_origins, (1, 2)).reshape((-1, 2, 3))
50
51    # pull the vectors for the faces we are going to sample from
52    tri_origins = tri_origins[face_index]
53    tri_vectors = tri_vectors[face_index]
54
55    # randomly generate two 0-1 scalar components to multiply edge vectors by
56    random_lengths = np.random.random((len(tri_vectors), 2, 1))
57
58    # points will be distributed on a quadrilateral if we use 2 0-1 samples
59    # if the two scalar components sum less than 1.0 the point will be
60    # inside the triangle, so we find vectors longer than 1.0 and
61    # transform them to be inside the triangle
62    random_test = random_lengths.sum(axis=1).reshape(-1) > 1.0
63    random_lengths[random_test] -= 1.0
64    random_lengths = np.abs(random_lengths)
65
66    # multiply triangle edge vectors by the random lengths and sum
67    sample_vector = (tri_vectors * random_lengths).sum(axis=1)
68
69    # finally, offset by the origin to generate
70    # (n,3) points in space on the triangle
71    samples = sample_vector + tri_origins
72
73    return samples, face_index
74
75
76def volume_mesh(mesh, count):
77    """
78    Use rejection sampling to produce points randomly
79    distributed in the volume of a mesh.
80
81
82    Parameters
83    ---------
84    mesh : trimesh.Trimesh
85      Geometry to sample
86    count : int
87      Number of points to return
88
89    Returns
90    ---------
91    samples : (n, 3) float
92      Points in the volume of the mesh where n <= count
93    """
94    points = (np.random.random((count, 3)) * mesh.extents) + mesh.bounds[0]
95    contained = mesh.contains(points)
96    samples = points[contained][:count]
97    return samples
98
99
100def volume_rectangular(extents,
101                       count,
102                       transform=None):
103    """
104    Return random samples inside a rectangular volume,
105    useful for sampling inside oriented bounding boxes.
106
107    Parameters
108    ----------
109    extents :   (3,) float
110      Side lengths of rectangular solid
111    count : int
112      Number of points to return
113    transform : (4, 4) float
114      Homogeneous transformation matrix
115
116    Returns
117    ---------
118    samples : (count, 3) float
119      Points in requested volume
120    """
121    samples = np.random.random((count, 3)) - .5
122    samples *= extents
123    if transform is not None:
124        samples = transformations.transform_points(samples,
125                                                   transform)
126    return samples
127
128
129def sample_surface_even(mesh, count, radius=None):
130    """
131    Sample the surface of a mesh, returning samples which are
132    VERY approximately evenly spaced. This is accomplished by
133    sampling and then rejecting pairs that are too close together.
134
135    Note that since it is using rejection sampling it may return
136    fewer points than requested (i.e. n < count). If this is the
137    case a log.warning will be emitted.
138
139    Parameters
140    ---------
141    mesh : trimesh.Trimesh
142      Geometry to sample the surface of
143    count : int
144      Number of points to return
145    radius : None or float
146      Removes samples below this radius
147
148    Returns
149    ---------
150    samples : (n, 3) float
151      Points in space on the surface of mesh
152    face_index : (n,) int
153      Indices of faces for each sampled point
154    """
155    from .points import remove_close
156
157    # guess radius from area
158    if radius is None:
159        radius = np.sqrt(mesh.area / (3 * count))
160
161    # get points on the surface
162    points, index = sample_surface(mesh, count * 3)
163
164    # remove the points closer than radius
165    points, mask = remove_close(points, radius)
166
167    # we got all the samples we expect
168    if len(points) >= count:
169        return points[:count], index[mask][:count]
170
171    # warn if we didn't get all the samples we expect
172    util.log.warning('only got {}/{} samples!'.format(
173        len(points), count))
174
175    return points, index[mask]
176
177
178def sample_surface_sphere(count):
179    """
180    Correctly pick random points on the surface of a unit sphere
181
182    Uses this method:
183    http://mathworld.wolfram.com/SpherePointPicking.html
184
185    Parameters
186    ----------
187    count : int
188      Number of points to return
189
190    Returns
191    ----------
192    points : (count, 3) float
193      Random points on the surface of a unit sphere
194    """
195    # get random values 0.0-1.0
196    u, v = np.random.random((2, count))
197    # convert to two angles
198    theta = np.pi * 2 * u
199    phi = np.arccos((2 * v) - 1)
200    # convert spherical coordinates to cartesian
201    points = util.spherical_to_vector(
202        np.column_stack((theta, phi)))
203    return points
204