Add unit tests for differentiable disk & rectangle

mitsuba-renderer · Nov 11, 2022 · b5d8c5d · b5d8c5d
1 parent d8bb5a8
commit b5d8c5d
Show file tree

Hide file tree

Showing 4 changed files with 237 additions and 8 deletions.
diff --git a/src/render/python/shape_v.cpp b/src/render/python/shape_v.cpp
@@ -122,7 +122,10 @@ MI_PY_EXPORT(Shape) {
         .def_method(Shape, is_mesh)
         .def_method(Shape, parameters_grad_enabled)
         .def_method(Shape, primitive_count)
-        .def_method(Shape, effective_primitive_count);
+        .def_method(Shape, effective_primitive_count)
+        .def("eval_parameterization", &Shape::eval_parameterization,
+             "uv"_a, "ray_flags"_a = +RayFlags::All, "active"_a = true,
+             D(Shape, eval_parameterization));
 
     bind_shape_generic<Shape *>(shape);
 
@@ -167,10 +170,7 @@ MI_PY_EXPORT(Shape) {
              }, D(Mesh, face_indices), "index"_a, "active"_a = true)
         .def("ray_intersect_triangle", &Mesh::ray_intersect_triangle,
              "index"_a, "ray"_a, "active"_a = true,
-             D(Mesh, ray_intersect_triangle))
-        .def("eval_parameterization", &Mesh::eval_parameterization,
-             "uv"_a, "ray_flags"_a = +RayFlags::All, "active"_a = true,
-             D(Mesh, eval_parameterization));
+             D(Mesh, ray_intersect_triangle));
 
     MI_PY_REGISTER_OBJECT("register_mesh", Mesh)
 }
diff --git a/src/shapes/tests/test_disk.py b/src/shapes/tests/test_disk.py
@@ -175,3 +175,122 @@ def test06_differentiable_surface_interaction_ray_backward(variants_all_ad_rgb):
     si = pi.compute_surface_interaction(ray)
     dr.backward(si.t)
     assert dr.allclose(dr.grad(ray.o), [0, 0, -1])
