jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def mean_vertex_error(vert, target, translate_align=False):
    if vert.shape[1] != target.shape[1]:
        min_ = min(vert.shape[1], target.shape[1])
        vert = vert[:, :min_, :]
        target = target[:, :min_, :]
    assert (len(vert.shape) == 3 or len(vert.shape) == 2) and vert.shape == target.shape
    if translate_align:
        vert = translation_alignment(vert, target)
    return torch.mean(torch.norm(vert - target, dim=-1)) * 100
//  [python]
def mean_vertex_error(vert, target, translate_align=False):
    if vert.shape[1] != target.shape[1]:
        min_ = min(vert.shape[1], target.shape[1])
        vert = vert[:, :min_, :]
        target = target[:, :min_, :]
    assert (len(vert.shape) == 3 or len(vert.shape) == 2) and vert.shape == target.shape
    if translate_align:
        vert = translation_alignment(vert, target)
    return torch.mean(torch.linalg.norm(vert - target, dim=-1)) * 100
//  [python]
def render_line_segment(a, b, size, distance='gauss', sigma=0.2,
                        normalize=False, widths=None):
    """
    a, b points defining the line segment
    widths: B x N
    outputs B x N x H x W
    """
    def sumprod(x, y, keepdim=True):
        return torch.sum(x * y, dim=-1, keepdim=keepdim)

    grid = torch.linspace(-1.0, 1.0, size, dtype=torch.float32, device=a.device)

    # FIXME: api different from numpy
    yv, xv = torch.meshgrid([grid, grid])
    # 1 x H x W x 2
    m = torch.cat([xv[..., None], yv[..., None]], dim=-1)[None, None]

    # B x N x 1 x 1 x 2
    a, b = a[:, :, None, None, :], b[:, :, None, None, :]
    t_min = sumprod(m - a, b - a) / \
        torch.max(sumprod(b - a, b - a), torch.tensor(1e-6, device=a.device))
    t_line = torch.clamp(t_min, 0.0, 1.0)

    # closest points on the line to every image pixel
    s = a + t_line * (b - a)

    # for rectangle
    if widths is not None:
        # get perpendicular unit vector
        u = b - a
        u = u[..., [1, 0]] * torch.Tensor([[[[[1, -1]]]]]).to(u.device)
        u = u / torch.norm(u, dim=-1, keepdim=True)

        t_min = sumprod(m - s, u) / \
            torch.max(sumprod(u, u), torch.tensor(1e-6, device=a.device))

        w = widths[..., None, None, None] / 2.0
        t_line = clamp(t_min, -w, w)

        s = s + t_line * u

    d = sumprod(s - m, s - m, keepdim=False)

    # normalize distancin
    if distance == 'gauss':
        d = torch.sqrt(d + 1e-6)
        d_norm = torch.exp(-d / (sigma ** 2))
    elif distance == 'norm':
        d = torch.sqrt(d + 1e-6)
        d_max = torch.sqrt(8)
        d_norm = (d_max - d) / d_max
    else:
        raise ValueError()

    if normalize:
        d_norm = d_norm / torch.sum(d_norm, (2, 3), keepdim=True)

    return d_norm
//  [python]
def render_line_segment(a, b, size, distance='gauss', sigma=0.2,
                        normalize=False, widths=None):
    """
    a, b points defining the line segment
    widths: B x N
    outputs B x N x H x W
    """
    def sumprod(x, y, keepdim=True):
        return torch.sum(x * y, dim=-1, keepdim=keepdim)

    grid = torch.linspace(-1.0, 1.0, size, dtype=torch.float32, device=a.device)

    # FIXME: api different from numpy
    yv, xv = torch.meshgrid([grid, grid])
    # 1 x H x W x 2
    m = torch.cat([xv[..., None], yv[..., None]], dim=-1)[None, None]

    # B x N x 1 x 1 x 2
    a, b = a[:, :, None, None, :], b[:, :, None, None, :]
    t_min = sumprod(m - a, b - a) / \
        torch.max(sumprod(b - a, b - a), torch.tensor(1e-6, device=a.device))
    t_line = torch.clamp(t_min, 0.0, 1.0)

    # closest points on the line to every image pixel
    s = a + t_line * (b - a)

    # for rectangle
    if widths is not None:
        # get perpendicular unit vector
        u = b - a
        u = u[..., [1, 0]] * torch.Tensor([[[[[1, -1]]]]]).to(u.device)
        u = u / torch.linalg.norm(u, dim=-1, keepdim=True)

        t_min = sumprod(m - s, u) / \
            torch.max(sumprod(u, u), torch.tensor(1e-6, device=a.device))

        w = widths[..., None, None, None] / 2.0
        t_line = torch.clamp(t_min, -w, w)

        s = s + t_line * u

    d = sumprod(s - m, s - m, keepdim=False)

    # normalize distancin
    if distance == 'gauss':
        d = torch.sqrt(d + 1e-6)
        d_norm = torch.exp(-d / (sigma ** 2))
    elif distance == 'norm':
        d = torch.sqrt(d + 1e-6)
        d_max = torch.sqrt(8)
        d_norm = (d_max - d) / d_max
    else:
        raise ValueError()

    if normalize:
        d_norm = d_norm / torch.sum(d_norm, (2, 3), keepdim=True)

    return d_norm
//  [python]
def calc_tri_areas(vertices, faces):
    v1 = vertices[faces[:, 0], :]
    v2 = vertices[faces[:, 1], :]
    v3 = vertices[faces[:, 2], :]
    
    v1 = v1 - v3
    v2 = v2 - v3
    return torch.norm(torch.cross(v1, v2, dim=1), dim=1) * .5
//  [python]
def calc_tri_areas(vertices, faces):
    v1 = vertices[faces[:, 0], :]
    v2 = vertices[faces[:, 1], :]
    v3 = vertices[faces[:, 2], :]
    
    v1 = v1 - v3
    v2 = v2 - v3
    return torch.linalg.norm(torch.cross(v1, v2, dim=1), dim=1) * .5
//  [python]
def compute_density(self, points):
        points_norm = torch.norm(points, dim=-1, keepdim=True)
        points_norm = torch.clamp(points_norm, min=1e-3)
        sigma = self.outer_nerf.density(torch.cat([points / points_norm, 1.0 / points_norm], -1))[..., 0]
        return sigma
//  [python]
def compute_density(self, points):
        points_norm = torch.linalg.norm(points, dim=-1, keepdim=True)
        points_norm = torch.clamp(points_norm, min=1e-3)
        sigma = self.outer_nerf.density(torch.cat([points / points_norm, 1.0 / points_norm], -1))[..., 0]
        return sigma
