How to use the pylops.utils.seismicevents.makeaxis function in pylops

def test_parabolic2d(par):
    """Create small dataset with a parabolic event and check that output
    contains the event apex at correct time and correct amplitude
    """
    # Data creation
    t0 = 50
    px = 0
    pxx = 1e-1
    amp = 0.6

    # Create axes
    t, _, x, _ = makeaxis(par)

    # Create data
    d, dwav = parabolic2d(x, t, t0, px, pxx, amp, np.ones(1))

    # Assert shape
    assert d.shape[0] == par['nx']
    assert d.shape[1] == par['nt']

    assert dwav.shape[0] == par['nx']
    assert dwav.shape[1] == par['nt']

    # Assert correct position of event
    assert_array_equal(d[par['nx']//2, t0], amp)

def test_linear3d(par):
    """Create small dataset with an horizontal event and check output
    contains event at correct time and correct amplitude
    """
    # Data creation
    v = 1
    t0 = 50
    theta = 0.
    phi = 0.
    amp = 0.6

    # Create axes
    t, _, x, y = makeaxis(par)

    # Create data
    d, dwav = linear3d(x, y, t, v, t0, theta, phi, amp, wav)

    #Assert shape
    assert d.shape[0] == par['ny']
    assert d.shape[1] == par['nx']
    assert d.shape[2] == par['nt']

    assert dwav.shape[0] == par['ny']
    assert dwav.shape[1] == par['nx']
    assert dwav.shape[2] == par['nt']

    # Assert correct position of event
    assert_array_equal(d[:, :, t0],
                       amp*np.ones((par['ny'], par['nx'])))

par['nt2'] = par['nt']
    v = 1500
    it0_m = 25
    t0_m = it0_m * par['dt']
    theta_m = 0
    phi_m = 0
    amp_m = 1.

    it0_G = np.array([25, 50, 75])
    t0_G = it0_G * par['dt']
    theta_G = (0, 0, 0)
    phi_G = (0, 0, 0)
    amp_G = (1., 0.6, 2.)

    # Create axis
    t, _, x, y = makeaxis(par)

    # Create wavelet
    wav = ricker(t[:41], f0=par['f0'])[0]

    # Generate model
    _, mwav = linear3d(x, x, t, v, t0_m, theta_m, phi_m, amp_m, wav)

    # Generate operator
    _, Gwav = linear3d(x, y, t, v, t0_G, theta_G, phi_G, amp_G, wav)

    # Add negative part to data and model
    if par['twosided']:
        mwav = np.concatenate((np.zeros((par['nx'], par['nx'], par['nt'] - 1)),
                               mwav), axis=-1)
        Gwav = np.concatenate((np.zeros((par['ny'], par['nx'], par['nt'] - 1)),
                               Gwav), axis=-1)

parmod = {'ox': -100, 'dx': 10, 'nx': 21,
          'oy': -50, 'dy': 10, 'ny': 11,
          'ot': 0, 'dt': 0.004, 'nt': 50,
          'f0': 40}

par1 = {'ny': 8, 'nx': 10, 'nt': 20,
        'dtype': 'float32'}  # even
par2 = {'ny': 9, 'nx': 11, 'nt': 21,
        'dtype': 'complex64'}  # odd

# deghosting params
vel_sep = 1000.0 # velocity at separation level
zrec = 20.0 # depth of receivers

# axes and wavelet
t, t2, x, y = makeaxis(parmod)
wav = ricker(t[:41], f0=parmod['f0'])[0]

@pytest.fixture(scope="module")
def create_data2D():
    """Create 2d dataset
    """
    t0_plus = np.array([0.02, 0.08])
    t0_minus = t0_plus + 0.04
    vrms = np.array([1400., 1800.])
    amp = np.array([1., -0.6])

    p2d_minus = hyperbolic2d(x, t, t0_minus, vrms, amp, wav)[1].T

    kx = np.fft.ifftshift(np.fft.fftfreq(parmod['nx'], parmod['dx']))
    freq = np.fft.rfftfreq(parmod['nt'], parmod['dt'])

'f0': 20, 'nfmax': 200}

t0_m = 0.2
vrms_m = 1100.0
amp_m = 1.0

t0_G = (0.2, 0.5, 0.7)
vrms_G = (1200., 1500., 2000.)
amp_G = (1., 0.6, 0.5)

# Taper
tap = taper3d(par['nt'], (par['ny'], par['nx']),
              (5, 5), tapertype='hanning')

# Create axis
t, t2, x, y = makeaxis(par)

# Create wavelet
wav = ricker(t[:41], f0=par['f0'])[0]

# Generate model
m, mwav = hyperbolic2d(x, t, t0_m, vrms_m, amp_m, wav)

# Generate operator
G, Gwav = np.zeros((par['ny'], par['nx'], par['nt'])), \
          np.zeros((par['ny'], par['nx'], par['nt']))
for iy, y0 in enumerate(y):
    G[iy], Gwav[iy] = hyperbolic2d(x-y0, t, t0_G, vrms_G, amp_G, wav)
G, Gwav = G*tap, Gwav*tap

# Add negative part to data and model
m = np.concatenate((np.zeros((par['nx'], par['nt']-1)), m), axis=-1)

'f0': 30}

