def test_evaluate_legendre_polynomial_preset(self):
        """Tests that evaluate_legendre_polynomial generates expected results."""
        tensor_size = np.random.randint(1, 3)
        tensor_shape = np.random.randint(1, 10, size=(tensor_size)).tolist()
        x = np.random.uniform(size=tensor_shape)

        with self.subTest(name="l_0_m_0"):
            l = tf.constant(0, shape=tensor_shape)
            m = tf.constant(0, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = np.ones_like(x)

            self.assertAllClose(pred, gt)

        with self.subTest(name="l_1_m_0"):
            l = tf.constant(1, shape=tensor_shape)
            m = tf.constant(0, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = x

            self.assertAllClose(pred, gt)

        with self.subTest(name="l_1_m_1"):
            l = tf.constant(1, shape=tensor_shape)
            m = tf.constant(1, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = -tf.sqrt(1.0 - x * x)

            self.assertAllClose(pred, gt)

        with self.subTest(name="l_2_m_0"):
            l = tf.constant(2, shape=tensor_shape)
            m = tf.constant(0, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = 0.5 * (3.0 * x * x - 1.0)

            self.assertAllClose(pred, gt)

        with self.subTest(name="l_2_m_1"):
            l = tf.constant(2, shape=tensor_shape)
            m = tf.constant(1, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = -3.0 * x * tf.sqrt(1.0 - x * x)

            self.assertAllClose(pred, gt)

        with self.subTest(name="l_2_m_2"):
            l = tf.constant(2, shape=tensor_shape)
            m = tf.constant(2, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = 3.0 * (1.0 - x * x)

            self.assertAllClose(pred, gt)

        with self.subTest(name="l_3_m_0"):
            l = tf.constant(3, shape=tensor_shape)
            m = tf.constant(0, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = 0.5 * x * (5.0 * x * x - 3.0)

            self.assertAllClose(pred, gt)

        with self.subTest(name="l_3_m_1"):
            l = tf.constant(3, shape=tensor_shape)
            m = tf.constant(1, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = 1.5 * (1.0 - 5.0 * x * x) * tf.sqrt(1.0 - x * x)

            self.assertAllClose(pred, gt)

        with self.subTest(name="l_3_m_2"):
            l = tf.constant(3, shape=tensor_shape)
            m = tf.constant(2, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = 15.0 * x * (1.0 - x * x)

            self.assertAllClose(pred, gt)

        with self.subTest(name="l_3_m_3"):
            l = tf.constant(3, shape=tensor_shape)
            m = tf.constant(3, shape=tensor_shape)
            pred = spherical_harmonics.evaluate_legendre_polynomial(l, m, x)
            gt = -15.0 * tf.pow(1.0 - x * x, 1.5)

            self.assertAllClose(pred, gt)
Exemplo n.º 2
0
 def evaluate_legendre_polynomial_fn(x):
     return spherical_harmonics.evaluate_legendre_polynomial(
         l_init, m_init, x)