jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def init_spatial_model(self):
        self.start_idx = 0
        self.Q_FMM = torch.tile(
            torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device), [self.n_freq, 1, 1]
        )
        self.G_NM = self.g_eps * torch.ones([self.n_source, self.n_mic], dtype=self.TYPE_FLOAT, device=self.device)
        for m in range(self.n_mic):
            self.G_NM[m % self.n_source, m] = 1

        if "circular" in self.init_SCM:
            pass
        elif "obs" in self.init_SCM:
            if hasattr(self, "XX_FTMM"):
                XX_FMM = self.XX_FTMM.sum(axis=1)
            else:
                XX_FMM = torch.einsum("fti, ftj -> fij", self.X_FTM, self.X_FTM.conj())
            w, v = torch.linalg.eig(XX_FMM)
            v = v[:, :, range(self.n_mic - 1, -1, -1)]
            self.Q_FMM = v.permute(0, 2, 1).conj()
        elif "twostep" == self.init_SCM:
            if self.n_iter_init >= self.n_iter:
                print(
                    "\n------------------------------------------------------------------\n"
                    f"Warning: n_iter_init must be smaller than n_iter (= {self.n_iter}).\n"
                    f"n_iter_init is changed from {self.n_iter_init} to {self.n_iter // 3}"
                    "\n------------------------------------------------------------------\n"
                )
                self.n_iter_init = self.n_iter // 3

            self.start_idx = self.n_iter_init

            separater_init = FastMNMF2(
                n_source=self.n_source,
                n_basis=2,
                device=self.device,
                init_SCM="circular",
                n_bit=self.n_bit,
                g_eps=self.g_eps,
            )
            separater_init.load_spectrogram(self.X_FTM, self.sample_rate)
            separater_init.solve(n_iter=self.start_idx, save_wav=False)

            self.Q_FMM = separater_init.Q_FMM
            self.G_NM = separater_init.G_NM
        else:
            raise ValueError("init_SCM should be circular, obs, or twostep.")

        self.G_NM /= self.G_NM.sum(axis=1)[:, None]
        self.normalize()
//  [python]
def init_spatial_model(self):
        self.start_idx = 0
        self.Q_FMM = torch.tile(
            torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device), [self.n_freq, 1, 1]
        )
        self.G_NM = self.g_eps * torch.ones([self.n_source, self.n_mic], dtype=self.TYPE_FLOAT, device=self.device)
        for m in range(self.n_mic):
            self.G_NM[m % self.n_source, m] = 1

        if "circular" in self.init_SCM:
            pass
        elif "obs" in self.init_SCM:
            if hasattr(self, "XX_FTMM"):
                XX_FMM = self.XX_FTMM.sum(axis=1)
            else:
                XX_FMM = torch.einsum("fti, ftj -> fij", self.X_FTM, self.X_FTM.conj())
            _, eig_vec_FMM = torch.linalg.eigh(XX_FMM)
            eig_vec_FMM = eig_vec_FMM[:, :, range(self.n_mic - 1, -1, -1)]
            self.Q_FMM = eig_vec_FMM.permute(0, 2, 1).conj()
        elif "twostep" == self.init_SCM:
            if self.n_iter_init >= self.n_iter:
                print(
                    "\n------------------------------------------------------------------\n"
                    f"Warning: n_iter_init must be smaller than n_iter (= {self.n_iter}).\n"
                    f"n_iter_init is changed from {self.n_iter_init} to {self.n_iter // 3}"
                    "\n------------------------------------------------------------------\n"
                )
                self.n_iter_init = self.n_iter // 3

            self.start_idx = self.n_iter_init

            separater_init = FastMNMF2(
                n_source=self.n_source,
                n_basis=2,
                device=self.device,
                init_SCM="circular",
                n_bit=self.n_bit,
                g_eps=self.g_eps,
            )
            separater_init.load_spectrogram(self.X_FTM, self.sample_rate)
            separater_init.solve(n_iter=self.start_idx, save_wav=False)

            self.Q_FMM = separater_init.Q_FMM
            self.G_NM = separater_init.G_NM
        else:
            raise ValueError("init_SCM should be circular, obs, or twostep.")

        self.G_NM /= self.G_NM.sum(axis=1)[:, None]
        self.normalize()
//  [python]
def softabs(hessian, a=1000.0):
    l, Q = torch.eig(hessian, eigenvectors=True)
    l = l[:, 0]
    return Q @ torch.diag(torch.div(l, torch.tanh(a * l))) @ Q.t()
//  [python]
def softabs(hessian, a=1000.0):
    l, Q = torch.linalg.eigh(hessian, "U")
    return Q @ torch.diag(torch.div(l, torch.tanh(a * l))) @ Q.t()
//  [python]
def average_quaternion_batch(Q, weights=None):
    """calculate the average quaternion of the multiple quaternions
    Args:
        Q (tensor): [B, num_quaternions, 4]
        weights (tensor, optional): [B, num_quaternions]. Defaults to None.

    Returns:
        oriented_q_avg: average quaternion, [B, 4]
    """
    
    if weights is None:
        weights = torch.ones((Q.shape[0], Q.shape[1]), device=Q.device) / Q.shape[1]
    A = torch.zeros((Q.shape[0], 4, 4), device=Q.device)
    weight_sum = torch.sum(weights, axis=-1)

    oriented_Q = ((Q[:, :, 0:1] > 0).float() - 0.5) * 2 * Q
    A = torch.einsum("abi,abk->abik", (oriented_Q, oriented_Q))
    A = torch.sum(torch.einsum("abij,ab->abij", (A, weights)), 1)
    A /= weight_sum.reshape(A.shape[0], -1).unsqueeze(-1).repeat(1, 4, 4)

    q_avg = torch.symeig(A, eigenvectors=True)[1][:, :, -1]
    oriented_q_avg = ((q_avg[:, 0:1] > 0).float() - 0.5) * 2 * q_avg
    return oriented_q_avg
//  [python]
def average_quaternion_batch(Q, weights=None):
    """calculate the average quaternion of the multiple quaternions
    Args:
        Q (tensor): [B, num_quaternions, 4]
        weights (tensor, optional): [B, num_quaternions]. Defaults to None.

    Returns:
        oriented_q_avg: average quaternion, [B, 4]
    """
    
    if weights is None:
        weights = torch.ones((Q.shape[0], Q.shape[1]), device=Q.device) / Q.shape[1]
    A = torch.zeros((Q.shape[0], 4, 4), device=Q.device)
    weight_sum = torch.sum(weights, axis=-1)

    oriented_Q = ((Q[:, :, 0:1] > 0).float() - 0.5) * 2 * Q
    A = torch.einsum("abi,abk->abik", (oriented_Q, oriented_Q))
    A = torch.sum(torch.einsum("abij,ab->abij", (A, weights)), 1)
    A /= weight_sum.reshape(A.shape[0], -1).unsqueeze(-1).repeat(1, 4, 4)

    q_avg = torch.linalg.eigh(A)[1][:, :, -1]
    oriented_q_avg = ((q_avg[:, 0:1] > 0).float() - 0.5) * 2 * q_avg
    return oriented_q_avg
