jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _calc_reciprocal_energy_matrix(self):
        # calculate reciprocal points
        recip = get_reciprocal_vectors(self.cell)
        shifts = get_shifts_within_cutoff(recip, self.cutoff_recip)  # (num_shifts, 3)
        ks_all = torch.matmul(shifts, recip)
        length_all = torch.linalg.norm(ks_all, dim=-1)  # (num_shifts, )

        # retrieve reciprocal points whose length are shorter than cutoff
        within_cutoff = (length_all > self.eps) & (length_all < self.cutoff_recip)
        ks = ks_all[within_cutoff]
        length = length_all[within_cutoff]
        # disps_ij[i, j, :] is displacement vector r_{ij}, (num_atoms, num_atoms, 3)
        disps_ij = self.pos[None, :, :] - self.pos[:, None, :]
        phases = torch.sum(ks[:, None, None, :] * disps_ij[None, :, :, :], dim=-1)

        e_recip_matrix_aug = (
            torch.cos(phases)
            * torch.exp(-0.5 * torch.square(self.eta * length[:, None, None]))
            / torch.square(length[:, None, None])
        )
        e_recip_matrix = (
            self.COULOMB_FACTOR
            * 4.0
            * math.pi
            / self.volume
            * torch.sum(e_recip_matrix_aug, dim=0)
        )
        return e_recip_matrix
//  [python]
def get_shifts_within_cutoff(cell, cutoff):
    """
    Return all shifts required to search for atoms within cutoff
    """
    device = cell.device

    # projected length for three planes
    proj = torch.zeros(3, device=device)
    nx = torch.cross(cell[1], cell[2])
    ny = torch.cross(cell[2], cell[0])
    nz = torch.cross(cell[0], cell[1])
    proj[0] = torch.dot(cell[0], nx / torch.norm(nx))
    proj[1] = torch.dot(cell[1], ny / torch.norm(ny))
    proj[2] = torch.dot(cell[2], nz / torch.norm(nz))

    shift = torch.ceil(cutoff / torch.abs(proj))

    grid = torch.cartesian_prod(
        torch.arange(-shift[0], shift[0] + 1, device=device),
        torch.arange(-shift[1], shift[1] + 1, device=device),
        torch.arange(-shift[2], shift[2] + 1, device=device),
    )

    return grid
//  [python]
def get_shifts_within_cutoff(cell, cutoff):
    """
    Return all shifts required to search for atoms within cutoff
    """
    device = cell.device

    # projected length for three planes
    proj = torch.zeros(3, device=device)
    nx = torch.cross(cell[1], cell[2])
    ny = torch.cross(cell[2], cell[0])
    nz = torch.cross(cell[0], cell[1])
    proj[0] = torch.dot(cell[0], nx / torch.linalg.norm(nx))
    proj[1] = torch.dot(cell[1], ny / torch.linalg.norm(ny))
    proj[2] = torch.dot(cell[2], nz / torch.linalg.norm(nz))

    shift = torch.ceil(cutoff / torch.abs(proj))

    grid = torch.cartesian_prod(
        torch.arange(-shift[0], shift[0] + 1, device=device),
        torch.arange(-shift[1], shift[1] + 1, device=device),
        torch.arange(-shift[2], shift[2] + 1, device=device),
    )

    return grid
//  [python]
def _calc_real_energy_matrix(self):
        """
        Calculate real-space-part energy in atomic unit
        """
        # calculate length between atoms `i` and `j` with `shift`
        shifts = get_shifts_within_cutoff(self.cell, self.cutoff_real)  # (num_shifts, 3)
        # disps_ij[i, j, :] is displacement vector r_{ij}
        disps_ij = self.pos[None, :, :] - self.pos[:, None, :]
        disps = disps_ij[None, :, :, :] + torch.matmul(shifts, self.cell)[:, None, None, :]
        distances_all = torch.norm(disps, dim=-1)  # (num_shifts, num_atoms, num_atoms)

        # retrieve pairs whose length are shorter than cutoff
        within_cutoff = (distances_all > self.eps) & (distances_all < self.cutoff_real)
        distances = distances_all[within_cutoff]

        e_real_matrix_aug = torch.zeros_like(distances_all)
        e_real_matrix_aug[within_cutoff] = torch.erfc(distances / (math.sqrt(2) * self.eta))
        if not self.point_charge:
            gammas_all = torch.sqrt(
                torch.square(self.sigmas[:, None]) + torch.square(self.sigmas[None, :])
            )
            gammas = torch.broadcast_to(gammas_all, distances_all.shape)[within_cutoff]
            e_real_matrix_aug[within_cutoff] -= torch.erfc(distances / (math.sqrt(2) * gammas))
        e_real_matrix_aug[within_cutoff] /= distances
        e_real_matrix = self.COULOMB_FACTOR * torch.sum(
            e_real_matrix_aug, dim=0
        )  # sum over shifts
        return e_real_matrix