//  [python]
def mean(self, points_in_manifold, **kwargs):
        """
        Computes the mean of the points points_in_manifold on the manifold according to the iterative likelihood maximization.
        """
        weights = kwargs.get('weights', None)
        init_mu = kwargs.get('init_mu', None)
        eps = kwargs.get('eps', 1e-5)
        max_iter = kwargs.get('max_iter', 50)
        return_projections = kwargs.get('return_projections', False)
        if isinstance(points_in_manifold, np.ndarray):
            points_in_manifold = torch.from_numpy(points_in_manifold).to(**self.tensor_args)
        if weights is None:
            weights = torch.ones((points_in_manifold.shape[0], 1), **self.tensor_args)  # batch size
        if init_mu is None:
            mu = points_in_manifold[0]
        else:
            mu = init_mu
        eps_of_mu = None
        k = 0
        for k in range(max_iter):
            mu_last = mu
            eps_of_mu = self.log_map(points_in_manifold, base=mu)
            delta = (eps_of_mu * weights).mean(0)
            mu = self.exp_map(delta, base=mu)
            if torch.norm(mu - mu_last) < eps:
                break
        if k == max_iter - 1:
            Manifold.logger.warning(f'Maximum of {max_iter} iterations reached!')
        if return_projections:
            return mu.squeeze(), eps_of_mu.squeeze()
        else:
            return mu.squeeze()
//  [python]
def mean(self, points_in_manifold, **kwargs):
        """
        Computes the mean of the points points_in_manifold on the manifold according to the iterative likelihood maximization.
        """
        weights = kwargs.get('weights', None)
        init_mu = kwargs.get('init_mu', None)
        eps = kwargs.get('eps', 1e-5)
        max_iter = kwargs.get('max_iter', 50)
        return_projections = kwargs.get('return_projections', False)
        if isinstance(points_in_manifold, np.ndarray):
            points_in_manifold = torch.from_numpy(points_in_manifold).to(**self.tensor_args)
        if weights is None:
            weights = torch.ones((points_in_manifold.shape[0], 1), **self.tensor_args)  # batch size
        if init_mu is None:
            mu = points_in_manifold[0]
        else:
            mu = init_mu
        eps_of_mu = None
        k = 0
        for k in range(max_iter):
            mu_last = mu
            eps_of_mu = self.log_map(points_in_manifold, base=mu)
            delta = (eps_of_mu * weights).mean(0)
            mu = self.exp_map(delta, base=mu)
            if torch.linalg.norm(mu - mu_last) < eps:
                break
        if k == max_iter - 1:
            Manifold.logger.warning(f'Maximum of {max_iter} iterations reached!')
        if return_projections:
            return mu.squeeze(), eps_of_mu.squeeze()
        else:
            return mu.squeeze()
//  [python]
def midpoint_triangulate(self, match, K_inv, P1, P2):
        # match: [b, 4, num] P1: [b, 3, 4]
        # Match is in the image coordinates. P1, P2 is camera parameters. [B, 3, 4] match: [B, M, 4]
        b, n = match.shape[0], match.shape[2]
        RT1 = K_inv.bmm(P1) # [b, 3, 4]
        RT2 = K_inv.bmm(P2)
        ones = torch.ones([b,1,n]).to(match.get_device())
        pts1 = torch.cat([match[:,:2,:], ones], 1)
        pts2 = torch.cat([match[:,2:,:], ones], 1)
        
        ray1_dir = (RT1[:,:,:3].transpose(1,2)).bmm(K_inv).bmm(pts1)# [b,3,n]
        ray1_dir = ray1_dir / (torch.norm(ray1_dir, dim=1, keepdim=True, p=2) + 1e-12)
        ray1_origin = (-1) * RT1[:,:,:3].transpose(1,2).bmm(RT1[:,:,3].unsqueeze(-1)) # [b, 3, 1]
        ray2_dir = (RT2[:,:,:3].transpose(1,2)).bmm(K_inv).bmm(pts2) # [b,3,n]
        ray2_dir = ray2_dir / (torch.norm(ray2_dir, dim=1, keepdim=True, p=2) + 1e-12)
        ray2_origin = (-1) * RT2[:,:,:3].transpose(1,2).bmm(RT2[:,:,3].unsqueeze(-1)) # [b, 3, 1]
    
        dir_cross = torch.cross(ray1_dir, ray2_dir, dim=1) # [b,3,n]
        denom = 1.0 / (torch.sum(dir_cross * dir_cross, dim=1, keepdim=True)+1e-12) # [b,1,n]
        origin_vec = (ray2_origin - ray1_origin).repeat(1,1,n) # [b,3,n]
        a1 = origin_vec.cross(ray2_dir, dim=1) # [b,3,n]
        a1 = torch.sum(a1 * dir_cross, dim=1, keepdim=True) * denom # [b,1,n]
        a2 = origin_vec.cross(ray1_dir, dim=1) # [b,3,n]
        a2 = torch.sum(a2 * dir_cross, dim=1, keepdim=True) * denom # [b,1,n]
        p1 = ray1_origin + a1 * ray1_dir
        p2 = ray2_origin + a2 * ray2_dir
        point = (p1 + p2) / 2.0 # [b,3,n]
        # Convert to homo coord to get consistent with other functions.
        point_homo = torch.cat([point, ones], dim=1).transpose(1,2) # [b,n,4]
        return point_homo
//  [python]
def midpoint_triangulate(self, match, K_inv, P1, P2):
        # match: [b, 4, num] P1: [b, 3, 4]
        # Match is in the image coordinates. P1, P2 is camera parameters. [B, 3, 4] match: [B, M, 4]
        b, n = match.shape[0], match.shape[2]
        RT1 = K_inv.bmm(P1) # [b, 3, 4]
        RT2 = K_inv.bmm(P2)
        ones = torch.ones([b,1,n]).to(match.get_device())
        pts1 = torch.cat([match[:,:2,:], ones], 1)
        pts2 = torch.cat([match[:,2:,:], ones], 1)
        
        ray1_dir = (RT1[:,:,:3].transpose(1,2)).bmm(K_inv).bmm(pts1)# [b,3,n]
        ray1_dir = ray1_dir / (torch.linalg.norm(ray1_dir, dim=1, keepdim=True, ord=2) + 1e-12)
        ray1_origin = (-1) * RT1[:,:,:3].transpose(1,2).bmm(RT1[:,:,3].unsqueeze(-1)) # [b, 3, 1]
        ray2_dir = (RT2[:,:,:3].transpose(1,2)).bmm(K_inv).bmm(pts2) # [b,3,n]
        ray2_dir = ray2_dir / (torch.linalg.norm(ray2_dir, dim=1, keepdim=True, ord=2) + 1e-12)
        ray2_origin = (-1) * RT2[:,:,:3].transpose(1,2).bmm(RT2[:,:,3].unsqueeze(-1)) # [b, 3, 1]
    
        dir_cross = torch.cross(ray1_dir, ray2_dir, dim=1) # [b,3,n]
        denom = 1.0 / (torch.sum(dir_cross * dir_cross, dim=1, keepdim=True)+1e-12) # [b,1,n]
        origin_vec = (ray2_origin - ray1_origin).repeat(1,1,n) # [b,3,n]
        a1 = origin_vec.cross(ray2_dir, dim=1) # [b,3,n]
        a1 = torch.sum(a1 * dir_cross, dim=1, keepdim=True) * denom # [b,1,n]
        a2 = origin_vec.cross(ray1_dir, dim=1) # [b,3,n]
        a2 = torch.sum(a2 * dir_cross, dim=1, keepdim=True) * denom # [b,1,n]
        p1 = ray1_origin + a1 * ray1_dir
        p2 = ray2_origin + a2 * ray2_dir
        point = (p1 + p2) / 2.0 # [b,3,n]
        # Convert to homo coord to get consistent with other functions.
        point_homo = torch.cat([point, ones], dim=1).transpose(1,2) # [b,n,4]
        return point_homo
