Gaussian quadrature

torchquad/integration/gaussian.py

@@ -1,5 +1,4 @@
 import numpy
-import scipy
 from autoray import numpy as anp
 from .grid_integrator import GridIntegrator
 
@@ -105,9 +104,15 @@ def _cached_points_and_weights(self, N):
             tuple: nodes and weights
         """
         root_args = (N, *self.root_args)
+        if not isinstance(N, int):
+            if hasattr(N, "item"):
+                root_args = (N.item(), *self.root_args)
+            else:
+                raise NotImplementedError(
+                    f"N {N} is not an int and lacks an `item` method"
+                )
         if root_args in self._cache:
             return self._cache[root_args]
-
         self._cache[root_args] = self.root_fn(*root_args)
         return self._cache[root_args]
 

torchquad/integration/grid_integrator.py

@@ -173,7 +173,11 @@ def get_jit_compiled_integrate(
                         self.calculate_grid, static_argnames=["N"]
                     )
                     self._jax_jit_calculate_result = jax.jit(
-                        self.calculate_result, static_argnames=["dim", "n_per_dim"]
+                        self.calculate_result,
+                        static_argnums=(
+                            1,
+                            2,
+                        ),  # dim and n_per_dim
                     )
                 jit_calculate_grid = self._jax_jit_calculate_grid
                 jit_calculate_result = self._jax_jit_calculate_result
@@ -207,7 +211,7 @@ def step1(integration_domain):
 
                 dim = int(integration_domain.shape[0])
 
-                def step3(function_values, hs):
+                def step3(function_values, hs, integration_domain):
                     return self.calculate_result(
                         function_values, dim, n_per_dim, hs, integration_domain
                     )
@@ -233,15 +237,17 @@ def step3(function_values, hs):
                 if function_values.requires_grad:
                     function_values = function_values.detach()
                     function_values.requires_grad = True
-                step3 = torch.jit.trace(step3, (function_values, hs))
+                step3 = torch.jit.trace(
+                    step3, (function_values, hs, integration_domain)
+                )
 
                 # Define a compiled integrate function
                 def compiled_integrate(fn, integration_domain):
                     grid_points, hs, _ = step1(integration_domain)
                     function_values, _ = self.evaluate_integrand(
                         fn, grid_points, weights=self._weights(n_per_dim, dim, backend)
                     )
-                    result = step3(function_values, hs)
+                    result = step3(function_values, hs, integration_domain)
                     return result
 
                 return compiled_integrate
@@ -253,7 +259,9 @@ def compiled_integrate(fn, integration_domain):
             def lazy_compiled_integrate(fn, integration_domain):
                 if compiled_func[0] is None:
                     compiled_func[0] = do_compile(fn)
-                return compiled_func[0](fn, integration_domain)
+                res = compiled_func[0](fn, integration_domain)
+                print(res)
+                return res
 
             return lazy_compiled_integrate
 

torchquad/tests/gauss_test.py

@@ -12,69 +12,59 @@
 def _run_gaussian_tests(backend, _precision):
     """Test the integrate function in integration.gaussian for the given backend."""
 
-    integrators = [GaussLegendre()]
+    gauss = GaussLegendre()
 
     # 1D Tests
     N = 60
 
-    for integrator in integrators:
-        ii = integrator
-        errors, funcs = compute_integration_test_errors(
-            ii.integrate,
-            {"N": N, "dim": 1},
-            integration_dim=1,
-            use_complex=True,
-            backend=backend,
-        )
-        print(
-            f"1D {integrator} Test passed. N: {N}, backend: {backend}, Errors: {errors}"
-        )
-        # Polynomials up to degree 1 can be integrated almost exactly with gaussian.
-        for err, test_function in zip(errors, funcs):
-            assert test_function.get_order() > 1 or err < 2e-11
-        for error in errors:
-            assert error < 1e-5
+    errors, funcs = compute_integration_test_errors(
+        gauss.integrate,
+        {"N": N, "dim": 1},
+        integration_dim=1,
+        use_complex=True,
+        backend=backend,
+    )
+    print(f"1D {gauss} Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+    # Polynomials up to degree 1 can be integrated almost exactly with gaussian.
+    for err, test_function in zip(errors, funcs):
+        assert test_function.get_order() > 1 or err < 2e-11
+    for error in errors:
+        assert error < 1e-5
 
     N = 2  # integration points, here 2 for order check (2 points should lead to almost 0 err for low order polynomials)
-    for integrator in integrators:
-        ii = integrator
-        errors, funcs = compute_integration_test_errors(
-            ii.integrate,
-            {"N": N, "dim": 1},
-            integration_dim=1,
-            use_complex=True,
-            backend=backend,
-        )
-        print(
-            f"1D {integrator} Test passed. N: {N}, backend: {backend}, Errors: {errors}"
-        )
-        # All polynomials up to degree = 1 should be 0
-        # If this breaks check if test functions in helper_functions changed.
-        for err, test_function in zip(errors, funcs):
-            assert test_function.get_order() > 1 or err < 1e-15
-        for error in errors[:2]:
-            assert error < 1e-15
+
+    errors, funcs = compute_integration_test_errors(
+        gauss.integrate,
+        {"N": N, "dim": 1},
+        integration_dim=1,
+        use_complex=True,
+        backend=backend,
+    )
+    print(f"1D {gauss} Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+    # All polynomials up to degree = 1 should be 0
+    # If this breaks check if test functions in helper_functions changed.
+    for err, test_function in zip(errors, funcs):
+        assert test_function.get_order() > 1 or err < 1e-15
+    for error in errors[:2]:
+        assert error < 1e-15
 
     # 3D Tests
     N = 60**3
-    for integrator in integrators:
-        ii = integrator
-        errors, funcs = compute_integration_test_errors(
-            ii.integrate,
-            {"N": N, "dim": 3},
-            integration_dim=3,
-            use_complex=True,
-            backend=backend,
-        )
-        print(
-            f"3D {integrator} Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+
+    errors, funcs = compute_integration_test_errors(
+        gauss.integrate,
+        {"N": N, "dim": 3},
+        integration_dim=3,
+        use_complex=True,
+        backend=backend,
+    )
+    print(f"3D {gauss} Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+    for err, test_function in zip(errors, funcs):
+        assert test_function.get_order() > 1 or (
+            err < 1e-12 if test_function.is_integrand_1d else err < 1e-11
         )
-        for err, test_function in zip(errors, funcs):
-            assert test_function.get_order() > 1 or (
-                err < 1e-12 if test_function.is_integrand_1d else err < 1e-11
-            )
-        for error in errors:
-            assert error < 6e-3
+    for error in errors:
+        assert error < 6e-3
 
     # Tensorflow crashes with an Op:StridedSlice UnimplementedError with 10
     # dimensions
@@ -84,22 +74,72 @@ def _run_gaussian_tests(backend, _precision):
 
     # 10D Tests
     N = (60**3) * 3
-    for integrator in integrators:
-        ii = integrator
+
+    errors, funcs = compute_integration_test_errors(
+        gauss.integrate,
+        {"N": N, "dim": 10},
+        integration_dim=10,
+        use_complex=True,
+        backend=backend,
+    )
+    print(f"10D {gauss} Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+    for err, test_function in zip(errors, funcs):
+        assert test_function.get_order() > 1 or err < 1e-11
+    for error in errors:
+        assert error < 7000
+
+    # JIT Tests
+    if backend != "numpy":
+        N = 60
+        jit_integrate = None
+
+        def integrate(*args, **kwargs):
+            # this function initializes the jit_integrate variable with a jit'ed integrate function
+            # which is then re-used on all other integrations (as is the point of JIT).
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = gauss.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.is_integrand_1d,
+        )
+
+        print(
+            f"1D Gaussian JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        # Polynomials up to degree 1 can be integrated almost exactly with gaussian.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > float("inf") or err < 2e-10
+        for error in errors:
+            assert error < 1e-5
+
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
+
         errors, funcs = compute_integration_test_errors(
-            ii.integrate,
-            {"N": N, "dim": 10},
-            integration_dim=10,
+            integrate,
+            {},
+            integration_dim=1,
             use_complex=True,
             backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
         )
         print(
-            f"10D {integrator} Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+            f"1D Gaussian JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
         )
         for err, test_function in zip(errors, funcs):
-            assert test_function.get_order() > 1 or err < 1e-11
+            assert test_function.get_order() > 1 or err < 2e-10
         for error in errors:
-            assert error < 7000
+            assert error < 1e-5
 
 
 test_integrate_numpy = setup_test_for_backend(_run_gaussian_tests, "numpy", "float64")