jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _cosine_distance(
    vectors: torch.Tensor, centroids: torch.Tensor, split_size: int
) -> Tuple[torch.Tensor, torch.Tensor]:
    if len(vectors.shape) != 2 or len(centroids.shape) != 2:
        raise ValueError(
            f"x and y must be 2-D matrix, got: vectors.shape={vectors.shape}"
            f", centroids.shape={centroids.shape}"
        )

    y2 = torch.norm(centroids.T, dim=0, keepdim=True)

    partitions = []
    distances = []

    for sub_vectors in torch.split(vectors, split_size):
        dists = _pairwise_cosine(sub_vectors, centroids, y2)
        part_ids = torch.argmin(dists, dim=1, keepdim=True)
        partitions.append(part_ids)
        distances.append(dists.take_along_dim(part_ids, dim=1))

    return torch.cat(partitions).reshape(-1), torch.cat(distances).reshape(-1)
//  [python]
def _cosine_distance(
    vectors: torch.Tensor, centroids: torch.Tensor, split_size: int
) -> Tuple[torch.Tensor, torch.Tensor]:
    if len(vectors.shape) != 2 or len(centroids.shape) != 2:
        raise ValueError(
            f"x and y must be 2-D matrix, got: vectors.shape={vectors.shape}"
            f", centroids.shape={centroids.shape}"
        )

    y2 = torch.linalg.norm(centroids.T, dim=0, keepdim=True)

    partitions = []
    distances = []

    for sub_vectors in torch.split(vectors, split_size):
        dists = _pairwise_cosine(sub_vectors, centroids, y2)
        part_ids = torch.argmin(dists, dim=1, keepdim=True)
        partitions.append(part_ids)
        distances.append(dists.take_along_dim(part_ids, dim=1))

    return torch.cat(partitions).reshape(-1), torch.cat(distances).reshape(-1)
//  [python]
def get_best_preds(
    forecasted_trajectory: torch.Tensor,
    gt_trajectory: torch.Tensor
) -> float:
    """Compute min Average Displacement Error.
    Args:
        forecasted_trajectory: Predicted trajectory with shape [k, fut_ts, 2]
        gt_trajectory: Ground truth trajectory with shape [fut_ts, 2]
        gt_fut_masks: Ground truth traj mask with shape (fut_ts)
    Returns:
        best_forecasted_trajectory: Predicted trajectory with shape [fut_ts, 2]
    """

    # [k, fut_ts]
    dist = torch.norm(gt_trajectory[None] - forecasted_trajectory, dim=-1)
    dist = dist[..., -1]
    dist[torch.isnan(dist)] = 0
    min_mode_idx = torch.argmin(dist, dim=-1)

    return forecasted_trajectory[min_mode_idx]
//  [python]
def get_best_preds(
    forecasted_trajectory: torch.Tensor,
    gt_trajectory: torch.Tensor
) -> float:
    """Compute min Average Displacement Error.
    Args:
        forecasted_trajectory: Predicted trajectory with shape [k, fut_ts, 2]
        gt_trajectory: Ground truth trajectory with shape [fut_ts, 2]
        gt_fut_masks: Ground truth traj mask with shape (fut_ts)
    Returns:
        best_forecasted_trajectory: Predicted trajectory with shape [fut_ts, 2]
    """

    # [k, fut_ts]
    dist = torch.linalg.norm(gt_trajectory[None] - forecasted_trajectory, dim=-1)
    dist = dist[..., -1]
    dist[torch.isnan(dist)] = 0
    min_mode_idx = torch.argmin(dist, dim=-1)

    return forecasted_trajectory[min_mode_idx]
//  [python]
def test_torch_normalize_vectors(cuda, grads_on):
    device = "cpu"
    if torch.cuda.is_available() and cuda:
        device = "cuda"
    for rr in [
        torch.cat([torch.ones((1, 3)), torch.zeros((1, 3))], dim=0).to(device),
        torch.ones((2, 3), device=device),
        torch.zeros((2, 3), device=device),
    ]:
        if grads_on:
            rr.requires_grad = True

        new_rr = _torch_normalize_vectors(rr)
        assert new_rr.shape == rr.shape
        assert all(
            [
                r in [0, pytest.approx(1)]
                for r in torch.norm(new_rr, dim=1).cpu().detach().numpy()
            ]
        )
//  [python]
def test_torch_normalize_vectors(cuda, grads_on):
    device = "cpu"
    if torch.cuda.is_available() and cuda:
        device = "cuda"
    for rr in [
        torch.cat([torch.ones((1, 3)), torch.zeros((1, 3))], dim=0).to(device),
        torch.ones((2, 3), device=device),
        torch.zeros((2, 3), device=device),
    ]:
        if grads_on:
            rr.requires_grad = True

        new_rr = _torch_normalize_vectors(rr)
        assert new_rr.shape == rr.shape
        assert all(
            [
                r in [0, pytest.approx(1)]
                for r in torch.linalg.norm(new_rr, axis=1).cpu().detach().numpy()
            ]
        )
//  [python]
def plt_func(u):
    return torch.norm(u, dim=1)
//  [python]
def plt_func(u):
    return torch.linalg.norm(u, dim=1)