//  [python]
def forward(self, x):
        """
        Generate SPD tensors from x

        Args:
            x (Tensor): Tensor to generate predictions for. Must have
                2d shape of form (:, input_shape). If symmetry='anisotropic',
                the expected
                `input_shape = sum([i for i in range(output_shape + 1)])`. If
                symmetry='orthotropic', then the expected `input_shape=9`.

        Returns:
            (Tensor): The predictions of the neural network. Will return
                shape (:, output_shape, output_shape)
        """
        x = torch.nan_to_num(x)  # we can't run torch.linalg.eig with NaNs!
        # init a placeholder tensor
        out = torch.zeros((x.shape[0], self.output_shape, self.output_shape),
                          dtype=x.dtype)
        out[:, self.inds_a, self.inds_b] = x
        out[:, self.inds_b, self.inds_a] = x

        # U, D, UT = torch.linalg.svd(out)  # SVD DOES NOT WORK! reason unknown
        D, U = torch.symeig(out,eigenvectors=True)
        D = self.positive_fun(D) + self.min_value
        if self.zero_eigvals:
            zeros_and_ones = torch.ones_like(D)
            # set the columns to zero
            zeros_and_ones[:, :self.n_zero_eigvals] = 0.0
            # zero out the smallest eigenvalues
            D = D * zeros_and_ones
        UT = U.inverse()  # don't transpose, need inverse!
        out = U @ torch.diag_embed(D) @ UT
        return out
//  [python]
def forward(self, x):
        """
        Generate SPD tensors from x

        Args:
            x (Tensor): Tensor to generate predictions for. Must have
                2d shape of form (:, input_shape). If symmetry='anisotropic',
                the expected
                `input_shape = sum([i for i in range(output_shape + 1)])`. If
                symmetry='orthotropic', then the expected `input_shape=9`.

        Returns:
            (Tensor): The predictions of the neural network. Will return
                shape (:, output_shape, output_shape)
        """
        x = torch.nan_to_num(x)  # we can't run torch.linalg.eig with NaNs!
        # init a placeholder tensor
        out = torch.zeros((x.shape[0], self.output_shape, self.output_shape),
                          dtype=x.dtype)
        out[:, self.inds_a, self.inds_b] = x
        out[:, self.inds_b, self.inds_a] = x

        # U, D, UT = torch.linalg.svd(out)  # SVD DOES NOT WORK! reason unknown
        D, U = torch.linalg.eigh(out)
        D = self.positive_fun(D) + self.min_value
        if self.zero_eigvals:
            zeros_and_ones = torch.ones_like(D)
            # set the columns to zero
            zeros_and_ones[:, :self.n_zero_eigvals] = 0.0
            # zero out the smallest eigenvalues
            D = D * zeros_and_ones
        UT = U.inverse()  # don't transpose, need inverse!
        out = U @ torch.diag_embed(D) @ UT
        return out
//  [python]
def opt_multiple_B_gauss_seidel(B_bar, w_bar, X_list, Y_list, lambda_1, w, rank_for_regressor, prev_B, max_iter=MAX_ITER_FOR_B, conv_criteria=CONV_CRITERIA):
    Y = change_to_matrix(Y_list)
    K = len(w_bar)
    # init B
    B = {}
    assert rank_for_regressor == 1
    if prev_B is None:
        raise NotImplementedError
    else:
        B = prev_B
    X_tilde = {}
    for k in range(K):
        X_tilde_list_k = [x @ w[k] for x in X_list]
        X_tilde[k] = change_to_matrix(X_tilde_list_k)
    curr_loss = 999999
    for ind_iter in range(max_iter):
        Y_check = {}
        X_check = {}
        for k in range(K):
            curr_sum_S = 0
            for k_prime in range(K):
                if k == k_prime:
                    pass
                else:
                    curr_sum_S += X_tilde[k_prime]@B[k_prime]@Herm(B[k_prime])   
            Y_check[k] = torch.cat( [Y - curr_sum_S, np.sqrt(lambda_1) * Herm(B_bar[k])], dim=0)
            X_check[k] = torch.cat( [X_tilde[k], np.sqrt(lambda_1) * Herm(B_bar[k])], dim=0)
            C = Herm(Y_check[k])@X_check[k] + Herm(X_check[k])@Y_check[k] - Herm(X_check[k])@X_check[k]
            L, Q = torch.eig(C, eigenvectors=True) 
            B[k] = Q[:, 0:rank_for_regressor] # Gauss-Seidel
        loss_tmp = inner_loss(B_bar, w_bar, X_list, Y_list, lambda_1, 0, B, w) #lambda_2 = 0 since we are fixing w here
        if loss_tmp<=conv_criteria:
            break 
        elif (loss_tmp <= curr_loss) and (torch.norm(curr_loss - loss_tmp) <= conv_criteria):
            break
        else:
            curr_loss = loss_tmp
    return B
//  [python]
def opt_multiple_B_gauss_seidel(B_bar, w_bar, X_list, Y_list, lambda_1, w, rank_for_regressor, prev_B, max_iter=MAX_ITER_FOR_B, conv_criteria=CONV_CRITERIA):
    Y = change_to_matrix(Y_list)
    K = len(w_bar)
    # init B
    B = {}
    assert rank_for_regressor == 1
    if prev_B is None:
        raise NotImplementedError
    else:
        B = prev_B
    X_tilde = {}
    for k in range(K):
        X_tilde_list_k = [x @ w[k] for x in X_list]
        X_tilde[k] = change_to_matrix(X_tilde_list_k)
    curr_loss = 999999
    for ind_iter in range(max_iter):
        Y_check = {}
        X_check = {}
        for k in range(K):
            curr_sum_S = 0
            for k_prime in range(K):
                if k == k_prime:
                    pass
                else:
                    curr_sum_S += X_tilde[k_prime]@B[k_prime]@Herm(B[k_prime])   
            Y_check[k] = torch.cat( [Y - curr_sum_S, np.sqrt(lambda_1) * Herm(B_bar[k])], dim=0)
            X_check[k] = torch.cat( [X_tilde[k], np.sqrt(lambda_1) * Herm(B_bar[k])], dim=0)
            C = Herm(Y_check[k])@X_check[k] + Herm(X_check[k])@Y_check[k] - Herm(X_check[k])@X_check[k]
            L, Q = torch.linalg.eigh(-C) 
            B[k] = Q[:, 0:rank_for_regressor] # Gauss-Seidel
        loss_tmp = inner_loss(B_bar, w_bar, X_list, Y_list, lambda_1, 0, B, w) #lambda_2 = 0 since we are fixing w here
        if loss_tmp<=conv_criteria:
            break 
        elif (loss_tmp <= curr_loss) and (torch.norm(curr_loss - loss_tmp) <= conv_criteria):
            break
        else:
            curr_loss = loss_tmp
    return B
