jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def update(self, minibatches, opt, sch):
        all_x = torch.cat([data[0].cuda().float() for data in minibatches])
        all_y = torch.cat([data[1].cuda().long() for data in minibatches])
        with torch.no_grad():
            all_x1 = torch.angle(torch.fft.fftn(all_x, dim=(2, 3)))
            tfea = self.teab(self.teaf(all_x1)).detach()

        all_z = self.bottleneck(self.featurizer(all_x))
        loss1 = F.cross_entropy(self.classifier(all_z), all_y)

        loss2 = F.mse_loss(all_z[:, :self.tfbd], tfea)*self.args.alpha
        if self.args.disttype == '2-norm':
            loss3 = -F.mse_loss(all_z[:, :self.tfbd],
                                all_z[:, self.tfbd:])*self.args.beta
        elif self.args.disttype == 'norm-2-norm':
            loss3 = -F.mse_loss(all_z[:, :self.tfbd]/torch.linalg.norm(all_z[:, :self.tfbd], dim=1, keepdim=True),
                                all_z[:, self.tfbd:]/torch.linalg.norm(all_z[:, self.tfbd:], dim=1, keepdim=True))*self.args.beta
        elif self.args.disttype == 'norm-1-norm':
            loss3 = -F.l1_loss(all_z[:, :self.tfbd]/torch.linalg.norm(all_z[:, :self.tfbd], dim=1, keepdim=True),
                               all_z[:, self.tfbd:]/torch.linalg.norm(all_z[:, self.tfbd:], dim=1, keepdim=True))*self.args.beta
        elif self.args.disttype == 'cos':
            loss3 = torch.mean(F.cosine_similarity(
                all_z[:, :self.tfbd], all_z[:, self.tfbd:]))*self.args.beta
        loss4 = 0
        if len(minibatches) > 1:
            for i in range(len(minibatches)-1):
                for j in range(i+1, len(minibatches)):
                    loss4 += self.coral(all_z[i*self.args.batch_size:(i+1)*self.args.batch_size, self.tfbd:],
                                        all_z[j*self.args.batch_size:(j+1)*self.args.batch_size, self.tfbd:])
            loss4 = loss4*2/(len(minibatches) *
                             (len(minibatches)-1))*self.args.lam
        else:
            loss4 = self.coral(all_z[:self.args.batch_size//2, self.tfbd:],
                               all_z[self.args.batch_size//2:, self.tfbd:])
            loss4 = loss4*self.args.lam

        loss = loss1+loss2+loss3+loss4
        opt.zero_grad()
        loss.backward()
        opt.step()
        if sch:
            sch.step()
        return {'class': loss1.item(), 'dist': (loss2).item(), 'exp': (loss3).item(), 'align': loss4.item(), 'total': loss.item()}
//  [python]
def get_rays(self):
        grid = create_meshgrid(self.image_height, self.image_width, normalized_coordinates=False)[0] + 0.5
        i, j = grid.unbind(-1)
        pts_view = torch.stack([(i-self.cx)/self.fl_x, (j-self.cy)/self.fl_y, torch.ones_like(i), torch.ones_like(i)], -1).to(self.data_device)
        c2w = torch.linalg.inv(self.world_view_transform.transpose(0, 1))
        pts_world =  pts_view @ c2w.T
        directions = pts_world[...,:3] - self.camera_center[None,None,:]
        return self.camera_center[None,None], directions / torch.norm(directions, dim=-1, keepdim=True)
//  [python]
def get_rays(self):
        grid = create_meshgrid(self.image_height, self.image_width, normalized_coordinates=False)[0] + 0.5
        i, j = grid.unbind(-1)
        pts_view = torch.stack([(i-self.cx)/self.fl_x, (j-self.cy)/self.fl_y, torch.ones_like(i), torch.ones_like(i)], -1).to(self.data_device)
        c2w = torch.linalg.inv(self.world_view_transform.transpose(0, 1))
        pts_world =  pts_view @ c2w.T
        directions = pts_world[...,:3] - self.camera_center[None,None,:]
        return self.camera_center[None,None], directions / torch.linalg.norm(directions, dim=-1, keepdim=True)
//  [python]
def spectrogram(self, x):
        n_fft, hop_length, win_length = self.resolution
        x = F.pad(
            x,
            (int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)),
            mode="reflect",
        )
        x = x.squeeze(1)
        x = torch.stft(
            x, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=False
        )  # [B, F, TT, 2]
        mag = torch.norm(x, p=2, dim=-1)  # [B, F, TT]

        return mag
//  [python]
def spectrogram(self, x):
        n_fft, hop_length, win_length = self.resolution
        x = F.pad(
            x,
            (int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)),
            mode="reflect",
        )
        x = x.squeeze(1)
        x = torch.stft(
            x, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=False
        )  # [B, F, TT, 2]
        mag = torch.linalg.norm(x, p=2, dim=-1)  # [B, F, TT]

        return mag
//  [python]
def group_pixels(center_points, offsets):
    """
    Gives each pixel in the image an instance id.
    Args:
        center_points: A Tensor of shape [K, 2] where K is the number of center points.
            The order of second dim is (y, x).
        offsets: A Tensor of shape [2, H, W] of raw offset output. The order of
            second dim is (offset_y, offset_x).
    Returns:
        A Tensor of shape [1, H, W] with values in range [1, K], which represents
            the center this pixel belongs to.
    """
    height, width = offsets.size()[1:]

    # Generates a coordinate map, where each location is the coordinate of
    # that location.
    y_coord, x_coord = torch.meshgrid(
        torch.arange(height, dtype=offsets.dtype, device=offsets.device),
        torch.arange(width, dtype=offsets.dtype, device=offsets.device),
    )
    coord = torch.cat((y_coord.unsqueeze(0), x_coord.unsqueeze(0)), dim=0)

    center_loc = coord + offsets
    center_loc = center_loc.flatten(1).T.unsqueeze_(0)  # [1, H*W, 2]
    center_points = center_points.unsqueeze(1)  # [K, 1, 2]

    # Distance: [K, H*W].
    distance = torch.norm(center_points - center_loc, dim=-1)

    # Finds center with minimum distance at each location, offset by 1, to
    # reserve id=0 for stuff.
    instance_id = torch.argmin(distance, dim=0).reshape((1, height, width)) + 1
    return instance_id
//  [python]
def group_pixels(center_points, offsets):
    """
    Gives each pixel in the image an instance id.
    Args:
        center_points: A Tensor of shape [K, 2] where K is the number of center points.
            The order of second dim is (y, x).
        offsets: A Tensor of shape [2, H, W] of raw offset output. The order of
            second dim is (offset_y, offset_x).
    Returns:
        A Tensor of shape [1, H, W] with values in range [1, K], which represents
            the center this pixel belongs to.
    """
    height, width = offsets.size()[1:]

    # Generates a coordinate map, where each location is the coordinate of
    # that location.
    y_coord, x_coord = torch.meshgrid(
        torch.arange(height, dtype=offsets.dtype, device=offsets.device),
        torch.arange(width, dtype=offsets.dtype, device=offsets.device),
    )
    coord = torch.cat((y_coord.unsqueeze(0), x_coord.unsqueeze(0)), dim=0)

    center_loc = coord + offsets
    center_loc = center_loc.flatten(1).T.unsqueeze_(0)  # [1, H*W, 2]
    center_points = center_points.unsqueeze(1)  # [K, 1, 2]

    # Distance: [K, H*W].
    distance = torch.linalg.norm(center_points - center_loc, dim=-1)

    # Finds center with minimum distance at each location, offset by 1, to
    # reserve id=0 for stuff.
    instance_id = torch.argmin(distance, dim=0).reshape((1, height, width)) + 1
    return instance_id
//  [python]
def get_perturbed_vectors(input_vectors, max_theta_radians):
    r"""
        middleschool trig rotation in n-dimensional space
        max_theta_radians : defines the maximum angle in radians to rotate input vectors by in a random direction.
    """
    random_angles = torch.rand_like(input_vectors[:,:,:,0], device=input_vectors.device, dtype=input_vectors.dtype) * max_theta_radians
    #The next three lines are an old family recipe for cooking up orthogonal vectors.
    random_vectors = torch.rand_like(input_vectors)
    projections = input_vectors * (torch.sum(random_vectors * input_vectors, dim=-1, keepdim=True) / 
                   torch.sum(input_vectors * input_vectors, dim=-1, keepdim=True))
    orthoshifts = random_vectors - projections
    #Delicious. Now sprinkle generously and knead them back into our base vectors
    target_shifts = (torch.norm(input_vectors, dim=-1, keepdim=True) 
                    * (orthoshifts / torch.norm(orthoshifts, dim=-1, keepdim=True)))
    results = (input_vectors * torch.cos(random_angles).unsqueeze(-1)) + (target_shifts * torch.sin(random_angles).unsqueeze(-1))
    return results
//  [python]
def get_perturbed_vectors(input_vectors, max_theta_radians):
    r"""
        middleschool trig rotation in n-dimensional space
        max_theta_radians : defines the maximum angle in radians to rotate input vectors by in a random direction.
    """
    random_angles = torch.rand_like(input_vectors[:,:,:,0], device=input_vectors.device, dtype=input_vectors.dtype) * max_theta_radians
    #The next three lines are an old family recipe for cooking up orthogonal vectors.
    random_vectors = torch.rand_like(input_vectors)
    projections = input_vectors * (torch.sum(random_vectors * input_vectors, dim=-1, keepdim=True) / 
                   torch.sum(input_vectors * input_vectors, dim=-1, keepdim=True))
    orthoshifts = random_vectors - projections
    #Delicious. Now sprinkle generously and knead them back into our base vectors
    target_shifts = (torch.linalg.norm(input_vectors, dim=-1, keepdim=True) 
                    * (orthoshifts / torch.linalg.norm(orthoshifts, dim=-1, keepdim=True)))
    results = (input_vectors * torch.cos(random_angles).unsqueeze(-1)) + (target_shifts * torch.sin(random_angles).unsqueeze(-1))
    return results
//  [python]
def forward(self, x_mag, y_mag):
        """Calculate forward propagation.
        Args:
            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
        Returns:
            Tensor: Spectral convergence loss value.
        """
        return torch.norm(y_mag - x_mag, p="fro") / torch.norm(y_mag, p="fro")
//  [python]
def forward(self, x_mag, y_mag):
        """Calculate forward propagation.
        Args:
            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
        Returns:
            Tensor: Spectral convergence loss value.
        """
        return torch.linalg.norm(y_mag - x_mag, ord="fro") / torch.linalg.norm(y_mag, ord="fro")
//  [python]
def batch_nn(in_vertex, query_vertex, batch_size = 1024, k = 1):
    topk_value = torch.zeros((len(query_vertex), k)).float().to(in_vertex.device)
    topk_idx = torch.zeros((len(query_vertex), k)).long().to(in_vertex.device)
    for i in range(0, len(query_vertex), batch_size):
        if i + batch_size >= len(query_vertex):
            batch_query_vertex = query_vertex[i:]
        else:
            batch_query_vertex = query_vertex[i: i + batch_size]

        diff = torch.norm(batch_query_vertex.unsqueeze(1) - in_vertex.unsqueeze(0), p = 2, dim = 2)
        tmp_topk_value, tmp_topk_idx = torch.topk(diff, k, dim = 1, largest = False)
        
        if i + batch_size >= len(query_vertex):
            topk_value[i:] = tmp_topk_value
            topk_idx[i:] = tmp_topk_idx
        else:
            topk_value[i: i + batch_size] = tmp_topk_value
            topk_idx[i: i + batch_size] = tmp_topk_idx
    return topk_value, topk_idx
//  [python]
def batch_nn(in_vertex, query_vertex, batch_size = 1024, k = 1):
    topk_value = torch.zeros((len(query_vertex), k)).float().to(in_vertex.device)
    topk_idx = torch.zeros((len(query_vertex), k)).long().to(in_vertex.device)
    for i in range(0, len(query_vertex), batch_size):
        if i + batch_size >= len(query_vertex):
            batch_query_vertex = query_vertex[i:]
        else:
            batch_query_vertex = query_vertex[i: i + batch_size]

        diff = torch.linalg.norm(batch_query_vertex.unsqueeze(1) - in_vertex.unsqueeze(0), ord = 2, dim = 2)
        tmp_topk_value, tmp_topk_idx = torch.topk(diff, k, dim = 1, largest = False)
        
        if i + batch_size >= len(query_vertex):
            topk_value[i:] = tmp_topk_value
            topk_idx[i:] = tmp_topk_idx
        else:
            topk_value[i: i + batch_size] = tmp_topk_value
            topk_idx[i: i + batch_size] = tmp_topk_idx
    return topk_value, topk_idx
//  [python]
def interpolate(xyzs, tris, neighbours, edge_mask):
    # Currently disabled
    return triangle2projection(tris)[0]
    N = xyzs.shape[0]
    factor = 4
    c_closests = calculate_centroid(tris)
    c_neighbours = calculate_centroid(neighbours)
    dc = torch.exp(-1 * torch.norm(xyzs - c_closests[0], dim=1, keepdim=True))
    dn = torch.exp(-factor * torch.norm(xyzs.repeat(1, 3).reshape(N * 3, 3) - c_neighbours[0], dim=1, keepdim=True)) * edge_mask
    distances = torch.cat([dc, dn.reshape(N, -1)], dim=1)
    triangles = torch.cat([triangle2projection(tris)[0][:, None, ...], triangle2projection(neighbours)[0][:, None, ...].reshape(N, -1, 4, 4)], dim=1)
    normalization = distances.sum(-1, keepdim=True)
    weights = distances / normalization

    return (triangles * weights[..., None, None]).sum(1)
//  [python]
def interpolate(xyzs, tris, neighbours, edge_mask):
    # Currently disabled
    return triangle2projection(tris)[0]
    N = xyzs.shape[0]
    factor = 4
    c_closests = calculate_centroid(tris)
    c_neighbours = calculate_centroid(neighbours)
    dc = torch.exp(-1 * torch.linalg.norm(xyzs - c_closests[0], dim=1, keepdim=True))
    dn = torch.exp(-factor * torch.linalg.norm(xyzs.repeat(1, 3).reshape(N * 3, 3) - c_neighbours[0], dim=1, keepdim=True)) * edge_mask
    distances = torch.cat([dc, dn.reshape(N, -1)], dim=1)
    triangles = torch.cat([triangle2projection(tris)[0][:, None, ...], triangle2projection(neighbours)[0][:, None, ...].reshape(N, -1, 4, 4)], dim=1)
    normalization = distances.sum(-1, keepdim=True)
    weights = distances / normalization

    return (triangles * weights[..., None, None]).sum(1)
//  [python]
def latent_blend(settings, a, b, t):
    """
    Interpolates two latent image representations according to the parameter t,
    where the interpolated vectors' magnitudes are also interpolated separately.
    The "detail_preservation" factor biases the magnitude interpolation towards
    the larger of the two magnitudes.
    """
    import torch

    # NOTE: We use inplace operations wherever possible.

    if len(t.shape) == 3:
        # [4][w][h] to [1][4][w][h]
        t2 = t.unsqueeze(0)
        # [4][w][h] to [1][1][w][h] - the [4] seem redundant.
        t3 = t[0].unsqueeze(0).unsqueeze(0)
    else:
        t2 = t
        t3 = t[:, 0][:, None]

    one_minus_t2 = 1 - t2
    one_minus_t3 = 1 - t3

    # Linearly interpolate the image vectors.
    a_scaled = a * one_minus_t2
    b_scaled = b * t2
    image_interp = a_scaled
    image_interp.add_(b_scaled)
    result_type = image_interp.dtype
    del a_scaled, b_scaled, t2, one_minus_t2

    # Calculate the magnitude of the interpolated vectors. (We will remove this magnitude.)
    # 64-bit operations are used here to allow large exponents.
    current_magnitude = torch.norm(image_interp, p=2, dim=1, keepdim=True).to(torch.float64).add_(0.00001)

    # Interpolate the powered magnitudes, then un-power them (bring them back to a power of 1).
    a_magnitude = torch.norm(a, p=2, dim=1, keepdim=True).to(torch.float64).pow_(
        settings.inpaint_detail_preservation) * one_minus_t3
    b_magnitude = torch.norm(b, p=2, dim=1, keepdim=True).to(torch.float64).pow_(
        settings.inpaint_detail_preservation) * t3
    desired_magnitude = a_magnitude
    desired_magnitude.add_(b_magnitude).pow_(1 / settings.inpaint_detail_preservation)
    del a_magnitude, b_magnitude, t3, one_minus_t3

    # Change the linearly interpolated image vectors' magnitudes to the value we want.
    # This is the last 64-bit operation.
    image_interp_scaling_factor = desired_magnitude
    image_interp_scaling_factor.div_(current_magnitude)
    image_interp_scaling_factor = image_interp_scaling_factor.to(result_type)
    image_interp_scaled = image_interp
    image_interp_scaled.mul_(image_interp_scaling_factor)
    del current_magnitude
    del desired_magnitude
    del image_interp
    del image_interp_scaling_factor
    del result_type

    return image_interp_scaled
//  [python]
def latent_blend(settings, a, b, t):
    """
    Interpolates two latent image representations according to the parameter t,
    where the interpolated vectors' magnitudes are also interpolated separately.
    The "detail_preservation" factor biases the magnitude interpolation towards
    the larger of the two magnitudes.
    """
    import torch

    # NOTE: We use inplace operations wherever possible.

    if len(t.shape) == 3:
        # [4][w][h] to [1][4][w][h]
        t2 = t.unsqueeze(0)
        # [4][w][h] to [1][1][w][h] - the [4] seem redundant.
        t3 = t[0].unsqueeze(0).unsqueeze(0)
    else:
        t2 = t
        t3 = t[:, 0][:, None]

    one_minus_t2 = 1 - t2
    one_minus_t3 = 1 - t3

    # Linearly interpolate the image vectors.
    a_scaled = a * one_minus_t2
    b_scaled = b * t2
    image_interp = a_scaled
    image_interp.add_(b_scaled)
    result_type = image_interp.dtype
    del a_scaled, b_scaled, t2, one_minus_t2

    # Calculate the magnitude of the interpolated vectors. (We will remove this magnitude.)
    # 64-bit operations are used here to allow large exponents.
    current_magnitude = torch.linalg.norm(image_interp, ord=2, dim=1, keepdim=True).to(torch.float64).add_(0.00001)

    # Interpolate the powered magnitudes, then un-power them (bring them back to a power of 1).
    a_magnitude = torch.linalg.norm(a, ord=2, dim=1, keepdim=True).to(torch.float64).pow_(
        settings.inpaint_detail_preservation) * one_minus_t3
    b_magnitude = torch.linalg.norm(b, ord=2, dim=1, keepdim=True).to(torch.float64).pow_(
        settings.inpaint_detail_preservation) * t3
    desired_magnitude = a_magnitude
    desired_magnitude.add_(b_magnitude).pow_(1 / settings.inpaint_detail_preservation)
    del a_magnitude, b_magnitude, t3, one_minus_t3

    # Change the linearly interpolated image vectors' magnitudes to the value we want.
    # This is the last 64-bit operation.
    image_interp_scaling_factor = desired_magnitude
    image_interp_scaling_factor.div_(current_magnitude)
    image_interp_scaling_factor = image_interp_scaling_factor.to(result_type)
    image_interp_scaled = image_interp
    image_interp_scaled.mul_(image_interp_scaling_factor)
    del current_magnitude
    del desired_magnitude
    del image_interp
    del image_interp_scaling_factor
    del result_type

    return image_interp_scaled
//  [python]
def forward(self, data: Data) -> Data:
        assert data.pos is not None
        assert data.edge_index is not None
        (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
        assert pos.dim() == 2 and pos.size(1) == 3

        cart = pos[col] - pos[row]

        rho = torch.norm(cart, p=2, dim=-1).view(-1, 1)

        theta = torch.atan2(cart[..., 1], cart[..., 0]).view(-1, 1)
        theta = theta + (theta < 0).type_as(theta) * (2 * PI)

        phi = torch.acos(cart[..., 2] / rho.view(-1)).view(-1, 1)

        if self.norm:
            rho = rho / (rho.max() if self.max is None else self.max)
            theta = theta / (2 * PI)
            phi = phi / PI

        spher = torch.cat([rho, theta, phi], dim=-1)

        if pseudo is not None and self.cat:
            pseudo = pseudo.view(-1, 1) if pseudo.dim() == 1 else pseudo
            data.edge_attr = torch.cat([pseudo, spher.type_as(pos)], dim=-1)
        else:
            data.edge_attr = spher

        return data
//  [python]
def forward(self, data: Data) -> Data:
        assert data.pos is not None
        assert data.edge_index is not None
        (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
        assert pos.dim() == 2 and pos.size(1) == 3

        cart = pos[col] - pos[row]

        rho = torch.linalg.norm(cart, ord=2, dim=-1).view(-1, 1)

        theta = torch.atan2(cart[..., 1], cart[..., 0]).view(-1, 1)
        theta = theta + (theta < 0).type_as(theta) * (2 * PI)

        phi = torch.acos(cart[..., 2] / rho.view(-1)).view(-1, 1)

        if self.norm:
            rho = rho / (rho.max() if self.max is None else self.max)
            theta = theta / (2 * PI)
            phi = phi / PI

        spher = torch.cat([rho, theta, phi], dim=-1)

        if pseudo is not None and self.cat:
            pseudo = pseudo.view(-1, 1) if pseudo.dim() == 1 else pseudo
            data.edge_attr = torch.cat([pseudo, spher.type_as(pos)], dim=-1)
        else:
            data.edge_attr = spher

        return data
//  [python]
def cal_angel(self, points, angel_k=3):
        r"""

        :param points: (B,N,2)
        :param angel_k: number of nearest neighbors
        :return: angle: the angel of each correspondence (B,N,N)
        """
        batch_size = points.shape[0]
        num_point = points.shape[1]
        points_dist = torch.norm((points[:, :, None, :] - points[:, None, :, :]), dim=-1)
        knn_indices = points_dist.topk(k=angel_k + 1, dim=2, largest=False)[1][:, :,
                      1:]  # (B,N,k)  k+1 : get out the itself indice
        # print(knn_indices)
        knn_indices = knn_indices.unsqueeze(3).expand(batch_size, num_point, angel_k, 2)  # (B,N,K,2)
        expanded_points = points.unsqueeze(1).expand(batch_size, num_point, num_point, 2)  # (B,N,K,2)
        knn_points = torch.gather(expanded_points, dim=2, index=knn_indices)  # (B,N,k,2)

        ref_vectors = knn_points - points.unsqueeze(2)  # (B, N, k, 2)
        anc_vectors = points.unsqueeze(1) - points.unsqueeze(2)  # (B, N, N, 2)
        ref_vectors = ref_vectors.unsqueeze(2).expand(batch_size, num_point, num_point, angel_k, 2)  # (B, N, N, k, 2)
        anc_vectors = anc_vectors.unsqueeze(3).expand(batch_size, num_point, num_point, angel_k, 2)  # (B, N, N, k, 2)
        zeros = torch.zeros((batch_size, num_point, num_point, angel_k, 1), device=ref_vectors.device)
        ref_vectors = torch.cat((ref_vectors, zeros), dim=-1)
        anc_vectors = torch.cat((anc_vectors, zeros), dim=-1)
        sin_values = torch.linalg.norm(torch.cross(ref_vectors, anc_vectors, dim=-1), dim=-1)  # (B, N, N, k)
        cos_values = torch.sum(ref_vectors * anc_vectors, dim=-1)  # (B, N, N, k)
        angles = torch.atan2(sin_values, cos_values)  # (B, N, N, k)
        angle, _ = torch.max(angles, dim=-1)
        return angle
//  [python]
def cal_angel(self, points, angel_k=3):
        r"""

        :param points: (B,N,2)
        :param angel_k: number of nearest neighbors
        :return: angle: the angel of each correspondence (B,N,N)
        """
        batch_size = points.shape[0]
        num_point = points.shape[1]
        points_dist = torch.linalg.norm((points[:, :, None, :] - points[:, None, :, :]), dim=-1)
        knn_indices = points_dist.topk(k=angel_k + 1, dim=2, largest=False)[1][:, :,
                      1:]  # (B,N,k)  k+1 : get out the itself indice
        # print(knn_indices)
        knn_indices = knn_indices.unsqueeze(3).expand(batch_size, num_point, angel_k, 2)  # (B,N,K,2)
        expanded_points = points.unsqueeze(1).expand(batch_size, num_point, num_point, 2)  # (B,N,K,2)
        knn_points = torch.gather(expanded_points, dim=2, index=knn_indices)  # (B,N,k,2)

        ref_vectors = knn_points - points.unsqueeze(2)  # (B, N, k, 2)
        anc_vectors = points.unsqueeze(1) - points.unsqueeze(2)  # (B, N, N, 2)
        ref_vectors = ref_vectors.unsqueeze(2).expand(batch_size, num_point, num_point, angel_k, 2)  # (B, N, N, k, 2)
        anc_vectors = anc_vectors.unsqueeze(3).expand(batch_size, num_point, num_point, angel_k, 2)  # (B, N, N, k, 2)
        zeros = torch.zeros((batch_size, num_point, num_point, angel_k, 1), device=ref_vectors.device)
        ref_vectors = torch.cat((ref_vectors, zeros), dim=-1)
        anc_vectors = torch.cat((anc_vectors, zeros), dim=-1)
        sin_values = torch.linalg.norm(torch.cross(ref_vectors, anc_vectors, dim=-1), dim=-1)  # (B, N, N, k)
        cos_values = torch.sum(ref_vectors * anc_vectors, dim=-1)  # (B, N, N, k)
        angles = torch.atan2(sin_values, cos_values)  # (B, N, N, k)
        angle, _ = torch.max(angles, dim=-1)
        return angle
//  [python]
def homo_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 3, 3]
            - src_keypts:    [bs, num_corr, 2]
            - tgt_keypts:    [bs, num_corr, 2]
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 3, 3]
        """
        assert initial_trans.shape[0] == 1
        if self.inlier_threshold == 2:  # for 3DMatch
            inlier_threshold_list = [4] * 5
        else:  # for KITTI
            inlier_threshold_list = [4] * 5

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:

            pred_position = torch.einsum('bnm,bmk->bnk', initial_trans[:, :2, :2],
                                         src_keypts.permute(0, 2, 1)) + initial_trans[:, :2, 2:3]  # [bs, num_corr, 3]

            pred_position_z = torch.einsum('bnm,bmk->bnk', initial_trans[:, 2:3, :2],
                                           src_keypts.permute(0, 2, 1)) + initial_trans[:, 2:3,
                                                                          2:3]  # [bs, num_corr, 3]
            pred_position = pred_position / pred_position_z
            pred_position = pred_position.permute(0, 2, 1)
            L2_dis = torch.norm(pred_position - tgt_keypts[:, :, :], dim=-1)  # [bs, num_corr]

            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                # update
                break
            else:
                previous_inlier_num = inlier_num
                initial_trans = kornia.geometry.homography.find_homography_dlt(src_keypts[:, pred_inlier, :],
                                                                               tgt_keypts[:, pred_inlier, :],
                                                                               weights=weights[:, pred_inlier])
        return initial_trans
//  [python]
def homo_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 3, 3]
            - src_keypts:    [bs, num_corr, 2]
            - tgt_keypts:    [bs, num_corr, 2]
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 3, 3]
        """
        assert initial_trans.shape[0] == 1
        if self.inlier_threshold == 2:  # for 3DMatch
            inlier_threshold_list = [4] * 5
        else:  # for KITTI
            inlier_threshold_list = [4] * 5

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:

            pred_position = torch.einsum('bnm,bmk->bnk', initial_trans[:, :2, :2],
                                         src_keypts.permute(0, 2, 1)) + initial_trans[:, :2, 2:3]  # [bs, num_corr, 3]

            pred_position_z = torch.einsum('bnm,bmk->bnk', initial_trans[:, 2:3, :2],
                                           src_keypts.permute(0, 2, 1)) + initial_trans[:, 2:3,
                                                                          2:3]  # [bs, num_corr, 3]
            pred_position = pred_position / pred_position_z
            pred_position = pred_position.permute(0, 2, 1)
            L2_dis = torch.linalg.norm(pred_position - tgt_keypts[:, :, :], dim=-1)  # [bs, num_corr]

            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                # update
                break
            else:
                previous_inlier_num = inlier_num
                initial_trans = kornia.geometry.homography.find_homography_dlt(src_keypts[:, pred_inlier, :],
                                                                               tgt_keypts[:, pred_inlier, :],
                                                                               weights=weights[:, pred_inlier])
        return initial_trans
//  [python]
def cal_trans(self, src_keypts, tgt_keypts):
        corr_pos = torch.cat((src_keypts, tgt_keypts), dim=-1)  # [bs, num_corr, 4]
        bs, num_corr = corr_pos.shape[0], corr_pos.shape[1]
        k = num_corr

        #################################
        # construct the spatial consistency matrix
        #################################
        src_dist = torch.norm((src_keypts[:, :, None, :] - src_keypts[:, None, :, :]), dim=-1)
        corr_compatibility = src_dist - torch.norm((tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :]),
                                                   dim=-1)
        corr_compatibility = torch.clamp(1.0 - corr_compatibility ** 2 / self.sigma_spat ** 2, min=0)  # bs,L,l

        #################################
        # Power iteratation to get the inlier probability
        #################################
        corr_compatibility[:, torch.arange(corr_compatibility.shape[1]), torch.arange(corr_compatibility.shape[1])] = 0

        total_weight = self.cal_leading_eigenvector(corr_compatibility, method='power')

        total_weight = total_weight.view([bs, k])
        total_weight = total_weight / (torch.sum(total_weight, dim=-1, keepdim=True) + 1e-6)
        #################################
        # calculate the transformation by weighted least-squares for each subsets in parallel
        #################################
        total_weight = total_weight.view([-1, k])

        src_knn, tgt_knn = src_keypts.view([-1, k, 2]), tgt_keypts.view([-1, k, 2])
        seed_as_center = False

        if seed_as_center:
            assert ("Not codes!")
        else:
            # not use seeds as neighborhood centers.
            seedwise_trans = self.rigid_transform_2d(src_knn, tgt_knn, total_weight)
            seedwise_trans = seedwise_trans.view([bs, 3, 3])

        #################################
        # calculate the inlier number for each hypothesis, and find the best transformation for each point cloud pair
        #################################
        pred_position = torch.einsum('bnm,bmk->bnk', seedwise_trans[:, :2, :2],
                                     src_keypts.permute(0, 2, 1)) + seedwise_trans[:, :2, 2:3]  # [bs, num_corr, 3]
        pred_position = pred_position.permute(0, 2, 1)
        L2_dis = torch.norm(pred_position - tgt_keypts[:, :, :], dim=-1)  # [bs, num_corr]
        seedwise_fitness = torch.mean((L2_dis < self.inlier_threshold).float(), dim=-1)  # [bs, 1]
        # seedwise_inlier_rmse = torch.sum(L2_dis * (L2_dis < config.inlier_threshold).float(), dim=1)
        batch_best_guess = seedwise_fitness

        # refine the pose by using all the inlier correspondences (done in the post-refinement step)
        final_trans = seedwise_trans
        final_labels = (L2_dis < self.inlier_threshold).float()

        return final_trans, final_labels