+
+
+def test07_differentiable_surface_interaction_ray_forward_follow_shape(variants_all_ad_rgb):
+    shape = mi.load_dict({'type' : 'disk'})
+    params = mi.traverse(shape)
+
+    # Test 00: With DetachShape and no moving rays, the output shouldn't produce
+    #          any gradients.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.1, -2), mi.Vector3f(0, 0, 1))
+
+    theta = mi.Float(0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.scale(1 + theta)
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All | mi.RayFlags.DetachShape)
+
+    dr.forward(theta)
+
+    assert dr.allclose(dr.grad(si.t), 0.0)
+    assert dr.allclose(dr.grad(si.p), 0.0)
+    assert dr.allclose(dr.grad(si.n), 0.0)
+    assert dr.allclose(dr.grad(si.uv), 0.0)
+
+    # Test 01: When the disk is inflating, the point will not move along the
+    #          ray. The normal isn't changing but the UVs are.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.1, -2), mi.Vector3f(0, 0, 1))
+
+    theta = mi.Float(0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.scale(1 + theta)
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All)
+
+    dr.forward(theta)
+
+    r = dr.norm(mi.Point2f(0.1, 0.1))
+    d_r = -r
+    d_v = 0
+    d_uv = mi.Vector2f(d_r, d_v)
+
+    assert dr.allclose(dr.grad(si.t), 0)
+    assert dr.allclose(dr.grad(si.p), 0)
+    assert dr.allclose(dr.grad(si.n), 0)
+    assert dr.allclose(dr.grad(si.uv), d_uv)
+
+    # Test 02: With FollowShape, any intersection point with translating disk
+    #          should move according to the translation. The normal and the
+    #          UVs should be static.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.1, -2.0), mi.Vector3f(0.0, 0.0, 1.0))
+
+    theta = mi.Float(0.0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.translate([theta, 0.0, 0.0])
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All | mi.RayFlags.FollowShape)
+
+    dr.forward(theta, dr.ADFlag.ClearNone)
+
+    assert dr.allclose(dr.grad(si.p), [1.0, 0.0, 0.0])
+    assert dr.allclose(dr.grad(si.n), 0.0)
+    assert dr.allclose(dr.grad(si.uv), 0.0)
+
+    # Test 03: With FollowShape, an off-center intersection point with a
+    #          rotating disk and its normal should follow the rotation speed
+    #          along the tangent direction. The UVs should be static.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.1, -2.0), mi.Vector3f(0.0, 0.0, 1.0))
+
+    theta = mi.Float(0.0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.rotate([0, 0, 1], 90 * theta)
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All | mi.RayFlags.FollowShape)
+
+    dr.forward(theta)
+
+    assert dr.allclose(dr.grad(si.p), [-dr.pi * 0.1 / 2, dr.pi * 0.1 / 2, 0.0])
+    assert dr.allclose(dr.grad(si.n), 0.0)
+    assert dr.allclose(dr.grad(si.uv), 0.0)
+
+    # Test 04: Without FollowShape, a disk that is only rotating shouldn't
+    #          produce any gradients for the intersection point and normal, but
+    #          for the UVs.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.1, -2.0), mi.Vector3f(0.0, 0.0, 1.0))
+
+    theta = mi.Float(0.0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.rotate([0, 0, 1], 90 * theta)
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All)
+
+    dr.forward(theta)
+
+    v = dr.atan2(0.1, 0.1) * dr.inv_two_pi
+    d_r = 0
+    d_v = -2 * v
+    d_uv = mi.Vector2f(d_r, d_v)
+
+    assert dr.allclose(dr.grad(si.p), 0.0)
+    assert dr.allclose(dr.grad(si.n), 0.0)
+    assert dr.allclose(dr.grad(si.uv), d_uv)
+
+
+def test08_eval_parameterization(variants_all_ad_rgb):
+    shape = mi.load_dict({'type' : 'disk'})
+    transform = mi.Transform4f().scale(0.2).translate([0.1, 0.2, 0.3]).rotate([1.0, 0, 0], 45)
+
+    si_before = shape.eval_parameterization(mi.Point2f(0.3, 0.6))
+
+    params = mi.traverse(shape)
+    params['to_world'] = transform
+    params.update()
+
+    si_after = shape.eval_parameterization(mi.Point2f(0.3, 0.6))
+    assert dr.allclose(si_before.uv, si_after.uv)
diff --git a/src/shapes/tests/test_rectangle.py b/src/shapes/tests/test_rectangle.py
@@ -171,3 +171,116 @@ def test07_differentiable_surface_interaction_ray_backward(variants_all_ad_rgb):
     si = pi.compute_surface_interaction(ray)
     dr.backward(si.t)
     assert dr.allclose(dr.grad(ray.o), [0, 0, -1])
