Many Simultaneous 2D Integrands allocating more memory than expected

[kovnerd]

I am trying to use Vegas to simultaneously perform many 2D integrations (something like Nvec**2 where Nvec is ~500). Using the straightforward approach yields a memory allocation error, where an array of shape (Nvec**2, Nvec**2) was failed to be allocated, despite the fact that my integrand is of shape (Nvec, Nvec). My best guess for this behavior is that vegas is trying to allocate memory for the covariance between each integrand, hence requiring a(Nvec**2, Nvec**2) array.
Is there a way to instruct Vegas integrators to not compute the covariance between all integrands? In principle, I do not require a rigorous error estimate of for these integrals, but I would like to share the integration grid between each of them. Below is an example cell from a JupyterNotebook where Nvec=20, producing an integration result of shape (20,20), and Nvec=556, producing the OOM traceback.

import numpy as np
import gvar as gv
import vegas


Nvec = 20
nu = np.linspace(0,14,num=Nvec)
@vegas.lbatchintegrand
def _inte(xy):
    x,y = xy[:,0],xy[:,1]
    nux,muy = np.multiply.outer(x-0.5, nu), np.multiply.outer(y-0.5, nu)
    kern = np.exp(-0.5*(x-y)**2)
    w = np.stack([np.multiply.outer(np.cos(nx),np.cos(my)) for nx,my in zip(nux,muy)],0)
    res = np.expand_dims(kern,(-1,-2))*w#should be of shape [Nxy, Nnu, Nnu]
    return res

integ = vegas.Integrator(map=[(0.,1.)]*2)
integ(_inte,  adapt=True)

result = integ(_inte, adapt=False)
print('result shape: ', result.shape)


Nvec = 556
nu = np.linspace(0,14,num=Nvec)
@vegas.lbatchintegrand
def _inte(xy):
    x,y = xy[:,0],xy[:,1]
    nux,muy = np.multiply.outer(x-0.5, nu), np.multiply.outer(y-0.5, nu)
    kern = np.exp(-0.5*(x-y)**2)
    w = np.stack([np.multiply.outer(np.cos(nx),np.cos(my)) for nx,my in zip(nux,muy)],0)
    res = np.expand_dims(kern,(-1,-2))*w#should be of shape [Nxy, Nnu, Nnu]
    return res

integ = vegas.Integrator(map=[(0.,1.)]*2)
integ(_inte,  adapt=True)

result = integ(_inte, adapt=False)
print('result shape: ', result.shape)

Output is

result shape:  (20, 20)
---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
Cell In[6], line 36
     33     return res
     35 integ = vegas.Integrator(map=[(0.,1.)]*2)
---> 36 integ(_inte,  adapt=True)
     38 result = integ(_inte, adapt=False)
     39 print('result shape: ', result.shape)

File src[/vegas/_vegas.pyx:2024](http://localhost:18889/vegas/_vegas.pyx#line=2023), in vegas._vegas.Integrator.__call__()

MemoryError: Unable to allocate 712. GiB for an array with shape (309136, 309136) and data type float64

[gplepage]

Yes, it does look like the memory needed for the covariance matrix is upsetting the computer. There is currently no way to turn the covariance matrices off but it looks like it might not be too difficult to make this an option. The code would still compute variances for each component of the integrand's output, but it would ignore correlations between different components. I will see what I can do but it could take a couple of weeks.

In the meantime one thing you could do for the big case is to first generate the grid by integrating just the res[:, 0, 0] component of your integrand. Replace the code right after you define _inte with:

@vegas.lbatchintegrand
def _inte00(xy):
    return np.array(_inte(xy)[:,0,0])

integ = vegas.Integrator(map=[(0.,1.)]*2)
w = integ(_inte00)
print('00 result =', w)
print()

Then use the grid from that integration to integrate the entire matrix, but using your own custom version of the vegas.Integrator (see the second half of https://vegas.readthedocs.io/en/latest/tutorial.html#vegas-as-a-random-number-generator in the documentation) that does not calculate correlations:

integral = 0.0
variance = 0.0
batch_f = _inte
nitn = 10
for itn in range(nitn):
    for x, wgt, hcube in integ.random_batch(yield_hcube=True):
        wgt_fx = wgt[:, None, None] * batch_f(x)
        # iterate over hypercubes: compute variance for each,
        #                          and accumulate for final result
        for i in range(hcube[0], hcube[-1] + 1):
            idx = (hcube == i)          # select array items for h-cube i
            nwf = np.sum(idx)           # number of points in h-cube i
            wf = wgt_fx[idx]
            sum_wf = np.sum(wf, axis=0)         # sum of wgt * f(x) for h-cube i
            sum_wf2 = np.sum(wf ** 2, axis=0)   # sum of (wgt * f(x)) ** 2
            integral += sum_wf
            variance += (sum_wf2 * nwf - sum_wf ** 2) / (nwf - 1.)
result = gv.gvar(integral / nitn, (variance / nitn / nitn) ** 0.5)

print(result.shape)
print(result[:5, :5])

This code, which is equivalent to what vegas does when adapt=False (but without correlations), gives the following output:

00 result = 0.92430(13)

(556, 556)
[[0.92442(12) 0.92440(12) 0.92433(12) 0.92421(12) 0.92404(12)]
 [0.92440(12) 0.92438(12) 0.92431(12) 0.92419(12) 0.92402(12)]
 [0.92433(12) 0.92431(12) 0.92423(12) 0.92412(12) 0.92395(12)]
 [0.92421(12) 0.92419(12) 0.92412(12) 0.92400(12) 0.92383(12)]
 [0.92404(12) 0.92402(12) 0.92395(12) 0.92383(12) 0.92366(12)]]

I believe this is what you want.