feat: forward integrand args via integrate(..., args=...)

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

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!")

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):

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

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
 

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

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):

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):

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,13 @@ 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):
-            return fn(x * domain_sizes + domain_starts) * domain_volume
+            return fn(x * domain_sizes + domain_starts, *extra_args) * domain_volume
 
         self._fn = transformed_integrand