Repository: scipy/scipy
Pull request: ENH: Add args argument for callable in quad_vec
#### Reference issue
Fixes #14412 

#### What does this implement/fix?
Provision to provide arguments for the callable function we wish to integrate with `scipy.integrate.quad_vec`.

#### Additional information
`scipy.integrate.quadrature` and `scipy.integrate.quad` both provide this argument. So it makes sense that it should be available in `quad_vec`.

Since `quad_vec` was added later, this got skipped somehow. Although when I dug up to find a good reason, I found the `args` argument to pass extra arguments was actually supported in the initial commit [here from PR #9642](https://github.com/scipy/scipy/pull/9642/commits/f48e866f9c6356423c0d648a5b015c373abf04ba), but got removed in the following [commit](https://github.com/scipy/scipy/pull/9642/commits/9fe95612b81a196887ac2bf60b890a717ee11feb) when the decision (see [comment #9642](https://github.com/scipy/scipy/pull/9642#issuecomment-451122360)) to make `quad_vec` a public function was taken.

Problem statement:
Extra function arguments to `scipy.integrate.quad_vec`
Hello, thanks for the work on scipy! 

I was looking at the docs/source for `scipy.integrate.quad_vec` and I noticed that there isn't a provision to provide arguments for the callable function we wish to integrate. For the `scipy.integrate.quadrature`, we have the `args` argument that we can pass to the function for such a behaviour. Is there a way to do that with `quad_vec`?

Thanks!

Hints from issue discussion:
Thanks for the suggestion. In the meantime, you can pass the parameters by using [functools.partial](https://docs.python.org/3/library/functools.html#functools.partial) for instance.
Ah thanks! I forgot about partial, for some reason. I'll do that for now. But it'll be great to have this feature built-in too in future!
@tupui I can send a PR to add this argument. Should be straightforward.

Commit message: TST: Add tests for args kwarg in quad_vec

Patch to review:
diff --git a/scipy/integrate/_quad_vec.py b/scipy/integrate/_quad_vec.py
index 0c25f8db6653..ff4b401fd951 100644
--- a/scipy/integrate/_quad_vec.py
+++ b/scipy/integrate/_quad_vec.py
@@ -101,8 +101,21 @@ def __repr__(self):
                                              for k in self.__keys))
 
 
+class _FunctionWrapper:
+    """
+    Object to wrap user function, allowing picklability
+    """
+    def __init__(self, f, args):
+        self.f = f
+        self.args = args
+
+    def __call__(self, x):
+        return self.f(x, *self.args)
+
+
 def quad_vec(f, a, b, epsabs=1e-200, epsrel=1e-8, norm='2', cache_size=100e6, limit=10000,
-             workers=1, points=None, quadrature=None, full_output=False):
+             workers=1, points=None, quadrature=None, full_output=False,
+             *, args=()):
     r"""Adaptive integration of a vector-valued function.
 
     Parameters
@@ -140,6 +153,8 @@ def quad_vec(f, a, b, epsabs=1e-200, epsrel=1e-8, norm='2', cache_size=100e6, li
         Default: 'gk21' for finite intervals and 'gk15' for (semi-)infinite
     full_output : bool, optional
         Return an additional ``info`` dictionary.
+    args : tuple, optional
+        Extra arguments to pass to function, if any.
 
     Returns
     -------
@@ -215,6 +230,13 @@ def quad_vec(f, a, b, epsabs=1e-200, epsrel=1e-8, norm='2', cache_size=100e6, li
     a = float(a)
     b = float(b)
 
+    if args:
+        if not isinstance(args, tuple):
+            args = (args,)
+
+        # create a wrapped function to allow the use of map and Pool.map
+        f = _FunctionWrapper(f, args)
+
     # Use simple transformations to deal with integrals over infinite
     # intervals.
     kwargs = dict(epsabs=epsabs,
diff --git a/scipy/integrate/tests/test__quad_vec.py b/scipy/integrate/tests/test__quad_vec.py
index 2e536118e104..4153f21167b9 100644
--- a/scipy/integrate/tests/test__quad_vec.py
+++ b/scipy/integrate/tests/test__quad_vec.py
@@ -87,6 +87,15 @@ def test_quad_vec_simple_inf(quadrature):
     assert_allclose(res, exact, rtol=0, atol=max(epsabs, 1.5 * err))
 
 
+def test_quad_vec_args():
+    f = lambda x, a: x * (x + a) * np.arange(3)
+    a = 2
+    exact = np.array([0, 4/3, 8/3])
+
+    res, err = quad_vec(f, 0, 1, args=(a,))
+    assert_allclose(res, exact, rtol=0, atol=1e-4)
+
+
 def _lorenzian(x):
     return 1 / (1 + x**2)
 
@@ -104,6 +113,25 @@ def test_quad_vec_pool():
         assert_allclose(res, np.pi, rtol=0, atol=1e-4)
 
 
+def _func_with_args(x, a):
+    return x * (x + a) * np.arange(3)
+
+
+def test_quad_vec_pool_args():
+    from multiprocessing.dummy import Pool
+
+    f = _func_with_args
+    a = 2
+    exact = np.array([0, 4/3, 8/3])
+
+    res, err = quad_vec(f, 0, 1, args=(a,), workers=4)
+    assert_allclose(res, exact, rtol=0, atol=1e-4)
+
+    with Pool(10) as pool:
+        res, err = quad_vec(f, 0, 1, args=(a,), workers=pool.map)
+        assert_allclose(res, exact, rtol=0, atol=1e-4)
+
+
 @quadrature_params
 def test_num_eval(quadrature):
     def f(x):


Code review comments:
diff_hunk: @@ -101,8 +101,21 @@ def __repr__(self):
                                              for k in self.__keys))
 
 
+class _FunctionWrapper:
path: scipy/integrate/_quad_vec.py
text: @user1:
I would put this in a common place like `scipy._lib.util` (or we create another file for multiprocessing with this function and `MapWrapper`). Because this is already used in `DifferentialEvolutionSolver`. It would be better to have just a single implementation of this.

@author:
Sure! How about a file inside _lib named `_multiprocessing_utils.py`?
I'll add both these functions there and make the necessary refactoring at other places/tests.

@user1:
LGTM, thanks for doing this.

@author:
I actually changed my mind later. I felt moving everything into `_multiprocessing_util.py` caused unnecessary refactoring. Moving this also meant to take out the tests for `MapWrapper` from `test__util` and create a new test file specific for these tests. All this did not make a lot of sense.

Instead, I just moved `_FunctionWrapper` inside `scipy._lib._util` which also serves the purpose and avoids polluting git blame.

diff_hunk: @@ -104,6 +113,25 @@ def test_quad_vec_pool():
         assert_allclose(res, np.pi, rtol=0, atol=1e-4)
 
 
+def _func_with_args(x, a):
+    return x * (x + a) * np.arange(3)
+
+
+def test_quad_vec_pool_args():
+    from multiprocessing.dummy import Pool
path: scipy/integrate/tests/test__quad_vec.py
text: @author:
Since it is now used in two test methods:`test_quad_vec_pool_args` and `test_quad_vec_pool`, @user1 should I move this import to the top of the file?

@user1:
Sure