Repository: google/jax
Pull request #22609 (2024-07-23 23:26:09)

Add `canonical` arg to Rotation.as_quat() and switch .inv() to use the quaternion conjugate.
Add `canonical` arg to Rotation.as_quat() and switch .inv() to use the quaternion conjugate.

This matches scipy behavior as of 1.11.

I also went through the tests and enabled a bunch of disabled tests which appear to pass now(?).



Patch:
diff --git a/jax/_src/scipy/spatial/transform.py b/jax/_src/scipy/spatial/transform.py
index a7366f538..46bd873bd 100644
--- a/jax/_src/scipy/spatial/transform.py
+++ b/jax/_src/scipy/spatial/transform.py
@@ -167,9 +167,12 @@ class Rotation(typing.NamedTuple):
     """Represent as rotation vectors."""
     return _as_rotvec(self.quat, degrees)
 
-  def as_quat(self) -> jax.Array:
+  def as_quat(self, canonical: bool=False) -> jax.Array:
     """Represent as quaternions."""
-    return self.quat
+    if canonical:
+      return _make_canonical(self.quat)
+    else:
+      return self.quat
 
   def inv(self):
     """Invert this rotation."""
@@ -415,7 +418,7 @@ def _from_mrp(mrp: jax.Array) -> jax.Array:
 
 @functools.partial(jnp.vectorize, signature='(n)->(n)')
 def _inv(quat: jax.Array) -> jax.Array:
-  return quat.at[3].set(-quat[3])
+  return quat * jnp.array([-1, -1, -1, 1], dtype=quat.dtype)
 
 
 @functools.partial(jnp.vectorize, signature='(n)->()')
@@ -439,3 +442,18 @@ def _normalize_quaternion(quat: jax.Array) -> jax.Array:
 @functools.partial(jnp.vectorize, signature='(n)->()')
 def _vector_norm(vector: jax.Array) -> jax.Array:
   return jnp.sqrt(jnp.dot(vector, vector))
+
+
+@functools.partial(jnp.vectorize, signature='(n)->(n)')
+def _make_canonical(quat: jax.Array) -> jax.Array:
+  is_neg = quat < 0
+  is_zero = quat == 0
+
+  neg = (
+      is_neg[3]
+      | (is_zero[3] & is_neg[0])
+      | (is_zero[3] & is_zero[0] & is_neg[1])
+      | (is_zero[3] & is_zero[0] & is_zero[1] & is_neg[2])
+  )
+
+  return jnp.where(neg, -quat, quat)


Test patch:
diff --git a/tests/scipy_spatial_test.py b/tests/scipy_spatial_test.py
index 928bac2d4..c02653dd1 100644
--- a/tests/scipy_spatial_test.py
+++ b/tests/scipy_spatial_test.py
@@ -50,10 +50,9 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype), rng(vector_shape, dtype),)
     jnp_fn = lambda q, v: jsp_Rotation.from_quat(q).apply(v, inverse=inverse)
-    # TODO(chrisflesher): re-enable this after accounting for sign degeneracy
-    # np_fn = lambda q, v: osp_Rotation.from_quat(q).apply(v, inverse=inverse).astype(dtype)  # HACK
+    np_fn = lambda q, v: osp_Rotation.from_quat(q).apply(v, inverse=inverse).astype(dtype)
     tol = 5e-2 if jtu.test_device_matches(['tpu']) else 1e-4
