jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def random_get_rays_at(self, img_idx, batch_size, pose = None):
        pose_cur = self.get_pose(img_idx, pose)
        
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
            
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)][:,None]     # batch_size, 1
        weight = self.weights[img_idx][(pixels_y, pixels_x)][:,None] # batch_size, 1
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(pose_cur[None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = pose_cur[None, :3, 3].expand(rays_v.shape) # batch_size, 3
        
        normal_sample = None
        if self.use_normal:
            normal_sample = self.normals[img_idx][(pixels_y, pixels_x)].cuda()

        planes_sample = None
        if self.use_planes:
            planes_sample = self.planes[img_idx][(pixels_y, pixels_x)].unsqueeze(-1).cuda()
        
        subplanes_sample = None
        if self.use_plane_offset_loss:
            subplanes_sample = self.subplanes[img_idx][(pixels_y, pixels_x)].unsqueeze(-1).cuda()

        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask, weight], dim=-1).cuda(), pixels_x, pixels_y, normal_sample, planes_sample, subplanes_sample
//  [python]
def vector_angle(x: torch.Tensor, y: torch.Tensor, dim: int, use_degree: bool = False):
    r"""Compute the angles between two set of 3D vectors.

    Args:
        x (Tensor): set of vectors (*, 3, *)
        y (Tensor): set of vectors (*, 3, *).
        dim (int): dimension index of the coordinates.
        use_degree (bool=False): If True, return angles in degree instead of rad.

    Returns:
        angles (Tensor): (*)
    """
    cross = torch.norm(torch.cross(x, y, dim=dim), dim=dim)  # (*, 3 *) x (*, 3, *) -> (*, 3, *) -> (*)
    dot = torch.sum(x * y, dim=dim)  # (*, 3 *) x (*, 3, *) -> (*)
    angles = torch.atan2(cross, dot)  # (*)
    if use_degree:
        angles = rad2deg(angles)
    return angles
//  [python]
def vector_angle(x: torch.Tensor, y: torch.Tensor, dim: int, use_degree: bool = False):
    r"""Compute the angles between two set of 3D vectors.

    Args:
        x (Tensor): set of vectors (*, 3, *)
        y (Tensor): set of vectors (*, 3, *).
        dim (int): dimension index of the coordinates.
        use_degree (bool=False): If True, return angles in degree instead of rad.

    Returns:
        angles (Tensor): (*)
    """
    cross = torch.linalg.norm(torch.cross(x, y, dim=dim), dim=dim)  # (*, 3 *) x (*, 3, *) -> (*, 3, *) -> (*)
    dot = torch.sum(x * y, dim=dim)  # (*, 3 *) x (*, 3, *) -> (*)
    angles = torch.atan2(cross, dot)  # (*)
    if use_degree:
        angles = rad2deg(angles)
    return angles
//  [python]
def step(self):
        for gr in self.param_groups:
            lr = gr['lr']
            infNorm = torch.scalar_tensor(0.0,device='cuda')
            for x in gr['params']: 
                if not x.grad is None:
                    infNorm = torch.max(infNorm,torch.norm(x.grad.flatten(),float('inf')))
            if torch.is_nonzero(infNorm):
                lr = lr / infNorm
                for x in gr['params']:
                    if not x.grad is None:
                        x -= lr * x.grad
//  [python]
def step(self):
        for gr in self.param_groups:
            lr = gr['lr']
            infNorm = torch.scalar_tensor(0.0,device='cuda')
            for x in gr['params']: 
                if not x.grad is None:
                    infNorm = torch.max(infNorm,torch.linalg.norm(x.grad.flatten(),ord=float('inf')))
            if torch.is_nonzero(infNorm):
                lr = lr / infNorm
                for x in gr['params']:
                    if not x.grad is None:
                        x -= lr * x.grad
//  [python]
def add_edge_features(graph, pos):
        """
        adds relative displacement & displacement norm as edge features
        """
        row, col = graph.edges()
        disp = torch.tensor(pos[row.long()] - pos[col.long()])
        disp_norm = torch.norm(disp, dim=-1, keepdim=True)
        graph.edata["x"] = torch.cat((disp, disp_norm), dim=1)
        return graph
//  [python]
def add_edge_features(graph, pos):
        """
        adds relative displacement & displacement norm as edge features
        """
        row, col = graph.edges()
        disp = torch.tensor(pos[row.long()] - pos[col.long()])
        disp_norm = torch.linalg.norm(disp, dim=-1, keepdim=True)
        graph.edata["x"] = torch.cat((disp, disp_norm), dim=1)
        return graph
//  [python]
def slerp(p0, p1, fract_mixing: float, adain=True):
    r""" Copied from lunarring/latentblending
    Helper function to correctly mix two random variables using spherical interpolation.
    The function will always cast up to float64 for sake of extra 4.
    Args:
        p0: 
            First tensor for interpolation
        p1: 
            Second tensor for interpolation
        fract_mixing: float 
            Mixing coefficient of interval [0, 1]. 
            0 will return in p0
            1 will return in p1
            0.x will return a mix between both preserving angular velocity.
    """
    if p0.dtype == torch.float16:
        recast_to = 'fp16'
    else:
        recast_to = 'fp32'

    p0 = p0.double()
    p1 = p1.double()

    if adain:
        mean1, std1 = calc_mean_std(p0)
        mean2, std2 = calc_mean_std(p1)
        mean = mean1 * (1 - fract_mixing) + mean2 * fract_mixing
        std = std1 * (1 - fract_mixing) + std2 * fract_mixing
        
    norm = torch.norm(p0) * torch.norm(p1)
    epsilon = 1e-7
    dot = torch.sum(p0 * p1) / norm
    dot = dot.clamp(-1+epsilon, 1-epsilon)

    theta_0 = torch.arccos(dot)
    sin_theta_0 = torch.sin(theta_0)
    theta_t = theta_0 * fract_mixing
    s0 = torch.sin(theta_0 - theta_t) / sin_theta_0
    s1 = torch.sin(theta_t) / sin_theta_0
    interp = p0*s0 + p1*s1

    if adain:
        interp = F.instance_norm(interp) * std + mean

    if recast_to == 'fp16':
        interp = interp.half()
    elif recast_to == 'fp32':
        interp = interp.float()

    return interp