//  [python]
def _calc_real_energy_matrix(self):
        """
        Calculate real-space-part energy in atomic unit
        """
        # calculate length between atoms `i` and `j` with `shift`
        shifts = get_shifts_within_cutoff(self.cell, self.cutoff_real)  # (num_shifts, 3)
        # disps_ij[i, j, :] is displacement vector r_{ij}
        disps_ij = self.pos[None, :, :] - self.pos[:, None, :]
        disps = disps_ij[None, :, :, :] + torch.matmul(shifts, self.cell)[:, None, None, :]
        distances_all = torch.linalg.norm(disps, dim=-1)  # (num_shifts, num_atoms, num_atoms)

        # retrieve pairs whose length are shorter than cutoff
        within_cutoff = (distances_all > self.eps) & (distances_all < self.cutoff_real)
        distances = distances_all[within_cutoff]

        e_real_matrix_aug = torch.zeros_like(distances_all)
        e_real_matrix_aug[within_cutoff] = torch.erfc(distances / (math.sqrt(2) * self.eta))
        if not self.point_charge:
            gammas_all = torch.sqrt(
                torch.square(self.sigmas[:, None]) + torch.square(self.sigmas[None, :])
            )
            gammas = torch.broadcast_to(gammas_all, distances_all.shape)[within_cutoff]
            e_real_matrix_aug[within_cutoff] -= torch.erfc(distances / (math.sqrt(2) * gammas))
        e_real_matrix_aug[within_cutoff] /= distances
        e_real_matrix = self.COULOMB_FACTOR * torch.sum(
            e_real_matrix_aug, dim=0
        )  # sum over shifts
        return e_real_matrix
//  [python]
def arap_deformation_loss(trajectory, node_radius=None, trajectory_rot=None, K=50, with_rot=True):
    init_pcl = trajectory[:, 0]
    radius = torch.norm(init_pcl.max(dim=0).values - init_pcl.min(dim=0).values) / 8
    fid = torch.randint(1, trajectory.shape[1], [])
    tar_pcl = trajectory[:, fid]

    N = init_pcl.shape[0]
    with torch.no_grad():
        radius = torch.norm(init_pcl.max(dim=0).values - init_pcl.min(dim=0).values) / 8
        device = init_pcl.device
        ii, jj, nn, weight = cal_connectivity_from_points(init_pcl, radius, K, trajectory=trajectory.detach(), node_radius=node_radius, mode='nn')
        L_opt = torch.eye(N).cuda()
        L_opt[ii, jj] = - weight[ii, nn]
    
    P = produce_edge_matrix_nfmt(init_pcl, (N, K, 3), ii, jj, nn, device=device)
    P_prime = produce_edge_matrix_nfmt(tar_pcl, (N, K, 3), ii, jj, nn, device=device)
    
    with torch.no_grad():
        D = torch.diag_embed(weight, dim1=1, dim2=2)
        S = torch.bmm(P.permute(0, 2, 1), torch.bmm(D, P_prime))
        U, sig, W = torch.svd(S)
        R = torch.bmm(W, U.permute(0, 2, 1))
        with torch.no_grad():
            # Need to flip the column of U corresponding to smallest singular value
            # for any det(Ri) <= 0
            entries_to_flip = torch.nonzero(torch.det(R) <= 0, as_tuple=False).flatten()  # idxs where det(R) <= 0
            if len(entries_to_flip) > 0:
                Umod = U.clone()
                cols_to_flip = torch.argmin(sig[entries_to_flip], dim=1)  # Get minimum singular value for each entry
                Umod[entries_to_flip, :, cols_to_flip] *= -1  # flip cols
                R[entries_to_flip] = torch.bmm(W[entries_to_flip], Umod[entries_to_flip].permute(0, 2, 1))
    arap_error = (weight[..., None] * (P_prime - torch.einsum('bxy,bky->bkx', R, P))).square().mean(dim=0).sum()

    if with_rot:
        init_rot = quaternion_to_matrix(trajectory_rot[:, 0])
        tar_rot = quaternion_to_matrix(trajectory_rot[:, fid])
        R_rot = torch.bmm(R, init_rot)
        rot_error = (R_rot - tar_rot).square().mean(dim=0).sum()
    else:
        rot_error = 0.

    return arap_error, rot_error * 1e2
//  [python]
def arap_deformation_loss(trajectory, node_radius=None, trajectory_rot=None, K=50, with_rot=True):
    init_pcl = trajectory[:, 0]
    radius = torch.linalg.norm(init_pcl.max(dim=0).values - init_pcl.min(dim=0).values) / 8
    fid = torch.randint(1, trajectory.shape[1], [])
    tar_pcl = trajectory[:, fid]

    N = init_pcl.shape[0]
    with torch.no_grad():
        radius = torch.linalg.norm(init_pcl.max(dim=0).values - init_pcl.min(dim=0).values) / 8
        device = init_pcl.device
        ii, jj, nn, weight = cal_connectivity_from_points(init_pcl, radius, K, trajectory=trajectory.detach(), node_radius=node_radius, mode='nn')
        L_opt = torch.eye(N).cuda()
        L_opt[ii, jj] = - weight[ii, nn]
    
    P = produce_edge_matrix_nfmt(init_pcl, (N, K, 3), ii, jj, nn, device=device)
    P_prime = produce_edge_matrix_nfmt(tar_pcl, (N, K, 3), ii, jj, nn, device=device)
    
    with torch.no_grad():
        D = torch.diag_embed(weight, dim1=1, dim2=2)
        S = torch.bmm(P.permute(0, 2, 1), torch.bmm(D, P_prime))
        U, sig, W = torch.svd(S)
        R = torch.bmm(W, U.permute(0, 2, 1))
        with torch.no_grad():
            # Need to flip the column of U corresponding to smallest singular value
            # for any det(Ri) <= 0
            entries_to_flip = torch.nonzero(torch.det(R) <= 0, as_tuple=False).flatten()  # idxs where det(R) <= 0
            if len(entries_to_flip) > 0:
                Umod = U.clone()
                cols_to_flip = torch.argmin(sig[entries_to_flip], dim=1)  # Get minimum singular value for each entry
                Umod[entries_to_flip, :, cols_to_flip] *= -1  # flip cols
                R[entries_to_flip] = torch.bmm(W[entries_to_flip], Umod[entries_to_flip].permute(0, 2, 1))
    arap_error = (weight[..., None] * (P_prime - torch.einsum('bxy,bky->bkx', R, P))).square().mean(dim=0).sum()

    if with_rot:
        init_rot = quaternion_to_matrix(trajectory_rot[:, 0])
        tar_rot = quaternion_to_matrix(trajectory_rot[:, fid])
        R_rot = torch.bmm(R, init_rot)
        rot_error = (R_rot - tar_rot).square().mean(dim=0).sum()
    else:
        rot_error = 0.

    return arap_error, rot_error * 1e2
