jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def gaussian_kernel(X, sigma2=None, Y=None, normalized=False, **ignored):
    if normalized:
        X = X / torch.norm(X, dim=1, keepdim=True)
        if Y is not None:
            Y = Y / torch.norm(Y, dim=1, keepdim=True)
    Dxx = compute_pdist_sq(X, Y)
    if sigma2 is None:
        sigma2 = Dxx.median()
    Kx = torch.exp(-Dxx / sigma2)
    return Kx
//  [python]
def gaussian_kernel(X, sigma2=None, Y=None, normalized=False, **ignored):
    if normalized:
        X = X / torch.linalg.norm(X, dim=1, keepdim=True)
        if Y is not None:
            Y = Y / torch.linalg.norm(Y, dim=1, keepdim=True)
    Dxx = compute_pdist_sq(X, Y)
    if sigma2 is None:
        sigma2 = Dxx.median()
    Kx = torch.exp(-Dxx / sigma2)
    return Kx
//  [python]
def _contains(self, points, params=Points.empty()):
        center, radius = self.domain._compute_center_and_radius(points.join(params), points.device)
        points = points[:, list(self.space.keys())].as_tensor
        norm = torch.norm(points - center, dim=1).reshape(-1, 1)
        return torch.isclose(norm[:, None], radius).reshape(-1, 1)
//  [python]
def _contains(self, points, params=Points.empty()):
        center, radius = self.domain._compute_center_and_radius(points.join(params), points.device)
        points = points[:, list(self.space.keys())].as_tensor
        norm = torch.linalg.norm(points - center, dim=1).reshape(-1, 1)
        return torch.isclose(norm[:, None], radius).reshape(-1, 1)
//  [python]
def vec_normalize(x: torch.Tensor) -> torch.Tensor:
    """
    normalize vectors at the last dimension
    """
    result = x / torch.norm(x, 2, -1, keepdim=True).clamp(1e-9)
    return result
//  [python]
def vec_normalize(x: torch.Tensor) -> torch.Tensor:
    """
    normalize vectors at the last dimension
    """
    result = x / torch.linalg.norm(x, 2, -1, keepdim=True).clamp(1e-9)
    return result
//  [python]
def rel_loss(self, v1, v2, v1_gt, v2_gt):
        rel_gt = v1.unsqueeze(1) - v2.unsqueeze(2)
        rel_gt = torch.norm(rel_gt, dim=-1)  # bs x V x V
        rel_pred = v1_gt.unsqueeze(1) - v2_gt.unsqueeze(2)
        rel_pred = torch.norm(rel_pred, dim=-1)  # bs x 21 x 21
        return self.L1Loss(rel_gt, rel_pred)
//  [python]
def rel_loss(self, v1, v2, v1_gt, v2_gt):
        rel_gt = v1.unsqueeze(1) - v2.unsqueeze(2)
        rel_gt = torch.linalg.norm(rel_gt, dim=-1)  # bs x V x V
        rel_pred = v1_gt.unsqueeze(1) - v2_gt.unsqueeze(2)
        rel_pred = torch.linalg.norm(rel_pred, dim=-1)  # bs x 21 x 21
        return self.L1Loss(rel_gt, rel_pred)
//  [python]
def edge_loss(self, verts_pred, verts_gt):
        edge_gt = verts_gt[:, self.faces]
        edge_gt = torch.stack([edge_gt[:, :, 0] - edge_gt[:, :, 1],
                               edge_gt[:, :, 1] - edge_gt[:, :, 2],
                               edge_gt[:, :, 2] - edge_gt[:, :, 0],
                               ], dim=2)  # B x F x 3 x 3
        edge_pred = verts_pred[:, self.faces]
        edge_pred = torch.stack([edge_pred[:, :, 0] - edge_pred[:, :, 1],
                                 edge_pred[:, :, 1] - edge_pred[:, :, 2],
                                 edge_pred[:, :, 2] - edge_pred[:, :, 0],
                                 ], dim=2)  # B x F x 3 x 3
        edge_length_gt = torch.norm(edge_gt, dim=-1)  # B x F x 3
        edge_length_pred = torch.norm(edge_pred, dim=-1)  # B x F x 3
        edge_length_loss = self.L1Loss(edge_length_pred, edge_length_gt)
        return edge_length_loss
//  [python]
def edge_loss(self, verts_pred, verts_gt):
        edge_gt = verts_gt[:, self.faces]
        edge_gt = torch.stack([edge_gt[:, :, 0] - edge_gt[:, :, 1],
                               edge_gt[:, :, 1] - edge_gt[:, :, 2],
                               edge_gt[:, :, 2] - edge_gt[:, :, 0],
                               ], dim=2)  # B x F x 3 x 3
        edge_pred = verts_pred[:, self.faces]
        edge_pred = torch.stack([edge_pred[:, :, 0] - edge_pred[:, :, 1],
                                 edge_pred[:, :, 1] - edge_pred[:, :, 2],
                                 edge_pred[:, :, 2] - edge_pred[:, :, 0],
                                 ], dim=2)  # B x F x 3 x 3
        edge_length_gt = torch.linalg.norm(edge_gt, dim=-1)  # B x F x 3
        edge_length_pred = torch.linalg.norm(edge_pred, dim=-1)  # B x F x 3
        edge_length_loss = self.L1Loss(edge_length_pred, edge_length_gt)
        return edge_length_loss