//  [python]
def slerp(p0, p1, fract_mixing: float, adain=True):
    r""" Copied from lunarring/latentblending
    Helper function to correctly mix two random variables using spherical interpolation.
    The function will always cast up to float64 for sake of extra 4.
    Args:
        p0: 
            First tensor for interpolation
        p1: 
            Second tensor for interpolation
        fract_mixing: float 
            Mixing coefficient of interval [0, 1]. 
            0 will return in p0
            1 will return in p1
            0.x will return a mix between both preserving angular velocity.
    """
    if p0.dtype == torch.float16:
        recast_to = 'fp16'
    else:
        recast_to = 'fp32'

    p0 = p0.double()
    p1 = p1.double()

    if adain:
        mean1, std1 = calc_mean_std(p0)
        mean2, std2 = calc_mean_std(p1)
        mean = mean1 * (1 - fract_mixing) + mean2 * fract_mixing
        std = std1 * (1 - fract_mixing) + std2 * fract_mixing
        
    norm = torch.linalg.norm(p0) * torch.linalg.norm(p1)
    epsilon = 1e-7
    dot = torch.sum(p0 * p1) / norm
    dot = dot.clamp(-1+epsilon, 1-epsilon)

    theta_0 = torch.arccos(dot)
    sin_theta_0 = torch.sin(theta_0)
    theta_t = theta_0 * fract_mixing
    s0 = torch.sin(theta_0 - theta_t) / sin_theta_0
    s1 = torch.sin(theta_t) / sin_theta_0
    interp = p0*s0 + p1*s1

    if adain:
        interp = F.instance_norm(interp) * std + mean

    if recast_to == 'fp16':
        interp = interp.half()
    elif recast_to == 'fp32':
        interp = interp.float()

    return interp
//  [python]
def loss_edgelen(deformed_verts, edge_pair, gt_edge_len):
    batch_size = deformed_verts.size(0)
    flow_edge_src = deformed_verts[torch.arange(batch_size).unsqueeze(1).cuda(),
                                   edge_pair[:, :, 0]]
    flow_edge_tar = deformed_verts[torch.arange(batch_size).unsqueeze(1).cuda(),
                                   edge_pair[:, :, 1]]

    pred_edge_len = torch.norm(flow_edge_tar - flow_edge_src,
                                      dim=2, keepdim=True).squeeze()
    rand_ind = np.random.choice(pred_edge_len.shape[1], 1000, replace=False)
    deformed_edge_len = torch.abs(gt_edge_len.squeeze()[:, rand_ind] - pred_edge_len[:, rand_ind]).mean()
    return deformed_edge_len
//  [python]
def loss_edgelen(deformed_verts, edge_pair, gt_edge_len):
    batch_size = deformed_verts.size(0)
    flow_edge_src = deformed_verts[torch.arange(batch_size).unsqueeze(1).cuda(),
                                   edge_pair[:, :, 0]]
    flow_edge_tar = deformed_verts[torch.arange(batch_size).unsqueeze(1).cuda(),
                                   edge_pair[:, :, 1]]

    pred_edge_len = torch.linalg.norm(flow_edge_tar - flow_edge_src,
                                      dim=2, keepdim=True).squeeze()
    rand_ind = np.random.choice(pred_edge_len.shape[1], 1000, replace=False)
    deformed_edge_len = torch.abs(gt_edge_len.squeeze()[:, rand_ind] - pred_edge_len[:, rand_ind]).mean()
    return deformed_edge_len
//  [python]
def loss_arap(nodal_g):
    # F(x) = x + U(x)
    nodal_g_deform = eye_like(nodal_g) + nodal_g
#     anlyt_grad_inner = torch.einsum("bpm,bmp->bmp", nodal_g_deform, nodal_g_deform)
    # anlyt_grad_inner = torch.bmm(nodal_g_deform.transpose(2,1), nodal_g_deform)
    anlyt_grad_inner = torch.matmul(nodal_g_deform.transpose(-1,-2), nodal_g_deform)
    cur_rigidity = torch.norm(torch.abs(anlyt_grad_inner - eye_like(anlyt_grad_inner) + 1e-7),
                                     p="fro", dim=(-1, -2))
    return cur_rigidity
//  [python]
def loss_arap(nodal_g):
    # F(x) = x + U(x)
    nodal_g_deform = eye_like(nodal_g) + nodal_g
#     anlyt_grad_inner = torch.einsum("bpm,bmp->bmp", nodal_g_deform, nodal_g_deform)
    # anlyt_grad_inner = torch.bmm(nodal_g_deform.transpose(2,1), nodal_g_deform)
    anlyt_grad_inner = torch.matmul(nodal_g_deform.transpose(-1,-2), nodal_g_deform)
    cur_rigidity = torch.linalg.norm(torch.abs(anlyt_grad_inner - eye_like(anlyt_grad_inner) + 1e-7),
                                     ord="fro", dim=(-1, -2))
    return cur_rigidity
