Feature: improve memory handling in Kirchhoff operator

[mrava87]

Motivation

This PR is mainly aimed at improving the memory efficiency of the Kirchhoff operator.

In the current implementation, we create an overall trav table (similar for amp in dynamic=True mode) of size nynxnz x nsnr. Assuming ns=nr this scales quadratically with the ns.

In the new implementation, we handle source and receiver tables independently (trav_srcs and trav_recs) and not merge them into a unique table. Assuming ns=nr this scales linearly with the ns.

Changes

The following changes are introduced:

• 4 new methods are created: _travsrcrec_kirch_matvec, _travsrcrec_kirch_rmatvec, _amp_kirch_matvec, and _amp_kirch_rmatvec, which are the counterparts of _trav_kirch_matvec, _trav_kirch_rmatvec, _ampsrcrec_kirch_matvec, and _ampsrcrec_kirch_rmatvec which used full tables
• trav and amp can now be provided as nd.arrays or tuples of 2 nd.arrays. In the second case, these Arte the src and rec tables. Internally the __init__ operator handles both cases to ensure backward compatibility.
• the order and number of outputs in _traveltime_table is changed. However, being this a private method we should not expect users to use it (aka users should not expect us not to make changes), so I believe this is fine.
• handled Homogenize imports of optional libraries #430

Performance tests

I run the following two performance tests using this code:

#!/usr/bin/env python
# coding: utf-8

import numpy as np
import matplotlib.pyplot as plt
import scipy as sp

from pylops.utils.wavelets import *
from pylops.waveeqprocessing.kirchhoff import Kirchhoff

## 2D layered in homogenous velocity

# Velocity Model
nx, nz = 301, 100
dx, dz = 4, 4
x, z = np.arange(nx)*dx, np.arange(nz)*dz
v0 = 1000 # initial velocity
kv = 0. # gradient
vel = np.outer(np.ones(nx), v0 +kv*z) 

# Reflectivity Model
refl = np.zeros((nx, nz))
#refl[nx//2, 50] = -1
refl[:, 10] = -1
refl[:, -10] = 0.5

# Receivers
nr = 101
rx = np.linspace(10*dx, (nx-10)*dx, nr)
rz = 20*np.ones(nr)
recs = np.vstack((rx, rz))
dr = recs[0,1]-recs[0,0]

# Sources
ns = 30
sx = np.linspace(dx*10, (nx-10)*dx, ns)
sz = 10*np.ones(ns)
sources = np.vstack((sx, sz))

# Model data and invert it
nt = 651
dt = 0.004
t = np.arange(nt)*dt
wav, wavt, wavc = ricker(t[:41], f0=20)

kop = Kirchhoff(z, x, t, sources, recs, v0, wav, wavc, mode='analytic', engine='numba')

d = kop * refl.ravel()
d = d.reshape(ns, nr, nt)

madj = kop.H * d.ravel()
madj = madj.reshape(nx, nz)

print('Done')

and equivalent code with old version of Kirchoff.

• dynamic=False
  
  

• dynamic=True
  
  

[cako]

This is definitely a good design change and I'm happy you kept the default for v2. We can then remove the old functionality in v3.