How to use the pylops.MatrixMult function in pylops

% linearization)

        G = [np.hstack((np.diag(G1[itheta] * np.ones(nt0)),
                        np.diag(G2[itheta] * np.ones(nt0)),
                        np.diag(G3[itheta] * np.ones(nt0))))
             for itheta in range(ntheta)]
        G = np.vstack(G).reshape(ntheta * nt0, 3 * nt0)

        # Create wavelet operator
        C = convmtx(wav, nt0)[:, len(wav) // 2:-len(wav) // 2 + 1]
        C = [C] * ntheta
        C = block_diag(*C)

        # Combine operators
        M = np.dot(C, np.dot(G, D))
        return MatrixMult(M, dims=spatdims)

    else:
        # Create wavelet operator
        Cop = Convolve1D(np.prod(np.array(dims)), h=wav,
                         offset=len(wav)//2, dir=0, dims=dims)

        # create AVO operator
        AVOop = AVOLinearModelling(theta, vsvp, spatdims=spatdims,
                                   linearization=linearization)

        # Create derivative operator
        dimsm = list(dims)
        dimsm[1] = AVOop.npars
        Dop = FirstDerivative(np.prod(np.array(dimsm)), dims=dimsm,
                              dir=0, sampling=1.)
        return Cop*AVOop*Dop

minv = lstsq(PPop.A, datar.reshape(nt0*ntheta,
                                                   nspatprod).squeeze(),
                             **kwargs_solver)[0]
            elif epsI is None and simultaneous:
                # solve unregularized equations simultaneously
                minv = lsqr(PPop, datar, **kwargs_solver)[0]
            elif epsI is not None:
                # create regularized normal equations
                PP = np.dot(PPop.A.T, PPop.A) + epsI * np.eye(nt0*nm)
                datarn = np.dot(PPop.A.T, datar.reshape(nt0*ntheta, nspatprod))
                if not simultaneous:
                    # solve regularized normal eqs. trace-by-trace
                    minv = lstsq(PP, datarn, **kwargs_solver)[0]
                else:
                    # solve regularized normal equations simultaneously
                    PPop_reg = MatrixMult(PP, dims=nspatprod)
                    minv = lsqr(PPop_reg, datarn.flatten(), **kwargs_solver)[0]
            #else:
            #    # create regularized normal eqs. and solve them simultaneously
            #    PP = np.dot(PPop.A.T, PPop.A) + epsI * np.eye(nt0*nm)
            #    datarn = PPop.A.T * datar.reshape(nt0*ntheta, nspatprod)
            #    PPop_reg = MatrixMult(PP, dims=ntheta*nspatprod)
            #    minv = lstsq(PPop_reg, datarn.flatten(), **kwargs_solver)[0]
        else:
            # solve unregularized normal equations simultaneously with lop
            minv = lsqr(PPop, datar, **kwargs_solver)[0]
    else:
        # Create Thicknov regularization
        if epsI is not None:
            if isinstance(epsI, (list, tuple)):
                if len(epsI) != nm:
                    raise ValueError('epsI must be a scalar or a list of'

def _PoststackLinearModelling(wav, nt0, spatdims=None, explicit=False,
                              sparse=False, _MatrixMult=MatrixMult,
                              _Convolve1D=Convolve1D,
                              _FirstDerivative=FirstDerivative,
                              args_MatrixMult={}, args_Convolve1D={},
                              args_FirstDerivative={}):
    """Post-stack linearized seismic modelling operator.

    Used to be able to provide operators from different libraries to
    PoststackLinearModelling. It operates in the same way as public method
    (PoststackLinearModelling) but has additional input parameters allowing
    passing a different operator and additional arguments to be passed to such
    operator.

    """
    if len(wav.shape) == 2 and wav.shape[0] != nt0:
        raise ValueError('Provide 1d wavelet or 2d wavelet composed of nt0 '
                         'wavelets')

epsI = np.sqrt(1e-4)

xne = \
    pylops.optimization.leastsquares.NormalEquationsInversion(Rop, [D2op], y,
                                                              epsI=epsI,
                                                              epsRs=[epsR],
                                                              returninfo=False,
                                                              **dict(maxiter=50))

###############################################################################
# Note that in case we have access to a fast implementation for the chain of
# forward and adjoint for the regularization operator
# (i.e., :math:`\nabla^T\nabla`), we can modify our call to
# :py:func:`pylops.optimization.leastsquares.NormalEquationsInversion` as
# follows:
ND2op = pylops.MatrixMult((D2op.H * D2op).tosparse()) # mimic fast D^T D

xne1 = \
    pylops.optimization.leastsquares.NormalEquationsInversion(Rop, [], y,
                                                              NRegs=[ND2op],
                                                              epsI=epsI,
                                                              epsNRs=[epsR],
                                                              returninfo=False,
                                                              **dict(maxiter=50))

###############################################################################
# We can do the same while using
# :py:func:`pylops.optimization.leastsquares.RegularizedInversion`
# which solves the following augmented problem
#
#   .. math::
#       \begin{bmatrix}

minv = lstsq(PPop.A, datar.reshape(nt0, nspatprod).squeeze(),
                             **kwargs_solver)[0]
            elif epsI is None and simultaneous:
                # solve unregularized equations simultaneously
                minv = lsqr(PPop, datar, **kwargs_solver)[0]
            elif epsI is not None:
                # create regularized normal equations
                PP = np.dot(PPop.A.T, PPop.A) + epsI * np.eye(nt0)
                datarn = np.dot(PPop.A.T, datar.reshape(nt0, nspatprod))
                if not simultaneous:
                    # solve regularized normal eqs. trace-by-trace
                    minv = lstsq(PP, datarn,
                                 **kwargs_solver)[0]
                else:
                    # solve regularized normal equations simultaneously
                    PPop_reg = MatrixMult(PP, dims=nspatprod)
                    minv = lsqr(PPop_reg, datar.flatten(), **kwargs_solver)[0]
            else:
                # create regularized normal eqs. and solve them simultaneously
                PP = np.dot(PPop.A.T, PPop.A) + epsI * np.eye(nt0)
                datarn = PPop.A.T * datar.reshape(nt0, nspatprod)
                PPop_reg = MatrixMult(PP, dims=nspatprod)
                minv = lstsq(PPop_reg.A, datarn.flatten(), **kwargs_solver)[0]
        else:
            # solve unregularized normal equations simultaneously with lop
            minv = lsqr(PPop, datar, **kwargs_solver)[0]
    else:
        if epsRL1 is None:
            # L2 inversion with spatial regularization
            if dims == 1:
                Regop = SecondDerivative(nt0, dtype=PPop.dtype)
            elif dims == 2:

G1, G2, G3 = fatti(theta, vsvp, n=nt0)
    else:
        logging.error('%s not an available linearization...',
                      linearization)
        raise NotImplementedError('%s not an available linearization...'
                                  % linearization)

    G = [np.hstack((np.diag(G1[itheta] * np.ones(nt0)),
                    np.diag(G2[itheta] * np.ones(nt0)),
                    np.diag(G3[itheta] * np.ones(nt0))))
         for itheta in range(ntheta)]
    G = np.vstack(G).reshape(ntheta * nt0, 3 * nt0)

    # Create infinite-reflectivity data
    M = np.dot(G, np.dot(D, m.T.flatten())).reshape(ntheta, nt0)
    Mconv = VStack([MatrixMult(convmtx(M[itheta], nwav)[wavc:-nwav+wavc+1])
                    for itheta in range(ntheta)])

    return Mconv

#           6  & 7 \\
#       \end{bmatrix} =
#       \begin{bmatrix}
#            0 &  5 &  0 & 10 \\
#            6 &  7 & 12 & 14 \\
#            0 & 15 &  0 & 20 \\
#           18 & 21 & 24 & 28 \\
#       \end{bmatrix}
A = np.array([[1, 2], [3, 4]])
B = np.array([[0, 5], [6, 7]])
AB = np.kron(A, B)

n1, m1 = A.shape
n2, m2 = B.shape

Aop = pylops.MatrixMult(A)
Bop = pylops.MatrixMult(B)

ABop = pylops.Kronecker(Aop, Bop)
x = np.ones(m1*m2)

y = AB.dot(x)
yop = ABop*x
xinv = ABop / yop

print('AB = \n', AB)

print('x = ', x)
print('y = ', y)
print('yop = ', yop)
print('xinv = ', x)

#               \end{bmatrix}
#       \begin{bmatrix}
#               \mathbf{x_1}  \\
#               \mathbf{x_2}  \\
#               \mathbf{x_3}
#       \end{bmatrix}
#
# and apply it using the same implementation of the
# :py:class:`pylops.MatrixMult` operator by simply telling the operator how we
# want the model to be organized through the ``dims`` input parameter.
A = np.array([[1., 2.], [4., 5.]])
x = np.array([[1., 1.],
              [2., 2.],
              [3., 3.]])

Aop = pylops.MatrixMult(A, dims=(3,), dtype='float64')
y = Aop*x.flatten()

xest, istop, itn, r1norm, r2norm = \
        lsqr(Aop, y, damp=1e-10, iter_lim=10, show=0)[0:5]
xest = xest.reshape(3, 2)

print('A= %s' % A)
print('x= %s' % x)
print('y= %s' % y)
print('lsqr solution xest= %s' % xest)