//  [python]
def log_rmat(r_mat: torch.Tensor) -> torch.Tensor:
    skew_mat = r_mat - r_mat.transpose(-1, -2)  # costruisco la amtrice skew
    sk_vec = skew2vec(skew_mat)
    s_angle = (sk_vec).norm(p=2, dim=-1) / 2
    c_angle = (torch.einsum("...ii", r_mat) - 1) / 2
    angle = torch.atan2(s_angle, c_angle)  # trovo l'angolo di rotazione [B]
    scale = angle / (2 * s_angle)  # [B]
    # if s_angle = 0, i.e. rotation by 0 or pi (180), we get NaNs
    # by definition, scale values are 0 if rotating by 0.
    # This also breaks down if rotating by pi, fix further down
    scale[angle == 0.0] = 0.0
    log_r_mat = scale[..., None, None] * skew_mat  # [B x 3 x 3]

    # Check for NaNs caused by 180deg rotations.
    nanlocs = log_r_mat[..., 0, 0].isnan()
    nanmats = r_mat[nanlocs]
    # We need to use an alternative way of finding the logarithm for nanmats,
    # Use eigendecomposition to discover axis of rotation.
    # By definition, these are symmetric, so use eigh.
    # NOTE: linalg.eig() isn't in torch 1.8,
    #       and torch.eig() doesn't do batched matrices
    eigval, eigvec = torch.symeig(nanmats)
    # Final eigenvalue == 1, might be slightly off because floats, but other two are -ve.
    # this *should* just be the last column if the docs for eigh are true.
    nan_axes = eigvec[..., -1, :]
    nan_angle = angle[nanlocs]
    nan_skew = vec2skew(nan_angle[..., None] * nan_axes)
    log_r_mat[nanlocs] = nan_skew
    return log_r_mat
//  [python]
def log_rmat(r_mat: torch.Tensor) -> torch.Tensor:
    skew_mat = r_mat - r_mat.transpose(-1, -2)  # costruisco la amtrice skew
    sk_vec = skew2vec(skew_mat)
    s_angle = (sk_vec).norm(p=2, dim=-1) / 2
    c_angle = (torch.einsum("...ii", r_mat) - 1) / 2
    angle = torch.atan2(s_angle, c_angle)  # trovo l'angolo di rotazione [B]
    scale = angle / (2 * s_angle)  # [B]
    # if s_angle = 0, i.e. rotation by 0 or pi (180), we get NaNs
    # by definition, scale values are 0 if rotating by 0.
    # This also breaks down if rotating by pi, fix further down
    scale[angle == 0.0] = 0.0
    log_r_mat = scale[..., None, None] * skew_mat  # [B x 3 x 3]

    # Check for NaNs caused by 180deg rotations.
    nanlocs = log_r_mat[..., 0, 0].isnan()
    nanmats = r_mat[nanlocs]
    # We need to use an alternative way of finding the logarithm for nanmats,
    # Use eigendecomposition to discover axis of rotation.
    # By definition, these are symmetric, so use eigh.
    # NOTE: linalg.eig() isn't in torch 1.8,
    #       and torch.eig() doesn't do batched matrices
    eigval, eigvec = torch.linalg.eigh(nanmats)
    # Final eigenvalue == 1, might be slightly off because floats, but other two are -ve.
    # this *should* just be the last column if the docs for eigh are true.
    nan_axes = eigvec[..., -1, :]
    nan_angle = angle[nanlocs]
    nan_skew = vec2skew(nan_angle[..., None] * nan_axes)
    log_r_mat[nanlocs] = nan_skew
    return log_r_mat