//  [python]
def compute_density_alpha(self, points, dists, dirs, nerf):
        norm = torch.norm(points, dim=-1, keepdim=True)
        points = torch.cat([points / norm, 1.0 / norm], -1)
        density, color = nerf(points, dirs)  # [...,1] [...,3]
        alpha = self.density_activation(density[..., 0], dists)
        color = linear_to_srgb(torch.exp(torch.clamp(color, max=5.0)))
        return alpha, color
//  [python]
def compute_density_alpha(self, points, dists, dirs, nerf):
        norm = torch.linalg.norm(points, dim=-1, keepdim=True)
        points = torch.cat([points / norm, 1.0 / norm], -1)
        density, color = nerf(points, dirs)  # [...,1] [...,3]
        alpha = self.density_activation(density[..., 0], dists)
        color = linear_to_srgb(torch.exp(torch.clamp(color, max=5.0)))
        return alpha, color
//  [python]
def ray_angle_filter(self, match, P1, P2, return_angle=False):
        # match: [b, 4, n] P: [B, 3, 4]
        b, n = match.shape[0], match.shape[2]
        K = P1[:,:,:3] # P1 with identity rotation and zero translation
        K_inv = torch.inverse(K)
        RT1 = K_inv.bmm(P1) # [b, 3, 4]
        RT2 = K_inv.bmm(P2)
        ones = torch.ones([b,1,n]).to(match.get_device())
        pts1 = torch.cat([match[:,:2,:], ones], 1)
        pts2 = torch.cat([match[:,2:,:], ones], 1)
        
        ray1_dir = (RT1[:,:,:3].transpose(1,2)).bmm(K_inv).bmm(pts1)# [b,3,n]
        ray1_dir = ray1_dir / (torch.norm(ray1_dir, dim=1, keepdim=True, p=2) + 1e-12)
        ray1_origin = (-1) * RT1[:,:,:3].transpose(1,2).bmm(RT1[:,:,3].unsqueeze(-1)) # [b, 3, 1]
        ray2_dir = (RT2[:,:,:3].transpose(1,2)).bmm(K_inv).bmm(pts2) # [b,3,n]
        ray2_dir = ray2_dir / (torch.norm(ray2_dir, dim=1, keepdim=True, p=2) + 1e-12)
        ray2_origin = (-1) * RT2[:,:,:3].transpose(1,2).bmm(RT2[:,:,3].unsqueeze(-1)) # [b, 3, 1]

        # We compute the angle betwwen vertical line from ray1 origin to ray2 and ray1.
        p1p2 = (ray1_origin - ray2_origin).repeat(1,1,n)
        verline = ray2_origin.repeat(1,1,n) + torch.sum(p1p2 * ray2_dir, dim=1, keepdim=True) * ray2_dir - ray1_origin.repeat(1,1,n) # [b,3,n]
        cosvalue = torch.sum(ray1_dir * verline, dim=1, keepdim=True)  / \
            ((torch.norm(ray1_dir, dim=1, keepdim=True, p=2) + 1e-12) * (torch.norm(verline, dim=1, keepdim=True, p=2) + 1e-12))# [b,1,n]

        mask = (cosvalue > 0.001).float() # we drop out angles less than 1' [b,1,n]
        flag = 0
        num = torch.min(torch.sum(mask, -1)).int()
        if num.cpu().detach().numpy() == 0:
            flag = 1
            filt_match = match[:,:,:100]
            if return_angle:
                return filt_match, flag, torch.zeros_like(mask).to(filt_match.get_device())
            else:
                return filt_match, flag
        nonzero_idx = []
        for i in range(b):
            idx = torch.nonzero(mask[i,0,:])[:num] # [num,1]
            nonzero_idx.append(idx)
        nonzero_idx = torch.stack(nonzero_idx, 0) # [b,num,1]
        filt_match = torch.gather(match.transpose(1,2), index=nonzero_idx.repeat(1,1,4), dim=1).transpose(1,2) # [b,4,num]
        if return_angle:
            return filt_match, flag, mask
        else:
            return filt_match, flag
//  [python]
def ray_angle_filter(self, match, P1, P2, return_angle=False):
        # match: [b, 4, n] P: [B, 3, 4]
        b, n = match.shape[0], match.shape[2]
        K = P1[:,:,:3] # P1 with identity rotation and zero translation
        K_inv = torch.inverse(K)
        RT1 = K_inv.bmm(P1) # [b, 3, 4]
        RT2 = K_inv.bmm(P2)
        ones = torch.ones([b,1,n]).to(match.get_device())
        pts1 = torch.cat([match[:,:2,:], ones], 1)
        pts2 = torch.cat([match[:,2:,:], ones], 1)
        
        ray1_dir = (RT1[:,:,:3].transpose(1,2)).bmm(K_inv).bmm(pts1)# [b,3,n]
        ray1_dir = ray1_dir / (torch.linalg.norm(ray1_dir, dim=1, keepdim=True) + 1e-12)
        ray1_origin = (-1) * RT1[:,:,:3].transpose(1,2).bmm(RT1[:,:,3].unsqueeze(-1)) # [b, 3, 1]
        ray2_dir = (RT2[:,:,:3].transpose(1,2)).bmm(K_inv).bmm(pts2) # [b,3,n]
        ray2_dir = ray2_dir / (torch.linalg.norm(ray2_dir, dim=1, keepdim=True) + 1e-12)
        ray2_origin = (-1) * RT2[:,:,:3].transpose(1,2).bmm(RT2[:,:,3].unsqueeze(-1)) # [b, 3, 1]

        # We compute the angle betwwen vertical line from ray1 origin to ray2 and ray1.
        p1p2 = (ray1_origin - ray2_origin).repeat(1,1,n)
        verline = ray2_origin.repeat(1,1,n) + torch.sum(p1p2 * ray2_dir, dim=1, keepdim=True) * ray2_dir - ray1_origin.repeat(1,1,n) # [b,3,n]
        cosvalue = torch.sum(ray1_dir * verline, dim=1, keepdim=True)  / \
            ((torch.linalg.norm(ray1_dir, dim=1, keepdim=True) + 1e-12) * (torch.linalg.norm(verline, dim=1, keepdim=True) + 1e-12))# [b,1,n]

        mask = (cosvalue > 0.001).float() # we drop out angles less than 1' [b,1,n]
        flag = 0
        num = torch.min(torch.sum(mask, -1)).int()
        if num.cpu().detach().numpy() == 0:
            flag = 1
            filt_match = match[:,:,:100]
            if return_angle:
                return filt_match, flag, torch.zeros_like(mask).to(filt_match.get_device())
            else:
                return filt_match, flag
        nonzero_idx = []
        for i in range(b):
            idx = torch.nonzero(mask[i,0,:])[:num] # [num,1]
            nonzero_idx.append(idx)
        nonzero_idx = torch.stack(nonzero_idx, 0) # [b,num,1]
        filt_match = torch.gather(match.transpose(1,2), index=nonzero_idx.repeat(1,1,4), dim=1).transpose(1,2) # [b,4,num]
        if return_angle:
            return filt_match, flag, mask
        else:
            return filt_match, flag
