How to use the sigpy.to_device function in sigpy

References:
        Martin Uecker, Peng Lai, Mark J. Murphy, Patrick Virtue, Michael Elad,
        John M. Pauly, Shreyas S. Vasanawala, and Michael Lustig
        ESPIRIT - An Eigenvalue Approach to Autocalibrating Parallel MRI:
        Where SENSE meets GRAPPA.
        Magnetic Resonance in Medicine, 71:990-1001 (2014)

    """
    img_ndim = ksp.ndim - 1
    num_coils = len(ksp)
    with sp.get_device(ksp):
        # Get calibration region
        calib_shape = [num_coils] + [calib_width] * img_ndim
        calib = sp.resize(ksp, calib_shape)
        calib = sp.to_device(calib, device)

    device = sp.Device(device)
    xp = device.xp
    with device:
        # Get calibration matrix
        kernel_shape = [num_coils] + [kernel_width] * img_ndim
        kernel_strides = [1] * (img_ndim + 1)
        mat = sp.array_to_blocks(calib, kernel_shape, kernel_strides)
        mat = mat.reshape([-1, sp.prod(kernel_shape)])

        # Perform SVD on calibration matrix
        _, S, VH = xp.linalg.svd(mat, full_matrices=False)
        VH = VH[S > thresh * S.max(), :]

        # Get kernels
        num_kernels = len(VH)

self.weights = sp.resize(
                    self.weights, ndim * [self.ksp_calib_width])

        else:
            self.img_shape = sp.estimate_shape(self.coord)
            calib_idx = np.amax(np.abs(self.coord), axis=-
                                1) < self.ksp_calib_width / 2

            self.coord = self.coord[calib_idx]
            self.y = self.y[:, calib_idx]

            if self.weights is not None:
                self.weights = self.weights[calib_idx]

        if self.weights is None:
            self.y = sp.to_device(self.y, self.device)
        else:
            self.y = sp.to_device(self.weights**0.5 * self.y, self.device)

        if self.coord is not None:
            self.coord = sp.to_device(self.coord, self.device)
        if self.weights is not None:
            self.weights = sp.to_device(self.weights, self.device)

        self.weights = _estimate_weights(self.y, self.weights, self.coord)

        if self.normalize:
            xp = self.device.xp
            with self.device:
                self.y = self.y / xp.linalg.norm(self.y)

where A is the Sense operator.

    Args:
        mps (array): sensitivity maps of shape [num_coils] + image shape.
        weights (array): k-space weights.
        coord (array): k-space coordinates of shape [...] + [ndim].
        lamda (float): regularization.

    Returns:
        array: k-space preconditioner of same shape as k-space.

    """
    dtype = mps.dtype

    if weights is not None:
        weights = sp.to_device(weights, device)

    device = sp.Device(device)
    xp = device.xp

    mps_shape = list(mps.shape)
    img_shape = mps_shape[1:]
    img2_shape = [i * 2 for i in img_shape]
    ndim = len(img_shape)

    scale = sp.prod(img2_shape)**1.5 / sp.prod(img_shape)
    with device:
        if coord is None:
            idx = (slice(None, None, 2), ) * ndim

            ones = xp.zeros(img2_shape, dtype=dtype)
            if weights is None:

calib_idx = np.amax(np.abs(self.coord), axis=-
                                1) < self.ksp_calib_width / 2

            self.coord = self.coord[calib_idx]
            self.y = self.y[:, calib_idx]

            if self.weights is not None:
                self.weights = self.weights[calib_idx]

        if self.weights is None:
            self.y = sp.to_device(self.y, self.device)
        else:
            self.y = sp.to_device(self.weights**0.5 * self.y, self.device)

        if self.coord is not None:
            self.coord = sp.to_device(self.coord, self.device)
        if self.weights is not None:
            self.weights = sp.to_device(self.weights, self.device)

        self.weights = _estimate_weights(self.y, self.weights, self.coord)

        if self.normalize:
            xp = self.device.xp
            with self.device:
                self.y = self.y / xp.linalg.norm(self.y)