//  [python]
def loss_nodal_rigidity(nodal_g, radius):
    const = 4 / 3 * np.pi * radius**3
    # F(x) = x + U(x)
    nodal_g_deform = eye_like(nodal_g) + nodal_g
    anlyt_grad_inner = torch.einsum(
        "bqnmp,bqnpl->bqnml",
        nodal_g_deform.transpose(-2, -1),
        nodal_g_deform.transpose(-2, -1),
    )
    cur_rigidity = const * torch.norm(
        torch.abs(anlyt_grad_inner - eye_like(anlyt_grad_inner) + 1e-7),
        p="fro",
        dim=(-1, -2),
        keepdim=True,
    )
    return cur_rigidity.sum(dim=2).mean()
//  [python]
def loss_nodal_rigidity(nodal_g, radius):
    const = 4 / 3 * np.pi * radius**3
    # F(x) = x + U(x)
    nodal_g_deform = eye_like(nodal_g) + nodal_g
    anlyt_grad_inner = torch.einsum(
        "bqnmp,bqnpl->bqnml",
        nodal_g_deform.transpose(-2, -1),
        nodal_g_deform.transpose(-2, -1),
    )
    cur_rigidity = const * torch.linalg.norm(
        torch.abs(anlyt_grad_inner - eye_like(anlyt_grad_inner) + 1e-7),
        ord="fro",
        dim=(-1, -2),
        keepdim=True,
    )
    return cur_rigidity.sum(dim=2).mean()
//  [python]
def power_iteration(mat_g: torch.Tensor, num_iters: int = 100) -> torch.Tensor:
    r"""Compute the maximum eigenvalue of matrix, for scaling.

        Mostly, power_iteration method is faster than torch.einval in case of the symmetric PSD matrix.
        Also, I removed the validation, error of singular value every iteration, so that boosting the speed.

    :param mat_g: torch.Tensor. the symmetric PSD matrix.
    :param num_iters: int. Number of iterations.
    """
    v = torch.randn(mat_g.shape[0], dtype=mat_g.dtype, device=mat_g.device)
    mat_v = torch.empty_like(v)

    for _ in range(num_iters):
        torch.mv(mat_g, v, out=mat_v)
        v = mat_v.div(torch.norm(mat_v))

    return (v.t() @ mat_g @ v).clamp_min_(1e-16)
//  [python]
def power_iteration(mat_g: torch.Tensor, num_iters: int = 100) -> torch.Tensor:
    r"""Compute the maximum eigenvalue of matrix, for scaling.

        Mostly, power_iteration method is faster than torch.einval in case of the symmetric PSD matrix.
        Also, I removed the validation, error of singular value every iteration, so that boosting the speed.

    :param mat_g: torch.Tensor. the symmetric PSD matrix.
    :param num_iters: int. Number of iterations.
    """
    v = torch.randn(mat_g.shape[0], dtype=mat_g.dtype, device=mat_g.device)
    mat_v = torch.empty_like(v)

    for _ in range(num_iters):
        torch.mv(mat_g, v, out=mat_v)
        v = mat_v.div(torch.linalg.norm(mat_v))

    return (v.t() @ mat_g @ v).clamp_min_(1e-16)
//  [python]
def unif_ball_sample(
    sigma: float,
    n_samples: int,
    n_walks: int,
    dim: int=3
):
    """
    source: https://en.wikipedia.org/wiki/N-sphere#Uniformly_at_random_within_the_n-ball
    """
    radius = sigma / 2.
    points = torch.normal(mean=0, std=1, size=(n_samples, n_walks, dim))
    points /= torch.norm(points, dim=-1, ord=2).reshape(n_samples, n_walks, 1)
    points *= torch.pow(torch.rand(size=(n_samples, n_walks, 1)), 1 / dim) * radius

    return points
//  [python]
def unif_ball_sample(
    sigma: float,
    n_samples: int,
    n_walks: int,
    dim: int=3
):
    """
    source: https://en.wikipedia.org/wiki/N-sphere#Uniformly_at_random_within_the_n-ball
    """
    radius = sigma / 2.
    points = torch.normal(mean=0, std=1, size=(n_samples, n_walks, dim))
    points /= torch.linalg.norm(points, dim=-1, ord=2).reshape(n_samples, n_walks, 1)
    points *= torch.pow(torch.rand(size=(n_samples, n_walks, 1)), 1 / dim) * radius

    return points
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze()  # batch_size, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape)  # batch_size, 3
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, :1]], dim=-1).cuda()
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze()  # batch_size, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape)  # batch_size, 3
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, :1]], dim=-1).cuda()
//  [python]
def magdev(tensor):
    magnitudes = torch.norm(tensor, dim=-1)
    std_deviation = torch.std(magnitudes).item()
    result = f"Standard deviation of magnitudes: {std_deviation}, which is {round(std_deviation/torch.mean(magnitudes, dim=-1).item(), 2)}x the mean"
    print(result)
    return result
//  [python]
def magdev(tensor):
    magnitudes = torch.linalg.norm(tensor, dim=-1)
    std_deviation = torch.std(magnitudes).item()
    result = f"Standard deviation of magnitudes: {std_deviation}, which is {round(std_deviation/torch.mean(magnitudes, dim=-1).item(), 2)}x the mean"
    print(result)
    return result