//  [python]
def forward(self, x):
    t = self.t.type_as(x)
    attention_weights = self.get_attention_weights(x)
    # create attention mask
    if self.training:
      with torch.no_grad():
        mean_att = attention_weights.mean(dim=1, keepdim=False)
        if self.opt['use_flux']:
          src_features = x[self.data_edge_index[0, :], :]
          dst_features = x[self.data_edge_index[1, :], :]
          delta = torch.norm(src_features-dst_features, dim=1)
          mean_att = mean_att * delta
        threshold = torch.quantile(mean_att, 1-self.opt['att_samp_pct'])
        mask = mean_att > threshold
        self.odefunc.edge_index = self.data_edge_index[:, mask.T]
        sampled_attention_weights = self.renormalise_attention(mean_att[mask])
        print('retaining {} of {} edges'.format(self.odefunc.edge_index.shape[1], self.data_edge_index.shape[1]))
        self.odefunc.attention_weights = sampled_attention_weights
    else:
      self.odefunc.edge_index = self.data_edge_index
      self.odefunc.attention_weights = attention_weights.mean(dim=1, keepdim=False)
    self.reg_odefunc.odefunc.edge_index, self.reg_odefunc.odefunc.edge_weight = self.odefunc.edge_index, self.odefunc.edge_weight
    self.reg_odefunc.odefunc.attention_weights = self.odefunc.attention_weights
    integrator = self.train_integrator if self.training else self.test_integrator

    reg_states = tuple(torch.zeros(x.size(0)).to(x) for i in range(self.nreg))

    func = self.reg_odefunc if self.training and self.nreg > 0 else self.odefunc
    state = (x,) + reg_states if self.training and self.nreg > 0 else x

    if self.opt["adjoint"] and self.training:
      state_dt = integrator(
        func, state, t,
        method=self.opt['method'],
        options={'step_size': self.opt['step_size']},
        adjoint_method=self.opt['adjoint_method'],
        adjoint_options={'step_size': self.opt['adjoint_step_size']},
        atol=self.atol,
        rtol=self.rtol,
        adjoint_atol=self.atol_adjoint,
        adjoint_rtol=self.rtol_adjoint)
    else:
      state_dt = integrator(
        func, state, t,
        method=self.opt['method'],
        options={'step_size': self.opt['step_size']},
        atol=self.atol,
        rtol=self.rtol)

    if self.training and self.nreg > 0:
      z = state_dt[0][1]
      reg_states = tuple(st[1] for st in state_dt[1:])
      return z, reg_states
    else:
      z = state_dt[1]
      return z
//  [python]
def forward(self, x):
    t = self.t.type_as(x)
    attention_weights = self.get_attention_weights(x)
    # create attention mask
    if self.training:
      with torch.no_grad():
        mean_att = attention_weights.mean(dim=1, keepdim=False)
        if self.opt['use_flux']:
          src_features = x[self.data_edge_index[0, :], :]
          dst_features = x[self.data_edge_index[1, :], :]
          delta = torch.linalg.norm(src_features-dst_features, dim=1)
          mean_att = mean_att * delta
        threshold = torch.quantile(mean_att, 1-self.opt['att_samp_pct'])
        mask = mean_att > threshold
        self.odefunc.edge_index = self.data_edge_index[:, mask.T]
        sampled_attention_weights = self.renormalise_attention(mean_att[mask])
        print('retaining {} of {} edges'.format(self.odefunc.edge_index.shape[1], self.data_edge_index.shape[1]))
        self.odefunc.attention_weights = sampled_attention_weights
    else:
      self.odefunc.edge_index = self.data_edge_index
      self.odefunc.attention_weights = attention_weights.mean(dim=1, keepdim=False)
    self.reg_odefunc.odefunc.edge_index, self.reg_odefunc.odefunc.edge_weight = self.odefunc.edge_index, self.odefunc.edge_weight
    self.reg_odefunc.odefunc.attention_weights = self.odefunc.attention_weights
    integrator = self.train_integrator if self.training else self.test_integrator

    reg_states = tuple(torch.zeros(x.size(0)).to(x) for i in range(self.nreg))

    func = self.reg_odefunc if self.training and self.nreg > 0 else self.odefunc
    state = (x,) + reg_states if self.training and self.nreg > 0 else x

    if self.opt["adjoint"] and self.training:
      state_dt = integrator(
        func, state, t,
        method=self.opt['method'],
        options={'step_size': self.opt['step_size']},
        adjoint_method=self.opt['adjoint_method'],
        adjoint_options={'step_size': self.opt['adjoint_step_size']},
        atol=self.atol,
        rtol=self.rtol,
        adjoint_atol=self.atol_adjoint,
        adjoint_rtol=self.rtol_adjoint)
    else:
      state_dt = integrator(
        func, state, t,
        method=self.opt['method'],
        options={'step_size': self.opt['step_size']},
        atol=self.atol,
        rtol=self.rtol)

    if self.training and self.nreg > 0:
      z = state_dt[0][1]
      reg_states = tuple(st[1] for st in state_dt[1:])
      return z, reg_states
    else:
      z = state_dt[1]
      return z
//  [python]
def _lipschitz_constant(W):
    #L = torch.linalg.norm(W, ord=2) ** 2
    WtW = torch.matmul(W.t(), W)
    # L = torch.linalg.eigvalsh(WtW)[-1]
    L = eigsh(WtW.detach().cpu().numpy(), k=1, which='LM',
              return_eigenvectors=False).item()

    if not np.isfinite(L):  # sometimes L is not finite because of potential cublas error.
        L = torch.norm(W, ord=2) ** 2
    return L
//  [python]
def _lipschitz_constant(W):
    #L = torch.linalg.norm(W, ord=2) ** 2
    WtW = torch.matmul(W.t(), W)
    # L = torch.linalg.eigvalsh(WtW)[-1]
    L = eigsh(WtW.detach().cpu().numpy(), k=1, which='LM',
              return_eigenvectors=False).item()

    if not np.isfinite(L):  # sometimes L is not finite because of potential cublas error.
        L = torch.linalg.norm(W, ord=2) ** 2
    return L
//  [python]
def gather_and_log_param_buffer_norms(model):
    with FSDP.state_dict_type(
        model,
        StateDictType.FULL_STATE_DICT,
        state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
    ):
        sd = model.state_dict()
        for k, v in sd.items():
            if dist.get_rank() == 0:
                print(k, torch.norm(v), v.sum())
        for n, m in model.named_buffers():
            if dist.get_rank() == 0:
                print(dist.get_rank(), n, torch.norm(m), m.sum())