def __init__(self, y, L, lamda=0.005,
                 mode='full', multi_channel=False, device=sp.cpu_device, **kwargs):
        self.y = sp.to_device(y, device)
        self.L = sp.to_device(L, device)
        self.lamda = lamda
        self.mode = mode
        self.multi_channel = multi_channel
        self.device = device

        self._get_params()
        self.A_R = sp.linop.ConvolveInput(self.R_shape, self.L,
                                          mode=self.mode,
                                          input_multi_channel=True,
                                          output_multi_channel=self.multi_channel)
        
        proxg_R = sp.prox.L1Reg(self.R_shape, lamda)
        super().__init__(self.A_R, self.y, proxg=proxg_R, **kwargs)

def __init__(self, y, mps, lamda,
                 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,
                        comm=comm, coil_batch_size=coil_batch_size)

        G = sp.linop.FiniteDifference(A.ishape)
        proxg = sp.prox.L1Reg(G.oshape, lamda)

        def g(x):
            device = sp.get_device(x)
            xp = device.xp
            with device:
                return lamda * xp.sum(xp.abs(x)).item()

        if comm is not None:

idx = []
        for i in range(self.ndim):
            if i == self.z:
                idx.append(slice(None, None, self.flips[i]))
            else:
                idx.append(self.slices[i])

        if idx:
            datav = to_device(self.data[idx], cpu_device)
        else:
            datav = to_device(self.data, cpu_device)

        # if self.z is not None:
        #     datav_dims = [self.z] + datav_dims
        coordv = to_device(self.coord, cpu_device)

        if self.mode == 'm':
            datav = np.abs(datav)
        elif self.mode == 'p':
            datav = np.angle(datav)
        elif self.mode == 'r':
            datav = np.real(datav)
        elif self.mode == 'i':
            datav = np.imag(datav)
        elif self.mode == 'l':
            eps = 1e-31
            datav = np.log(np.abs(datav) + eps)

        datav = datav.ravel()

        if self.axsc is None:

idx = []
        for i in range(self.ndim):
            if i == self.z:
                idx.append(slice(None, None, self.flips[i]))
            else:
                idx.append(self.slices[i])

        idx = tuple(idx)
        if idx:
            datav = sp.to_device(self.data[idx])
        else:
            datav = sp.to_device(self.data)

        # if self.z is not None:
        #     datav_dims = [self.z] + datav_dims
        coordv = sp.to_device(self.coord)

        if self.mode == 'm':
            datav = np.abs(datav)
        elif self.mode == 'p':
            datav = np.angle(datav)
        elif self.mode == 'r':
            datav = np.real(datav)
        elif self.mode == 'i':
            datav = np.imag(datav)
        elif self.mode == 'l':
            eps = 1e-31
            datav = np.log(np.abs(datav) + eps)

        datav = datav.ravel()
        if self.vmin is None:
            if datav.min() == datav.max():

self.coord = self.coord[calib_idx]
            self.y = self.y[:, calib_idx]

            if self.weights is not None:
                self.weights = self.weights[calib_idx]

        if self.weights is None:
            self.y = sp.to_device(self.y, self.device)
        else:
            self.y = sp.to_device(self.weights**0.5 * self.y, self.device)

        if self.coord is not None:
            self.coord = sp.to_device(self.coord, self.device)
        if self.weights is not None:
            self.weights = sp.to_device(self.weights, self.device)

        self.weights = _estimate_weights(self.y, self.weights, self.coord)

        if self.normalize:
            xp = self.device.xp
            with self.device:
                self.y = self.y / xp.linalg.norm(self.y)

"""Apply hard pulse rotation to input magnetization.

    Args:
        input (array): magnetization array.
        b1 (complex float): complex B1 value in radian.

    Returns:
        array: magnetization array after hard pulse rotation,
            in representation consistent with input.

    """
    p = to_density_matrix(input)

    device = sp.get_device(p)
    with device:
        b1 = sp.to_device(b1, device)
        p = _exp(b1) @ p @ _exp(-b1)

    if is_bloch_vector(input):
        return to_bloch_vector(p)
    else:
        return p