(feat): add JIT integrator tests

docs/source/tutorial.rst

@@ -586,7 +586,7 @@ Compiling the integrate method
 ``````````````````````````````
 
 To speed up the quadrature in situations where it is executed often with the
-same number of points ``N`` and dimensionality ``dim``,
+same number of points ``N``, dimensionality ``dim``, and shape of the ``integrand``,
 we can JIT-compile the performance-relevant parts of the integrate method:
 
 .. code:: ipython3

torchquad/tests/boole_test.py

@@ -61,6 +61,59 @@ def _run_boole_tests(backend, _precision):
     # for error in errors:
     # assert error < 5e-9
 
+    # JIT Tests
+    if backend != "numpy":
+        N = 401
+        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 = bl.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 Boole JIT Test passed for 1D integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D Boole JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
 
 test_integrate_numpy = setup_test_for_backend(_run_boole_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_boole_tests, "torch", "float64")

torchquad/tests/helper_functions.py

@@ -324,6 +324,7 @@ def compute_integration_test_errors(
     use_complex,
     backend,
     use_multi_dim_integrand=True,
+    filter_test_functions=lambda x: x,
 ):
     """Computes errors on all test functions for given dimension and integrator.
 
@@ -334,7 +335,7 @@ def compute_integration_test_errors(
         use_complex (Boolean): If True, skip complex example functions.
         backend (string): Numerical backend for the example functions.
         use_multi_dim_integrand (bool, optional): Whether or not to allow for a multi-dimensional integrand i.e an array of integrands
-
+        filter_test_functions (function, optional): function for filtering which test functions to run
     Returns:
         (list, list): Absolute errors on all example functions and the chosen
             example functions
@@ -344,8 +345,9 @@ def compute_integration_test_errors(
 
     # Compute integration errors on the chosen functions and remember those
     # functions
-    for test_function in get_test_functions(
-        integration_dim, backend, use_multi_dim_integrand
+    for test_function in filter(
+        filter_test_functions,
+        get_test_functions(integration_dim, backend, use_multi_dim_integrand),
     ):
         if not use_complex and test_function.is_complex:
             continue

torchquad/tests/monte_carlo_test.py

@@ -82,6 +82,68 @@ def _run_monte_carlo_tests(backend, _precision):
     for error in errors:
         assert error < 26
 
+    # JIT Tests
+    if backend != "numpy":
+        N = 100000
+        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 = mc.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 MC JIT Test passed for 1D integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 0 or err < 1e-14
+
+        # If this breaks check if test functions in helper_functions changed.
+        for error in errors[:3]:
+            assert error < 1e-2
+
+        assert errors[3] < 0.5
+        assert errors[4] < 32.0
+
+        for error in errors[6:10]:
+            assert error < 2e-2
+
+        for error in errors[10:]:
+            assert error < 35.0
+
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D MC 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() > 0 or err < 1e-14
+
 
 test_integrate_numpy = setup_test_for_backend(
     _run_monte_carlo_tests, "numpy", "float32"

torchquad/tests/simpson_test.py

@@ -89,6 +89,62 @@ def _run_simpson_tests(backend, _precision):
     for error in errors:
         assert error < 5e-9
 
+    # JIT Tests
+    if backend != "numpy":
+        N = 100001
+        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 = simp.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 Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 3 or (
+                err < 3e-11 if test_function.is_integrand_1d else err < 6e-10
+            )  # errors add up if the integrand is higher dimensional
+        for error in errors:
+            assert error < 1e-7
+
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D Simpson 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() > 3 or (
+                err < 3e-11 if test_function.is_integrand_1d else err < 6e-10
+            )  # errors add up if the integrand is higher dimensional
+        for error in errors:
+            assert error < 1e-7
+
 
 test_integrate_numpy = setup_test_for_backend(_run_simpson_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_simpson_tests, "torch", "float64")

torchquad/tests/trapezoid_test.py

@@ -86,6 +86,62 @@ def _run_trapezoid_tests(backend, _precision):
     for error in errors:
         assert error < 7000
 
+    # JIT Tests
+    if backend != "numpy":
+        N = 100000
+        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 = tp.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 Trapezoid JIT Test passed for 1D integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 1 or (
+                err < 2e-11 if test_function.is_integrand_1d else err < 5e-10
+            )  # errors add up if the integrand is higher dimensional
+        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(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D Trapezoid 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 < 2e-11 if test_function.is_integrand_1d else err < 5e-10
+            )  # errors add up if the integrand is higher dimensional
+        for error in errors:
+            assert error < 1e-5
+
 
 test_integrate_numpy = setup_test_for_backend(_run_trapezoid_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_trapezoid_tests, "torch", "float64")