//  [python]
def cal_trans(self, src_keypts, tgt_keypts):
        corr_pos = torch.cat((src_keypts, tgt_keypts), dim=-1)  # [bs, num_corr, 4]
        bs, num_corr = corr_pos.shape[0], corr_pos.shape[1]
        k = num_corr

        #################################
        # construct the spatial consistency matrix
        #################################
        src_dist = torch.linalg.norm((src_keypts[:, :, None, :] - src_keypts[:, None, :, :]), dim=-1)
        corr_compatibility = src_dist - torch.linalg.norm((tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :]),
                                                   dim=-1)
        corr_compatibility = torch.clamp(1.0 - corr_compatibility ** 2 / self.sigma_spat ** 2, min=0)  # bs,L,l

        #################################
        # Power iteratation to get the inlier probability
        #################################
        corr_compatibility[:, torch.arange(corr_compatibility.shape[1]), torch.arange(corr_compatibility.shape[1])] = 0

        total_weight = self.cal_leading_eigenvector(corr_compatibility, method='power')

        total_weight = total_weight.view([bs, k])
        total_weight = total_weight / (torch.sum(total_weight, dim=-1, keepdim=True) + 1e-6)
        #################################
        # calculate the transformation by weighted least-squares for each subsets in parallel
        #################################
        total_weight = total_weight.view([-1, k])

        src_knn, tgt_knn = src_keypts.view([-1, k, 2]), tgt_keypts.view([-1, k, 2])
        seed_as_center = False

        if seed_as_center:
            assert ("Not codes!")
        else:
            # not use seeds as neighborhood centers.
            seedwise_trans = self.rigid_transform_2d(src_knn, tgt_knn, total_weight)
            seedwise_trans = seedwise_trans.view([bs, 3, 3])

        #################################
        # calculate the inlier number for each hypothesis, and find the best transformation for each point cloud pair
        #################################
        pred_position = torch.einsum('bnm,bmk->bnk', seedwise_trans[:, :2, :2],
                                     src_keypts.permute(0, 2, 1)) + seedwise_trans[:, :2, 2:3]  # [bs, num_corr, 3]
        pred_position = pred_position.permute(0, 2, 1)
        L2_dis = torch.linalg.norm(pred_position - tgt_keypts[:, :, :], dim=-1)  # [bs, num_corr]
        seedwise_fitness = torch.mean((L2_dis < self.inlier_threshold).float(), dim=-1)  # [bs, 1]
        # seedwise_inlier_rmse = torch.sum(L2_dis * (L2_dis < config.inlier_threshold).float(), dim=1)
        batch_best_guess = seedwise_fitness

        # refine the pose by using all the inlier correspondences (done in the post-refinement step)
        final_trans = seedwise_trans
        final_labels = (L2_dis < self.inlier_threshold).float()

        return final_trans, final_labels
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 3, 3]
            - src_keypts:    [bs, num_corr, 2]
            - tgt_keypts:    [bs, num_corr, 2]
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 3, 3]
        """
        assert initial_trans.shape[0] == 1
        if self.inlier_threshold == 5:  # for 3DMatch
            inlier_threshold_list = [5] * 20
        else:  # for KITTI
            inlier_threshold_list = [8] * 20

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = self.transform(src_keypts, initial_trans)
            L2_dis = torch.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                # update
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = self.rigid_transform_2d(
                A=src_keypts[:, pred_inlier, :],
                B=tgt_keypts[:, pred_inlier, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier],
                # weights=((1-L2_dis/inlier_threshold)**2)[:, pred_inlier]
            )
        return initial_trans
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 3, 3]
            - src_keypts:    [bs, num_corr, 2]
            - tgt_keypts:    [bs, num_corr, 2]
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 3, 3]
        """
        assert initial_trans.shape[0] == 1
        if self.inlier_threshold == 5:  # for 3DMatch
            inlier_threshold_list = [5] * 20
        else:  # for KITTI
            inlier_threshold_list = [8] * 20

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = self.transform(src_keypts, initial_trans)
            L2_dis = torch.linalg.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                # update
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = self.rigid_transform_2d(
                A=src_keypts[:, pred_inlier, :],
                B=tgt_keypts[:, pred_inlier, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier],
                # weights=((1-L2_dis/inlier_threshold)**2)[:, pred_inlier]
            )
        return initial_trans
