jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_compatibility(src_keypts, tgt_keypts, with_dist_corr_compatibility, dist_sigma=None):
    # [n,3], [n,3]
    #################################
    # Step1: corr_compatibily
    #################################
    src_dist = torch.linalg.norm((src_keypts.unsqueeze(-2) - src_keypts.unsqueeze(-3)), dim=-1) # [n1, n1]
    tgt_dist = torch.linalg.norm((tgt_keypts.unsqueeze(-2) - tgt_keypts.unsqueeze(-3)), dim=-1) # [n1, n1]
    
    corr_compatibility = torch.min(torch.stack((src_dist/(tgt_dist+1e-12),tgt_dist/(src_dist+1e-12)), dim=-1), dim=-1).values + torch.eye(src_dist.shape[-1]).to(src_dist)
    corr_compatibility = torch.clamp((corr_compatibility / 0.9)**2, min=0)
    
    #################################
    # Step2: dist_compatibily
    #################################
    if with_dist_corr_compatibility:
        dist_corr_compatibility = (src_dist + tgt_dist)/2
        dist_corr_compatibility = torch.clamp(1.0 - dist_corr_compatibility ** 2 / dist_sigma ** 2, min=0) # [n1, n1]
        corr_compatibility *= dist_corr_compatibility
        corr_compatibility /= corr_compatibility.max(dim=-1, keepdim=True).values.max(dim=-1, keepdim=True).values
    return corr_compatibility
//  [python]
def get_quaternion(v0, v1):
    # v0: N, D (normalized)
    # v1: N, D (normalized)

    cross = torch.cross(v0, v1)  # N, 3, cross = sin_theta * unit_vecter
    axis = F.normalize(cross, dim=-1)
    sin_theta = torch.norm(cross, dim=-1) 
    cos_theta = torch.bmm(v0.unsqueeze(-2), v1.unsqueeze(-1)).squeeze() # 900, # COS(theta): cosine_similarity [-1, 1 ]
    theta = torch.atan(sin_theta / cos_theta)
    # print((sin_theta / cos_theta).min(), (sin_theta / cos_theta).max())
    
    # print(theta, theta.min(), theta.max())
    # print(torch.arcsin(sin_theta), torch.arcsin(sin_theta).min(), torch.arcsin(sin_theta).max())
    # print(torch.arccos(cos_theta), torch.arccos(cos_theta).min(), torch.arccos(cos_theta).max())
    
    sin_theta2 = torch.sin(theta/2)
    cos_theta2 = torch.cos(theta/2)
    
    # quarternion = torch.cat([axis, theta.unsqueeze(-1)], dim=-1) # 900, 4
    quarternion = torch.cat([cos_theta2.unsqueeze(-1), axis * sin_theta2.unsqueeze(-1)], dim=-1)
 
    return quarternion
//  [python]
def get_quaternion(v0, v1):
    # v0: N, D (normalized)
    # v1: N, D (normalized)

    cross = torch.cross(v0, v1)  # N, 3, cross = sin_theta * unit_vecter
    axis = F.normalize(cross, dim=-1)
    sin_theta = torch.linalg.norm(cross, dim=-1) 
    cos_theta = torch.bmm(v0.unsqueeze(-2), v1.unsqueeze(-1)).squeeze() # 900, # COS(theta): cosine_similarity [-1, 1 ]
    theta = torch.atan(sin_theta / cos_theta)
    # print((sin_theta / cos_theta).min(), (sin_theta / cos_theta).max())
    
    # print(theta, theta.min(), theta.max())
    # print(torch.arcsin(sin_theta), torch.arcsin(sin_theta).min(), torch.arcsin(sin_theta).max())
    # print(torch.arccos(cos_theta), torch.arccos(cos_theta).min(), torch.arccos(cos_theta).max())
    
    sin_theta2 = torch.sin(theta/2)
    cos_theta2 = torch.cos(theta/2)
    
    # quarternion = torch.cat([axis, theta.unsqueeze(-1)], dim=-1) # 900, 4
    quarternion = torch.cat([cos_theta2.unsqueeze(-1), axis * sin_theta2.unsqueeze(-1)], dim=-1)
 
    return quarternion
//  [python]
def sample_inputs(input_types, batch_size, dtype, rng):
    dev = "cuda:0" if torch.cuda.is_available() else "cpu"
    if isinstance(batch_size, int):
        batch_size = (batch_size,)

    def _sample(input_type):
        type_str, param = input_type

        def _quat_sample():
            q = torch.rand(*batch_size, param, device=dev, dtype=dtype, generator=rng)
            return q / torch.norm(q, dim=-1, keepdim=True)

        sample_fns = {
            "tangent": lambda: torch.rand(
                *batch_size, param, device=dev, dtype=dtype, generator=rng
            ),
            "group": lambda: param.rand(
                *batch_size, device=dev, generator=rng, dtype=dtype
            ),
            "quat": lambda: _quat_sample(),
            "matrix": lambda: torch.rand(
                (*batch_size,) + param, device=dev, generator=rng, dtype=dtype
            ),
        }
        return sample_fns[type_str]()

    return tuple(_sample(type_str) for type_str in input_types)
//  [python]
def sample_inputs(input_types, batch_size, dtype, rng):
    dev = "cuda:0" if torch.cuda.is_available() else "cpu"
    if isinstance(batch_size, int):
        batch_size = (batch_size,)

    def _sample(input_type):
        type_str, param = input_type

        def _quat_sample():
            q = torch.rand(*batch_size, param, device=dev, dtype=dtype, generator=rng)
            return q / torch.linalg.norm(q, dim=-1, keepdim=True)

        sample_fns = {
            "tangent": lambda: torch.rand(
                *batch_size, param, device=dev, dtype=dtype, generator=rng
            ),
            "group": lambda: param.rand(
                *batch_size, device=dev, generator=rng, dtype=dtype
            ),
            "quat": lambda: _quat_sample(),
            "matrix": lambda: torch.rand(
                (*batch_size,) + param, device=dev, generator=rng, dtype=dtype
            ),
        }
        return sample_fns[type_str]()

    return tuple(_sample(type_str) for type_str in input_types)
//  [python]
def pc_normalize_torch(point):
    #point:[n,3]
    assert len(point.size()) == 2
    assert point.size(1) == 3
    # normalize the point and face
    with torch.no_grad():
        avg = torch.mean(point.t(), dim = 1)
        normed_point = point - avg[np.newaxis, :]
        scale = torch.max(torch.norm(normed_point, dim = 1), dim = 0)[0]
        normed_point = normed_point / scale[np.newaxis, np.newaxis]
    return normed_point
