docstr: Check the rays_directions_world output of raysamplers.
docstr_tok: ["Check","the","rays_directions_world","output","of","raysamplers","."]
code: def _check_raysampler_ray_directions(
      self, cameras, raysampler, ray_bundle
  ):
      """
      Check the rays_directions_world output of raysamplers.
      """
  
      batch_size = cameras.R.shape[0]
      n_pts_per_ray = ray_bundle.lengths.shape[-1]
      spatial_size = ray_bundle.xys.shape[1:-1]
      n_rays_per_image = spatial_size.numel()
  
      # obtain the ray points in world coords
      rays_points_world = cameras.unproject_points(
          torch.cat(
              (
                  ray_bundle.xys.view(
                      batch_size, n_rays_per_image, 1, 2
                  ).expand(
                      batch_size, n_rays_per_image, n_pts_per_ray, 2
                  ),
                  ray_bundle.lengths.view(
                      batch_size, n_rays_per_image, n_pts_per_ray, 1
                  ),
              ),
              dim=-1,
          ).view(batch_size, -1, 3)
      ).view(batch_size, -1, n_pts_per_ray, 3)
  
      # reshape to common testing size
      rays_directions_world_normed = torch.nn.functional.normalize(
          ray_bundle.directions.view(batch_size, -1, 3), dim=-1
      )
  
      # check that the l2-normed difference of all consecutive planes
      # of points in world coords matches ray_directions_world
      rays_directions_world_ = torch.nn.functional.normalize(
          rays_points_world[:, :, -1:] - rays_points_world[:, :, :-1],
          dim=-1,
      )
      self.assertClose(
          rays_directions_world_normed[:, :, None].expand_as(
              rays_directions_world_
          ),
          rays_directions_world_,
          atol=1e-4,
      )
  
      # check the ray directions rotated using camera rotation matrix
      # match the ray directions of a camera with trivial extrinsics
      cameras_trivial_extrinsic = cameras.clone()
      cameras_trivial_extrinsic.R = eyes(
          N=batch_size,
          dim=3,
          dtype=cameras.R.dtype,
          device=cameras.device,
      )
      cameras_trivial_extrinsic.T = torch.zeros_like(cameras.T)
  
      # make sure we get the same random rays in case we call the
      # MonteCarloRaysampler twice below
      with torch.random.fork_rng(
          devices=range(torch.cuda.device_count())
      ):
          torch.random.manual_seed(42)
          ray_bundle_world_fix_seed = raysampler(cameras=cameras)
          torch.random.manual_seed(42)
          ray_bundle_camera_fix_seed = raysampler(
              cameras=cameras_trivial_extrinsic
          )
  
      rays_directions_camera_fix_seed_ = Rotate(
          cameras.R, device=cameras.R.device
      ).transform_points(
          ray_bundle_world_fix_seed.directions.view(batch_size, -1, 3)
      )
  
      self.assertClose(
          rays_directions_camera_fix_seed_,
          ray_bundle_camera_fix_seed.directions.view(batch_size, -1, 3),
          atol=1e-5,
      )
  
imports: import unittest
  from typing import Callable
  import torch
  from pytorch3d.common.compat import meshgrid_ij
  from pytorch3d.ops import eyes
  from pytorch3d.renderer import MonteCarloRaysampler, MultinomialRaysampler, NDCGridRaysampler, NDCMultinomialRaysampler
  from pytorch3d.renderer.cameras import FoVOrthographicCameras, FoVPerspectiveCameras, OrthographicCameras, PerspectiveCameras
  from pytorch3d.renderer.implicit.raysampling import _jiggle_within_stratas, _safe_multinomial
  from pytorch3d.renderer.implicit.utils import ray_bundle_to_ray_points, ray_bundle_variables_to_ray_points
  from pytorch3d.transforms import Rotate
  from .common_testing import TestCaseMixin
  from .test_cameras import init_random_cameras
label_desc: Use image node_id 8 for calling the TestRaysampling obj's underlying member method code with example usage: obj._check_raysampler_ray_directions(cameras, raysampler, ray_bundle) without return types