How to use the pylops.Restriction function in pylops

ax2.plot(np.arange(-N, N + 1) * dt, diags.T, '--r', lw=1)
ax2.plot(np.arange(-N, N + 1) * dt, diag_ave, 'k', lw=4)
ax2.set_title("Averaged covariance 'filter'")

###############################################################################
# Let's define now the sampling operator as well as create our covariance
# matrices in terms of linear operators. This may not be strictly necessary
# here but shows how even Bayesian-type of inversion can very easily scale to
# large model and data spaces.

# Sampling operator
perc_subsampling = 0.2
ntsub = int(np.round(nt*perc_subsampling))
iava = np.sort(np.random.permutation(np.arange(nt))[:ntsub])
iava[-1] = nt-1 # assume we have the last sample to avoid instability
Rop = pylops.Restriction(nt, iava, dtype='float64')

# Covariance operators
Cm_op = \
    pylops.signalprocessing.Convolve1D(nt, diag_ave, offset=N)
Cd_op = sigmad**2 * pylops.Identity(ntsub)

###############################################################################
# We model now our data and add noise that respects our prior definition
n = np.random.normal(0, sigmad, nt)
y = Rop * x
yn = Rop * (x + n)
ymask = Rop.mask(x)
ynmask = Rop.mask(x + n)

###############################################################################
# First we apply the Bayesian inversion equation

wav = inputdata['wav'][::2]
wav_c = np.argmax(wav)
x = np.apply_along_axis(convolve, 1, x, wav, mode='full')
x = x[:, wav_c:][:, :par['nt']]

# gain
gain = np.tile((taxis**2)[:, np.newaxis], (1, par['nx'])).T
x = x*gain

# subsampling locations
perc_subsampling = 0.5
Nsub = int(np.round(par['nx']*perc_subsampling))
iava = np.sort(np.random.permutation(np.arange(par['nx']))[:Nsub])

# restriction operator
Rop = pylops.Restriction(par['nx']*par['nt'], iava,
                         dims=(par['nx'], par['nt']),
                         dir=0, dtype='float64')

y = Rop*x.flatten()
xadj = Rop.H*y.flatten()

y = y.reshape(Nsub, par['nt'])
xadj = xadj.reshape(par['nx'], par['nt'])

# apply mask
ymask = Rop.mask(x.flatten())

# sliding windows with radon transform
dx = par['dx']
nwins = 4
nwin = 27

axs[0].set_title('Data(frequency domain)')
axs[1].plot(t, x, 'k', LineWidth=2)
axs[1].set_title('Data(time domain)')
axs[1].axis('tight')

###############################################################################
# We now define the locations at which the signal will be sampled.

# subsampling locations
perc_subsampling = 0.2
Nsub = int(np.round(N*perc_subsampling))

iava = np.sort(np.random.permutation(np.arange(N))[:Nsub])

# Create restriction operator
Rop = pylops.Restriction(N, iava, dtype='float64')

y = Rop*x
ymask = Rop.mask(x)

# Visualize data
fig = plt.figure(figsize=(12, 4))
plt.plot(t, x, 'k', lw=3)
plt.plot(t, x, '.k', ms=20, label='all samples')
plt.plot(t, ymask, '.g', ms=15, label='available samples')
plt.legend()
plt.title('Data restriction')

###############################################################################
# To start let's consider the simplest *'solver'*, i.e., *least-square inversion
# without regularization*. We aim here to minimize the following cost function:
#

N = 200
dt = 0.004
t = np.arange(N)*dt
x = np.zeros(N)

for freq, amp in zip(freqs, amps):
    x = x + amp*np.sin(2*np.pi*freq*t)

# subsampling locations
perc_subsampling = 0.4
Nsub = int(np.round(N*perc_subsampling))

iava = np.sort(np.random.permutation(np.arange(N))[:Nsub])

# Create restriction operator
Rop = pylops.Restriction(N, iava, dtype='float64')

y = Rop*x
ymask = Rop.mask(x)

# Visualize data
fig = plt.figure(figsize=(15, 5))
plt.plot(t, x, 'k', lw=3)
plt.plot(t, x, '.k', ms=20, label='all samples')
plt.plot(t, ymask, '.g', ms=15, label='available samples')
plt.legend()
plt.title('Data restriction')

###############################################################################
# Back to our first cost function, this can be very easily implemented using the
# :math:`/` for PyLops operators, which will automatically call the
# :py:func:`scipy.sparse.linalg.lsqr` with some default parameters.