//  [python]
def compute_norm(self, grad: Tensor):
        denom = (self.preconditioner / self.bias_correction2).sqrt().add_(self.epsilon)
        adagrad_nrm = torch.norm(grad / denom)
        return adagrad_nrm
//  [python]
def compute_norm(self, grad: Tensor):
        denom = (self.preconditioner / self.bias_correction2).sqrt().add_(self.epsilon)
        adagrad_nrm = torch.linalg.norm(grad / denom)
        return adagrad_nrm
//  [python]
def compute_norm(self, grad: Tensor) -> Tensor:
        return torch.norm(self.precondition(grad))
//  [python]
def compute_norm(self, grad: Tensor) -> Tensor:
        return torch.linalg.norm(self.precondition(grad))
//  [python]
def forward(self, positions):
        def contract(x):
            x = x * self.scale_factor
            
            mag = torch.norm(x, p=self.order, dim=-1)
            mask = mag >= 1
            x_new = x.clone()
            x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])

            return x_new

        if isinstance(positions, Gaussians):
            means = contract(positions.mean.clone())

            if self.scale_factor != 1.0:
                raise NotImplementedError('Scaling of gaussian covariances is not implemented.')
            
            # FIXME: covariances are not handled for now!
            contract = lambda x: (2 - (1 / torch.norm(x, p=self.order, dim=-1, keepdim=True))) * (
                x / torch.norm(x, p=self.order, dim=-1, keepdim=True)
            )
            jc_means = vmap(jacrev(contract))(positions.mean.view(-1, positions.mean.shape[-1]))
            jc_means = jc_means.view(list(positions.mean.shape) + [positions.mean.shape[-1]])

            # Only update covariances on positions outside the unit sphere
            mag = positions.mean.norm(dim=-1)
            mask = mag >= 1
            cov = positions.cov.clone()
            cov[mask] = jc_means[mask] @ positions.cov[mask] @ torch.transpose(jc_means[mask], -2, -1)

            return Gaussians(mean=means, cov=cov)

        return contract(positions)
//  [python]
def forward(self, positions):
        def contract(x):
            x = x * self.scale_factor
            
            mag = torch.linalg.norm(x, ord=self.order, dim=-1)
            mask = mag >= 1
            x_new = x.clone()
            x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])

            return x_new

        if isinstance(positions, Gaussians):
            means = contract(positions.mean.clone())

            if self.scale_factor != 1.0:
                raise NotImplementedError('Scaling of gaussian covariances is not implemented.')
            
            # FIXME: covariances are not handled for now!
            contract = lambda x: (2 - (1 / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True))) * (
                x / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True)
            )
            jc_means = vmap(jacrev(contract))(positions.mean.view(-1, positions.mean.shape[-1]))
            jc_means = jc_means.view(list(positions.mean.shape) + [positions.mean.shape[-1]])

            # Only update covariances on positions outside the unit sphere
            mag = positions.mean.norm(dim=-1)
            mask = mag >= 1
            cov = positions.cov.clone()
            cov[mask] = jc_means[mask] @ positions.cov[mask] @ torch.transpose(jc_means[mask], -2, -1)

            return Gaussians(mean=means, cov=cov)

        return contract(positions)
//  [python]
def direction_norm(self, grad: Tensor) -> Tensor:
        return torch.norm(grad)
//  [python]
def direction_norm(self, grad: Tensor) -> Tensor:
        return torch.linalg.norm(grad)
//  [python]
def compute_accel_norm(joints):
    """
    :param joints (T, J, 3)
    """
    vel = joints[1:] - joints[:-1]  # (T-1, J, 3)
    acc = vel[1:] - vel[:-1]  # (T-2, J, 3)
    return torch.norm(acc, dim=-1)
//  [python]
def compute_accel_norm(joints):
    """
    :param joints (T, J, 3)
    """
    vel = joints[1:] - joints[:-1]  # (T-1, J, 3)
    acc = vel[1:] - vel[:-1]  # (T-2, J, 3)
    return torch.linalg.norm(acc, dim=-1)
//  [python]
def precondition(self, grad: Tensor) -> Tensor:

        preconditioned_grad = grad.clone()

        # iterate over all dimensions
        for k, _ in enumerate(self.dims):
            preconditioner = self.preconditioners[k]
            inv_preconditioner = self.inv_preconditioners[k]
            preconditioner_type = self.preconditioner_types[k]

            # handle diagonal case while retaining dims
            if self.diagonal_threshold is not None:

                # precondition in diagonal case
                if preconditioner_type == PreconditionerType.DIAGONAL:
                    denom = (preconditioner / self.bias_correction2).add_(self.epsilon)
                    preconditioned_grad.div_(denom.pow(-1 / (2 * self.order))[(None,) * k + (...,) + (None,) * (self.order - k - 1)])

                # precondition in full-matrix case
                else:
                    gradient_idx = [*range(1, self.order + 1)]
                    matrix_product_idx = deepcopy(gradient_idx)
                    matrix_product_idx[k] = 0
                    preconditioned_grad = torch.einsum(
                        inv_preconditioner,
                        [0, k + 1],
                        preconditioned_grad,
                        gradient_idx,
                        matrix_product_idx,
                    )

            # more efficient but transposes grad continually
            else:
                preconditioned_grad = torch.tensordot(preconditioned_grad, inv_preconditioner, [[0], [0]])

        # apply grafting
        if self.grafting_type != GraftingType.NONE:
            grafting_norm = self.grafting.direction_norm(grad)
            preconditioned_grad = preconditioned_grad * grafting_norm / (torch.norm(preconditioned_grad) + 1e-16)

        return preconditioned_grad