//  [python]
def axis_angle_add(theta, roll_axis, alpha):
    """Composition of two axis-angle rotations (PyTorch version)
    Args:
        theta: N x 3
        roll_axis: N x 3
        alph: N x 1
    Returns:
        equivalent axis-angle of the composition
    """
    alpha = alpha / 2.

    l2norm = torch.norm(theta + 1e-8, p=2, dim=1)
    angle = torch.unsqueeze(l2norm, -1)

    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    b_cos = torch.cos(angle).cpu()
    b_sin = torch.sin(angle).cpu()

    a_cos = torch.cos(alpha)
    a_sin = torch.sin(alpha)

    dot_mm = torch.sum(normalized * roll_axis, dim=1, keepdim=True)
    cross_mm = torch.zeros_like(normalized)
    cross_mm[:, 0] = roll_axis[:, 1] * normalized[:, 2] - roll_axis[:, 2] * normalized[:, 1]
    cross_mm[:, 1] = roll_axis[:, 2] * normalized[:, 0] - roll_axis[:, 0] * normalized[:, 2]
    cross_mm[:, 2] = roll_axis[:, 0] * normalized[:, 1] - roll_axis[:, 1] * normalized[:, 0]

    c_cos = a_cos * b_cos - a_sin * b_sin * dot_mm
    c_sin_n = a_sin * b_cos * roll_axis + a_cos * b_sin * normalized + a_sin * b_sin * cross_mm

    c_angle = 2 * torch.acos(c_cos)
    c_sin = torch.sin(c_angle * 0.5)
    c_n = (c_angle / c_sin) * c_sin_n

    return c_n
//  [python]
def axis_angle_add(theta, roll_axis, alpha):
    """Composition of two axis-angle rotations (PyTorch version)
    Args:
        theta: N x 3
        roll_axis: N x 3
        alph: N x 1
    Returns:
        equivalent axis-angle of the composition
    """
    alpha = alpha / 2.

    l2norm = torch.linalg.norm(theta + 1e-8, ord=2, dim=1)
    angle = torch.unsqueeze(l2norm, -1)

    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    b_cos = torch.cos(angle).cpu()
    b_sin = torch.sin(angle).cpu()

    a_cos = torch.cos(alpha)
    a_sin = torch.sin(alpha)

    dot_mm = torch.sum(normalized * roll_axis, dim=1, keepdim=True)
    cross_mm = torch.zeros_like(normalized)
    cross_mm[:, 0] = roll_axis[:, 1] * normalized[:, 2] - roll_axis[:, 2] * normalized[:, 1]
    cross_mm[:, 1] = roll_axis[:, 2] * normalized[:, 0] - roll_axis[:, 0] * normalized[:, 2]
    cross_mm[:, 2] = roll_axis[:, 0] * normalized[:, 1] - roll_axis[:, 1] * normalized[:, 0]

    c_cos = a_cos * b_cos - a_sin * b_sin * dot_mm
    c_sin_n = a_sin * b_cos * roll_axis + a_cos * b_sin * normalized + a_sin * b_sin * cross_mm

    c_angle = 2 * torch.acos(c_cos)
    c_sin = torch.sin(c_angle * 0.5)
    c_n = (c_angle / c_sin) * c_sin_n

    return c_n
//  [python]
def _warn_ray_tracing(self, centers):
        centers = centers.reshape([-1,3])
        distance = torch.norm(centers,dim=-1) + 1.0
        max_dist = torch.max(distance).cpu().numpy()
        if max_dist>10.0:
            print(f'warning!!! the max distance from the camera is {max_dist:.4f}, which is beyond 10.0 for the ray tracer')
//  [python]
def _warn_ray_tracing(self, centers):
        centers = centers.reshape([-1,3])
        distance = torch.linalg.norm(centers,dim=-1) + 1.0
        max_dist = torch.max(distance).cpu().numpy()
        if max_dist>10.0:
            print(f'warning!!! the max distance from the camera is {max_dist:.4f}, which is beyond 10.0 for the ray tracer')
//  [python]
def compute_validation_info(self, z_vals, rays_o, rays_d, weights, human_poses, step):
        depth = torch.sum(weights * z_vals, -1, keepdim=True)  # rn,
        points = depth * rays_d + rays_o  # rn,3
        gradients = self.sdf_network.gradient(points)  # rn,3
        inner_mask = torch.norm(points, dim=-1, keepdim=True) <= 1.0
        outputs = {
            'depth': depth,  # rn,1
            'normal': ((F.normalize(gradients, dim=-1) + 1.0) * 0.5) * inner_mask,  # rn,3
        }

        feature_vector = self.sdf_network(points)[..., 1:]  # rn,f
        _, occ_info, inter_results = self.color_network(points, gradients, -F.normalize(rays_d, dim=-1), feature_vector, human_poses, inter_results=True, step=step)
        _, occ_prob, _ = get_intersection(self.sdf_inter_fun, self.deviation_network, points, occ_info['reflective'], sn0=128, sn1=9)  # pn,sn-1
        occ_prob_gt = torch.sum(occ_prob, dim=-1, keepdim=True)
        outputs['occ_prob_gt'] = occ_prob_gt
        for k, v in inter_results.items(): inter_results[k] = v * inner_mask
        outputs.update(inter_results)
        return outputs
//  [python]
def compute_validation_info(self, z_vals, rays_o, rays_d, weights, human_poses, step):
        depth = torch.sum(weights * z_vals, -1, keepdim=True)  # rn,
        points = depth * rays_d + rays_o  # rn,3
        gradients = self.sdf_network.gradient(points)  # rn,3
        inner_mask = torch.linalg.norm(points, dim=-1, keepdim=True) <= 1.0
        outputs = {
            'depth': depth,  # rn,1
            'normal': ((F.normalize(gradients, dim=-1) + 1.0) * 0.5) * inner_mask,  # rn,3
        }

        feature_vector = self.sdf_network(points)[..., 1:]  # rn,f
        _, occ_info, inter_results = self.color_network(points, gradients, -F.normalize(rays_d, dim=-1), feature_vector, human_poses, inter_results=True, step=step)
        _, occ_prob, _ = get_intersection(self.sdf_inter_fun, self.deviation_network, points, occ_info['reflective'], sn0=128, sn1=9)  # pn,sn-1
        occ_prob_gt = torch.sum(occ_prob, dim=-1, keepdim=True)
        outputs['occ_prob_gt'] = occ_prob_gt
        for k, v in inter_results.items(): inter_results[k] = v * inner_mask
        outputs.update(inter_results)
        return outputs