//  [python]
def rotation_matrix(
    a: Float[Tensor, "3"], b: Float[Tensor, "3"]
) -> Float[Tensor, "3 3"]:
    """Compute the rotation matrix that rotates vector a to vector b.

    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    """
    a = a / t.norm(a)
    b = b / t.norm(b)
    v = t.cross(a, b)
    c = t.dot(a, b)
    # If vectors are exactly opposite, we add a little noise to one of them
    if c < -1 + 1e-8:
        eps = (torch.rand(3) - 0.5) * 0.01
        return rotation_matrix(a + eps, b)
    s = t.norm(v)
    skew_sym_mat = torch.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    return (
        torch.eye(3)
        + skew_sym_mat
        + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
    )
//  [python]
def rotation_matrix(
    a: Float[Tensor, "3"], b: Float[Tensor, "3"]
) -> Float[Tensor, "3 3"]:
    """Compute the rotation matrix that rotates vector a to vector b.

    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    """
    a = a / torch.linalg.norm(a)
    b = b / torch.linalg.norm(b)
    v = torch.linalg.cross(a, b)
    c = torch.dot(a, b)
    # If vectors are exactly opposite, we add a little noise to one of them
    if c < -1 + 1e-8:
        eps = (torch.rand(3) - 0.5) * 0.01
        return rotation_matrix(a + eps, b)
    s = torch.linalg.norm(v)
    skew_sym_mat = torch.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    return (
        torch.eye(3)
        + skew_sym_mat
        + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
    )
//  [python]
def compute_delta_angle(y_pred: torch.Tensor, y_gt: torch.Tensor, rotate_type: RotateType):
    if rotate_type == RotateType.Vec6d:
        delta_angle = torch.abs(diff_angle_between_vec6d(y_pred, y_gt))
    elif rotate_type == RotateType.SVD9d or rotate_type == RotateType.Matrix:
        delta_angle = torch.abs(diff_angle_bewteen_matrix(y_pred, y_gt))
    elif rotate_type == RotateType.AxisAngle:
        delta_angle = torch.norm(y_pred - y_gt, dim=-1)
    elif rotate_type == RotateType.Quaternion:
        raise NotImplementedError
    else:
        raise NotImplementedError

    return delta_angle
//  [python]
def compute_delta_angle(y_pred: torch.Tensor, y_gt: torch.Tensor, rotate_type: RotateType):
    if rotate_type == RotateType.Vec6d:
        delta_angle = torch.abs(diff_angle_between_vec6d(y_pred, y_gt))
    elif rotate_type == RotateType.SVD9d or rotate_type == RotateType.Matrix:
        delta_angle = torch.abs(diff_angle_bewteen_matrix(y_pred, y_gt))
    elif rotate_type == RotateType.AxisAngle:
        delta_angle = torch.linalg.norm(y_pred - y_gt, dim=-1)
    elif rotate_type == RotateType.Quaternion:
        raise NotImplementedError
    else:
        raise NotImplementedError

    return delta_angle
//  [python]
def normalize_vec6d(x: torch.Tensor):  # TODO: Test
    assert x.shape[-2:] == (3, 2)
    x: torch.Tensor = x / torch.norm(x, dim=-2, keepdims=True)

    first_col: torch.Tensor = x[..., 0].contiguous()
    second_col: torch.Tensor = x[..., 1].contiguous()
    last_col: torch.Tensor = torch.cross(first_col, second_col, dim=-1)
    last_col: torch.Tensor = last_col / torch.norm(last_col, dim=-1, keepdims=True)

    second_col: torch.Tensor = torch.cross(-first_col, last_col, dim=-1)
    second_col: torch.Tensor = second_col / torch.norm(second_col, dim=-1, keepdims=True)

    return first_col, second_col, last_col
//  [python]
def normalize_vec6d(x: torch.Tensor):  # TODO: Test
    assert x.shape[-2:] == (3, 2)
    x: torch.Tensor = x / torch.linalg.norm(x, dim=-2, keepdims=True)

    first_col: torch.Tensor = x[..., 0].contiguous()
    second_col: torch.Tensor = x[..., 1].contiguous()
    last_col: torch.Tensor = torch.cross(first_col, second_col, dim=-1)
    last_col: torch.Tensor = last_col / torch.linalg.norm(last_col, dim=-1, keepdims=True)

    second_col: torch.Tensor = torch.cross(-first_col, last_col, dim=-1)
    second_col: torch.Tensor = second_col / torch.linalg.norm(second_col, dim=-1, keepdims=True)

    return first_col, second_col, last_col
//  [python]
def calc_distogram(pos: torch.Tensor, min_bin: float, max_bin: float, num_bins: int):
    # pos: [*, L, 3]
    dists_2d = torch.norm(
        pos[..., :, None, :] - pos[..., None, :, :], dim=-1
    )[..., None]
    lower = torch.linspace(
        min_bin,
        max_bin,
        num_bins,
        device=pos.device)
    upper = torch.cat([lower[1:], lower.new_tensor([1e8])], dim=-1)
    distogram = ((dists_2d > lower) * (dists_2d < upper)).type(pos.dtype)
    return distogram