+
+
+def test08_differentiable_surface_interaction_ray_forward_follow_shape(variants_all_ad_rgb):
+    shape = mi.load_dict({'type' : 'rectangle'})
+    params = mi.traverse(shape)
+
+    # Test 00: With DetachShape and no moving rays, the output shouldn't produce
+    #          any gradients.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.1, -2), mi.Vector3f(0, 0, 1))
+
+    theta = mi.Float(0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.scale(1 + theta)
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All | mi.RayFlags.DetachShape)
+
+    dr.forward(theta)
+
+    assert dr.allclose(dr.grad(si.t), 0.0)
+    assert dr.allclose(dr.grad(si.p), 0.0)
+    assert dr.allclose(dr.grad(si.n), 0.0)
+    assert dr.allclose(dr.grad(si.uv), 0.0)
+
+    # Test 01: When the rectangle is stretched, the point will not move along
+    #          the ray. The normal isn't changing but the UVs are.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.2, -2), mi.Vector3f(0, 0, 1))
+
+    theta = mi.Float(0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.scale([1 + theta, 1 + 2 * theta, 0])
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All)
+
+    dr.forward(theta)
+
+    d_uv = mi.Point2f(-0.1 / 2, -0.2)
+
+    assert dr.allclose(dr.grad(si.t), 0)
+    assert dr.allclose(dr.grad(si.p), 0)
+    assert dr.allclose(dr.grad(si.n), 0)
+    assert dr.allclose(dr.grad(si.uv), d_uv)
+
+    # Test 02: With FollowShape, any intersection point with translating rectangle
+    #          should move according to the translation. The normal and the
+    #          UVs should be static.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.1, -2.0), mi.Vector3f(0.0, 0.0, 1.0))
+
+    theta = mi.Float(0.0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.translate([theta, 0.0, 0.0])
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All | mi.RayFlags.FollowShape)
+
+    dr.forward(theta, dr.ADFlag.ClearNone)
+
+    assert dr.allclose(dr.grad(si.p), [1.0, 0.0, 0.0])
+    assert dr.allclose(dr.grad(si.n), 0.0)
+    assert dr.allclose(dr.grad(si.uv), 0.0)
+
+    # Test 03: With FollowShape, an off-center intersection point with a
+    #          rotating rectangle and its normal should follow the rotation speed
+    #          along the tangent direction. The UVs should be static.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.1, -2.0), mi.Vector3f(0.0, 0.0, 1.0))
+
+    theta = mi.Float(0.0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.rotate([0, 0, 1], 90 * theta)
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All | mi.RayFlags.FollowShape)
+
+    dr.forward(theta)
+
+    assert dr.allclose(dr.grad(si.p), [-dr.pi * 0.1 / 2, dr.pi * 0.1 / 2, 0.0])
+    assert dr.allclose(dr.grad(si.n), 0.0)
+    assert dr.allclose(dr.grad(si.uv), 0.0)
+
+    # Test 04: Without FollowShape, a rectangle that is only rotating shouldn't
+    #          produce any gradients for the intersection point and normal, but
+    #          for the UVs.
+
+    ray = mi.Ray3f(mi.Vector3f(0.1, 0.1, -2.0), mi.Vector3f(0.0, 0.0, 1.0))
+
+    theta = mi.Float(0.0)
+    dr.enable_grad(theta)
+    params['to_world'] = mi.Transform4f.rotate([0, 0, 1], 90 * theta)
+    params.update()
+    si = shape.ray_intersect(ray, mi.RayFlags.All)
+
+    dr.forward(theta)
+
+    d_uv = [dr.pi * 0.1 / 4, -dr.pi * 0.1 / 4]
+
+    assert dr.allclose(dr.grad(si.p), 0.0)
+    assert dr.allclose(dr.grad(si.n), 0.0)
+    assert dr.allclose(dr.grad(si.uv), d_uv)
+
+
+def test09_eval_parameterization(variants_all_ad_rgb):
+    shape = mi.load_dict({'type' : 'rectangle'})
+    transform = mi.Transform4f().scale(0.2).translate([0.1, 0.2, 0.3]).rotate([1.0, 0, 0], 45)
+
+    si_before = shape.eval_parameterization(mi.Point2f(0.3, 0.6))
+
+    params = mi.traverse(shape)
+    params['to_world'] = transform
+    params.update()
+
+    si_after = shape.eval_parameterization(mi.Point2f(0.3, 0.6))
+    assert dr.allclose(si_before.uv, si_after.uv)
diff --git a/src/shapes/tests/test_sphere.py b/src/shapes/tests/test_sphere.py
@@ -310,9 +310,6 @@ def test08_differentiable_surface_interaction_ray_forward_follow_shape(variants_
     assert dr.allclose(dr.grad(si.n), 0.0)
     assert dr.allclose(dr.grad(si.uv), [0.0, -0.5])
 
-
-
-
 def test09_si_singularity(variants_all_rgb):
     scene = mi.load_dict({"type" : "scene", 's': { 'type': 'sphere' }})
     ray = mi.Ray3f([0, 0, -1], [0, 0, 1])