//  [python]
def _damping_factor_loss_fn(
    damping_factor: torch.Tensor,
    gradient: torch.Tensor,
    hessian: torch.Tensor,
    trust_radius: float,
) -> torch.Tensor:
    """Computes the squared difference between the target trust radius and the step size
    proposed by the Levenberg–Marquardt solver.

    This is used when finding the optimal damping factor.

    Args:
        damping_factor: The damping factor with ``shape=(1,)``.
        gradient: The gradient with ``shape=(n,)``.
        hessian: The hessian with ``shape=(n, n)``.
        trust_radius: The target trust radius.

    Returns:
        The squared difference.
    """
    dx, _ = _solver(damping_factor, gradient, hessian)
    dx_norm = torch.norm(dx)

    _LOGGER.debug(
        f"finding trust radius: length {dx_norm:.4e} (target {trust_radius:.4e})"
    )

    return (dx_norm - trust_radius) ** 2
//  [python]
def _damping_factor_loss_fn(
    damping_factor: torch.Tensor,
    gradient: torch.Tensor,
    hessian: torch.Tensor,
    trust_radius: float,
) -> torch.Tensor:
    """Computes the squared difference between the target trust radius and the step size
    proposed by the Levenberg–Marquardt solver.

    This is used when finding the optimal damping factor.

    Args:
        damping_factor: The damping factor with ``shape=(1,)``.
        gradient: The gradient with ``shape=(n,)``.
        hessian: The hessian with ``shape=(n, n)``.
        trust_radius: The target trust radius.

    Returns:
        The squared difference.
    """
    dx, _ = _solver(damping_factor, gradient, hessian)
    dx_norm = torch.linalg.norm(dx)

    _LOGGER.debug(
        f"finding trust radius: length {dx_norm:.4e} (target {trust_radius:.4e})"
    )

    return (dx_norm - trust_radius) ** 2
//  [python]
def _torch_normalize_vectors(rr):
    """Normalize surface vertices."""
    norm = torch.norm(rr, dim=1, keepdim=True)
    mask = norm > 0
    norm[~mask] = 1  # in case norm is zero, divide by 1
    new_rr = rr / norm
    return new_rr
//  [python]
def _torch_normalize_vectors(rr):
    """Normalize surface vertices."""
    norm = torch.linalg.norm(rr, axis=1, keepdim=True)
    mask = norm > 0
    norm[~mask] = 1  # in case norm is zero, divide by 1
    new_rr = rr / norm
    return new_rr
//  [python]
def get_grid(basis: torch.Tensor) -> tuple[torch.Tensor, ...]:
    """Return grid containing linear combinations of basis vectors sufficient to cover Wigner-Seitz cell"""
    a, b, c = basis[:, 0], basis[:, 1], basis[:, 2]
    # Find maximum radius of sphere within parallelopiped from basis lattice vectors
    Rmax = np.max(
        [torch.norm(_vec) for _vec in [a, b, c, a + b, b + c, a + c, a + b + c]])
    # This sphere should be within shape_min in each direction
    # Therefore shape_min >= Rmax * (R^-1 or G^-1)
    shape_min = (Rmax * (torch.inverse(basis).norm(dim=0)))
    shape_min = torch.ceil(shape_min)
    iGridX = np.arange(-shape_min[0], shape_min[0] + 1)
    iGridY = np.arange(-shape_min[1], shape_min[1] + 1)
    iGridZ = np.arange(-shape_min[2], shape_min[2] + 1)
    # Reduce grid for any orthogonal basis vector
    if np.dot(a, b) == 0 == np.dot(a, c):
        iGridX = np.arange(-1., 2.)
    if np.dot(b, a) == 0 == np.dot(b, c):
        iGridY = np.arange(-1., 2.)
    if np.dot(c, a) == 0 == np.dot(c, b):
        iGridZ = np.arange(-1., 2.)
    iGrid = np.array([x.flatten() for x in np.meshgrid(iGridX, iGridY, iGridZ)]).T
    iGrid = iGrid[np.where((iGrid ** 2).sum(axis=1) > 0)[0]]  # eliminate origin
    return torch.tensor(iGrid), torch.tensor(iGrid) @ basis.T
//  [python]
def get_grid(basis: torch.Tensor) -> tuple[torch.Tensor, ...]:
    """Return grid containing linear combinations of basis vectors sufficient to cover Wigner-Seitz cell"""
    a, b, c = basis[:, 0], basis[:, 1], basis[:, 2]
    # Find maximum radius of sphere within parallelopiped from basis lattice vectors
    Rmax = np.max(
        [torch.linalg.norm(_vec) for _vec in [a, b, c, a + b, b + c, a + c, a + b + c]])
    # This sphere should be within shape_min in each direction
    # Therefore shape_min >= Rmax * (R^-1 or G^-1)
    shape_min = (Rmax * (torch.linalg.inv(basis).norm(dim=0)))
    shape_min = torch.ceil(shape_min)
    iGridX = np.arange(-shape_min[0], shape_min[0] + 1)
    iGridY = np.arange(-shape_min[1], shape_min[1] + 1)
    iGridZ = np.arange(-shape_min[2], shape_min[2] + 1)
    # Reduce grid for any orthogonal basis vector
    if np.dot(a, b) == 0 == np.dot(a, c):
        iGridX = np.arange(-1., 2.)
    if np.dot(b, a) == 0 == np.dot(b, c):
        iGridY = np.arange(-1., 2.)
    if np.dot(c, a) == 0 == np.dot(c, b):
        iGridZ = np.arange(-1., 2.)
    iGrid = np.array([x.flatten() for x in np.meshgrid(iGridX, iGridY, iGridZ)]).T
    iGrid = iGrid[np.where((iGrid ** 2).sum(axis=1) > 0)[0]]  # eliminate origin
    return torch.tensor(iGrid), torch.tensor(iGrid) @ basis.T
