jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def gen_rays_at(self, img_idx, pose = None, resolution_level=1):
        pose_cur = self.get_pose(img_idx, pose)

        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(pose_cur[None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = pose_cur[None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def normalize_scene(poses, n_images, bounds):
    # TODO: vdalidate the effectiveness of this function
    cam_pos = poses[:, :, 3].clone()
    center_ = cam_pos.mean(dim=0, keepdim=False)
    bias = cam_pos - center_.unsqueeze(0)
    radius_ = torch.norm(bias, p=2, dim=-1, keepdim=False).max().item()
    cam_pos = (cam_pos - center_.unsqueeze(0)) / radius_
    poses[:, :, 3] = cam_pos
    bounds = (bounds / radius_)
    return poses, bounds, center_, radius_
//  [python]
def normalize_scene(poses, n_images, bounds):
    # TODO: vdalidate the effectiveness of this function
    cam_pos = poses[:, :, 3].clone()
    center_ = cam_pos.mean(dim=0, keepdim=False)
    bias = cam_pos - center_.unsqueeze(0)
    radius_ = torch.linalg.norm(bias, ord=2, dim=-1, keepdim=False).max().item()
    cam_pos = (cam_pos - center_.unsqueeze(0)) / radius_
    poses[:, :, 3] = cam_pos
    bounds = (bounds / radius_)
    return poses, bounds, center_, radius_
//  [python]
def optimizer_based_smpl(self, init_kps3d,
                             kps3d_from_smpl) -> torch.Tensor:

        optim_kps3d = nn.Parameter(init_kps3d)
        optimizer = optim.Adam([optim_kps3d], lr=self.lr)
        last_loss = torch.tensor(0.)

        for step in range(self.steps):
            optimizer.zero_grad()
            loss = torch.norm(kps3d_from_smpl - optim_kps3d, p=2)

            if torch.abs(last_loss - loss) < 1e-3:
                break
            else:
                last_loss = loss
            loss.backward()
            optimizer.step()
        return optim_kps3d.detach()
//  [python]
def optimizer_based_smpl(self, init_kps3d,
                             kps3d_from_smpl) -> torch.Tensor:

        optim_kps3d = nn.Parameter(init_kps3d)
        optimizer = optim.Adam([optim_kps3d], lr=self.lr)
        last_loss = torch.tensor(0.)

        for step in range(self.steps):
            optimizer.zero_grad()
            loss = torch.linalg.norm(kps3d_from_smpl - optim_kps3d, ord=2)

            if torch.abs(last_loss - loss) < 1e-3:
                break
            else:
                last_loss = loss
            loss.backward()
            optimizer.step()
        return optim_kps3d.detach()
//  [python]
def extract_depth_maps(self, sdf_network, con_volume, intrinsics, c2ws, H, W, near, far):
        """
        extract depth maps from the density volume
        :param con_volume: [1, 1+C, dX, dY, dZ]  can by con_volume or sdf_volume
        :param c2ws: [B, 4, 4]
        :param H:
        :param W:
        :param near:
        :param far:
        :return:
        """
        device = con_volume.device
        batch_size = intrinsics.shape[0]

        with torch.no_grad():
            ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
                                    torch.linspace(0, W - 1, W))  # pytorch's meshgrid has indexing='ij'
            p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3

            intrinsics_inv = torch.inverse(intrinsics)

            p = p.view(-1, 3).float().to(device)  # N_rays, 3
            p = torch.matmul(intrinsics_inv[:, None, :3, :3], p[:, :, None]).squeeze()  # Batch, N_rays, 3
            rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # Batch, N_rays, 3
            rays_v = torch.matmul(c2ws[:, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # Batch, N_rays, 3
            rays_o = c2ws[:, None, :3, 3].expand(rays_v.shape)  # Batch, N_rays, 3
            rays_d = rays_v

        rays_o = rays_o.contiguous().view(-1, 3)
        rays_d = rays_d.contiguous().view(-1, 3)

        ################## - sphere tracer to extract depth maps               ######################
        depth_masks_sphere, depth_maps_sphere = self.ray_tracer.extract_depth_maps(
            rays_o, rays_d,
            near[None, :].repeat(rays_o.shape[0], 1),
            far[None, :].repeat(rays_o.shape[0], 1),
            sdf_network, con_volume
        )

        depth_maps = depth_maps_sphere.view(batch_size, 1, H, W)
        depth_masks = depth_masks_sphere.view(batch_size, 1, H, W)

        depth_maps = torch.where(depth_masks, depth_maps,
                                 torch.zeros_like(depth_masks.float()).to(device))  # fill invalid pixels by 0

        return depth_maps, depth_masks
//  [python]
def extract_depth_maps(self, sdf_network, con_volume, intrinsics, c2ws, H, W, near, far):
        """
        extract depth maps from the density volume
        :param con_volume: [1, 1+C, dX, dY, dZ]  can by con_volume or sdf_volume
        :param c2ws: [B, 4, 4]
        :param H:
        :param W:
        :param near:
        :param far:
        :return:
        """
        device = con_volume.device
        batch_size = intrinsics.shape[0]

        with torch.no_grad():
            ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
                                    torch.linspace(0, W - 1, W))  # pytorch's meshgrid has indexing='ij'
            p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3

            intrinsics_inv = torch.inverse(intrinsics)

            p = p.view(-1, 3).float().to(device)  # N_rays, 3
            p = torch.matmul(intrinsics_inv[:, None, :3, :3], p[:, :, None]).squeeze()  # Batch, N_rays, 3
            rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # Batch, N_rays, 3
            rays_v = torch.matmul(c2ws[:, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # Batch, N_rays, 3
            rays_o = c2ws[:, None, :3, 3].expand(rays_v.shape)  # Batch, N_rays, 3
            rays_d = rays_v

        rays_o = rays_o.contiguous().view(-1, 3)
        rays_d = rays_d.contiguous().view(-1, 3)

        ################## - sphere tracer to extract depth maps               ######################
        depth_masks_sphere, depth_maps_sphere = self.ray_tracer.extract_depth_maps(
            rays_o, rays_d,
            near[None, :].repeat(rays_o.shape[0], 1),
            far[None, :].repeat(rays_o.shape[0], 1),
            sdf_network, con_volume
        )

        depth_maps = depth_maps_sphere.view(batch_size, 1, H, W)
        depth_masks = depth_masks_sphere.view(batch_size, 1, H, W)

        depth_maps = torch.where(depth_masks, depth_maps,
                                 torch.zeros_like(depth_masks.float()).to(device))  # fill invalid pixels by 0

        return depth_maps, depth_masks
//  [python]
def norm(self, a, axis=None, keepdims=False):
        return torch.norm(a, dim=axis, keepdim=keepdims)
//  [python]
def norm(self, a, axis=None, keepdims=False):
        return torch.linalg.norm(a, dim=axis, keepdims=keepdims)
//  [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, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(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], dim=-1).cuda(), pixels_x, pixels_y, normal_sample, planes_sample, subplanes_sample
//  [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, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(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], dim=-1).cuda(), pixels_x, pixels_y, normal_sample, planes_sample, subplanes_sample
//  [python]
def generate_circle_by_vectors(t, C, r, n, u):
    n = n/n.norm()
    u = u/u.norm()
    P_circle = r*torch.cos(t)[:,None]*u + r*torch.sin(t)[:,None]*torch.cross(n,u) + C
    return P_circle
//  [python]
def generate_circle_by_vectors(t, C, r, n, u):
    n = n/torch.linalg.norm(n)
    u = u/torch.linalg.norm(u)
    P_circle = r*torch.cos(t)[:,None]*u + r*torch.sin(t)[:,None]*torch.cross(n,u) + C
    return P_circle
//  [python]
def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
        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[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]
        key_rots = [rot_0, rot_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)
//  [python]
def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
        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[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]
        key_rots = [rot_0, rot_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)
//  [python]
def poly_r_to_theta(poly: Tensor, image_points: Tensor,
                    principal_point: Tensor) -> Tensor:
    '''
    Convert polynom that projects radius to `z` to polynom
    that projects incident angle`theta` to radius
    `(x, y, z)` is the view vector with angle`theta` to `(0, 0, 1)`
    radius is the length of `(x, y)` vector
    '''
    r = torch.norm(image_points - principal_point, dim=-1)
    z = polyval(r, poly)
    theta = torch.arctan(r / z)
    from numpy.polynomial import Polynomial

    # leave out absolute term
    degree = torch.arange(len(poly) - 1) + 1
    poly_theta = Polynomial.fit(theta.cpu().flatten().numpy(), r.cpu(
    ).flatten().numpy(), deg=degree.tolist(), domain=(-1, 1)).coef
    return torch.from_numpy(poly_theta).to(poly)
//  [python]
def poly_r_to_theta(poly: Tensor, image_points: Tensor,
                    principal_point: Tensor) -> Tensor:
    '''
    Convert polynom that projects radius to `z` to polynom
    that projects incident angle`theta` to radius
    `(x, y, z)` is the view vector with angle`theta` to `(0, 0, 1)`
    radius is the length of `(x, y)` vector
    '''
    r = torch.linalg.norm(image_points - principal_point, dim=-1)
    z = polyval(r, poly)
    theta = torch.arctan(r / z)
    from numpy.polynomial import Polynomial

    # leave out absolute term
    degree = torch.arange(len(poly) - 1) + 1
    poly_theta = Polynomial.fit(theta.cpu().flatten().numpy(), r.cpu(
    ).flatten().numpy(), deg=degree.tolist(), domain=(-1, 1)).coef
    return torch.from_numpy(poly_theta).to(poly)
//  [python]
def angle_between(u, v, n=None):
    if n is None:
        return torch.arctan2(torch.norm(torch.cross(u,v)), torch.dot(u,v))
    else:
        return torch.arctan2(torch.dot(n,torch.cross(u,v)), torch.dot(u,v))
//  [python]
def angle_between(u, v, n=None):
    if n is None:
        return torch.arctan2(torch.linalg.norm(torch.cross(u,v)), torch.dot(u,v))
    else:
        return torch.arctan2(torch.dot(n,torch.cross(u,v)), torch.dot(u,v))
//  [python]
def render_core_outside(self,
                            rays_o,
                            rays_d,
                            z_vals,
                            sample_dist,
                            nerf,
                            background_rgb=None):
        batch_size, n_samples = z_vals.shape

        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat(
            [dists,
             torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
        mid_z_vals = z_vals + dists * 0.5

        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[
            ..., :, None]  # batch_size, n_samples, 3

        dis_to_center = torch.norm(
            pts, p=2, dim=-1, keepdim=True).clip(1.0, 1e10)
        pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center],
                        dim=-1)  # batch_size, n_samples, 4

        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)

        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
        dirs = dirs.reshape(-1, 3)

        density, sampled_color = nerf(pts, dirs)
        alpha = 1.0 - torch.exp(
            -F.softplus(density.reshape(batch_size, n_samples)) * dists)
        alpha = alpha.reshape(batch_size, n_samples)
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1),
            -1)[:, :-1]
        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
        color = (weights[:, :, None] * sampled_color).sum(dim=1)
        if background_rgb is not None:
            color = color + background_rgb * (
                1.0 - weights.sum(dim=-1, keepdim=True))

        return {
            'color': color,
            'sampled_color': sampled_color,
            'alpha': alpha,
            'weights': weights,
        }
//  [python]
def render_core_outside(self,
                            rays_o,
                            rays_d,
                            z_vals,
                            sample_dist,
                            nerf,
                            background_rgb=None):
        batch_size, n_samples = z_vals.shape

        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat(
            [dists,
             torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
        mid_z_vals = z_vals + dists * 0.5

        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[
            ..., :, None]  # batch_size, n_samples, 3

        dis_to_center = torch.linalg.norm(
            pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
        pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center],
                        dim=-1)  # batch_size, n_samples, 4

        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)

        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
        dirs = dirs.reshape(-1, 3)

        density, sampled_color = nerf(pts, dirs)
        alpha = 1.0 - torch.exp(
            -F.softplus(density.reshape(batch_size, n_samples)) * dists)
        alpha = alpha.reshape(batch_size, n_samples)
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1),
            -1)[:, :-1]
        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
        color = (weights[:, :, None] * sampled_color).sum(dim=1)
        if background_rgb is not None:
            color = color + background_rgb * (
                1.0 - weights.sum(dim=-1, keepdim=True))

        return {
            'color': color,
            'sampled_color': sampled_color,
            'alpha': alpha,
            'weights': weights,
        }