# linear events
v = 1500
t0 = [0.1, 0.2, 0.3]
theta = [0, 0, 0]
amp = [1., -2, 0.5]

# parabolic event
tp0 = [0.13]
px = [0]
pxx = [5e-7]
ampp = [0.7]

# create axis
taxis, taxis2, xaxis, yaxis = pylops.utils.seismicevents.makeaxis(par)

# create wavelet
wav = ricker(taxis[:41], f0=par['f0'])[0]

# generate model
y = pylops.utils.seismicevents.linear2d(xaxis, taxis, v, t0,
                                        theta, amp, wav)[1] + \
    pylops.utils.seismicevents.parabolic2d(xaxis, taxis, tp0,
                                           px, pxx, ampp, wav)[1]

###############################################################################
# We can now create the :py:class:`pylops.signalprocessing.Radon2D` operator.
# We also apply its adjoint to the data to obtain a representation of those
# 3 linear events overlapping to a parabolic event in the Radon domain.
# Similarly, we feed the operator to a sparse solver like
# :py:class:`pylops.optimization.sparsity.FISTA` to obtain a sparse

'f0': 20, 'nfmax': 200}

t0_m = [0.2]
vrms_m = [700.]
amp_m = [1.]

t0_G = [0.2, 0.5, 0.7]
vrms_G = [800., 1200., 1500.]
amp_G = [1., 0.6, 0.5]

# Taper
tap = taper3d(par['nt'], [par['ny'], par['nx']],
              (5, 5), tapertype='hanning')

# Create axis
t, t2, x, y = makeaxis(par)

# Create wavelet
wav = ricker(t[:41], f0=par['f0'])[0]

# Generate model
m, mwav = hyperbolic2d(x, t, t0_m, vrms_m, amp_m, wav)

# Generate operator
G, Gwav = np.zeros((par['ny'], par['nx'], par['nt'])), \
          np.zeros((par['ny'], par['nx'], par['nt']))
for iy, y0 in enumerate(y):
    G[iy], Gwav[iy] = hyperbolic2d(x-y0, t, t0_G, vrms_G, amp_G, wav)
G, Gwav = G*tap, Gwav*tap

# Add negative part to data and model
m = np.concatenate((np.zeros((par['nx'], par['nt']-1)), m), axis=-1)

###############################################################################
# Let's start by creating an 2-dimensional array of size :math:`n_x \times n_t`
# composed of 3 parabolic events
par = {'ox':-140, 'dx':2, 'nx':140,
       'ot':0, 'dt':0.004, 'nt':200,
       'f0': 20}

v = 1500
t0 = [0.2, 0.4, 0.5]
px = [0, 0, 0]
pxx = [1e-5, 5e-6, 1e-20]
amp = [1., -2, 0.5]

# Create axis
t, t2, x, y = pylops.utils.seismicevents.makeaxis(par)

# Create wavelet
wav = pylops.utils.wavelets.ricker(t[:41], f0=par['f0'])[0]

# Generate model
_, data = pylops.utils.seismicevents.parabolic2d(x, t, t0, px,
                                                 pxx, amp, wav)

###############################################################################
# We want to divide this 2-dimensional data into small overlapping
# patches in the spatial direction and apply the adjoint of the
# :py:class:`pylops.signalprocessing.Radon2D` operator to each patch. This is
# done by simply using the adjoint of the
# :py:class:`pylops.signalprocessing.Sliding2D` operator
nwins = 5
winsize = 36

import matplotlib.pyplot as plt

import pylops

plt.close('all')

############################################
# Let's first define the time and space axes as well as some auxiliary input
# parameters that we will use to create a Ricker wavelet
par = {'ox':-200, 'dx':2, 'nx':201,
       'oy':-100, 'dy':2, 'ny':101,
       'ot':0, 'dt':0.004, 'nt':501,
       'f0': 20, 'nfmax': 210}

# Create axis
t, t2, x, y = pylops.utils.seismicevents.makeaxis(par)

# Create wavelet
wav = pylops.utils.wavelets.ricker(np.arange(41) * par['dt'],
                                   f0=par['f0'])[0]

############################################
# We want to create a 2d data with a number of crossing linear events using the
# :py:func:`pylops.utils.seismicevents.linear2d` routine.
v = 1500
t0 = [0.2, 0.7, 1.6]
theta = [40, 0, -60]
amp = [1., 0.6, -2.]

mlin, mlinwav = pylops.utils.seismicevents.linear2d(x, t, v, t0,
                                                    theta, amp, wav)