Refactor BernsteinPoly test for prior evaluation

Author: heptaflar

test/b/test_b_jax.py

@@ -202,37 +202,8 @@ def test_from_lookup_table(self):
         )
         y = BernsteinGrid(x).eval(c)
 
-        f = BernsteinPoly.from_lookup_table(k, x, y, non_negative=True)
-        b = f.prior()
-        self.assertEqual(c.shape, b.shape)
-        self.assertFalse(np.any(b < 0.0))
-        self.assertAlmostEqual(c[0, 0, 0], b[0, 0, 0])
-        self.assertAlmostEqual(c[0, 0, 1], b[0, 0, 1])
-        self.assertAlmostEqual(c[0, 0, 2], b[0, 0, 2])
-        self.assertAlmostEqual(c[0, 1, 0], b[0, 1, 0])
-        self.assertAlmostEqual(c[0, 1, 1], b[0, 1, 1])
-        self.assertAlmostEqual(c[0, 1, 2], b[0, 1, 2])
-        self.assertAlmostEqual(c[0, 2, 0], b[0, 2, 0])
-        self.assertAlmostEqual(c[0, 2, 1], b[0, 2, 1])
-        self.assertAlmostEqual(c[0, 2, 2], b[0, 2, 2])
-        self.assertAlmostEqual(c[1, 0, 0], b[1, 0, 0])
-        self.assertAlmostEqual(c[1, 0, 1], b[1, 0, 1])
-        self.assertAlmostEqual(c[1, 0, 2], b[1, 0, 2])
-        self.assertAlmostEqual(c[1, 1, 0], b[1, 1, 0])
-        self.assertAlmostEqual(c[1, 1, 1], b[1, 1, 1])
-        self.assertAlmostEqual(c[1, 1, 2], b[1, 1, 2])
-        self.assertAlmostEqual(c[1, 2, 0], b[1, 2, 0])
-        self.assertAlmostEqual(c[1, 2, 1], b[1, 2, 1])
-        self.assertAlmostEqual(c[1, 2, 2], b[1, 2, 2])
-        self.assertAlmostEqual(c[2, 0, 0], b[2, 0, 0])
-        self.assertAlmostEqual(c[2, 0, 1], b[2, 0, 1])
-        self.assertAlmostEqual(c[2, 0, 2], b[2, 0, 2])
-        self.assertAlmostEqual(c[2, 1, 0], b[2, 1, 0])
-        self.assertAlmostEqual(c[2, 1, 1], b[2, 1, 1])
-        self.assertAlmostEqual(c[2, 1, 2], b[2, 1, 2])
-        self.assertAlmostEqual(c[2, 2, 0], b[2, 2, 0])
-        self.assertAlmostEqual(c[2, 2, 1], b[2, 2, 1])
-        self.assertAlmostEqual(c[2, 2, 2], b[2, 2, 2])
+        f = BernsteinPoly.from_lookup_table(k, x, y)
+        self.assertTrue(np.allclose(BernsteinGrid(x).eval(f.prior())), y)
 
 
 class BSolveTest(unittest.TestCase):