jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(
        self,
        mean: Tensor,
        cov: Tensor,
        seed: Optional[int] = None,
        inv_transform: bool = False,
    ) -> None:
        r"""Engine for qMC sampling from a multivariate Normal `N(\mu, \Sigma)`.

        Args:
            mean: The mean vector.
            cov: The covariance matrix.
            seed: The seed with which to seed the random number generator of the
                underlying SobolEngine.
            inv_transform: If True, use inverse transform instead of Box-Muller.
        """
        # validate inputs
        if not cov.shape[0] == cov.shape[1]:
            raise ValueError("Covariance matrix is not square.")
        if not mean.shape[0] == cov.shape[0]:
            raise ValueError("Dimension mismatch between mean and covariance.")
        if not torch.allclose(cov, cov.transpose(-1, -2)):
            raise ValueError("Covariance matrix is not symmetric.")
        self._mean = mean
        self._normal_engine = NormalQMCEngine(
            d=mean.shape[0], seed=seed, inv_transform=inv_transform
        )
        # compute Cholesky decomp; if it fails, do the eigendecomposition
        try:
            self._corr_matrix = torch.linalg.cholesky(cov).transpose(-1, -2)
        except RuntimeError:
            eigval, eigvec = torch.linalg.eigh(cov)
            tol = 1e-8 if eigval.dtype == torch.double else 1e-6
            if torch.any(eigval < -tol):
                raise ValueError("Covariance matrix not PSD.")
            eigval_root = eigval.clamp_min(0.0).sqrt()
            self._corr_matrix = (eigvec * eigval_root).transpose(-1, -2)
//  [python]
def norm_sender_normalized(pyr, num_scale=2, num_bands=6, blksz=3, eps=1e-12):
    r"""Normalize pyramid with local spatial neighbor and band neighbor"""
    border = blksz // 2
    guardband = 16
    subbands = []
    for si in range(num_scale):
        for bi in range(num_bands):
            idx = si * num_bands + bi
            current_band = pyr[idx]

            N = blksz ** 2

            # 3x3 window pixels
            tmp = F.unfold(current_band.unsqueeze(1), 3, stride=1)
            tmp = tmp.transpose(1, 2)
            b, hw = tmp.shape[:2]
            # parent pixels
            parent_idx = idx + num_bands
            if parent_idx < len(pyr):
                tmp_parent = pyr[parent_idx]
                tmp_parent = imresize(tmp_parent, sizes=current_band.shape[-2:])
                tmp_parent = tmp_parent[:, border:-border, border:-border].reshape(
                    b, hw, 1
                )
                tmp = torch.cat((tmp, tmp_parent), dim=-1)
                N += 1
            # neighbor band pixels
            for ni in range(num_bands):
                if ni != bi:
                    ni_idx = si * num_bands + ni
                    tmp_nei = pyr[ni_idx]
                    tmp_nei = tmp_nei[:, border:-border, border:-border].reshape(
                        b, hw, 1
                    )
                    tmp = torch.cat((tmp, tmp_nei), dim=-1)
            C_x = tmp.transpose(1, 2) @ tmp / tmp.shape[1]
            # correct possible negative eigenvalue
            L, Q = torch.symeig(C_x, eigenvectors=True)
            L_pos = L * (L > 0)
            L_pos_sum = L_pos.sum(dim=1, keepdim=True)
            L = (
                L_pos
                * L.sum(dim=1, keepdim=True)
                / (L_pos_sum + (L_pos_sum == 0).float())
            )
            C_x = Q @ torch.diag_embed(L) @ Q.transpose(1, 2)

            o_c = current_band[:, border:-border, border:-border]
            b, h, w = o_c.shape
            o_c = o_c.reshape(b, hw)
            o_c = o_c - o_c.mean(dim=1, keepdim=True)

            if hasattr(torch.linalg, "lstsq"):
                tmp_y = (
                    torch.linalg.lstsq(
                        C_x.transpose(1, 2), tmp.transpose(1, 2)
                    ).solution.transpose(1, 2)
                    * tmp
                    / N
                )
            else:
                warn(
                    "For numerical stability, we use torch.linal.lstsq to calculate matrix inverse for PyTorch > 1.9.0. The results might be slightly different if you use older version of PyTorch."
                )
                tmp_y = (tmp @ torch.linalg.pinv(C_x)) * tmp / N

            z = tmp_y.sum(dim=2).sqrt()
            mask = z != 0
            g_c = o_c * mask / (z * mask + eps)
            g_c = g_c.reshape(b, h, w)

            gb = int(guardband / (2 ** (si)))
            g_c = g_c[:, gb:-gb, gb:-gb]
            g_c = g_c - g_c.mean(dim=(1, 2), keepdim=True)
            subbands.append(g_c)

    return subbands