//  [python]
def calc_distogram(pos: torch.Tensor, min_bin: float, max_bin: float, num_bins: int):
    # pos: [*, L, 3]
    dists_2d = torch.linalg.norm(
        pos[..., :, None, :] - pos[..., None, :, :], axis=-1
    )[..., None]
    lower = torch.linspace(
        min_bin,
        max_bin,
        num_bins,
        device=pos.device)
    upper = torch.cat([lower[1:], lower.new_tensor([1e8])], dim=-1)
    distogram = ((dists_2d > lower) * (dists_2d < upper)).type(pos.dtype)
    return distogram
//  [python]
def quat_to_rotvec(q: torch.Tensor, do_normalize: bool = False):  # Test OK
    """
    Modified from scipy.spatial.transform.Rotation.as_rotvec
    Convert quaternion to rot vec
    return: rotvec
    """
    assert q.shape[-1] == 4 and q.shape[0] > 0
    if do_normalize:
        q = quat_normalize(q)

    quat: torch.Tensor = flip_quat_by_w(q)
    angle: torch.Tensor = torch.as_tensor(2.0) * torch.atan2(torch.norm(quat[:, :3], dim=1), quat[:, 3])

    eps = 1e-3
    small_angle: torch.Tensor = torch.as_tensor(angle <= eps)
    large_angle: torch.Tensor = torch.as_tensor(~small_angle)

    scale_small = scale_large = None
    if torch.any(small_angle):
        scale_small = (2 + angle ** 2 / 12 + 7 * angle ** 4 / 2880)
    if torch.any(large_angle):
        scale_large = angle / torch.sin(angle / 2)
    if scale_small is None:
        scale = scale_large
    elif scale_large is None:
        scale = scale_small
    else:
        scale = torch.where(small_angle,scale_small, scale_large)
        

    rotvec: torch.Tensor = scale[:, None] * quat[:, :3]

    return rotvec.view(-1) if q.shape == (4,) else rotvec
//  [python]
def quat_to_rotvec(q: torch.Tensor, do_normalize: bool = False):  # Test OK
    """
    Modified from scipy.spatial.transform.Rotation.as_rotvec
    Convert quaternion to rot vec
    return: rotvec
    """
    assert q.shape[-1] == 4 and q.shape[0] > 0
    if do_normalize:
        q = quat_normalize(q)

    quat: torch.Tensor = flip_quat_by_w(q)
    angle: torch.Tensor = torch.as_tensor(2.0) * torch.atan2(torch.linalg.norm(quat[:, :3], dim=1), quat[:, 3])

    eps = 1e-3
    small_angle: torch.Tensor = torch.as_tensor(angle <= eps)
    large_angle: torch.Tensor = torch.as_tensor(~small_angle)

    scale_small = scale_large = None
    if torch.any(small_angle):
        scale_small = (2 + angle ** 2 / 12 + 7 * angle ** 4 / 2880)
    if torch.any(large_angle):
        scale_large = angle / torch.sin(angle / 2)
    if scale_small is None:
        scale = scale_large
    elif scale_large is None:
        scale = scale_small
    else:
        scale = torch.where(small_angle,scale_small, scale_large)
        

    rotvec: torch.Tensor = scale[:, None] * quat[:, :3]

    return rotvec.view(-1) if q.shape == (4,) else rotvec
//  [python]
def quat_from_rotvec(rotvec: torch.Tensor) -> torch.Tensor:
    """
    Modified from scipy.spatial.transform.Rotation.from_rotvec() method
    Convert rotvec to quaternion

    return: quaternion in torch.Tensor
    """
    if rotvec.ndim not in [1, 2] or rotvec.shape[-1] != 3:
        raise ValueError("Expected `rot_vec` to have shape (3,) or (N, 3), got {}".format(rotvec.shape))

    if rotvec.shape == (3,):
        rotvec = rotvec[None, :]

    norms: torch.Tensor = torch.norm(rotvec, dim=1)
    small_angle = torch.as_tensor(norms <= 1e-3)
    large_angle = torch.as_tensor(~small_angle)

    # for hack..
    # small_angle = torch.as_tensor(norm > -10000)
    # large_angle = torch.as_tensor(~small_angle)

    scale_small = scale_large = None
    if torch.any(small_angle):
        scale_small = (0.5 - norms ** 2 / 48 + norms ** 4 / 3840)
    if torch.any(large_angle):
        scale_large = (torch.sin(norms / 2) / norms)

    if scale_small is None:
        scale = scale_large
    elif scale_large is None:
        scale = scale_small
    else:
        scale = torch.where(small_angle, scale_small, scale_large)
        # scale = CatWithMask.apply(scale_small, scale_large, small_angle)

    quat_xyz = scale[:, None] * rotvec  # [..., :3]
    quat_w = torch.cos(torch.as_tensor(0.5) * norms)[..., None]  # [..., 3, None]
    quat = torch.cat([quat_xyz, quat_w], dim=-1)

    return quat.view(-1) if rotvec.size == 3 else quat