//  [python]
def precondition(self, grad: Tensor) -> Tensor:

        preconditioned_grad = grad.clone()

        # iterate over all dimensions
        for k, _ in enumerate(self.dims):
            preconditioner = self.preconditioners[k]
            inv_preconditioner = self.inv_preconditioners[k]
            preconditioner_type = self.preconditioner_types[k]

            # handle diagonal case while retaining dims
            if self.diagonal_threshold is not None:

                # precondition in diagonal case
                if preconditioner_type == PreconditionerType.DIAGONAL:
                    denom = (preconditioner / self.bias_correction2).add_(self.epsilon)
                    preconditioned_grad.div_(denom.pow(-1 / (2 * self.order))[(None,) * k + (...,) + (None,) * (self.order - k - 1)])

                # precondition in full-matrix case
                else:
                    gradient_idx = [*range(1, self.order + 1)]
                    matrix_product_idx = deepcopy(gradient_idx)
                    matrix_product_idx[k] = 0
                    preconditioned_grad = torch.einsum(
                        inv_preconditioner,
                        [0, k + 1],
                        preconditioned_grad,
                        gradient_idx,
                        matrix_product_idx,
                    )

            # more efficient but transposes grad continually
            else:
                preconditioned_grad = torch.tensordot(preconditioned_grad, inv_preconditioner, [[0], [0]])

        # apply grafting
        if self.grafting_type != GraftingType.NONE:
            grafting_norm = self.grafting.direction_norm(grad)
            preconditioned_grad = preconditioned_grad * grafting_norm / (torch.linalg.norm(preconditioned_grad) + 1e-16)

        return preconditioned_grad
//  [python]
def forward(self, positions):
        def _contract(x):
            mag = torch.norm(x, p=self.order, dim=-1)
            mask = mag >= 1
            x_new = x.clone()
            x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])

            return x_new

        def contract(x):
            x = x * self.scale_factor
            if self.aabb is None:
                mag = torch.norm(x, p=self.order, dim=-1)
                mask = mag >= 1
            else:
                in_aabb_mask = ((x >= self.aabb[0]) & (x <= self.aabb[1])).all(-1)  # whether including = ?
                mask = ~in_aabb_mask
            x_new = x.clone()
            x_new[mask] = x[mask] / self.alpha
            return _contract(x_new)

        if isinstance(positions, Gaussians):
            means = contract(positions.mean.clone())
            # FIXME: covariances are not handled for now!
            if self.scale_factor != 1.0:
                raise NotImplementedError('Scaling of gaussian covariances is not implemented.')
            contract = lambda x: (2 - (1 / torch.norm(x, p=self.order, dim=-1, keepdim=True))) * (
                x / torch.norm(x, p=self.order, dim=-1, keepdim=True)
            )
            jc_means = vmap(jacrev(contract))(positions.mean.view(-1, positions.mean.shape[-1]))
            jc_means = jc_means.view(list(positions.mean.shape) + [positions.mean.shape[-1]])

            # Only update covariances on positions outside the unit sphere
            mag = positions.mean.norm(dim=-1)
            mask = mag >= 1
            cov = positions.cov.clone()
            cov[mask] = jc_means[mask] @ positions.cov[mask] @ torch.transpose(jc_means[mask], -2, -1)

            return Gaussians(mean=means, cov=cov)

        return contract(positions)
//  [python]
def forward(self, positions):
        def _contract(x):
            mag = torch.linalg.norm(x, ord=self.order, dim=-1)
            mask = mag >= 1
            x_new = x.clone()
            x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])

            return x_new

        def contract(x):
            x = x * self.scale_factor
            if self.aabb is None:
                mag = torch.linalg.norm(x, ord=self.order, dim=-1)
                mask = mag >= 1
            else:
                in_aabb_mask = ((x >= self.aabb[0]) & (x <= self.aabb[1])).all(-1)  # whether including = ?
                mask = ~in_aabb_mask
            x_new = x.clone()
            x_new[mask] = x[mask] / self.alpha
            return _contract(x_new)

        if isinstance(positions, Gaussians):
            means = contract(positions.mean.clone())
            # FIXME: covariances are not handled for now!
            if self.scale_factor != 1.0:
                raise NotImplementedError('Scaling of gaussian covariances is not implemented.')
            contract = lambda x: (2 - (1 / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True))) * (
                x / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True)
            )
            jc_means = vmap(jacrev(contract))(positions.mean.view(-1, positions.mean.shape[-1]))
            jc_means = jc_means.view(list(positions.mean.shape) + [positions.mean.shape[-1]])

            # Only update covariances on positions outside the unit sphere
            mag = positions.mean.norm(dim=-1)
            mask = mag >= 1
            cov = positions.cov.clone()
            cov[mask] = jc_means[mask] @ positions.cov[mask] @ torch.transpose(jc_means[mask], -2, -1)

            return Gaussians(mean=means, cov=cov)

        return contract(positions)