//  [python]
def norm_sender_normalized(pyr, num_scale=2, num_bands=6, blksz=3, eps=1e-12):
    r"""Normalize pyramid with local spatial neighbor and band neighbor"""
    border = blksz // 2
    guardband = 16
    subbands = []
    for si in range(num_scale):
        for bi in range(num_bands):
            idx = si * num_bands + bi
            current_band = pyr[idx]

            N = blksz ** 2

            # 3x3 window pixels
            tmp = F.unfold(current_band.unsqueeze(1), 3, stride=1)
            tmp = tmp.transpose(1, 2)
            b, hw = tmp.shape[:2]
            # parent pixels
            parent_idx = idx + num_bands
            if parent_idx < len(pyr):
                tmp_parent = pyr[parent_idx]
                tmp_parent = imresize(tmp_parent, sizes=current_band.shape[-2:])
                tmp_parent = tmp_parent[:, border:-border, border:-border].reshape(
                    b, hw, 1
                )
                tmp = torch.cat((tmp, tmp_parent), dim=-1)
                N += 1
            # neighbor band pixels
            for ni in range(num_bands):
                if ni != bi:
                    ni_idx = si * num_bands + ni
                    tmp_nei = pyr[ni_idx]
                    tmp_nei = tmp_nei[:, border:-border, border:-border].reshape(
                        b, hw, 1
                    )
                    tmp = torch.cat((tmp, tmp_nei), dim=-1)
            C_x = tmp.transpose(1, 2) @ tmp / tmp.shape[1]
            # correct possible negative eigenvalue
            L, Q = torch.linalg.eigh(C_x)
            L_pos = L * (L > 0)
            L_pos_sum = L_pos.sum(dim=1, keepdim=True)
            L = (
                L_pos
                * L.sum(dim=1, keepdim=True)
                / (L_pos_sum + (L_pos_sum == 0).float())
            )
            C_x = Q @ torch.diag_embed(L) @ Q.transpose(1, 2)

            o_c = current_band[:, border:-border, border:-border]
            b, h, w = o_c.shape
            o_c = o_c.reshape(b, hw)
            o_c = o_c - o_c.mean(dim=1, keepdim=True)

            if hasattr(torch.linalg, "lstsq"):
                tmp_y = (
                    torch.linalg.lstsq(
                        C_x.transpose(1, 2), tmp.transpose(1, 2)
                    ).solution.transpose(1, 2)
                    * tmp
                    / N
                )
            else:
                warn(
                    "For numerical stability, we use torch.linal.lstsq to calculate matrix inverse for PyTorch > 1.9.0. The results might be slightly different if you use older version of PyTorch."
                )
                tmp_y = (tmp @ torch.linalg.pinv(C_x)) * tmp / N

            z = tmp_y.sum(dim=2).sqrt()
            mask = z != 0
            g_c = o_c * mask / (z * mask + eps)
            g_c = g_c.reshape(b, h, w)

            gb = int(guardband / (2 ** (si)))
            g_c = g_c[:, gb:-gb, gb:-gb]
            g_c = g_c - g_c.mean(dim=(1, 2), keepdim=True)
            subbands.append(g_c)

    return subbands
//  [python]
def __init__(self, n_qubits, 
                       P_idxs, 
                       M, 
                       rho_init=None, 
                       map_method='chol_h', 
                       P_proj=1.5):
        """
        Args:
            n_qubits (int): The number of qubits.
            P_idxs (tensor): Index of the POVM used for measurement, Not all measurements 
                are necessarily used.
            M (tensor): The POVM, size (K, 2, 2).
            rho_init (tensor): If None, initialize the parameters randomly, and vice versa with rho.
            map_method (str): State-mapping method, include ['chol', 'chol_h', 'proj_F', 'proj_S', 'proj_A'].
            P_proj (float): P order.
        """
        super(SNN_nn, self).__init__()

        self.N = n_qubits
        self.P_idxs = P_idxs
        self.M = M
        self.device = M.device
        self.map_method = map_method
        self.P_proj = P_proj

        params = torch.randn(4**n_qubits, requires_grad=True)

        if rho_init is not None:  # chol_h
            eigenvalues, eigenvecs = torch.eig(rho_init)
            eigenvalues_sqrt = torch.sqrt(torch.abs(eigenvalues[:,0]))
            A = eigenvecs * eigenvalues_sqrt[:,None]
            H = torch.matmul(A, eigenvecs.T.conj())
            L = torch.tril(H)
            Y = L.real + L.imag.T
            params = Y

        self.params = nn.Parameter(params)
//  [python]
def __init__(self, n_qubits, 
                       P_idxs, 
                       M, 
                       rho_init=None, 
                       map_method='chol_h', 
                       P_proj=1.5):
        """
        Args:
            n_qubits (int): The number of qubits.
            P_idxs (tensor): Index of the POVM used for measurement, Not all measurements 
                are necessarily used.
            M (tensor): The POVM, size (K, 2, 2).
            rho_init (tensor): If None, initialize the parameters randomly, and vice versa with rho.
            map_method (str): State-mapping method, include ['chol', 'chol_h', 'proj_F', 'proj_S', 'proj_A'].
            P_proj (float): P order.
        """
        super(SNN_nn, self).__init__()

        self.N = n_qubits
        self.P_idxs = P_idxs
        self.M = M
        self.device = M.device
        self.map_method = map_method
        self.P_proj = P_proj

        params = torch.randn(4**n_qubits, requires_grad=True)

        if rho_init is not None:  # chol_h
            eigenvalues, eigenvecs = torch.linalg.eigh(rho_init)
            eigenvalues_sqrt = torch.sqrt(torch.abs(eigenvalues))
            A = eigenvecs * eigenvalues_sqrt
            H = torch.matmul(A, eigenvecs.T.conj())
            L = torch.tril(H)
            Y = L.real + L.imag.T
            params = Y

        self.params = nn.Parameter(params)
//  [python]
def get_code_stats(code, device):
	code = code.to(torch.float64).to(device)
	n = code.size(0)
	m = code.size(2)
	s = torch.zeros(m, dtype = torch.float64, device = device)
	sxy = torch.zeros(m, m, dtype = torch.float64, device = device)
	sv = torch.zeros(m, dtype = torch.float64, device = device)

	for code_batch in code.split(1024):
		mean, std = code_batch.unbind(1)
		s.add_(mean.sum(0))
		sxy.add_(mean.t() @ mean)
		sv.add_(std.pow(2).sum(0))

	ex = s.cpu() / n
	cov = sxy.cpu() / n - ex.unsqueeze(0) * ex.unsqueeze(1)
	ev = sv.cpu() / n

	full_var = torch.diag(cov) + ev
	full_std = full_var.sqrt()
	var_ratio = torch.diag(cov) / full_var

	normalized_cov = cov / (full_std.unsqueeze(0) * full_std.unsqueeze(1))

	eigval, eigvec = torch.symeig(normalized_cov, eigenvectors=True)

	stats = {
		'mean': ex.to(torch.float32),
		'std': full_std.to(torch.float32),
		'var_ratio': var_ratio.to(torch.float32),
		'eigval': eigval.to(torch.float32),
		'eigvec': eigvec.to(torch.float32)
	}
	return stats