//  [python]
def _matrix_and_concentrations(
        self,
        img: torch.Tensor,
        mask: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Gets the H&E stain matrix and concentrations for a given image.

        Args:
            img (torch.Tensor): Image (RGB uint8) with dimensions W, H, C.
            mask (bool): Mask white pixels (255) during calculation.
                Defaults to False.
            standardize (bool): Perform brightness standardization.
                Defaults to True.

        Returns:
            A tuple containing

                torch.Tensor: H&E stain matrix, shape = (3, 2)

                torch.Tensor: Concentrations of individual stains
        """
        img = img.reshape((-1, 3))

        if mask:
            ones = torch.all(img == 255, dim=1)

        if self.standardize:
            img = standardize_brightness(img, mask=mask)

        # Calculate optical density.
        OD = -torch.log((img.to(torch.float32) + 1) / self.Io)

        # Remove transparent pixles.
        if mask:
            ODhat = OD[~ (torch.any(OD < self.beta, dim=1) | ones)]
        else:
            ODhat = OD[~torch.any(OD < self.beta, dim=1)]

        # Compute eigenvectors.
        eigvals, eigvecs = torch.symeig(torch.cov(ODhat.T), eigenvectors=True)

        # Project on the plane spanned by the eigenvectors corresponding
        # to the two largest eigenvalues.
        That = ODhat.matmul(eigvecs[:, 1:3])

        phi = torch.atan2(That[:, 1], That[:, 0])

        minPhi = torch.quantile(phi, self.alpha / 100)
        maxPhi = torch.quantile(phi, 1 - self.alpha / 100)

        vMin = torch.tensordot(eigvecs[:, 1:3], T(torch.stack((torch.cos(minPhi), torch.sin(minPhi)))), dims=[[-1], [-1]])
        vMax = torch.tensordot(eigvecs[:, 1:3], T(torch.stack((torch.cos(maxPhi), torch.sin(maxPhi)))), dims=[[-1], [-1]])

        # Ensure the vector corresponding to hematoxylin is first, eosin second.
        if vMin[0] > vMax[0]:
            HE = torch.stack((vMin, vMax)).T
        else:
            HE = torch.stack((vMax, vMin)).T

        if mask:
            OD = OD[~ ones]

        # Rows correspond to channels (RGB), columns to OD values.
        Y = torch.reshape(OD, (-1, 3)).T

        # Determine concentrations of the individual stains.
        C = torch.linalg.lstsq(HE, Y, rcond=None)[0]

        return HE, C
//  [python]
def _matrix_and_concentrations(
        self,
        img: torch.Tensor,
        mask: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Gets the H&E stain matrix and concentrations for a given image.

        Args:
            img (torch.Tensor): Image (RGB uint8) with dimensions W, H, C.
            mask (bool): Mask white pixels (255) during calculation.
                Defaults to False.
            standardize (bool): Perform brightness standardization.
                Defaults to True.

        Returns:
            A tuple containing

                torch.Tensor: H&E stain matrix, shape = (3, 2)

                torch.Tensor: Concentrations of individual stains
        """
        img = img.reshape((-1, 3))

        if mask:
            ones = torch.all(img == 255, dim=1)

        if self.standardize:
            img = standardize_brightness(img, mask=mask)

        # Calculate optical density.
        OD = -torch.log((img.to(torch.float32) + 1) / self.Io)

        # Remove transparent pixles.
        if mask:
            ODhat = OD[~ (torch.any(OD < self.beta, dim=1) | ones)]
        else:
            ODhat = OD[~torch.any(OD < self.beta, dim=1)]

        # Compute eigenvectors.
        eigvals, eigvecs = torch.linalg.eigh(torch.cov(ODhat.T))

        # Project on the plane spanned by the eigenvectors corresponding
        # to the two largest eigenvalues.
        That = ODhat.matmul(eigvecs[:, 1:3])

        phi = torch.atan2(That[:, 1], That[:, 0])

        minPhi = torch.quantile(phi, self.alpha / 100)
        maxPhi = torch.quantile(phi, 1 - self.alpha / 100)

        vMin = torch.tensordot(eigvecs[:, 1:3], T(torch.stack((torch.cos(minPhi), torch.sin(minPhi)))), dims=[[-1], [-1]])
        vMax = torch.tensordot(eigvecs[:, 1:3], T(torch.stack((torch.cos(maxPhi), torch.sin(maxPhi)))), dims=[[-1], [-1]])

        # Ensure the vector corresponding to hematoxylin is first, eosin second.
        if vMin[0] > vMax[0]:
            HE = torch.stack((vMin, vMax)).T
        else:
            HE = torch.stack((vMax, vMin)).T

        if mask:
            OD = OD[~ ones]

        # Rows correspond to channels (RGB), columns to OD values.
        Y = torch.reshape(OD, (-1, 3)).T

        # Determine concentrations of the individual stains.
        C = torch.linalg.lstsq(HE, Y, rcond=None)[0]

        return HE, C
//  [python]
def DLT_postprocess(data, logits, Ks, relu):
    if relu:
        weights = torch.relu(torch.tanh(logits))
    else:
        log_ng = F.logsigmoid(logits )       
        weights = torch.exp(log_ng)
    
    x_shp = data.shape
    
    # Make input data (num_img_pair x num_corr x 4), b*5*n, (u,v,x,y,z)
    xx = data.view(x_shp[0], x_shp[1], 5).permute(0, 2, 1)
    ones = torch.ones_like(xx[:, 2])
    zeros = torch.zeros_like(xx[:, 2])
    Xw = xx[:, 2]
    Yw = xx[:, 3]
    Zw = xx[:, 4]
    u = xx[:, 0]
    v = xx[:, 1]

    X_up = torch.stack([
        Xw, Yw, Zw, ones,
        zeros, zeros, zeros, zeros,
        -u * Xw, -u * Yw, -u*Zw , -u
    ], dim=1).permute(0, 2, 1)
    X_bottom = torch.stack([
        zeros, zeros, zeros, zeros,
        Xw, Yw, Zw, ones,
        -v * Xw, -v * Yw, -v * Zw, -v
    ], dim=1).permute(0, 2, 1)
   
    X = torch.cat([X_up, X_bottom], dim=1)
    
    w_cat = torch.cat([weights, weights], dim=-1)
    wX = w_cat.unsqueeze(-1) * X
    XwX = torch.matmul(X.permute(0, 2, 1), wX)

    w, v = torch.symeig(XwX, eigenvectors=True)
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 3, 4))
    R_hat = e_hat[:,:,:3]
    t_hat = e_hat[:,:,3]

    # U,D,V = torch.svd(R_hat)
    U, D, Vh = torch.linalg.svd(R_hat, full_matrices=False)
    c = torch.reciprocal(D.sum(dim=-1)/3.0).unsqueeze(-1)
    
    # Verify the sign of c, using the point with highest weight
    _, indice = torch.max(weights, dim=-1)
    Point_verify = xx[torch.arange(x_shp[0]), 2:, indice]
    ones = torch.ones([x_shp[0], 1]).to(device)
    Point_verify = torch.cat([Point_verify, ones], dim=-1).unsqueeze(1)

    criterion = c * (torch.matmul(Point_verify, e_hat[:,-1].view(x_shp[0], 4, 1)).squeeze(-1))
    inv_c = -1*c
    c = torch.where(criterion > 0, c, inv_c)
    
    R_hat = torch.matmul(U, Vh.conj()) * torch.sign(c.unsqueeze(-1))
    t_hat = c * t_hat
    
    #(b,3,4)
    e_post = torch.cat([R_hat, t_hat.unsqueeze(-1)], dim=-1)
    return e_post