//  [python]
def gather_and_log_param_buffer_norms(model):
    with FSDP.state_dict_type(
        model,
        StateDictType.FULL_STATE_DICT,
        state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
    ):
        sd = model.state_dict()
        for k, v in sd.items():
            if dist.get_rank() == 0:
                print(k, torch.linalg.norm(v), v.sum())
        for n, m in model.named_buffers():
            if dist.get_rank() == 0:
                print(dist.get_rank(), n, torch.linalg.norm(m), m.sum())
//  [python]
def field2bild(points, field, uncertainty, title, box_size, ang_pix):
    f = open(title + '.bild', 'a')
    color = 25
    cols = torch.norm(points - field, dim=1)
    points = points.detach().cpu().numpy()
    points = (points + 0.5) * box_size * ang_pix
    uncertainty = uncertainty*box_size*ang_pix
    cols = torch.round((uncertainty/20 * 45)).long()
    cos = torch.ones_like(cols)
    field = field.detach().cpu().numpy()
    field = (field + 0.5) * box_size * ang_pix
    y = np.concatenate([points, field], 1)
    #cols = torch.round(cols / torch.max(cols) * 65).long()
    for k in range(y.shape[0]):

        f.write('%s %.18g %.18g %.18g\n' % ('.color', 0.0, 0.0, 0.0))
        f.write("%s %.18g %.18g %.18g %.18g %.18g %.18g %.18g\n %.18g\n %.18g\n" % (
            '.arrow', y[k, 0], y[k, 1], y[k, 2], y[k, 3], y[k, 4], y[k, 5], 0.5, 1, 0.5))
    #f.write('%s %.18g\n' % ('.transparency', 0.6))
    # for k in range(y.shape[0]):
    #    f.write('%s %.18g\n' % ('.color', cols[k]))
    #    f.write("%s %.18g %.18g %.18g %.18g\n" % (
    #        '.sphere', y[k, 3], y[k, 4], y[k, 5], uncertainty[k]))

    f.close()
//  [python]
def field2bild(points, field, uncertainty, title, box_size, ang_pix):
    f = open(title + '.bild', 'a')
    color = 25
    cols = torch.linalg.norm(points - field, dim=1)
    points = points.detach().cpu().numpy()
    points = (points + 0.5) * box_size * ang_pix
    uncertainty = uncertainty*box_size*ang_pix
    cols = torch.round((uncertainty/20 * 45)).long()
    cos = torch.ones_like(cols)
    field = field.detach().cpu().numpy()
    field = (field + 0.5) * box_size * ang_pix
    y = np.concatenate([points, field], 1)
    #cols = torch.round(cols / torch.max(cols) * 65).long()
    for k in range(y.shape[0]):

        f.write('%s %.18g %.18g %.18g\n' % ('.color', 0.0, 0.0, 0.0))
        f.write("%s %.18g %.18g %.18g %.18g %.18g %.18g %.18g\n %.18g\n %.18g\n" % (
            '.arrow', y[k, 0], y[k, 1], y[k, 2], y[k, 3], y[k, 4], y[k, 5], 0.5, 1, 0.5))
    #f.write('%s %.18g\n' % ('.transparency', 0.6))
    # for k in range(y.shape[0]):
    #    f.write('%s %.18g\n' % ('.color', cols[k]))
    #    f.write("%s %.18g %.18g %.18g %.18g\n" % (
    #        '.sphere', y[k, 3], y[k, 4], y[k, 5], uncertainty[k]))

    f.close()