//  [python]
def get_code_stats(code, device):
	code = code.to(torch.float64).to(device)
	n = code.size(0)
	m = code.size(2)
	s = torch.zeros(m, dtype = torch.float64, device = device)
	sxy = torch.zeros(m, m, dtype = torch.float64, device = device)
	sv = torch.zeros(m, dtype = torch.float64, device = device)

	for code_batch in code.split(1024):
		mean, std = code_batch.unbind(1)
		s.add_(mean.sum(0))
		sxy.add_(mean.t() @ mean)
		sv.add_(std.pow(2).sum(0))

	ex = s.cpu() / n
	cov = sxy.cpu() / n - ex.unsqueeze(0) * ex.unsqueeze(1)
	ev = sv.cpu() / n

	full_var = torch.diag(cov) + ev
	full_std = full_var.sqrt()
	var_ratio = torch.diag(cov) / full_var

	normalized_cov = cov / (full_std.unsqueeze(0) * full_std.unsqueeze(1))

	eigval, eigvec = torch.linalg.eigh(normalized_cov)

	stats = {
		'mean': ex.to(torch.float32),
		'std': full_std.to(torch.float32),
		'var_ratio': var_ratio.to(torch.float32),
		'eigval': eigval.to(torch.float32),
		'eigvec': eigvec.to(torch.float32)
	}
	return stats
//  [python]
def get_cost_pytorch(manifold, C, y, k=3):
    """Create autograd cost function for a set of covariance matrices C

    Parameters
    ----------
    manifold; Manifold
        Target manifold for optimization
    C: array, shape (2 * n_samples, dim, dim)
        Covariance matrices to decompose in subspaces
    y: array, shape (2 * n_samples)
        Class label (-1 or +1) associated with C
    k: int
        nearest neighbors consider for optimization

    Returns
    -------
    cost: function
        Cost function for pymanopt Problem
    """
    n_samples, _, _ = C.shape
    distances, neighbors = nearest_neighbors(C)
    nn_ = torch.from_numpy(kneighbors(k, neighbors))
    tri = torch.tril(torch.ones(n_samples, n_samples), diagonal=-1)
    line_idx, col_idx = torch.nonzero(tri, as_tuple=True)
    C_, y_ = torch.from_numpy(C), torch.from_numpy(y)

    @pymanopt.function.pytorch(manifold)
    def cost_pytorch(U):
        cost = 0.0
        for i, j in zip(line_idx, col_idx):
            if nn_[i, j] != 0:
                ei, ev = torch.symeig(
                    torch.matmul(torch.matmul(torch.transpose(U, 1, 0), C_[i]), U),
                    eigenvectors=True,
                )
                W = torch.matmul(
                    torch.matmul(ev, torch.diag(1.0 / torch.sqrt(ei))),
                    torch.transpose(ev, 1, 0),
                )
                C = torch.matmul(
                    torch.matmul(
                        torch.matmul(torch.matmul(W, torch.transpose(U, 1, 0)), C_[j]),
                        U,
                    ),
                    W,
                )
                ei = torch.symeig(C, eigenvectors=False)
                sign = y_[i] * y_[j]
                cost += sign * torch.sum(torch.log(ei) ** 2)
        return cost

    return cost_pytorch
//  [python]
def get_cost_pytorch(manifold, C, y, k=3):
    """Create autograd cost function for a set of covariance matrices C

    Parameters
    ----------
    manifold; Manifold
        Target manifold for optimization
    C: array, shape (2 * n_samples, dim, dim)
        Covariance matrices to decompose in subspaces
    y: array, shape (2 * n_samples)
        Class label (-1 or +1) associated with C
    k: int
        nearest neighbors consider for optimization

    Returns
    -------
    cost: function
        Cost function for pymanopt Problem
    """
    n_samples, _, _ = C.shape
    distances, neighbors = nearest_neighbors(C)
    nn_ = torch.from_numpy(kneighbors(k, neighbors))
    tri = torch.tril(torch.ones(n_samples, n_samples), diagonal=-1)
    line_idx, col_idx = torch.nonzero(tri, as_tuple=True)
    C_, y_ = torch.from_numpy(C), torch.from_numpy(y)

    @pymanopt.function.pytorch(manifold)
    def cost_pytorch(U):
        cost = 0.0
        for i, j in zip(line_idx, col_idx):
            if nn_[i, j] != 0:
                ei, ev = torch.linalg.eigh(
                    torch.matmul(torch.matmul(torch.transpose(U, 1, 0), C_[i]), U),
                    UPLO="U",
                )
                W = torch.matmul(
                    torch.matmul(ev, torch.diag(1.0 / torch.sqrt(ei))),
                    torch.transpose(ev, 1, 0),
                )
                C = torch.matmul(
                    torch.matmul(
                        torch.matmul(torch.matmul(W, torch.transpose(U, 1, 0)), C_[j]),
                        U,
                    ),
                    W,
                )
                ei = torch.linalg.eigvalsh(C, UPLO="U")
                sign = y_[i] * y_[j]
                cost += sign * torch.sum(torch.log(ei) ** 2)
        return cost

    return cost_pytorch
//  [python]
def match_histogram(target_tensor, source_tensor, mode="avg"):
    if mode == "False":
        return target_tensor

    backup = target_tensor.clone()
    try:
        if mode == "avg":
            elementwise = True
            random_frame = False
        else:
            elementwise = False
            random_frame = True

        if not isinstance(source_tensor, list):
            source_tensor = [source_tensor]

        output_tensor = torch.zeros_like(target_tensor)
        for source in source_tensor:
            target = target_tensor.permute(0, 3, 2, 1)  # Function expects b,w,h,c
            source = source.permute(0, 3, 2, 1)  # Function expects b,w,h,c
            if elementwise:
                source = source.mean(0).unsqueeze(0)
            if random_frame:
                source = source[np.random.randint(0, source.shape[0])].unsqueeze(0)

            matched_tensor = torch.zeros_like(target)
            for idx in range(target.shape[0] if elementwise else 1):
                frame = target[idx].unsqueeze(0) if elementwise else target
                _, t, Ct = get_histogram(frame + 1e-3 * torch.randn(size=frame.shape, device=frame.device))
                mu_s, _, Cs = get_histogram(
                    source.to(frame) + 1e-3 * torch.randn(size=source.shape, device=frame.device)
                )

                # PCA
                eva_t, eve_t = torch.symeig(Ct, eigenvectors=True)
                Et = torch.sqrt(torch.diagflat(eva_t))
                Et[Et != Et] = 0  # Convert nan to 0
                Qt = torch.mm(torch.mm(eve_t, Et), eve_t.T)

                eva_s, eve_s = torch.symeig(Cs, eigenvectors=True)
                Es = torch.sqrt(torch.diagflat(eva_s))
                Es[Es != Es] = 0  # Convert nan to 0
                Qs = torch.mm(torch.mm(eve_s, Es), eve_s.T)

                ts = torch.mm(torch.mm(Qs, torch.inverse(Qt)), t)

                match = ts.reshape(*frame.permute(0, 3, 1, 2).shape).permute(0, 2, 3, 1)
                match += mu_s

                if elementwise:
                    matched_tensor[idx] = match
                else:
                    matched_tensor = match
            output_tensor += matched_tensor.permute(0, 3, 2, 1) / len(source_tensor)
    except RuntimeError:
        traceback.print_exc()
        print("Skipping histogram matching...")
        output_tensor = backup
    return output_tensor.clamp(
        min([s.min().item() for s in source_tensor]), max([s.max().item() for s in source_tensor])
    )
