|
- """Transformations for 3D coordinates."""
-
- import mindspore.numpy as mnp
-
-
- def vecs_sub(v1, v2):
- """Computes v1 - v2."""
- return v1 - v2
-
-
- def vecs_robust_norm(v, epsilon=1e-8):
- """Computes norm of vectors 'v'."""
-
- return mnp.sqrt(mnp.sum(mnp.square(v), axis=-1) + epsilon)
-
-
- def vecs_robust_normalize(v, epsilon=1e-8):
- """Normalizes vectors 'v'."""
-
- norms = vecs_robust_norm(v, epsilon)
- return v / norms[..., None]
-
-
- def vecs_dot_vecs(v1, v2):
- """Dot product of vectors 'v1' and 'v2'."""
- return mnp.sum(v1 * v2, axis=-1)
-
-
- def vecs_cross_vecs(v1, v2):
- """Cross product of vectors 'v1' and 'v2'."""
-
- return mnp.concatenate(((v1[..., 1] * v2[..., 2] - v1[..., 2] * v2[..., 1])[..., None],
- (v1[..., 2] * v2[..., 0] - v1[..., 0] * v2[..., 2])[..., None],
- (v1[..., 0] * v2[..., 1] - v1[..., 1] * v2[..., 0])[..., None]), axis=-1)
-
-
- def rots_from_two_vecs(e0_unnormalized, e1_unnormalized):
- """Create rotation matrices from unnormalized vectors for the x and y-axes."""
-
- # Normalize the unit vector for the x-axis, e0.
- e0 = vecs_robust_normalize(e0_unnormalized)
-
- # make e1 perpendicular to e0.
- c = vecs_dot_vecs(e1_unnormalized, e0)
- e1 = e1_unnormalized - c[..., None] * e0
- e1 = vecs_robust_normalize(e1)
-
- # Compute e2 as cross product of e0 and e1.
- e2 = vecs_cross_vecs(e0, e1)
-
- rots = mnp.concatenate(
- (mnp.concatenate([e0[..., 0][None, ...], e1[..., 0][None, ...], e2[..., 0][None, ...]], axis=0)[None, ...],
- mnp.concatenate([e0[..., 1][None, ...], e1[..., 1][None, ...], e2[..., 1][None, ...]], axis=0)[None, ...],
- mnp.concatenate([e0[..., 2][None, ...], e1[..., 2][None, ...], e2[..., 2][None, ...]], axis=0)[None, ...]),
- axis=0)
- return rots
-
-
- def rigids_from_3_points(
- point_on_neg_x_axis, # shape (...)
- origin, # shape (...)
- point_on_xy_plane, # shape (...)
- ): # shape (...)
- """Create Rigids from 3 points. """
-
- m = rots_from_two_vecs(
- e0_unnormalized=vecs_sub(origin, point_on_neg_x_axis),
- e1_unnormalized=vecs_sub(point_on_xy_plane, origin))
- return m, origin
-
-
- def invert_rots(m):
- """Computes inverse of rotations 'm'."""
-
- return mnp.transpose(m, (1, 0, 2, 3, 4))
-
-
- def rots_mul_vecs(m, v):
- """Apply rotations 'm' to vectors 'v'."""
-
- return mnp.concatenate(((m[0][0] * v[..., 0] + m[0][1] * v[..., 1] + m[0][2] * v[..., 2])[..., None],
- (m[1][0] * v[..., 0] + m[1][1] * v[..., 1] + m[1][2] * v[..., 2])[..., None],
- (m[2][0] * v[..., 0] + m[2][1] * v[..., 1] + m[2][2] * v[..., 2])[..., None]), axis=-1)
-
-
- def invert_rigids(rot, trans):
- """Computes group inverse of rigid transformations 'r'."""
-
- inv_rots = invert_rots(rot)
- t = rots_mul_vecs(inv_rots, trans)
- inv_trans = -t
- return inv_rots, inv_trans
-
-
- def vecs_add(v1, v2):
- """Add two vectors 'v1' and 'v2'."""
-
- return v1 + v2
-
-
- def rigids_mul_vecs(rot, trans, v):
- """Apply rigid transforms 'r' to points 'v'."""
-
- return vecs_add(rots_mul_vecs(rot, v), trans)
|