###############################################################################
# To start we import a 2d image and define our restriction operator to irregularly and randomly
# sample the image for 30% of the entire grid
im = np.load('../testdata/python.npy')[:, :, 0]

Nz, Nx = im.shape
N = Nz * Nx

# Subsample signal
perc_subsampling = 0.2

Nsub2d = int(np.round(N*perc_subsampling))
iava = np.sort(np.random.permutation(np.arange(N))[:Nsub2d])

# Create operators and data
Rop = pylops.Restriction(N, iava, dtype='float64')
D2op = pylops.Laplacian((Nz, Nx), weights=(1, 1), dtype='float64')

x = im.flatten()
y = Rop * x
y1 = Rop.mask(x)

###############################################################################
# We will now use two different routines from our optimization toolbox
# to estimate our original image in the regular grid.

xcg_reg_lop = \
    pylops.optimization.leastsquares.NormalEquationsInversion(Rop, [D2op], y,
                                                              epsRs=[np.sqrt(0.1)],
                                                              returninfo=False,
                                                              **dict(maxiter=200))

subax.set_ylim([0.5, -0.5])

###############################################################################
# Finally we show how the :py:class:`pylops.Restriction` is not limited to
# one dimensional signals but can be applied to sample locations of a specific
# axis of a multi-dimensional array.
# subsampling locations
nx, nt = 100, 50

x = np.random.normal(0, 1, (nx, nt))

perc_subsampling = 0.4
nxsub = int(np.round(nx*perc_subsampling))
iava = np.sort(np.random.permutation(np.arange(nx))[:nxsub])

Rop = pylops.Restriction(nx*nt, iava, dims=(nx, nt), dir=0, dtype='float64')
y = (Rop*x.ravel()).reshape(nxsub, nt)
ymask = Rop.mask(x)

fig, axs = plt.subplots(1, 3, figsize=(10, 5))
axs[0].imshow(x.T, cmap='gray')
axs[0].set_title('Model')
axs[0].axis('tight')
axs[1].imshow(y.T, cmap='gray')
axs[1].set_title('Data')
axs[1].axis('tight')
axs[2].imshow(ymask.T, cmap='gray')
axs[2].set_title('Masked model')
axs[2].axis('tight')

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

# model
_, x = linear2d(xaxis, taxis, v, t0_m, theta_m, amp_m, wav)

###############################################################################
# We can now define the spatial locations along which the data has been
# sampled. In this specific example we will assume that we have access only to
# 40% of the 'original' locations.
perc_subsampling = 0.6
nxsub = int(np.round(par['nx']*perc_subsampling))

iava = np.sort(np.random.permutation(np.arange(par['nx']))[:nxsub])

# restriction operator
Rop = pylops.Restriction(par['nx']*par['nt'],
                         iava, dims=(par['nx'], par['nt']),
                         dir=0, dtype='float64')

# data
y = Rop*x.ravel()
y = y.reshape(nxsub, par['nt'])

# mask
ymask = Rop.mask(x.flatten())

# inverse
xinv = Rop / y.ravel()
xinv = xinv.reshape(par['nx'], par['nt'])

fig, axs = plt.subplots(1, 2, sharey=True, figsize=(5, 4))
axs[0].imshow(x.T, cmap='gray', vmin=-2, vmax=2,

dims = (nrec, dimsd[1])
    else:
        dimsd = data.shape
        dims = (nrec[0], nrec[1], dimsd[2])

    # create restriction/interpolation operator
    if iava.dtype == float:
        Rop = Interp(np.prod(dims), iava, dims=dims,
                     dir=0, kind='linear', dtype=dtype)
        if ndims == 3 and iava1 is not None:
            dims1 = (len(iava), nrec[1], dimsd[2])
            Rop1 = Interp(np.prod(dims1), iava1, dims=dims1,
                          dir=1, kind='linear', dtype=dtype)
            Rop = Rop1*Rop
    else:
        Rop = Restriction(np.prod(dims), iava, dims=dims,
                          dir=0, dtype=dtype)
        if ndims == 3 and iava1 is not None:
            dims1 = (len(iava), nrec[1], dimsd[2])
            Rop1 = Restriction(np.prod(dims1), iava1, dims=dims1,
                               dir=1, dtype=dtype)
            Rop = Rop1*Rop

    # create other operators for inversion
    if kind == 'spatial':
        prec = False
        dotcflag = 0
        if ndims == 3 and iava1 is not None:
            Regop = Laplacian(dims=dims, dirs=(0, 1), dtype=dtype)
        else:
            Regop = SecondDerivative(np.prod(dims),
                                     dims=(dims), dir=0,