-    # self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=tol)
+    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=tol)
     self._CompileAndCheck(jnp_fn, args_maker, tol=tol)
 
   @jtu.sample_product(
@@ -66,7 +65,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda q: jsp_Rotation.from_quat(q).as_euler(seq=seq, degrees=degrees)
-    np_fn = lambda q: osp_Rotation.from_quat(q).as_euler(seq=seq, degrees=degrees).astype(dtype)  # HACK
+    np_fn = lambda q: osp_Rotation.from_quat(q).as_euler(seq=seq, degrees=degrees).astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -78,7 +77,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda q: jsp_Rotation.from_quat(q).as_matrix()
-    np_fn = lambda q: osp_Rotation.from_quat(q).as_matrix().astype(dtype)  # HACK
+    np_fn = lambda q: osp_Rotation.from_quat(q).as_matrix().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -90,7 +89,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda q: jsp_Rotation.from_quat(q).as_mrp()
-    np_fn = lambda q: osp_Rotation.from_quat(q).as_mrp().astype(dtype)  # HACK
+    np_fn = lambda q: osp_Rotation.from_quat(q).as_mrp().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -103,7 +102,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda q: jsp_Rotation.from_quat(q).as_rotvec(degrees=degrees)
-    np_fn = lambda q: osp_Rotation.from_quat(q).as_rotvec(degrees=degrees).astype(dtype)  # HACK
+    np_fn = lambda q: osp_Rotation.from_quat(q).as_rotvec(degrees=degrees).astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -114,8 +113,22 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
   def testRotationAsQuat(self, shape, dtype):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
-    jnp_fn = lambda q: jsp_Rotation.from_quat(jnp.where(jnp.sum(q, axis=0) > 0, q, -q)).as_quat()
-    np_fn = lambda q: osp_Rotation.from_quat(onp.where(jnp.sum(q, axis=0) > 0, q, -q)).as_quat().astype(dtype)  # HACK
+    jnp_fn = lambda q: jsp_Rotation.from_quat(q).as_quat()
+    np_fn = lambda q: osp_Rotation.from_quat(q).as_quat().astype(dtype)
+    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
+    self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
+
+  @jtu.sample_product(
+    dtype=float_dtypes,
+    shape=[(4,), (num_samples, 4)],
+  )
+  def testRotationAsQuatCanonical(self, shape, dtype):
+    if scipy_version < (1, 11, 0):
+      self.skipTest("Scipy 1.11.0 added the `canonical` arg.")
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: (rng(shape, dtype),)
+    jnp_fn = lambda q: jsp_Rotation.from_quat(q).as_quat(canonical=True)
+    np_fn = lambda q: osp_Rotation.from_quat(q).as_quat(canonical=True).astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -130,7 +143,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype), rng(other_shape, dtype),)
     jnp_fn = lambda q, o: jsp_Rotation.concatenate([jsp_Rotation.from_quat(q), jsp_Rotation.from_quat(o)]).as_rotvec()
-    np_fn = lambda q, o: osp_Rotation.concatenate([osp_Rotation.from_quat(q), osp_Rotation.from_quat(o)]).as_rotvec().astype(dtype)  # HACK
+    np_fn = lambda q, o: osp_Rotation.concatenate([osp_Rotation.from_quat(q), osp_Rotation.from_quat(o)]).as_rotvec().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -143,7 +156,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda q: jsp_Rotation.from_quat(jnp.where(jnp.sum(q, axis=0) > 0, q, -q))[indexer].as_quat()
-    np_fn = lambda q: osp_Rotation.from_quat(onp.where(onp.sum(q, axis=0) > 0, q, -q))[indexer].as_quat().astype(dtype)  # HACK
+    np_fn = lambda q: osp_Rotation.from_quat(onp.where(onp.sum(q, axis=0) > 0, q, -q))[indexer].as_quat().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -159,7 +172,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     shape = (size, len(seq))
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda a: jsp_Rotation.from_euler(seq, a, degrees).as_rotvec()
-    np_fn = lambda a: osp_Rotation.from_euler(seq, a, degrees).as_rotvec().astype(dtype)  # HACK
+    np_fn = lambda a: osp_Rotation.from_euler(seq, a, degrees).as_rotvec().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -171,9 +184,8 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda m: jsp_Rotation.from_matrix(m).as_rotvec()
-    # TODO(chrisflesher): re-enable this after accounting for sign degeneracy
-    # np_fn = lambda m: osp_Rotation.from_matrix(m).as_rotvec().astype(dtype)  # HACK
-    # self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
+    np_fn = lambda m: osp_Rotation.from_matrix(m).as_rotvec().astype(dtype)
+    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
   @jtu.sample_product(
@@ -184,7 +196,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda m: jsp_Rotation.from_mrp(m).as_rotvec()
-    np_fn = lambda m: osp_Rotation.from_mrp(m).as_rotvec().astype(dtype)  # HACK
+    np_fn = lambda m: osp_Rotation.from_mrp(m).as_rotvec().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -196,7 +208,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda r: jsp_Rotation.from_rotvec(r).as_rotvec()
-    np_fn = lambda r: osp_Rotation.from_rotvec(r).as_rotvec().astype(dtype)  # HACK
+    np_fn = lambda r: osp_Rotation.from_rotvec(r).as_rotvec().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -207,7 +219,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
   def testRotationIdentity(self, num, dtype):
     args_maker = lambda: (num,)
     jnp_fn = lambda n: jsp_Rotation.identity(n, dtype).as_rotvec()
-    np_fn = lambda n: osp_Rotation.identity(n).as_rotvec().astype(dtype)  # HACK
+    np_fn = lambda n: osp_Rotation.identity(n).as_rotvec().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -232,7 +244,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype), jnp.abs(rng(shape[0], dtype)) if rng_weights else None)
     jnp_fn = lambda q, w: jsp_Rotation.from_quat(q).mean(w).as_rotvec()