//  [python]
def compute_occ_loss(self, occ_info, points, sdf, gradients, dirs, step):
        if step < self.cfg['occ_loss_step']: return torch.zeros(1)

        occ_prob = occ_info['occ_prob']
        reflective = occ_info['reflective']

        # select a subset for occ loss
        # note we only apply occ loss on the surface
        inner_mask = torch.norm(points, dim=-1) < 0.999  # pn
        sdf_mask = torch.abs(sdf) < self.cfg['occ_sdf_thresh']
        normal_mask = torch.sum(gradients * dirs, -1) < 0  # pn
        mask = (inner_mask & normal_mask & sdf_mask)

        if torch.sum(mask) > self.cfg['occ_loss_max_pn']:
            indices = torch.nonzero(mask)[:, 0]  # npn
            idx = torch.randperm(indices.shape[0], device='cuda')  # npn
            indices = indices[idx[:self.cfg['occ_loss_max_pn']]]  # max_pn
            mask_new = torch.zeros_like(mask)
            mask_new[indices] = 1
            mask = mask_new

        if torch.sum(mask) > 0:
            inter_dist, inter_prob, inter_sdf = get_intersection(self.sdf_inter_fun, self.deviation_network, points[mask], reflective[mask], sn0=64, sn1=16)  # pn,sn-1
            occ_prob_gt = torch.sum(inter_prob, -1, keepdim=True)
            return F.l1_loss(occ_prob[mask], occ_prob_gt)
        else:
            return torch.zeros(1)
//  [python]
def compute_occ_loss(self, occ_info, points, sdf, gradients, dirs, step):
        if step < self.cfg['occ_loss_step']: return torch.zeros(1)

        occ_prob = occ_info['occ_prob']
        reflective = occ_info['reflective']

        # select a subset for occ loss
        # note we only apply occ loss on the surface
        inner_mask = torch.linalg.norm(points, dim=-1) < 0.999  # pn
        sdf_mask = torch.abs(sdf) < self.cfg['occ_sdf_thresh']
        normal_mask = torch.sum(gradients * dirs, -1) < 0  # pn
        mask = (inner_mask & normal_mask & sdf_mask)

        if torch.sum(mask) > self.cfg['occ_loss_max_pn']:
            indices = torch.nonzero(mask)[:, 0]  # npn
            idx = torch.randperm(indices.shape[0], device='cuda')  # npn
            indices = indices[idx[:self.cfg['occ_loss_max_pn']]]  # max_pn
            mask_new = torch.zeros_like(mask)
            mask_new[indices] = 1
            mask = mask_new

        if torch.sum(mask) > 0:
            inter_dist, inter_prob, inter_sdf = get_intersection(self.sdf_inter_fun, self.deviation_network, points[mask], reflective[mask], sn0=64, sn1=16)  # pn,sn-1
            occ_prob_gt = torch.sum(inter_prob, -1, keepdim=True)
            return F.l1_loss(occ_prob[mask], occ_prob_gt)
        else:
            return torch.zeros(1)
//  [python]
def gram_schmidt_qr(A):
    m, n = A.shape
    Q = torch.zeros(m, n)
    R = torch.zeros(n, n)

    for j in range(n):
        v = A[:, j]
        for i in range(j):
            R[i, j] = torch.dot(Q[:, i], A[:, j])
            v -= R[i, j] * Q[:, i]
        R[j, j] = torch.norm(v)
        Q[:, j] = v / R[j, j]

    return Q, R
//  [python]
def gram_schmidt_qr(A):
    m, n = A.shape
    Q = torch.zeros(m, n)
    R = torch.zeros(n, n)

    for j in range(n):
        v = A[:, j]
        for i in range(j):
            R[i, j] = torch.dot(Q[:, i], A[:, j])
            v -= R[i, j] * Q[:, i]
        R[j, j] = torch.linalg.norm(v)
        Q[:, j] = v / R[j, j]

    return Q, R
//  [python]
def update_M(self, obs, act, reward, next_obs, next_act):
        with torch.no_grad():
            phi_t = self.phi(obs, act)
            phi_tp1 = self.phi(next_obs, next_act)

        pred_tp1, pred_r = self.M(phi_t, None)
        reward_loss = F.mse_loss(pred_r, reward)

        BC_loss = (torch.norm(pred_tp1 - self.gamma * phi_tp1, dim=1) ** 2).mean()
        M_update_loss = self.cfg.reward_weight * reward_loss + BC_loss

        self.M_opt.zero_grad(set_to_none=True)
        M_update_loss.backward()
        self.M_opt.step()

        return {
            "M/BC_loss": BC_loss,
            "M/reward_loss": reward_loss,
            "M/combined_loss": M_update_loss,
        }
//  [python]
def update_M(self, obs, act, reward, next_obs, next_act):
        with torch.no_grad():
            phi_t = self.phi(obs, act)
            phi_tp1 = self.phi(next_obs, next_act)

        pred_tp1, pred_r = self.M(phi_t, None)
        reward_loss = F.mse_loss(pred_r, reward)

        BC_loss = (torch.linalg.norm(pred_tp1 - self.gamma * phi_tp1, dim=1) ** 2).mean()
        M_update_loss = self.cfg.reward_weight * reward_loss + BC_loss

        self.M_opt.zero_grad(set_to_none=True)
        M_update_loss.backward()
        self.M_opt.step()

        return {
            "M/BC_loss": BC_loss,
            "M/reward_loss": reward_loss,
            "M/combined_loss": M_update_loss,
        }
//  [python]
u_pred_up_data = u_pred_up_data.flatten().cpu()

u_exact_up = np.sqrt(np.sqrt(x**2+y0**2))*np.sqrt((1-x/np.sqrt(x**2+y0**2))/2)
u_exact_up = torch.from_numpy(u_exact_up)
error_up_data = torch.abs(u_pred_up_data - u_exact_up) # get the error in every points
error_up_data = torch.norm(error_up_data)/torch.norm(u_exact_up)
//  [python]
u_pred_up_data = u_pred_up_data.flatten().cpu()

u_exact_up = np.sqrt(np.sqrt(x**2+y0**2))*np.sqrt((1-x/np.sqrt(x**2+y0**2))/2)
u_exact_up = torch.from_numpy(u_exact_up)
error_up_data = torch.abs(u_pred_up_data - u_exact_up) # get the error in every points
error_up_data = torch.linalg.norm(error_up_data)/torch.linalg.norm(u_exact_up)
//  [python]
def update_phi(self, obs, act, reward, next_obs, next_act):
        phi_t = self.phi(obs, act)
        phi_tp1 = self.phi(next_obs, next_act)

        pred_tp1, pred_r = self.M(phi_t, None)

        # Design Loss
        cov = torch.bmm(phi_t.unsqueeze(-1), phi_t.unsqueeze(-2)).mean(dim=0)
        design_loss = -self.log_det(cov)

        # Reward Loss
        reward_loss = F.mse_loss(pred_r, reward)

        # Transition Loss
        BC_loss = (torch.norm(pred_tp1 - self.gamma * phi_tp1, dim=1) ** 2).mean()

        # Combined Loss
        phi_loss = (
            self.cfg.reward_weight * reward_loss
            + BC_loss
            + self.cfg.design_weight * design_loss
        )

        self.phi_opt.zero_grad(set_to_none=True)
        phi_loss.backward()
        torch.nn.utils.clip_grad_norm_(
            self.phi.parameters(), self.cfg.phi_grad_norm_clip
        )
        self.phi_opt.step()

        return {
            "phi/design_loss": design_loss,
            "phi/reward_loss": reward_loss,
            "phi/BC_loss": BC_loss,
            "phi/combined_loss": phi_loss,
        }