//  [python]
def print_params_min_max_norm(optimizer, iteration):
    """Print min, max, and norm of all parameters."""
    index = 0
    rank = torch.distributed.get_rank()
    string = "iteration, rank, index, tensor-model-parallel, min, max, norm\n"
    optimizer_ = optimizer.optimizer
    for param_group in optimizer_.param_groups:
        for param in param_group["params"]:
            index += 1
            min_ = param.data.min()
            max_ = param.data.max()
            norm = param.data.norm()
            string += "{:7d}, {:4d}, {:4d}, {:2d}, ".format(
                iteration, rank, index, int(param.tensor_model_parallel)
            )
            string += "{:.6E}, {:.6E}, {:.6E}\n".format(min_, max_, norm)
    print(string, flush=True)
//  [python]
def print_params_min_max_norm(optimizer, iteration):
    """Print min, max, and norm of all parameters."""
    index = 0
    rank = torch.distributed.get_rank()
    string = "iteration, rank, index, tensor-model-parallel, min, max, norm\n"
    optimizer_ = optimizer.optimizer
    for param_group in optimizer_.param_groups:
        for param in param_group["params"]:
            index += 1
            min_ = param.data.min()
            max_ = param.data.max()
            norm = torch.linalg.norm(param.data)
            string += "{:7d}, {:4d}, {:4d}, {:2d}, ".format(
                iteration, rank, index, int(param.tensor_model_parallel)
            )
            string += "{:.6E}, {:.6E}, {:.6E}\n".format(min_, max_, norm)
    print(string, flush=True)
//  [python]
def update_preconditioners(self, grad: Tensor) -> None:

        # iterate over all dimensions
        for k, dim in enumerate(self.dims):
            preconditioner = self.preconditioners[k]
            preconditioner_type = self.preconditioner_types[k]

            # update preconditioners (diagonal case)
            if preconditioner_type == PreconditionerType.DIAGONAL:

                # Adagrad accumulation
                if self.beta2 == 1.0:
                    preconditioner.add_(
                        grad.transpose(0, k).contiguous().view(dim, -1).norm(dim=1).pow(2)
                    )

                # Exponential moving average
                else:
                    preconditioner.mul_(self.beta2).add_(
                        grad.transpose(0, k).contiguous().view(dim, -1).norm(dim=1).pow(2),
                        alpha=1 - self.beta2,
                    )

            # update preconditioners (full-matrix case)
            else:

                # Adagrad accumulation
                if self.beta2 == 1.0:
                    contract_idx = [*range(k)] + [*range(k + 1, self.order)]
                    preconditioner.add_(
                        torch.tensordot(
                            grad,
                            grad,
                            dims=(contract_idx, contract_idx),
                        ).to(self.dtype)
                    )

                # Exponential moving average
                else:
                    contract_idx = [*range(k)] + [*range(k + 1, self.order)]
                    preconditioner.mul_(self.beta2).add_(
                        torch.tensordot(
                            grad,
                            grad,
                            dims=(contract_idx, contract_idx),
                        ).to(self.dtype),
                        alpha=1 - self.beta2,
                    )

        # update grafting method
        if self.grafting_type != GraftingType.NONE:
            self.grafting.update_preconditioners(grad)

        self.num_updates += 1
        if self.use_bias_correction and self.beta2 < 1.0:
            self.bias_correction2 = 1.0 - self.beta2**self.num_updates
//  [python]
def update_preconditioners(self, grad: Tensor) -> None:

        # iterate over all dimensions
        for k, dim in enumerate(self.dims):
            preconditioner = self.preconditioners[k]
            preconditioner_type = self.preconditioner_types[k]

            # update preconditioners (diagonal case)
            if preconditioner_type == PreconditionerType.DIAGONAL:

                # Adagrad accumulation
                if self.beta2 == 1.0:
                    preconditioner.add_(
                        torch.linalg.norm(
                            grad.transpose(0, k).contiguous().view(dim, -1),
                            dim=1,
                        ).pow(2)
                    )

                # Exponential moving average
                else:
                    preconditioner.mul_(self.beta2).add_(
                        torch.linalg.norm(
                            grad.transpose(0, k).contiguous().view(dim, -1),
                            dim=1,
                        ).pow(2),
                        alpha=1 - self.beta2,
                    )

            # update preconditioners (full-matrix case)
            else:

                # Adagrad accumulation
                if self.beta2 == 1.0:
                    contract_idx = [*range(k)] + [*range(k + 1, self.order)]
                    preconditioner.add_(
                        torch.tensordot(
                            grad,
                            grad,
                            dims=(contract_idx, contract_idx),
                        ).to(self.dtype)
                    )

                # Exponential moving average
                else:
                    contract_idx = [*range(k)] + [*range(k + 1, self.order)]
                    preconditioner.mul_(self.beta2).add_(
                        torch.tensordot(
                            grad,
                            grad,
                            dims=(contract_idx, contract_idx),
                        ).to(self.dtype),
                        alpha=1 - self.beta2,
                    )

        # update grafting method
        if self.grafting_type != GraftingType.NONE:
            self.grafting.update_preconditioners(grad)

        self.num_updates += 1
        if self.use_bias_correction and self.beta2 < 1.0:
            self.bias_correction2 = 1.0 - self.beta2**self.num_updates