//  [python]
def _(
    a: torch.Tensor,
    ord: Optional[Union[str, float, int]] = None,
    axis: Optional[Union[int, Tuple[int, ...]]] = None,
    keepdims: bool = False,
) -> torch.Tensor:
    return t.norm(a, ord=ord, dim=axis, keepdim=keepdims)
//  [python]
def _(
    a: torch.Tensor,
    ord: Optional[Union[str, float, int]] = None,
    axis: Optional[Union[int, Tuple[int, ...]]] = None,
    keepdims: bool = False,
) -> torch.Tensor:
    return torch.linalg.norm(a, ord=ord, dim=axis, keepdims=keepdims)
//  [python]
def normalize(X, p=2):
    if isinstance(X, torch.Tensor):
        norm = torch.norm(X.flatten(1), p=p, dim=1)
        norm = norm.clip(min=1e-8)
        for _ in range(len(X.shape)-1):
            norm = norm.unsqueeze(-1)
        return X / norm
    elif isinstance(X, np.ndarray):
        norm = np.linalg.norm(X.reshape(X.shape[0], -1), ord=p, axis=1)
        norm = np.clip(norm, a_min=1e-8, a_max=None)
        for _ in range(len(X.shape)-1):
            norm = np.expand_dims(norm, -1)
        return X / norm
    else:
        raise TypeError('Input array not instances of torch.Tensor or np.ndarray')
//  [python]
def normalize(X, p=2):
    if isinstance(X, torch.Tensor):
        norm = torch.linalg.norm(X.flatten(1), ord=p, axis=1)
        norm = norm.clip(min=1e-8)
        for _ in range(len(X.shape)-1):
            norm = norm.unsqueeze(-1)
        return X / norm
    elif isinstance(X, np.ndarray):
        norm = np.linalg.norm(X.reshape(X.shape[0], -1), ord=p, axis=1)
        norm = np.clip(norm, a_min=1e-8, a_max=None)
        for _ in range(len(X.shape)-1):
            norm = np.expand_dims(norm, -1)
        return X / norm
    else:
        raise TypeError('Input array not instances of torch.Tensor or np.ndarray')
//  [python]
def pixel2cam(pixel_coord, f, c, T_=None):
    x = (pixel_coord[..., 0] - c[0]) / f[0] * pixel_coord[..., 2]
    y = (pixel_coord[..., 1] - c[1]) / f[1] * pixel_coord[..., 2]
    z = pixel_coord[..., 2]
    try:
        cam_coord = np.concatenate((x[...,None], y[...,None], z[...,None]), -1)
    except:
        cam_coord = torch.cat((x[...,None], y[...,None], z[...,None]), -1)
        
    if T_ is not None: # MANO space와 scale과 wrist를 맞추고자
        ratio = torch.norm(T_[:,9] - T_[:,0], dim=-1) / torch.norm(cam_coord[:,:,9] - cam_coord[:,:,0], dim=-1)
        cam_coord = cam_coord * ratio[:,:,None,None]  # template, m
        cam_coord = cam_coord - cam_coord[:, :, :1] + T_[:,:1]
    return cam_coord
//  [python]
def pixel2cam(pixel_coord, f, c, T_=None):
    x = (pixel_coord[..., 0] - c[0]) / f[0] * pixel_coord[..., 2]
    y = (pixel_coord[..., 1] - c[1]) / f[1] * pixel_coord[..., 2]
    z = pixel_coord[..., 2]
    try:
        cam_coord = np.concatenate((x[...,None], y[...,None], z[...,None]), -1)
    except:
        cam_coord = torch.cat((x[...,None], y[...,None], z[...,None]), -1)
        
    if T_ is not None: # MANO space와 scale과 wrist를 맞추고자
        ratio = torch.linalg.norm(T_[:,9] - T_[:,0], dim=-1) / torch.linalg.norm(cam_coord[:,:,9] - cam_coord[:,:,0], dim=-1)
        cam_coord = cam_coord * ratio[:,:,None,None]  # template, m
        cam_coord = cam_coord - cam_coord[:, :, :1] + T_[:,:1]
    return cam_coord
//  [python]
def _has_converged(
    dx: torch.Tensor,
    loss_history: list[torch.Tensor],
    gradient: torch.Tensor,
    step_quality: float,
    config: LevenbergMarquardtConfig,
) -> bool:
    """Check whether the optimization has converged.

    Args:
        dx: The current step.
        loss_history: The loss history.
        gradient: The current gradient.
        step_quality: The quality of the current step.
        config: The optimizer config.

    Returns:
        Whether the optimization has converged.
    """
    if config.n_convergence_criteria == 0:
        return False

    if step_quality <= config.quality_threshold_low:
        # don't converge on low quality steps
        return False

    grad_norm = torch.norm(gradient)
    grad_converged = grad_norm < config.convergence_gradient

    step_norm = torch.norm(dx)
    step_converged = 0.0 <= step_norm < config.convergence_step

    loss_std = (
        torch.inf
        if len(loss_history) == 0
        else torch.std(torch.stack(loss_history[-config.n_convergence_steps :]))
    )
    loss_converged = (
        len(loss_history) >= config.n_convergence_steps
        and loss_std < config.convergence_loss
    )

    if grad_converged:
        _LOGGER.info(f"gradient norm is converged: ({grad_norm:.2e})")
    if step_converged:
        _LOGGER.info(f"step size is converged: ({step_norm:.2e})")
    if loss_converged:
        _LOGGER.info(f"loss is converged: ({loss_std:.2e})")

    return bool(
        sum((grad_converged, step_converged, loss_converged))
        >= config.n_convergence_criteria
    )
