def test_rotate_zonal_harmonics_random(self): """Tests the outputs of test_rotate_zonal_harmonics.""" dtype = tf.float64 max_band = 2 zonal_coeffs = tf.constant(np.random.uniform(-1.0, 1.0, size=[3]), dtype=dtype) tensor_size = np.random.randint(3) tensor_shape = np.random.randint(1, 10, size=(tensor_size)).tolist() theta = tf.constant(np.random.uniform(0.0, np.pi, size=tensor_shape + [1]), dtype=dtype) phi = tf.constant(np.random.uniform(0.0, 2.0 * np.pi, size=tensor_shape + [1]), dtype=dtype) rotated_zonal_coeffs = spherical_harmonics.rotate_zonal_harmonics( zonal_coeffs, theta, phi) zonal_coeffs = spherical_harmonics.tile_zonal_coefficients( zonal_coeffs) l, m = spherical_harmonics.generate_l_m_permutations(max_band) l = tf.broadcast_to(l, tensor_shape + l.shape.as_list()) m = tf.broadcast_to(m, tensor_shape + m.shape.as_list()) theta_zero = tf.constant(0.0, shape=tensor_shape + [1], dtype=dtype) phi_zero = tf.constant(0.0, shape=tensor_shape + [1], dtype=dtype) gt = zonal_coeffs * spherical_harmonics.evaluate_spherical_harmonics( l, m, theta_zero, phi_zero) gt = tf.reduce_sum(input_tensor=gt, axis=-1) pred = rotated_zonal_coeffs * spherical_harmonics.evaluate_spherical_harmonics( l, m, theta + theta_zero, phi + phi_zero) pred = tf.reduce_sum(input_tensor=pred, axis=-1) self.assertAllClose(gt, pred)
def test_rotate_zonal_harmonics_jacobian_random(self): """Tests the jacobian of rotate_zonal_harmonics.""" dtype = tf.float64 tensor_size = np.random.randint(3) tensor_shape = np.random.randint(1, 5, size=(tensor_size)).tolist() theta_init = np.random.uniform(0.0, np.pi, size=tensor_shape + [1]) phi_init = np.random.uniform(0.0, 2.0 * np.pi, size=tensor_shape + [1]) zonal_coeffs = tf.convert_to_tensor(value=np.random.uniform(-1.0, 1.0, size=[3]), dtype=dtype) theta = tf.convert_to_tensor(value=theta_init) phi = tf.convert_to_tensor(value=phi_init) y = spherical_harmonics.rotate_zonal_harmonics(zonal_coeffs, theta, phi) self.assert_jacobian_is_correct(theta, theta_init, y) self.assert_jacobian_is_correct(phi, phi_init, y)
def test_rotate_zonal_harmonics_jacobian_random(self): """Tests the jacobian of rotate_zonal_harmonics.""" dtype = tf.float64 tensor_size = np.random.randint(3) tensor_shape = np.random.randint(1, 5, size=(tensor_size)).tolist() theta_init = np.random.uniform(0.0, np.pi, size=tensor_shape + [1]) phi_init = np.random.uniform(0.0, 2.0 * np.pi, size=tensor_shape + [1]) zonal_coeffs = tf.convert_to_tensor(value=np.random.uniform(-1.0, 1.0, size=[3]), dtype=dtype) # Wrap these in identities because some assert_* ops look at the constant # tensor value and mark it as unfeedable. theta = tf.identity(tf.convert_to_tensor(value=theta_init)) phi = tf.identity(tf.convert_to_tensor(value=phi_init)) y = spherical_harmonics.rotate_zonal_harmonics(zonal_coeffs, theta, phi) self.assert_jacobian_is_correct(theta, theta_init, y) self.assert_jacobian_is_correct(phi, phi_init, y)
def rotate_zonal_harmonics_fn(theta, phi): return spherical_harmonics.rotate_zonal_harmonics( zonal_coeffs, theta, phi)