//  [python]
def DLT_postprocess(data, logits, Ks, relu):
    if relu:
        weights = torch.relu(torch.tanh(logits))
    else:
        log_ng = F.logsigmoid(logits )       
        weights = torch.exp(log_ng)
    
    x_shp = data.shape
    
    # Make input data (num_img_pair x num_corr x 4), b*5*n, (u,v,x,y,z)
    xx = data.view(x_shp[0], x_shp[1], 5).permute(0, 2, 1)
    ones = torch.ones_like(xx[:, 2])
    zeros = torch.zeros_like(xx[:, 2])
    Xw = xx[:, 2]
    Yw = xx[:, 3]
    Zw = xx[:, 4]
    u = xx[:, 0]
    v = xx[:, 1]

    X_up = torch.stack([
        Xw, Yw, Zw, ones,
        zeros, zeros, zeros, zeros,
        -u * Xw, -u * Yw, -u*Zw , -u
    ], dim=1).permute(0, 2, 1)
    X_bottom = torch.stack([
        zeros, zeros, zeros, zeros,
        Xw, Yw, Zw, ones,
        -v * Xw, -v * Yw, -v * Zw, -v
    ], dim=1).permute(0, 2, 1)
   
    X = torch.cat([X_up, X_bottom], dim=1)
    
    w_cat = torch.cat([weights, weights], dim=-1)
    wX = w_cat.unsqueeze(-1) * X
    XwX = torch.matmul(X.permute(0, 2, 1), wX)

    _, v = torch.linalg.eigh(XwX, UPLO='U')
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 3, 4))
    R_hat = e_hat[:,:,:3]
    t_hat = e_hat[:,:,3]

    # U,D,V = torch.svd(R_hat)
    U, D, Vh = torch.linalg.svd(R_hat, full_matrices=False)
    c = torch.reciprocal(D.sum(dim=-1)/3.0).unsqueeze(-1)
    
    # Verify the sign of c, using the point with highest weight
    _, indice = torch.max(weights, dim=-1)
    Point_verify = xx[torch.arange(x_shp[0]), 2:, indice]
    ones = torch.ones([x_shp[0], 1]).to(device)
    Point_verify = torch.cat([Point_verify, ones], dim=-1).unsqueeze(1)

    criterion = c * (torch.matmul(Point_verify, e_hat[:,-1].view(x_shp[0], 4, 1)).squeeze(-1))
    inv_c = -1*c
    c = torch.where(criterion > 0, c, inv_c)
    
    R_hat = torch.matmul(U, Vh.conj()) * torch.sign(c.unsqueeze(-1))
    t_hat = c * t_hat
    
    #(b,3,4)
    e_post = torch.cat([R_hat, t_hat.unsqueeze(-1)], dim=-1)
    return e_post
//  [python]
def spec_temp_batch(x, y, eps, allowable_failure_rate=.5):
    emp_cov_eigvals, emp_cov_eigvecs = np.linalg.eigh(x)
    y_hats, epsilons, successful_indices, num_iter_reweighting_refinements_completed = \
        spectral_templates_batch(emp_cov_batch=x, emp_cov_eigvecs_batch=emp_cov_eigvecs,
                                 allowable_failures=int(allowable_failure_rate*len(x)))

    # remove runs which did not complete successfully
    y_hats, epsilons, num_iter_reweighting_refinements_completed = \
        list(compress(y_hats, successful_indices)), list(compress(epsilons, successful_indices)), \
        list(compress(num_iter_reweighting_refinements_completed, successful_indices))
    y_hat = np.stack(y_hats, axis=0)
    y = y[successful_indices, :, :]

    y_hat = zero_diagonals(torch.tensor(y_hat))
    return y_hat
//  [python]
def spec_temp_batch(x, y, eps, allowable_failure_rate=.5):
    emp_cov_eigvals, emp_cov_eigvecs = eigh(x)
    y_hats, epsilons, successful_indices, num_iter_reweighting_refinements_completed = \
        spectral_templates_batch(emp_cov_batch=x, emp_cov_eigvecs_batch=emp_cov_eigvecs.numpy(),
                                 allowable_failures=int(allowable_failure_rate*len(x)))

    # remove runs which did not complete successfully
    y_hats, epsilons, num_iter_reweighting_refinements_completed = \
        list(compress(y_hats, successful_indices)), list(compress(epsilons, successful_indices)), \
        list(compress(num_iter_reweighting_refinements_completed, successful_indices))
    y_hat = np.stack(y_hats, axis=0)
    y = y[successful_indices, :, :]

    y_hat = zero_diagonals(torch.tensor(y_hat))
    return y_hat
    """
    metrics, threshold, regressions = find_best_performances(torch.tensor(y), torch.tensor(y_hat), threshold, regressions)

    metrics['failure_rate'] = 1 - len(successful_indices) / len(x)

    print(f'\n### Prediction Stats: using eps {eps:.5f}, threshold {threshold} ###')
    print(f'\tpred edge weights: max: {np.nanmax(y_hat):.3f}, median {np.nanmedian(y_hat):.3f}, mean {np.nanmean(y_hat):.3f}')
    print(f'\tave sparsity of true graphs: {edge_density(y).mean():.3f}')
    print(f'\tave sparsity of (i) raw pred: {edge_density(y_hat).mean():.3f}, (ii) pred > {threshold:.7f}:  {edge_density((y_hat > threshold) + 0.0).mean():.3f}')
    print(f'\tepsilons: mean: {np.nanmean(epsilons):.3f}, std: {np.nanstd(epsilons):.3f} ... epsilons {epsilons}')
    print(f"\tfailure rate: {100*metrics['failure_rate']:.4f}%")
    print(f'### Performance ###')
    print(f"\terrors {metrics['error']*100:.5f}%, using found best threshold: {threshold:.7f}")
    print(f"\tmses:  raw: {metrics['raw_mse']:.5f}, ols: {metrics['ols_mse']:.5f}, scaling_mse: {metrics['ols_no_intercept_mse']:.5f}")
    print('\t# iterative refines', num_iter_reweighting_refinements_completed)
    return metrics, threshold, regressions
    """
//  [python]
def exact(params, closure, L, delta=1e-8, testing=False):
    df = tuple_to_vec.tuple_to_vector(
        torch.autograd.grad(closure(), list(params), create_graph=True))
    H = derivatives.flat_hessian(df, list(params))

    c = df.detach().to(torch.double)
    A = H.detach().to(torch.double)

    if c.dim() != 1:
        raise ValueError(f"`c` must be a vector, but c = {c}")

    if A.dim() > 2:
        raise ValueError(f"`A` must be a matrix, but A = {A}")

    if c.size()[0] != A.size()[0]:
        raise ValueError("`c` and `A` mush have the same 1st dimension")

    if (A.t() - A).max() > 0.1:
        raise ValueError("`A` is not symmetric")

    T, U = torch.symeig(A)
    ct = U.t().mv(c)

    def inv(T, L, tau):
        return (T + L / 2 * tau).reciprocal()

    def dual(tau):
        return L / 12 * tau.pow(3) + 1 / 2 * \
               inv(T, L, tau).mul(ct.square()).sum()

    tau_best = line_search.ray_line_search(
        dual,
        left_point=torch.tensor([0.]),
        middle_point=torch.tensor([2.]),
        delta=delta)

    invert = inv(T, L, tau_best)
    x = -U.mv(invert.mul(ct).type_as(U))

    if testing and (c + L / 2 * x.norm() * x + A.mv(x)).abs().max().item() >= 0.01:
        raise ValueError('obtained `x` is not optimal')

    return tuple_to_vec.rollup_vector(x, list(params))
