How to use the sigpy.mri.linop.Sense function in sigpy

def test_stspa_3d_nonexplicit(self):
        nz = 3
        target, sens = self.problem_3d(3, nz)
        dim = target.shape[0]

        g, k1, t, s = rf.spiral_arch(0.24, dim, 4e-6, 200, 0.035)
        k1 = k1 / dim

        k1 = rf.stack_of(k1, nz, 0.1)
        A = sp.mri.linop.Sense(sens, k1, weights=None, tseg=None,
                               ishape=target.shape).H

        pulses = sp.mri.rf.stspa(target, sens, st=None,
                                 coord=k1,
                                 dt=4e-6, max_iter=30, alpha=10, tol=1E-3,
                                 phase_update_interval=200, explicit=False)

        npt.assert_array_almost_equal(A*pulses, target, 1E-3)

def test_stspa_spiral(self):

        target, sens = self.problem_2d(8)

        fov = 0.55
        gts = 6.4e-6
        gslew = 190
        gamp = 40
        R = 1
        dx = 0.025  # in m
        # construct a trajectory
        g, k, t, s = rf.spiral_arch(fov / R, dx, gts, gslew, gamp)

        A = linop.Sense(sens, coord=k, ishape=target.shape).H

        pulses = rf.stspa(target, sens, k, dt=4e-6, alpha=1,
                          b0=None, st=None,
                          explicit=False, max_iter=100, tol=1E-4)

        npt.assert_array_almost_equal(A*pulses, target, 1E-3)

def test_noncart_sense_model(self):
        img_shape = [16, 16]
        mps_shape = [8, 16, 16]

        img = sp.randn(img_shape, dtype=np.complex)
        mps = sp.randn(mps_shape, dtype=np.complex)

        y, x = np.mgrid[:16, :16]
        coord = np.stack([np.ravel(y - 8), np.ravel(x - 8)], axis=1)
        coord = coord.astype(np.float)

        A = linop.Sense(mps, coord=coord)
        check_linop_adjoint(A, dtype=np.complex)
        npt.assert_allclose(sp.fft(img * mps, axes=[-1, -2]).ravel(),
                            (A * img).ravel(), atol=0.1, rtol=0.1)

def test_sense_model_batch(self):
        img_shape = [16, 16]
        mps_shape = [8, 16, 16]

        img = sp.randn(img_shape, dtype=np.complex)
        mps = sp.randn(mps_shape, dtype=np.complex)

        mask = np.zeros(img_shape, dtype=np.complex)
        mask[::2, ::2] = 1.0

        A = linop.Sense(mps, coil_batch_size=1)
        check_linop_adjoint(A, dtype=np.complex)
        npt.assert_allclose(sp.fft(img * mps, axes=[-1, -2]),
                            A * img)

def test_sense_model(self):
        img_shape = [16, 16]
        mps_shape = [8, 16, 16]

        img = sp.randn(img_shape, dtype=np.complex)
        mps = sp.randn(mps_shape, dtype=np.complex)

        mask = np.zeros(img_shape)
        mask[::2, ::2] = 1.0

        A = linop.Sense(mps)

        check_linop_adjoint(A, dtype=np.complex)

        npt.assert_allclose(sp.fft(img * mps, axes=[-1, -2]),
                            A * img)

def test_stspa_radial(self):
        # makes dim*dim*2 trajectory
        traj = sp.mri.radial((self.sens.shape[1], self.sens.shape[1], 2),
                             self.img_shape, golden=True, dtype=np.float)
        # reshape to be Nt*2 trajectory
        traj = np.reshape(traj, [traj.shape[0] * traj.shape[1], 2])

        A = linop.Sense(self.sens, coord=traj,
                        weights=None, ishape=self.target.shape).H

        pulses = rf.stspa(self.target, self.sens, traj, self.dt, alpha=1,
                          B0=None, pinst=float('inf'), pavg=float('inf'),
                          explicit=False, max_iter=100, tol=1E-4)

        npt.assert_array_almost_equal(A * pulses, self.target, 1E-3)

def test_sense_model_with_comm(self):
            img_shape = [16, 16]
            mps_shape = [8, 16, 16]
            comm = sp.Communicator()

            img = sp.randn(img_shape, dtype=np.complex)
            mps = sp.randn(mps_shape, dtype=np.complex)
            comm.allreduce(img)
            comm.allreduce(mps)
            ksp = sp.fft(img * mps, axes=[-1, -2])

            A = linop.Sense(mps[comm.rank::comm.size], comm=comm)

            npt.assert_allclose(A.H(ksp[comm.rank::comm.size]), np.sum(
                sp.ifft(ksp, axes=[-1, -2]) * mps.conjugate(), 0))

def test_noncart_sense_model_batch(self):
        img_shape = [16, 16]
        mps_shape = [8, 16, 16]

        img = sp.randn(img_shape, dtype=np.complex)
        mps = sp.randn(mps_shape, dtype=np.complex)

        y, x = np.mgrid[:16, :16]
        coord = np.stack([np.ravel(y - 8), np.ravel(x - 8)], axis=1)
        coord = coord.astype(np.float)

        A = linop.Sense(mps, coord=coord, coil_batch_size=1)
        check_linop_adjoint(A, dtype=np.complex)
        npt.assert_allclose(sp.fft(img * mps, axes=[-1, -2]).ravel(),
                            (A * img).ravel(), atol=0.1, rtol=0.1)

def __init__(self, y, mps, lamda=0, weights=None,
                 coord=None, device=sp.cpu_device, coil_batch_size=None,
                 comm=None, show_pbar=True, **kwargs):
        weights = _estimate_weights(y, weights, coord)
        if weights is not None:
            y = sp.to_device(y * weights**0.5, device=device)
        else:
            y = sp.to_device(y, device=device)

        A = linop.Sense(mps, coord=coord, weights=weights,
                        coil_batch_size=coil_batch_size, comm=comm)

        if comm is not None:
            show_pbar = show_pbar and comm.rank == 0

        super().__init__(A, y, lamda=lamda, show_pbar=show_pbar, **kwargs)

Args:
        target (array): desired magnetization profile.
        sens (array): sensitivity maps.
        mask (array): kspace sampling pattern
        coord (array): coordinates for noncartesian trajectories
        max_iter (int): max number of iterations
        tol (float): allowable error

    References:
        Grissom, W., Yip, C., Zhang, Z., Stenger, V. A., Fessler, J. A. & Noll, D. C.
        (2006).
        Spatial Domain Method for the Design of RF Pulses in Multicoil Parallel Excitation.
        Magnetic resonance in medicine, 56, 620-629.
    """

    A = linop.Sense(sens, coord, weights=mask, ishape=target.shape).H

    pulses = np.zeros(sens.shape, np.complex)

    alg_method = sp.alg.ConjugateGradient(A.H*A,A.H*target,pulses,P=None, max_iter=max_iter,tol=tol)
    while not alg_method.done():
        alg_method.update()

    return pulses