//  [python]
def quat_from_rotvec(rotvec: torch.Tensor) -> torch.Tensor:
    """
    Modified from scipy.spatial.transform.Rotation.from_rotvec() method
    Convert rotvec to quaternion

    return: quaternion in torch.Tensor
    """
    if rotvec.ndim not in [1, 2] or rotvec.shape[-1] != 3:
        raise ValueError("Expected `rot_vec` to have shape (3,) or (N, 3), got {}".format(rotvec.shape))

    if rotvec.shape == (3,):
        rotvec = rotvec[None, :]

    norms: torch.Tensor = torch.linalg.norm(rotvec, axis=1)
    small_angle = torch.as_tensor(norms <= 1e-3)
    large_angle = torch.as_tensor(~small_angle)

    # for hack..
    # small_angle = torch.as_tensor(norm > -10000)
    # large_angle = torch.as_tensor(~small_angle)

    scale_small = scale_large = None
    if torch.any(small_angle):
        scale_small = (0.5 - norms ** 2 / 48 + norms ** 4 / 3840)
    if torch.any(large_angle):
        scale_large = (torch.sin(norms / 2) / norms)

    if scale_small is None:
        scale = scale_large
    elif scale_large is None:
        scale = scale_small
    else:
        scale = torch.where(small_angle, scale_small, scale_large)
        # scale = CatWithMask.apply(scale_small, scale_large, small_angle)

    quat_xyz = scale[:, None] * rotvec  # [..., :3]
    quat_w = torch.cos(torch.as_tensor(0.5) * norms)[..., None]  # [..., 3, None]
    quat = torch.cat([quat_xyz, quat_w], dim=-1)

    return quat.view(-1) if rotvec.size == 3 else quat
//  [python]
def loss_function(self, recon_x, x, z):
        recon_loss = F.mse_loss(
            recon_x.reshape(x.shape[0], -1), x.reshape(x.shape[0], -1), reduction="none"
        ).sum(dim=-1)

        gen_reg_loss = self._compute_gp(recon_x, x)

        embedding_loss = 0.5 * torch.norm(z, dim=-1) ** 2

        return (
            (
                recon_loss
                + self.model_config.embedding_weight * embedding_loss
                + self.model_config.reg_weight * gen_reg_loss
            ).mean(dim=0),
            (recon_loss).mean(dim=0),
            (gen_reg_loss).mean(dim=0),
            (embedding_loss).mean(dim=0),
        )
//  [python]
def loss_function(self, recon_x, x, z):
        recon_loss = F.mse_loss(
            recon_x.reshape(x.shape[0], -1), x.reshape(x.shape[0], -1), reduction="none"
        ).sum(dim=-1)

        gen_reg_loss = self._compute_gp(recon_x, x)

        embedding_loss = 0.5 * torch.linalg.norm(z, dim=-1) ** 2

        return (
            (
                recon_loss
                + self.model_config.embedding_weight * embedding_loss
                + self.model_config.reg_weight * gen_reg_loss
            ).mean(dim=0),
            (recon_loss).mean(dim=0),
            (gen_reg_loss).mean(dim=0),
            (embedding_loss).mean(dim=0),
        )
//  [python]
def compute_A_geodesic(x_0,
                       x_1,
                       A,
                       t,
                       lambd # Strength of penalty on missing x_1
                      ):
    """Compute an A-geodesic x(t) between x_0 and x_1 parametrized over t \in [0,1].
    Report points x(t) at user-supplied times t.
    The last time t must be 1."""
    
    # Initialize the model for velocity field v
    space_dims = 2
    v = VelocityField(space_dims).to(device)
    
    # Initialize optimizer
    lr = 1e-3
    weight_decay = 0
    opt = torch.optim.AdamW(v.parameters(), lr=lr, weight_decay=weight_decay)
    
    n_epochs = 500
    losses = []
    # Train the velocity field v to push x_0 to x_1 along an A-geodesic
    for epoch in range(n_epochs):
        opt.zero_grad()
        x_t = odeint(v, x_0, t, method='dopri5')
        target_penalty = lambd*torch.norm(x_t[-1] - x_1)
        # Compute kinetic energy loss
        A_at_x = A(x_t)
        v_at_x = torch.unsqueeze(v(t, x_t), dim=2)
        #kinetic_energy = torch.mean(torch.sum(torch.squeeze(v_at_x * (A_at_x @ v_at_x)), dim=1))
        kinetic_energy = torch.mean(torch.sum(torch.squeeze(v_at_x * torch.linalg.solve(A_at_x, v_at_x)), dim=1)) # v(t,x(t))^T @ A^-1(x(t)) @ v(t,x(t))
        #print(kinetic_energy)
        # Compute total loss and backprop
        loss = kinetic_energy + target_penalty
        losses.append(loss.detach().cpu())
        loss.backward()
        opt.step()
    
    # Report points x(t) for final geodesic
    return odeint(v, x_0, t), losses