//  [python]
def _has_converged(
    dx: torch.Tensor,
    loss_history: list[torch.Tensor],
    gradient: torch.Tensor,
    step_quality: float,
    config: LevenbergMarquardtConfig,
) -> bool:
    """Check whether the optimization has converged.

    Args:
        dx: The current step.
        loss_history: The loss history.
        gradient: The current gradient.
        step_quality: The quality of the current step.
        config: The optimizer config.

    Returns:
        Whether the optimization has converged.
    """
    if config.n_convergence_criteria == 0:
        return False

    if step_quality <= config.quality_threshold_low:
        # don't converge on low quality steps
        return False

    grad_norm = torch.linalg.norm(gradient)
    grad_converged = grad_norm < config.convergence_gradient

    step_norm = torch.linalg.norm(dx)
    step_converged = 0.0 <= step_norm < config.convergence_step

    loss_std = (
        torch.inf
        if len(loss_history) == 0
        else torch.std(torch.stack(loss_history[-config.n_convergence_steps :]))
    )
    loss_converged = (
        len(loss_history) >= config.n_convergence_steps
        and loss_std < config.convergence_loss
    )

    if grad_converged:
        _LOGGER.info(f"gradient norm is converged: ({grad_norm:.2e})")
    if step_converged:
        _LOGGER.info(f"step size is converged: ({step_norm:.2e})")
    if loss_converged:
        _LOGGER.info(f"loss is converged: ({loss_std:.2e})")

    return bool(
        sum((grad_converged, step_converged, loss_converged))
        >= config.n_convergence_criteria
    )
//  [python]
def solve(self, k0, x, y, theta):
        """
        Args:
            k0, x, y, theta: torch.Size([batch_size])
        """

        end_state_constraint = torch.stack([x, y, theta], dim=1)

        radius = torch.hypot(x, y)
        solution = torch.stack([
            torch.zeros_like(radius),
            torch.zeros_like(radius),
            radius,
        ], dim=1)
        end_state = self.get_end_state(k0, solution)
        cost = None
        for i in range(self.max_iter):
            delta_state = end_state_constraint - end_state

            cost = torch.norm(delta_state, dim=1)
            if cost.max().item() <= self.threshold:
                break
                
            J = self.calculate_Jacobian(k0, solution)
            delta_solution = -torch.bmm(torch.inverse(J), delta_state.unsqueeze(2))
            solution -= self.step_size * delta_solution.squeeze(2)

            end_state = self.get_end_state(k0, solution)

        param = self.solution2param(k0, solution)
        info = Data(solution=solution, cost=cost)
        return param, info
//  [python]
def solve(self, k0, x, y, theta):
        """
        Args:
            k0, x, y, theta: torch.Size([batch_size])
        """

        end_state_constraint = torch.stack([x, y, theta], dim=1)

        radius = torch.hypot(x, y)
        solution = torch.stack([
            torch.zeros_like(radius),
            torch.zeros_like(radius),
            radius,
        ], dim=1)
        end_state = self.get_end_state(k0, solution)
        cost = None
        for i in range(self.max_iter):
            delta_state = end_state_constraint - end_state

            cost = torch.linalg.norm(delta_state, dim=1)
            if cost.max().item() <= self.threshold:
                break
                
            J = self.calculate_Jacobian(k0, solution)
            delta_solution = -torch.bmm(torch.inverse(J), delta_state.unsqueeze(2))
            solution -= self.step_size * delta_solution.squeeze(2)

            end_state = self.get_end_state(k0, solution)

        param = self.solution2param(k0, solution)
        info = Data(solution=solution, cost=cost)
        return param, info
//  [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 mask_sigma(sigma,
               pts,
               mask_radius,
               mask_rad_norm,
               invert_mask=False):
  """Compute and apply mask based on scene bounds."""
  masks = []

  if mask_rad_norm is not None:
    if mask_rad_norm == "inf":
      bounds = (torch.max(pts.abs(), -1)[0] <= mask_radius)
      masks.append(bounds)
    else:
      norm = torch.norm(pts, p=mask_rad_norm, dim=-1)
      bounds = norm <= mask_radius
      masks.append(bounds)

  if masks:
    mask = functools.reduce(torch.logical_and, masks)
    if invert_mask:
      mask = torch.logical_not(mask)
    return sigma * mask
  else:
    raise ValueError

  return sigma
//  [python]
def mask_sigma(sigma,
               pts,
               mask_radius,
               mask_rad_norm,
               invert_mask=False):
  """Compute and apply mask based on scene bounds."""
  masks = []

  if mask_rad_norm is not None:
    if mask_rad_norm == "inf":
      bounds = (torch.max(pts.abs(), -1)[0] <= mask_radius)
      masks.append(bounds)
    else:
      norm = torch.linalg.norm(pts, ord=mask_rad_norm, axis=-1)
      bounds = norm <= mask_radius
      masks.append(bounds)

  if masks:
    mask = functools.reduce(torch.logical_and, masks)
    if invert_mask:
      mask = torch.logical_not(mask)
    return sigma * mask
  else:
    raise ValueError

  return sigma
//  [python]
def norm(array: torch.tensor, *args, **kwargs) -> float:
        return torch.norm(array, *args, **kwargs)
//  [python]
def norm(array: torch.tensor, *args, **kwargs) -> float:
        return torch.linalg.norm(array, *args, **kwargs)