//  [python]
def exact(params, closure, L, delta=1e-8, testing=False):
    df = tuple_to_vec.tuple_to_vector(
        torch.autograd.grad(closure(), list(params), create_graph=True))
    H = derivatives.flat_hessian(df, list(params))

    c = df.detach().to(torch.double)
    A = H.detach().to(torch.double)

    if c.dim() != 1:
        raise ValueError(f"`c` must be a vector, but c = {c}")

    if A.dim() > 2:
        raise ValueError(f"`A` must be a matrix, but A = {A}")

    if c.size()[0] != A.size()[0]:
        raise ValueError("`c` and `A` mush have the same 1st dimension")

    if (A.t() - A).max() > 0.1:
        raise ValueError("`A` is not symmetric")

    T, U = torch.linalg.eigh(A)
    ct = U.t().mv(c)

    def inv(T, L, tau):
        return (T + L / 2 * tau).reciprocal()

    def dual(tau):
        return L / 12 * tau.pow(3) + 1 / 2 * \
               inv(T, L, tau).mul(ct.square()).sum()

    tau_best = line_search.ray_line_search(
        dual,
        left_point=torch.tensor([0.]),
        middle_point=torch.tensor([2.]),
        delta=delta)

    invert = inv(T, L, tau_best)
    x = -U.mv(invert.mul(ct).type_as(U))

    if testing and (c + L / 2 * x.norm() * x + A.mv(x)).abs().max().item() >= 0.01:
        raise ValueError('obtained `x` is not optimal')

    return tuple_to_vec.rollup_vector(x, list(params))
//  [python]
def spectral_templates_batch(emp_cov_batch: np.ndarray, emp_cov_eigvecs_batch: np.ndarray, allowable_failures: int):
    if emp_cov_eigvecs_batch is None:
        emp_cov_eigvecs_batch = eigh(emp_cov_batch)[1]

    y_hats, epsilons = [], []
    successful_runs, num_iter_reweighting_refinements_completed = [], []
    for i in range(len(emp_cov_batch)):
        try:
            print(f'##########')
            print(f'{i}th sample - start binary search')
            start = time.time()
            y_hat, epsilon, a = spectral_templates(emp_cov=emp_cov_batch[i], emp_cov_eigenvectors=emp_cov_eigvecs_batch[i])
            y_hats.append(y_hat); epsilons.append(epsilon); successful_runs.append(True); num_iter_reweighting_refinements_completed.append(a)
            print(f'{i}th sample complete: {time.time() - start:.3f}s')
        except Exception as e:
            print(f'{i}th sample - exception: {e}')
            y_hats.append(None); epsilons.append(None); successful_runs.append(False); num_iter_reweighting_refinements_completed.append(None)
        num_failures = sum([f is None for f in y_hats])
        if num_failures > allowable_failures:
            raise ValueError(f'this batch has failed {num_failures} times. More than the {allowable_failures} allowable failures.')

    print(f'batch completed with {np.sum(successful_runs)}/{len(emp_cov_batch)} completing succesfully ')
    return y_hats, epsilons, successful_runs, num_iter_reweighting_refinements_completed
//  [python]
def get_pca(self):
        cov_symm = (self.cov + self.cov.T) / 2
        eigvals, eigvecs = torch.symeig(cov_symm)
        return eigvals, eigvecs
//  [python]
def get_pca(self):
        cov_symm = (self.cov + self.cov.T) / 2
        eigvals, eigvecs = torch.linalg.eigh(cov_symm)
        return eigvals, eigvecs
//  [python]
def batch_matrix_sqrt_inv(C, epsilon=None, debug=False):
    """ Compute the inverse square root of a batch of matrices.
    Args:
        C: (batch, d, d) Tensor containing a batch of matrices to compute the inverse square root.
        epsilon: (float) Small constant to add to the eigenvalues to avoid numerical instability.
    Returns:
        C_inv_sqrt: (batch, d, d) Tensor containing the inverse square root of the input matrices.
    """
    dim = C.shape[-1]
    # Perform batched eigenvalue decomposition
    eigenvalues, Q = torch.symeig(C, eigenvectors=True)  # C = Q @ Lambda @ Q^T
    cond_num = torch.abs(eigenvalues[:, -1] / eigenvalues[:, 0])
    if epsilon is None:
        epsilon = torch.finfo(eigenvalues.dtype).eps

    Lambda_inv_sqrt = 1.0 / torch.sqrt(eigenvalues)

    # Perform batched matrix multiplication to compute C^(-1/2) = Q @ 1/sqrt(Lambda) @ Q^T
    C_inv_sqrt = Q @ torch.diag_embed(Lambda_inv_sqrt) @ Q.mT

    # Determine if some matrices are defective.
    max_eigval = eigenvalues[:, -1]
    min_eigval = eigenvalues[:, 0]
    relevant_eigvals_mask = eigenvalues > (max_eigval * epsilon).unsqueeze(-1)
    ranks = relevant_eigvals_mask.sum(dim=-1)
    defective_mask = ranks < dim
    # Recompute defective matrix inversion and sqrt by ignoring defective eigenvalues and eigenvectors
    if defective_mask.any():
        # Eigvals are sorted in ascending order, so we can just remove the first `dim-rank` dimensions.
        unique_def_ranks = torch.unique(ranks[defective_mask])
        for rank in unique_def_ranks:
            matrix_mask = defective_mask & (ranks == rank)  # Select defective matrices of rank `rank`
            Q_low_rank = Q[matrix_mask, :, -rank:]  # Select the last `rank` eigenvectors
            eig_low_rank = eigenvalues[matrix_mask, -rank:]  # Select the last `rank` eigenvalues
            Lambda_inv_sqrt_low_rank = 1.0 / torch.sqrt(eig_low_rank)
            C_inv_sqrt[matrix_mask] = Q_low_rank @ torch.diag_embed(Lambda_inv_sqrt_low_rank) @ Q_low_rank.mT

    if debug:
        C_inv_true = torch.linalg.pinv(C, hermitian=True)
        C_inv_pred = C_inv_sqrt @ C_inv_sqrt
        assert torch.allclose(C_inv_true, C_inv_pred, atol=1e-5), \
            f"Matrix norm error {torch.linalg.matrix_norm(C_inv_true - C_inv_pred, ord='fro', dim=(-2, -1))}"

    return C_inv_sqrt, cond_num