//  [python]
def sam(noise: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
    """
    Measure spectral similarity using spectral angle mapper. Result is in radians.

    Parameters
    ----------
    noise: torch.Tensor
    reference: torch.Tensor

    Returns
    -------
    SAM: torch.Tensor
        spectral similarity in radians. If inputs are matrices, matrices are returned.
        i.e. a MxN is returned if a MxNxC is given
    """
    if len(noise.shape) == 1:
        # case 1: vectors -> easy
        numer = torch.dot(noise, reference)
        denom = torch.norm(noise) * torch.norm(reference)
    else:
        # case 2: matrices -> return a MxN if MxNxC is given
        numer = torch.sum(noise * reference, dim=-1)
        denom = torch.norm(noise, dim=-1) * torch.norm(reference, dim=-1)
    eps = torch.finfo(denom.dtype).eps
    return torch.arccos(numer / (denom + eps))
//  [python]
def sam(noise: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
    """
    Measure spectral similarity using spectral angle mapper. Result is in radians.

    Parameters
    ----------
    noise: torch.Tensor
    reference: torch.Tensor

    Returns
    -------
    SAM: torch.Tensor
        spectral similarity in radians. If inputs are matrices, matrices are returned.
        i.e. a MxN is returned if a MxNxC is given
    """
    if len(noise.shape) == 1:
        # case 1: vectors -> easy
        numer = torch.dot(noise, reference)
        denom = torch.linalg.norm(noise) * torch.linalg.norm(reference)
    else:
        # case 2: matrices -> return a MxN if MxNxC is given
        numer = torch.sum(noise * reference, dim=-1)
        denom = torch.linalg.norm(noise, dim=-1) * torch.linalg.norm(reference, dim=-1)
    eps = torch.finfo(denom.dtype).eps
    return torch.arccos(numer / (denom + eps))
//  [python]
def lookat_matrix(forward_vector, up_vector):
    # Z will point forwards, towards the box center
    new_z = forward_vector / torch.norm(forward_vector, dim=-1, keepdim=True)
    # Get the X (right direction) as the cross of forward and up.
    new_x = torch.cross(new_z, up_vector)
    # Get alternative X by rotating the new Z around the old Y by 90 degrees
    # in case lookdir happens to align with the up vector and the above cross product is zero.
    new_x_alt = torch.stack([new_z[:, 2], torch.zeros_like(new_z[:, 2]), -new_z[:, 0]], dim=1)
    new_x = torch.where(torch.norm(new_x, dim=-1, keepdim=True) == 0, new_x_alt, new_x)
    new_x = new_x / torch.norm(new_x, dim=-1, keepdim=True)
    # Complete the right-handed coordinate system to get Y
    new_y = torch.cross(new_z, new_x)
    # Stack the axis vectors to get the rotation matrix
    return torch.stack([new_x, new_y, new_z], dim=1)
//  [python]
def lookat_matrix(forward_vector, up_vector):
    # Z will point forwards, towards the box center
    new_z = forward_vector / torch.linalg.norm(forward_vector, dim=-1, keepdim=True)
    # Get the X (right direction) as the cross of forward and up.
    new_x = torch.linalg.cross(new_z, up_vector)
    # Get alternative X by rotating the new Z around the old Y by 90 degrees
    # in case lookdir happens to align with the up vector and the above cross product is zero.
    new_x_alt = torch.stack([new_z[:, 2], torch.zeros_like(new_z[:, 2]), -new_z[:, 0]], dim=1)
    new_x = torch.where(torch.linalg.norm(new_x, dim=-1, keepdim=True) == 0, new_x_alt, new_x)
    new_x = new_x / torch.linalg.norm(new_x, dim=-1, keepdim=True)
    # Complete the right-handed coordinate system to get Y
    new_y = torch.linalg.cross(new_z, new_x)
    # Stack the axis vectors to get the rotation matrix
    return torch.stack([new_x, new_y, new_z], dim=1)
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1)  # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1)  # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def modify_conformer_torsion_angles_batch(pos, edge_index, mask_rotate, torsion_updates):
    pos = pos + 0
    for idx_edge, e in enumerate(edge_index):
        u, v = e[0], e[1]

        # check if need to reverse the edge, v should be connected to the part that gets rotated
        assert not mask_rotate[idx_edge, u]
        assert mask_rotate[idx_edge, v]

        rot_vec = pos[:, u] - pos[:, v]  # convention: positive rotation if pointing inwards
        rot_mat = axis_angle_to_matrix(
            rot_vec / torch.norm(rot_vec, dim=-1, keepdims=True) * torsion_updates[:, idx_edge:idx_edge + 1])

        pos[:, mask_rotate[idx_edge]] = torch.bmm(pos[:, mask_rotate[idx_edge]] - pos[:, v:v + 1], torch.transpose(rot_mat, 1, 2)) + pos[:, v:v + 1]

    return pos
//  [python]
def modify_conformer_torsion_angles_batch(pos, edge_index, mask_rotate, torsion_updates):
    pos = pos + 0
    for idx_edge, e in enumerate(edge_index):
        u, v = e[0], e[1]

        # check if need to reverse the edge, v should be connected to the part that gets rotated
        assert not mask_rotate[idx_edge, u]
        assert mask_rotate[idx_edge, v]

        rot_vec = pos[:, u] - pos[:, v]  # convention: positive rotation if pointing inwards
        rot_mat = axis_angle_to_matrix(
            rot_vec / torch.linalg.norm(rot_vec, dim=-1, keepdims=True) * torsion_updates[:, idx_edge:idx_edge + 1])

        pos[:, mask_rotate[idx_edge]] = torch.bmm(pos[:, mask_rotate[idx_edge]] - pos[:, v:v + 1], torch.transpose(rot_mat, 1, 2)) + pos[:, v:v + 1]

    return pos
//  [python]
def test_initialize_lengthscale(self):
        kernel = CategoricalKernel()
        kernel.initialize(lengthscale=1)
        actual_value = torch.tensor(1.0).view_as(kernel.lengthscale)
        self.assertLess(torch.norm(kernel.lengthscale - actual_value), 1e-5)
//  [python]
def test_initialize_lengthscale(self):
        kernel = CategoricalKernel()
        kernel.initialize(lengthscale=1)
        actual_value = torch.tensor(1.0).view_as(kernel.lengthscale)
        self.assertLess(torch.linalg.norm(kernel.lengthscale - actual_value), 1e-5)
//  [python]
def test_initialize_lengthscale_batch(self):
        kernel = CategoricalKernel(batch_shape=torch.Size([2]))
        ls_init = torch.tensor([1.0, 2.0])
        kernel.initialize(lengthscale=ls_init)
        actual_value = ls_init.view_as(kernel.lengthscale)
        self.assertLess(torch.norm(kernel.lengthscale - actual_value), 1e-5)
//  [python]
def test_initialize_lengthscale_batch(self):
        kernel = CategoricalKernel(batch_shape=torch.Size([2]))
        ls_init = torch.tensor([1.0, 2.0])
        kernel.initialize(lengthscale=ls_init)
        actual_value = ls_init.view_as(kernel.lengthscale)
        self.assertLess(torch.linalg.norm(kernel.lengthscale - actual_value), 1e-5)
//  [python]
def calculate_mpjpe(gt_joints, pred_joints):
    """
    gt_joints: num_poses x num_joints(22) x 3
    pred_joints: num_poses x num_joints(22) x 3
    (obtained from recover_from_ric())
    """
    assert gt_joints.shape == pred_joints.shape, f"GT shape: {gt_joints.shape}, pred shape: {pred_joints.shape}"

    # Align by root (pelvis)
    pelvis = gt_joints[:, [0]].mean(1)
    gt_joints = gt_joints - torch.unsqueeze(pelvis, dim=1)
    pelvis = pred_joints[:, [0]].mean(1)
    pred_joints = pred_joints - torch.unsqueeze(pelvis, dim=1)

    # Compute MPJPE
    mpjpe = torch.norm(pred_joints - gt_joints, dim=-1) # num_poses x num_joints=22
    mpjpe_seq = mpjpe.mean(-1) # num_poses

    return mpjpe_seq
//  [python]
def calculate_mpjpe(gt_joints, pred_joints):
    """
    gt_joints: num_poses x num_joints(22) x 3
    pred_joints: num_poses x num_joints(22) x 3
    (obtained from recover_from_ric())
    """
    assert gt_joints.shape == pred_joints.shape, f"GT shape: {gt_joints.shape}, pred shape: {pred_joints.shape}"

    # Align by root (pelvis)
    pelvis = gt_joints[:, [0]].mean(1)
    gt_joints = gt_joints - torch.unsqueeze(pelvis, dim=1)
    pelvis = pred_joints[:, [0]].mean(1)
    pred_joints = pred_joints - torch.unsqueeze(pelvis, dim=1)

    # Compute MPJPE
    mpjpe = torch.linalg.norm(pred_joints - gt_joints, dim=-1) # num_poses x num_joints=22
    mpjpe_seq = mpjpe.mean(-1) # num_poses

    return mpjpe_seq
//  [python]
def gen_one_ray_at(self, img_idx, x, y):
        """

        Parameters
        ----------
        img_idx :
        x : for width
        y : for height

        Returns
        -------

        """
        image = np.uint8(self.images_np[img_idx] * 256)
        image = cv.resize(image, fx=self.downsample_factor, fy=self.downsample_factor, dsize=None)
        image2 = cv.circle(image, (x, y), radius=10, color=(0, 0, 255), thickness=-1)

        pixels_x = torch.Tensor([x]).long()
        pixels_y = torch.Tensor([y]).long()
        color = self.images[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        mask = (self.masks[img_idx][(pixels_y, pixels_x)] > 0).to(torch.float32)  # 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(-1)  # 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(-1)  # 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.cuda(), rays_v.cuda(), color.cuda(), mask[:, :1].cuda()],
                         dim=-1), image2
//  [python]
def gen_one_ray_at(self, img_idx, x, y):
        """

        Parameters
        ----------
        img_idx :
        x : for width
        y : for height

        Returns
        -------

        """
        image = np.uint8(self.images_np[img_idx] * 256)
        image = cv.resize(image, fx=self.downsample_factor, fy=self.downsample_factor, dsize=None)
        image2 = cv.circle(image, (x, y), radius=10, color=(0, 0, 255), thickness=-1)

        pixels_x = torch.Tensor([x]).long()
        pixels_y = torch.Tensor([y]).long()
        color = self.images[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        mask = (self.masks[img_idx][(pixels_y, pixels_x)] > 0).to(torch.float32)  # 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(-1)  # 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(-1)  # 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.cuda(), rays_v.cuda(), color.cuda(), mask[:, :1].cuda()],
                         dim=-1), image2
//  [python]
def inside_sphere(self):
        self_grid_xyz = torch.stack(torch.meshgrid(
            torch.linspace(self.xyz_min[0], self.xyz_max[0], self.density.shape[2]),
            torch.linspace(self.xyz_min[1], self.xyz_max[1], self.density.shape[3]),
            torch.linspace(self.xyz_min[2], self.xyz_max[2], self.density.shape[4]),
        ), -1)
        sphere_mask = (torch.norm(self_grid_xyz, p=2, dim=-1, keepdim=True) < 1.0).reshape(*self.density.shape)
        self.density[~sphere_mask] = -100
//  [python]
def inside_sphere(self):
        self_grid_xyz = torch.stack(torch.meshgrid(
            torch.linspace(self.xyz_min[0], self.xyz_max[0], self.density.shape[2]),
            torch.linspace(self.xyz_min[1], self.xyz_max[1], self.density.shape[3]),
            torch.linspace(self.xyz_min[2], self.xyz_max[2], self.density.shape[4]),
        ), -1)
        sphere_mask = (torch.linalg.norm(self_grid_xyz, ord=2, dim=-1, keepdim=True) < 1.0).reshape(*self.density.shape)
        self.density[~sphere_mask] = -100
//  [python]
def newton_inv(self, A):
        lam = torch.norm(A @ A.T)
        alpha = 1/lam
        inv = alpha * A.T 
        eye = torch.eye(A.shape[0])
        
        for i in range(self.n_newton_steps):
            inv = inv @ (2*eye-A@inv)

        return inv
//  [python]
def newton_inv(self, A):
        lam = torch.linalg.norm(A @ A.T)
        alpha = 1/lam
        inv = alpha * A.T 
        eye = torch.eye(A.shape[0])
        
        for i in range(self.n_newton_steps):
            inv = inv @ (2*eye-A@inv)

        return inv
//  [python]
def test_joint_state():
    tensor_args = TensorDeviceType()
    j_names = ["j0", "j1", "j2", "j3"]
    loc_j = ["j4", "jb"]
    final_j = ["jb", "j0", "j1", "j2", "j3", "j4"]

    position = tensor_args.to_device([i for i in range(len(j_names))])

    loc_position = tensor_args.to_device([i + len(j_names) for i in range(len(loc_j))])

    js_1 = JointState.from_position(position, joint_names=j_names)
    js_lock = JointState.from_position(loc_position, loc_j)

    final_js = js_1.get_augmented_joint_state(final_j, js_lock)
    assert final_js.joint_names == final_j
    assert (
        t.norm(final_js.position - tensor_args.to_device([5, 0, 1, 2, 3, 4])).item()
        < 1e-8
    )
//  [python]
def test_joint_state():
    tensor_args = TensorDeviceType()
    j_names = ["j0", "j1", "j2", "j3"]
    loc_j = ["j4", "jb"]
    final_j = ["jb", "j0", "j1", "j2", "j3", "j4"]

    position = tensor_args.to_device([i for i in range(len(j_names))])

    loc_position = tensor_args.to_device([i + len(j_names) for i in range(len(loc_j))])

    js_1 = JointState.from_position(position, joint_names=j_names)
    js_lock = JointState.from_position(loc_position, loc_j)

    final_js = js_1.get_augmented_joint_state(final_j, js_lock)
    assert final_js.joint_names == final_j
    assert (
        torch.linalg.norm(final_js.position - tensor_args.to_device([5, 0, 1, 2, 3, 4])).item()
        < 1e-8
    )
//  [python]
def angles_deg_and_distances(xy_coords):
    distances = torch.norm(xy_coords, dim=1)
    angular = direction_to_angle_deg(xy_coords[:, 0], xy_coords[:, 1])
    return distances, angular
//  [python]
def angles_deg_and_distances(xy_coords):
    distances = torch.linalg.norm(xy_coords, axis=1)
    angular = direction_to_angle_deg(xy_coords[:, 0], xy_coords[:, 1])
    return distances, angular
//  [python]
def batch_proj(x, line):
    # x:[batch,3], line:[batch,N,3]
    line_length = line.shape[-2]
    batch_dim = x.ndim - 1
    if isinstance(x, torch.Tensor):
        delta = line[..., 0:2] - torch.unsqueeze(x[..., 0:2], dim=-2).repeat(
            *([1] * batch_dim), line_length, 1
        )
        dis = torch.norm(delta, dim=-1)
        idx0 = torch.argmin(dis, dim=-1)
        idx = idx0.view(*line.shape[:-2], 1, 1).repeat(
            *([1] * (batch_dim + 1)), line.shape[-1]
        )
        line_min = torch.squeeze(torch.gather(line, -2, idx), dim=-2)
        dx = x[..., None, 0] - line[..., 0]
        dy = x[..., None, 1] - line[..., 1]
        delta_y = -dx * torch.sin(line_min[..., None, 2]) + dy * torch.cos(
            line_min[..., None, 2]
        )
        delta_x = dx * torch.cos(line_min[..., None, 2]) + dy * torch.sin(
            line_min[..., None, 2]
        )

        delta_psi = angle_wrap(x[..., 2] - line_min[..., 2])

        return (
            delta_x,
            delta_y,
            torch.unsqueeze(delta_psi, dim=-1),
        )

    elif isinstance(x, np.ndarray):
        delta = line[..., 0:2] - np.repeat(
            x[..., np.newaxis, 0:2], line_length, axis=-2
        )
        dis = np.linalg.norm(delta, axis=-1)
        idx0 = np.argmin(dis, axis=-1)
        idx = idx0.reshape(*line.shape[:-2], 1, 1).repeat(line.shape[-1], axis=-1)
        line_min = np.squeeze(np.take_along_axis(line, idx, axis=-2), axis=-2)
        dx = x[..., None, 0] - line[..., 0]
        dy = x[..., None, 1] - line[..., 1]
        delta_y = -dx * np.sin(line_min[..., None, 2]) + dy * np.cos(
            line_min[..., None, 2]
        )
        delta_x = dx * np.cos(line_min[..., None, 2]) + dy * np.sin(
            line_min[..., None, 2]
        )
        delta_psi = angle_wrap(x[..., 2] - line_min[..., 2])
        return (
            delta_x,
            delta_y,
            np.expand_dims(delta_psi, axis=-1),
        )
//  [python]
def batch_proj(x, line):
    # x:[batch,3], line:[batch,N,3]
    line_length = line.shape[-2]
    batch_dim = x.ndim - 1
    if isinstance(x, torch.Tensor):
        delta = line[..., 0:2] - torch.unsqueeze(x[..., 0:2], dim=-2).repeat(
            *([1] * batch_dim), line_length, 1
        )
        dis = torch.linalg.norm(delta, axis=-1)
        idx0 = torch.argmin(dis, dim=-1)
        idx = idx0.view(*line.shape[:-2], 1, 1).repeat(
            *([1] * (batch_dim + 1)), line.shape[-1]
        )
        line_min = torch.squeeze(torch.gather(line, -2, idx), dim=-2)
        dx = x[..., None, 0] - line[..., 0]
        dy = x[..., None, 1] - line[..., 1]
        delta_y = -dx * torch.sin(line_min[..., None, 2]) + dy * torch.cos(
            line_min[..., None, 2]
        )
        delta_x = dx * torch.cos(line_min[..., None, 2]) + dy * torch.sin(
            line_min[..., None, 2]
        )

        delta_psi = angle_wrap(x[..., 2] - line_min[..., 2])

        return (
            delta_x,
            delta_y,
            torch.unsqueeze(delta_psi, dim=-1),
        )

    elif isinstance(x, np.ndarray):
        delta = line[..., 0:2] - np.repeat(
            x[..., np.newaxis, 0:2], line_length, axis=-2
        )
        dis = np.linalg.norm(delta, axis=-1)
        idx0 = np.argmin(dis, axis=-1)
        idx = idx0.reshape(*line.shape[:-2], 1, 1).repeat(line.shape[-1], axis=-1)
        line_min = np.squeeze(np.take_along_axis(line, idx, axis=-2), axis=-2)
        dx = x[..., None, 0] - line[..., 0]
        dy = x[..., None, 1] - line[..., 1]
        delta_y = -dx * np.sin(line_min[..., None, 2]) + dy * np.cos(
            line_min[..., None, 2]
        )
        delta_x = dx * np.cos(line_min[..., None, 2]) + dy * np.sin(
            line_min[..., None, 2]
        )
        delta_psi = angle_wrap(x[..., 2] - line_min[..., 2])
        return (
            delta_x,
            delta_y,
            np.expand_dims(delta_psi, axis=-1),
        )
//  [python]
def test_batch_joint_state():
    tensor_args = TensorDeviceType()
    j_names = ["j0", "j1", "j2", "j3"]
    loc_j = ["j4", "jb"]
    final_j = ["jb", "j0", "j1", "j2", "j3", "j4"]

    # $position = tensor_args.to_device([i for i in range(len(j_names))])
    position = torch.zeros((10, len(j_names)), device=tensor_args.device, dtype=tensor_args.dtype)
    for i in range(len(j_names)):
        position[:, i] = i

    loc_position = tensor_args.to_device([i + len(j_names) for i in range(len(loc_j))])

    js_1 = JointState.from_position(position, joint_names=j_names)
    js_lock = JointState.from_position(loc_position, loc_j)

    final_js = js_1.get_augmented_joint_state(final_j, js_lock)
    assert final_js.joint_names == final_j
    assert (
        torch.norm(
            final_js.position - tensor_args.to_device([5, 0, 1, 2, 3, 4]).unsqueeze(0)
        ).item()
        < 1e-8
    )
//  [python]
def test_batch_joint_state():
    tensor_args = TensorDeviceType()
    j_names = ["j0", "j1", "j2", "j3"]
    loc_j = ["j4", "jb"]
    final_j = ["jb", "j0", "j1", "j2", "j3", "j4"]

    # $position = tensor_args.to_device([i for i in range(len(j_names))])
    position = torch.zeros((10, len(j_names)), device=tensor_args.device, dtype=tensor_args.dtype)
    for i in range(len(j_names)):
        position[:, i] = i

    loc_position = tensor_args.to_device([i + len(j_names) for i in range(len(loc_j))])

    js_1 = JointState.from_position(position, joint_names=j_names)
    js_lock = JointState.from_position(loc_position, loc_j)

    final_js = js_1.get_augmented_joint_state(final_j, js_lock)
    assert final_js.joint_names == final_j
    assert (
        torch.linalg.norm(
            final_js.position - tensor_args.to_device([5, 0, 1, 2, 3, 4]).unsqueeze(0)
        ).item()
        < 1e-8
    )
//  [python]
def get_spherical_coords(X):
    """
    Returns the spherical coordinates of 3D points written in cartesian coordinates
    :param X: (Tensor) Tensor with shape (N, 3) that represents 3D points in cartesian coordinates.
    :return: (3-tuple of Tensors) r_x, elev_x and azim_x are Tensors with shape (N) that corresponds
    to radius, elevation and azimuths of all 3D points.
    """
    r_x = torch.norm(X, dim=1)

    elev_x = torch.asin(X[:, 1] / r_x)  # between -pi/2 and pi/2
    elev_x[X[:, 1] / r_x <= -1] = -np.pi / 2
    elev_x[X[:, 1] / r_x >= 1] = np.pi / 2

    azim_x = torch.acos(X[:, 2] / (r_x * torch.cos(elev_x)))
    azim_x[X[:, 2] / (r_x * torch.cos(elev_x)) <= -1] = np.pi
    azim_x[X[:, 2] / (r_x * torch.cos(elev_x)) >= 1] = 0.
    azim_x[X[:, 0] < 0] *= -1

    return r_x, elev_x, azim_x
//  [python]
def get_spherical_coords(X):
    """
    Returns the spherical coordinates of 3D points written in cartesian coordinates
    :param X: (Tensor) Tensor with shape (N, 3) that represents 3D points in cartesian coordinates.
    :return: (3-tuple of Tensors) r_x, elev_x and azim_x are Tensors with shape (N) that corresponds
    to radius, elevation and azimuths of all 3D points.
    """
    r_x = torch.linalg.norm(X, dim=1)

    elev_x = torch.asin(X[:, 1] / r_x)  # between -pi/2 and pi/2
    elev_x[X[:, 1] / r_x <= -1] = -np.pi / 2
    elev_x[X[:, 1] / r_x >= 1] = np.pi / 2

    azim_x = torch.acos(X[:, 2] / (r_x * torch.cos(elev_x)))
    azim_x[X[:, 2] / (r_x * torch.cos(elev_x)) <= -1] = np.pi
    azim_x[X[:, 2] / (r_x * torch.cos(elev_x)) >= 1] = 0.
    azim_x[X[:, 0] < 0] *= -1

    return r_x, elev_x, azim_x
//  [python]
def _forpdns(img, t: torch.Tensor, rowsp, colsp, bandsp, sigma, diag_s, c1, c12u, c12v, yp_m):
        # signal based decomposition
        x1 = img.reshape((rowsp * colsp, bandsp))
        # x1 = x1.repeat(torch.pow(sigma, 2), (rowsp * colsp, 1)) * x1
        sure = torch.zeros((t.shape[0], 2), dtype=img.dtype, device=img.device)
        for i, lam in enumerate(t):
            sig_lam_eye = torch.diag(1.0 / (1.0 / (lam * diag_s) + 1.0))
            majc = c1 - c12u @ (sig_lam_eye @ c12v)
            xp_m = (majc @ x1.T).T
            norm1 = torch.norm(yp_m - xp_m, 2)
            sure[i, 0] = norm1 / (rowsp * colsp) + 2 * torch.trace(majc)
            if i > 1 and sure[i, 0] > sure[i - 1, 0]:
                break
        opt_params = t[torch.argmin(sure)]
        xp = xp_m.reshape((rowsp, colsp, bandsp))
        return xp, sigma, opt_params, sure
//  [python]
def _forpdns(img, t: torch.Tensor, rowsp, colsp, bandsp, sigma, diag_s, c1, c12u, c12v, yp_m):
        # signal based decomposition
        x1 = img.reshape((rowsp * colsp, bandsp))
        # x1 = x1.repeat(torch.pow(sigma, 2), (rowsp * colsp, 1)) * x1
        sure = torch.zeros((t.shape[0], 2), dtype=img.dtype, device=img.device)
        for i, lam in enumerate(t):
            sig_lam_eye = torch.diag(1.0 / (1.0 / (lam * diag_s) + 1.0))
            majc = c1 - c12u @ (sig_lam_eye @ c12v)
            xp_m = (majc @ x1.T).T
            norm1 = torch.linalg.norm(yp_m - xp_m, ord=2)
            sure[i, 0] = norm1 / (rowsp * colsp) + 2 * torch.trace(majc)
            if i > 1 and sure[i, 0] > sure[i - 1, 0]:
                break
        opt_params = t[torch.argmin(sure)]
        xp = xp_m.reshape((rowsp, colsp, bandsp))
        return xp, sigma, opt_params, sure
//  [python]
def quat_to_rotvec(quat: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    # 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 quat_to_rotvec(quat: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    # 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 check(self, value):
        tol = (
            torch.finfo(value.dtype).eps * value.size(-1) * 10
        )  # 10 is an adjustable fudge factor
        row_norm = torch.norm(value.detach(), dim=-1)
        unit_row_norm = (row_norm - 1.0).abs().le(tol).all(dim=-1)
        return _LowerCholesky().check(value) & unit_row_norm
//  [python]
def check(self, value):
        tol = (
            torch.finfo(value.dtype).eps * value.size(-1) * 10
        )  # 10 is an adjustable fudge factor
        row_norm = torch.linalg.norm(value.detach(), dim=-1)
        unit_row_norm = (row_norm - 1.0).abs().le(tol).all(dim=-1)
        return _LowerCholesky().check(value) & unit_row_norm
//  [python]
def test_compute_pair_vector(self, graph_Mo):
        s1, g1, state1 = graph_Mo
        lattice = torch.tensor(s1.lattice.matrix, dtype=matgl.float_th).unsqueeze(dim=0)
        g1.edata["pbc_offshift"] = torch.matmul(g1.edata["pbc_offset"], lattice[0])
        g1.ndata["pos"] = g1.ndata["frac_coords"] @ lattice[0]
        bv, bd = compute_pair_vector_and_distance(g1)
        g1.edata["bond_vec"] = bv
        d = torch.norm(g1.edata["bond_vec"], dim=1)

        _, _, _, d2 = s1.get_neighbor_list(r=5.0)

        np.testing.assert_array_almost_equal(np.sort(d), np.sort(d2))
//  [python]
def test_compute_pair_vector(self, graph_Mo):
        s1, g1, state1 = graph_Mo
        lattice = torch.tensor(s1.lattice.matrix, dtype=matgl.float_th).unsqueeze(dim=0)
        g1.edata["pbc_offshift"] = torch.matmul(g1.edata["pbc_offset"], lattice[0])
        g1.ndata["pos"] = g1.ndata["frac_coords"] @ lattice[0]
        bv, bd = compute_pair_vector_and_distance(g1)
        g1.edata["bond_vec"] = bv
        d = torch.linalg.norm(g1.edata["bond_vec"], axis=1)

        _, _, _, d2 = s1.get_neighbor_list(r=5.0)

        np.testing.assert_array_almost_equal(np.sort(d), np.sort(d2))
//  [python]
def normalization_inverse_sqrt_dist_centered(x_in_world, view_cell_center, max_depth):
    localized = x_in_world - view_cell_center
    local = torch.sqrt(torch.norm(localized, dim=-1))
    res = localized / (math.sqrt(max_depth) * local[..., None])
    return res
//  [python]
def normalization_inverse_sqrt_dist_centered(x_in_world, view_cell_center, max_depth):
    localized = x_in_world - view_cell_center
    local = torch.sqrt(torch.linalg.norm(localized, dim=-1))
    res = localized / (math.sqrt(max_depth) * local[..., None])
    return res
//  [python]
def norm_matrix(x):
    return torch.norm(x)
//  [python]
def norm_matrix(x):
    return torch.linalg.norm(x)
//  [python]
def jac_mag(self):
        """Magnitude of jacobian of objective function at current iteration."""
        if self._g_mag is None:
            self._g_mag = t.norm(self.jac)
        return self._g_mag
//  [python]
def jac_mag(self):
        """Magnitude of jacobian of objective function at current iteration."""
        if self._g_mag is None:
            self._g_mag = norm(self.jac)
        return self._g_mag
//  [python]
def test_compute_pair_vector_for_molecule(self, graph_CH4):
        s2, g2, state2 = graph_CH4
        lattice = torch.tensor(np.identity(3), dtype=matgl.float_th).unsqueeze(dim=0)
        g2.edata["pbc_offshift"] = torch.matmul(g2.edata["pbc_offset"], lattice[0])
        g2.ndata["pos"] = g2.ndata["frac_coords"] @ lattice[0]
        bv, bd = compute_pair_vector_and_distance(g2)
        g2.edata["bond_vec"] = bv
        d = torch.norm(g2.edata["bond_vec"], dim=1)

        d2 = np.array(
            [
                1.089,
                1.089,
                1.089,
                1.089,
                1.089,
                1.089,
                1.089,
                1.089,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
            ]
        )

        np.testing.assert_array_almost_equal(np.sort(d), np.sort(d2))
//  [python]
def test_compute_pair_vector_for_molecule(self, graph_CH4):
        s2, g2, state2 = graph_CH4
        lattice = torch.tensor(np.identity(3), dtype=matgl.float_th).unsqueeze(dim=0)
        g2.edata["pbc_offshift"] = torch.matmul(g2.edata["pbc_offset"], lattice[0])
        g2.ndata["pos"] = g2.ndata["frac_coords"] @ lattice[0]
        bv, bd = compute_pair_vector_and_distance(g2)
        g2.edata["bond_vec"] = bv
        d = torch.linalg.norm(g2.edata["bond_vec"], axis=1)

        d2 = np.array(
            [
                1.089,
                1.089,
                1.089,
                1.089,
                1.089,
                1.089,
                1.089,
                1.089,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
                1.77833,
            ]
        )

        np.testing.assert_array_almost_equal(np.sort(d), np.sort(d2))
//  [python]
def Jr(self, x):
        """
        Right jacobian of so(3)
        """
        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.lshape+(3, 3))
        Jr = I - (1-theta.cos())/theta**2 * K + (theta - theta.sin())/theta**3 * K@K
        return torch.where(theta>torch.finfo(theta.dtype).eps, Jr, I)
//  [python]
def Jr(self, x):
        """
        Right jacobian of so(3)
        """
        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.lshape+(3, 3))
        Jr = I - (1-theta.cos())/theta**2 * K + (theta - theta.sin())/theta**3 * K@K
        return torch.where(theta>torch.finfo(theta.dtype).eps, Jr, I)
//  [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.norm(p, p=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