//  [python]
def match_histogram(target_tensor, source_tensor, mode="avg"):
    if mode == "False":
        return target_tensor

    backup = target_tensor.clone()
    try:
        if mode == "avg":
            elementwise = True
            random_frame = False
        else:
            elementwise = False
            random_frame = True

        if not isinstance(source_tensor, list):
            source_tensor = [source_tensor]

        output_tensor = torch.zeros_like(target_tensor)
        for source in source_tensor:
            target = target_tensor.permute(0, 3, 2, 1)  # Function expects b,w,h,c
            source = source.permute(0, 3, 2, 1)  # Function expects b,w,h,c
            if elementwise:
                source = source.mean(0).unsqueeze(0)
            if random_frame:
                source = source[np.random.randint(0, source.shape[0])].unsqueeze(0)

            matched_tensor = torch.zeros_like(target)
            for idx in range(target.shape[0] if elementwise else 1):
                frame = target[idx].unsqueeze(0) if elementwise else target
                _, t, Ct = get_histogram(frame + 1e-3 * torch.randn(size=frame.shape, device=frame.device))
                mu_s, _, Cs = get_histogram(
                    source.to(frame) + 1e-3 * torch.randn(size=source.shape, device=frame.device)
                )

                # PCA
                eva_t, eve_t = torch.linalg.eigh(Ct, UPLO="U")
                Et = torch.sqrt(torch.diagflat(eva_t))
                Et[Et != Et] = 0  # Convert nan to 0
                Qt = torch.mm(torch.mm(eve_t, Et), eve_t.T)

                eva_s, eve_s = torch.linalg.eigh(Cs, UPLO="U")
                Es = torch.sqrt(torch.diagflat(eva_s))
                Es[Es != Es] = 0  # Convert nan to 0
                Qs = torch.mm(torch.mm(eve_s, Es), eve_s.T)

                ts = torch.mm(torch.mm(Qs, torch.inverse(Qt)), t)

                match = ts.reshape(*frame.permute(0, 3, 1, 2).shape).permute(0, 2, 3, 1)
                match += mu_s

                if elementwise:
                    matched_tensor[idx] = match
                else:
                    matched_tensor = match
            output_tensor += matched_tensor.permute(0, 3, 2, 1) / len(source_tensor)
    except RuntimeError:
        traceback.print_exc()
        print("Skipping histogram matching...")
        output_tensor = backup
    return output_tensor.clamp(
        min([s.min().item() for s in source_tensor]), max([s.max().item() for s in source_tensor])
    )
//  [python]
def Eigen2xywhTheda_gpu(input_vector):  # [:, xywhtheda] [a1, a2, cx; a3, a4, cy]

    # 2D
    if len(input_vector.shape) == 2:
        output = torch.zeros((input_vector.shape[0], 5))
        input = torch.zeros((input_vector.shape[0], 2, 3))   # Vector to Matrix
        input[:, 0, 2] = input_vector[:, 0]  # cx
        input[:, 1, 2] = input_vector[:, 1]  # cy
        input[:, 0, 0] = input_vector[:, 2]  # a1
        input[:, 0, 1] = input_vector[:, 4]  # a2
        input[:, 1, 0] = input_vector[:, 4]  # a3
        input[:, 1, 1] = input_vector[:, 3]  # a4

        output[:, 0] = input[:, 0, 2]  # cx
        output[:, 1] = input[:, 1, 2]  # cy

        values, vectors = torch.symeig(input[:, :, :2], eigenvectors=True)  # np.where
        a = values.argmax(axis=1)
        output[:, 2] = values[torch.arange(len(a)), a]
        b = 1 - a
        output[:, 3] = values[torch.arange(len(b)), b]

        theda_vector = vectors[torch.arange(len(a)), :, a]
        theda = torch.arctan2(theda_vector[:, 1], theda_vector[:, 0]) + pi  # 0 ~ 2*pi
        theda[theda >= pi] -= pi
        theda[theda == pi] -= pi  # [0, 2*pi] to [0, pi)

        output_theda = pi - theda
        output_theda[output_theda == pi] -= pi  # (0, pi] to [0, pi)
        output[:, 4] = output_theda

    elif len(input_vector.shape) == 3:
        output_3 = torch.zeros((input_vector.shape[0], input_vector.shape[1], 5))
        for index in range(input_vector.shape[0]):
            output = torch.zeros((input_vector.shape[1], 5))
            input = torch.zeros((input_vector.shape[1], 2, 3))   # Vector to Matrix
            input[:, 0, 2] = input_vector[index, :, 0]  # cx
            input[:, 1, 2] = input_vector[index, :, 1]  # cy
            input[:, 0, 0] = input_vector[index, :, 2]  # a1
            input[:, 0, 1] = input_vector[index, :, 4]  # a2
            input[:, 1, 0] = input_vector[index, :, 4]  # a3
            input[:, 1, 1] = input_vector[index, :, 3]  # a4

            output[:, 0] = input[:, 0, 2]  # cx
            output[:, 1] = input[:, 1, 2]  # cy

            values, vectors = torch.symeig(input[:, :, :2], eigenvectors=True)  # np.where
            a = values.argmax(axis=1)
            output[:, 2] = values[torch.arange(a.size), a]

            b = 1 - a
            output[:, 3] = values[torch.arange(b.size), b]

            theda_vector = vectors[torch.arange(a.size), :, a]
            theda = torch.arctan2(theda_vector[:, 1], theda_vector[:, 0]) + pi  # 0 ~ 2*pi
            theda[theda >= pi] -= pi
            theda[theda == pi] -= pi  # [0, 2*pi] to [0, pi)

            output_theda = pi - theda
            output_theda[output_theda == pi] -= pi  # (0, pi] to [0, pi)
            output[:, 4] = output_theda

            output_3[index, :, :] = output
        output = output_3

    output[..., 4] = output[..., 4] / pi * 180
    return output