//  [python]
def batch_matrix_sqrt_inv(C, epsilon=None, debug=False):
    """ Compute the inverse square root of a batch of matrices.
    Args:
        C: (batch, d, d) Tensor containing a batch of matrices to compute the inverse square root.
        epsilon: (float) Small constant to add to the eigenvalues to avoid numerical instability.
    Returns:
        C_inv_sqrt: (batch, d, d) Tensor containing the inverse square root of the input matrices.
    """
    dim = C.shape[-1]
    # Perform batched eigenvalue decomposition
    eigenvalues, Q = torch.linalg.eigh(C)  # C = Q @ Lambda @ Q^T
    cond_num = torch.abs(eigenvalues[:, -1] / eigenvalues[:, 0])
    if epsilon is None:
        epsilon = torch.finfo(eigenvalues.dtype).eps

    Lambda_inv_sqrt = 1.0 / torch.sqrt(eigenvalues)

    # Perform batched matrix multiplication to compute C^(-1/2) = Q @ 1/sqrt(Lambda) @ Q^T
    C_inv_sqrt = Q @ torch.diag_embed(Lambda_inv_sqrt) @ Q.mT

    # Determine if some matrices are defective.
    max_eigval = eigenvalues[:, -1]
    min_eigval = eigenvalues[:, 0]
    relevant_eigvals_mask = eigenvalues > (max_eigval * epsilon).unsqueeze(-1)
    ranks = relevant_eigvals_mask.sum(dim=-1)
    defective_mask = ranks < dim
    # Recompute defective matrix inversion and sqrt by ignoring defective eigenvalues and eigenvectors
    if defective_mask.any():
        # Eigvals are sorted in ascending order, so we can just remove the first `dim-rank` dimensions.
        unique_def_ranks = torch.unique(ranks[defective_mask])
        for rank in unique_def_ranks:
            matrix_mask = defective_mask & (ranks == rank)  # Select defective matrices of rank `rank`
            Q_low_rank = Q[matrix_mask, :, -rank:]  # Select the last `rank` eigenvectors
            eig_low_rank = eigenvalues[matrix_mask, -rank:]  # Select the last `rank` eigenvalues
            Lambda_inv_sqrt_low_rank = 1.0 / torch.sqrt(eig_low_rank)
            C_inv_sqrt[matrix_mask] = Q_low_rank @ torch.diag_embed(Lambda_inv_sqrt_low_rank) @ Q_low_rank.mT

    if debug:
        C_inv_true = torch.linalg.pinv(C, hermitian=True)
        C_inv_pred = C_inv_sqrt @ C_inv_sqrt
        assert torch.allclose(C_inv_true, C_inv_pred, atol=1e-5), \
            f"Matrix norm error {torch.linalg.matrix_norm(C_inv_true - C_inv_pred, ord='fro', dim=(-2, -1))}"

    return C_inv_sqrt, cond_num
//  [python]
def torch_minimize_kernel(method, objective, init, **kwargs):
    import torch

    kernel_init = init @ init.T * .1
    dim = init.shape[1]

    def kernel_objective(K, *args, **kwargs):
        with torch.no_grad():
            D, U = torch.symeig(K)
            D = D[-dim:].clamp(min=0)
            K[:] = U[:, -dim:].mm(D.diag()).mm(U[:, -dim:].transpose(0, 1))
        return objective(K, *args, **kwargs)

    result = torch_minimize(method, kernel_objective, kernel_init, **kwargs)
    U, s, _ = np.linalg.svd(result.x)

    result.x = U[:, :dim] @ np.sqrt(np.diag(s[:dim]))
    return result
//  [python]
def torch_minimize_kernel(method, objective, init, **kwargs):
    import torch

    kernel_init = init @ init.T * .1
    dim = init.shape[1]

    def kernel_objective(K, *args, **kwargs):
        with torch.no_grad():
            D, U = torch.linalg.eigh(K)
            D = D[-dim:].clamp(min=0)
            K[:] = U[:, -dim:].mm(D.diag()).mm(U[:, -dim:].transpose(0, 1))
        return objective(K, *args, **kwargs)

    result = torch_minimize(method, kernel_objective, kernel_init, **kwargs)
    U, s, _ = np.linalg.svd(result.x)

    result.x = U[:, :dim] @ np.sqrt(np.diag(s[:dim]))
    return result
//  [python]
def weighted_8points(x_in, logits):
    # x_in: batch * 1 * N * 4
    mask = logits[:, 0, :, 0]
    weights = logits[:, 1, :, 0]

    mask = torch.sigmoid(mask)
    weights = torch.exp(weights) * mask
    weights = weights / (torch.sum(weights, dim=-1, keepdim=True) + 1e-5)

    x_shp = x_in.shape
    x_in = x_in.squeeze(1)

    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1).contiguous()

    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1).contiguous()
    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X
    XwX = torch.matmul(X.permute(0, 2, 1).contiguous(), wX)

    # Recover essential matrix from self-adjoing eigen
    e, v = torch.symeig(XwX, eigenvectors=True)

    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))

    # Make unit norm just in case
    e_hat = e_hat / torch.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def weighted_8points(x_in, logits):
    # x_in: batch * 1 * N * 4
    mask = logits[:, 0, :, 0]
    weights = logits[:, 1, :, 0]

    mask = torch.sigmoid(mask)
    weights = torch.exp(weights) * mask
    weights = weights / (torch.sum(weights, dim=-1, keepdim=True) + 1e-5)

    x_shp = x_in.shape
    x_in = x_in.squeeze(1)

    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1).contiguous()

    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1).contiguous()
    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X
    XwX = torch.matmul(X.permute(0, 2, 1).contiguous(), wX)

    # Recover essential matrix from self-adjoing eigen
    _, v = torch.linalg.eigh(XwX)

    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))

    # Make unit norm just in case
    e_hat = e_hat / torch.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def normal_loss(points,k=6, epsilon=0.001):
    points = torch.transpose(points,1,2)
    batch_size = points.size(0)
    points_mean = points.mean(1,keepdim=True) 
    points_centered = points - points_mean
    knn_indices, _ = get_nearest_neighbors_indices_batch(points_centered.detach().cpu().numpy(),points_centered.detach().cpu().numpy(),k)
    knn_indices = torch.LongTensor(knn_indices).to(points.device)
    knn_points = []
    for b in range(batch_size):   
        knn_points.append(points[b,knn_indices[b]].unsqueeze(0))
    
    knn_points = torch.cat(knn_points,0)
    central_diff = knn_points
    per_pt_cov = central_diff.unsqueeze(4) * central_diff.unsqueeze(3)
    covariances = per_pt_cov.mean(2)    
    curvatures, local_coord_frames = torch.symeig(covariances,eigenvectors=True) 
    vertex_normal = local_coord_frames[:, :, :, 0]

    knn_normals = []
    for b in range(batch_size):   
        knn_normals.append(vertex_normal[b,knn_indices[b]].unsqueeze(0))
    
    knn_normals = torch.cat(knn_normals,0)
    loss = l1_loss(knn_normals[:,:,1:,:]-knn_normals[:,:,0,:].unsqueeze(2),torch.zeros_like(knn_normals[:,:,1:,:]))
    return loss
