def test03_depth_packet_stairs(variant_packet_rgb): from mitsuba.core import Ray3f as Ray3fX, Properties from mitsuba.render import Scene if mitsuba.core.MTS_ENABLE_EMBREE: pytest.skip("EMBREE enabled") props = Properties("scene") props["_unnamed_0"] = create_stairs_packet(11) scene = Scene(props) mitsuba.set_variant("scalar_rgb") from mitsuba.core import Ray3f, Vector3f n = 4 inv_n = 1.0 / (n - 1) rays = Ray3fX.zero(n * n) d = [0, 0, -1] wavelengths = [] for x in range(n): for y in range(n): o = Vector3f(x * inv_n, y * inv_n, 2) o = o * 0.999 + 0.0005 rays[x * n + y] = Ray3f(o, d, 0, 100, 0.5, wavelengths) res_naive = scene.ray_intersect_naive(rays) res = scene.ray_intersect(rays) res_shadow = scene.ray_test(rays) # TODO: spot-check (here, we only check consistency) assert ek.all(res_shadow == res.is_valid()) compare_results(res_naive, res, atol=1e-6)
def test03_ray_intersect_instance(variants_all_rgb): from mitsuba.core import xml, Ray3f, ScalarVector3f, ScalarTransform4f as T """Check that we can the correct instance pointer when tracing a ray""" scene = xml.load_dict({ 'type': 'scene', 'group_0': { 'type': 'shapegroup', 'shape': { 'type': 'rectangle' } }, 'instance_00': { 'type': 'instance', "group": { "type": "ref", "id": "group_0" }, 'to_world': T.translate([-0.5, -0.5, 0.0]) * T.scale(0.5) }, 'instance_01': { 'type': 'instance', "group": { "type": "ref", "id": "group_0" }, 'to_world': T.translate([-0.5, 0.5, 0.0]) * T.scale(0.5) }, 'instance_10': { 'type': 'instance', "group": { "type": "ref", "id": "group_0" }, 'to_world': T.translate([0.5, -0.5, 0.0]) * T.scale(0.5) }, 'shape': { 'type': 'rectangle', 'to_world': T.translate([0.5, 0.5, 0.0]) * T.scale(0.5) } }) ray = Ray3f([-0.5, -0.5, -12], [0.0, 0.0, 1.0], 0.0, []) pi = scene.ray_intersect_preliminary(ray) assert '[0.5, 0, 0, -0.5]' in str(pi) assert '[0, 0.5, 0, -0.5]' in str(pi) ray = Ray3f([-0.5, 0.5, -12], [0.0, 0.0, 1.0], 0.0, []) pi = scene.ray_intersect_preliminary(ray) assert '[0.5, 0, 0, -0.5]' in str(pi) assert '[0, 0.5, 0, 0.5]' in str(pi) ray = Ray3f([0.5, -0.5, -12], [0.0, 0.0, 1.0], 0.0, []) pi = scene.ray_intersect_preliminary(ray) assert '[0.5, 0, 0, 0.5]' in str(pi) assert '[0, 0.5, 0, -0.5]' in str(pi) ray = Ray3f([0.5, 0.5, -12], [0.0, 0.0, 1.0], 0.0, []) pi = scene.ray_intersect_preliminary(ray) assert 'instance = nullptr' in str(pi) or 'instance = [nullptr]' in str(pi)
def test03_ray_intersect(variant_scalar_rgb): from mitsuba.core import xml, Ray3f, Transform4f for r in [1, 2, 4]: for l in [1, 5]: s = xml.load_dict({ "type": "scene", "foo": { "type": "cylinder", "to_world": Transform4f.scale((r, r, l)) } }) # grid size n = 10 for x in ek.linspace(Float, -1, 1, n): for z in ek.linspace(Float, -1, 1, n): x = 1.1 * r * x z = 1.1 * l * z ray = Ray3f(o=[x, -10, z], d=[0, 1, 0], time=0.0, wavelengths=[]) si_found = s.ray_test(ray) si = s.ray_intersect(ray) assert si_found == si.is_valid() assert si_found == ek.allclose(si.p[0]**2 + si.p[1]**2, r**2) if si_found: ray = Ray3f(o=[x, -10, z], d=[0, 1, 0], time=0.0, wavelengths=[]) si = s.ray_intersect(ray) ray_u = Ray3f(ray) ray_v = Ray3f(ray) eps = 1e-4 ray_u.o += si.dp_du * eps ray_v.o += si.dp_dv * eps si_u = s.ray_intersect(ray_u) si_v = s.ray_intersect(ray_v) dn = si.shape.normal_derivative(si, True, True) if si_u.is_valid(): dp_du = (si_u.p - si.p) / eps dn_du = (si_u.n - si.n) / eps assert ek.allclose(dp_du, si.dp_du, atol=2e-2) assert ek.allclose(dn_du, dn[0], atol=2e-2) if si_v.is_valid(): dp_dv = (si_v.p - si.p) / eps dn_dv = (si_v.n - si.n) / eps assert ek.allclose(dp_dv, si.dp_dv, atol=2e-2) assert ek.allclose(dn_dv, dn[1], atol=2e-2)
def test03_ray_intersect(variant_scalar_rgb): from mitsuba.core import xml, Ray3f, Transform4f # Scalar scene = xml.load_dict({ "type": "scene", "foo": { "type": "rectangle", "to_world": Transform4f.scale((2.0, 0.5, 1.0)) } }) n = 15 coords = ek.linspace(Float, -1, 1, n) rays = [ Ray3f(o=[a, a, 5], d=[0, 0, -1], time=0.0, wavelengths=[]) for a in coords ] si_scalar = [] valid_count = 0 for i in range(n): its_found = scene.ray_test(rays[i]) si = scene.ray_intersect(rays[i]) assert its_found == (abs(coords[i]) <= 0.5) assert si.is_valid() == its_found si_scalar.append(si) valid_count += its_found assert valid_count == 7 try: mitsuba.set_variant('packet_rgb') from mitsuba.core import xml, Ray3f as Ray3fX except ImportError: pytest.skip("packet_rgb mode not enabled") # Packet scene_p = xml.load_dict({ "type": "scene", "foo": { "type": "rectangle", "to_world": Transform4f.scale((2.0, 0.5, 1.0)) } }) packet = Ray3fX.zero(n) for i in range(n): packet[i] = rays[i] si_p = scene_p.ray_intersect(packet) its_found_p = scene_p.ray_test(packet) assert ek.all(si_p.is_valid() == its_found_p) for i in range(n): assert ek.allclose(si_p.t[i], si_scalar[i].t)
def test03_ray_intersect(variant_scalar_rgb): from mitsuba.core import Ray3f if mitsuba.core.MTS_ENABLE_EMBREE: pytest.skip("EMBREE enabled") for r in [1, 2, 4]: for l in [1, 5]: s = example_scene((r, r, l)) # grid size n = 10 xx = ek.linspace(Float, -1, 1, n) zz = ek.linspace(Float, -1, 1, n) for x in xx: for z in zz: x = 1.1*r*x z = 1.1*l*z ray = Ray3f(o=[x, -10, z], d=[0, 1, 0], time=0.0, wavelengths=[]) si_found = s.ray_test(ray) si = s.ray_intersect(ray) assert si_found == si.is_valid() assert si_found == ek.allclose(si.p[0]**2 + si.p[1]**2, r**2) if si_found: ray = Ray3f(o=[x, -10, z], d=[0, 1, 0], time=0.0, wavelengths=[]) si = s.ray_intersect(ray) ray_u = Ray3f(ray) ray_v = Ray3f(ray) eps = 1e-4 ray_u.o += si.dp_du * eps ray_v.o += si.dp_dv * eps si_u = s.ray_intersect(ray_u) si_v = s.ray_intersect(ray_v) dn = si.shape.normal_derivative(si, True, True) if si_u.is_valid(): dp_du = (si_u.p - si.p) / eps dn_du = (si_u.n - si.n) / eps assert ek.allclose(dp_du, si.dp_du, atol=2e-2) assert ek.allclose(dn_du, dn[0], atol=2e-2) if si_v.is_valid(): dp_dv = (si_v.p - si.p) / eps dn_dv = (si_v.n - si.n) / eps assert ek.allclose(dp_dv, si.dp_dv, atol=2e-2) assert ek.allclose(dn_dv, dn[1], atol=2e-2)
def test04_differentiable_surface_interaction_ray_forward( variant_gpu_autodiff_rgb): from mitsuba.core import xml, Ray3f, Vector3f, UInt32 shape = xml.load_dict({'type': 'disk'}) ray = Ray3f(Vector3f(0.1, -0.2, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) pi = shape.ray_intersect_preliminary(ray) ek.set_requires_gradient(ray.o) ek.set_requires_gradient(ray.d) # If the ray origin is shifted along the x-axis, so does si.p si = pi.compute_surface_interaction(ray) ek.forward(ray.o.x) assert ek.allclose(ek.gradient(si.p), [1, 0, 0]) # If the ray origin is shifted along the y-axis, so does si.p si = pi.compute_surface_interaction(ray) ek.forward(ray.o.y) assert ek.allclose(ek.gradient(si.p), [0, 1, 0]) # If the ray origin is shifted along the z-axis, so does si.t si = pi.compute_surface_interaction(ray) ek.forward(ray.o.z) assert ek.allclose(ek.gradient(si.t), -1) # If the ray direction is shifted along the x-axis, so does si.p si = pi.compute_surface_interaction(ray) ek.forward(ray.d.x) assert ek.allclose(ek.gradient(si.p), [10, 0, 0]) # If the ray origin is shifted toward the center of the disk, so does si.uv.x ray = Ray3f(Vector3f(0.9999999, 0.0, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) ek.set_requires_gradient(ray.o) si = shape.ray_intersect(ray) ek.forward(ray.o.x) assert ek.allclose(ek.gradient(si.uv), [1, 0]) # If the ray origin is shifted tangent to the disk, si.uv.y moves by 1 / (2pi) si = shape.ray_intersect(ray) ek.forward(ray.o.y) assert ek.allclose(ek.gradient(si.uv), [0, 0.5 / ek.pi], atol=1e-5) # If the ray origin is shifted tangent to the disk, si.dp_dv will also have a component is x si = shape.ray_intersect(ray) ek.forward(ray.o.y) assert ek.allclose(ek.gradient(si.dp_dv), [-1, 0, 0])
def test03_ray_intersect_transform(variant_scalar_rgb): if mitsuba.core.MTS_ENABLE_EMBREE: pytest.skip("EMBREE enabled") from mitsuba.core import Ray3f for r in [1, 3]: s = example_scene(radius=r, extra="""<transform name="to_world"> <rotate y="1.0" angle="30"/> <translate x="0.0" y="1.0" z="0.0"/> </transform>""") # grid size n = 21 inv_n = 1.0 / n for x in range(n): for y in range(n): x_coord = r * (2 * (x * inv_n) - 1) y_coord = r * (2 * (y * inv_n) - 1) ray = Ray3f(o=[x_coord, y_coord + 1, -8], d=[0.0, 0.0, 1.0], time=0.0, wavelengths=[]) si_found = s.ray_test(ray) assert si_found == (x_coord ** 2 + y_coord ** 2 <= r * r) \ or ek.abs(x_coord ** 2 + y_coord ** 2 - r * r) < 1e-8 if si_found: ray = Ray3f(o=[x_coord, y_coord + 1, -8], d=[0.0, 0.0, 1.0], time=0.0, wavelengths=[]) si = s.ray_intersect(ray) ray_u = Ray3f(ray) ray_v = Ray3f(ray) eps = 1e-4 ray_u.o += si.dp_du * eps ray_v.o += si.dp_dv * eps si_u = s.ray_intersect(ray_u) si_v = s.ray_intersect(ray_v) if si_u.is_valid(): du = (si_u.uv - si.uv) / eps assert ek.allclose(du, [1, 0], atol=2e-2) if si_v.is_valid(): dv = (si_v.uv - si.uv) / eps assert ek.allclose(dv, [0, 1], atol=2e-2)
def test14_differentiable_surface_interaction_ray_backward( variant_gpu_autodiff_rgb): from mitsuba.core import xml, Ray3f, Vector3f, UInt32 scene = xml.load_string(''' <scene version="2.0.0"> <shape type="obj" id="rect"> <string name="filename" value="resources/data/common/meshes/rectangle.obj"/> </shape> </scene> ''') ray = Ray3f(Vector3f(-0.3, -0.4, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) pi = scene.ray_intersect_preliminary(ray) ek.set_requires_gradient(ray.o) # If si.p is shifted along the x-axis, so does the ray origin si = pi.compute_surface_interaction(ray) ek.backward(si.p.x) assert ek.allclose(ek.gradient(ray.o), [1, 0, 0]) # If si.t is changed, so does the ray origin along the z-axis si = pi.compute_surface_interaction(ray) ek.backward(si.t) assert ek.allclose(ek.gradient(ray.o), [0, 0, -1])
def test04_sample_direct(variant_scalar_rgb): from mitsuba.core.xml import load_string from mitsuba.core import Ray3f from mitsuba.render import Interaction3f if mitsuba.core.MTS_ENABLE_EMBREE: pytest.skip("EMBREE enabled") sphere = load_string('<shape type="sphere" version="2.0.0"/>') def sample_cone(sample, cos_theta_max): cos_theta = (1 - sample[1]) + sample[1] * cos_theta_max sin_theta = ek.sqrt(1 - cos_theta * cos_theta) phi = 2 * ek.pi * sample[0] s, c = ek.sin(phi), ek.cos(phi) return [c * sin_theta, s * sin_theta, cos_theta] it = Interaction3f.zero() it.p = [0, 0, -3] it.t = 0 sin_cone_angle = 1.0 / it.p[2] cos_cone_angle = ek.sqrt(1 - sin_cone_angle**2) for xi_1 in ek.linspace(Float, 0, 1, 10): for xi_2 in ek.linspace(Float, 1e-3, 1 - 1e-3, 10): sample = sphere.sample_direction(it, [xi_2, 1 - xi_1]) d = sample_cone([xi_1, xi_2], cos_cone_angle) its = sphere.ray_intersect(Ray3f(it.p, d, 0, [])) assert ek.allclose(d, sample.d, atol=1e-5, rtol=1e-5) assert ek.allclose(its.t, sample.dist, atol=1e-5, rtol=1e-5) assert ek.allclose(its.p, sample.p, atol=1e-5, rtol=1e-5)
def test04_sample_direct(variant_scalar_rgb): from mitsuba.core import xml, Ray3f from mitsuba.render import Interaction3f sphere = xml.load_dict({"type": "sphere"}) def sample_cone(sample, cos_theta_max): cos_theta = (1 - sample[1]) + sample[1] * cos_theta_max sin_theta = ek.sqrt(1 - cos_theta * cos_theta) phi = 2 * ek.pi * sample[0] s, c = ek.sin(phi), ek.cos(phi) return [c * sin_theta, s * sin_theta, cos_theta] it = Interaction3f.zero() it.p = [0, 0, -3] it.t = 0 sin_cone_angle = 1.0 / it.p[2] cos_cone_angle = ek.sqrt(1 - sin_cone_angle**2) for xi_1 in ek.linspace(Float, 0, 1, 10): for xi_2 in ek.linspace(Float, 1e-3, 1 - 1e-3, 10): sample = sphere.sample_direction(it, [xi_2, 1 - xi_1]) d = sample_cone([xi_1, xi_2], cos_cone_angle) its = sphere.ray_intersect(Ray3f(it.p, d, 0, [])) assert ek.allclose(d, sample.d, atol=1e-5, rtol=1e-5) assert ek.allclose(its.t, sample.dist, atol=1e-5, rtol=1e-5) assert ek.allclose(its.p, sample.p, atol=1e-5, rtol=1e-5)
def test01_depth_scalar_stairs(variant_scalar_rgb): from mitsuba.core import Ray3f from mitsuba.render import SurfaceInteraction3f if mitsuba.core.MTS_ENABLE_EMBREE: pytest.skip("EMBREE enabled") n_steps = 20 scene = make_synthetic_scene(n_steps) n = 128 inv_n = 1.0 / (n - 1) wavelengths = [] for x in range(n - 1): for y in range(n - 1): o = [x * inv_n, y * inv_n, 2] d = [0, 0, -1] r = Ray3f(o, d, 0.5, wavelengths) r.mint = 0 r.maxt = 100 res_naive = scene.ray_intersect_naive(r) res = scene.ray_intersect(r) res_shadow = scene.ray_test(r) step_idx = ek.floor((y * inv_n) * n_steps) assert ek.all(res_shadow) assert ek.all(res_shadow == res_naive.is_valid()) expected = SurfaceInteraction3f() expected.t = 2.0 - (step_idx / n_steps) compare_results(res_naive, expected, atol=1e-9) compare_results(res_naive, res)
def test03_ray_intersect_transform(variant_scalar_rgb): from mitsuba.core import xml, Ray3f, Transform4f for r in [1, 3]: s = xml.load_dict({ "type": "sphere", "radius": r, "to_world": Transform4f.translate([0, 1, 0]) * Transform4f.rotate([0, 1, 0], 30.0) }) # grid size n = 21 inv_n = 1.0 / n for x in range(n): for y in range(n): x_coord = r * (2 * (x * inv_n) - 1) y_coord = r * (2 * (y * inv_n) - 1) ray = Ray3f(o=[x_coord, y_coord + 1, -8], d=[0.0, 0.0, 1.0], time=0.0, wavelengths=[]) si_found = s.ray_test(ray) assert si_found == (x_coord ** 2 + y_coord ** 2 <= r * r) \ or ek.abs(x_coord ** 2 + y_coord ** 2 - r * r) < 1e-8 if si_found: si = s.ray_intersect(ray) ray_u = Ray3f(ray) ray_v = Ray3f(ray) eps = 1e-4 ray_u.o += si.dp_du * eps ray_v.o += si.dp_dv * eps si_u = s.ray_intersect(ray_u) si_v = s.ray_intersect(ray_v) if si_u.is_valid(): du = (si_u.uv - si.uv) / eps assert ek.allclose(du, [1, 0], atol=2e-2) if si_v.is_valid(): dv = (si_v.uv - si.uv) / eps assert ek.allclose(dv, [0, 1], atol=2e-2)
def test02_ray_intersect_transform(variant_scalar_rgb, shape): from mitsuba.core import Ray3f, ScalarVector3f from mitsuba.render import HitComputeFlags trans = ScalarVector3f([0, 1, 0]) angle = 15 for scale in [0.5, 2.7]: s, s_inst = example_scene(shape, scale, trans, angle) # grid size n = 21 inv_n = 1.0 / n for x in range(n): for y in range(n): x_coord = scale * (2 * (x * inv_n) - 1) y_coord = scale * (2 * (y * inv_n) - 1) ray = Ray3f(o=ScalarVector3f([x_coord, y_coord, -12]) + trans, d=[0.0, 0.0, 1.0], time=0.0, wavelengths=[]) si_found = s.ray_test(ray) si_found_inst = s_inst.ray_test(ray) assert si_found == si_found_inst for dn_flags in [ HitComputeFlags.dNGdUV, HitComputeFlags.dNSdUV ]: if si_found: si = s.ray_intersect(ray, HitComputeFlags.All | dn_flags) si_inst = s_inst.ray_intersect( ray, HitComputeFlags.All | dn_flags) assert si.prim_index == si_inst.prim_index assert si.instance is None assert si_inst.instance is not None assert ek.allclose(si.t, si_inst.t, atol=2e-2) assert ek.allclose(si.time, si_inst.time, atol=2e-2) assert ek.allclose(si.p, si_inst.p, atol=2e-2) assert ek.allclose(si.dp_du, si_inst.dp_du, atol=2e-2) assert ek.allclose(si.dp_dv, si_inst.dp_dv, atol=2e-2) assert ek.allclose(si.uv, si_inst.uv, atol=2e-2) assert ek.allclose(si.wi, si_inst.wi, atol=2e-2) if ek.norm(si.dn_du) > 0.0 and ek.norm(si.dn_dv) > 0.0: assert ek.allclose(si.dn_du, si_inst.dn_du, atol=2e-2) assert ek.allclose(si.dn_dv, si_inst.dn_dv, atol=2e-2)
def test10_ray_intersect_preliminary(variants_all_rgb): if 'packet' in mitsuba.variant(): pytest.skip( "pi.compute_surface_interaction isn't bound for packet modes") from mitsuba.core import xml, Ray3f, Vector3f, UInt32 from mitsuba.render import HitComputeFlags scene = xml.load_string(''' <scene version="2.0.0"> <shape type="obj"> <string name="filename" value="resources/data/common/meshes/rectangle.obj"/> </shape> </scene> ''') ray = Ray3f(Vector3f(-0.3, -0.3, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) pi = scene.ray_intersect_preliminary(ray) assert ek.allclose(pi.t, 10) assert pi.prim_index == 0 assert ek.allclose(pi.prim_uv, [0.35, 0.3]) si = pi.compute_surface_interaction(ray) assert ek.allclose(si.t, 10) assert ek.allclose(si.p, [-0.3, -0.3, 0.0]) assert ek.allclose(si.uv, [0.35, 0.35]) assert ek.allclose(si.dp_du, [2.0, 0.0, 0.0]) assert ek.allclose(si.dp_dv, [0.0, 2.0, 0.0]) ray = Ray3f(Vector3f(0.3, 0.3, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) pi = scene.ray_intersect_preliminary(ray) assert ek.allclose(pi.t, 10) assert pi.prim_index == 1 assert ek.allclose(pi.prim_uv, [0.3, 0.35]) si = pi.compute_surface_interaction(ray) assert ek.allclose(si.t, 10) assert ek.allclose(si.p, [0.3, 0.3, 0.0]) assert ek.allclose(si.uv, [0.65, 0.65]) assert ek.allclose(si.dp_du, [2.0, 0.0, 0.0]) assert ek.allclose(si.dp_dv, [0.0, 2.0, 0.0])
def test05_differentiable_surface_interaction_ray_forward( variant_gpu_autodiff_rgb): from mitsuba.core import xml, Ray3f, Vector3f, UInt32 shape = xml.load_dict({'type': 'sphere'}) ray = Ray3f(Vector3f(0.0, -10.0, 0.0), Vector3f(0.0, 1.0, 0.0), 0, []) pi = shape.ray_intersect_preliminary(ray) ek.set_requires_gradient(ray.o) ek.set_requires_gradient(ray.d) # If the ray origin is shifted along the x-axis, so does si.p si = pi.compute_surface_interaction(ray) ek.forward(ray.o.x) assert ek.allclose(ek.gradient(si.p), [1, 0, 0]) # If the ray origin is shifted along the z-axis, so does si.p si = pi.compute_surface_interaction(ray) ek.forward(ray.o.z) assert ek.allclose(ek.gradient(si.p), [0, 0, 1]) # If the ray origin is shifted along the y-axis, so does si.t si = pi.compute_surface_interaction(ray) ek.forward(ray.o.y) assert ek.allclose(ek.gradient(si.t), -1) # If the ray direction is shifted along the x-axis, so does si.p si = pi.compute_surface_interaction(ray) ek.forward(ray.d.x) assert ek.allclose(ek.gradient(si.p), [9, 0, 0]) # If the ray origin is shifted tangent to the sphere (azimuth), so si.uv.x move by 1 / 2pi ek.set_requires_gradient(ray.o) si = shape.ray_intersect(ray) ek.forward(ray.o.x) assert ek.allclose(ek.gradient(si.uv), [1 / (2.0 * ek.pi), 0]) # If the ray origin is shifted tangent to the sphere (inclination), so si.uv.y move by 2 / 2pi ek.set_requires_gradient(ray.o) si = shape.ray_intersect(ray) ek.forward(ray.o.z) assert ek.allclose(ek.gradient(si.uv), [0, -2 / (2.0 * ek.pi)]) # # If the ray origin is shifted along the x-axis, so does si.n ek.set_requires_gradient(ray.o) si = shape.ray_intersect(ray) ek.forward(ray.o.x) assert ek.allclose(ek.gradient(si.n), [1, 0, 0]) # # If the ray origin is shifted along the z-axis, so does si.n ek.set_requires_gradient(ray.o) si = shape.ray_intersect(ray) ek.forward(ray.o.z) assert ek.allclose(ek.gradient(si.n), [0, 0, 1])
def sample_ray( self, time, sample1, # wavelength sample2, # pos sample3, # dir active): wavelengths, spec_weight = self.m_intensity.sample( SurfaceInteraction3f(), ek.arange(sample1), active) trafo = self.m_world_transform.eval(ref.time) ray = Ray3f(trafo * Point3f(0), warp.square_to_uniform_sphere(sample3), time, wavelengths) # print(spec_weight.class_().name()) return (ray, spec_weight * 4.0 * Pi)
def test01_ray_intersect(variant_scalar_rgb, shape): from mitsuba.core import Ray3f from mitsuba.render import HitComputeFlags s, s_inst = example_scene(shape) # grid size n = 21 inv_n = 1.0 / n for x in range(n): for y in range(n): x_coord = (2 * (x * inv_n) - 1) y_coord = (2 * (y * inv_n) - 1) ray = Ray3f(o=[x_coord, y_coord + 1, -8], d=[0.0, 0.0, 1.0], time=0.0, wavelengths=[]) si_found = s.ray_test(ray) si_found_inst = s_inst.ray_test(ray) assert si_found == si_found_inst if si_found: si = s.ray_intersect( ray, HitComputeFlags.All | HitComputeFlags.dNSdUV) si_inst = s_inst.ray_intersect( ray, HitComputeFlags.All | HitComputeFlags.dNSdUV) assert si.prim_index == si_inst.prim_index assert si.instance is None assert si_inst.instance is not None assert ek.allclose(si.t, si_inst.t, atol=2e-2) assert ek.allclose(si.time, si_inst.time, atol=2e-2) assert ek.allclose(si.p, si_inst.p, atol=2e-2) assert ek.allclose(si.sh_frame.n, si_inst.sh_frame.n, atol=2e-2) assert ek.allclose(si.dp_du, si_inst.dp_du, atol=2e-2) assert ek.allclose(si.dp_dv, si_inst.dp_dv, atol=2e-2) assert ek.allclose(si.uv, si_inst.uv, atol=2e-2) assert ek.allclose(si.wi, si_inst.wi, atol=2e-2) if ek.norm(si.dn_du) > 0.0 and ek.norm(si.dn_dv) > 0.0: assert ek.allclose(si.dn_du, si_inst.dn_du, atol=2e-2) assert ek.allclose(si.dn_dv, si_inst.dn_dv, atol=2e-2)
def sample_ray( self, time, sample1, # wavelength sample2, # pos sample3, # dir active): ps = self.m_shape.sample_position(time, sample2, active) local = warp.square_to_cosine_hemisphere(sample3) si = SurfaceInteraction3f(ps, 0) wavelengths, spec_weight = self.m_radiance.sample( si, ek.arange(sample1), active) ray = Ray3f(ps.p, Frame3f(ps.n).to_world(local), time, wavelengths) return (ray, spec_weight * self.m_area_times_pi)
def test12_differentiable_surface_interaction_automatic( variant_gpu_autodiff_rgb): from mitsuba.core import xml, Ray3f, Vector3f, UInt32 from mitsuba.render import HitComputeFlags scene = xml.load_string(''' <scene version="2.0.0"> <shape type="obj" id="rect"> <string name="filename" value="resources/data/common/meshes/rectangle.obj"/> </shape> </scene> ''') ray = Ray3f(Vector3f(-0.3, -0.3, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) pi = scene.ray_intersect_preliminary(ray) # si should not be attached if not necessary si = pi.compute_surface_interaction(ray) assert not ek.requires_gradient(si.t) assert not ek.requires_gradient(si.p) # si should be attached if ray is attached ek.set_requires_gradient(ray.o) si = pi.compute_surface_interaction(ray) assert ek.requires_gradient(si.t) assert ek.requires_gradient(si.p) # si should not be attached if falgs says so ek.set_requires_gradient(ray.o) si = pi.compute_surface_interaction(ray, HitComputeFlags.NonDifferentiable) assert not ek.requires_gradient(si.t) assert not ek.requires_gradient(si.p) # si should be attached if shape parameters are attached params = traverse(scene) shape_param_key = 'rect.vertex_positions_buf' ek.set_requires_gradient(params[shape_param_key]) params.set_dirty(shape_param_key) params.update() ek.set_requires_gradient(ray.o, False) si = pi.compute_surface_interaction(ray) assert ek.requires_gradient(si.t) assert ek.requires_gradient(si.p)
def test05_differentiable_surface_interaction_ray_backward(variant_gpu_autodiff_rgb): from mitsuba.core import xml, Ray3f, Vector3f, UInt32 shape = xml.load_dict({'type' : 'cylinder'}) ray = Ray3f(Vector3f(0.0, -10.0, 0.0), Vector3f(0.0, 1.0, 0.0), 0, []) pi = shape.ray_intersect_preliminary(ray) ek.set_requires_gradient(ray.o) # If si.p is shifted along the x-axis, so does the ray origin si = pi.compute_surface_interaction(ray) ek.backward(si.p.x) assert ek.allclose(ek.gradient(ray.o), [1, 0, 0]) # If si.t is changed, so does the ray origin along the z-axis si = pi.compute_surface_interaction(ray) ek.backward(si.t) assert ek.allclose(ek.gradient(ray.o), [0, -1, 0])
def test02_depth_scalar_bunny(variant_scalar_rgb): from mitsuba.core import Ray3f, BoundingBox3f from mitsuba.core.xml import load_string from mitsuba.render import SurfaceInteraction3f if mitsuba.core.MTS_ENABLE_EMBREE: pytest.skip("EMBREE enabled") scene = load_string(""" <scene version="0.5.0"> <shape type="ply"> <string name="filename" value="resources/data/ply/bunny_lowres.ply"/> </shape> </scene> """) b = scene.bbox() n = 100 inv_n = 1.0 / (n - 1) wavelengths = [] for x in range(n): for y in range(n): o = [ b.min[0] * (1 - x * inv_n) + b.max[0] * x * inv_n, b.min[1] * (1 - y * inv_n) + b.max[1] * y * inv_n, b.min[2] ] d = [0, 0, 1] r = Ray3f(o, d, 0.5, wavelengths) r.mint = 0 r.maxt = 100 res_naive = scene.ray_intersect_naive(r) res = scene.ray_intersect(r) res_shadow = scene.ray_test(r) assert ek.all(res_shadow == res_naive.is_valid()) compare_results(res_naive, res)
def test05_differentiable_surface_interaction_ray_forward( variant_gpu_autodiff_rgb): from mitsuba.core import xml, Ray3f, Vector3f, UInt32 shape = xml.load_dict({'type': 'rectangle'}) ray = Ray3f(Vector3f(-0.3, -0.3, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) pi = shape.ray_intersect_preliminary(ray) ek.set_requires_gradient(ray.o) ek.set_requires_gradient(ray.d) # If the ray origin is shifted along the x-axis, so does si.p si = pi.compute_surface_interaction(ray) ek.forward(ray.o.x) assert ek.allclose(ek.gradient(si.p), [1, 0, 0]) # If the ray origin is shifted along the y-axis, so does si.p si = pi.compute_surface_interaction(ray) ek.forward(ray.o.y) assert ek.allclose(ek.gradient(si.p), [0, 1, 0]) # If the ray origin is shifted along the x-axis, so does si.uv si = pi.compute_surface_interaction(ray) ek.forward(ray.o.x) assert ek.allclose(ek.gradient(si.uv), [0.5, 0]) # If the ray origin is shifted along the z-axis, so does si.t si = pi.compute_surface_interaction(ray) ek.forward(ray.o.z) assert ek.allclose(ek.gradient(si.t), -1) # If the ray direction is shifted along the x-axis, so does si.p si = pi.compute_surface_interaction(ray) ek.forward(ray.d.x) assert ek.allclose(ek.gradient(si.p), [10, 0, 0])
def test05_sample_half_wave_world(variant_scalar_mono_polarized): from mitsuba.core import Ray3f, Spectrum from mitsuba.core.xml import load_string from mitsuba.render import BSDFContext, TransportMode from mitsuba.render.mueller import stokes_basis, rotate_mueller_basis_collinear def spectrum_from_stokes(v): res = Spectrum(0.0) for i in range(4): res[i, 0] = v[i] return res # Test polarized implementation. Special case of delta = 180˚, also known # as a half-wave plate. (In world coordinates.) # # Following "Polarized Light - Fundamentals and Applications" by Edward Collett # Chapter 5.3: # # Case 1 & 2) Switch between diagonal linear polarization states (-45˚ & + 45˚) # Case 3 & 4) Switch circular polarization direction linear_pos = spectrum_from_stokes([1, 0, +1, 0]) linear_neg = spectrum_from_stokes([1, 0, -1, 0]) circular_right = spectrum_from_stokes([1, 0, 0, +1]) circular_left = spectrum_from_stokes([1, 0, 0, -1]) # Incident direction forward = [0, -1, 0] ctx = BSDFContext() ctx.mode = TransportMode.Importance ray = Ray3f([0, 100, 0], forward, 0.0, 0.0) # Build scene with given polarizer rotation angle scene_str = """<scene version='2.0.0'> <shape type="rectangle"> <bsdf type="retarder"> <spectrum name="delta" value="180"/> </bsdf> <transform name="to_world"> <rotate x="1" y="0" z="0" angle="-90"/> <!-- Rotate s.t. it is no longer aligned with local coordinates --> </transform> </shape> </scene>""" scene = load_string(scene_str) # Intersect ray against geometry si = scene.ray_intersect(ray) bs, M_local = si.bsdf().sample(ctx, si, 0.0, [0.0, 0.0]) M_world = si.to_world_mueller(M_local, -si.wi, bs.wo) # Make sure we are measuring w.r.t. to the optical table M = rotate_mueller_basis_collinear(M_world, forward, stokes_basis(forward), [-1, 0, 0]) # Case 1) assert ek.allclose(M @ linear_pos, linear_neg, atol=1e-3) # Case 2) assert ek.allclose(M @ linear_neg, linear_pos, atol=1e-3) # Case 3) assert ek.allclose(M @ circular_right, circular_left, atol=1e-3) # Case 4) assert ek.allclose(M @ circular_left, circular_right, atol=1e-3)
def test03_sample_world(variant_scalar_mono_polarized): from mitsuba.core import Ray3f, Spectrum from mitsuba.core.xml import load_string from mitsuba.render import BSDFContext, TransportMode from mitsuba.render.mueller import stokes_basis, rotate_mueller_basis_collinear def spectrum_from_stokes(v): res = Spectrum(0.0) for i in range(4): res[i,0] = v[i] return res # Test polarized implementation, version in world coordinates. This involves # additional coordinate changes to the local reference frame and back. # # The polarizer is rotated to different angles and hit with fully # unpolarized light (Stokes vector [1, 0, 0, 0]). # We then test if the outgoing Stokes vector corresponds to the expected # rotation of linearly polarized light. # Incident direction forward = [0, 1, 0] stokes_in = spectrum_from_stokes([1, 0, 0, 0]) ctx = BSDFContext() ctx.mode = TransportMode.Importance ray = Ray3f([0, -100, 0], forward, 0.0, 0.0) # Polarizer rotation angles angles = [0, 90, +45, -45] # Expected outgoing Stokes vector expected_states = [spectrum_from_stokes([0.5, 0.5, 0, 0]), spectrum_from_stokes([0.5, -0.5, 0, 0]), spectrum_from_stokes([0.5, 0, +0.5, 0]), spectrum_from_stokes([0.5, 0, -0.5, 0])] for k in range(len(angles)): angle = angles[k] expected = expected_states[k] # Build scene with given polarizer rotation angle scene_str = """<scene version='2.0.0'> <shape type="rectangle"> <bsdf type="polarizer"/> <transform name="to_world"> <rotate x="0" y="0" z="1" angle="{}"/> <!-- Rotate around optical axis --> <rotate x="1" y="0" z="0" angle="-90"/> <!-- Rotate s.t. it is no longer aligned with local coordinates --> </transform> </shape> </scene>""".format(angle) scene = load_string(scene_str) # Intersect ray against geometry si = scene.ray_intersect(ray) bs, M_local = si.bsdf().sample(ctx, si, 0.0, [0.0, 0.0]) M_world = si.to_world_mueller(M_local, -si.wi, bs.wo) # Make sure we are measuring w.r.t. to the optical table M = rotate_mueller_basis_collinear(M_world, forward, stokes_basis(forward), [-1, 0, 0]) stokes_out = M @ stokes_in assert ek.allclose(stokes_out, expected, atol=1e-3)
def test03_ray_intersect(variant_scalar_rgb): if mitsuba.core.MTS_ENABLE_EMBREE: pytest.skip("EMBREE enabled") from mitsuba.core.xml import load_string from mitsuba.core import Ray3f # Scalar scene = load_string("""<scene version="2.0.0"> <shape type="rectangle"> <transform name="to_world"> <scale x="2" y="0.5" z="1"/> </transform> </shape> </scene>""") n = 15 coords = ek.linspace(Float, -1, 1, n) rays = [ Ray3f(o=[a, a, 5], d=[0, 0, -1], time=0.0, wavelengths=[]) for a in coords ] si_scalar = [] valid_count = 0 for i in range(n): # print(rays[i]) its_found = scene.ray_test(rays[i]) si = scene.ray_intersect(rays[i]) assert its_found == (abs(coords[i]) <= 0.5) assert si.is_valid() == its_found si_scalar.append(si) valid_count += its_found assert valid_count == 7 try: mitsuba.set_variant('packet_rgb') from mitsuba.core.xml import load_string as load_string_packet from mitsuba.core import Ray3f as Ray3fX except ImportError: pytest.skip("packet_rgb mode not enabled") # Packet scene_p = load_string_packet("""<scene version="2.0.0"> <shape type="rectangle"> <transform name="to_world"> <scale x="2" y="0.5" z="1"/> </transform> </shape> </scene>""") packet = Ray3fX.zero(n) for i in range(n): packet[i] = rays[i] si_p = scene_p.ray_intersect(packet) its_found_p = scene_p.ray_test(packet) assert ek.all(si_p.is_valid() == its_found_p) for i in range(n): assert ek.allclose(si_p.t[i], si_scalar[i].t)
def test16_differentiable_surface_interaction_params_backward( variant_gpu_autodiff_rgb): from mitsuba.core import xml, Float, Ray3f, Vector3f, UInt32, Transform4f scene = xml.load_string(''' <scene version="2.0.0"> <shape type="obj" id="rect"> <string name="filename" value="resources/data/common/meshes/rectangle.obj"/> </shape> </scene> ''') params = traverse(scene) vertex_pos_key = 'rect.vertex_positions_buf' vertex_normals_key = 'rect.vertex_normals_buf' vertex_texcoords_key = 'rect.vertex_texcoords_buf' ek.set_requires_gradient(params[vertex_pos_key]) ek.set_requires_gradient(params[vertex_normals_key]) ek.set_requires_gradient(params[vertex_texcoords_key]) params.set_dirty(vertex_pos_key) params.set_dirty(vertex_normals_key) params.set_dirty(vertex_texcoords_key) params.update() # Hit the upper right corner of the rectancle (the 4th vertex) ray = Ray3f(Vector3f(0.99999, 0.99999, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) pi = scene.ray_intersect_preliminary(ray) # --------------------------------------- # Test vertex posistions # If si.t changes, so the 4th vertex should move along the z-axis si = pi.compute_surface_interaction(ray) ek.backward(si.t) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], atol=1e-5) # If si.p moves along the z-axis, so does the 4th vertex si = pi.compute_surface_interaction(ray) ek.backward(si.p.z) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], atol=1e-5) # To increase si.dp_du along the x-axis, we need to strech the upper edge of the rectangle si = pi.compute_surface_interaction(ray) ek.backward(si.dp_du.x) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [0, 0, 0, -1, 0, 0, 0, 0, 0, 1, 0, 0], atol=1e-5) # To increase si.dp_du along the y-axis, we need to transform the rectangle into a trapezoid si = pi.compute_surface_interaction(ray) ek.backward(si.dp_du.y) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [0, 0, 0, 0, -1, 0, 0, 0, 0, 0, 1, 0], atol=1e-5) # To increase si.dp_dv along the x-axis, we need to transform the rectangle into a trapezoid si = pi.compute_surface_interaction(ray) ek.backward(si.dp_dv.x) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [-1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], atol=1e-5) # To increase si.dp_dv along the y-axis, we need to strech the right edge of the rectangle si = pi.compute_surface_interaction(ray) ek.backward(si.dp_dv.y) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [0, -1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0], atol=1e-5) # To increase si.n along the x-axis, we need to rotate the right edge around the y axis si = pi.compute_surface_interaction(ray) ek.backward(si.n.x) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [0, 0, 0, 0, 0, 0.5, 0, 0, 0, 0, 0, -0.5], atol=1e-5) # To increase si.n along the y-axis, we need to rotate the top edge around the x axis si = pi.compute_surface_interaction(ray) ek.backward(si.n.y) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [0, 0, 0.5, 0, 0, 0, 0, 0, 0, 0, 0, -0.5], atol=1e-5) # To increase si.sh_frame.n along the x-axis, we need to rotate the right edge around the y axis params.set_dirty(vertex_pos_key) params.update() si = pi.compute_surface_interaction(ray) ek.backward(si.sh_frame.n.x) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [0, 0, 0, 0, 0, 0.5, 0, 0, 0, 0, 0, -0.5], atol=1e-5) # To increase si.sh_frame.n along the y-axis, we need to rotate the top edge around the x axis params.set_dirty(vertex_pos_key) params.update() si = pi.compute_surface_interaction(ray) ek.backward(si.sh_frame.n.y) assert ek.allclose(ek.gradient(params[vertex_pos_key]), [0, 0, 0.5, 0, 0, 0, 0, 0, 0, 0, 0, -0.5], atol=1e-5) # --------------------------------------- # Test vertex texcoords # To increase si.uv along the x-axis, we need to move the uv of the 4th vertex along the x-axis si = pi.compute_surface_interaction(ray) ek.backward(si.uv.x) assert ek.allclose(ek.gradient(params[vertex_texcoords_key]), [0, 0, 0, 0, 0, 0, 1, 0], atol=1e-5) # To increase si.uv along the y-axis, we need to move the uv of the 4th vertex along the y-axis si = pi.compute_surface_interaction(ray) ek.backward(si.uv.y) assert ek.allclose(ek.gradient(params[vertex_texcoords_key]), [0, 0, 0, 0, 0, 0, 0, 1], atol=1e-5) # To increase si.dp_du along the x-axis, we need to shrink the uv along the top edge of the rectangle si = pi.compute_surface_interaction(ray) ek.backward(si.dp_du.x) assert ek.allclose(ek.gradient(params[vertex_texcoords_key]), [0, 0, 2, 0, 0, 0, -2, 0], atol=1e-5) # To increase si.dp_du along the y-axis, we need to shrink the uv along the right edge of the rectangle si = pi.compute_surface_interaction(ray) ek.backward(si.dp_dv.y) assert ek.allclose(ek.gradient(params[vertex_texcoords_key]), [0, 2, 0, 0, 0, 0, 0, -2], atol=1e-5)
def test15_differentiable_surface_interaction_params_forward( variant_gpu_autodiff_rgb): from mitsuba.core import xml, Float, Ray3f, Vector3f, UInt32, Transform4f # Convert flat array into a vector of arrays (will be included in next enoki release) def ravel(buf, dim=3): idx = dim * UInt32.arange(ek.slices(buf) // dim) if dim == 2: return Vector2f(ek.gather(buf, idx), ek.gather(buf, idx + 1)) elif dim == 3: return Vector3f(ek.gather(buf, idx), ek.gather(buf, idx + 1), ek.gather(buf, idx + 2)) # Return contiguous flattened array (will be included in next enoki release) def unravel(source, target, dim=3): idx = UInt32.arange(ek.slices(source)) for i in range(dim): ek.scatter(target, source[i], dim * idx + i) scene = xml.load_string(''' <scene version="2.0.0"> <shape type="obj" id="rect"> <string name="filename" value="resources/data/common/meshes/rectangle.obj"/> </shape> </scene> ''') params = traverse(scene) shape_param_key = 'rect.vertex_positions_buf' positions_buf = params[shape_param_key] positions_initial = ravel(positions_buf) # Create differential parameter to be optimized diff_vector = Vector3f(0.0) ek.set_requires_gradient(diff_vector) # Apply the transformation to mesh vertex position and update scene def apply_transformation(trasfo): trasfo = trasfo(diff_vector) new_positions = trasfo.transform_point(positions_initial) unravel(new_positions, params[shape_param_key]) params.set_dirty(shape_param_key) params.update() # --------------------------------------- # Test translation ray = Ray3f(Vector3f(-0.2, -0.3, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) pi = scene.ray_intersect_preliminary(ray) # # If the vertices are shifted along z-axis, so does si.t apply_transformation(lambda v: Transform4f.translate(v)) si = pi.compute_surface_interaction(ray) ek.forward(diff_vector.z) assert ek.allclose(ek.gradient(si.t), 1) # If the vertices are shifted along z-axis, so does si.p apply_transformation(lambda v: Transform4f.translate(v)) si = pi.compute_surface_interaction(ray) ek.forward(diff_vector.z) assert ek.allclose(ek.gradient(si.p), [0.0, 0.0, 1.0]) # If the vertices are shifted along x-axis, so does si.uv (times 0.5) apply_transformation(lambda v: Transform4f.translate(v)) si = pi.compute_surface_interaction(ray) ek.forward(diff_vector.x) assert ek.allclose(ek.gradient(si.uv), [-0.5, 0.0]) # If the vertices are shifted along y-axis, so does si.uv (times 0.5) apply_transformation(lambda v: Transform4f.translate(v)) si = pi.compute_surface_interaction(ray) ek.forward(diff_vector.y) assert ek.allclose(ek.gradient(si.uv), [0.0, -0.5]) # --------------------------------------- # Test rotation ray = Ray3f(Vector3f(-0.99999, -0.99999, -10.0), Vector3f(0.0, 0.0, 1.0), 0, []) pi = scene.ray_intersect_preliminary(ray) # If the vertices are rotated around the center, so does si.uv (times 0.5) apply_transformation(lambda v: Transform4f.rotate([0, 0, 1], v.x)) si = pi.compute_surface_interaction(ray) ek.forward(diff_vector.x) du = 0.5 * ek.sin(2 * ek.pi / 360.0) assert ek.allclose(ek.gradient(si.uv), [-du, du], atol=1e-6)
def test03_ray_intersect(variant_scalar_rgb): from mitsuba.core import xml, Ray3f, Vector3f, Transform4f from mitsuba.render import HitComputeFlags for r in [1, 3, 5]: for translate in [ Vector3f([0.0, 0.0, 0.0]), Vector3f([1.0, -5.0, 0.0]) ]: s = xml.load_dict({ "type": "scene", "foo": { "type": "disk", "to_world": Transform4f.translate(translate) * Transform4f.scale( (r, r, 1.0)) } }) # grid size n = 10 for x in ek.linspace(Float, -1, 1, n): for y in ek.linspace(Float, -1, 1, n): x = 1.1 * r * (x - translate[0]) y = 1.1 * r * (y - translate[1]) ray = Ray3f(o=[x, y, -10], d=[0, 0, 1], time=0.0, wavelengths=[]) si_found = s.ray_test(ray) assert si_found == (x**2 + y**2 <= r * r) if si_found: ray = Ray3f(o=[x, y, -10], d=[0, 0, 1], time=0.0, wavelengths=[]) si = s.ray_intersect( ray, HitComputeFlags.All | HitComputeFlags.dNSdUV) ray_u = Ray3f(ray) ray_v = Ray3f(ray) eps = 1e-4 ray_u.o += si.dp_du * eps ray_v.o += si.dp_dv * eps si_u = s.ray_intersect(ray_u) si_v = s.ray_intersect(ray_v) if si_u.is_valid(): dp_du = (si_u.p - si.p) / eps dn_du = (si_u.n - si.n) / eps assert ek.allclose(dp_du, si.dp_du, atol=2e-2) assert ek.allclose(dn_du, si.dn_du, atol=2e-2) if si_v.is_valid(): dp_dv = (si_v.p - si.p) / eps dn_dv = (si_v.n - si.n) / eps assert ek.allclose(dp_dv, si.dp_dv, atol=2e-2) assert ek.allclose(dn_dv, si.dn_dv, atol=2e-2)