//  [python]
def rodrigues_rot2(P, n1, theta):
    '''
    Rotate points P wrt axis k by theta radians
    '''
    
    # If P is only 1d array (coords of single point), fix it to be matrix
    if P.ndim == 1:
        P = P[None,...]
    
    k = torch.cross(P, n1.unsqueeze(0))
    k = k/torch.norm(k)
    
    # Compute rotated points
    P_rot = torch.zeros((len(P),3))
    for i in range(len(P)):
        P_rot[i] = P[i]*torch.cos(theta[i]) + torch.cross(k[i],P[i])*torch.sin(theta[i]) + k[i]*torch.dot(k[i],P[i])*(1-torch.cos(theta[i]))

    return P_rot
//  [python]
def rodrigues_rot2(P, n1, theta):
    '''
    Rotate points P wrt axis k by theta radians
    '''
    
    # If P is only 1d array (coords of single point), fix it to be matrix
    if P.ndim == 1:
        P = P[None,...]
    
    k = torch.cross(P, n1.unsqueeze(0))
    k = k/torch.linalg.norm(k)
    
    # Compute rotated points
    P_rot = torch.zeros((len(P),3))
    for i in range(len(P)):
        P_rot[i] = P[i]*torch.cos(theta[i]) + torch.cross(k[i],P[i])*torch.sin(theta[i]) + k[i]*torch.dot(k[i],P[i])*(1-torch.cos(theta[i]))

    return P_rot