//  [python]
def pc_normalize_torch(point):
    #point:[n,3]
    assert len(point.size()) == 2
    assert point.size(1) == 3
    # normalize the point and face
    with torch.no_grad():
        avg = torch.mean(point.t(), dim = 1)
        normed_point = point - avg[np.newaxis, :]
        scale = torch.max(torch.linalg.norm(normed_point, dim = 1), dim = 0)[0]
        normed_point = normed_point / scale[np.newaxis, np.newaxis]
    return normed_point
//  [python]
def _safe_norm(v):
    """Compute the norm of a vector, checking for finite values."""
    if not torch.isfinite(v).all():
        return torch.tensor(float("inf"), dtype=v.dtype, device=v.device)
    return torch.norm(v)
//  [python]
def _safe_norm(v):
    """Compute the norm of a vector, checking for finite values."""
    if not torch.isfinite(v).all():
        return torch.tensor(float("inf"), dtype=v.dtype, device=v.device)
    return torch.linalg.norm(v)
//  [python]
def farthest_points_sample(obj_points, num_points):
    assert obj_points.size(1) == 3
    b,_,n = obj_points.size()

    selected = torch.randint(obj_points.size(2), [obj_points.size(0),1]).cuda()
    dists = torch.full([obj_points.size(0), obj_points.size(2)], fill_value = np.inf).cuda()

    for _ in range(num_points - 1):
        dists = torch.min(dists, torch.norm(obj_points - torch.gather(obj_points, 2, selected[:,-1].unsqueeze(1).unsqueeze(2).expand(b,3,1)), dim = 1))
        selected = torch.cat([selected, torch.argmax(dists, dim=1, keepdim=True)], dim = 1)
    res_points = torch.gather(obj_points, 2, selected.unsqueeze(1).expand(b, 3, num_points))

    return res_points
//  [python]
def farthest_points_sample(obj_points, num_points):
    assert obj_points.size(1) == 3
    b,_,n = obj_points.size()

    selected = torch.randint(obj_points.size(2), [obj_points.size(0),1]).cuda()
    dists = torch.full([obj_points.size(0), obj_points.size(2)], fill_value = np.inf).cuda()

    for _ in range(num_points - 1):
        dists = torch.min(dists, torch.linalg.norm(obj_points - torch.gather(obj_points, 2, selected[:,-1].unsqueeze(1).unsqueeze(2).expand(b,3,1)), dim = 1))
        selected = torch.cat([selected, torch.argmax(dists, dim=1, keepdim=True)], dim = 1)
    res_points = torch.gather(obj_points, 2, selected.unsqueeze(1).expand(b, 3, num_points))

    return res_points