//  [python]
def test_interval_grid_sampling_with_n():
    I = Interval(R1('x'), 0, 1)
    points = I.sample_grid(n=10)
    assert points.as_tensor.shape == (10, 1)
    assert all(I._contains(points))
    for i in range(8):
        dist_1 = torch.norm(points.as_tensor[i+1] - points.as_tensor[i])
        dist_2 = torch.norm(points.as_tensor[i+1] - points.as_tensor[i+2])
        assert torch.isclose(dist_1, dist_2)
//  [python]
def test_interval_grid_sampling_with_n():
    I = Interval(R1('x'), 0, 1)
    points = I.sample_grid(n=10)
    assert points.as_tensor.shape == (10, 1)
    assert all(I._contains(points))
    for i in range(8):
        dist_1 = torch.linalg.norm(points.as_tensor[i+1] - points.as_tensor[i])
        dist_2 = torch.linalg.norm(points.as_tensor[i+1] - points.as_tensor[i+2])
        assert torch.isclose(dist_1, dist_2)
//  [python]
def pre_process(
        self,
        feature_stack: torch.Tensor,
        max_length: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Pre-process the CNN features.

        Args:
          feature_stack (torch.Tensor): Features extracted from CNN
          max_length (Tensor | None): Used to unit normalize the feature_stack vector. If ``max_len`` is not
            provided, the length is calculated from the ``feature_stack``. Defaults to None.

        Returns:
            (Tuple): Stacked features and length
        """
        if max_length is None:
            max_length = torch.max(torch.norm(feature_stack, p=2, dim=1))

        if self.feature_scaling_method == FeatureScalingMethod.NORM:
            feature_stack /= torch.norm(feature_stack, p=2, dim=1)[:, None]
        elif self.feature_scaling_method == FeatureScalingMethod.SCALE:
            feature_stack /= max_length
        else:
            msg = "Unknown pre-processing mode. Available modes are: Normalized and Scale."
            raise RuntimeError(msg)
        return feature_stack, max_length
//  [python]
def pre_process(
        self,
        feature_stack: torch.Tensor,
        max_length: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Pre-process the CNN features.

        Args:
          feature_stack (torch.Tensor): Features extracted from CNN
          max_length (Tensor | None): Used to unit normalize the feature_stack vector. If ``max_len`` is not
            provided, the length is calculated from the ``feature_stack``. Defaults to None.

        Returns:
            (Tuple): Stacked features and length
        """
        if max_length is None:
            max_length = torch.max(torch.linalg.norm(feature_stack, ord=2, dim=1))

        if self.feature_scaling_method == FeatureScalingMethod.NORM:
            feature_stack /= torch.linalg.norm(feature_stack, ord=2, dim=1)[:, None]
        elif self.feature_scaling_method == FeatureScalingMethod.SCALE:
            feature_stack /= max_length
        else:
            msg = "Unknown pre-processing mode. Available modes are: Normalized and Scale."
            raise RuntimeError(msg)
        return feature_stack, max_length
//  [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.type(torch.float32)
        if rot_mats is not None:
            rot_mats = rot_mats.type(torch.float32)

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

        self._rot_mats = rot_mats
        self._quats = quats
//  [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.type(torch.float32)
        if rot_mats is not None:
            rot_mats = rot_mats.type(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 get_rotvec(self, eps=1e-6) -> torch.Tensor:
        """
        Return the underlying axis-angle rotation vector.

        Follow's scipy's implementation:
        https://github.com/scipy/scipy/blob/HEAD/scipy/spatial/transform/_rotation.pyx#L1385-L1402

        Returns:
            The stored rotation as a axis-angle vector.
        """
        quat = self.get_quats()
        # w > 0 to ensure 0 <= angle <= pi
        flip = (quat[..., :1] < 0).float()
        quat = (-1 * quat) * flip + (1 - flip) * quat

        angle = 2 * torch.atan2(torch.norm(quat[..., 1:], dim=-1), quat[..., 0])

        angle2 = angle * angle
        small_angle_scales = 2 + angle2 / 12 + 7 * angle2 * angle2 / 2880
        large_angle_scales = angle / torch.sin(angle / 2 + eps)

        small_angles = (angle <= 1e-3).float()
        rot_vec_scale = (
            small_angle_scales * small_angles + (1 - small_angles) * large_angle_scales
        )
        rot_vec = rot_vec_scale[..., None] * quat[..., 1:]
        return rot_vec
//  [python]
def get_rotvec(self, eps=1e-6) -> torch.Tensor:
        """
        Return the underlying axis-angle rotation vector.

        Follow's scipy's implementation:
        https://github.com/scipy/scipy/blob/HEAD/scipy/spatial/transform/_rotation.pyx#L1385-L1402

        Returns:
            The stored rotation as a axis-angle vector.
        """
        quat = self.get_quats()
        # w > 0 to ensure 0 <= angle <= pi
        flip = (quat[..., :1] < 0).float()
        quat = (-1 * quat) * flip + (1 - flip) * quat

        angle = 2 * torch.atan2(torch.linalg.norm(quat[..., 1:], dim=-1), quat[..., 0])

        angle2 = angle * angle
        small_angle_scales = 2 + angle2 / 12 + 7 * angle2 * angle2 / 2880
        large_angle_scales = angle / torch.sin(angle / 2 + eps)

        small_angles = (angle <= 1e-3).float()
        rot_vec_scale = (
            small_angle_scales * small_angles + (1 - small_angles) * large_angle_scales
        )
        rot_vec = rot_vec_scale[..., None] * quat[..., 1:]
        return rot_vec
//  [python]
def quanterion_between_two_vectors(vector_a: torch.Tensor, vector_b: torch.Tensor):
    """
    Quaternion between two vectors

    Args:
        vector_a: N x 3
        vector_b: N x 3
    Returns:
            torch.Tensor: N x 3 x 3
    """
    vector_a_norm = vector_a / torch.norm(vector_a, dim=1)
    vector_b_norm = vector_b / torch.norm(vector_b, dim=1)
    half = vector_a_norm + vector_b_norm
    half = half / torch.norm(half, dim=1)
    w = torch.sum(vector_a_norm * half, dim=1)
    xyz = torch.cross(vector_a_norm, half, dim=1)
    quaternion = torch.cat([w.unsqueeze(-1), xyz], dim=-1)
    return quaternion
//  [python]
def quanterion_between_two_vectors(vector_a: torch.Tensor, vector_b: torch.Tensor):
    """
    Quaternion between two vectors

    Args:
        vector_a: N x 3
        vector_b: N x 3
    Returns:
            torch.Tensor: N x 3 x 3
    """
    vector_a_norm = vector_a / torch.linalg.norm(vector_a, dim=1)
    vector_b_norm = vector_b / torch.linalg.norm(vector_b, dim=1)
    half = vector_a_norm + vector_b_norm
    half = half / torch.linalg.norm(half, dim=1)
    w = torch.sum(vector_a_norm * half, dim=1)
    xyz = torch.cross(vector_a_norm, half, dim=1)
    quaternion = torch.cat([w.unsqueeze(-1), xyz], dim=-1)
    return quaternion
//  [python]
def eval_springs(x, v, indices, rest, ke, kd, f):
    i = indices[:, 0]
    j = indices[:, 1]

    xi = x[i]
    xj = x[j]

    vi = v[i]
    vj = v[j]

    xij = xi - xj
    vij = vi - vj

    l = torch.norm(xij, dim=1)
    l_inv = 1.0 / l

    # normalized spring direction
    dir = (xij.T * l_inv).T

    c = l - rest
    dcdt = torch.sum(dir * vij, axis=1)

    # damping based on relative velocity.
    fs = dir.T * (ke * c + kd * dcdt)

    f.index_add_(dim=0, index=i, source=-fs.T)
    f.index_add_(dim=0, index=j, source=fs.T)
//  [python]
def eval_springs(x, v, indices, rest, ke, kd, f):
    i = indices[:, 0]
    j = indices[:, 1]

    xi = x[i]
    xj = x[j]

    vi = v[i]
    vj = v[j]

    xij = xi - xj
    vij = vi - vj

    l = torch.linalg.norm(xij, axis=1)
    l_inv = 1.0 / l

    # normalized spring direction
    dir = (xij.T * l_inv).T

    c = l - rest
    dcdt = torch.sum(dir * vij, axis=1)

    # damping based on relative velocity.
    fs = dir.T * (ke * c + kd * dcdt)

    f.index_add_(dim=0, index=i, source=-fs.T)
    f.index_add_(dim=0, index=j, source=fs.T)
//  [python]
def rotation_matrix_between_two_vectors(vector_a: torch.Tensor, vector_b: torch.Tensor):
    """
    Calculate rotation matrix between two vectors

    Args:
        vector_a: N x 3
        vector_b: N x 3
    Returns:
            torch.Tensor: N x 3 x 3
    """
    vector_a_norm = vector_a / torch.norm(vector_a, dim=1, keepdim=True)
    vector_b_norm = vector_b / torch.norm(vector_b, dim=1, keepdim=True)
    v = torch.cross(vector_a_norm, vector_b_norm, dim=1)
    cos = torch.sum(vector_a_norm * vector_b_norm, dim=1)
    sin = torch.norm(v, dim=1)
    kmat = torch.zeros([vector_a_norm.shape[0], 3, 3])
    kmat[:, 0, 1] = -v[:, 2]
    kmat[:, 0, 2] = v[:, 1]
    kmat[:, 1, 0] = v[:, 2]
    kmat[:, 1, 2] = -v[:, 0]
    kmat[:, 2, 0] = -v[:, 1]
    kmat[:, 2, 1] = v[:, 0]
    rotation_matrix = torch.eye(3) + kmat + torch.sum(kmat * kmat, dim=(1, 2)) * ((1 - cos) / (sin ** 2))
    return rotation_matrix
//  [python]
def rotation_matrix_between_two_vectors(vector_a: torch.Tensor, vector_b: torch.Tensor):
    """
    Calculate rotation matrix between two vectors

    Args:
        vector_a: N x 3
        vector_b: N x 3
    Returns:
            torch.Tensor: N x 3 x 3
    """
    vector_a_norm = vector_a / torch.linalg.norm(vector_a, dim=1)
    vector_b_norm = vector_b / torch.linalg.norm(vector_b, dim=1)
    v = torch.cross(vector_a_norm, vector_b_norm, dim=1)
    cos = torch.sum(vector_a_norm * vector_b_norm, dim=1)
    sin = torch.linalg.norm(v, dim=1)
    kmat = torch.zeros([vector_a_norm.shape[0], 3, 3])
    kmat[:, 0, 1] = -v[:, 2]
    kmat[:, 0, 2] = v[:, 1]
    kmat[:, 1, 0] = v[:, 2]
    kmat[:, 1, 2] = -v[:, 0]
    kmat[:, 2, 0] = -v[:, 1]
    kmat[:, 2, 1] = v[:, 0]
    rotation_matrix = torch.eye(3) + kmat + torch.sum(kmat * kmat, dim=(1, 2)) * ((1 - cos) / (sin ** 2))
    return rotation_matrix