//  [python]
def compute_A_geodesic(x_0,
                       x_1,
                       A,
                       t,
                       lambd # Strength of penalty on missing x_1
                      ):
    """Compute an A-geodesic x(t) between x_0 and x_1 parametrized over t \in [0,1].
    Report points x(t) at user-supplied times t.
    The last time t must be 1."""
    
    # Initialize the model for velocity field v
    space_dims = 2
    v = VelocityField(space_dims).to(device)
    
    # Initialize optimizer
    lr = 1e-3
    weight_decay = 0
    opt = torch.optim.AdamW(v.parameters(), lr=lr, weight_decay=weight_decay)
    
    n_epochs = 500
    losses = []
    # Train the velocity field v to push x_0 to x_1 along an A-geodesic
    for epoch in range(n_epochs):
        opt.zero_grad()
        x_t = odeint(v, x_0, t, method='dopri5')
        target_penalty = lambd*torch.linalg.norm(x_t[-1] - x_1)
        # Compute kinetic energy loss
        A_at_x = A(x_t)
        v_at_x = torch.unsqueeze(v(t, x_t), dim=2)
        #kinetic_energy = torch.mean(torch.sum(torch.squeeze(v_at_x * (A_at_x @ v_at_x)), dim=1))
        kinetic_energy = torch.mean(torch.sum(torch.squeeze(v_at_x * torch.linalg.solve(A_at_x, v_at_x)), dim=1)) # v(t,x(t))^T @ A^-1(x(t)) @ v(t,x(t))
        #print(kinetic_energy)
        # Compute total loss and backprop
        loss = kinetic_energy + target_penalty
        losses.append(loss.detach().cpu())
        loss.backward()
        opt.step()
    
    # Report points x(t) for final geodesic
    return odeint(v, x_0, t), losses