//  [python]
def update_phi(self, obs, act, reward, next_obs, next_act):
        phi_t = self.phi(obs, act)
        phi_tp1 = self.phi(next_obs, next_act)

        pred_tp1, pred_r = self.M(phi_t, None)

        # Design Loss
        cov = torch.bmm(phi_t.unsqueeze(-1), phi_t.unsqueeze(-2)).mean(dim=0)
        design_loss = -self.log_det(cov)

        # Reward Loss
        reward_loss = F.mse_loss(pred_r, reward)

        # Transition Loss
        BC_loss = (torch.linalg.norm(pred_tp1 - self.gamma * phi_tp1, dim=1) ** 2).mean()

        # Combined Loss
        phi_loss = (
            self.cfg.reward_weight * reward_loss
            + BC_loss
            + self.cfg.design_weight * design_loss
        )

        self.phi_opt.zero_grad(set_to_none=True)
        phi_loss.backward()
        torch.nn.utils.clip_grad_norm_(
            self.phi.parameters(), self.cfg.phi_grad_norm_clip
        )
        self.phi_opt.step()

        return {
            "phi/design_loss": design_loss,
            "phi/reward_loss": reward_loss,
            "phi/BC_loss": BC_loss,
            "phi/combined_loss": phi_loss,
        }
//  [python]
def RBF(x):
    
    d_total_t = torch.tensor(d_total, device='cuda').unsqueeze(1)
    
    w_t = torch.tensor(w,  device='cuda')

    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1)
    # distance matrix for K matrix
    R2 = torch.norm(d_total_t - x_l, dim=2)
    y = torch.mm(torch.exp(-gama*R2.T), w_t)
    return y
//  [python]
def RBF(x):
    
    d_total_t = torch.tensor(d_total, device='cuda').unsqueeze(1)
    
    w_t = torch.tensor(w,  device='cuda')

    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1)
    # distance matrix for K matrix
    R2 = torch.linalg.norm(d_total_t - x_l, dim=2)
    y = torch.mm(torch.exp(-gama*R2.T), w_t)
    return y
//  [python]
def RBF(x):
    d_total_t = torch.from_numpy(d_total).unsqueeze(1).cuda()
    w_t = torch.from_numpy(w).cuda()
    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1) 
    # 得到R2大矩阵
    R2 = torch.norm(d_total_t - x_l, dim=2)
    y = torch.mm(torch.exp(-gama*R2.T), w_t)
    return y
//  [python]
def RBF(x):
    d_total_t = torch.from_numpy(d_total).unsqueeze(1).cuda()
    w_t = torch.from_numpy(w).cuda()
    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1) 
    # 得到R2大矩阵
    R2 = torch.linalg.norm(d_total_t - x_l, dim=2)
    y = torch.mm(torch.exp(-gama*R2.T), w_t)
    return y
//  [python]
def evaluate():
    N_test = 100
    x = np.linspace(-1, 1, N_test)
    y = np.linspace(-1, 1, N_test)
    x, y = np.meshgrid(x, y)
    xy_test = np.stack((x.flatten(), y.flatten()),1)
    xy_test = torch.from_numpy(xy_test).cuda() # to the model for the prediction
    
    u_pred = pred(xy_test)
    u_pred = u_pred.data
    u_pred = u_pred.reshape(N_test, N_test)
    

    u_exact = np.zeros(x.shape)
    u_exact[y>0] = np.sqrt(np.sqrt(x[y>0]**2+y[y>0]**2))*np.sqrt((1-x[y>0]/np.sqrt(x[y>0]**2+y[y>0]**2))/2)
    u_exact[y<0] = -np.sqrt(np.sqrt(x[y<0]**2+y[y<0]**2))*np.sqrt((1-x[y<0]/np.sqrt(x[y<0]**2+y[y<0]**2))/2)

    u_exact = u_exact.reshape(x.shape)
    u_exact = torch.tensor(u_exact)
    error = torch.abs(u_pred - u_exact) # get the error in every points
    error_t = torch.norm(error)/torch.norm(u_exact) # get the total relative L2 error
    return error_t
//  [python]
def evaluate():
    N_test = 100
    x = np.linspace(-1, 1, N_test)
    y = np.linspace(-1, 1, N_test)
    x, y = np.meshgrid(x, y)
    xy_test = np.stack((x.flatten(), y.flatten()),1)
    xy_test = torch.from_numpy(xy_test).cuda() # to the model for the prediction
    
    u_pred = pred(xy_test)
    u_pred = u_pred.data
    u_pred = u_pred.reshape(N_test, N_test)
    

    u_exact = np.zeros(x.shape)
    u_exact[y>0] = np.sqrt(np.sqrt(x[y>0]**2+y[y>0]**2))*np.sqrt((1-x[y>0]/np.sqrt(x[y>0]**2+y[y>0]**2))/2)
    u_exact[y<0] = -np.sqrt(np.sqrt(x[y<0]**2+y[y<0]**2))*np.sqrt((1-x[y<0]/np.sqrt(x[y<0]**2+y[y<0]**2))/2)

    u_exact = u_exact.reshape(x.shape)
    u_exact = torch.tensor(u_exact)
    error = torch.abs(u_pred - u_exact) # get the error in every points
    error_t = torch.linalg.norm(error)/torch.linalg.norm(u_exact) # get the total relative L2 error
    return error_t