//  [python]
def rodrigues_rot(P, n0, n1):
    
    # If P is only 1d array (coords of single point), fix it to be matrix
    if P.ndim == 1:
        P = P[None,...]
    
    # Get vector of rotation k and angle theta
    n0 = n0/torch.norm(n0)
    n1 = n1/torch.norm(n1)
    k = torch.cross(n0,n1)
    k = k/torch.norm(k)
    theta = torch.arccos(torch.dot(n0,n1))
    
    # Compute rotated points
    P_rot = torch.zeros((len(P),3))
    for i in range(len(P)):
        P_rot[i] = P[i]*torch.cos(theta) + torch.cross(k,P[i])*torch.sin(theta) + k*torch.dot(k,P[i])*(1-torch.cos(theta))

    return P_rot
//  [python]
def rodrigues_rot(P, n0, n1):
    
    # If P is only 1d array (coords of single point), fix it to be matrix
    if P.ndim == 1:
        P = P[None,...]
    
    # Get vector of rotation k and angle theta
    n0 = n0/torch.linalg.norm(n0)
    n1 = n1/torch.linalg.norm(n1)
    k = torch.cross(n0,n1)
    k = k/torch.linalg.norm(k)
    theta = torch.arccos(torch.dot(n0,n1))
    
    # Compute rotated points
    P_rot = torch.zeros((len(P),3))
    for i in range(len(P)):
        P_rot[i] = P[i]*torch.cos(theta) + torch.cross(k,P[i])*torch.sin(theta) + k*torch.dot(k,P[i])*(1-torch.cos(theta))

    return P_rot
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
        radius = torch.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        prev_cos_val = torch.cat(
            [torch.zeros([batch_size, 1]).to(self.device), cos_val[:, :-1]],
            dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

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

        z_samples = sample_pdf(
            z_vals, weights, n_importance, det=True,
            device=self.device).detach()
        return z_samples
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        prev_cos_val = torch.cat(
            [torch.zeros([batch_size, 1]).to(self.device), cos_val[:, :-1]],
            dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

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

        z_samples = sample_pdf(
            z_vals, weights, n_importance, det=True,
            device=self.device).detach()
        return z_samples
//  [python]
def get_angles(target_camera: PerspectiveCameras, context_cameras: PerspectiveCameras, centroid=None):
    '''
    Get angles between cameras wrt a centroid

    Args:
        target_camera (Pytorch3D Camera): a camera object with a single camera
        context_cameras (PyTorch3D Camera): a camera object

    Returns:
        theta_deg (Tensor): a tensor containing angles in degrees
    '''
    a1 = target_camera.get_camera_center()
    b1 = context_cameras.get_camera_center()

    a = a1 - centroid.unsqueeze(0)
    a = a.expand(len(context_cameras), -1)
    b = b1 - centroid.unsqueeze(0)

    ab_dot = (a*b).sum(dim=-1)
    theta = torch.acos((ab_dot)/(torch.norm(a, dim=-1) * torch.norm(b, dim=-1)))
    theta_deg = theta * 180 / math.pi
    
    return theta_deg
//  [python]
def get_angles(target_camera: PerspectiveCameras, context_cameras: PerspectiveCameras, centroid=None):
    '''
    Get angles between cameras wrt a centroid

    Args:
        target_camera (Pytorch3D Camera): a camera object with a single camera
        context_cameras (PyTorch3D Camera): a camera object

    Returns:
        theta_deg (Tensor): a tensor containing angles in degrees
    '''
    a1 = target_camera.get_camera_center()
    b1 = context_cameras.get_camera_center()

    a = a1 - centroid.unsqueeze(0)
    a = a.expand(len(context_cameras), -1)
    b = b1 - centroid.unsqueeze(0)

    ab_dot = (a*b).sum(dim=-1)
    theta = torch.acos((ab_dot)/(torch.linalg.norm(a, dim=-1) * torch.linalg.norm(b, dim=-1)))
    theta_deg = theta * 180 / math.pi
    
    return theta_deg
//  [python]
def enable_row_pruning(self, ratio, method):
        # Here, we support two cases: L1 norm based pruning and topk based pruning
        self.row_pruning_ratio = ratio
        self.row_pruning_method = method

        if method == 'l1':
            # compute the l1 norm of each column
            weight_norm = torch.norm(self.weight.data, p=1, dim=1)
            mask = TopKBinarizer.apply(weight_norm, self.row_pruning_ratio, False)
            mask = mask.view(-1, 1)
            mask = mask.to(self.weight.device)
        elif method == 'topk':
            self.row_mask_scores = nn.Parameter(torch.Tensor(self.weight.size(0), 1))
            self.row_mask_scores.data = self.row_mask_scores.data.to(self.weight.device)
            init.kaiming_uniform_(self.row_mask_scores, a=math.sqrt(5))
            mask = None
        else:
            raise NotImplementedError

        self.register_buffer('row_pruning_mask', mask)
//  [python]
def enable_row_pruning(self, ratio, method):
        # Here, we support two cases: L1 norm based pruning and topk based pruning
        self.row_pruning_ratio = ratio
        self.row_pruning_method = method

        if method == 'l1':
            # compute the l1 norm of each column
            weight_norm = torch.linalg.norm(self.weight.data, ord=1, dim=1)
            mask = TopKBinarizer.apply(weight_norm, self.row_pruning_ratio, False)
            mask = mask.view(-1, 1)
            mask = mask.to(self.weight.device)
        elif method == 'topk':
            self.row_mask_scores = nn.Parameter(torch.Tensor(self.weight.size(0), 1))
            self.row_mask_scores.data = self.row_mask_scores.data.to(self.weight.device)
            init.kaiming_uniform_(self.row_mask_scores, a=math.sqrt(5))
            mask = None
        else:
            raise NotImplementedError

        self.register_buffer('row_pruning_mask', mask)
//  [python]
def _log_prob(loc, scale, value):
    # # if self._validate_args:
    # #     self._validate_sample(value)
    # # compute the variance
    # # var = (scale ** 2)
    # var = scale * scale
    # log_scale = math.log(scale) #if isinstance(scale, ) else scale.log()
    # return -((value - loc) ** 2) / (2 * var) - log_scale - math.log(math.sqrt(2 * math.pi))

    """
    value : B X d
    """
    loss = []
    for v in value:
        loss.append(const_k + 0.5 * (t.norm(v) ** 2))
    loss = torch.tensor(loss, requires_grad=True).reshape(value.shape[0], -1).to(device)
    return loss
//  [python]
def _log_prob(loc, scale, value):
    # # if self._validate_args:
    # #     self._validate_sample(value)
    # # compute the variance
    # # var = (scale ** 2)
    # var = scale * scale
    # log_scale = math.log(scale) #if isinstance(scale, ) else scale.log()
    # return -((value - loc) ** 2) / (2 * var) - log_scale - math.log(math.sqrt(2 * math.pi))

    """
    value : B X d
    """
    loss = []
    for v in value:
        loss.append(const_k + 0.5 * (torch.linalg.norm(v) ** 2))
    loss = torch.tensor(loss, requires_grad=True).reshape(value.shape[0], -1).to(device)
    return loss
//  [python]
def enable_channel_pruning(self, ratio, method):
        # Here, we support two cases: L1 norm based pruning and topk based pruning
        self.channel_pruning_ratio = ratio
        self.channel_pruning_method = method

        if method == 'l1':
            # compute the l1 norm of each conv2d kernel (the last three dimension)
            weight_norm = torch.norm(self.weight.data, p=1, dim=[1, 2, 3])
            mask = TopKBinarizer.apply(weight_norm, self.channel_pruning_ratio, False)
            mask = mask.view(-1, 1, 1, 1)
            mask = mask.to(self.weight.device)
        elif method == 'topk':
            self.channel_mask_scores = nn.Parameter(torch.Tensor(self.weight.size(0), 1, 1, 1))
            self.channel_mask_scores.data = self.channel_mask_scores.data.to(self.weight.device)
            init.kaiming_uniform_(self.channel_mask_scores, a=math.sqrt(5))
            mask = None
        else:
            raise NotImplementedError

        self.register_buffer('channel_pruning_mask', mask)
//  [python]
def enable_channel_pruning(self, ratio, method):
        # Here, we support two cases: L1 norm based pruning and topk based pruning
        self.channel_pruning_ratio = ratio
        self.channel_pruning_method = method

        if method == 'l1':
            # compute the l1 norm of each conv2d kernel (the last three dimension)
            weight_norm = torch.linalg.norm(self.weight.data, ord=1, dim=[1, 2, 3])
            mask = TopKBinarizer.apply(weight_norm, self.channel_pruning_ratio, False)
            mask = mask.view(-1, 1, 1, 1)
            mask = mask.to(self.weight.device)
        elif method == 'topk':
            self.channel_mask_scores = nn.Parameter(torch.Tensor(self.weight.size(0), 1, 1, 1))
            self.channel_mask_scores.data = self.channel_mask_scores.data.to(self.weight.device)
            init.kaiming_uniform_(self.channel_mask_scores, a=math.sqrt(5))
            mask = None
        else:
            raise NotImplementedError

        self.register_buffer('channel_pruning_mask', mask)
//  [python]
def perpendicular_component(
    a: Tensor | SameMergeSpace,
    b: Tensor | SameMergeSpace,
    **kwargs,
) -> Tensor | SameMergeSpace:
    norm_a = torch.norm(a)
    res = b - a * (a / norm_a * (b / norm_a)).sum()
    if res.isnan().any():
        return torch.zeros_like(a)
    return res
//  [python]
def perpendicular_component(
    a: Tensor | SameMergeSpace,
    b: Tensor | SameMergeSpace,
    **kwargs,
) -> Tensor | SameMergeSpace:
    norm_a = torch.linalg.norm(a)
    res = b - a * (a / norm_a * (b / norm_a)).sum()
    if res.isnan().any():
        return torch.zeros_like(a)
    return res
//  [python]
def test_sphere_random_sampling_with_n():
    C = Sphere(R3('x'), [0, 0, 0.0], 4)
    points = C.sample_random_uniform(n=40)
    assert points.as_tensor.shape == (40, 3)
    assert all(torch.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_sphere_random_sampling_with_n():
    C = Sphere(R3('x'), [0, 0, 0.0], 4)
    points = C.sample_random_uniform(n=40)
    assert points.as_tensor.shape == (40, 3)
    assert all(torch.linalg.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_sphere_random_sampling_with_d():
    C = Sphere(R3('x'), [0, 0, 0.0], 5)
    points = C.sample_random_uniform(d=13)
    assert all(t.norm(points.as_tensor, dim=1) <= 5.0)
//  [python]
def test_sphere_random_sampling_with_d():
    C = Sphere(R3('x'), [0, 0, 0.0], 5)
    points = C.sample_random_uniform(d=13)
    assert all(torch.linalg.norm(points.as_tensor, dim=1) <= 5.0)
//  [python]
def forward(self, q, k, v, need_weights=True):
        """Calculate attention.

        :param q: Query of shape (n_features,).
        :param k: Keys of shape (n_keys, n_features).
        :param v: Values of shape (n_keys, n_values_features).
        """
        eps = 1e-9
        if not self.classical:
            mk = (k - torch.mean(k, axis=1).view(-1, 1)) / (torch.std(k, axis=1).view(-1, 1) + eps)
            mq = (q - torch.mean(q)) / (torch.std(q) + eps)
        else:  # self.classical
            mk = k / torch.clamp(torch.norm(k, axis=0), min=eps)
            mq = q / torch.clamp(torch.norm(q, axis=0), min=eps)

        scores = mk @ mq
        weights = scores
        if self.use_softmax:
            weights = self.softmax(weights)
        weighted_values = weights @ v
        if need_weights:
            return weighted_values, weights

        return weighted_values
//  [python]
def forward(self, q, k, v, need_weights=True):
        """Calculate attention.

        :param q: Query of shape (n_features,).
        :param k: Keys of shape (n_keys, n_features).
        :param v: Values of shape (n_keys, n_values_features).
        """
        eps = 1e-9
        if not self.classical:
            mk = (k - torch.mean(k, axis=1).view(-1, 1)) / (torch.std(k, axis=1).view(-1, 1) + eps)
            mq = (q - torch.mean(q)) / (torch.std(q) + eps)
        else:  # self.classical
            mk = k / torch.clamp(torch.linalg.norm(k, axis=0), min=eps)
            mq = q / torch.clamp(torch.linalg.norm(q, axis=0), min=eps)

        scores = mk @ mq
        weights = scores
        if self.use_softmax:
            weights = self.softmax(weights)
        weighted_values = weights @ v
        if need_weights:
            return weighted_values, weights

        return weighted_values
//  [python]
def test_sphere_grid_sampling_with_n():
    C = Sphere(R3('x'), [0, 0, 0], 4)
    points = C.sample_grid(n=40)
    assert points.as_tensor.shape == (40, 3)
    assert all(torch.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_sphere_grid_sampling_with_n():
    C = Sphere(R3('x'), [0, 0, 0], 4)
    points = C.sample_grid(n=40)
    assert points.as_tensor.shape == (40, 3)
    assert all(torch.linalg.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_sphere_grid_sampling_with_to_small_n():
    C = Sphere(R3('x'), [0, 0, 0], 4)
    points = C.sample_grid(n=4)
    assert points.as_tensor.shape == (4, 3)
    assert all(torch.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_sphere_grid_sampling_with_to_small_n():
    C = Sphere(R3('x'), [0, 0, 0], 4)
    points = C.sample_grid(n=4)
    assert points.as_tensor.shape == (4, 3)
    assert all(torch.linalg.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_sphere_grid_sampling_with_d():
    C = Sphere(R3('x'), [0, 0, 0], 4)
    points = C.sample_grid(d=1)
    assert all(t.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_sphere_grid_sampling_with_d():
    C = Sphere(R3('x'), [0, 0, 0], 4)
    points = C.sample_grid(d=1)
    assert all(torch.linalg.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_sphere_boundary_random_sampling_with_d():
    C = Sphere(R3('x'), [0, 0, 0], 4).boundary
    points = C.sample_random_uniform(d=15)
    assert all(torch.isclose(points.as_tensor.norm(dim=1), torch.tensor(4.0)))
//  [python]
def test_sphere_boundary_random_sampling_with_d():
    C = Sphere(R3('x'), [0, 0, 0], 4).boundary
    points = C.sample_random_uniform(d=15)
    assert all(torch.isclose(torch.linalg.norm(points.as_tensor, dim=1), torch.tensor(4.0)))
//  [python]
def test_sphere_boundary_random_sampling_with_n():
    C = Sphere(R3('x'), [0, 0, 0], 2).boundary
    points = C.sample_random_uniform(n=10)
    assert points.as_tensor.shape == (10, 3)
    assert all(torch.isclose(torch.norm(points.as_tensor, dim=1), torch.tensor(2.0)))
//  [python]
def test_sphere_boundary_random_sampling_with_n():
    C = Sphere(R3('x'), [0, 0, 0], 2).boundary
    points = C.sample_random_uniform(n=10)
    assert points.as_tensor.shape == (10, 3)
    assert all(torch.isclose(torch.linalg.norm(points.as_tensor, dim=1), torch.tensor(2.0)))
//  [python]
def test_sphere_boundary_grid_sampling_with_n():
    C = Sphere(R3('x'), [0, 0, 0], 1).boundary
    points = C.sample_grid(n=30)
    assert points.as_tensor.shape == (30, 3)
    assert all(torch.isclose(torch.norm(points.as_tensor, dim=1), torch.tensor(1.0)))
//  [python]
def test_sphere_boundary_grid_sampling_with_n():
    C = Sphere(R3('x'), [0, 0, 0], 1).boundary
    points = C.sample_grid(n=30)
    assert points.as_tensor.shape == (30, 3)
    assert all(torch.isclose(torch.linalg.norm(points.as_tensor, dim=1), torch.tensor(1.0)))
//  [python]
def test_sphere_boundary_grid_sampling_with_d():
    C = Sphere(R3('x'), [0, 0, 0], 4).boundary
    points = C.sample_grid(d=13)
    assert all(torch.isclose(torch.norm(points.as_tensor, dim=1), torch.tensor(4.0)))
//  [python]
def test_sphere_boundary_grid_sampling_with_d():
    C = Sphere(R3('x'), [0, 0, 0], 4).boundary
    points = C.sample_grid(d=13)
    assert all(torch.isclose(torch.linalg.norm(points.as_tensor, dim=1), torch.tensor(4.0)))
//  [python]
def frobenius_norm(self, M):
        return torch.norm(M)
//  [python]
def frobenius_norm(self, M):
        return torch.linalg.norm(M, ord='fro')
//  [python]
def __init__(self, D, mu=None, lmbda=None, verbose=False):
        super(R_pca, self).__init__()
        self.D = D
        self.S = torch.zeros(self.D.shape)
        self.Y = torch.zeros(self.D.shape)

        if mu is not None: self.mu = mu
        else: self.mu = torch.prod(torch.tensor(self.D.shape)) / (4*torch.norm(self.D, ord=1))
            
        self.mu_inv = 1 / self.mu

        if lmbda is not None: self.lmbda = lmbda
        else: self.lmbda = 1 / torch.sqrt(torch.tensor(self.D.shape).max())
            
        self.verbose = verbose
//  [python]
def __init__(self, D, mu=None, lmbda=None, verbose=False):
        super(R_pca, self).__init__()
        self.D = D
        self.S = torch.zeros(self.D.shape)
        self.Y = torch.zeros(self.D.shape)

        if mu is not None: self.mu = mu
        else: self.mu = torch.prod(torch.tensor(self.D.shape)) / (4*torch.linalg.norm(self.D, ord=1))
            
        self.mu_inv = 1 / self.mu

        if lmbda is not None: self.lmbda = lmbda
        else: self.lmbda = 1 / torch.sqrt(torch.tensor(self.D.shape).max())
            
        self.verbose = verbose
//  [python]
def loss(self, M, L):
        return self.lr*(torch.norm(L, ord='nuc') + self.lmbda*Norm(self.S, ord=1) + self.inv_mu*Norm(M-L-self.S, ord='fro'))
//  [python]
def loss(self, M, L):
        return self.lr*(torch.linalg.matrix_norm(L, ord='nuc') + self.lmbda*Norm(self.S, ord=1) + self.inv_mu*Norm(M-L-self.S, ord='fro'))
//  [python]
def norm(x, ord=None, axis=None, keepdims=False):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = cast(x, dtype)
    return torch.norm(x, p=ord, dim=axis, keepdim=keepdims)
//  [python]
def norm(x, ord=None, axis=None, keepdims=False):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = cast(x, dtype)
    return torch.linalg.norm(x, ord=ord, dim=axis, keepdim=keepdims)
//  [python]
def my_step(self, source_batch, batch_idx, test=False):
        vertex_loss = torch.tensor(0.0, device=self.device)
        jacobian_loss = torch.tensor(0.0, device=self.device)

        # sanity checking the poisson solve, getting back GT vertices from GT jacobians. This is not used in this training.
        # GTT = batches.get_batch(0).poisson.jacobians_from_vertices(pred_V[0])
        # 		GT_V = batches.get_batch(0).poisson.solve_poisson(GTT)
        source = source_batch[0]
        target = source_batch[1]

        pred_V, pred_J, pred_J_restricted = self.predict_map(source, target)



        GT_V, GT_J, GT_J_restricted = self.get_gt_map(source, target)


        if UNIT_TEST_POISSON_SOLVE:
            success = self.check_map(source, target, GT_J, GT_V) < 0.0001
            print(self.check_map(source, target, GT_J, GT_V))
            assert(success), f"UNIT_TEST_POISSON_SOLVE FAILED!! {self.check_map(source, target, GT_J, GT_V)}"

        # compute losses

        if not self.args.no_vertex_loss:
            # predict jacobians close as possible to GT
            pred_V = pred_V - pred_V.mean(1, keepdim=True) + GT_V.mean(1, keepdim=True) 
            vertex_loss += self.mse(pred_V, GT_V)

        if not self.args.no_jacobian_loss:
            if self.args.normalize_jac_loss:
                norms = torch.norm(GT_J_restricted, dim=(2, 3))
                norms = 1 / torch.clamp(norms, 1e-2, 1e2).unsqueeze(-1).unsqueeze(-1)
                # comp = GT_J_restricted[i]/norms
                jacobian_loss += self.mse(GT_J_restricted / norms, pred_J_restricted / norms)
            else:
                jacobian_loss += self.mse(GT_J_restricted, pred_J_restricted)

        loss = (self.args.vertex_loss_weight * vertex_loss + jacobian_loss)
        loss = loss.type_as(GT_V)
        val_loss = loss
        if self.verbose:
            print(
                f"batch of {target.get_vertices().shape[0]:d} source <--> target pairs, each mesh {target.get_vertices().shape[1]:d} vertices, {source.get_source_triangles().shape[1]:d} faces")
        ret = {
            "target_V": target.get_vertices().detach(),
            "source_V": source.get_vertices().detach(),
            "loss": loss,
            "val_loss": val_loss,
            "vertex_loss": vertex_loss.detach(),
            "jacobian_loss": jacobian_loss.detach(),
            "pred_V": pred_V.detach(),
            "T": source.get_source_triangles(),
            'source_ind': source.source_ind,
            'target_inds': target.target_inds
            # "source_samples": batches.get_batch(0).get_loaded_data(True,"samples").detach()
        }

        if self.args.test:
            ret['pred_J_R'] = pred_J_restricted.detach()
            ret['target_J_R'] = GT_J_restricted.detach()
            ret['pred_time'] = pred_time
        return ret
//  [python]
def my_step(self, source_batch, batch_idx, test=False):
        vertex_loss = torch.tensor(0.0, device=self.device)
        jacobian_loss = torch.tensor(0.0, device=self.device)

        # sanity checking the poisson solve, getting back GT vertices from GT jacobians. This is not used in this training.
        # GTT = batches.get_batch(0).poisson.jacobians_from_vertices(pred_V[0])
        # 		GT_V = batches.get_batch(0).poisson.solve_poisson(GTT)
        source = source_batch[0]
        target = source_batch[1]

        pred_V, pred_J, pred_J_restricted = self.predict_map(source, target)



        GT_V, GT_J, GT_J_restricted = self.get_gt_map(source, target)


        if UNIT_TEST_POISSON_SOLVE:
            success = self.check_map(source, target, GT_J, GT_V) < 0.0001
            print(self.check_map(source, target, GT_J, GT_V))
            assert(success), f"UNIT_TEST_POISSON_SOLVE FAILED!! {self.check_map(source, target, GT_J, GT_V)}"

        # compute losses

        if not self.args.no_vertex_loss:
            # predict jacobians close as possible to GT
            pred_V = pred_V - pred_V.mean(1, keepdim=True) + GT_V.mean(1, keepdim=True) 
            vertex_loss += self.mse(pred_V, GT_V)

        if not self.args.no_jacobian_loss:
            if self.args.normalize_jac_loss:
                norms = torch.linalg.norm(GT_J_restricted, dim=(2, 3))
                norms = 1 / torch.clamp(norms, 1e-2, 1e2).unsqueeze(-1).unsqueeze(-1)
                # comp = GT_J_restricted[i]/norms
                jacobian_loss += self.mse(GT_J_restricted / norms, pred_J_restricted / norms)
            else:
                jacobian_loss += self.mse(GT_J_restricted, pred_J_restricted)

        loss = (self.args.vertex_loss_weight * vertex_loss + jacobian_loss)
        loss = loss.type_as(GT_V)
        val_loss = loss
        if self.verbose:
            print(
                f"batch of {target.get_vertices().shape[0]:d} source <--> target pairs, each mesh {target.get_vertices().shape[1]:d} vertices, {source.get_source_triangles().shape[1]:d} faces")
        ret = {
            "target_V": target.get_vertices().detach(),
            "source_V": source.get_vertices().detach(),
            "loss": loss,
            "val_loss": val_loss,
            "vertex_loss": vertex_loss.detach(),
            "jacobian_loss": jacobian_loss.detach(),
            "pred_V": pred_V.detach(),
            "T": source.get_source_triangles(),
            'source_ind': source.source_ind,
            'target_inds': target.target_inds
            # "source_samples": batches.get_batch(0).get_loaded_data(True,"samples").detach()
        }

        if self.args.test:
            ret['pred_J_R'] = pred_J_restricted.detach()
            ret['target_J_R'] = GT_J_restricted.detach()
            ret['pred_time'] = pred_time
        return ret
//  [python]
def loss(self, x_des: Pose, q: torch.Tensor, q_in: torch.Tensor):
        kin_state = self._robot_world.get_kinematics(q)
        distance = self._robot_world.pose_distance(x_des, kin_state.ee_pose)
        d_sdf = self._robot_world.collision_constraint(kin_state.link_spheres_tensor.unsqueeze(1))
        d_self = self._robot_world.self_collision_cost(kin_state.link_spheres_tensor.unsqueeze(1))
        loss = 0.1 * torch.norm(q_in - q, dim=-1) + distance + 100.0 * (d_self + d_sdf)
        return loss
//  [python]
def loss(self, x_des: Pose, q: torch.Tensor, q_in: torch.Tensor):
        kin_state = self._robot_world.get_kinematics(q)
        distance = self._robot_world.pose_distance(x_des, kin_state.ee_pose)
        d_sdf = self._robot_world.collision_constraint(kin_state.link_spheres_tensor.unsqueeze(1))
        d_self = self._robot_world.self_collision_cost(kin_state.link_spheres_tensor.unsqueeze(1))
        loss = 0.1 * torch.linalg.norm(q_in - q, dim=-1) + distance + 100.0 * (d_self + d_sdf)
        return loss