//  [python]
def _filter_indices(
        self, positions: torch.Tensor, neighbor_idx: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """
        Routine for filtering out pair indices and offets due to the buffer region, which would otherwise slow down
        the calculators.

        Args:
            positions (torch.Tensor): Tensor of the Cartesian atom positions.
            neighbor_idx (dict(str, torch.Tensor)): Dictionary containing pair indices and offets

        Returns:
            dict(str, torch.Tensor): Dictionary containing updated pair indices and offets
        """
        offsets = neighbor_idx[properties.offsets]
        idx_i = neighbor_idx[properties.idx_i]
        idx_j = neighbor_idx[properties.idx_j]

        Rij = positions[idx_j] - positions[idx_i] + offsets
        d_ij = torch.norm(Rij, dim=1)
        d_ij_filter = d_ij <= self.cutoff

        neighbor_idx[properties.idx_i] = neighbor_idx[properties.idx_i][d_ij_filter]
        neighbor_idx[properties.idx_j] = neighbor_idx[properties.idx_j][d_ij_filter]
        neighbor_idx[properties.offsets] = neighbor_idx[properties.offsets][
            d_ij_filter, :
        ]

        return neighbor_idx
//  [python]
def _filter_indices(
        self, positions: torch.Tensor, neighbor_idx: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """
        Routine for filtering out pair indices and offets due to the buffer region, which would otherwise slow down
        the calculators.

        Args:
            positions (torch.Tensor): Tensor of the Cartesian atom positions.
            neighbor_idx (dict(str, torch.Tensor)): Dictionary containing pair indices and offets

        Returns:
            dict(str, torch.Tensor): Dictionary containing updated pair indices and offets
        """
        offsets = neighbor_idx[properties.offsets]
        idx_i = neighbor_idx[properties.idx_i]
        idx_j = neighbor_idx[properties.idx_j]

        Rij = positions[idx_j] - positions[idx_i] + offsets
        d_ij = torch.linalg.norm(Rij, dim=1)
        d_ij_filter = d_ij <= self.cutoff

        neighbor_idx[properties.idx_i] = neighbor_idx[properties.idx_i][d_ij_filter]
        neighbor_idx[properties.idx_j] = neighbor_idx[properties.idx_j][d_ij_filter]
        neighbor_idx[properties.offsets] = neighbor_idx[properties.offsets][
            d_ij_filter, :
        ]

        return neighbor_idx
//  [python]
def get_rays(self, H, W, K, R, T):
        K = K.float()
        R = R.float()
        T = T.float()[:, None]
        # pose = torch.cat([R, T], dim=1)
        pose = torch.eye(4, device=R.device).float()
        pose[:3, :4] = torch.cat([R, T], dim=1)
        # pose = torch.cat([pose, torch.tensor([[0, 0, 0, 1]]).to(pose.device)], dim=0)
        pose = torch.linalg.inv(pose)

        l = 1
        tx = torch.linspace(0, W - 1, W // l, device=pose.device)
        ty = torch.linspace(0, H - 1, H // l, device=pose.device)
        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(
            torch.linalg.inv(K)[None, None, :, :], p[:, :, :, None]
        ).squeeze()  # W, H, 3
        rays_v_norm = torch.norm(
            p, p=2, dim=-1, keepdim=True
        )  # TODO: used in sdfstudio
        rays_v = p / rays_v_norm  # W, H, 3
        rays_v = torch.matmul(
            pose[None, None, :3, :3], rays_v[:, :, :, None]
        ).squeeze()  # W, H, 3
        rays_o = pose[None, None, :3, 3].expand(rays_v.shape)  # W, H, 3

        # normalize ray directions
        rays_v = F.normalize(rays_v, dim=-1)
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def get_rays(self, H, W, K, R, T):
        K = K.float()
        R = R.float()
        T = T.float()[:, None]
        # pose = torch.cat([R, T], dim=1)
        pose = torch.eye(4, device=R.device).float()
        pose[:3, :4] = torch.cat([R, T], dim=1)
        # pose = torch.cat([pose, torch.tensor([[0, 0, 0, 1]]).to(pose.device)], dim=0)
        pose = torch.linalg.inv(pose)

        l = 1
        tx = torch.linspace(0, W - 1, W // l, device=pose.device)
        ty = torch.linspace(0, H - 1, H // l, device=pose.device)
        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(
            torch.linalg.inv(K)[None, None, :, :], p[:, :, :, None]
        ).squeeze()  # W, H, 3
        rays_v_norm = torch.linalg.norm(
            p, ord=2, dim=-1, keepdim=True
        )  # TODO: used in sdfstudio
        rays_v = p / rays_v_norm  # W, H, 3
        rays_v = torch.matmul(
            pose[None, None, :3, :3], rays_v[:, :, :, None]
        ).squeeze()  # W, H, 3
        rays_o = pose[None, None, :3, 3].expand(rays_v.shape)  # W, H, 3

        # normalize ray directions
        rays_v = F.normalize(rays_v, dim=-1)
        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.
        """
        level = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // level)
        ty = torch.linspace(0, self.H - 1, self.H // level)
        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.
        """
        level = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // level)
        ty = torch.linspace(0, self.H - 1, self.H // level)
        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 gen_rays_at_camera(self, pose, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        level = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // level)
        ty = torch.linspace(0, self.H - 1, self.H // level)
        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[0, 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(pose[:3, :3], rays_v[:, :, :,
                                                   None]).squeeze()  # W, H, 3
        rays_o = pose[: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_camera(self, pose, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        level = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // level)
        ty = torch.linspace(0, self.H - 1, self.H // level)
        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[0, 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(pose[:3, :3], rays_v[:, :, :,
                                                   None]).squeeze()  # W, H, 3
        rays_o = pose[:3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [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])  # bs
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])  # bs
        color = self.images[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]  # batch_size, 3

        depth = self.depths[img_idx][(pixels_y, pixels_x)]  # 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(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], depth[:, None]],
            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])  # bs
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])  # bs
        color = self.images[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]  # batch_size, 3

        depth = self.depths[img_idx][(pixels_y, pixels_x)]  # 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(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], depth[:, None]],
            dim=-1).cuda()
//  [python]
def forward(
        cls,
        embeds: Float[Tensor, "bs num_samples num_classes"],
        weights: Float[Tensor, "bs num_samples 1"],
    ) -> Float[Tensor, "bs num_classes"]:
        """Calculate semantics along the ray."""
        output = torch.sum(weights * embeds, dim=-2)
        output = output / torch.norm(output, dim=-1, keepdim=True)
        return output
//  [python]
def forward(
        cls,
        embeds: Float[Tensor, "bs num_samples num_classes"],
        weights: Float[Tensor, "bs num_samples 1"],
    ) -> Float[Tensor, "bs num_classes"]:
        """Calculate semantics along the ray."""
        output = torch.sum(weights * embeds, dim=-2)
        output = output / torch.linalg.norm(output, dim=-1, keepdim=True)
        return output
//  [python]
def after_batch(self, ligs_coords_pred, ligs_coords, batch_indices):
        cutoff = 5
        centroid_distances = []
        for lig_coords_pred, lig_coords in zip(ligs_coords_pred, ligs_coords):
            centroid_distances.append(torch.norm(lig_coords_pred.mean(dim=0) - lig_coords.mean(dim=0)))
        centroid_distances = torch.tensor(centroid_distances)
        above_cutoff = torch.tensor(batch_indices)[torch.where(centroid_distances > cutoff)[0]]
        if isinstance(self.sampler, HardSampler):
            self.sampler.add_hard_indices(above_cutoff.tolist())
//  [python]
def after_batch(self, ligs_coords_pred, ligs_coords, batch_indices):
        cutoff = 5
        centroid_distances = []
        for lig_coords_pred, lig_coords in zip(ligs_coords_pred, ligs_coords):
            centroid_distances.append(torch.linalg.norm(lig_coords_pred.mean(dim=0) - lig_coords.mean(dim=0)))
        centroid_distances = torch.tensor(centroid_distances)
        above_cutoff = torch.tensor(batch_indices)[torch.where(centroid_distances > cutoff)[0]]
        if isinstance(self.sampler, HardSampler):
            self.sampler.add_hard_indices(above_cutoff.tolist())
//  [python]
def gen_random_rays_at_mask(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])  # bs
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])  # bs
        color = self.images[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]  # batch_size, 3

        depth = self.depths[img_idx][(pixels_y, pixels_x)]  # 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(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], depth[:, None]],
            dim=-1).cuda()
//  [python]
def gen_random_rays_at_mask(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])  # bs
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])  # bs
        color = self.images[img_idx][(pixels_y, pixels_x)]  # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]  # batch_size, 3

        depth = self.depths[img_idx][(pixels_y, pixels_x)]  # 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(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], depth[:, None]],
            dim=-1).cuda()