//  [python]
def Eigen2xywhTheda_gpu(input_vector):  # [:, xywhtheda] [a1, a2, cx; a3, a4, cy]

    # 2D
    if len(input_vector.shape) == 2:
        output = torch.zeros((input_vector.shape[0], 5))
        input = torch.zeros((input_vector.shape[0], 2, 3))   # Vector to Matrix
        input[:, 0, 2] = input_vector[:, 0]  # cx
        input[:, 1, 2] = input_vector[:, 1]  # cy
        input[:, 0, 0] = input_vector[:, 2]  # a1
        input[:, 0, 1] = input_vector[:, 4]  # a2
        input[:, 1, 0] = input_vector[:, 4]  # a3
        input[:, 1, 1] = input_vector[:, 3]  # a4

        output[:, 0] = input[:, 0, 2]  # cx
        output[:, 1] = input[:, 1, 2]  # cy

        values, vectors = torch.linalg.eigh(input[:, :, :2])  # np.where
        a = values.argmax(axis=1)
        output[:, 2] = values[torch.arange(len(a)), a]
        b = 1 - a
        output[:, 3] = values[torch.arange(len(b)), b]

        theda_vector = vectors[torch.arange(len(a)), :, a]
        theda = torch.arctan2(theda_vector[:, 1], theda_vector[:, 0]) + pi  # 0 ~ 2*pi
        theda[theda >= pi] -= pi
        theda[theda == pi] -= pi  # [0, 2*pi] to [0, pi)

        output_theda = pi - theda
        output_theda[output_theda == pi] -= pi  # (0, pi] to [0, pi)
        output[:, 4] = output_theda

    elif len(input_vector.shape) == 3:
        output_3 = torch.zeros((input_vector.shape[0], input_vector.shape[1], 5))
        for index in range(input_vector.shape[0]):
            output = torch.zeros((input_vector.shape[1], 5))
            input = torch.zeros((input_vector.shape[1], 2, 3))   # Vector to Matrix
            input[:, 0, 2] = input_vector[index, :, 0]  # cx
            input[:, 1, 2] = input_vector[index, :, 1]  # cy
            input[:, 0, 0] = input_vector[index, :, 2]  # a1
            input[:, 0, 1] = input_vector[index, :, 4]  # a2
            input[:, 1, 0] = input_vector[index, :, 4]  # a3
            input[:, 1, 1] = input_vector[index, :, 3]  # a4

            output[:, 0] = input[:, 0, 2]  # cx
            output[:, 1] = input[:, 1, 2]  # cy

            values, vectors = torch.linalg.eigh(input[:, :, :2])  # np.where
            a = values.argmax(axis=1)
            output[:, 2] = values[torch.arange(a.size), a]

            b = 1 - a
            output[:, 3] = values[torch.arange(b.size), b]

            theda_vector = vectors[torch.arange(a.size), :, a]
            theda = torch.arctan2(theda_vector[:, 1], theda_vector[:, 0]) + pi  # 0 ~ 2*pi
            theda[theda >= pi] -= pi
            theda[theda == pi] -= pi  # [0, 2*pi] to [0, pi)

            output_theda = pi - theda
            output_theda[output_theda == pi] -= pi  # (0, pi] to [0, pi)
            output[:, 4] = output_theda

            output_3[index, :, :] = output
        output = output_3

    output[..., 4] = output[..., 4] / pi * 180
    return output
//  [python]
def dm_to_mixture_of_state(dm: torch.Tensor, atol=1e-10):
    """
    Args:
        q_device: The state vector to take the partial trace over.
        keep_indices: Which indices to take the partial trace of the
            state_vector on.
        atol: The tolerance for determining that a factored state is pure.
    """
    size = dm.size()
    batched = len(size) % 2 == 1
    if batched:
        dims = int(np.log2(np.prod(size) / size[0]) / 2)

        eigvals, eigvecs = torch.symeig(dm)
        result_list = []
        for batch in range(size[0]):
            mixture = tuple(
                zip(
                    eigvals[batch],
                    [vec.reshape([2] * dims) for vec in eigvecs[batch].T],
                )
            )
            result_list.append(
                tuple([(float(p[0]), p[1]) for p in mixture if p[0] > 1e-10])
            )
        return result_list
    else:

        dims = int(np.log2(np.prod(size)) / 2)

        eigvals, eigvecs = torch.symeig(dm)
        mixture = tuple(zip(eigvals, [vec.reshape([2] * dims) for vec in eigvecs.T]))
        return tuple([(float(p[0]), p[1]) for p in mixture if p[0] > 1e-10])
//  [python]
def dm_to_mixture_of_state(dm: torch.Tensor, atol=1e-10):
    """
    Args:
        q_device: The state vector to take the partial trace over.
        keep_indices: Which indices to take the partial trace of the
            state_vector on.
        atol: The tolerance for determining that a factored state is pure.
    """
    size = dm.size()
    batched = len(size) % 2 == 1
    if batched:
        dims = int(np.log2(np.prod(size) / size[0]) / 2)

        eigvals, eigvecs = torch.linalg.eigh(dm)
        result_list = []
        for batch in range(size[0]):
            mixture = tuple(
                zip(
                    eigvals[batch],
                    [vec.reshape([2] * dims) for vec in eigvecs[batch].T],
                )
            )
            result_list.append(
                tuple([(float(p[0]), p[1]) for p in mixture if p[0] > 1e-10])
            )
        return result_list
    else:

        dims = int(np.log2(np.prod(size)) / 2)

        eigvals, eigvecs = torch.linalg.eigh(dm)
        mixture = tuple(zip(eigvals, [vec.reshape([2] * dims) for vec in eigvecs.T]))
        return tuple([(float(p[0]), p[1]) for p in mixture if p[0] > 1e-10])
