def test_multi_parabola_1d(self):
"""Tests specular reflection for a multiple parabolic reflectors with broadcasting.
Function: reflector_problem.raytracing.reflections_laws.
specular_reflection(incident_rays, normals)
"""
N = 128
incident_rays_angles = torch.linspace(0, pi, N)
incident_rays = to_unit_vector(incident_rays_angles)
N_parabola = 10
directions = to_unit_vector(
torch.linspace(9 * pi / 8, 15 * pi / 8, N_parabola))
C = 1
parabola_gradients = C / (1 - torch.sum(incident_rays.view(1, N, 2) * directions.view(N_parabola, 1, 2), dim=-1, keepdim=True))**2 *\
(gradient_to_normal(incident_rays.view(1, N, 2) - directions.view(N_parabola, 1, 2)))
parabola_gradients_normalized = normalize_vector(parabola_gradients)
parabola_normals = gradient_to_normal(parabola_gradients_normalized)
reflected_rays = specular_reflection(incident_rays, parabola_normals)
self.assertEqual(reflected_rays.shape, (N_parabola, N, 2))
self.assertEqual(incident_rays.dtype, reflected_rays.dtype)
self.assertTrue(
torch.allclose(reflected_rays, directions.view(N_parabola, 1, 2)))
def test_to_unit_vector_norm(self):
"""Verify the vectors returned by the to_unit_vector() function are unit vectors.
Function: reflector_problem.raytracing.utils.to_unit_vector(angles)
"""
angles = torch.linspace(0, 2 * pi, 1000)
unit_vectors = utils.to_unit_vector(angles)
self.assertTrue(
torch.allclose(unit_vectors.norm(dim=-1), torch.Tensor([1.])))
self.assertEqual(unit_vectors.shape, (1000, 2))
self.assertEqual(angles.dtype, unit_vectors.dtype)
def test_to_unit_vector_norm_batch(self):
"""Verify the vectors returned by the to_unit_vector() function are unit vectors for input with multiple batch dimensions.
Function: reflector_problem.raytracing.utils.to_unit_vector(angles)
"""
B1, B2, B3, B4 = 4, 5, 6, 7 # Batch dimensions
angles = torch.linspace(0, 2 * pi,
B1 * B2 * B3 * B4).view(B1, B2, B3, B4)
unit_vectors = utils.to_unit_vector(angles)
self.assertTrue(
torch.allclose(unit_vectors.norm(dim=-1), torch.Tensor([1.])))
self.assertEqual(unit_vectors.shape, (B1, B2, B3, B4, 2))
self.assertEqual(angles.dtype, unit_vectors.dtype)
def test_parabola_1d(self):
"""Tests specular reflection for a parabolic reflector.
Function: reflector_problem.raytracing.reflections_laws.
specular_reflection(incident_rays, normals)
"""
N = 128
incident_rays_angles = torch.linspace(0, pi, N)
incident_rays = to_unit_vector(incident_rays_angles)
direction = to_unit_vector(torch.Tensor([5 * pi / 3]))
C = 1
parabola_gradients = C / (1 - torch.sum(incident_rays * direction.view(1, 2), dim=-1, keepdim=True))**2 *\
(gradient_to_normal(incident_rays - direction.view(1, 2)))
parabola_gradients_normalized = normalize_vector(parabola_gradients)
parabola_normals = gradient_to_normal(parabola_gradients_normalized)
reflected_rays = specular_reflection(incident_rays, parabola_normals)
self.assertEqual(reflected_rays.shape, (N, 2))
self.assertEqual(incident_rays.dtype, reflected_rays.dtype)
self.assertTrue(torch.allclose(reflected_rays, direction.view(-1, 2)))
def test_naive_rw(self):
"""Tests the NaiveRayWeighter class
Class: reflector_problem.raytracing.ray_weighter.NaiveRayWeighter
"""
gaussian = torch.distributions.Normal(loc=3 * pi / 2, scale=pi / 10)
source_description = SourceDescription(
lambda x: gaussian.log_prob(x).exp(), gaussian.cdf)
ray_weighter = NaiveRayWeighter(source_description)
N = 2048
angles = torch.linspace(pi, 2 * pi, N)
weights = ray_weighter.compute_weights(utils.to_unit_vector(angles))
self.assertEqual(weights.shape, (N, ))
self.assertEqual(weights.dtype, angles.dtype)
self.assertTrue(torch.allclose(torch.sum(weights), torch.Tensor([1.])))
def test_flat_broadcasting_1d(self):
"""Tests specular reflection for a flat reflector
Function: reflector_problem.raytracing.reflections_laws.
specular_reflection(incident_rays, normals)
"""
N = 128
incident_rays_angles = torch.linspace(0, pi, N)
incident_rays = to_unit_vector(incident_rays_angles)
flat_normals = torch.Tensor([0., 1.]).view(1, 2)
reflected_rays = specular_reflection(incident_rays, flat_normals)
self.assertEqual(incident_rays.shape, reflected_rays.shape)
self.assertEqual(incident_rays.dtype, reflected_rays.dtype)
self.assertTrue(
torch.allclose(incident_rays[..., 0], reflected_rays[..., 0]))
self.assertTrue(
torch.allclose(incident_rays[..., 1], -reflected_rays[..., 1]))
def test_flat_weirdshape(self):
"""Tests specular reflection for a flat reflector, with weird shapes (to check batch and broadcasting)
Function: reflector_problem.raytracing.reflections_laws.
specular_reflection(incident_rays, normals)
"""
B1, B2, B3, B4 = 5, 2, 3, 7
incident_rays_angles = torch.linspace(0, pi, B1 * 1 * B3 * 1).view(
B1, 1, B3, 1)
incident_rays = to_unit_vector(incident_rays_angles)
flat_normals = torch.Tensor([0., 1.]).view(1, 1, 1, 1,
2).repeat(1, B2, 1, B4, 1)
reflected_rays = specular_reflection(incident_rays, flat_normals)
self.assertEqual(reflected_rays.shape, (B1, B2, B3, B4, 2))
self.assertEqual(incident_rays.dtype, reflected_rays.dtype)
self.assertTrue(
torch.allclose(incident_rays[..., 0], reflected_rays[..., 0]))
self.assertTrue(
torch.allclose(incident_rays[..., 1], -reflected_rays[..., 1]))
def test_compute_reflector(self):
"""Tests the compute_reflector function
Function: reflector_problem.raytracing.reflector_utils.compute_reflector(unit_vector_support, potential)
"""
N_discr = 128
angles = torch.linspace(0, 2 * pi, N_discr)
N_heights = 16
heights = torch.linspace(1, 5, N_heights)
unit_vectors = utils.to_unit_vector(angles).view(-1, 1, 2)
potential = torch.ones((N_discr, 1, 1), dtype=angles.dtype)
potentials = potential * torch.log(heights.view(1, -1, 1))
reflector_points = reflector_utils.compute_reflector(
unit_vectors, potentials)
self.assertEqual(reflector_points.shape, (N_discr, N_heights, 2))
self.assertEqual(reflector_points.dtype, unit_vectors.dtype)
self.assertTrue(
torch.allclose(reflector_points.norm(dim=-1),
heights.type_as(unit_vectors)))