//  [python]
def forward(self, rgbs: Tensor, densities: Tensor, points_3d: Tensor) -> Tensor:
        r"""Renders 3D regularly sampled points along rays.

        Args:
            rgbs: RGB values of points along rays :math:`(*, N, 3)`
            densities: Volume densities of points along rays :math:`(*, N)`
            points_3d: 3D points along rays :math:`(*, N, 3)`

        Returns:
            Rendered RGB values for each ray :math:`(*, 3)`
        """
        KORNIA_CHECK_SHAPE(rgbs, ["*", "N", "3"])
        KORNIA_CHECK_SHAPE(densities, ["*", "N"])
        KORNIA_CHECK_SHAPE(points_3d, ["*", "N", "3"])

        num_ray_points: int = points_3d.shape[-2]

        points_3d = points_3d.reshape(-1, num_ray_points, 3)  # (*, N, 3)

        delta_3d = points_3d[0, 1, :] - points_3d[0, 0, :]  # (*, 3)
        delta = torch.norm(delta_3d, dim=-1)

        alpha = 1 - torch.exp(-1.0 * densities * delta)  # (*, N)

        return self._render(alpha, rgbs)
//  [python]
def forward(self, rgbs: Tensor, densities: Tensor, points_3d: Tensor) -> Tensor:
        r"""Renders 3D regularly sampled points along rays.

        Args:
            rgbs: RGB values of points along rays :math:`(*, N, 3)`
            densities: Volume densities of points along rays :math:`(*, N)`
            points_3d: 3D points along rays :math:`(*, N, 3)`

        Returns:
            Rendered RGB values for each ray :math:`(*, 3)`
        """
        KORNIA_CHECK_SHAPE(rgbs, ["*", "N", "3"])
        KORNIA_CHECK_SHAPE(densities, ["*", "N"])
        KORNIA_CHECK_SHAPE(points_3d, ["*", "N", "3"])

        num_ray_points: int = points_3d.shape[-2]

        points_3d = points_3d.reshape(-1, num_ray_points, 3)  # (*, N, 3)

        delta_3d = points_3d[0, 1, :] - points_3d[0, 0, :]  # (*, 3)
        delta = torch.linalg.norm(delta_3d, dim=-1)

        alpha = 1 - torch.exp(-1.0 * densities * delta)  # (*, N)

        return self._render(alpha, rgbs)
//  [python]
def lookAt(eye, at, up):
    a = eye - at
    w = a / t.norm(a)
    u = torch.cross(up, w)
    u = u / t.norm(u)
    v = torch.cross(w, u)
    translate = torch.tensor([[1, 0, 0, -eye[0]], 
                              [0, 1, 0, -eye[1]], 
                              [0, 0, 1, -eye[2]], 
                              [0, 0, 0, 1]], dtype=eye.dtype, device=eye.device)
    rotate = torch.tensor([[u[0], u[1], u[2], 0], 
                           [v[0], v[1], v[2], 0], 
                           [w[0], w[1], w[2], 0], 
                           [0, 0, 0, 1]], dtype=eye.dtype, device=eye.device)
    return rotate @ translate
//  [python]
def lookAt(eye, at, up):
    a = eye - at
    w = a / torch.linalg.norm(a)
    u = torch.cross(up, w)
    u = u / torch.linalg.norm(u)
    v = torch.cross(w, u)
    translate = torch.tensor([[1, 0, 0, -eye[0]], 
                              [0, 1, 0, -eye[1]], 
                              [0, 0, 1, -eye[2]], 
                              [0, 0, 0, 1]], dtype=eye.dtype, device=eye.device)
    rotate = torch.tensor([[u[0], u[1], u[2], 0], 
                           [v[0], v[1], v[2], 0], 
                           [w[0], w[1], w[2], 0], 
                           [0, 0, 0, 1]], dtype=eye.dtype, device=eye.device)
    return rotate @ translate
//  [python]
def forward(self, input, target):
        """Forward method for loss function.

        :param torch.Tensor input: Input tensor from real data.
        :param torch.Tensor target: Model tensor output.
        :return: Loss evaluation.
        :rtype: torch.Tensor
        """
        loss = torch.norm((input - target), p=self.p, dim=-1)
        if self.relative:
            loss = loss / torch.norm(input, p=self.p, dim=-1)
        return self._reduction(loss)
//  [python]
def forward(self, input, target):
        """Forward method for loss function.

        :param torch.Tensor input: Input tensor from real data.
        :param torch.Tensor target: Model tensor output.
        :return: Loss evaluation.
        :rtype: torch.Tensor
        """
        loss = torch.linalg.norm((input - target), ord=self.p, dim=-1)
        if self.relative:
            loss = loss / torch.linalg.norm(input, ord=self.p, dim=-1)
        return self._reduction(loss)
//  [python]
def cosine_sample(N, size=None):
    # construct local frame
    N = N/t.norm(N)

    dx0 = torch.tensor([0, N[2], -N[1]], dtype=N.dtype, device=N.device)
    dx1 = torch.tensor([-N[2], 0, N[0]], dtype=N.dtype, device=N.device)

    dx = torch.where(dot(dx0, dx0) > dot(dx1, dx1), dx0, dx1)
    #dx = dx0 if np.dot(dx0,dx0) > np.dot(dx1,dx1) else dx1
    dx = dx / t.norm(dx)
    dy = torch.cross(N,dx)
    dy = dy / t.norm(dy)

    # cosine sampling in local frame
    if size is None:
        phi = 2.0 * np.pi * np.random.uniform()
        s = np.random.uniform()
    else:
        phi = 2.0 * np.pi * torch.rand(*size, 1, dtype=N.dtype, device=N.device)
        s = torch.rand(*size, 1, dtype=N.dtype, device=N.device)
    costheta = np.sqrt(s)
    sintheta = np.sqrt(1.0 - s)

    # cartesian vector in local space
    x = np.cos(phi)*sintheta
    y = np.sin(phi)*sintheta
    z = costheta

    # local to world
    return dx*x + dy*y + N*z