//  [python]
def cal_trans_weight(self, src_keypts, tgt_keypts, weight):
        # weight:[bs,L]
        corr_pos = torch.cat((src_keypts, tgt_keypts), dim=-1)  # [bs, num_corr, 4]
        bs, num_corr = corr_pos.shape[0], corr_pos.shape[1]
        k = num_corr

        #################################
        # construct the spatial consistency matrix
        #################################
        src_dist = torch.norm((src_keypts[:, :, None, :] - src_keypts[:, None, :, :]), dim=-1)
        corr_compatibility = src_dist - torch.norm((tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :]),
                                                   dim=-1)
        corr_compatibility = torch.clamp(1.0 - corr_compatibility ** 2 / self.sigma_spat ** 2, min=0)  # bs,L,l

        #################################
        # Power iteratation to get the inlier probability
        #################################
        corr_compatibility[:, torch.arange(corr_compatibility.shape[1]), torch.arange(corr_compatibility.shape[1])] = 0

        total_weight = self.cal_leading_eigenvector(corr_compatibility, method='power')

        total_weight = total_weight.view([bs, k]) * weight
        total_weight = total_weight / (torch.sum(total_weight, dim=-1, keepdim=True) + 1e-6)
        #################################
        # calculate the transformation by weighted least-squares for each subsets in parallel
        #################################
        total_weight = total_weight.view([-1, k])

        src_knn, tgt_knn = src_keypts.view([-1, k, 2]), tgt_keypts.view([-1, k, 2])
        seed_as_center = False

        if seed_as_center:
            assert ("Not codes!")
        else:
            # not use seeds as neighborhood centers.
            seedwise_trans = self.rigid_transform_2d(src_knn, tgt_knn, total_weight)
            seedwise_trans = seedwise_trans.view([bs, 3, 3])

        #################################
        # calculate the inlier number for each hypothesis, and find the best transformation for each point cloud pair
        #################################
        pred_position = torch.einsum('bnm,bmk->bnk', seedwise_trans[:, :2, :2],
                                     src_keypts.permute(0, 2, 1)) + seedwise_trans[:, :2, 2:3]  # [bs, num_corr, 3]
        pred_position_z = torch.einsum('bnm,bmk->bnk', seedwise_trans[:, 2:3, :2],
                                       src_keypts.permute(0, 2, 1)) + seedwise_trans[:, 2:3, 2:3]  # [bs, num_corr, 3]
        pred_position = pred_position / pred_position_z

        pred_position = pred_position.permute(0, 2, 1)
        L2_dis = torch.norm(pred_position - tgt_keypts[:, :, :], dim=-1)  # [bs, num_corr]
        seedwise_fitness = torch.mean((L2_dis < self.inlier_threshold).float(), dim=-1)  # [bs, 1]
        # seedwise_inlier_rmse = torch.sum(L2_dis * (L2_dis < config.inlier_threshold).float(), dim=1)
        batch_best_guess = seedwise_fitness

        # refine the pose by using all the inlier correspondences (done in the post-refinement step)
        final_trans = seedwise_trans
        final_labels = (L2_dis < self.inlier_threshold).float()

        return final_trans, final_labels