//  [python]
def test_solve_lyapunov_spectral():
    d = 10
    tril_ind = torch.tril_indices(d, d)
    # setting up some spd matrices
    Q = torch.zeros(d, d)
    Q[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
    Q = Q @ Q.t()
    K = torch.zeros(d, d)
    K[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
    K = K @ K.t() + torch.eye(d) * 1e-3
    # solving the lyapunov equation using eig
    D, R = torch.eig(K, eigenvectors=True)
    st_time = timeit.default_timer()
    X = solve_lyapunov_spectral(D, R, Q)
    fast_time = timeit.default_timer() - st_time
    Q_solve = K @ X.t() + X @ K.t()
    with torch.no_grad():
        npt.assert_array_almost_equal(Q_solve, Q, decimal=4)
    # check it's actually fast
    st_time = timeit.default_timer()
    # naive solve
    I = torch.eye(d)
    LHS = torch.kron(I, K) + torch.kron(K, I)
    vecX = torch.linalg.solve(LHS, Q.t().ravel())
    slow_time = timeit.default_timer() - st_time
    assert slow_time > fast_time
//  [python]
def test_solve_lyapunov_spectral():
    d = 10
    tril_ind = torch.tril_indices(d, d)
    # setting up some spd matrices
    Q = torch.zeros(d, d)
    Q[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
    Q = Q @ Q.t()
    K = torch.zeros(d, d)
    K[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
    K = K @ K.t() + torch.eye(d) * 1e-3
    # solving the lyapunov equation using eig
    D, R = torch.linalg.eigh(K)
    st_time = timeit.default_timer()
    X = solve_lyapunov_spectral(D, R, Q)
    fast_time = timeit.default_timer() - st_time
    Q_solve = K @ X.t() + X @ K.t()
    with torch.no_grad():
        npt.assert_array_almost_equal(Q_solve, Q, decimal=4)
    # check it's actually fast
    st_time = timeit.default_timer()
    # naive solve
    I = torch.eye(d)
    LHS = torch.kron(I, K) + torch.kron(K, I)
    vecX = torch.linalg.solve(LHS, Q.t().ravel())
    slow_time = timeit.default_timer() - st_time
    assert slow_time > fast_time
//  [python]
def test_solve_lyapunov_spectral_double_batch():
    d = 10
    m = 3
    tril_ind = torch.tril_indices(d, d)
    # setting up some spd matrices
    K_in = []
    for j in range(m):
        K = init_spd(d) + torch.eye(d) * 1e-3
        K_in.append(K)
    Q_in = torch.stack([torch.stack([init_spd(d) for i in range(3)]) for i in range(5)])
    K_in = torch.stack(K_in)
    # solving the lyapunov equation using eig
    D, R = torch.eig(K_in, eigenvectors=True)
    st_time = timeit.default_timer()
    X = solve_lyapunov_spectral(D, R, Q_in)
    fast_time = timeit.default_timer() - st_time
    Q_solve = K_in @ X.transpose(-2, -1) + X @ K_in.transpose(-2, -1)
    with torch.no_grad():
        npt.assert_array_almost_equal(
            Q_solve - Q_in, torch.zeros(5, 3, d, d), decimal=4
        )
//  [python]
def test_solve_lyapunov_spectral_double_batch():
    d = 10
    m = 3
    tril_ind = torch.tril_indices(d, d)
    # setting up some spd matrices
    K_in = []
    for j in range(m):
        K = init_spd(d) + torch.eye(d) * 1e-3
        K_in.append(K)
    Q_in = torch.stack([torch.stack([init_spd(d) for i in range(3)]) for i in range(5)])
    K_in = torch.stack(K_in)
    # solving the lyapunov equation using eig
    D, R = torch.linalg.eigh(K_in)
    st_time = timeit.default_timer()
    X = solve_lyapunov_spectral(D, R, Q_in)
    fast_time = timeit.default_timer() - st_time
    Q_solve = K_in @ X.transpose(-2, -1) + X @ K_in.transpose(-2, -1)
    with torch.no_grad():
        npt.assert_array_almost_equal(
            Q_solve - Q_in, torch.zeros(5, 3, d, d), decimal=4
        )
//  [python]
def test_solve_lyapunov_spectral_batch():
    d = 10
    m = 3
    tril_ind = torch.tril_indices(d, d)
    # setting up some spd matrices
    Q_in = []
    K_in = []
    for j in range(m):
        Q = torch.zeros(d, d)
        Q[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
        Q = Q @ Q.t()
        Q_in.append(Q)
        K = torch.zeros(d, d)
        K[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
        K = K @ K.t() + torch.eye(d) * 1e-3
        K_in.append(K)
    Q_in = torch.stack(Q_in)
    K_in = torch.stack(K_in)
    # solving the lyapunov equation using eig
    D, R = torch.eig(K_in, eigenvectors=True)
    st_time = timeit.default_timer()
    X = solve_lyapunov_spectral(D, R, Q_in)
    fast_time = timeit.default_timer() - st_time
    Q_solve = K_in @ X.transpose(-2, -1) + X @ K_in.transpose(-2, -1)
    with torch.no_grad():
        npt.assert_array_almost_equal(Q_solve, Q_in, decimal=4)
//  [python]
def test_solve_lyapunov_spectral_batch():
    d = 10
    m = 3
    tril_ind = torch.tril_indices(d, d)
    # setting up some spd matrices
    Q_in = []
    K_in = []
    for j in range(m):
        Q = torch.zeros(d, d)
        Q[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
        Q = Q @ Q.t()
        Q_in.append(Q)
        K = torch.zeros(d, d)
        K[tril_ind[0], tril_ind[1]] = torch.randn(d * (d + 1) // 2)
        K = K @ K.t() + torch.eye(d) * 1e-3
        K_in.append(K)
    Q_in = torch.stack(Q_in)
    K_in = torch.stack(K_in)
    # solving the lyapunov equation using eig
    D, R = torch.linalg.eigh(K_in)
    st_time = timeit.default_timer()
    X = solve_lyapunov_spectral(D, R, Q_in)
    fast_time = timeit.default_timer() - st_time
    Q_solve = K_in @ X.transpose(-2, -1) + X @ K_in.transpose(-2, -1)
    with torch.no_grad():
        npt.assert_array_almost_equal(Q_solve, Q_in, decimal=4)
//  [python]
def ortho_matrix(O: torch.Tensor, use_cholesky: bool = True) -> torch.Tensor:
    """Return orthonormalization matrix of a basis.
    The basis is specified by its overlap matrix or metric, `O`.

    Parameters
    ----------
    O
        Overlap matrix / metric (Hermitian, positive definite) in last
        two dimensions, and batched over any preceding dimensions
    use_cholesky
        If True, use Cholesky decomposition followed by a triangular solve,
        essentially amounting to Gram-Schmidt orthonormalization.
        If False, return the symmetric orthonormalization matrix calculated
        by diagonalizing O, which may be more stable, but may be an order of
        magnitude slower than the default Cholesky method
    """
    assert O.shape[-2] == O.shape[-1]  # check square
    if use_cholesky:
        # Gram-Schmidt orthonormalization matrix:
        identity = torch.eye(O.shape[-1], device=O.device, dtype=O.dtype).view(
            (1,) * (len(O.shape) - 2) + O.shape[-2:]
        )
        return dagger(
            torch.linalg.solve_triangular(
                torch.linalg.cholesky(O), identity, upper=False
            )
        )
    else:
        # Symmetric orthonormalization matrix:
        lbda, V = torch.symeig(O)
        return V @ ((1.0 / torch.sqrt(lbda))[..., None] * dagger(V))
//  [python]
def ortho_matrix(O: torch.Tensor, use_cholesky: bool = True) -> torch.Tensor:
    """Return orthonormalization matrix of a basis.
    The basis is specified by its overlap matrix or metric, `O`.

    Parameters
    ----------
    O
        Overlap matrix / metric (Hermitian, positive definite) in last
        two dimensions, and batched over any preceding dimensions
    use_cholesky
        If True, use Cholesky decomposition followed by a triangular solve,
        essentially amounting to Gram-Schmidt orthonormalization.
        If False, return the symmetric orthonormalization matrix calculated
        by diagonalizing O, which may be more stable, but may be an order of
        magnitude slower than the default Cholesky method
    """
    assert O.shape[-2] == O.shape[-1]  # check square
    if use_cholesky:
        # Gram-Schmidt orthonormalization matrix:
        identity = torch.eye(O.shape[-1], device=O.device, dtype=O.dtype).view(
            (1,) * (len(O.shape) - 2) + O.shape[-2:]
        )
        return dagger(
            torch.linalg.solve_triangular(
                torch.linalg.cholesky(O), identity, upper=False
            )
        )
    else:
        # Symmetric orthonormalization matrix:
        lbda, V = torch.linalg.eigh(O)
        return V @ ((1.0 / torch.sqrt(lbda))[..., None] * dagger(V))
//  [python]
def forward(self, S_vec, return_grad=False):
        S = torch.zeros((*S_vec.shape[:-1], self.d, self.d))
        S[..., self.ind[0], self.ind[1]] = -S_vec
        S[..., self.ind[1], self.ind[0]] = S_vec
        D, V = torch.linalg.eig(1.0j * S)  # j * S is Hermitian
        D = D.real / 1.0j
        exp_D = torch.exp(D)
        V_trns_conj = V.transpose(-2, -1).conj()
        exp_S = (V * exp_D.unsqueeze(-2) @ V_trns_conj).real
        if not return_grad:
            return exp_S
        else:
            # https://www.janmagnus.nl/wips/expo-23.pdf
            # todo: can we write Delta in terms of a kronecker product?
            Delta = torch.zeros_like(V)
            diff = (exp_D[..., self.ind[0]] - exp_D[..., self.ind[1]]) / (
                D[..., self.ind[0]] - D[..., self.ind[1]]
            )
            Delta[..., self.ind[0], self.ind[1]] = diff
            Delta[..., self.ind[1], self.ind[0]] = diff
            Delta[..., self.d_ind, self.d_ind] = exp_D
            if S.dim() == 2:
                R = torch.kron(V.conj(), V)
                R_inv_dupl = (
                    torch.kron(V.transpose(-2, -1), V_trns_conj) @ self.dupl_mat
                )
            elif S.dim() == 3:
                # batch wise kronecker product
                R = torch.einsum("ijk,ilm->ijlkm", V.conj(), V).view(
                    -1, self.d ** 2, self.d ** 2
                )
                R_inv_dupl = (
                    torch.einsum(
                        "ijk,ilm->ijlkm", V.transpose(-2, -1), V_trns_conj
                    ).view(-1, self.d ** 2, self.d ** 2)
                    @ self.dupl_mat
                )
            else:
                raise ValueError("S must be 2 or 3 dimensional")
            grad_exp_S = (
                R * Delta.flatten(start_dim=-2, end_dim=-1).unsqueeze(-2) @ (R_inv_dupl)
            )
            return exp_S, grad_exp_S.real
//  [python]
def forward(self, S_vec, return_grad=False):
        S = torch.zeros((*S_vec.shape[:-1], self.d, self.d))
        S[..., self.ind[0], self.ind[1]] = -S_vec
        S[..., self.ind[1], self.ind[0]] = S_vec
        D, V = torch.linalg.eigh(1.0j * S)  # j * S is Hermitian
        D = D / 1.0j
        exp_D = torch.exp(D)
        V_trns_conj = V.transpose(-2, -1).conj()
        exp_S = (V * exp_D.unsqueeze(-2) @ V_trns_conj).real
        if not return_grad:
            return exp_S
        else:
            # https://www.janmagnus.nl/wips/expo-23.pdf
            # todo: can we write Delta in terms of a kronecker product?
            Delta = torch.zeros_like(V)
            diff = (exp_D[..., self.ind[0]] - exp_D[..., self.ind[1]]) / (
                D[..., self.ind[0]] - D[..., self.ind[1]]
            )
            Delta[..., self.ind[0], self.ind[1]] = diff
            Delta[..., self.ind[1], self.ind[0]] = diff
            Delta[..., self.d_ind, self.d_ind] = exp_D
            if S.dim() == 2:
                R = torch.kron(V.conj(), V)
                R_inv_dupl = (
                    torch.kron(V.transpose(-2, -1), V_trns_conj) @ self.dupl_mat
                )
            elif S.dim() == 3:
                # batch wise kronecker product
                R = torch.einsum("ijk,ilm->ijlkm", V.conj(), V).view(
                    -1, self.d ** 2, self.d ** 2
                )
                R_inv_dupl = (
                    torch.einsum(
                        "ijk,ilm->ijlkm", V.transpose(-2, -1), V_trns_conj
                    ).view(-1, self.d ** 2, self.d ** 2)
                    @ self.dupl_mat
                )
            else:
                raise ValueError("S must be 2 or 3 dimensional")
            grad_exp_S = (
                R * Delta.flatten(start_dim=-2, end_dim=-1).unsqueeze(-2) @ (R_inv_dupl)
            )
            return exp_S, grad_exp_S.real
//  [python]
def eighg(
    H: Waitable[torch.Tensor],
    O: torch.Tensor,
    use_cholesky: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Solve Hermitian generalized eigenvalue problem.
    Specifically, find `E` and `V` that satisfy `H` @ `V` = `O` @ `V` @ `E`.

    Parameters
    ----------
    H
        Set of complex Hermitian (in last two dimensions) matrices
        to diagonalize, all dimensions before last two are batched over.
        Pending communication on H may be overlapped with orthogonalization of O,
        and is hence taken as a Waitable Tensor to facilitate this.
    O
        Corresponding overlap (metric) matrices, with same size as H.
        Must additionally be positive-definite.
    use_cholesky
        See :meth:`qimpy.mpi.ortho_matrix`

    Returns
    -------
    E : torch.Tensor
        Real eigenvalues (shape = H.shape[:-1])
    V : torch.Tensor
        Eigenvectors (same shape as H and O)
    """
    # Start orthogonoalization:
    rc.compute_stream_wait_current()
    with torch.cuda.stream(rc.compute_stream):
        U = ortho_matrix(O, use_cholesky)
    # Finish pending communication on H (if any)
    Hresult = H.wait()
    # Finish orthogonalization:
    rc.current_stream_wait_compute()
    # Diagonalize:
    E, V = torch.symeig(dagger(U) @ (Hresult @ U), eigenvectors=True)
    V = U @ V  # transform eigenvectors back to original basis
    return E, V
//  [python]
def eighg(
    H: Waitable[torch.Tensor],
    O: torch.Tensor,
    use_cholesky: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Solve Hermitian generalized eigenvalue problem.
    Specifically, find `E` and `V` that satisfy `H` @ `V` = `O` @ `V` @ `E`.

    Parameters
    ----------
    H
        Set of complex Hermitian (in last two dimensions) matrices
        to diagonalize, all dimensions before last two are batched over.
        Pending communication on H may be overlapped with orthogonalization of O,
        and is hence taken as a Waitable Tensor to facilitate this.
    O
        Corresponding overlap (metric) matrices, with same size as H.
        Must additionally be positive-definite.
    use_cholesky
        See :meth:`qimpy.mpi.ortho_matrix`

    Returns
    -------
    E : torch.Tensor
        Real eigenvalues (shape = H.shape[:-1])
    V : torch.Tensor
        Eigenvectors (same shape as H and O)
    """
    # Start orthogonoalization:
    rc.compute_stream_wait_current()
    with torch.cuda.stream(rc.compute_stream):
        U = ortho_matrix(O, use_cholesky)
    # Finish pending communication on H (if any)
    Hresult = H.wait()
    # Finish orthogonalization:
    rc.current_stream_wait_compute()
    # Diagonalize:
    E, V = torch.linalg.eigh(dagger(U) @ (Hresult @ U))
    V = U @ V  # transform eigenvectors back to original basis
    return E, V
//  [python]
def test_diag_laplacian_builder_produces_valid_laplacian(graph_id, d, normalised):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = DiagLaplacianBuilder(size, data.edge_index, d=d, normalised=normalised)
    maps = torch.FloatTensor(size=(data.edge_index.size(1), d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    assert torch.isfinite(L[1]).all()

    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, V = torch.symeig(L, eigenvectors=True)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps)
    if normalised:
        assert torch.all(Q <= 2. + eps)

    L_dense = build_dense_laplacian(size, data.edge_index, maps.to(dtype=torch.float64), d=d,
                                    normalised=normalised, diagonal_maps=True)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_diag_laplacian_builder_produces_valid_laplacian(graph_id, d, normalised):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = DiagLaplacianBuilder(size, data.edge_index, d=d, normalised=normalised)
    maps = torch.FloatTensor(size=(data.edge_index.size(1), d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    assert torch.isfinite(L[1]).all()

    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, V = torch.linalg.eigh(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps)
    if normalised:
        assert torch.all(Q <= 2. + eps)

    L_dense = build_dense_laplacian(size, data.edge_index, maps.to(dtype=torch.float64), d=d,
                                    normalised=normalised, diagonal_maps=True)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_diag_laplacian_builder_with_fixed_maps_produces_valid_laplacian(graph_id, d, normalised):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = DiagLaplacianBuilder(size, data.edge_index, d=d, normalised=normalised, add_hp=True, add_lp=True)
    maps = torch.FloatTensor(size=(data.edge_index.size(1), d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    assert torch.isfinite(L[1]).all()

    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*(d+2), size*(d+2)), dtype=torch.float64).to_dense()
    Q, V = torch.symeig(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps)
    if normalised:
        assert torch.all(Q <= 2. + eps)

    # Append fixed submaps to the restriction maps.
    values = [1.0, -1.0]
    L_dense = build_dense_laplacian(size, data.edge_index, maps.to(dtype=torch.float64), d=d,
                                    normalised=normalised, diagonal_maps=True, values=values)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_diag_laplacian_builder_with_fixed_maps_produces_valid_laplacian(graph_id, d, normalised):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = DiagLaplacianBuilder(size, data.edge_index, d=d, normalised=normalised, add_hp=True, add_lp=True)
    maps = torch.FloatTensor(size=(data.edge_index.size(1), d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    assert torch.isfinite(L[1]).all()

    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*(d+2), size*(d+2)), dtype=torch.float64).to_dense()
    Q, V = torch.linalg.eigh(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps)
    if normalised:
        assert torch.all(Q <= 2. + eps)

    # Append fixed submaps to the restriction maps.
    values = [1.0, -1.0]
    L_dense = build_dense_laplacian(size, data.edge_index, maps.to(dtype=torch.float64), d=d,
                                    normalised=normalised, diagonal_maps=True, values=values)
    assert torch.allclose(L, L_dense, atol=eps)