//  [python]
def cosine_sample(N, size=None):
    # construct local frame
    N = N/torch.linalg.norm(N)

    dx0 = torch.tensor([0, N[2], -N[1]], dtype=N.dtype, device=N.device)
    dx1 = torch.tensor([-N[2], 0, N[0]], dtype=N.dtype, device=N.device)

    dx = torch.where(dot(dx0, dx0) > dot(dx1, dx1), dx0, dx1)
    #dx = dx0 if np.dot(dx0,dx0) > np.dot(dx1,dx1) else dx1
    dx = dx / torch.linalg.norm(dx)
    dy = torch.cross(N,dx)
    dy = dy / torch.linalg.norm(dy)

    # cosine sampling in local frame
    if size is None:
        phi = 2.0 * np.pi * np.random.uniform()
        s = np.random.uniform()
    else:
        phi = 2.0 * np.pi * torch.rand(*size, 1, dtype=N.dtype, device=N.device)
        s = torch.rand(*size, 1, dtype=N.dtype, device=N.device)
    costheta = np.sqrt(s)
    sintheta = np.sqrt(1.0 - s)

    # cartesian vector in local space
    x = np.cos(phi)*sintheta
    y = np.sin(phi)*sintheta
    z = costheta

    # local to world
    return dx*x + dy*y + N*z
//  [python]
def distance_to_reference_trajectory(pred_centroid: torch.Tensor, ref_traj: torch.Tensor) -> torch.Tensor:
    """ Computes the distance from the predicted centroid to the closest waypoint in the reference trajectory.

    :param pred_centroid: predicted centroid tensor, size: [batch_size, 2]
    :type pred_centroid: torch.Tensor, float
    :param ref_traj: reference trajectory tensor, size: [batch_size, num_timestamps, 2]
    :type ref_traj: torch.Tensor, float
    :return: closest distance between the predicted centroid and the reference trajectory, size: [batch_size,]
    :rtype: torch.Tensor, float
    """
    # [batch_size, 2]
    assert pred_centroid.dim() == 2
    # [batch_size, num_timestamps, 2]
    assert ref_traj.dim() == 3

    # [batch_size,]
    euclidean_distance = torch.norm(pred_centroid.unsqueeze(1) - ref_traj, p=2, dim=-1)
    return torch.amin(euclidean_distance, dim=1)
//  [python]
def distance_to_reference_trajectory(pred_centroid: torch.Tensor, ref_traj: torch.Tensor) -> torch.Tensor:
    """ Computes the distance from the predicted centroid to the closest waypoint in the reference trajectory.

    :param pred_centroid: predicted centroid tensor, size: [batch_size, 2]
    :type pred_centroid: torch.Tensor, float
    :param ref_traj: reference trajectory tensor, size: [batch_size, num_timestamps, 2]
    :type ref_traj: torch.Tensor, float
    :return: closest distance between the predicted centroid and the reference trajectory, size: [batch_size,]
    :rtype: torch.Tensor, float
    """
    # [batch_size, 2]
    assert pred_centroid.dim() == 2
    # [batch_size, num_timestamps, 2]
    assert ref_traj.dim() == 3

    # [batch_size,]
    euclidean_distance = torch.linalg.norm(pred_centroid.unsqueeze(1) - ref_traj, ord=2, dim=-1)
    return torch.amin(euclidean_distance, dim=1)
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        q = torch.stack([(pixels_x/self.W-0.5)*2, (pixels_y/self.H-0.5)*2, torch.zeros_like(pixels_y)], dim=-1) # W, H, 3

        v = torch.stack([torch.zeros_like(pixels_y), torch.zeros_like(pixels_y), torch.ones_like(pixels_y)], dim=-1) # W, H, 3

        # orthogonal projection
        rays_v = v / torch.norm(v, p=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3].cuda(), rays_v[:, :, :, None].cuda()).squeeze()  # W, H, 3
        
        rays_o = torch.matmul(self.pose_all[img_idx, None, None, :3, :3].cuda(), q[:, :, :, None].cuda()).squeeze()  # W, H, 3
        rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape).cuda() + rays_o  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        q = torch.stack([(pixels_x/self.W-0.5)*2, (pixels_y/self.H-0.5)*2, torch.zeros_like(pixels_y)], dim=-1) # W, H, 3

        v = torch.stack([torch.zeros_like(pixels_y), torch.zeros_like(pixels_y), torch.ones_like(pixels_y)], dim=-1) # W, H, 3

        # orthogonal projection
        rays_v = v / torch.linalg.norm(v, ord=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3].cuda(), rays_v[:, :, :, None].cuda()).squeeze()  # W, H, 3
        
        rays_o = torch.matmul(self.pose_all[img_idx, None, None, :3, :3].cuda(), q[:, :, :, None].cuda()).squeeze()  # W, H, 3
        rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape).cuda() + rays_o  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def calculate_csm_inati_iter_prototype(im, niter=5):

    im = im[:,None]
    ncha = len(im)
    D_sum = im.sum(axis=(1, 2, 3))
    v = 1/torch.norm(D_sum)
    D_sum *= v
    R = 0
    for cha in range(ncha):
        R += torch.conj(D_sum[cha]) * im[cha,...]

    eps = torch.finfo(im.real.dtype).eps * torch.abs(im).mean()
    eps.to(im)

    for itr in range(niter):
        
        R = torch.conj(R)

        coil_map = im * R[None, ...]

        coil_map_conv = coil_map

        D = coil_map_conv * torch.conj(coil_map_conv)
        R = D.sum(axis=0)
        R = torch.sqrt(R) + eps
        R = 1/R
        coil_map = coil_map_conv * R[np.newaxis, ...]
        D = im * torch.conj(coil_map)
        R = D.sum(axis=0)
        D = coil_map * R[None, ...]

        D_sum = D.sum(axis=(1, 2, 3))

        v = 1/torch.norm(D_sum)
        D_sum *= v

        imT = 0
        for cha in range(ncha):
            imT += torch.conj(D_sum[cha]) * coil_map[cha, ...]

        magT = torch.abs(imT) + eps
        imT /= magT
        R = R * imT
        imT = torch.conj(imT)
        coil_map = coil_map * imT[None, ...]

    coil_map = coil_map[:, 0, :, :]

    return coil_map