//  [python]
def PointToLineSegmentDis(P,A,B):
    #点到线段的距离， 点的size为1，线段时多条
    AB = B - A
    AB_norm = AB / (torch.norm(AB,dim=-1, keepdim=True)+1e-8)

    PA = A - P[None,:]
    BP = P[None,:] - B
    s = torch.multiply(PA,AB_norm).sum(-1)
    t = torch.multiply(BP, AB_norm).sum(-1)
    z = torch.zeros(s.shape[0]).to(P.device).float()

    tmp = torch.vstack([s,t,z])
    h = torch.max(tmp,dim=0)[0]

    AP = P - A
    cross_result = torch.cross(AB_norm,AP,dim=-1)
    c = torch.norm(cross_result,dim=-1)

    return torch.sqrt(c**2+h**2)
//  [python]
def PointToLineSegmentDis(P,A,B):
    #点到线段的距离， 点的size为1，线段时多条
    AB = B - A
    AB_norm = AB / (torch.linalg.norm(AB,dim=-1, keepdim=True)+1e-8)

    PA = A - P[None,:]
    BP = P[None,:] - B
    s = torch.multiply(PA,AB_norm).sum(-1)
    t = torch.multiply(BP, AB_norm).sum(-1)
    z = torch.zeros(s.shape[0]).to(P.device).float()

    tmp = torch.vstack([s,t,z])
    h = torch.max(tmp,dim=0)[0]

    AP = P - A
    cross_result = torch.cross(AB_norm,AP,dim=-1)
    c = torch.linalg.norm(cross_result,dim=-1)

    return torch.sqrt(c**2+h**2)