//  [python]
def normal_loss(points,k=6, epsilon=0.001):
    points = torch.transpose(points,1,2)
    batch_size = points.size(0)
    points_mean = points.mean(1,keepdim=True) 
    points_centered = points - points_mean
    knn_indices, _ = get_nearest_neighbors_indices_batch(points_centered.detach().cpu().numpy(),points_centered.detach().cpu().numpy(),k)
    knn_indices = torch.LongTensor(knn_indices).to(points.device)
    knn_points = []
    for b in range(batch_size):   
        knn_points.append(points[b,knn_indices[b]].unsqueeze(0))
    
    knn_points = torch.cat(knn_points,0)
    central_diff = knn_points
    per_pt_cov = central_diff.unsqueeze(4) * central_diff.unsqueeze(3)
    covariances = per_pt_cov.mean(2)    
    curvatures, local_coord_frames = torch.linalg.eigh(covariances) 
    vertex_normal = local_coord_frames[:, :, :, 0]

    knn_normals = []
    for b in range(batch_size):   
        knn_normals.append(vertex_normal[b,knn_indices[b]].unsqueeze(0))
    
    knn_normals = torch.cat(knn_normals,0)
    loss = l1_loss(knn_normals[:,:,1:,:]-knn_normals[:,:,0,:].unsqueeze(2),torch.zeros_like(knn_normals[:,:,1:,:]))
    return loss
//  [python]
def run_batch(wandb, dataloader, threshold=None, regressions=None):
    batch = next(iter(dataloader))
    eps, allowable_failure_rate = wandb.config.eps, wandb.config.allowable_failure_rate
    x, y, _, _, _ = batch

    emp_cov_eigvals, emp_cov_eigvecs = torch.linalg.eigh(x)
    y_hats, epsilons, successful_indices, num_iter_reweighting_refinements_completed = \
        spectral_templates_batch(emp_cov_batch=x, emp_cov_eigvecs_batch=emp_cov_eigvecs.numpy(),
                                 allowable_failures=int(allowable_failure_rate*len(x)))

    # remove runs which did not complete successfully
    y_hats, epsilons, num_iter_reweighting_refinements_completed = \
        list(compress(y_hats, successful_indices)), list(compress(epsilons, successful_indices)), \
        list(compress(num_iter_reweighting_refinements_completed, successful_indices))
    y_hat = np.stack(y_hats, axis=0)
    y = y[successful_indices, :, :]

    y_hat = zero_diagonals(torch.tensor(y_hat))
    metrics, threshold, regressions = find_best_performances(torch.tensor(y), torch.tensor(y_hat), threshold, regressions)

    metrics['failure_rate'] = 1 - len(successful_indices) / len(x)

    print(f'\n### Prediction Stats: using eps {eps:.5f}, threshold {threshold} ###')
    print(f'\tpred edge weights: max: {np.nanmax(y_hat):.3f}, median {np.nanmedian(y_hat):.3f}, mean {np.nanmean(y_hat):.3f}')
    print(f'\tave sparsity of true graphs: {edge_density(y).mean():.3f}')
    print(f'\tave sparsity of (i) raw pred: {edge_density(y_hat).mean():.3f}, (ii) pred > {threshold:.7f}:  {edge_density((y_hat > threshold) + 0.0).mean():.3f}')
    print(f'\tepsilons: mean: {np.nanmean(epsilons):.3f}, std: {np.nanstd(epsilons):.3f} ... epsilons {epsilons}')
    print(f"\tfailure rate: {100*metrics['failure_rate']:.4f}%")
    print(f'### Performance ###')
    print(f"\terrors {metrics['error']*100:.5f}%, using found best threshold: {threshold:.7f}")
    print(f"\tmses:  raw: {metrics['raw_mse']:.5f}, ols: {metrics['ols_mse']:.5f}, scaling_mse: {metrics['ols_no_intercept_mse']:.5f}")
    print('\t# iterative refines', num_iter_reweighting_refinements_completed)
    return metrics, threshold, regressions
//  [python]
def run_batch(wandb, dataloader, threshold=None, regressions=None):
    batch = next(iter(dataloader))
    eps, allowable_failure_rate = wandb.config.eps, wandb.config.allowable_failure_rate
    x, y, _, _, _ = batch

    emp_cov_eigvals, emp_cov_eigvecs = eigh(x)
    y_hats, epsilons, successful_indices, num_iter_reweighting_refinements_completed = \
        spectral_templates_batch(emp_cov_batch=x, emp_cov_eigvecs_batch=emp_cov_eigvecs.numpy(),
                                 allowable_failures=int(allowable_failure_rate*len(x)))

    # remove runs which did not complete successfully
    y_hats, epsilons, num_iter_reweighting_refinements_completed = \
        list(compress(y_hats, successful_indices)), list(compress(epsilons, successful_indices)), \
        list(compress(num_iter_reweighting_refinements_completed, successful_indices))
    y_hat = np.stack(y_hats, axis=0)
    y = y[successful_indices, :, :]

    y_hat = zero_diagonals(torch.tensor(y_hat))
    metrics, threshold, regressions = find_best_performances(torch.tensor(y), torch.tensor(y_hat), threshold, regressions)

    metrics['failure_rate'] = 1 - len(successful_indices) / len(x)

    print(f'\n### Prediction Stats: using eps {eps:.5f}, threshold {threshold} ###')
    print(f'\tpred edge weights: max: {np.nanmax(y_hat):.3f}, median {np.nanmedian(y_hat):.3f}, mean {np.nanmean(y_hat):.3f}')
    print(f'\tave sparsity of true graphs: {edge_density(y).mean():.3f}')
    print(f'\tave sparsity of (i) raw pred: {edge_density(y_hat).mean():.3f}, (ii) pred > {threshold:.7f}:  {edge_density((y_hat > threshold) + 0.0).mean():.3f}')
    print(f'\tepsilons: mean: {np.nanmean(epsilons):.3f}, std: {np.nanstd(epsilons):.3f} ... epsilons {epsilons}')
    print(f"\tfailure rate: {100*metrics['failure_rate']:.4f}%")
    print(f'### Performance ###')
    print(f"\terrors {metrics['error']*100:.5f}%, using found best threshold: {threshold:.7f}")
    print(f"\tmses:  raw: {metrics['raw_mse']:.5f}, ols: {metrics['ols_mse']:.5f}, scaling_mse: {metrics['ols_no_intercept_mse']:.5f}")
    print('\t# iterative refines', num_iter_reweighting_refinements_completed)
    return metrics, threshold, regressions
//  [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.symeig(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)