//  [python]
def calculate_csm_inati_iter_prototype(im, niter=5):

    im = im[:,None]
    ncha = len(im)
    D_sum = im.sum(axis=(1, 2, 3))
    v = 1/torch.linalg.norm(D_sum)
    D_sum *= v
    R = 0
    for cha in range(ncha):
        R += torch.conj(D_sum[cha]) * im[cha,...]

    eps = torch.finfo(im.real.dtype).eps * torch.abs(im).mean()
    eps.to(im)

    for itr in range(niter):
        
        R = torch.conj(R)

        coil_map = im * R[None, ...]

        coil_map_conv = coil_map

        D = coil_map_conv * torch.conj(coil_map_conv)
        R = D.sum(axis=0)
        R = torch.sqrt(R) + eps
        R = 1/R
        coil_map = coil_map_conv * R[np.newaxis, ...]
        D = im * torch.conj(coil_map)
        R = D.sum(axis=0)
        D = coil_map * R[None, ...]

        D_sum = D.sum(axis=(1, 2, 3))

        v = 1/torch.linalg.norm(D_sum)
        D_sum *= v

        imT = 0
        for cha in range(ncha):
            imT += torch.conj(D_sum[cha]) * coil_map[cha, ...]

        magT = torch.abs(imT) + eps
        imT /= magT
        R = R * imT
        imT = torch.conj(imT)
        coil_map = coil_map * imT[None, ...]

    coil_map = coil_map[:, 0, :, :]

    return coil_map
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size]).cpu() 
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size]).cpu() 
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        normal = self.normals_world[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        
        q = torch.stack([(pixels_x / self.W-0.5)*2, (pixels_y / self.H-0.5)*2, torch.zeros_like(pixels_y)], dim=-1).float()  # batch_size, 3
        v = torch.stack([torch.zeros_like(pixels_y), torch.zeros_like(pixels_y), torch.ones_like(pixels_y)], dim=-1).float()

        # q = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float() # bsz, 3
        # q = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], q[:, :, None]).squeeze()  # bsz, 3
        # q[:, 2] = 0
        
        rays_v = v / torch.norm(v, p=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        
        rays_o = torch.matmul(self.pose_all[img_idx, None, :3, :3], q[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) + rays_o # batch_size, 3
        
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, None], normal], dim=-1).cuda()
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size]).cpu() 
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size]).cpu() 
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        normal = self.normals_world[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        
        q = torch.stack([(pixels_x / self.W-0.5)*2, (pixels_y / self.H-0.5)*2, torch.zeros_like(pixels_y)], dim=-1).float()  # batch_size, 3
        v = torch.stack([torch.zeros_like(pixels_y), torch.zeros_like(pixels_y), torch.ones_like(pixels_y)], dim=-1).float()

        # q = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float() # bsz, 3
        # q = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], q[:, :, None]).squeeze()  # bsz, 3
        # q[:, 2] = 0
        
        rays_v = v / torch.linalg.norm(v, ord=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        
        rays_o = torch.matmul(self.pose_all[img_idx, None, :3, :3], q[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) + rays_o # batch_size, 3
        
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, None], normal], dim=-1).cuda()
//  [python]
def _encode_prompt(self, prompt):
        # get prompt text embeddings
        text_inputs = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids

        if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
            removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :])
            logger.warning(
                "The following part of your input was truncated because CLIP can only handle sequences up to"
                f" {self.tokenizer.model_max_length} tokens: {removed_text}"
            )
            text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length]
        prompt_embeds = self.clip.get_text_features(text_input_ids.to(self.device))
        prompt_embeds = prompt_embeds / torch.norm(prompt_embeds, dim=-1, keepdim=True)
        prompt_embeds = prompt_embeds[:, None, :]
        return prompt_embeds
//  [python]
def _encode_prompt(self, prompt):
        # get prompt text embeddings
        text_inputs = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids

        if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
            removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :])
            logger.warning(
                "The following part of your input was truncated because CLIP can only handle sequences up to"
                f" {self.tokenizer.model_max_length} tokens: {removed_text}"
            )
            text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length]
        prompt_embeds = self.clip.get_text_features(text_input_ids.to(self.device))
        prompt_embeds = prompt_embeds / torch.linalg.norm(prompt_embeds, dim=-1, keepdim=True)
        prompt_embeds = prompt_embeds[:, None, :]
        return prompt_embeds
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape

        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        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_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        # ----------------------------------------------------------------------------------------------------------
        # Use min value of [ cos, prev_cos ]
        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
        # robust when meeting situations like below:
        #
        # SDF
        # ^
        # |\          -----x----...
        # | \        /
        # |  x      x
        # |---\----/-------------> 0 level
        # |    \  /
        # |     \/
        # |
        # ----------------------------------------------------------------------------------------------------------
        prev_cos_val = torch.cat([torch.zeros([batch_size, 1], device=sdf.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], device=alpha.device), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()

        return z_samples