//  [python]
def train_data(Nb, Nf):
    '''
    generate essential boundary points and domain points 
    '''
    

    xu = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  np.ones([int(Nb/4),1])])
    xd = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  -np.ones([int(Nb/4),1])])
    xl = np.hstack([-np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xr = np.hstack([np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xcrack = np.hstack([-np.random.rand(int(Nb/4),1), np.zeros([int(Nb/4),1])])
    # add corner points when attribute random points 
    xc1 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]]) 
    xc2 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    xc3 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    Xb = np.concatenate((xu, xd, xl, xr, xcrack, xc1, xc2, xc3)) 

    Xb = torch.tensor(Xb, device='cuda') 
    Xf = torch.rand(Nf, 2)*2-1
    
    Xf1 = Xf[(Xf[:, 1]>0) & (torch.norm(Xf, dim=1)>=delta)] # points in the up region
    Xf2 = Xf[(Xf[:, 1]<0) & (torch.norm(Xf, dim=1)>=delta)] # points in the down region 
    
    Xf1 = torch.tensor(Xf1, requires_grad=True, device='cuda')
    Xf2 = torch.tensor(Xf2,  requires_grad=True, device='cuda')
    
    return Xb, Xf1, Xf2
//  [python]
def train_data(Nb, Nf):
    '''
    generate essential boundary points and domain points 
    '''
    

    xu = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  np.ones([int(Nb/4),1])])
    xd = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  -np.ones([int(Nb/4),1])])
    xl = np.hstack([-np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xr = np.hstack([np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xcrack = np.hstack([-np.random.rand(int(Nb/4),1), np.zeros([int(Nb/4),1])])
    # add corner points when attribute random points 
    xc1 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]]) 
    xc2 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    xc3 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    Xb = np.concatenate((xu, xd, xl, xr, xcrack, xc1, xc2, xc3)) 

    Xb = torch.tensor(Xb, device='cuda') 
    Xf = torch.rand(Nf, 2)*2-1
    
    Xf1 = Xf[(Xf[:, 1]>0) & (torch.linalg.norm(Xf, dim=1)>=delta)] # points in the up region
    Xf2 = Xf[(Xf[:, 1]<0) & (torch.linalg.norm(Xf, dim=1)>=delta)] # points in the down region 
    
    Xf1 = torch.tensor(Xf1, requires_grad=True, device='cuda')
    Xf2 = torch.tensor(Xf2,  requires_grad=True, device='cuda')
    
    return Xb, Xf1, Xf2
//  [python]
def render_core(self, rays_o, rays_d, z_vals, human_poses, cos_anneal_ratio=0.0, step=None, is_train=True):
        batch_size, n_samples = z_vals.shape

        # section length in original space
        dists = z_vals[..., 1:] - z_vals[..., :-1]  # rn,sn-1
        dists = torch.cat([dists, dists[..., -1:]], -1)  # rn,sn
        mid_z_vals = z_vals + dists * 0.5

        points = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * mid_z_vals.unsqueeze(-1)
        inner_mask = torch.norm(points, dim=-1) <= 1.0
        outer_mask = ~inner_mask

        dirs = rays_d.unsqueeze(-2).expand(batch_size, n_samples, 3)
        human_poses_pt = human_poses.unsqueeze(-3).expand(batch_size, n_samples, 3, 4)
        dirs = F.normalize(dirs, dim=-1)
        alpha, sampled_color = torch.zeros(batch_size, n_samples), torch.zeros(batch_size, n_samples, 3)

        if torch.sum(outer_mask) > 0:
            alpha[outer_mask], sampled_color[outer_mask] = self.compute_density_alpha(points[outer_mask], dists[outer_mask], -dirs[outer_mask], self.outer_nerf)

        if torch.sum(inner_mask) > 0:
            alpha[inner_mask], gradients, feature_vector, inv_s, sdf = self.compute_sdf_alpha(points[inner_mask], dists[inner_mask], dirs[inner_mask], cos_anneal_ratio, step)
            sampled_color[inner_mask], occ_info = self.color_network(points[inner_mask], gradients, -dirs[inner_mask], feature_vector, human_poses_pt[inner_mask], step=step)
            # Eikonal loss
            gradient_error = (torch.linalg.norm(gradients, ord=2, dim=-1) - 1.0) ** 2
        else:
            gradient_error = torch.zeros(1)

        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[..., :-1]  # rn,sn
        color = (sampled_color * weights[..., None]).sum(dim=1)

        outputs = {
            'ray_rgb': color,  # rn,3
            'gradient_error': gradient_error,  # rn
        }

        if torch.sum(inner_mask) > 0:
            outputs['std'] = torch.mean(1 / inv_s)
        else:
            outputs['std'] = torch.zeros(1)

        if step < 1000:
            mask = torch.norm(points, dim=-1) < 1.2
            outputs['sdf_pts'] = points[mask]
            outputs['sdf_vals'] = self.sdf_network.sdf(points[mask])[..., 0]

        if self.cfg['apply_occ_loss']:
            # occlusion loss
            if torch.sum(inner_mask) > 0:
                outputs['loss_occ'] = self.compute_occ_loss(occ_info, points[inner_mask], sdf, gradients, dirs[inner_mask], step)
            else:
                outputs['loss_occ'] = torch.zeros(1)

        if not is_train:
            outputs.update(self.compute_validation_info(z_vals, rays_o, rays_d, weights, human_poses, step))

        return outputs
//  [python]
def render_core(self, rays_o, rays_d, z_vals, human_poses, cos_anneal_ratio=0.0, step=None, is_train=True):
        batch_size, n_samples = z_vals.shape

        # section length in original space
        dists = z_vals[..., 1:] - z_vals[..., :-1]  # rn,sn-1
        dists = torch.cat([dists, dists[..., -1:]], -1)  # rn,sn
        mid_z_vals = z_vals + dists * 0.5

        points = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * mid_z_vals.unsqueeze(-1)
        inner_mask = torch.linalg.norm(points, dim=-1) <= 1.0
        outer_mask = ~inner_mask

        dirs = rays_d.unsqueeze(-2).expand(batch_size, n_samples, 3)
        human_poses_pt = human_poses.unsqueeze(-3).expand(batch_size, n_samples, 3, 4)
        dirs = F.normalize(dirs, dim=-1)
        alpha, sampled_color = torch.zeros(batch_size, n_samples), torch.zeros(batch_size, n_samples, 3)

        if torch.sum(outer_mask) > 0:
            alpha[outer_mask], sampled_color[outer_mask] = self.compute_density_alpha(points[outer_mask], dists[outer_mask], -dirs[outer_mask], self.outer_nerf)

        if torch.sum(inner_mask) > 0:
            alpha[inner_mask], gradients, feature_vector, inv_s, sdf = self.compute_sdf_alpha(points[inner_mask], dists[inner_mask], dirs[inner_mask], cos_anneal_ratio, step)
            sampled_color[inner_mask], occ_info = self.color_network(points[inner_mask], gradients, -dirs[inner_mask], feature_vector, human_poses_pt[inner_mask], step=step)
            # Eikonal loss
            gradient_error = (torch.linalg.norm(gradients, ord=2, dim=-1) - 1.0) ** 2
        else:
            gradient_error = torch.zeros(1)

        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[..., :-1]  # rn,sn
        color = (sampled_color * weights[..., None]).sum(dim=1)

        outputs = {
            'ray_rgb': color,  # rn,3
            'gradient_error': gradient_error,  # rn
        }

        if torch.sum(inner_mask) > 0:
            outputs['std'] = torch.mean(1 / inv_s)
        else:
            outputs['std'] = torch.zeros(1)

        if step < 1000:
            mask = torch.linalg.norm(points, dim=-1) < 1.2
            outputs['sdf_pts'] = points[mask]
            outputs['sdf_vals'] = self.sdf_network.sdf(points[mask])[..., 0]

        if self.cfg['apply_occ_loss']:
            # occlusion loss
            if torch.sum(inner_mask) > 0:
                outputs['loss_occ'] = self.compute_occ_loss(occ_info, points[inner_mask], sdf, gradients, dirs[inner_mask], step)
            else:
                outputs['loss_occ'] = torch.zeros(1)

        if not is_train:
            outputs.update(self.compute_validation_info(z_vals, rays_o, rays_d, weights, human_poses, step))

        return outputs
//  [python]
def psi_update(psi, alpha, tres, sigma, eta_psi, lamb, modified=False):
    if modified:
        dpsi = eta(1, eta_0=eta_psi)*grad_psi_modified(psi, tres, alpha, sigma, lamb) # (D,K)
        new_psi = psi + dpsi
    else:
        dpsi = eta(1, eta_0=eta_psi)*grad_psi(tres, alpha, sigma) # (D,K)
        new_psi = f.normalize(psi + dpsi, dim=0)
#     if torch.isnan(new_psi).any():
#         print("Infinity in new_psi. Setting dpsi = 0.")
#         new_psi = psi.clone()
    dpsi_length = torch.mean(torch.norm(new_psi - psi, dim=0))
    psi = new_psi.clone()
    return psi, dpsi_length
//  [python]
def psi_update(psi, alpha, tres, sigma, eta_psi, lamb, modified=False):
    if modified:
        dpsi = eta(1, eta_0=eta_psi)*grad_psi_modified(psi, tres, alpha, sigma, lamb) # (D,K)
        new_psi = psi + dpsi
    else:
        dpsi = eta(1, eta_0=eta_psi)*grad_psi(tres, alpha, sigma) # (D,K)
        new_psi = f.normalize(psi + dpsi, dim=0)
#     if torch.isnan(new_psi).any():
#         print("Infinity in new_psi. Setting dpsi = 0.")
#         new_psi = psi.clone()
    dpsi_length = torch.mean(torch.linalg.norm(new_psi - psi, dim=0))
    psi = new_psi.clone()
    return psi, dpsi_length
//  [python]
def w_update_riemann(optimizer, x, w, q, res, sigma, modified=False):
    """
    optimizer - geoopt.optim.RiemannianAdam/geoopt.optim.RiemannianSGD
    psi, w - geoopt.ManifoldParameter
    """
    w_old = w.clone()
    optimizer.zero_grad()
    if modified:
        w.grad = -grad_w_modified(x, res, w, q, sigma)
    else:
        w.grad = -grad_w(x, res, w, q, sigma)
    optimizer.step()
    dw_norm = torch.norm(w - w_old)
    return w, dw_norm
//  [python]
def w_update_riemann(optimizer, x, w, q, res, sigma, modified=False):
    """
    optimizer - geoopt.optim.RiemannianAdam/geoopt.optim.RiemannianSGD
    psi, w - geoopt.ManifoldParameter
    """
    w_old = w.clone()
    optimizer.zero_grad()
    if modified:
        w.grad = -grad_w_modified(x, res, w, q, sigma)
    else:
        w.grad = -grad_w(x, res, w, q, sigma)
    optimizer.step()
    dw_norm = torch.linalg.norm(w - w_old)
    return w, dw_norm
//  [python]
def _l2normalize(self, x, eps=1e-12):
        r"""Function to calculate the ``L2 Normalized`` form of a Tensor
        Args:
            x (torch.Tensor): Tensor which needs to be normalized.
            eps (float, optional): A small value needed to avoid infinite values.
        Returns:
            Normalized form of the tensor ``x``.
        """
        return x / (torch.norm(x) + eps)
//  [python]
def _l2normalize(self, x, eps=1e-12):
        r"""Function to calculate the ``L2 Normalized`` form of a Tensor
        Args:
            x (torch.Tensor): Tensor which needs to be normalized.
            eps (float, optional): A small value needed to avoid infinite values.
        Returns:
            Normalized form of the tensor ``x``.
        """
        return x / (torch.linalg.norm(x) + eps)
//  [python]
def train_data(Nb, Nf):
    xu = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  np.ones([int(Nb/4),1])])
    xd = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  -np.ones([int(Nb/4),1])])
    xl = np.hstack([-np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xr = np.hstack([np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xcrack = np.hstack([-np.random.rand(int(Nb/4),1), np.zeros([int(Nb/4),1])])
    xc1 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]]) 
    xc2 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    xc3 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    Xb = np.concatenate((xu, xd, xl, xr, xcrack, xc1, xc2, xc3))
    Xb = torch.tensor(Xb, device='cuda')   
    Xf = torch.rand(Nf, 2)*2-1

    Xf1 = Xf[(Xf[:, 1]>0) & (torch.norm(Xf, dim=1)>=delta)]
    Xf2 = Xf[(Xf[:, 1]<0) & (torch.norm(Xf, dim=1)>=delta)] 
    
    Xf1 = torch.tensor(Xf1, requires_grad=True, device='cuda')
    Xf2 = torch.tensor(Xf2, requires_grad=True, device='cuda')
    
    return Xb, Xf1, Xf2
//  [python]
def train_data(Nb, Nf):
    xu = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  np.ones([int(Nb/4),1])])
    xd = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  -np.ones([int(Nb/4),1])])
    xl = np.hstack([-np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xr = np.hstack([np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xcrack = np.hstack([-np.random.rand(int(Nb/4),1), np.zeros([int(Nb/4),1])])
    xc1 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]]) 
    xc2 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    xc3 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    Xb = np.concatenate((xu, xd, xl, xr, xcrack, xc1, xc2, xc3))
    Xb = torch.tensor(Xb, device='cuda')   
    Xf = torch.rand(Nf, 2)*2-1

    Xf1 = Xf[(Xf[:, 1]>0) & (torch.linalg.norm(Xf, dim=1)>=delta)]
    Xf2 = Xf[(Xf[:, 1]<0) & (torch.linalg.norm(Xf, dim=1)>=delta)] 
    
    Xf1 = torch.tensor(Xf1, requires_grad=True, device='cuda')
    Xf2 = torch.tensor(Xf2, requires_grad=True, device='cuda')
    
    return Xb, Xf1, Xf2
//  [python]
def train_data(Nb, Nf):
    xu = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  np.ones([int(Nb/4),1])])
    xd = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  -np.ones([int(Nb/4),1])])
    xl = np.hstack([-np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xr = np.hstack([np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xcrack = np.hstack([-np.random.rand(int(Nb/4),1), np.zeros([int(Nb/4),1])])

    xc1 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]]) 
    xc2 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    xc3 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    Xb = np.concatenate((xu, xd, xl, xr, xcrack, xc1, xc2, xc3)) 

    Xb = torch.tensor(Xb, device='cuda') 
    
    Xf = torch.rand(Nf, 2)*2-1
    

    Xf1 = Xf[(Xf[:, 1]>0) & (torch.norm(Xf, dim=1)>=delta)] 
    Xf2 = Xf[(Xf[:, 1]<0) & (torch.norm(Xf, dim=1)>=delta)]  
    
    Xf1 = torch.tensor(Xf1, requires_grad=True, device='cuda')
    Xf2 = torch.tensor(Xf2, requires_grad=True, device='cuda')
    
    return Xb, Xf1, Xf2