//  [python]
def poincare_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.PoincareBall]
    ex = torch.randn(*shape, dtype=torch.float64) / 3
    ev = torch.randn(*shape, dtype=torch.float64) / 3
    x = torch.tanh(torch.norm(ex)) * ex / torch.norm(ex)
    ex = x.clone()
    v = ev.clone()
    manifold = geoopt.PoincareBall().to(dtype=torch.float64)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.PoincareBallExact().to(dtype=torch.float64)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def poincare_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.PoincareBall]
    ex = torch.randn(*shape, dtype=torch.float64) / 3
    ev = torch.randn(*shape, dtype=torch.float64) / 3
    x = torch.tanh(torch.linalg.norm(ex)) * ex / torch.linalg.norm(ex)
    ex = x.clone()
    v = ev.clone()
    manifold = geoopt.PoincareBall().to(dtype=torch.float64)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.PoincareBallExact().to(dtype=torch.float64)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def sphere_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.Sphere]
    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    x = ex / torch.norm(ex)
    v = ev - (x @ ev) * x

    manifold = geoopt.Sphere()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.SphereExact()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def sphere_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.Sphere]
    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    x = ex / torch.linalg.norm(ex)
    v = ev - (x @ ev) * x

    manifold = geoopt.Sphere()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.SphereExact()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def training_step(self, batch, batch_idx):
        batch.pos = batch.positions

        emb_2d = batch.emb_2d
        # emb_3d = self.encoder_3d(batch, True)
        _, emb_3d = self.encoder_3d(batch.x, batch.pos, batch.radius_edge_index, batch.batch, True)

        # spatial graphs
        batch.edge_index = batch.radius_edge_index
        edge_attr = torch.exp( - torch.norm(batch.pos[batch.edge_index[0]] - batch.pos[batch.edge_index[1]], dim=1) )
        batch.edge_attr = torch.einsum('i,j->ij', edge_attr, torch.linspace(1, 5, 16).to(edge_attr.device))
        batch.n_edges = batch.n_edges_radius

        atom_pred_list, extra_output = self.decoder_2dto3d(batch, emb_3d)

        gt_poss = tgeom.utils.unbatch(batch.pos, batch.batch)
        atom_pred_list = tgeom.utils.unbatch(atom_pred_list[-1], batch.batch)

        loss = 0
        for pred_pos, gt_pos in zip(atom_pred_list, gt_poss):
            pred_pos = self.kbash_alignment(pred_pos, gt_pos)
            _loss = (pred_pos - gt_pos).pow(2).mean()
            loss += _loss

            # print(_loss)
            # pdb.set_trace()
        loss /= len(atom_pred_list)

        # avoiding collapse
        emb_3d_norm = emb_3d.norm(dim=1)
        sim_matrix = torch.einsum('ik,jk->ij', emb_3d, emb_3d) / torch.einsum('i,j->ij', emb_3d_norm, emb_3d_norm)
        sim_matrix = torch.exp(sim_matrix)
        pos_sim = sim_matrix[range(emb_3d.shape[0]), range(emb_3d.shape[0])]
        loss_neg_contrast = torch.log( (sim_matrix.sum(dim=1) - pos_sim) ).mean()
        loss += loss_neg_contrast * 0.1

        # loss, loss_dict = self.decoder_2dto3d.compute_loss(atom_pred_list, extra_output, batch)
        self.log('train_loss', loss, batch_size=batch.num_graphs, on_epoch=True, sync_dist=True)

        return loss