-    np_fn = lambda q, w: osp_Rotation.from_quat(q).mean(w).as_rotvec().astype(dtype)  # HACK
+    np_fn = lambda q, w: osp_Rotation.from_quat(q).mean(w).as_rotvec().astype(dtype)
     tol = 5e-3  # 1e-4 too tight for TF32
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=tol)
     self._CompileAndCheck(jnp_fn, args_maker, tol=tol)
@@ -247,7 +259,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype), rng(other_shape, dtype))
     jnp_fn = lambda q, o: (jsp_Rotation.from_quat(q) * jsp_Rotation.from_quat(o)).as_rotvec()
-    np_fn = lambda q, o: (osp_Rotation.from_quat(q) * osp_Rotation.from_quat(o)).as_rotvec().astype(dtype)  # HACK
+    np_fn = lambda q, o: (osp_Rotation.from_quat(q) * osp_Rotation.from_quat(o)).as_rotvec().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -259,7 +271,21 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(shape, dtype),)
     jnp_fn = lambda q: jsp_Rotation.from_quat(q).inv().as_rotvec()
-    np_fn = lambda q: osp_Rotation.from_quat(q).inv().as_rotvec().astype(dtype)  # HACK
+    np_fn = lambda q: osp_Rotation.from_quat(q).inv().as_rotvec().astype(dtype)
+    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
+    self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
+
+  @jtu.sample_product(
+    dtype=float_dtypes,
+    shape=[(4,), (num_samples, 4)],
+  )
+  def testRotationInvConjugate(self, shape, dtype):
+    if scipy_version < (1, 11, 0):
+      self.skipTest("Scipy prior to 1.11.0 used a negative conjugate.")
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: (rng(shape, dtype),)
+    jnp_fn = lambda q: jsp_Rotation.from_quat(q).inv().as_quat()
+    np_fn = lambda q: osp_Rotation.from_quat(q).inv().as_quat().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
@@ -297,7 +323,7 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
     args_maker = lambda: (rng(shape, dtype),)
     times = jnp.arange(shape[0], dtype=dtype)
     jnp_fn = lambda q: jsp_Slerp.init(times, jsp_Rotation.from_quat(q))(compute_times).as_rotvec()
-    np_fn = lambda q: osp_Slerp(times, osp_Rotation.from_quat(q))(compute_times).as_rotvec().astype(dtype)  # HACK
+    np_fn = lambda q: osp_Slerp(times, osp_Rotation.from_quat(q))(compute_times).as_rotvec().astype(dtype)
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)