//  [python]
def __init__(self, args):

        HVAE.__init__(self, args)
        # defines the Neural net to compute the metric

        # first layer
        self.metric_fc1 = nn.Linear(self.input_dim*self.n_channels, args.metric_fc)

        # diagonal
        self.metric_fc21 = nn.Linear(args.metric_fc, self.latent_dim)
        # remaining coefficients
        k = int(self.latent_dim * (self.latent_dim - 1) / 2)
        self.metric_fc22 = nn.Linear(args.metric_fc, k)

        self.T = nn.Parameter(torch.Tensor([args.temperature]), requires_grad=False)
        self.lbd = nn.Parameter(
            torch.Tensor([args.regularization]), requires_grad=False
        )

        # this is used to store the matrices and centroids throughout trainning for
        # further use in metric update (L is the cholesky decomposition of M)
        self.M = []
        self.centroids = []

        # define a starting metric (gamma_i = 0 & L = I_d)
        def G(z):
            return (
                torch.eye(self.latent_dim, device=self.device).unsqueeze(0)
                * torch.exp(-torch.norm(z.unsqueeze(1), dim=-1) ** 2)
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(self.device)

        self.G = G
//  [python]
def __init__(self, args):

        HVAE.__init__(self, args)
        # defines the Neural net to compute the metric

        # first layer
        self.metric_fc1 = nn.Linear(self.input_dim*self.n_channels, args.metric_fc)

        # diagonal
        self.metric_fc21 = nn.Linear(args.metric_fc, self.latent_dim)
        # remaining coefficients
        k = int(self.latent_dim * (self.latent_dim - 1) / 2)
        self.metric_fc22 = nn.Linear(args.metric_fc, k)

        self.T = nn.Parameter(torch.Tensor([args.temperature]), requires_grad=False)
        self.lbd = nn.Parameter(
            torch.Tensor([args.regularization]), requires_grad=False
        )

        # this is used to store the matrices and centroids throughout trainning for
        # further use in metric update (L is the cholesky decomposition of M)
        self.M = []
        self.centroids = []

        # define a starting metric (gamma_i = 0 & L = I_d)
        def G(z):
            return (
                torch.eye(self.latent_dim, device=self.device).unsqueeze(0)
                * torch.exp(-torch.linalg.norm(z.unsqueeze(1), dim=-1) ** 2)
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(self.device)

        self.G = G
//  [python]
def forward(self, z: Tensor, pos: Tensor,
                batch: OptTensor = None) -> Tensor:
        r"""Forward pass.

        Args:
            z (torch.Tensor): Atomic number of each atom with shape
                :obj:`[num_atoms]`.
            pos (torch.Tensor): Coordinates of each atom with shape
                :obj:`[num_atoms, 3]`.
            batch (torch.Tensor, optional): Batch indices assigning each atom
                to a separate molecule with shape :obj:`[num_atoms]`.
                (default: :obj:`None`)
        """
        batch = torch.zeros_like(z) if batch is None else batch

        h = self.embedding(z)
        edge_index, edge_weight = self.interaction_graph(pos, batch)
        edge_attr = self.distance_expansion(edge_weight)

        for interaction in self.interactions:
            h = h + interaction(h, edge_index, edge_weight, edge_attr)

        h = self.lin1(h)
        h = self.act(h)
        h = self.lin2(h)

        if self.dipole:
            # Get center of mass.
            mass = self.atomic_mass[z].view(-1, 1)
            M = self.sum_aggr(mass, batch, dim=0)
            c = self.sum_aggr(mass * pos, batch, dim=0) / M
            h = h * (pos - c.index_select(0, batch))

        if not self.dipole and self.mean is not None and self.std is not None:
            h = h * self.std + self.mean

        if not self.dipole and self.atomref is not None:
            h = h + self.atomref(z)

        out = self.readout(h, batch, dim=0)

        if self.dipole:
            out = torch.norm(out, dim=-1, keepdim=True)

        if self.scale is not None:
            out = self.scale * out

        return out
//  [python]
def forward(self, z: Tensor, pos: Tensor,
                batch: OptTensor = None) -> Tensor:
        r"""Forward pass.

        Args:
            z (torch.Tensor): Atomic number of each atom with shape
                :obj:`[num_atoms]`.
            pos (torch.Tensor): Coordinates of each atom with shape
                :obj:`[num_atoms, 3]`.
            batch (torch.Tensor, optional): Batch indices assigning each atom
                to a separate molecule with shape :obj:`[num_atoms]`.
                (default: :obj:`None`)
        """
        batch = torch.zeros_like(z) if batch is None else batch

        h = self.embedding(z)
        edge_index, edge_weight = self.interaction_graph(pos, batch)
        edge_attr = self.distance_expansion(edge_weight)

        for interaction in self.interactions:
            h = h + interaction(h, edge_index, edge_weight, edge_attr)

        h = self.lin1(h)
        h = self.act(h)
        h = self.lin2(h)

        if self.dipole:
            # Get center of mass.
            mass = self.atomic_mass[z].view(-1, 1)
            M = self.sum_aggr(mass, batch, dim=0)
            c = self.sum_aggr(mass * pos, batch, dim=0) / M
            h = h * (pos - c.index_select(0, batch))

        if not self.dipole and self.mean is not None and self.std is not None:
            h = h * self.std + self.mean

        if not self.dipole and self.atomref is not None:
            h = h + self.atomref(z)

        out = self.readout(h, batch, dim=0)

        if self.dipole:
            out = torch.linalg.norm(out, dim=-1, keepdim=True)

        if self.scale is not None:
            out = self.scale * out

        return out
//  [python]
def farthest_points_normal_sample(obj_points, obj_normal, num_points):
    assert obj_points.size(1) == 3
    assert obj_points.size(2) == obj_normal.size(2)
    b,_,n = obj_points.size()

    selected = torch.randint(obj_points.size(2), [obj_points.size(0),1]).cuda()
    dists = torch.full([obj_points.size(0), obj_points.size(2)], fill_value = np.inf).cuda()

    for _ in range(num_points - 1):
        dists = torch.min(dists, torch.norm(obj_points - torch.gather(obj_points, 2, selected[:,-1].unsqueeze(1).unsqueeze(2).expand(b,3,1)), dim = 1))
        selected = torch.cat([selected, torch.argmax(dists, dim=1, keepdim=True)], dim = 1)
    res_points = torch.gather(obj_points, 2, selected.unsqueeze(1).expand(b, 3, num_points))
    res_normal = torch.gather(obj_normal, 2, selected.unsqueeze(1).expand(b, 3, num_points))

    return res_points, res_normal
//  [python]
def farthest_points_normal_sample(obj_points, obj_normal, num_points):
    assert obj_points.size(1) == 3
    assert obj_points.size(2) == obj_normal.size(2)
    b,_,n = obj_points.size()

    selected = torch.randint(obj_points.size(2), [obj_points.size(0),1]).cuda()
    dists = torch.full([obj_points.size(0), obj_points.size(2)], fill_value = np.inf).cuda()

    for _ in range(num_points - 1):
        dists = torch.min(dists, torch.linalg.norm(obj_points - torch.gather(obj_points, 2, selected[:,-1].unsqueeze(1).unsqueeze(2).expand(b,3,1)), dim = 1))
        selected = torch.cat([selected, torch.argmax(dists, dim=1, keepdim=True)], dim = 1)
    res_points = torch.gather(obj_points, 2, selected.unsqueeze(1).expand(b, 3, num_points))
    res_normal = torch.gather(obj_normal, 2, selected.unsqueeze(1).expand(b, 3, num_points))

    return res_points, res_normal
//  [python]
def update_metric(self):
        """
        As soon as the model has seen all the data points (i.e. at the end of 1 loop)
        we update the final metric function using \mu(x_i) as centroids
        """
        # convert to 1 big tensor
        self.M_tens = torch.cat(self.M)
        self.centroids_tens = torch.cat(self.centroids)

        # define new metric
        def G(z):
            return torch.inverse(
                (
                    self.M_tens.unsqueeze(0)
                    * torch.exp(
                        -torch.norm(
                            self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                        )
                        ** 2
                        / (self.T ** 2)
                    )
                    .unsqueeze(-1)
                    .unsqueeze(-1)
                ).sum(dim=1)
                + self.lbd * torch.eye(self.latent_dim).to(self.device)
            )

        def G_inv(z):
            return (
                self.M_tens.unsqueeze(0)
                * torch.exp(
                    -torch.norm(
                        self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                    )
                    ** 2
                    / (self.T ** 2)
                )
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(self.device)

        self.G = G
        self.G_inv = G_inv
        self.M = []
        self.centroids = []
//  [python]
def update_metric(self):
        """
        As soon as the model has seen all the data points (i.e. at the end of 1 loop)
        we update the final metric function using \mu(x_i) as centroids
        """
        # convert to 1 big tensor
        self.M_tens = torch.cat(self.M)
        self.centroids_tens = torch.cat(self.centroids)

        # define new metric
        def G(z):
            return torch.inverse(
                (
                    self.M_tens.unsqueeze(0)
                    * torch.exp(
                        -torch.linalg.norm(
                            self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                        )
                        ** 2
                        / (self.T ** 2)
                    )
                    .unsqueeze(-1)
                    .unsqueeze(-1)
                ).sum(dim=1)
                + self.lbd * torch.eye(self.latent_dim).to(self.device)
            )

        def G_inv(z):
            return (
                self.M_tens.unsqueeze(0)
                * torch.exp(
                    -torch.linalg.norm(
                        self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                    )
                    ** 2
                    / (self.T ** 2)
                )
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(self.device)

        self.G = G
        self.G_inv = G_inv
        self.M = []
        self.centroids = []
//  [python]
def clip_grad_norm_(tensor, max_norm):
    grad_norm = item(torch.norm(tensor))
    if grad_norm > max_norm > 0:
        clip_coef = max_norm / (grad_norm + 1e-6)
        tensor.mul_(clip_coef)
    return grad_norm
//  [python]
def clip_grad_norm_(tensor, max_norm):
    grad_norm = item(torch.linalg.norm(tensor))
    if grad_norm > max_norm > 0:
        clip_coef = max_norm / (grad_norm + 1e-6)
        tensor.mul_(clip_coef)
    return grad_norm
//  [python]
def z2polar(edges):
    z = edges.dst["xy"] - edges.src["xy"]
    rho = th.norm(z, dim=-1, p=2)
    x, y = z.unbind(dim=-1)
    phi = th.atan2(y, x)
    return {"u": th.cat([rho.unsqueeze(-1), phi.unsqueeze(-1)], -1)}
//  [python]
def z2polar(edges):
    z = edges.dst["xy"] - edges.src["xy"]
    rho = th.linalg.norm(z, dim=-1, ord=2)
    x, y = z.unbind(dim=-1)
    phi = th.atan2(y, x)
    return {"u": th.cat([rho.unsqueeze(-1), phi.unsqueeze(-1)], -1)}
//  [python]
def __getitem__(self, index):
        # Separate filename and text
        if self.n_speakers > 1:
            audiopath, *extra, text, speaker = self.audiopaths_and_text[index]
            speaker = int(speaker)
        else:
            audiopath, *extra, text = self.audiopaths_and_text[index]
            speaker = None

        mel = self.get_mel(audiopath)
        text = self.get_text(text)
        pitch = self.get_pitch(index, mel.size(-1))
        energy = torch.norm(mel.float(), dim=0, p=2)
        attn_prior = self.get_prior(index, mel.shape[1], text.shape[0])

        assert pitch.size(-1) == mel.size(-1)

        # No higher formants?
        if len(pitch.size()) == 1:
            pitch = pitch[None, :]

        return (text, mel, len(text), pitch, energy, speaker, attn_prior,
                audiopath)
//  [python]
def __getitem__(self, index):
        # Separate filename and text
        if self.n_speakers > 1:
            audiopath, *extra, text, speaker = self.audiopaths_and_text[index]
            speaker = int(speaker)
        else:
            audiopath, *extra, text = self.audiopaths_and_text[index]
            speaker = None

        mel = self.get_mel(audiopath)
        text = self.get_text(text)
        pitch = self.get_pitch(index, mel.size(-1))
        energy = torch.linalg.norm(mel.float(), dim=0, ord=2)
        attn_prior = self.get_prior(index, mel.shape[1], text.shape[0])

        assert pitch.size(-1) == mel.size(-1)

        # No higher formants?
        if len(pitch.size()) == 1:
            pitch = pitch[None, :]

        return (text, mel, len(text), pitch, energy, speaker, attn_prior,
                audiopath)
//  [python]
def _forward_one_pair(
        self, embedder: nn.Module, mesh_name_1: str, mesh_name_2: str
    ) -> torch.Tensor:
        """
        Do a forward pass with a selected pair of meshes
        Args:
            embedder (nn.Module): module that computes vertex embeddings for different meshes
            mesh_name_1 (str): first mesh name
            mesh_name_2 (str): second mesh name
        Return:
            Tensor containing the loss value
        """
        embeddings_1, geodists_1 = self._get_embeddings_and_geodists_for_mesh(embedder, mesh_name_1)
        embeddings_2, geodists_2 = self._get_embeddings_and_geodists_for_mesh(embedder, mesh_name_2)
        sim_matrix_12 = embeddings_1.mm(embeddings_2.T)

        c_12 = F.softmax(sim_matrix_12 / self.temperature, dim=1)
        c_21 = F.softmax(sim_matrix_12.T / self.temperature, dim=1)
        c_11 = c_12.mm(c_21)
        c_22 = c_21.mm(c_12)

        loss_cycle_11 = torch.norm(geodists_1 * c_11, p=self.norm_p)
        loss_cycle_22 = torch.norm(geodists_2 * c_22, p=self.norm_p)

        return loss_cycle_11 + loss_cycle_22
//  [python]
def _forward_one_pair(
        self, embedder: nn.Module, mesh_name_1: str, mesh_name_2: str
    ) -> torch.Tensor:
        """
        Do a forward pass with a selected pair of meshes
        Args:
            embedder (nn.Module): module that computes vertex embeddings for different meshes
            mesh_name_1 (str): first mesh name
            mesh_name_2 (str): second mesh name
        Return:
            Tensor containing the loss value
        """
        embeddings_1, geodists_1 = self._get_embeddings_and_geodists_for_mesh(embedder, mesh_name_1)
        embeddings_2, geodists_2 = self._get_embeddings_and_geodists_for_mesh(embedder, mesh_name_2)
        sim_matrix_12 = embeddings_1.mm(embeddings_2.T)

        c_12 = F.softmax(sim_matrix_12 / self.temperature, dim=1)
        c_21 = F.softmax(sim_matrix_12.T / self.temperature, dim=1)
        c_11 = c_12.mm(c_21)
        c_22 = c_21.mm(c_12)

        loss_cycle_11 = torch.linalg.norm(geodists_1 * c_11, ord=self.norm_p)
        loss_cycle_22 = torch.linalg.norm(geodists_2 * c_22, ord=self.norm_p)

        return loss_cycle_11 + loss_cycle_22
//  [python]
def _sample(self):
        gan = self.gan

        self.pos = self.direction * self.velocity + self.pos# * self.mask + (1-self.mask) * self.origin
        self.gan.latent.z = self.pos
        self.steps += 1
        if self.steps % 60 == 0:
            self.direction = -self.direction
        if (self.steps - 30) % 180 == 0:
            self.index+=1
            print("Index=",self.index)
            self.direction = self.eigvec[:, self.index].unsqueeze(0)
            self.direction = self.direction / torch.norm(self.direction)

        g = gan.generator.forward(self.pos)
        return [
            ('generator', g)
        ]
//  [python]
def _sample(self):
        gan = self.gan

        self.pos = self.direction * self.velocity + self.pos# * self.mask + (1-self.mask) * self.origin
        self.gan.latent.z = self.pos
        self.steps += 1
        if self.steps % 60 == 0:
            self.direction = -self.direction
        if (self.steps - 30) % 180 == 0:
            self.index+=1
            print("Index=",self.index)
            self.direction = self.eigvec[:, self.index].unsqueeze(0)
            self.direction = self.direction / torch.linalg.norm(self.direction)

        g = gan.generator.forward(self.pos)
        return [
            ('generator', g)
        ]
//  [python]
def _update_metrics(self):
        # time for which the command was executed
        max_command_time = self.cfg.resampling_time_range[1]
        max_command_step = max_command_time / self._env.step_dt
        # logs data
        self.metrics["error_vel_xy"] += (
            torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_lin_vel_b[:, :2], dim=-1) / max_command_step
        )
        self.metrics["error_vel_yaw"] += (
            torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b[:, 2]) / max_command_step
        )
//  [python]
def _update_metrics(self):
        # time for which the command was executed
        max_command_time = self.cfg.resampling_time_range[1]
        max_command_step = max_command_time / self._env.step_dt
        # logs data
        self.metrics["error_vel_xy"] += (
            torch.linalg.norm(self.vel_command_b[:, :2] - self.robot.data.root_lin_vel_b[:, :2], dim=-1) / max_command_step
        )
        self.metrics["error_vel_yaw"] += (
            torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b[:, 2]) / max_command_step
        )
//  [python]
def backward(ctx, grad_output):
        M, M_sqrt_inv, I = ctx.saved_tensors
        M_inv = M_sqrt_inv.bmm(M_sqrt_inv)
        grad_lya = - M_inv.bmm(grad_output).bmm(M_inv)
        norm_sqrt_inv = torch.norm(M_sqrt_inv)
        b = M_sqrt_inv / norm_sqrt_inv
        c = grad_lya / norm_sqrt_inv
        for i in range(8):
            # In case you might terminate the iteration by checking convergence
            # if th.norm(b-I)<1e-4:
            #    break
            b_2 = b.bmm(b)
            c = 0.5 * (c.bmm(3.0 * I - b_2) - b_2.bmm(c) + b.bmm(c).bmm(b))
            b = 0.5 * b.bmm(3.0 * I - b_2)
        grad_input = 0.5 * c
        return grad_input
//  [python]
def backward(ctx, grad_output):
        M, M_sqrt_inv, I = ctx.saved_tensors
        M_inv = M_sqrt_inv.bmm(M_sqrt_inv)
        grad_lya = - M_inv.bmm(grad_output).bmm(M_inv)
        norm_sqrt_inv = torch.linalg.norm(M_sqrt_inv)
        b = M_sqrt_inv / norm_sqrt_inv
        c = grad_lya / norm_sqrt_inv
        for i in range(8):
            # In case you might terminate the iteration by checking convergence
            # if th.norm(b-I)<1e-4:
            #    break
            b_2 = b.bmm(b)
            c = 0.5 * (c.bmm(3.0 * I - b_2) - b_2.bmm(c) + b.bmm(c).bmm(b))
            b = 0.5 * b.bmm(3.0 * I - b_2)
        grad_input = 0.5 * c
        return grad_input
//  [python]
def compute_twist_ratation(rotation_matrix, twist_axis):
    '''
    Compute the twist component of given rotation and twist axis
    https://stackoverflow.com/questions/3684269/component-of-a-quaternion-rotation-around-an-axis
    Parameters
    ----------
    rotation_matrix : Tensor (B, 3, 3,)
        The rotation to convert
    twist_axis : Tensor (B, 3,)
        The twist axis
    Returns
    -------
    Tensor (B, 3, 3)
        The twist rotation
    '''
    quaternion = rotation_matrix_to_quaternion(rotation_matrix)

    twist_axis = twist_axis / (torch.norm(twist_axis, dim=1, keepdim=True) + 1e-9)

    projection = torch.einsum('bi,bi->b', twist_axis, quaternion[:, 1:]).unsqueeze(-1) * twist_axis

    twist_quaternion = torch.cat([quaternion[:, 0:1], projection], dim=1)
    twist_quaternion = twist_quaternion / (torch.norm(twist_quaternion, dim=1, keepdim=True) + 1e-9)

    twist_rotation = quaternion_to_rotation_matrix(twist_quaternion)

    # twist_angle0 = quaternion_to_angle(twist_quaternion)

    twist_aa = quaternion_to_angle_axis(twist_quaternion)

    # twist_rotation2 = batch_rodrigues(twist_aa)

    # print('twist_rotation2', torch.dist(twist_rotation2, twist_rotation))

    twist_angle = torch.sum(twist_aa, dim=1, keepdim=True) / torch.sum(twist_axis, dim=1, keepdim=True)

    # print('norm', torch.dist(torch.norm(twist_axis * twist_angle, dim=1), torch.norm(twist_aa, dim=1)))
    # exit()


    # i = 1
    # for i in range(len(twist_rotation2)):
    #     print('twist aa', twist_angle[i], torch.norm(twist_axis * twist_angle, dim=1)[i], torch.norm(twist_aa, dim=1)[i], torch.dist(batch_rodrigues(twist_axis * twist_angle)[i], batch_rodrigues(twist_aa)[i]))
    #     # print('twist aa', batch_rodrigues(-twist_axis * twist_angle)[i], batch_rodrigues(twist_aa)[i])
    # exit()

    return twist_rotation, twist_angle
//  [python]
def compute_twist_ratation(rotation_matrix, twist_axis):
    '''
    Compute the twist component of given rotation and twist axis
    https://stackoverflow.com/questions/3684269/component-of-a-quaternion-rotation-around-an-axis
    Parameters
    ----------
    rotation_matrix : Tensor (B, 3, 3,)
        The rotation to convert
    twist_axis : Tensor (B, 3,)
        The twist axis
    Returns
    -------
    Tensor (B, 3, 3)
        The twist rotation
    '''
    quaternion = rotation_matrix_to_quaternion(rotation_matrix)

    twist_axis = twist_axis / (torch.linalg.norm(twist_axis, dim=1, keepdim=True) + 1e-9)

    projection = torch.einsum('bi,bi->b', twist_axis, quaternion[:, 1:]).unsqueeze(-1) * twist_axis

    twist_quaternion = torch.cat([quaternion[:, 0:1], projection], dim=1)
    twist_quaternion = twist_quaternion / (torch.linalg.norm(twist_quaternion, dim=1, keepdim=True) + 1e-9)

    twist_rotation = quaternion_to_rotation_matrix(twist_quaternion)

    # twist_angle0 = quaternion_to_angle(twist_quaternion)

    twist_aa = quaternion_to_angle_axis(twist_quaternion)

    # twist_rotation2 = batch_rodrigues(twist_aa)

    # print('twist_rotation2', torch.dist(twist_rotation2, twist_rotation))

    twist_angle = torch.sum(twist_aa, dim=1, keepdim=True) / torch.sum(twist_axis, dim=1, keepdim=True)

    # print('norm', torch.dist(torch.linalg.norm(twist_axis * twist_angle, dim=1), torch.linalg.norm(twist_aa, dim=1)))
    # exit()


    # i = 1
    # for i in range(len(twist_rotation2)):
    #     print('twist aa', twist_angle[i], torch.linalg.norm(twist_axis * twist_angle, dim=1)[i], torch.linalg.norm(twist_aa, dim=1)[i], torch.dist(batch_rodrigues(twist_axis * twist_angle)[i], batch_rodrigues(twist_aa)[i]))
    #     # print('twist aa', batch_rodrigues(-twist_axis * twist_angle)[i], batch_rodrigues(twist_aa)[i])
    # exit()

    return twist_rotation, twist_angle
//  [python]
def calculate_l1_norm(f):
    f_norm = torch.norm(f, p=2, dim=-1, keepdim=True)
    f = f / (f_norm + 1e-9)
    return f
//  [python]
def calculate_l1_norm(f):
    f_norm = torch.linalg.norm(f, ord=2, dim=-1, keepdim=True)
    f = f / (f_norm + 1e-9)
    return f
//  [python]
def eigs_AD(mat, lr=1e-2, v0=None, it_time=500, tol=1e-15):
    """
    :param mat: 带分解矩阵
    :param lr: 初始学习率
    :param v0: 初始向量
    :param it_time: 最大迭代次数
    :param tol: 收敛阈值
    :return lm: 最大本征值
    :return v0: 最大本征向量
    """
    # 初始化及预处理v0
    if v0 is None:
        v0 = tc.randn(mat.shape[0], )
    v0 /= v0.norm()
    v0.requires_grad = True
    # 建立优化器，这里使用Adam优化器，可自动动态控制学习率
    optimizer = tc.optim.Adam([v0], lr=lr)
    v1 = copy.deepcopy(v0.data)
    # 进入主循环
    for t in range(it_time):
        # 计算损失函数
        f = -v0.matmul(mat).matmul(v0) / v0.matmul(v0)
        f.backward()  # 进行反向传播
        optimizer.step()  # 更新参数
        optimizer.zero_grad()  # 清空梯度
        conv = tc.norm(v0.data - v1) / v1.numel()  # 计算收敛性
        if conv < tol:
            break
        else:
            v1 = copy.deepcopy(v0.data)
    with tc.no_grad():
        v0 = v0 / tc.norm(v0)  # 归一化，计算本征向量
        lm = v0.matmul(mat).max() / v0.max()  # 计算本征值
    return lm, v0
//  [python]
def eigs_AD(mat, lr=1e-2, v0=None, it_time=500, tol=1e-15):
    """
    :param mat: 带分解矩阵
    :param lr: 初始学习率
    :param v0: 初始向量
    :param it_time: 最大迭代次数
    :param tol: 收敛阈值
    :return lm: 最大本征值
    :return v0: 最大本征向量
    """
    # 初始化及预处理v0
    if v0 is None:
        v0 = tc.randn(mat.shape[0], )
    v0 /= tc.linalg.norm(v0)
    v0.requires_grad = True
    # 建立优化器，这里使用Adam优化器，可自动动态控制学习率
    optimizer = tc.optim.Adam([v0], lr=lr)
    v1 = copy.deepcopy(v0.data)
    # 进入主循环
    for t in range(it_time):
        # 计算损失函数
        f = -v0.matmul(mat).matmul(v0) / v0.matmul(v0)
        f.backward()  # 进行反向传播
        optimizer.step()  # 更新参数
        optimizer.zero_grad()  # 清空梯度
        conv = tc.linalg.norm(v0.data - v1) / v1.numel()  # 计算收敛性
        if conv < tol:
            break
        else:
            v1 = copy.deepcopy(v0.data)
    with tc.no_grad():
        v0 = v0 / tc.linalg.norm(v0)  # 归一化，计算本征向量
        lm = v0.matmul(mat).max() / v0.max()  # 计算本征值
    return lm, v0
//  [python]
def g(F):
    FF = F.t() @ F
    FF_norm = torch.norm(FF, p='fro')
    return (1/(FF_norm+epsilon_F)) * FF
//  [python]
def g(F):
    FF = F.t() @ F
    FF_norm = torch.linalg.norm(FF, ord='fro')
    return (1/(FF_norm+epsilon_F)) * FF
//  [python]
def calc_distance_vertically(_pc_scene, predicted_quads):

    # TODO: Use Tensor instead of np.ndarray

    pc_scene = _pc_scene.cuda()
    pc_center = torch.mean(pc_scene, dim=0)  # To find the inner side of the point clouds
    distance = 10.0 * torch.ones((pc_scene.shape[0],), dtype=torch.float).cuda()

    idx_array = torch.ones((pc_scene.shape[0],), dtype=torch.float).cuda()

    # Calculate distances
    for predict_idx, _predicted_quad in enumerate(predicted_quads):
        if isinstance(_predicted_quad, np.ndarray):
            _predicted_quad = torch.tensor(_predicted_quad)

        predicted_quad = _predicted_quad.cuda()
        # TODO: convert params predicted_quad to torch.Tensor before and after nms to form a complete compute map
        quad_center = torch.mean(predicted_quad, dim=0)
        predicted_quad_norm = torch.cross(predicted_quad[1] - predicted_quad[0], predicted_quad[2] - predicted_quad[0])
        predicted_quad_norm = predicted_quad_norm / torch.norm(predicted_quad_norm)

        if torch.dot(pc_center - quad_center, predicted_quad_norm) > 0:
            predicted_quad_norm = -predicted_quad_norm  # Make inner distance < 0 and outsider > 0

        vertical_distance = (pc_scene - quad_center) @ predicted_quad_norm
        vertical_distance = vertical_distance.float()
        # vertical_distance = torch.bmm((pc_scene - quad_center).view(num_points, 1, 3),
        #                               predicted_quad_norm.repeat(num_points, 1).view(num_points, 3, 1)) \
        #     .view(num_points, )

        # Use Dynamic programming to find the "nearest" quad with minimum absolute error :)
        mask = torch.abs(vertical_distance) < torch.abs(distance)
        distance[mask] = vertical_distance[mask]
        idx_array[mask] = predict_idx

    return distance, idx_array
//  [python]
def calc_distance_vertically(_pc_scene, predicted_quads):

    # TODO: Use Tensor instead of np.ndarray

    pc_scene = _pc_scene.cuda()
    pc_center = torch.mean(pc_scene, dim=0)  # To find the inner side of the point clouds
    distance = 10.0 * torch.ones((pc_scene.shape[0],), dtype=torch.float).cuda()

    idx_array = torch.ones((pc_scene.shape[0],), dtype=torch.float).cuda()

    # Calculate distances
    for predict_idx, _predicted_quad in enumerate(predicted_quads):
        if isinstance(_predicted_quad, np.ndarray):
            _predicted_quad = torch.tensor(_predicted_quad)

        predicted_quad = _predicted_quad.cuda()
        # TODO: convert params predicted_quad to torch.Tensor before and after nms to form a complete compute map
        quad_center = torch.mean(predicted_quad, dim=0)
        predicted_quad_norm = torch.cross(predicted_quad[1] - predicted_quad[0], predicted_quad[2] - predicted_quad[0])
        predicted_quad_norm = predicted_quad_norm / torch.linalg.norm(predicted_quad_norm)

        if torch.dot(pc_center - quad_center, predicted_quad_norm) > 0:
            predicted_quad_norm = -predicted_quad_norm  # Make inner distance < 0 and outsider > 0

        vertical_distance = (pc_scene - quad_center) @ predicted_quad_norm
        vertical_distance = vertical_distance.float()
        # vertical_distance = torch.bmm((pc_scene - quad_center).view(num_points, 1, 3),
        #                               predicted_quad_norm.repeat(num_points, 1).view(num_points, 3, 1)) \
        #     .view(num_points, )

        # Use Dynamic programming to find the "nearest" quad with minimum absolute error :)
        mask = torch.abs(vertical_distance) < torch.abs(distance)
        distance[mask] = vertical_distance[mask]
        idx_array[mask] = predict_idx

    return distance, idx_array
//  [python]
def get_params(model, id_data, dev="cuda"):
    """Calculate params for vim and SIRC on training set."""

    print(f"eval on: {id_data.name} training set")

    # this is being done on the training set this time
    labels, logits, features = get_outputs_labels(
        model, id_data.train_loader, dev=dev
    )

    # get final fc layer 
    if config["model"]["model_type"] == "resnet101v2": # special case
        W = model.head.conv.weight.cpu().detach().squeeze()
        b = model.head.conv.bias.cpu().detach().squeeze()
    else:
        try:
            W = model.state_dict()["fc.weight"].detach().clone().cpu()
            b = model.state_dict()["fc.bias"].detach().clone().cpu()
        except:
            W = model.state_dict()["classifier.weight"].detach().clone().cpu()
            b = model.state_dict()["classifier.bias"].detach().clone().cpu()

    u = -torch.linalg.pinv(W) @ b

    # size of subspace
    # pretty much just a heuristic
    D = 1000 if features.shape[-1] > 1500 else 512 
    if features.shape[-1] < 512:
        D = features.shape[-1]//2
    centered_feats = features - u
    U = torch.linalg.eigh(
        centered_feats.T@centered_feats
    ).eigenvectors.flip(-1)
    R = U[:,D:] # eigenvectors in columns
    assert R.shape[0] == features.shape[-1]
    vlogits = torch.norm(
        (R.T @ centered_feats.unsqueeze(dim=-1)).squeeze(dim=-1), 
        p=2, dim=-1
    )
    alpha = logits.max(dim=-1).values.mean() / vlogits.mean()
    alpha = alpha.item()

    vim_params_dict = {
        "alpha": alpha,
        "u":  u,
        "R": R
    }

    centered_feats = F.normalize(centered_feats, dim=-1)
    U = torch.linalg.eigh(
        centered_feats.T@centered_feats
    ).eigenvectors.flip(-1) # rev order to descending eigenvalue
    R = U[:,D:] # eigenvectors in columns
    assert R.shape[0] == features.shape[-1]
    vlogits = torch.norm(
        (R.T @ centered_feats.unsqueeze(dim=-1)).squeeze(dim=-1), 
        p=2, dim=-1
    )
    alpha = logits.max(dim=-1).values.mean() / vlogits.mean()
    alpha = alpha.item()
    vim_params_dict.update({
        "norm_alpha": alpha,
        "norm_R": R
    })

    # knn params
    knn_params_dict = {
        # these are already shuffled by the trainloader
        "features": features[:n_samples] 
    }

    # training data stats
    uncs = uncertainties(
        logits, 
        features=features, 
        vim_params=vim_params_dict,
        knn_params=knn_params_dict
    )


    for k,v in uncs.items():
        if torch.isnan(v).any():
            print(f"{k} contains nan values")
        if torch.isinf(v).any():
            print(f"{k} contains inf values")
    stats = metric_stats(uncs)
    return vim_params_dict, knn_params_dict, stats
//  [python]
def get_params(model, id_data, dev="cuda"):
    """Calculate params for vim and SIRC on training set."""

    print(f"eval on: {id_data.name} training set")

    # this is being done on the training set this time
    labels, logits, features = get_outputs_labels(
        model, id_data.train_loader, dev=dev
    )

    # get final fc layer 
    if config["model"]["model_type"] == "resnet101v2": # special case
        W = model.head.conv.weight.cpu().detach().squeeze()
        b = model.head.conv.bias.cpu().detach().squeeze()
    else:
        try:
            W = model.state_dict()["fc.weight"].detach().clone().cpu()
            b = model.state_dict()["fc.bias"].detach().clone().cpu()
        except:
            W = model.state_dict()["classifier.weight"].detach().clone().cpu()
            b = model.state_dict()["classifier.bias"].detach().clone().cpu()

    u = -torch.linalg.pinv(W) @ b

    # size of subspace
    # pretty much just a heuristic
    D = 1000 if features.shape[-1] > 1500 else 512 
    if features.shape[-1] < 512:
        D = features.shape[-1]//2
    centered_feats = features - u
    U = torch.linalg.eigh(
        centered_feats.T@centered_feats
    ).eigenvectors.flip(-1)
    R = U[:,D:] # eigenvectors in columns
    assert R.shape[0] == features.shape[-1]
    vlogits = torch.linalg.norm(
        (R.T @ centered_feats.unsqueeze(dim=-1)).squeeze(dim=-1), 
        ord=2, dim=-1
    )
    alpha = logits.max(dim=-1).values.mean() / vlogits.mean()
    alpha = alpha.item()

    vim_params_dict = {
        "alpha": alpha,
        "u":  u,
        "R": R
    }

    centered_feats = F.normalize(centered_feats, dim=-1)
    U = torch.linalg.eigh(
        centered_feats.T@centered_feats
    ).eigenvectors.flip(-1) # rev order to descending eigenvalue
    R = U[:,D:] # eigenvectors in columns
    assert R.shape[0] == features.shape[-1]
    vlogits = torch.linalg.norm(
        (R.T @ centered_feats.unsqueeze(dim=-1)).squeeze(dim=-1), 
        ord=2, dim=-1
    )
    alpha = logits.max(dim=-1).values.mean() / vlogits.mean()
    alpha = alpha.item()
    vim_params_dict.update({
        "norm_alpha": alpha,
        "norm_R": R
    })

    # knn params
    knn_params_dict = {
        # these are already shuffled by the trainloader
        "features": features[:n_samples] 
    }

    # training data stats
    uncs = uncertainties(
        logits, 
        features=features, 
        vim_params=vim_params_dict,
        knn_params=knn_params_dict
    )


    for k,v in uncs.items():
        if torch.isnan(v).any():
            print(f"{k} contains nan values")
        if torch.isinf(v).any():
            print(f"{k} contains inf values")
    stats = metric_stats(uncs)
    return vim_params_dict, knn_params_dict, stats
//  [python]
def _get_pnorm_after_step(self, module):
        optim = torch.optim.SGD(module.parameters(), lr=0.01)
        loss = self._compute_output(module).sum()
        loss.backward()
        optim.step()
        return torch.norm(torch.stack([p.detach().norm() for p in module.parameters()]))
//  [python]
def _get_pnorm_after_step(self, module):
        optim = torch.optim.SGD(module.parameters(), lr=0.01)
        loss = self._compute_output(module).sum()
        loss.backward()
        optim.step()
        return torch.linalg.norm(torch.stack([torch.linalg.norm(p.detach()) for p in module.parameters()]))
//  [python]
def inn_iter(f, z_init, threshold, max_iter, name='Forward'):
        z_prev, z = z_init, f(z_init)
        nstep = 0
        while nstep < max_iter:
            z_prev, z = z, f(z)
            abs_diff = torch.norm(z_prev - z).item()
            if abs_diff < threshold:
                break
            nstep += 1
        if nstep == max_iter:
            print(f'step {nstep}, not converged, abs_diff: {abs_diff}')
        return z, abs_diff
//  [python]
def inn_iter(f, z_init, threshold, max_iter, name='Forward'):
        z_prev, z = z_init, f(z_init)
        nstep = 0
        while nstep < max_iter:
            z_prev, z = z, f(z)
            abs_diff = torch.linalg.norm(z_prev - z).item()
            if abs_diff < threshold:
                break
            nstep += 1
        if nstep == max_iter:
            print(f'step {nstep}, not converged, abs_diff: {abs_diff}')
        return z, abs_diff
//  [python]
def calc_distance_from_center(_pc_scene, predicted_quads, lambda_l=0):
    pc_scene = _pc_scene.cuda()
    pc_center = torch.mean(pc_scene, dim=0)
    distance = 10.0 * torch.ones((pc_scene.shape[0],), dtype=torch.float).cuda()

    # Calculate distances per quad
    for _predicted_quad in predicted_quads:
        predicted_quad = torch.tensor(_predicted_quad).cuda()
        # TODO: convert params predicted_quad to torch.Tensor before and after nms to form a complete compute map
        quad_center = torch.mean(predicted_quad, dim=0)
        predicted_quad_norm = torch.cross(predicted_quad[1] - predicted_quad[0], predicted_quad[2] - predicted_quad[0])
        predicted_quad_norm = predicted_quad_norm / torch.norm(predicted_quad_norm)

        if torch.dot(pc_center - quad_center, predicted_quad_norm) > 0:
            predicted_quad_norm = -predicted_quad_norm  # Make inner distance < 0 and outsider > 0


        vertical_distance = (pc_scene - quad_center) @ predicted_quad_norm

        # Calc parallel loss: penalty those points whose projection lying outside the quads
        # Norm: (x, y ,z) -> (-y, x, z)
        parallel_norm1 = predicted_quad[1] - predicted_quad[0]
        parallel_norm1 /= torch.norm(parallel_norm1)
        parallel_norm2 = torch.cross(predicted_quad_norm, parallel_norm1)
        parallel_norm2 /= torch.norm(parallel_norm2)
        limit1_p = torch.dot(predicted_quad[1] - quad_center, parallel_norm1)
        limit2_p = torch.dot(predicted_quad[2] - quad_center, parallel_norm2)

        parallel_distance1 = torch.relu(torch.abs((pc_scene - quad_center) @ parallel_norm1) - torch.abs(limit1_p))
        parallel_distance2 = torch.relu(torch.abs((pc_scene - quad_center) @ parallel_norm2) - torch.abs(limit2_p))
        cond = parallel_distance1 > parallel_distance2
        parallel_distance = torch.where(cond, parallel_distance1, parallel_distance2)

        new_distance = lambda_l * torch.abs(vertical_distance) + (1-lambda_l) * parallel_distance
        # new_distance = parallel_distance
        distance[torch.abs(new_distance) < torch.abs(distance)] = new_distance[torch.abs(new_distance) < torch.abs(distance)]

    return distance
//  [python]
def calc_distance_from_center(_pc_scene, predicted_quads, lambda_l=0):
    pc_scene = _pc_scene.cuda()
    pc_center = torch.mean(pc_scene, dim=0)
    distance = 10.0 * torch.ones((pc_scene.shape[0],), dtype=torch.float).cuda()

    # Calculate distances per quad
    for _predicted_quad in predicted_quads:
        predicted_quad = torch.tensor(_predicted_quad).cuda()
        # TODO: convert params predicted_quad to torch.Tensor before and after nms to form a complete compute map
        quad_center = torch.mean(predicted_quad, dim=0)
        predicted_quad_norm = torch.cross(predicted_quad[1] - predicted_quad[0], predicted_quad[2] - predicted_quad[0])
        predicted_quad_norm = predicted_quad_norm / torch.linalg.norm(predicted_quad_norm)

        if torch.dot(pc_center - quad_center, predicted_quad_norm) > 0:
            predicted_quad_norm = -predicted_quad_norm  # Make inner distance < 0 and outsider > 0


        vertical_distance = (pc_scene - quad_center) @ predicted_quad_norm

        # Calc parallel loss: penalty those points whose projection lying outside the quads
        # Norm: (x, y ,z) -> (-y, x, z)
        parallel_norm1 = predicted_quad[1] - predicted_quad[0]
        parallel_norm1 /= torch.linalg.norm(parallel_norm1)
        parallel_norm2 = torch.cross(predicted_quad_norm, parallel_norm1)
        parallel_norm2 /= torch.linalg.norm(parallel_norm2)
        limit1_p = torch.dot(predicted_quad[1] - quad_center, parallel_norm1)
        limit2_p = torch.dot(predicted_quad[2] - quad_center, parallel_norm2)

        parallel_distance1 = torch.relu(torch.abs((pc_scene - quad_center) @ parallel_norm1) - torch.abs(limit1_p))
        parallel_distance2 = torch.relu(torch.abs((pc_scene - quad_center) @ parallel_norm2) - torch.abs(limit2_p))
        cond = parallel_distance1 > parallel_distance2
        parallel_distance = torch.where(cond, parallel_distance1, parallel_distance2)

        new_distance = lambda_l * torch.abs(vertical_distance) + (1-lambda_l) * parallel_distance
        # new_distance = parallel_distance
        distance[torch.abs(new_distance) < torch.abs(distance)] = new_distance[torch.abs(new_distance) < torch.abs(distance)]

    return distance