From 48ea3ac3ede4a184fd2091e27418f1cdece333d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20G=C3=B3mez?= Date: Fri, 24 Jul 2026 18:42:34 +0200 Subject: [PATCH 1/2] feat(integrators): forward integrand args via integrate(..., args=...) evaluate_integrand already accepted args and called fn(points, *args), but the public integrate() methods never exposed it. Threads an args parameter through every integrator (Trapezoid, Simpson, Boole, GaussLegendre, MonteCarlo, VEGAS) so parametric integrands can be integrated without a lambda wrapper. - args defaults to None (fn(points)), so existing calls are unchanged. - Applies to the eager integrate() path; the JIT-compiled path is unchanged. Re-implements the stale PR #188 and closes #187. Co-Authored-By: Dan Barzilay --- CHANGELOG.md | 3 + tests/args_test.py | 113 +++++++++++++++++++++++ torchquad/integration/boole.py | 5 +- torchquad/integration/gaussian.py | 5 +- torchquad/integration/grid_integrator.py | 5 +- torchquad/integration/monte_carlo.py | 4 +- torchquad/integration/simpson.py | 5 +- torchquad/integration/trapezoid.py | 5 +- torchquad/integration/vegas.py | 5 +- 9 files changed, 138 insertions(+), 12 deletions(-) create mode 100644 tests/args_test.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 56d086f2..51f03258 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,9 @@ The 0.6 line is a modernization and credibility release: modern tooling, honest packaging, and closing long-open fixed issues. ### Added +- `args` argument on every integrator's `integrate()` — extra parameters are + forwarded to the integrand as `fn(points, *args)`, so parametric integrands no + longer need a lambda wrapper (#187, #188). - Optional-dependency extras: `dev`, `docs`, and CPU-convenience backend extras `torch`, `jax`, `tensorflow`, `all`. - `release_testing/` suite — slower end-to-end checks run against the latest diff --git a/tests/args_test.py b/tests/args_test.py new file mode 100644 index 00000000..ed37b570 --- /dev/null +++ b/tests/args_test.py @@ -0,0 +1,113 @@ +"""Tests for passing extra integrand arguments via ``args`` (issues #187, #188). + +Every integrator forwards ``args`` to the integrand as ``fn(points, *args)``, so a +parametric integrand can be integrated without wrapping it in a lambda. +""" + +import numpy as np + +from torchquad.integration.trapezoid import Trapezoid +from torchquad.integration.simpson import Simpson +from torchquad.integration.boole import Boole +from torchquad.integration.gaussian import GaussLegendre +from torchquad.integration.monte_carlo import MonteCarlo +from torchquad.integration.vegas import VEGAS +from helper_functions import setup_test_for_backend + +_ALPHA = 3.0 +_BETA = 2.0 +_DOMAIN = [[0.0, 1.0]] +# integral over [0, 1] of (alpha * x + beta) is alpha/2 + beta. +_EXPECTED = _ALPHA * 0.5 + _BETA + + +def _to_float(result): + """Convert a scalar backend tensor (possibly on GPU) to a Python float.""" + if hasattr(result, "cpu"): + result = result.cpu() + return float(np.asarray(result)) + + +def _parametric(x, alpha, beta): + """A parametric integrand exercising multi-argument ``*args`` unpacking.""" + return alpha * x[:, 0] + beta + + +def _args_deterministic_test(backend, dtype_name=None): + """Grid integrators must forward args; the integrand here is exact for them.""" + for integrator_cls, N in [(Trapezoid, 101), (Simpson, 101), (Boole, 101), (GaussLegendre, 32)]: + result = _to_float( + integrator_cls().integrate( + _parametric, dim=1, N=N, integration_domain=_DOMAIN, args=(_ALPHA, _BETA) + ) + ) + assert abs(result - _EXPECTED) < 1e-9, ( + f"{integrator_cls.__name__} with args gave {result}, expected {_EXPECTED}" + ) + + +def _args_default_none_test(backend, dtype_name=None): + """The default args=None must behave exactly like binding the parameters up front.""" + with_args = _to_float( + Simpson().integrate( + _parametric, dim=1, N=101, integration_domain=_DOMAIN, args=(_ALPHA, _BETA) + ) + ) + bound = _to_float( + Simpson().integrate( + lambda x: _parametric(x, _ALPHA, _BETA), dim=1, N=101, integration_domain=_DOMAIN + ) + ) + assert with_args == bound + + +def _args_monte_carlo_test(backend, dtype_name=None): + """MonteCarlo must forward args.""" + result = _to_float( + MonteCarlo().integrate( + _parametric, dim=1, N=10000, integration_domain=_DOMAIN, seed=0, args=(_ALPHA, _BETA) + ) + ) + assert abs(result - _EXPECTED) < 0.05, f"MonteCarlo with args gave {result}" + + +def _args_vegas_test(backend, dtype_name=None): + """VEGAS must forward args (numpy and torch only).""" + result = _to_float( + VEGAS().integrate( + _parametric, dim=1, N=20000, integration_domain=_DOMAIN, seed=0, args=(_ALPHA, _BETA) + ) + ) + assert abs(result - _EXPECTED) < 0.05, f"VEGAS with args gave {result}" + + +test_args_deterministic_numpy = setup_test_for_backend(_args_deterministic_test, "numpy", "float64") +test_args_deterministic_torch = setup_test_for_backend(_args_deterministic_test, "torch", "float64") +test_args_deterministic_tensorflow = setup_test_for_backend( + _args_deterministic_test, "tensorflow", "float64" +) +test_args_deterministic_jax = setup_test_for_backend(_args_deterministic_test, "jax", "float64") + +test_args_default_none_numpy = setup_test_for_backend(_args_default_none_test, "numpy", "float64") +test_args_default_none_torch = setup_test_for_backend(_args_default_none_test, "torch", "float64") + +test_args_monte_carlo_numpy = setup_test_for_backend(_args_monte_carlo_test, "numpy", "float64") +test_args_monte_carlo_torch = setup_test_for_backend(_args_monte_carlo_test, "torch", "float64") +test_args_monte_carlo_tensorflow = setup_test_for_backend( + _args_monte_carlo_test, "tensorflow", "float64" +) +test_args_monte_carlo_jax = setup_test_for_backend(_args_monte_carlo_test, "jax", "float64") + +# VEGAS supports numpy and torch only. +test_args_vegas_numpy = setup_test_for_backend(_args_vegas_test, "numpy", "float64") +test_args_vegas_torch = setup_test_for_backend(_args_vegas_test, "torch", "float64") + + +if __name__ == "__main__": + for _backend in ["numpy", "torch", "tensorflow", "jax"]: + _args_deterministic_test(_backend) + _args_monte_carlo_test(_backend) + _args_default_none_test("numpy") + for _backend in ["numpy", "torch"]: + _args_vegas_test(_backend) + print("All args tests passed!") diff --git a/torchquad/integration/boole.py b/torchquad/integration/boole.py index a04ae96a..0158cd29 100644 --- a/torchquad/integration/boole.py +++ b/torchquad/integration/boole.py @@ -11,7 +11,7 @@ class Boole(NewtonCotes): def __init__(self): super().__init__() - def integrate(self, fn, dim, N=None, integration_domain=None, backend=None): + def integrate(self, fn, dim, N=None, integration_domain=None, backend=None, args=None): """Integrates the passed function on the passed domain using Boole's rule. Args: @@ -20,11 +20,12 @@ def integrate(self, fn, dim, N=None, integration_domain=None, backend=None): N (int, optional): Total number of sample points to use for the integration. N has to be such that N^(1/dim) - 1 % 4 == 0. Defaults to 5 points per dimension if None is given. integration_domain (list or backend tensor, optional): Integration domain, e.g. [[-1,1],[0,1]]. Defaults to [-1,1]^dim. It can also determine the numerical backend. backend (string, optional): Numerical backend. Defaults to integration_domain's backend if it is a tensor and otherwise to the backend from the latest call to set_up_backend or "torch" for backwards compatibility. + args (list or tuple, optional): Extra arguments passed to the integrand as ``fn(points, *args)``. Defaults to None. Returns: backend-specific number: Integral value """ - return super().integrate(fn, dim, N, integration_domain, backend) + return super().integrate(fn, dim, N, integration_domain, backend, args=args) @staticmethod def _apply_composite_rule(cur_dim_areas, dim, hs, domain): diff --git a/torchquad/integration/gaussian.py b/torchquad/integration/gaussian.py index a626ced1..3e57c60f 100644 --- a/torchquad/integration/gaussian.py +++ b/torchquad/integration/gaussian.py @@ -24,7 +24,7 @@ def __init__(self): self._root_args = () self._cache = {} - def integrate(self, fn, dim, N=8, integration_domain=None, backend=None): + def integrate(self, fn, dim, N=8, integration_domain=None, backend=None, args=None): """Integrates the passed function on the passed domain using a Gaussian rule (Gauss-Legendre on [-1,1] as a default). Args: @@ -33,11 +33,12 @@ def integrate(self, fn, dim, N=8, integration_domain=None, backend=None): N (int, optional): Total number of sample points to use for the integration. Should be odd. Defaults to 3 points per dimension if None is given. integration_domain (list or backend tensor, optional): Integration domain, e.g. [[-1,1],[0,1]]. Defaults to [-1,1]^dim. It also determines the numerical backend if possible. backend (string, optional): Numerical backend. This argument is ignored if the backend can be inferred from integration_domain. Defaults to the backend from the latest call to set_up_backend or "torch" for backwards compatibility. + args (list or tuple, optional): Extra arguments passed to the integrand as ``fn(points, *args)``. Defaults to None. Returns: backend-specific number: Integral value """ - return super().integrate(fn, dim, N, integration_domain, backend) + return super().integrate(fn, dim, N, integration_domain, backend, args=args) def _weights(self, N, dim, backend, requires_grad=False): """return the weights, broadcast across the dimensions, generated from the polynomial of choice diff --git a/torchquad/integration/grid_integrator.py b/torchquad/integration/grid_integrator.py index 6ed71cc9..acd5e67f 100644 --- a/torchquad/integration/grid_integrator.py +++ b/torchquad/integration/grid_integrator.py @@ -29,7 +29,7 @@ def f(integration_domain, N, requires_grad=False, backend=None): def _weights(self, N, dim, backend, requires_grad=False): return None - def integrate(self, fn, dim, N, integration_domain, backend): + def integrate(self, fn, dim, N, integration_domain, backend, args=None): """Integrate the passed function on the passed domain using a Composite Newton Cotes rule. The argument meanings are explained in more detail in the sub-classes. @@ -39,6 +39,7 @@ def integrate(self, fn, dim, N, integration_domain, backend): N (int): Total number of sample points to use for the integration. integration_domain (list or backend tensor): Integration domain, e.g. [[-1,1],[0,1]]. It can also determine the numerical backend. backend (string): Numerical backend. Ignored if it can be inferred from integration_domain. + args (list or tuple, optional): Extra arguments passed to the integrand as ``fn(points, *args)``. Defaults to None. Returns: float: integral value @@ -55,7 +56,7 @@ def integrate(self, fn, dim, N, integration_domain, backend): logger.debug("Evaluating integrand on the grid.") function_values, num_points = self.evaluate_integrand( - fn, grid_points, weights=self._weights(n_per_dim, dim, backend) + fn, grid_points, weights=self._weights(n_per_dim, dim, backend), args=args ) self._nr_of_fevals = num_points diff --git a/torchquad/integration/monte_carlo.py b/torchquad/integration/monte_carlo.py index 9d57f728..85e1d254 100644 --- a/torchquad/integration/monte_carlo.py +++ b/torchquad/integration/monte_carlo.py @@ -26,6 +26,7 @@ def integrate( seed=None, rng=None, backend=None, + args=None, ): """Integrates the passed function on the passed domain using vanilla Monte Carlo Integration. @@ -37,6 +38,7 @@ def integrate( seed (int, optional): Random number generation seed to the sampling point creation, only set if provided. Defaults to None. rng (RNG, optional): An initialised RNG; this can be used when compiling the function for Tensorflow backend (string, optional): Numerical backend. Defaults to integration_domain's backend if it is a tensor and otherwise to the backend from the latest call to set_up_backend or "torch" for backwards compatibility. + args (list or tuple, optional): Extra arguments passed to the integrand as ``fn(points, *args)``. Defaults to None. Returns: backend-specific number: Integral value @@ -51,7 +53,7 @@ def integrate( integration_domain = _setup_integration_domain(dim, integration_domain, backend) sample_points = self.calculate_sample_points(N, integration_domain, seed, rng) logger.debug("Evaluating integrand") - function_values, self._nr_of_fevals = self.evaluate_integrand(fn, sample_points) + function_values, self._nr_of_fevals = self.evaluate_integrand(fn, sample_points, args=args) return self.calculate_result(function_values, integration_domain) @expand_func_values_and_squeeze_integral diff --git a/torchquad/integration/simpson.py b/torchquad/integration/simpson.py index 4c73ab52..5785806b 100644 --- a/torchquad/integration/simpson.py +++ b/torchquad/integration/simpson.py @@ -11,7 +11,7 @@ class Simpson(NewtonCotes): def __init__(self): super().__init__() - def integrate(self, fn, dim, N=None, integration_domain=None, backend=None): + def integrate(self, fn, dim, N=None, integration_domain=None, backend=None, args=None): """Integrates the passed function on the passed domain using Simpson's rule. Args: @@ -20,11 +20,12 @@ def integrate(self, fn, dim, N=None, integration_domain=None, backend=None): N (int, optional): Total number of sample points to use for the integration. Should be odd. Defaults to 3 points per dimension if None is given. integration_domain (list or backend tensor, optional): Integration domain, e.g. [[-1,1],[0,1]]. Defaults to [-1,1]^dim. It can also determine the numerical backend. backend (string, optional): Numerical backend. Defaults to integration_domain's backend if it is a tensor and otherwise to the backend from the latest call to set_up_backend or "torch" for backwards compatibility. + args (list or tuple, optional): Extra arguments passed to the integrand as ``fn(points, *args)``. Defaults to None. Returns: backend-specific number: Integral value """ - return super().integrate(fn, dim, N, integration_domain, backend) + return super().integrate(fn, dim, N, integration_domain, backend, args=args) @staticmethod def _apply_composite_rule(cur_dim_areas, dim, hs, domain): diff --git a/torchquad/integration/trapezoid.py b/torchquad/integration/trapezoid.py index b9387988..3acdcd37 100644 --- a/torchquad/integration/trapezoid.py +++ b/torchquad/integration/trapezoid.py @@ -9,7 +9,7 @@ class Trapezoid(NewtonCotes): def __init__(self): super().__init__() - def integrate(self, fn, dim, N=1000, integration_domain=None, backend=None): + def integrate(self, fn, dim, N=1000, integration_domain=None, backend=None, args=None): """Integrates the passed function on the passed domain using the trapezoid rule. Args: @@ -18,11 +18,12 @@ def integrate(self, fn, dim, N=1000, integration_domain=None, backend=None): N (int, optional): Total number of sample points to use for the integration. Defaults to 1000. integration_domain (list or backend tensor, optional): Integration domain, e.g. [[-1,1],[0,1]]. Defaults to [-1,1]^dim. It can also determine the numerical backend. backend (string, optional): Numerical backend. Defaults to integration_domain's backend if it is a tensor and otherwise to the backend from the latest call to set_up_backend or "torch" for backwards compatibility. + args (list or tuple, optional): Extra arguments passed to the integrand as ``fn(points, *args)``. Defaults to None. Returns: backend-specific number: Integral value """ - return super().integrate(fn, dim, N, integration_domain, backend) + return super().integrate(fn, dim, N, integration_domain, backend, args=args) @staticmethod def _apply_composite_rule(cur_dim_areas, dim, hs, domain): diff --git a/torchquad/integration/vegas.py b/torchquad/integration/vegas.py index bafdb549..c1abbc16 100644 --- a/torchquad/integration/vegas.py +++ b/torchquad/integration/vegas.py @@ -41,6 +41,7 @@ def integrate( max_iterations=20, use_warmup=True, backend=None, + args=None, ): """Integrates the passed function on the passed domain using VEGAS. @@ -60,6 +61,7 @@ def integrate( max_iterations (int, optional): Maximum number of vegas iterations to perform. The number of performed iterations is usually lower than this value because the number of sample points per iteration increases every fifth iteration. Defaults to 20. use_warmup (bool, optional): If True, execute a warmup to initialize the vegas map. Defaults to True. backend (string, optional): Numerical backend. "jax" and "tensorflow" are unsupported. Defaults to integration_domain's backend if it is a tensor and otherwise to the backend from the latest call to set_up_backend or "torch" for backwards compatibility. + args (list or tuple, optional): Extra arguments passed to the integrand as ``fn(points, *args)``. Defaults to None. Raises: ValueError: If the integration_domain or backend argument is invalid @@ -105,9 +107,10 @@ def integrate( domain_starts = integration_domain[:, 0] domain_sizes = integration_domain[:, 1] - domain_starts domain_volume = anp.prod(domain_sizes) + extra_args = () if args is None else args def transformed_integrand(x): - return fn(x * domain_sizes + domain_starts) * domain_volume + return fn(x * domain_sizes + domain_starts, *extra_args) * domain_volume self._fn = transformed_integrand From 1d97cac17b6ebdf44fa90590ee670c246b954357 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20G=C3=B3mez?= Date: Fri, 24 Jul 2026 18:46:35 +0200 Subject: [PATCH 2/2] feat(integrators): note VEGAS args normalization is intentionally local --- torchquad/integration/vegas.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/torchquad/integration/vegas.py b/torchquad/integration/vegas.py index c1abbc16..8332288c 100644 --- a/torchquad/integration/vegas.py +++ b/torchquad/integration/vegas.py @@ -107,6 +107,9 @@ def integrate( domain_starts = integration_domain[:, 0] domain_sizes = integration_domain[:, 1] - domain_starts domain_volume = anp.prod(domain_sizes) + # VEGAS bakes the integrand into a closure and evaluates it via _eval + # rather than passing args through evaluate_integrand, so it repeats the + # None -> () normalization here. Keep the two in sync if either changes. extra_args = () if args is None else args def transformed_integrand(x):