//  [python]
def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
        """
        Interpolate pose between two cameras.
        """
        level = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // level)
        ty = torch.linspace(0, self.H - 1, self.H // level)
        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[0, 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
        trans = self.pose_all[idx_0, :3, 3] * (
            1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
        pose_0 = np.linalg.inv(pose_0)
        pose_1 = np.linalg.inv(pose_1)
        rot_0 = pose_0[:3, :3]
        rot_1 = pose_1[:3, :3]
        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
        key_times = [0, 1]
        slerp = Slerp(key_times, rots)
        rot = slerp(ratio)
        pose = np.diag([1.0, 1.0, 1.0, 1.0])
        pose = pose.astype(np.float32)
        pose[:3, :3] = rot.as_matrix()
        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
        pose = np.linalg.inv(pose)
        rot = torch.from_numpy(pose[:3, :3]).cuda()
        trans = torch.from_numpy(pose[:3, 3]).cuda()
        rays_v = torch.matmul(rot[None, None, :3, :3],
                              rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = trans[None, None, :3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1), pose
//  [python]
def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
        """
        Interpolate pose between two cameras.
        """
        level = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // level)
        ty = torch.linspace(0, self.H - 1, self.H // level)
        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[0, 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
        trans = self.pose_all[idx_0, :3, 3] * (
            1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
        pose_0 = np.linalg.inv(pose_0)
        pose_1 = np.linalg.inv(pose_1)
        rot_0 = pose_0[:3, :3]
        rot_1 = pose_1[:3, :3]
        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
        key_times = [0, 1]
        slerp = Slerp(key_times, rots)
        rot = slerp(ratio)
        pose = np.diag([1.0, 1.0, 1.0, 1.0])
        pose = pose.astype(np.float32)
        pose[:3, :3] = rot.as_matrix()
        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
        pose = np.linalg.inv(pose)
        rot = torch.from_numpy(pose[:3, :3]).cuda()
        trans = torch.from_numpy(pose[:3, 3]).cuda()
        rays_v = torch.matmul(rot[None, None, :3, :3],
                              rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = trans[None, None, :3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1), pose
//  [python]
def gen_rays_across(self, idx_0, idx_1, ratio, resolution_level=1):
        """
        Interpolate pose between two cameras.
        """
        level = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // level)
        ty = torch.linspace(0, self.H - 1, self.H // level)
        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[0, 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
        trans = self.pose_all[idx_0, :3, 3] * (
            1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
        pose_0 = np.linalg.inv(pose_0)
        pose_1 = np.linalg.inv(pose_1)
        rot_0 = pose_0[:3, :3]
        rot_1 = pose_1[:3, :3]
        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
        key_times = [0, 1]
        slerp = Slerp(key_times, rots)
        rot = slerp(ratio)
        pose = np.diag([1.0, 1.0, 1.0, 1.0])
        pose = pose.astype(np.float32)
        pose[:3, :3] = rot.as_matrix()
        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
        pose = np.linalg.inv(pose)
        rot = torch.from_numpy(pose[:3, :3]).cuda()
        trans = torch.from_numpy(pose[:3, 3]).cuda()
        rays_v = torch.matmul(rot[None, None, :3, :3],
                              rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = trans[None, None, :3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1), pose
//  [python]
def gen_rays_across(self, idx_0, idx_1, ratio, resolution_level=1):
        """
        Interpolate pose between two cameras.
        """
        level = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // level)
        ty = torch.linspace(0, self.H - 1, self.H // level)
        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[0, 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
        trans = self.pose_all[idx_0, :3, 3] * (
            1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
        pose_0 = np.linalg.inv(pose_0)
        pose_1 = np.linalg.inv(pose_1)
        rot_0 = pose_0[:3, :3]
        rot_1 = pose_1[:3, :3]
        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
        key_times = [0, 1]
        slerp = Slerp(key_times, rots)
        rot = slerp(ratio)
        pose = np.diag([1.0, 1.0, 1.0, 1.0])
        pose = pose.astype(np.float32)
        pose[:3, :3] = rot.as_matrix()
        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
        pose = np.linalg.inv(pose)
        rot = torch.from_numpy(pose[:3, :3]).cuda()
        trans = torch.from_numpy(pose[:3, 3]).cuda()
        rays_v = torch.matmul(rot[None, None, :3, :3],
                              rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = trans[None, None, :3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1), pose
//  [python]
def pointSegDis(xyz, xyzSeg):
    # 点到直线距离
    #xyzSeg : B,N,2,3
    # xyz: B,S,3
    _,S,_=xyz.shape
    _,N,_,_=xyzSeg.shape
    AB = xyzSeg[:,:,1,:] - xyzSeg[:,:,0,:] # B,N,3
    P = xyz.unsqueeze(2).repeat((1,1,N,1)) # B,S,1,3
    A = xyzSeg[:,:,0,:].unsqueeze(1).repeat((1,S,1,1)) # B,1,N,3
    AP = P-A
    AB1 = AB.unsqueeze(1).repeat((1,S,1,1)) # B,S,N,3
    AB_norm = torch.norm(AB,dim=-1) # B,N
    y = torch.cross(AB1,AP,dim=-1) # B,S,N,3
    y = torch.norm(y,dim=-1) # B,S,N
    AB_norm = AB_norm.unsqueeze(1).repeat((1,S,1))
    dis = y/(AB_norm+1e-6)

    D = torch.ones_like(dis).float()*100000
    dis = torch.where(AB_norm==0, D, dis)
    return dis
//  [python]
def pointSegDis(xyz, xyzSeg):
    # 点到直线距离
    #xyzSeg : B,N,2,3
    # xyz: B,S,3
    _,S,_=xyz.shape
    _,N,_,_=xyzSeg.shape
    AB = xyzSeg[:,:,1,:] - xyzSeg[:,:,0,:] # B,N,3
    P = xyz.unsqueeze(2).repeat((1,1,N,1)) # B,S,1,3
    A = xyzSeg[:,:,0,:].unsqueeze(1).repeat((1,S,1,1)) # B,1,N,3
    AP = P-A
    AB1 = AB.unsqueeze(1).repeat((1,S,1,1)) # B,S,N,3
    AB_norm = torch.linalg.norm(AB,dim=-1) # B,N
    y = torch.cross(AB1,AP,dim=-1) # B,S,N,3
    y = torch.linalg.norm(y,dim=-1) # B,S,N
    AB_norm = AB_norm.unsqueeze(1).repeat((1,S,1))
    dis = y/(AB_norm+1e-6)

    D = torch.ones_like(dis).float()*100000
    dis = torch.where(AB_norm==0, D, dis)
    return dis
//  [python]
def weighted_svd_decomposition(
    A: torch.Tensor,
    W: torch.Tensor,
    heuristic: Optional[str],
    randomized: bool,
    num_ranks: int,
    num_oversampling: int = 5,
    normalize: bool = False,
    reduce_before_sqrt: bool = True,  # seems to be better empirically
) -> Tuple[torch.Tensor, torch.Tensor]:
    if A.ndim != 2 or W.ndim != 2:
        raise ValueError(f"Expected 2D Matrix, but got {A.ndim} and {W.ndim}.")
    if A.shape != W.shape:
        raise ValueError(f"Expected A.shape == W.shape, but got {A.shape} and {W.shape}.")

    if heuristic is None:
        heuristic = "two-sided"
    if heuristic not in ["none", "row", "col", "two-sided"]:
        raise ValueError(f"Heuristic {heuristic} not supported.")

    if normalize is True:
        W = W / torch.norm(W)

    if reduce_before_sqrt is True:
        # (A.shape[0], 1)
        W1 = torch.sqrt(torch.mean(W, dim=1, keepdim=True))
        # (1, A.shape[1])
        W2 = torch.sqrt(torch.mean(W, dim=0, keepdim=True))
    else:
        # (A.shape[0], 1)
        W1 = torch.mean(torch.sqrt(W), dim=1, keepdim=True)
        # (1, A.shape[1])
        W2 = torch.mean(torch.sqrt(W), dim=0, keepdim=True)

    if heuristic == "none":
        A_tilde = A
    elif heuristic == "row":
        A_tilde = W1 * A
    elif heuristic == "col":
        A_tilde = A * W2
    elif heuristic == "two-sided":
        A_tilde = W1 * A * W2
    else:
        raise ValueError

    L1_tilde, L2_tilde = svd_decomposition(
        A_tilde,
        randomized=randomized,
        num_ranks=num_ranks,
        num_oversampling=num_oversampling)

    if heuristic == "none":
        L1 = L1_tilde
        L2 = L2_tilde
    elif heuristic == "row":
        L1 = L1_tilde / W1
        L2 = L2_tilde
    elif heuristic == "col":
        L1 = L1_tilde
        L2 = L2_tilde / W2
    elif heuristic == "two-sided":
        L1 = L1_tilde / W1
        L2 = L2_tilde / W2
    else:
        raise ValueError

    return L1, L2