diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst index d2164e21..c9d034d1 100644 --- a/docs/source/tutorial.rst +++ b/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 diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py index beeeace3..aaea3f1a 100644 --- a/torchquad/tests/boole_test.py +++ b/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") diff --git a/torchquad/tests/helper_functions.py b/torchquad/tests/helper_functions.py index c0dcc884..88a8ee84 100644 --- a/torchquad/tests/helper_functions.py +++ b/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 diff --git a/torchquad/tests/monte_carlo_test.py b/torchquad/tests/monte_carlo_test.py index ab624bb0..a8b48896 100644 --- a/torchquad/tests/monte_carlo_test.py +++ b/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" diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py index 28172c56..b9865ac8 100644 --- a/torchquad/tests/simpson_test.py +++ b/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") diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py index 0d224d1a..086cbdb7 100644 --- a/torchquad/tests/trapezoid_test.py +++ b/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")