//  [python]
def training_step(self, batch, batch_idx):
        batch.pos = batch.positions

        emb_2d = batch.emb_2d
        # emb_3d = self.encoder_3d(batch, True)
        _, emb_3d = self.encoder_3d(batch.x, batch.pos, batch.radius_edge_index, batch.batch, True)

        # spatial graphs
        batch.edge_index = batch.radius_edge_index
        edge_attr = torch.exp( - torch.linalg.norm(batch.pos[batch.edge_index[0]] - batch.pos[batch.edge_index[1]], dim=1) )
        batch.edge_attr = torch.einsum('i,j->ij', edge_attr, torch.linspace(1, 5, 16).to(edge_attr.device))
        batch.n_edges = batch.n_edges_radius

        atom_pred_list, extra_output = self.decoder_2dto3d(batch, emb_3d)

        gt_poss = tgeom.utils.unbatch(batch.pos, batch.batch)
        atom_pred_list = tgeom.utils.unbatch(atom_pred_list[-1], batch.batch)

        loss = 0
        for pred_pos, gt_pos in zip(atom_pred_list, gt_poss):
            pred_pos = self.kbash_alignment(pred_pos, gt_pos)
            _loss = (pred_pos - gt_pos).pow(2).mean()
            loss += _loss

            # print(_loss)
            # pdb.set_trace()
        loss /= len(atom_pred_list)

        # avoiding collapse
        emb_3d_norm = emb_3d.norm(dim=1)
        sim_matrix = torch.einsum('ik,jk->ij', emb_3d, emb_3d) / torch.einsum('i,j->ij', emb_3d_norm, emb_3d_norm)
        sim_matrix = torch.exp(sim_matrix)
        pos_sim = sim_matrix[range(emb_3d.shape[0]), range(emb_3d.shape[0])]
        loss_neg_contrast = torch.log( (sim_matrix.sum(dim=1) - pos_sim) ).mean()
        loss += loss_neg_contrast * 0.1

        # loss, loss_dict = self.decoder_2dto3d.compute_loss(atom_pred_list, extra_output, batch)
        self.log('train_loss', loss, batch_size=batch.num_graphs, on_epoch=True, sync_dist=True)

        return loss
//  [python]
def sphere_subspace_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.Sphere]
    subspace = torch.rand(shape[-1], 2, dtype=torch.float64)

    Q, _ = geoopt.linalg.batch_linalg.qr(subspace, "reduced")
    P = Q @ Q.t()

    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    x = (ex @ P.t()) / torch.norm(ex @ P.t())
    v = (ev - (x @ ev) * x) @ P.t()

    manifold = geoopt.Sphere(intersection=subspace)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.SphereExact(intersection=subspace)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def sphere_subspace_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.Sphere]
    subspace = torch.rand(shape[-1], 2, dtype=torch.float64)

    Q, _ = geoopt.linalg.batch_linalg.qr(subspace, "reduced")
    P = Q @ Q.t()

    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    x = (ex @ P.t()) / torch.linalg.norm(ex @ P.t())
    v = (ev - (x @ ev) * x) @ P.t()

    manifold = geoopt.Sphere(intersection=subspace)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.SphereExact(intersection=subspace)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case