Best Python code snippet using lisa_python
quat_affine_test.py
Source:quat_affine_test.py
...23r2t = quat_affine.rot_list_to_tensor24v2t = quat_affine.vec_list_to_tensor25q2r = lambda q: r2t(quat_affine.quat_to_rot(q))26class QuatAffineTest(absltest.TestCase):27 def _assert_check(self, to_check, tol=1e-5):28 for k, (correct, generated) in to_check.items():29 if VERBOSE:30 logging.info(k)31 logging.info('Correct %s', correct)32 logging.info('Predicted %s', generated)33 self.assertLess(np.max(np.abs(correct - generated)), tol)34 def test_conversion(self):35 quat = jnp.array([-2., 5., -1., 4.])36 rotation = jnp.array([37 [0.26087, 0.130435, 0.956522],38 [-0.565217, -0.782609, 0.26087],39 [0.782609, -0.608696, -0.130435]])40 translation = jnp.array([1., -3., 4.])41 point = jnp.array([0.7, 3.2, -2.9])42 a = quat_affine.QuatAffine(quat, translation, unstack_inputs=True)43 true_new_point = jnp.matmul(rotation, point[:, None])[:, 0] + translation44 self._assert_check({45 'rot': (rotation, r2t(a.rotation)),46 'trans': (translation, v2t(a.translation)),47 'point': (true_new_point,48 v2t(a.apply_to_point(jnp.moveaxis(point, -1, 0)))),49 # Because of the double cover, we must be careful and compare rotations50 'quat': (q2r(a.quaternion),51 q2r(quat_affine.rot_to_quat(a.rotation))),52 })53 def test_double_cover(self):54 """Test that -q is the same rotation as q."""55 rng = jax.random.PRNGKey(42)56 keys = jax.random.split(rng)57 q = jax.random.normal(keys[0], (2, 4))58 trans = jax.random.normal(keys[1], (2, 3))59 a1 = quat_affine.QuatAffine(q, trans, unstack_inputs=True)60 a2 = quat_affine.QuatAffine(-q, trans, unstack_inputs=True)61 self._assert_check({62 'rot': (r2t(a1.rotation),63 r2t(a2.rotation)),64 'trans': (v2t(a1.translation),65 v2t(a2.translation)),66 })67 def test_homomorphism(self):68 rng = jax.random.PRNGKey(42)69 keys = jax.random.split(rng, 4)70 vec_q1 = jax.random.normal(keys[0], (2, 3))71 q1 = jnp.concatenate([72 jnp.ones_like(vec_q1)[:, :1],73 vec_q1], axis=-1)74 q2 = jax.random.normal(keys[1], (2, 4))75 t1 = jax.random.normal(keys[2], (2, 3))76 t2 = jax.random.normal(keys[3], (2, 3))77 a1 = quat_affine.QuatAffine(q1, t1, unstack_inputs=True)78 a2 = quat_affine.QuatAffine(q2, t2, unstack_inputs=True)79 a21 = a2.pre_compose(jnp.concatenate([vec_q1, t1], axis=-1))80 rng, key = jax.random.split(rng)81 x = jax.random.normal(key, (2, 3))82 new_x = a21.apply_to_point(jnp.moveaxis(x, -1, 0))83 new_x_apply2 = a2.apply_to_point(a1.apply_to_point(jnp.moveaxis(x, -1, 0)))84 self._assert_check({85 'quat': (q2r(quat_affine.quat_multiply(a2.quaternion, a1.quaternion)),86 q2r(a21.quaternion)),87 'rot': (jnp.matmul(r2t(a2.rotation), r2t(a1.rotation)),88 r2t(a21.rotation)),89 'point': (v2t(new_x_apply2),90 v2t(new_x)),91 'inverse': (x, v2t(a21.invert_point(new_x))),92 })93 def test_batching(self):94 """Test that affine applies batchwise."""95 rng = jax.random.PRNGKey(42)96 keys = jax.random.split(rng, 3)97 q = jax.random.uniform(keys[0], (5, 2, 4))98 t = jax.random.uniform(keys[1], (2, 3))99 x = jax.random.uniform(keys[2], (5, 1, 3))100 a = quat_affine.QuatAffine(q, t, unstack_inputs=True)101 y = v2t(a.apply_to_point(jnp.moveaxis(x, -1, 0)))102 y_list = []103 for i in range(5):104 for j in range(2):105 a_local = quat_affine.QuatAffine(q[i, j], t[j],106 unstack_inputs=True)107 y_local = v2t(a_local.apply_to_point(jnp.moveaxis(x[i, 0], -1, 0)))108 y_list.append(y_local)109 y_combine = jnp.reshape(jnp.stack(y_list, axis=0), (5, 2, 3))110 self._assert_check({111 'batch': (y_combine, y),112 'quat': (q2r(a.quaternion),113 q2r(quat_affine.rot_to_quat(a.rotation))),114 })115 def assertAllClose(self, a, b, rtol=1e-06, atol=1e-06):116 self.assertTrue(np.allclose(a, b, rtol=rtol, atol=atol))117 def assertAllEqual(self, a, b):118 self.assertTrue(np.all(np.array(a) == np.array(b)))119if __name__ == '__main__':...
Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Get 100 minutes of automation test minutes FREE!!