def tricubic_sample_3d(volume, optical):
    """
    tricubic sampling; can guarantee continuous gradient  (interpolation border)
    :param volume: [B, C, ID, IH, IW]
    :param optical: [B, D, H, W, 3]
    :param sample_num:
    :return:
    """

    @torch.no_grad()
    def get_shifts(x):
        x1 = -1 * (1 + x - torch.floor(x))
        x2 = -1 * (x - torch.floor(x))
        x3 = torch.floor(x) + 1 - x
        x4 = torch.floor(x) + 2 - x

        return torch.stack([x1, x2, x3, x4], dim=-1)  # (B,d,h,w,4)

    N, C, ID, IH, IW = volume.shape
    _, D, H, W, _ = optical.shape

    device = volume.device

    ix = optical[..., 0]
    iy = optical[..., 1]
    iz = optical[..., 2]

    ix = ((ix + 1) / 2) * (IW - 1)  # (B,d,h,w)
    iy = ((iy + 1) / 2) * (IH - 1)
    iz = ((iz + 1) / 2) * (ID - 1)

    ix = ix.view(-1)
    iy = iy.view(-1)
    iz = iz.view(-1)

    with torch.no_grad():
        shifts_x = get_shifts(ix).view(-1, 4)  # (B*d*h*w,4)
        shifts_y = get_shifts(iy).view(-1, 4)
        shifts_z = get_shifts(iz).view(-1, 4)

        perm_weights = torch.ones([N * D * H * W, 4 * 4 * 4]).long().to(device)
        perm = torch.cumsum(perm_weights, dim=-1) - 1  # (B*d*h*w,64)

        perm_z = perm // 16  # [N*D*H*W, num]
        perm_y = (perm - perm_z * 16) // 4
        perm_x = (perm - perm_z * 16 - perm_y * 4)

        shifts_x = torch.gather(shifts_x, 1, perm_x)  # [N*D*H*W, num]
        shifts_y = torch.gather(shifts_y, 1, perm_y)
        shifts_z = torch.gather(shifts_z, 1, perm_z)

        ix_target = (ix[:, None] + shifts_x).long()  # [N*D*H*W, num]
        iy_target = (iy[:, None] + shifts_y).long()
        iz_target = (iz[:, None] + shifts_z).long()

        torch.clamp(ix_target, 0, IW - 1, out=ix_target)
        torch.clamp(iy_target, 0, IH - 1, out=iy_target)
        torch.clamp(iz_target, 0, ID - 1, out=iz_target)

    local_dist_x = ix - ix_target[:, 1]  # ! attention here is [:, 1]
    local_dist_y = iy - iy_target[:, 1 + 4]
    local_dist_z = iz - iz_target[:, 1 + 16]

    local_dist_x = local_dist_x.view(N, 1, D * H * W).repeat(1, C, 1).view(-1)
    local_dist_y = local_dist_y.view(N, 1, D * H * W).repeat(1, C, 1).view(-1)
    local_dist_z = local_dist_z.view(N, 1, D * H * W).repeat(1, C, 1).view(-1)

    # ! attention: IW is correct
    idx_target = iz_target * ID ** 2 + iy_target * IW + ix_target  # [N*D*H*W, num]

    volume = volume.view(N, C, ID * IH * IW)

    out = torch.gather(volume, 2,
                       idx_target.view(N, 1, D * H * W * 64).repeat(1, C, 1))
    out = out.view(N * C * D * H * W, 4, 4, 4)

    # - tricubic_interpolate() is a bit faster than tricubic_interpolate_batch()
    final = tricubic_interpolate(out, local_dist_x, local_dist_y, local_dist_z).view(N, C, D, H, W)  # [N,C,D,H,W]

    return final

    tricubic sampling; can guarantee continuous gradient  (interpolation border)
    :param volume: [B, C, ID, IH, IW]
    :param optical: [B, D, H, W, 3]
    :param sample_num:
    :return:
    
threestudio/models/geometry/grid_sampler.py