jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(
        self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
        Args:
            rot_mats:
                A [*, 3, 3] rotation matrix tensor. Mutually exclusive with quats
            quats:
                A [*, 4] quaternion. Mutually exclusive with rot_mats. If normalize_quats is not True, must be a unit
                quaternion
            normalize_quats:
                If quats is specified, whether to normalize quats
        """
        if (rot_mats is None and quats is None) or (
            rot_mats is not None and quats is not None
        ):
            raise ValueError("Exactly one input argument must be specified")

        if (rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or (
            quats is not None and quats.shape[-1] != 4
        ):
            raise ValueError("Incorrectly shaped rotation matrix or quaternion")

        # Force full-precision
        if quats is not None:
            quats = quats.to(dtype=torch.float32)
        if rot_mats is not None:
            rot_mats = rot_mats.to(dtype=torch.float32)

        if quats is not None and normalize_quats:
            quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def torch_find_closest_on_surface_full(othersRefList, refN, pointN, resource_folder):
    """
    full implementation of cortical projection using the balloon inflation algorithm described in
    https://doi.org/10.1016/j.neuroimage.2005.01.018
    :param othersRefList: the refN x pointN x 3 transformed sensor locations (into ref brains)
    :param refN: number of reference brains
    :param pointN: number of sensors
    :return:
    """
    device = othersRefList.device
    otherRefCList = torch.empty((refN, pointN, 3), dtype=torch.float).to(device)
    # otherRefCList = np.empty((1, refN), dtype=object)
    for i in range(refN):
        my_str = resource_folder+"/MNI_templates/xyzall{}.npy".format(str(i + 1))
        XYZ = load_raw_MNI_data(my_str, i, resource_folder=resource_folder)
        XYZ = torch.FloatTensor(XYZ).to(othersRefList.device).to(device)
        projectionListC = torch.ones((pointN, 3)).to(device)
        for j in range(pointN):
            P = othersRefList[i, j, :3]
            PP = torch.broadcast_to(P, XYZ.shape)
            D = torch.norm(XYZ - PP, dim=1)
            top = round(XYZ.shape[0] * 0.05)  # select 5% of data (original paper selects 1000 points)
            XYZtop = XYZ[torch.topk(D, largest=False, k=top).indices, :]
            Nclose = 200
            XYZclose = XYZ[torch.topk(D, largest=False, k=Nclose).indices, :]
            PNear = torch.mean(XYZclose, dim=0)  # select mean of closest 200 points

            # Line between P and PNear
            p1 = P.unsqueeze(0)
            p2 = PNear.unsqueeze(0)
            p3 = XYZtop
            # cross product the line with the point and normalize gives exactly distance from line
            distance_from_line = torch.norm(torch.cross(torch.broadcast_to(p2 - p1, p3.shape), p3 - p1) / torch.norm(p2-p1), dim=1)
            det = 0
            rodR = 0
            while det == 0:
                rodR += 1
                Iless2 = torch.where(distance_from_line <= rodR)
                rod = XYZtop[Iless2]
                det = torch.sum(rod ** 2)

            # Find brain surface points on the vicinity of P (l 862)
            PPB = torch.broadcast_to(P, rod.shape)
            VicD = torch.norm(rod - PPB, dim=1)
            NVic = 3
            if rod.shape[0] < NVic:
                NVic = rod.shape[0]
            IVicD = torch.argsort(VicD)
            NIVicD = IVicD[0:NVic]
            Vic = rod[NIVicD, :]
            CP = torch.mean(Vic, dim=0)

            projectionListC[j, :] = CP
        otherRefCList[i] = projectionListC
    return otherRefCList
//  [python]
def torch_find_closest_on_surface_full(othersRefList, refN, pointN, resource_folder):
    """
    full implementation of cortical projection using the balloon inflation algorithm described in
    https://doi.org/10.1016/j.neuroimage.2005.01.018
    :param othersRefList: the refN x pointN x 3 transformed sensor locations (into ref brains)
    :param refN: number of reference brains
    :param pointN: number of sensors
    :return:
    """
    device = othersRefList.device
    otherRefCList = torch.empty((refN, pointN, 3), dtype=torch.float).to(device)
    # otherRefCList = np.empty((1, refN), dtype=object)
    for i in range(refN):
        my_str = resource_folder+"/MNI_templates/xyzall{}.npy".format(str(i + 1))
        XYZ = load_raw_MNI_data(my_str, i, resource_folder=resource_folder)
        XYZ = torch.FloatTensor(XYZ).to(othersRefList.device).to(device)
        projectionListC = torch.ones((pointN, 3)).to(device)
        for j in range(pointN):
            P = othersRefList[i, j, :3]
            PP = torch.broadcast_to(P, XYZ.shape)
            D = torch.linalg.norm(XYZ - PP, dim=1)
            top = round(XYZ.shape[0] * 0.05)  # select 5% of data (original paper selects 1000 points)
            XYZtop = XYZ[torch.topk(D, largest=False, k=top).indices, :]
            Nclose = 200
            XYZclose = XYZ[torch.topk(D, largest=False, k=Nclose).indices, :]
            PNear = torch.mean(XYZclose, dim=0)  # select mean of closest 200 points

            # Line between P and PNear
            p1 = P.unsqueeze(0)
            p2 = PNear.unsqueeze(0)
            p3 = XYZtop
            # cross product the line with the point and normalize gives exactly distance from line
            distance_from_line = torch.linalg.norm(torch.cross(torch.broadcast_to(p2 - p1, p3.shape), p3 - p1) / torch.linalg.norm(p2-p1), dim=1)
            det = 0
            rodR = 0
            while det == 0:
                rodR += 1
                Iless2 = torch.where(distance_from_line <= rodR)
                rod = XYZtop[Iless2]
                det = torch.sum(rod ** 2)

            # Find brain surface points on the vicinity of P (l 862)
            PPB = torch.broadcast_to(P, rod.shape)
            VicD = torch.linalg.norm(rod - PPB, dim=1)
            NVic = 3
            if rod.shape[0] < NVic:
                NVic = rod.shape[0]
            IVicD = torch.argsort(VicD)
            NIVicD = IVicD[0:NVic]
            Vic = rod[NIVicD, :]
            CP = torch.mean(Vic, dim=0)

            projectionListC[j, :] = CP
        otherRefCList[i] = projectionListC
    return otherRefCList
//  [python]
def get_dislf(seq, xyz, mask):
    L = seq.shape[0]
    resolved_cys_mask = ((seq==ChemData().aa2num['CYS']) * mask[:,5]).nonzero().squeeze(-1)  # cys[5]=='sg'
    sgs = xyz[resolved_cys_mask,5]
    ii,jj = torch.triu_indices(sgs.shape[0],sgs.shape[0],1)
    d_sg_sg = torch.norm(sgs[ii,:]-sgs[jj,:], dim=-1)
    is_dslf = (d_sg_sg>1.7)*(d_sg_sg<2.3)

    dslf = []
    for i in is_dslf.nonzero():
        dslf.append( (
            resolved_cys_mask[ii[i]].item(),
            resolved_cys_mask[jj[i]].item(),
        ) )
    return dslf
//  [python]
def get_dislf(seq, xyz, mask):
    L = seq.shape[0]
    resolved_cys_mask = ((seq==ChemData().aa2num['CYS']) * mask[:,5]).nonzero().squeeze(-1)  # cys[5]=='sg'
    sgs = xyz[resolved_cys_mask,5]
    ii,jj = torch.triu_indices(sgs.shape[0],sgs.shape[0],1)
    d_sg_sg = torch.linalg.norm(sgs[ii,:]-sgs[jj,:], dim=-1)
    is_dslf = (d_sg_sg>1.7)*(d_sg_sg<2.3)

    dslf = []
    for i in is_dslf.nonzero():
        dslf.append( (
            resolved_cys_mask[ii[i]].item(),
            resolved_cys_mask[jj[i]].item(),
        ) )
    return dslf
//  [python]
def get_rays(directions: torch.Tensor,
             c2w: torch.Tensor,
             ndc: bool,
             ndc_near: float = 1.0,
             intrinsics: Optional[Intrinsics] = None,
             normalize_rd: bool = True):
    """Get ray origin and normalized directions in world coordinate for all pixels in one image.
    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems

    Args:
        directions:
        c2w:
        ndc:
        ndc_near:
        intrinsics:
        normalize_rd:

    Returns:

    """
    directions = directions.view(-1, 3)  # [n_rays, 3]
    if len(c2w.shape) == 2:
        c2w = c2w[None, ...]
    rd = (directions[:, None, :] * c2w[:, :3, :3]).sum(dim=-1)
    ro = torch.broadcast_to(c2w[:, :3, 3], directions.shape)
    if ndc:
        assert intrinsics is not None, "intrinsics must not be None when NDC active."
        ro, rd = ndc_rays_blender(
            intrinsics=intrinsics, near=ndc_near, rays_o=ro, rays_d=rd)
    if normalize_rd:
        rd /= torch.norm(rd, dim=-1, keepdim=True)
    return ro, rd
//  [python]
def get_rays(directions: torch.Tensor,
             c2w: torch.Tensor,
             ndc: bool,
             ndc_near: float = 1.0,
             intrinsics: Optional[Intrinsics] = None,
             normalize_rd: bool = True):
    """Get ray origin and normalized directions in world coordinate for all pixels in one image.
    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems

    Args:
        directions:
        c2w:
        ndc:
        ndc_near:
        intrinsics:
        normalize_rd:

    Returns:

    """
    directions = directions.view(-1, 3)  # [n_rays, 3]
    if len(c2w.shape) == 2:
        c2w = c2w[None, ...]
    rd = (directions[:, None, :] * c2w[:, :3, :3]).sum(dim=-1)
    ro = torch.broadcast_to(c2w[:, :3, 3], directions.shape)
    if ndc:
        assert intrinsics is not None, "intrinsics must not be None when NDC active."
        ro, rd = ndc_rays_blender(
            intrinsics=intrinsics, near=ndc_near, rays_o=ro, rays_d=rd)
    if normalize_rd:
        rd /= torch.linalg.norm(rd, dim=-1, keepdim=True)
    return ro, rd
//  [python]
def forward(self, x: Dict[str, torch.Tensor], render_kwargs: Dict[str, str]):
        if not self.training or not ((self.cur_iter % self.frequency) == 0) or self.cur_iter >= self.stop_iter:
            return x

        rays = x[self.rays_name]

        points = x[self.in_points_field]
        points = points.view(rays.shape[0], -1, points.shape[-1])

        dists = x[self.in_distances_field]
        dists = dists.view(rays.shape[0], -1, dists.shape[-1])

        # Get offset
        diffs = points[..., 1:, :] - points[..., :-1, :]

        offset = diffs.new_zeros(diffs.shape[0], diffs.shape[1], self.random_per_sample) \
            + torch.linspace(
                0.0, 1.0 - 1.0 / self.random_per_sample, self.random_per_sample, device=diffs.device
            ).view(1, 1, -1)
        offset = offset + torch.rand_like(offset) / self.random_per_sample

        # Add offset
        points = torch.cat(
            [
                points[:, :-1, None, :] + offset.unsqueeze(-1) * diffs.unsqueeze(-2),
                points[:, -1:, None, :].repeat(1, 1, self.random_per_sample, 1)
            ],
            dim=1
        ).view(points.shape[0], -1, 3)
        dists = torch.norm(points - rays[..., None, :3], dim=-1)

        x[self.out_points_field] = points
        x[self.out_distances_field] = dists

        # Update outputs
        for key in x.keys():
            if key not in ['points', 'distances', 'rays']:
                x[key] = x[key].view(
                    points.shape[0], -1, 1, x[key].shape[-1]
                ).repeat(1, 1, self.random_per_sample, 1)
                x[key] = x[key].view(points.shape[0], -1, x[key].shape[-1])

        return x
//  [python]
def forward(self, x: Dict[str, torch.Tensor], render_kwargs: Dict[str, str]):
        if not self.training or not ((self.cur_iter % self.frequency) == 0) or self.cur_iter >= self.stop_iter:
            return x

        rays = x[self.rays_name]

        points = x[self.in_points_field]
        points = points.view(rays.shape[0], -1, points.shape[-1])

        dists = x[self.in_distances_field]
        dists = dists.view(rays.shape[0], -1, dists.shape[-1])

        # Get offset
        diffs = points[..., 1:, :] - points[..., :-1, :]

        offset = diffs.new_zeros(diffs.shape[0], diffs.shape[1], self.random_per_sample) \
            + torch.linspace(
                0.0, 1.0 - 1.0 / self.random_per_sample, self.random_per_sample, device=diffs.device
            ).view(1, 1, -1)
        offset = offset + torch.rand_like(offset) / self.random_per_sample

        # Add offset
        points = torch.cat(
            [
                points[:, :-1, None, :] + offset.unsqueeze(-1) * diffs.unsqueeze(-2),
                points[:, -1:, None, :].repeat(1, 1, self.random_per_sample, 1)
            ],
            dim=1
        ).view(points.shape[0], -1, 3)
        dists = torch.linalg.norm(points - rays[..., None, :3], dim=-1)

        x[self.out_points_field] = points
        x[self.out_distances_field] = dists

        # Update outputs
        for key in x.keys():
            if key not in ['points', 'distances', 'rays']:
                x[key] = x[key].view(
                    points.shape[0], -1, 1, x[key].shape[-1]
                ).repeat(1, 1, self.random_per_sample, 1)
                x[key] = x[key].view(points.shape[0], -1, x[key].shape[-1])

        return x
//  [python]
def compute_point_ray_distance(
        points: torch.Tensor,
        ray_origins: torch.Tensor,
        ray_directions: torch.Tensor,
) -> T.Dict[str, T.Any]:
    """
    Compute the distance between each point to each ray.

    Args:
        points:
            (*, n, 3)
        ray_origins:
            (*, m, 3)
        ray_directions:
            (*, m, 3)

    Returns:
        dists: (*, m, n) distance between each point to each ray
        projections:  (*, m, n, 3) the projected points on the ray
        ts: (*, m, n) length on ray (can be negative)
    """

    points = points.unsqueeze(-3)  # (*, 1, n, 3)
    ray_origins = ray_origins.unsqueeze(-2)  # (*, m, 1, 3)
    ray_directions = ray_directions.unsqueeze(-2)  # (*, m, 1, 3)
    dv = points - ray_origins  # (*, m, n, 3)

    ts = (dv * ray_directions).sum(dim=-1, keepdim=True)  # (*, m, n, 1)  projected length on ray (can be negative)
    projections = ray_origins + ts * ray_directions  # (*, m, n, 3)
    dists = torch.norm(points - projections, dim=-1)  # (*, m, n)

    return dict(
        dists=dists,  # (*, m, n)
        projections=projections,  # (*, m, n, 3)
        ts=ts.squeeze(-1),  # (*, m, n)
    )
//  [python]
def compute_point_ray_distance(
        points: torch.Tensor,
        ray_origins: torch.Tensor,
        ray_directions: torch.Tensor,
) -> T.Dict[str, T.Any]:
    """
    Compute the distance between each point to each ray.

    Args:
        points:
            (*, n, 3)
        ray_origins:
            (*, m, 3)
        ray_directions:
            (*, m, 3)

    Returns:
        dists: (*, m, n) distance between each point to each ray
        projections:  (*, m, n, 3) the projected points on the ray
        ts: (*, m, n) length on ray (can be negative)
    """

    points = points.unsqueeze(-3)  # (*, 1, n, 3)
    ray_origins = ray_origins.unsqueeze(-2)  # (*, m, 1, 3)
    ray_directions = ray_directions.unsqueeze(-2)  # (*, m, 1, 3)
    dv = points - ray_origins  # (*, m, n, 3)

    ts = (dv * ray_directions).sum(dim=-1, keepdim=True)  # (*, m, n, 1)  projected length on ray (can be negative)
    projections = ray_origins + ts * ray_directions  # (*, m, n, 3)
    dists = torch.linalg.norm(points - projections, ord=2, dim=-1)  # (*, m, n)

    return dict(
        dists=dists,  # (*, m, n)
        projections=projections,  # (*, m, n, 3)
        ts=ts.squeeze(-1),  # (*, m, n)
    )
//  [python]
def rodrigues_alignment_matrix(src_vectors, tgt_vectors):
    r"""Compute the Rodrigues rotation matrix aligning source vectors to target vectors.

    Args:
        src_vectors (Tensor): source vectors (*, 3)
        tgt_vectors (Tensor): target vectors (*, 3)

    Returns:
        rotations (Tensor): rotation matrix (*, 3, 3)
    """
    input_shape = src_vectors.shape
    src_vectors = src_vectors.view(-1, 3)  # (B, 3)
    tgt_vectors = tgt_vectors.view(-1, 3)  # (B, 3)

    # compute axes
    src_vectors = F.normalize(src_vectors, dim=-1, p=2)  # (B, 3)
    tgt_vectors = F.normalize(tgt_vectors, dim=-1, p=2)  # (B, 3)
    src_skews = skew_symmetric_matrix(src_vectors)  # (B, 3, 3)
    axes = torch.matmul(src_skews,
                        tgt_vectors.unsqueeze(-1)).squeeze(-1)  # (B, 3)

    # compute rodrigues rotation matrix
    sin_values = torch.norm(axes, dim=-1)  # (B,)
    cos_values = (src_vectors * tgt_vectors).sum(dim=-1)  # (B,)
    axes = F.normalize(axes, dim=-1, p=2)  # (B, 3)
    skews = skew_symmetric_matrix(axes)  # (B, 3, 3)
    eyes = torch.eye(3).cuda().unsqueeze(0).expand_as(skews)  # (B, 3, 3)
    sin_values = sin_values.view(-1, 1, 1)
    cos_values = cos_values.view(-1, 1, 1)
    rotations = eyes + sin_values * skews + (1.0 - cos_values) * torch.matmul(
        skews, skews)

    # handle opposite direction
    sin_values = sin_values.view(-1)
    cos_values = cos_values.view(-1)
    masks = torch.logical_and(torch.eq(sin_values, 0.0),
                              torch.lt(cos_values, 0.0))
    rotations[masks] *= -1

    output_shape = input_shape[:-1] + (3, 3)
    rotations = rotations.view(*output_shape)

    return rotations
//  [python]
def rodrigues_alignment_matrix(src_vectors, tgt_vectors):
    r"""Compute the Rodrigues rotation matrix aligning source vectors to target vectors.

    Args:
        src_vectors (Tensor): source vectors (*, 3)
        tgt_vectors (Tensor): target vectors (*, 3)

    Returns:
        rotations (Tensor): rotation matrix (*, 3, 3)
    """
    input_shape = src_vectors.shape
    src_vectors = src_vectors.view(-1, 3)  # (B, 3)
    tgt_vectors = tgt_vectors.view(-1, 3)  # (B, 3)

    # compute axes
    src_vectors = F.normalize(src_vectors, dim=-1, p=2)  # (B, 3)
    tgt_vectors = F.normalize(tgt_vectors, dim=-1, p=2)  # (B, 3)
    src_skews = skew_symmetric_matrix(src_vectors)  # (B, 3, 3)
    axes = torch.matmul(src_skews,
                        tgt_vectors.unsqueeze(-1)).squeeze(-1)  # (B, 3)

    # compute rodrigues rotation matrix
    sin_values = torch.linalg.norm(axes, dim=-1)  # (B,)
    cos_values = (src_vectors * tgt_vectors).sum(dim=-1)  # (B,)
    axes = F.normalize(axes, dim=-1, p=2)  # (B, 3)
    skews = skew_symmetric_matrix(axes)  # (B, 3, 3)
    eyes = torch.eye(3).cuda().unsqueeze(0).expand_as(skews)  # (B, 3, 3)
    sin_values = sin_values.view(-1, 1, 1)
    cos_values = cos_values.view(-1, 1, 1)
    rotations = eyes + sin_values * skews + (1.0 - cos_values) * torch.matmul(
        skews, skews)

    # handle opposite direction
    sin_values = sin_values.view(-1)
    cos_values = cos_values.view(-1)
    masks = torch.logical_and(torch.eq(sin_values, 0.0),
                              torch.lt(cos_values, 0.0))
    rotations[masks] *= -1

    output_shape = input_shape[:-1] + (3, 3)
    rotations = rotations.view(*output_shape)

    return rotations
//  [python]
def mouth_open_degree(scaled_mouth_keypoints):
        """
        计算张嘴程度
        :param scaled_mouth_keypoints: 按嘴部框范围标准化后的关键点坐标
        :return:
        """
        up_mouth_keypoints = scaled_mouth_keypoints[13:16]
        down_mouth_keypoints = scaled_mouth_keypoints[17:20]
        # 计算嘴对应点之间的距离
        mouth_distance = torch.norm(up_mouth_keypoints - down_mouth_keypoints, dim=1)
        return mouth_distance
//  [python]
def mouth_open_degree(scaled_mouth_keypoints):
        """
        计算张嘴程度
        :param scaled_mouth_keypoints: 按嘴部框范围标准化后的关键点坐标
        :return:
        """
        up_mouth_keypoints = scaled_mouth_keypoints[13:16]
        down_mouth_keypoints = scaled_mouth_keypoints[17:20]
        # 计算嘴对应点之间的距离
        mouth_distance = torch.linalg.norm(up_mouth_keypoints - down_mouth_keypoints, axis=1)
        return mouth_distance
//  [python]
def eigh(
    A: pt.Tensor,
    B: Optional[pt.Tensor] = None,
    eps: Optional[float] = 1e-15,
    use_eigh_cpu: Optional[bool] = False,
) -> Tuple[pt.Tensor, pt.Tensor]:
    """
    Eigenvalue decomposition of a complex Hermitian symmetric matrix
    """

    # for 2x2, use specialized routine
    if A.shape[-1] == 2:
        return eigh_2x2(A, B=B, eps=eps)

    if B is not None:
        assert A.shape == B.shape, "A and B should have the same size"

        # generalized eigenvalue decomposition
        Bval, Bvec = eigh_wrapper(B, use_cpu=use_eigh_cpu)
        Bval_sqrt = pt.sqrt(pt.clamp(Bval, min=1e-5))
        B_sqrt_inv = Bvec @ (pt.reciprocal(Bval_sqrt[..., :, None]) * hermite(Bvec))

        Linv_A_LHinv = B_sqrt_inv @ A @ hermite(B_sqrt_inv)

        eigenvalues, eigenvectors = eigh_wrapper(Linv_A_LHinv, use_cpu=use_eigh_cpu)

        eigenvectors = hermite(B_sqrt_inv) @ eigenvectors

        # normalize the eigenvectors
        norm = t.norm(eigenvectors, dim=-2, keepdim=True)
        eigenvectors = eigenvectors / pt.clamp(norm, min=1e-5)

    else:
        # regular eigenvalue decomposition
        eigenvalues, eigenvectors = eigh_wrapper(A, use_cpu=use_eigh_cpu)

    return_type_eigh = collections.namedtuple(
        "return_type_ptiva_linalg_eigh", ("eigenvalues", "eigenvectors")
    )

    return return_type_eigh(eigenvalues, eigenvectors)
//  [python]
def eigh(
    A: pt.Tensor,
    B: Optional[pt.Tensor] = None,
    eps: Optional[float] = 1e-15,
    use_eigh_cpu: Optional[bool] = False,
) -> Tuple[pt.Tensor, pt.Tensor]:
    """
    Eigenvalue decomposition of a complex Hermitian symmetric matrix
    """

    # for 2x2, use specialized routine
    if A.shape[-1] == 2:
        return eigh_2x2(A, B=B, eps=eps)

    if B is not None:
        assert A.shape == B.shape, "A and B should have the same size"

        # generalized eigenvalue decomposition
        Bval, Bvec = eigh_wrapper(B, use_cpu=use_eigh_cpu)
        Bval_sqrt = pt.sqrt(pt.clamp(Bval, min=1e-5))
        B_sqrt_inv = Bvec @ (pt.reciprocal(Bval_sqrt[..., :, None]) * hermite(Bvec))

        Linv_A_LHinv = B_sqrt_inv @ A @ hermite(B_sqrt_inv)

        eigenvalues, eigenvectors = eigh_wrapper(Linv_A_LHinv, use_cpu=use_eigh_cpu)

        eigenvectors = hermite(B_sqrt_inv) @ eigenvectors

        # normalize the eigenvectors
        norm = pt.linalg.norm(eigenvectors, dim=-2, keepdim=True)
        eigenvectors = eigenvectors / pt.clamp(norm, min=1e-5)

    else:
        # regular eigenvalue decomposition
        eigenvalues, eigenvectors = eigh_wrapper(A, use_cpu=use_eigh_cpu)

    return_type_eigh = collections.namedtuple(
        "return_type_ptiva_linalg_eigh", ("eigenvalues", "eigenvectors")
    )

    return return_type_eigh(eigenvalues, eigenvectors)
//  [python]
def calculate_acceletation(motionloader, device, xyz):
    # for now even if it is not xyz, the acceleration is one the euclidian/pose
    outfeat = "output_xyz" if xyz else "output"

    sum_acc = 0
    num_acc = 0
    for batch in motionloader:
        motion = batch[outfeat].permute(0, 3, 1, 2)
        bs, num_frames, njoints, nfeats = motion.shape
        
        velocity = motion[:, 1:] - motion[:, :-1]
        acceleration = velocity[:, 1:] - velocity[:, :-1]
        acceleration_normed = torch.norm(acceleration, dim=3)

        lengths = batch["lengths"]
        mask = lengths_to_mask(lengths - 2)  # because acceleration

        usefull_accs_n = acceleration_normed[mask]
        sum_acc += usefull_accs_n.sum().item()
        num_acc += np.prod(usefull_accs_n.shape)

    acceleration = sum_acc/num_acc
    return acceleration
//  [python]
def calculate_acceletation(motionloader, device, xyz):
    # for now even if it is not xyz, the acceleration is one the euclidian/pose
    outfeat = "output_xyz" if xyz else "output"

    sum_acc = 0
    num_acc = 0
    for batch in motionloader:
        motion = batch[outfeat].permute(0, 3, 1, 2)
        bs, num_frames, njoints, nfeats = motion.shape
        
        velocity = motion[:, 1:] - motion[:, :-1]
        acceleration = velocity[:, 1:] - velocity[:, :-1]
        acceleration_normed = torch.linalg.norm(acceleration, axis=3)

        lengths = batch["lengths"]
        mask = lengths_to_mask(lengths - 2)  # because acceleration

        usefull_accs_n = acceleration_normed[mask]
        sum_acc += usefull_accs_n.sum().item()
        num_acc += np.prod(usefull_accs_n.shape)

    acceleration = sum_acc/num_acc
    return acceleration
//  [python]
def test_dot_accuracy():
    y = torch.FloatTensor([[1, 0, 1], [0, -1, 0], [0.5, 0.5, 0.5]])
    y_pred = torch.FloatTensor([[0.3, -0.2, 0.5], [0.8, 0.5, 0.1], [-0.3, -0.4, 0.2]])

    expected_accuracy = 0
    for crt_y, crt_y_pred in zip(y, y_pred):
        crt_y = crt_y / torch.norm(crt_y)
        crt_y_pred = crt_y_pred / torch.norm(crt_y_pred)
        expected_accuracy += 0.5 * (1 + torch.dot(crt_y, crt_y_pred))

    expected_accuracy /= len(y)
    ns = SimpleNamespace(y=y, y_pred=y_pred)
    accuracy = dot_accuracy(ns, None)

    assert pytest.approx(accuracy) == expected_accuracy
//  [python]
def test_dot_accuracy():
    y = torch.FloatTensor([[1, 0, 1], [0, -1, 0], [0.5, 0.5, 0.5]])
    y_pred = torch.FloatTensor([[0.3, -0.2, 0.5], [0.8, 0.5, 0.1], [-0.3, -0.4, 0.2]])

    expected_accuracy = 0
    for crt_y, crt_y_pred in zip(y, y_pred):
        crt_y = crt_y / torch.linalg.norm(crt_y)
        crt_y_pred = crt_y_pred / torch.linalg.norm(crt_y_pred)
        expected_accuracy += 0.5 * (1 + torch.dot(crt_y, crt_y_pred))

    expected_accuracy /= len(y)
    ns = SimpleNamespace(y=y, y_pred=y_pred)
    accuracy = dot_accuracy(ns, None)

    assert pytest.approx(accuracy) == expected_accuracy
//  [python]
def lookat_matrix(source_pos, target_pos, up):
    """
    IMPORTANT: USES RIGHT UP BACK XYZ CONVENTION
    :param source_pos (*, 3)
    :param target_pos (*, 3)
    :param up (3,)
    """
    *dims, _ = source_pos.shape
    up = up.reshape(*(1,) * len(dims), 3)
    up = up / torch.norm(up, dim=-1, keepdim=True)
    back = normalize(target_pos - source_pos)
    right = normalize(torch.cross(up, back))
    up = normalize(torch.cross(back, right))
    R = torch.stack([right, up, back], dim=-1)
    return make_4x4_pose(R, source_pos)
//  [python]
def lookat_matrix(source_pos, target_pos, up):
    """
    IMPORTANT: USES RIGHT UP BACK XYZ CONVENTION
    :param source_pos (*, 3)
    :param target_pos (*, 3)
    :param up (3,)
    """
    *dims, _ = source_pos.shape
    up = up.reshape(*(1,) * len(dims), 3)
    up = up / torch.linalg.norm(up, dim=-1, keepdim=True)
    back = normalize(target_pos - source_pos)
    right = normalize(torch.linalg.cross(up, back))
    up = normalize(torch.linalg.cross(back, right))
    R = torch.stack([right, up, back], dim=-1)
    return make_4x4_pose(R, source_pos)
//  [python]
def compute_lookat_box(bb_min, bb_max, intrins):
    """
    The center and distance to a scene with bb_min, bb_max
    to place a camera with given intrinsics
    :param bb_min (3,)
    :param bb_max (3,)
    :param intrinsics, (fx, fy, cx, cy) of camera
    :param view_angle (optional) viewing angle in radians (elevation)
    """
    fx, fy, cx, cy = intrins
    bb_min, bb_max = torch.tensor(bb_min), torch.tensor(bb_max)
    center = 0.5 * (bb_min + bb_max)
    size = torch.norm(bb_max - bb_min)
    cam_dist = np.sqrt(fx**2 + fy**2) / np.sqrt(cx**2 + cy**2)
    cam_dist = 0.75 * size * cam_dist
    return center, cam_dist
//  [python]
def compute_lookat_box(bb_min, bb_max, intrins):
    """
    The center and distance to a scene with bb_min, bb_max
    to place a camera with given intrinsics
    :param bb_min (3,)
    :param bb_max (3,)
    :param intrinsics, (fx, fy, cx, cy) of camera
    :param view_angle (optional) viewing angle in radians (elevation)
    """
    fx, fy, cx, cy = intrins
    bb_min, bb_max = torch.tensor(bb_min), torch.tensor(bb_max)
    center = 0.5 * (bb_min + bb_max)
    size = torch.linalg.norm(bb_max - bb_min)
    cam_dist = np.sqrt(fx**2 + fy**2) / np.sqrt(cx**2 + cy**2)
    cam_dist = 0.75 * size * cam_dist
    return center, cam_dist
//  [python]
def so3_Jl(x):
    K = vec2skew(x)
    theta = torch.norm(x, dim=-1, keepdim=True).unsqueeze(-1)
    theta2 = theta**2
    I = torch.eye(3, device=x.device, dtype=x.dtype).expand(x.shape[:-1]+(3, 3))
    idx = (theta > torch.finfo(theta.dtype).eps)
    coef1 = torch.zeros_like(theta, requires_grad=False)
    coef1[idx] = (1-theta[idx].cos())/theta2[idx]
    coef1[~idx] = 0.5 - (1.0/24.0) * theta2[~idx]

    coef2 = torch.zeros_like(theta, requires_grad=False)
    coef2[idx] = (theta[idx] - theta[idx].sin()) / (theta[idx] * theta2[idx])
    coef2[~idx] = 1.0/6.0 - (1.0/120) * theta2[~idx]
    return (I + coef1 * K + coef2 * (K@K))
//  [python]
def so3_Jl(x):
    K = vec2skew(x)
    theta = torch.linalg.norm(x, dim=-1, keepdim=True).unsqueeze(-1)
    theta2 = theta**2
    I = torch.eye(3, device=x.device, dtype=x.dtype).expand(x.shape[:-1]+(3, 3))
    idx = (theta > torch.finfo(theta.dtype).eps)
    coef1 = torch.zeros_like(theta, requires_grad=False)
    coef1[idx] = (1-theta[idx].cos())/theta2[idx]
    coef1[~idx] = 0.5 - (1.0/24.0) * theta2[~idx]

    coef2 = torch.zeros_like(theta, requires_grad=False)
    coef2[idx] = (theta[idx] - theta[idx].sin()) / (theta[idx] * theta2[idx])
    coef2[~idx] = 1.0/6.0 - (1.0/120) * theta2[~idx]
    return (I + coef1 * K + coef2 * (K@K))
//  [python]
def so3_Jl_inv(x):
    K = vec2skew(x)
    theta = torch.norm(x, dim=-1, keepdim=True).unsqueeze(-1)
    I = torch.eye(3, device=x.device, dtype=x.dtype).expand(x.shape[:-1]+(3, 3))
    idx = (theta > torch.finfo(theta.dtype).eps)
    coef2 = torch.zeros_like(theta, requires_grad=False)
    theta_idx = theta[idx]
    theta_half_idx, theta2_idx = 0.5 * theta_idx, theta_idx * theta_idx
    coef2[idx] = (1.0 - theta_idx * theta_half_idx.cos() / (2.0 * theta_half_idx.sin())) / theta2_idx
    coef2[~idx] = 1.0 / 12.0
    return (I - 0.5 * K + coef2 * (K @ K))
//  [python]
def so3_Jl_inv(x):
    K = vec2skew(x)
    theta = torch.linalg.norm(x, dim=-1, keepdim=True).unsqueeze(-1)
    I = torch.eye(3, device=x.device, dtype=x.dtype).expand(x.shape[:-1]+(3, 3))
    idx = (theta > torch.finfo(theta.dtype).eps)
    coef2 = torch.zeros_like(theta, requires_grad=False)
    theta_idx = theta[idx]
    theta_half_idx, theta2_idx = 0.5 * theta_idx, theta_idx * theta_idx
    coef2[idx] = (1.0 - theta_idx * theta_half_idx.cos() / (2.0 * theta_half_idx.sin())) / theta2_idx
    coef2[~idx] = 1.0 / 12.0
    return (I - 0.5 * K + coef2 * (K @ K))
//  [python]
def set_angle(self, rot_index, angle, coord):
        angle = angle + 1e-10 # add epsilon to avoid NaN gradients
        idx1, idx2 = self.rot_edges[rot_index]
        # idx1 = self.rot_edges[rot_index, 0]
        # idx2 = self.rot_edges[rot_index, 1]
        mask = self.rot_masks[rot_index]
        if (isinstance(angle, float) or angle.shape == tuple()) and len(coord.shape) == 2:
            ax = coord[idx2] - coord[idx1]
            ax_norm = ax/torch.norm(ax)
            rot = roma.rotvec_to_rotmat(ax_norm*angle)
            to_rotate = coord # [mask]
            rotated = torch.einsum('ij,bj->bi', rot, (to_rotate - coord[idx2])) + coord[idx2]
            # ret = coord.clone()
            # ret[mask] = rotated
            ret = coord*(~mask.unsqueeze(-1)) + rotated*mask.unsqueeze(-1)
            return ret
        elif len(angle.shape) == 1 and len(coord.shape) == 3:
            ax = coord[:,idx2] - coord[:,idx1]
            ax_norm = ax/torch.norm(ax, dim=1, keepdim=True)
            rotvec = torch.einsum('t,td->td', angle, ax_norm)
            rot = roma.rotvec_to_rotmat(rotvec)
            to_rotate = coord[:,mask]
            center = coord[:,idx2].unsqueeze(1)
            rotated = torch.einsum('tij,tbj->tbi', rot, (to_rotate - center)) + center
            ret = coord.clone()
            ret[:,mask] = rotated
            return ret
        elif len(angle.shape) == 1 and len(coord.shape) == 2:
            ax = coord[idx2] - coord[idx1]
            ax_norm = ax/torch.norm(ax)
            rotvec = torch.einsum('t,d->td', angle, ax_norm)
            rot = roma.rotvec_to_rotmat(rotvec)
            to_rotate = coord[mask]
            center = coord[idx2]
            rotated = torch.einsum('tij,bj->tbi', rot, (to_rotate - center)) + center
            ret = coord.expand(len(angle), -1, -1).clone()
            ret[:, mask] = rotated
            return ret
        else:
            raise NotImplementedError
//  [python]
def set_angle(self, rot_index, angle, coord):
        angle = angle + 1e-10 # add epsilon to avoid NaN gradients
        idx1, idx2 = self.rot_edges[rot_index]
        # idx1 = self.rot_edges[rot_index, 0]
        # idx2 = self.rot_edges[rot_index, 1]
        mask = self.rot_masks[rot_index]
        if (isinstance(angle, float) or angle.shape == tuple()) and len(coord.shape) == 2:
            ax = coord[idx2] - coord[idx1]
            ax_norm = ax/torch.linalg.norm(ax)
            rot = roma.rotvec_to_rotmat(ax_norm*angle)
            to_rotate = coord # [mask]
            rotated = torch.einsum('ij,bj->bi', rot, (to_rotate - coord[idx2])) + coord[idx2]
            # ret = coord.clone()
            # ret[mask] = rotated
            ret = coord*(~mask.unsqueeze(-1)) + rotated*mask.unsqueeze(-1)
            return ret
        elif len(angle.shape) == 1 and len(coord.shape) == 3:
            ax = coord[:,idx2] - coord[:,idx1]
            ax_norm = ax/torch.linalg.norm(ax, axis=1, keepdim=True)
            rotvec = torch.einsum('t,td->td', angle, ax_norm)
            rot = roma.rotvec_to_rotmat(rotvec)
            to_rotate = coord[:,mask]
            center = coord[:,idx2].unsqueeze(1)
            rotated = torch.einsum('tij,tbj->tbi', rot, (to_rotate - center)) + center
            ret = coord.clone()
            ret[:,mask] = rotated
            return ret
        elif len(angle.shape) == 1 and len(coord.shape) == 2:
            ax = coord[idx2] - coord[idx1]
            ax_norm = ax/torch.linalg.norm(ax)
            rotvec = torch.einsum('t,d->td', angle, ax_norm)
            rot = roma.rotvec_to_rotmat(rotvec)
            to_rotate = coord[mask]
            center = coord[idx2]
            rotated = torch.einsum('tij,bj->tbi', rot, (to_rotate - center)) + center
            ret = coord.expand(len(angle), -1, -1).clone()
            ret[:, mask] = rotated
            return ret
        else:
            raise NotImplementedError
//  [python]
def mesh_developable_reg(mesh):

    verts = mesh.vertices
    tris = mesh.faces

    device = verts.device
    V = verts.shape[0]
    F = tris.shape[0]

    POS_EPS = 1e-6
    REL_EPS = 1e-6

    def normalize(vecs):
        return vecs / (torch.norm(vecs, dim=-1, keepdim=True) + POS_EPS)

    tri_pos = verts[tris]

    vert_normal_covariance_sum = torch.zeros((V, 9), device=device)
    vert_area = torch.zeros(V, device=device)
    vert_degree = torch.zeros(V, dtype=torch.int32, device=device)

    for iC in range(3):  # loop over three corners of each triangle

        # gather tri verts
        pRoot = tri_pos[:, iC, :]
        pA = tri_pos[:, (iC + 1) % 3, :]
        pB = tri_pos[:, (iC + 2) % 3, :]

        # compute the corner angle & normal
        vA = pA - pRoot
        vAn = normalize(vA)
        vB = pB - pRoot
        vBn = normalize(vB)
        area_normal = torch.cross(vA, vB, dim=-1)
        face_area = 0.5 * torch.norm(area_normal, dim=-1)
        normal = normalize(area_normal)
        corner_angle = torch.acos(torch.clamp(torch.sum(vAn * vBn, dim=-1), min=-1., max=1.))

        # add up the contribution to the covariance matrix
        outer = normal[:, :, None] @ normal[:, None, :]
        contrib = corner_angle[:, None] * outer.reshape(-1, 9)

        # scatter the result to the appropriate matrices
        vert_normal_covariance_sum = torch_scatter.scatter_add(src=contrib,
                                                               index=tris[:, iC],
                                                               dim=-2,
                                                               out=vert_normal_covariance_sum)

        vert_area = torch_scatter.scatter_add(src=face_area / 3.,
                                              index=tris[:, iC],
                                              dim=-1,
                                              out=vert_area)

        vert_degree = torch_scatter.scatter_add(src=torch.ones(F, dtype=torch.int32, device=device),
                                                index=tris[:, iC],
                                                dim=-1,
                                                out=vert_degree)

    # The energy is the smallest eigenvalue of the outer-product matrix
    vert_normal_covariance_sum = vert_normal_covariance_sum.reshape(
        -1, 3, 3)  # reshape to a batch of matrices
    vert_normal_covariance_sum = vert_normal_covariance_sum + torch.eye(
        3, device=device)[None, :, :] * REL_EPS

    min_eigvals = torch.min(torch.eig(vert_normal_covariance_sum, eigenvectors=False)[0].abs(), dim=-1).values

    # Mask out degree-3 vertices
    vert_area = torch.where(vert_degree == 3, torch.tensor(0, dtype=vert_area.dtype,device=vert_area.device), vert_area)

    # Adjust the vertex area weighting so it is unit-less, and 1 on average
    vert_area = vert_area * (V / torch.sum(vert_area, dim=-1, keepdim=True))

    return vert_area * min_eigvals
//  [python]
def mesh_developable_reg(mesh):

    verts = mesh.vertices
    tris = mesh.faces

    device = verts.device
    V = verts.shape[0]
    F = tris.shape[0]

    POS_EPS = 1e-6
    REL_EPS = 1e-6

    def normalize(vecs):
        return vecs / (torch.linalg.norm(vecs, dim=-1, keepdim=True) + POS_EPS)

    tri_pos = verts[tris]

    vert_normal_covariance_sum = torch.zeros((V, 9), device=device)
    vert_area = torch.zeros(V, device=device)
    vert_degree = torch.zeros(V, dtype=torch.int32, device=device)

    for iC in range(3):  # loop over three corners of each triangle

        # gather tri verts
        pRoot = tri_pos[:, iC, :]
        pA = tri_pos[:, (iC + 1) % 3, :]
        pB = tri_pos[:, (iC + 2) % 3, :]

        # compute the corner angle & normal
        vA = pA - pRoot
        vAn = normalize(vA)
        vB = pB - pRoot
        vBn = normalize(vB)
        area_normal = torch.linalg.cross(vA, vB, dim=-1)
        face_area = 0.5 * torch.linalg.norm(area_normal, dim=-1)
        normal = normalize(area_normal)
        corner_angle = torch.acos(torch.clamp(torch.sum(vAn * vBn, dim=-1), min=-1., max=1.))

        # add up the contribution to the covariance matrix
        outer = normal[:, :, None] @ normal[:, None, :]
        contrib = corner_angle[:, None] * outer.reshape(-1, 9)

        # scatter the result to the appropriate matrices
        vert_normal_covariance_sum = torch_scatter.scatter_add(src=contrib,
                                                               index=tris[:, iC],
                                                               dim=-2,
                                                               out=vert_normal_covariance_sum)

        vert_area = torch_scatter.scatter_add(src=face_area / 3.,
                                              index=tris[:, iC],
                                              dim=-1,
                                              out=vert_area)

        vert_degree = torch_scatter.scatter_add(src=torch.ones(F, dtype=torch.int32, device=device),
                                                index=tris[:, iC],
                                                dim=-1,
                                                out=vert_degree)

    # The energy is the smallest eigenvalue of the outer-product matrix
    vert_normal_covariance_sum = vert_normal_covariance_sum.reshape(
        -1, 3, 3)  # reshape to a batch of matrices
    vert_normal_covariance_sum = vert_normal_covariance_sum + torch.eye(
        3, device=device)[None, :, :] * REL_EPS

    min_eigvals = torch.min(torch.linalg.eigvals(vert_normal_covariance_sum).abs(), dim=-1).values

    # Mask out degree-3 vertices
    vert_area = torch.where(vert_degree == 3, torch.tensor(0, dtype=vert_area.dtype,device=vert_area.device), vert_area)

    # Adjust the vertex area weighting so it is unit-less, and 1 on average
    vert_area = vert_area * (V / torch.sum(vert_area, dim=-1, keepdim=True))

    return vert_area * min_eigvals
//  [python]
def forward(
        self, V: torch.Tensor, k: int, max_iter: int, tol: float
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """

        :param V:
        :type V: torch.Tensor
        :param k:
        :type k: int
        :param max_iter:
        :type max_iter: int
        :param tol:
        :type tol: float
        :return:
        :rtype: Tuple[torch.Tensor, torch.Tensor]
        """
        W, H = self.initialize(V, k)  # NNDSVD-A init
        init_err = torch.norm(V - torch.mm(W, H))
        prev_err = init_err

        for i in range(1, max_iter + 1):
            W, H = self.step(V, W, H)

            if i % 10 == 0:  # check rel err stopping criterion
                err = torch.norm(V - torch.mm(W, H))
                if (prev_err - err) / init_err < tol:
                    break
                prev_err = err

            # TODO: convergence warning for case i == max_iter:

        return W, H
//  [python]
def forward(
        self, V: torch.Tensor, k: int, max_iter: int, tol: float
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """

        :param V:
        :type V: torch.Tensor
        :param k:
        :type k: int
        :param max_iter:
        :type max_iter: int
        :param tol:
        :type tol: float
        :return:
        :rtype: Tuple[torch.Tensor, torch.Tensor]
        """
        W, H = self.initialize(V, k)  # NNDSVD-A init
        init_err = torch.linalg.norm(V - torch.mm(W, H))
        prev_err = init_err

        for i in range(1, max_iter + 1):
            W, H = self.step(V, W, H)

            if i % 10 == 0:  # check rel err stopping criterion
                err = torch.linalg.norm(V - torch.mm(W, H))
                if (prev_err - err) / init_err < tol:
                    break
                prev_err = err

            # TODO: convergence warning for case i == max_iter:

        return W, H
//  [python]
def calcQ(x):
    tau, phi = x[..., :3], x[..., 3:]
    Tau, Phi = vec2skew(tau), vec2skew(phi)
    theta = torch.norm(phi, dim=-1, keepdim=True).unsqueeze(-1)
    theta2 = theta**2
    theta4 = theta2**2
    idx = (theta > torch.finfo(theta.dtype).eps)
    # coef1
    coef1 = torch.zeros_like(theta, requires_grad=False)
    coef1[idx] = (theta[idx] - theta[idx].sin()) / (theta2[idx] * theta[idx])
    coef1[~idx] = 1.0 / 6.0 - (1.0 / 120.0) * theta2[~idx]
    # coef2
    coef2 = torch.zeros_like(theta, requires_grad=False)
    coef2[idx] = (theta2[idx] + 2 * theta[idx].cos() - 2) / (2 * theta4[idx])
    coef2[~idx] = 1.0 / 24.0 - (1.0 / 720.0) * theta2[~idx]
    # coef3
    coef3 = torch.zeros_like(theta, requires_grad=False)
    coef3[idx] = (2 * theta[idx] - 3 * theta[idx].sin() + theta[idx] * theta[idx].cos()) / (2 * theta4[idx] * theta[idx])
    coef3[~idx] = 1.0 / 120.0 - (1.0 / 2520.0) * theta2[~idx]
    Q = 0.5 * Tau + coef1 * (Phi@Tau + Tau@Phi + Phi@Tau@Phi) + \
        coef2 * (Phi@Phi@Tau + Tau@Phi@Phi - 3*Phi@Tau@Phi) + coef3 * (Phi@Tau@Phi@Phi + Phi@Phi@Tau@Phi)
    return Q
//  [python]
def calcQ(x):
    tau, phi = x[..., :3], x[..., 3:]
    Tau, Phi = vec2skew(tau), vec2skew(phi)
    theta = torch.linalg.norm(phi, dim=-1, keepdim=True).unsqueeze(-1)
    theta2 = theta**2
    theta4 = theta2**2
    idx = (theta > torch.finfo(theta.dtype).eps)
    # coef1
    coef1 = torch.zeros_like(theta, requires_grad=False)
    coef1[idx] = (theta[idx] - theta[idx].sin()) / (theta2[idx] * theta[idx])
    coef1[~idx] = 1.0 / 6.0 - (1.0 / 120.0) * theta2[~idx]
    # coef2
    coef2 = torch.zeros_like(theta, requires_grad=False)
    coef2[idx] = (theta2[idx] + 2 * theta[idx].cos() - 2) / (2 * theta4[idx])
    coef2[~idx] = 1.0 / 24.0 - (1.0 / 720.0) * theta2[~idx]
    # coef3
    coef3 = torch.zeros_like(theta, requires_grad=False)
    coef3[idx] = (2 * theta[idx] - 3 * theta[idx].sin() + theta[idx] * theta[idx].cos()) / (2 * theta4[idx] * theta[idx])
    coef3[~idx] = 1.0 / 120.0 - (1.0 / 2520.0) * theta2[~idx]
    Q = 0.5 * Tau + coef1 * (Phi@Tau + Tau@Phi + Phi@Tau@Phi) + \
        coef2 * (Phi@Phi@Tau + Tau@Phi@Phi - 3*Phi@Tau@Phi) + coef3 * (Phi@Tau@Phi@Phi + Phi@Phi@Tau@Phi)
    return Q
//  [python]
def align_loss(x, y, alpha=2):
    return torch.norm(x - y, p=2, dim=1).pow(alpha).mean()
//  [python]
def align_loss(x, y, alpha=2):
    return torch.linalg.norm(x - y, ord=2, dim=1).pow(alpha).mean()
//  [python]
def coord2cross(x, edge_index, batch_mask, norm_constant=1):

    mean = unsorted_segment_sum(x, batch_mask,
                                num_segments=batch_mask.max() + 1,
                                normalization_factor=None,
                                aggregation_method='mean')
    row, col = edge_index
    cross = torch.cross(x[row]-mean[batch_mask[row]],
                        x[col]-mean[batch_mask[col]], dim=1)
    norm = torch.norm(cross, dim=1, keepdim=True)
    cross = cross / (norm + norm_constant)
    return cross
//  [python]
def coord2cross(x, edge_index, batch_mask, norm_constant=1):

    mean = unsorted_segment_sum(x, batch_mask,
                                num_segments=batch_mask.max() + 1,
                                normalization_factor=None,
                                aggregation_method='mean')
    row, col = edge_index
    cross = torch.cross(x[row]-mean[batch_mask[row]],
                        x[col]-mean[batch_mask[col]], dim=1)
    norm = torch.linalg.norm(cross, dim=1, keepdim=True)
    cross = cross / (norm + norm_constant)
    return cross
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        device = rays_o.device
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        # ----------------------------------------------------------------------------------------------------------
        # Use min value of [ cos, prev_cos ]
        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
        # robust when meeting situations like below:
        #
        # SDF
        # ^
        # |\          -----x----...
        # | \        /
        # |  x      x
        # |---\----/-------------> 0 level
        # |    \  /
        # |     \/
        # |
        # ----------------------------------------------------------------------------------------------------------
        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]).to(device), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples