Created
May 8, 2020 18:09
-
-
Save mikedh/49be00831a56c76fc736c23e0f8d0e2b to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import os | |
| import trimesh | |
| import numpy as np | |
| tol_norm = 1e-4 | |
| def test_align(align): | |
| """ | |
| Test an "align two vectors" function | |
| """ | |
| is_rigid = trimesh.transformations.is_rigid | |
| # start with some edge cases and make sure the transform works | |
| target = np.array([0, 0, -1], dtype=np.float64) | |
| vectors = np.vstack(( | |
| trimesh.unitize(np.random.random((1000, 3)) - .5), | |
| [-target, target], | |
| trimesh.util.generate_basis(target), | |
| [[7.12106798e-07, -7.43194705e-08, 1.00000000e+00], | |
| [0, 0, -1], | |
| [1e-4, 1e-4, -1]])) | |
| for vector in vectors: | |
| T, a = align(vector, target, return_angle=True) | |
| assert is_rigid(T) | |
| assert np.isclose(np.linalg.det(T), 1.0) | |
| # rotate vector with transform | |
| check = np.dot(T[:3, :3], vector) | |
| # compare to target vector | |
| norm = np.linalg.norm(check - target) | |
| assert norm < tol_norm | |
| # these vectors should be perpendicular and zero | |
| angles = [align(i, target, return_angle=True)[1] | |
| for i in trimesh.util.generate_basis(target)] | |
| assert np.allclose( | |
| angles, [np.pi / 2, np.pi / 2, 0.0]) | |
| # generate angles from 0 to 180 degrees | |
| angles = np.linspace(0.0, np.pi / 1e7, 10000) | |
| # generate on- plane vectors | |
| vectors = np.column_stack((np.cos(angles), | |
| np.sin(angles), | |
| np.zeros(len(angles)))) | |
| T = align([0, 0, -1], [-1e-17, 1e-17, 1]) | |
| assert np.isclose(np.linalg.det(T), 1.0) | |
| T = align([0, 0, -1], [-1e-4, 1e-4, 1]) | |
| assert np.isclose(np.linalg.det(T), 1.0) | |
| vector_1 = np.array([7.12106798e-07, -7.43194705e-08, 1.00000000e+00]) | |
| vector_2 = np.array([0, 0, -1]) | |
| T, angle = align(vector_1, vector_2, return_angle=True) | |
| assert np.isclose(np.linalg.det(T), 1.0) | |
| def generate_basis(vector): | |
| vector = np.array(vector) | |
| u = np.linalg.svd(vector[:, np.newaxis])[0] | |
| if np.linalg.det(u) < 0.0: | |
| u[:, -1] *= -1 | |
| return u | |
| def align_svd(a, b, return_angle=False): | |
| """ | |
| Find the rotation matrix that transforms one 3D vector | |
| to another. | |
| Parameters | |
| ------------ | |
| a : (3,) float | |
| Unit vector | |
| b : (3,) float | |
| Unit vector | |
| return_angle : bool | |
| Return the angle between vectors or not | |
| Returns | |
| ------------- | |
| matrix : (4, 4) float | |
| Homogenous transform to rotate from `a` to `b` | |
| angle : float | |
| If `return_angle` angle in radians between `a` and `b` | |
| """ | |
| a = np.array(a, dtype=np.float64) | |
| b = np.array(b, dtype=np.float64) | |
| if a.shape != (3,) or b.shape != (3,): | |
| raise ValueError('vectors must be (3,)!') | |
| # unitize vectors in place | |
| a /= np.linalg.norm(a) | |
| b /= np.linalg.norm(b) | |
| # find the SVD of the two vectors | |
| au = np.linalg.svd(a.reshape((-1, 1)))[0] | |
| bu = np.linalg.svd(b.reshape((-1, 1)))[0] | |
| if np.linalg.det(au) < 0: | |
| au[:, -1] *= -1 | |
| if np.linalg.det(bu) < 0: | |
| bu[:, -1] *= -1 | |
| # put rotation into homogenous transformation | |
| matrix = np.eye(4) | |
| matrix[:3, :3] = bu.dot(au.T) | |
| if return_angle: | |
| # projection of a onto b | |
| dot = np.dot(a, b) | |
| # clip to avoid floating point error | |
| angle = np.arccos(np.clip(dot, -1.0, 1.0)) | |
| if dot < -1e-5: | |
| angle += np.pi | |
| return matrix, angle | |
| return matrix | |
| if __name__ == '__main__': | |
| import cProfile | |
| import pstats | |
| import io | |
| import time | |
| pr = cProfile.Profile() | |
| pr.enable() | |
| tic = [time.time()] | |
| test_align(align_svd) | |
| tic.append(time.time()) | |
| print(r'\n\n\`align_svd`\n\n') | |
| # ... do something ... | |
| pr.disable() | |
| s = io.StringIO() | |
| sortby = 'cumulative' | |
| ps = pstats.Stats(pr, stream=s).sort_stats(sortby) | |
| ps.print_stats() | |
| print(s.getvalue()) | |
| pr = cProfile.Profile() | |
| pr.enable() | |
| tic.append(time.time()) | |
| test_align(trimesh.geometry.align_vectors) | |
| tic.append(time.time()) | |
| print(r'\n\n\`trimesh.geometry.align_vectors`\n\n') | |
| pr.disable() | |
| s = io.StringIO() | |
| sortby = 'cumulative' | |
| ps = pstats.Stats(pr, stream=s).sort_stats(sortby) | |
| ps.print_stats() | |
| print(s.getvalue()) | |
| print('`align_svd`: {}s\n`trimesh.geometry.align_vectors`: {}s\n'.format( | |
| *np.diff(tic)[[0, -1]])) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Results: