def test12_hmax_fwd(m): x = m.Float(1, 2, 8, 5, 8) ek.enable_grad(x) y = ek.hmax_async(x) ek.forward(x) assert len(y) == 1 and ek.allclose(y[0], 8) assert ek.allclose(ek.grad(y), [2]) # Approximation
def test52_scatter_fwd_eager(m): with EagerMode(): x = m.Float(4.0) ek.enable_grad(x) ek.set_grad(x, 1) values = x * x * ek.linspace(m.Float, 1, 4, 4) idx = 2 * ek.arange(m.UInt32, 4) buf = ek.zero(m.Float, 10) ek.scatter(buf, values, idx) assert ek.grad_enabled(buf) ref = [16.0, 0.0, 32.0, 0.0, 48.0, 0.0, 64.0, 0.0, 0.0, 0.0] assert ek.allclose(buf, ref) grad = ek.grad(buf) ref_grad = [8.0, 0.0, 16.0, 0.0, 24.0, 0.0, 32.0, 0.0, 0.0, 0.0] assert ek.allclose(grad, ref_grad) # Overwrite first value with non-diff value, resulting gradient entry should be 0 y = m.Float(3) idx = m.UInt32(0) ek.scatter(buf, y, idx) ref = [3.0, 0.0, 32.0, 0.0, 48.0, 0.0, 64.0, 0.0, 0.0, 0.0] assert ek.allclose(buf, ref) ek.forward(x) grad = ek.grad(buf) ref_grad = [0.0, 0.0, 16.0, 0.0, 24.0, 0.0, 32.0, 0.0, 0.0, 0.0] assert ek.allclose(grad, ref_grad)
def test_55_diffloop_simple_fwd(m, no_record): fi, fo = m.Float(1, 2, 3), m.Float(0, 0, 0) ek.enable_grad(fi) loop = m.Loop("MyLoop", lambda: fo) while loop(fo < 10): fo += fi ek.forward(fi) assert ek.grad(fo) == m.Float(10, 5, 4)
def test_57_diffloop_masking_rev(m, no_record): fo = ek.zero(m.Float, 10) fi = m.Float(1, 2) i = m.UInt32(0, 5) ek.enable_grad(fi) loop = m.Loop("MyLoop", lambda: i) while loop(i < 5): ek.scatter_reduce(ek.ReduceOp.Add, fo, fi, i) i += 1 ek.forward(fi) assert fo == m.Float(1, 1, 1, 1, 1, 0, 0, 0, 0, 0) assert ek.grad(fo) == m.Float(1, 1, 1, 1, 1, 0, 0, 0, 0, 0)
def test45_diff_loop(m): def mcint(a, b, f, sample_count=100000): rng = m.PCG32() i = m.UInt32(0) result = m.Float(0) l = m.Loop(i, rng, result) while l.cond(i < sample_count): result += f(ek.lerp(a, b, rng.next_float32())) i += 1 return result * (b - a) / sample_count class EllipticK(ek.CustomOp): # --- Internally used utility methods --- # Integrand of the 'K' function def K(self, x, m_): return ek.rsqrt(1 - m_ * ek.sqr(ek.sin(x))) # Derivative of the above with respect to 'm' def dK(self, x, m_): m_ = m.Float(m_) # Convert 'm' to differentiable type ek.enable_grad(m_) y = self.K(x, m_) ek.forward(m_) return ek.grad(y) # Monte Carlo integral of dK, used in forward/reverse pass def eval_grad(self): return mcint(a=0, b=ek.Pi / 2, f=lambda x: self.dK(x, self.m_)) # --- CustomOp interface --- def eval(self, m_): self.m_ = m_ # Stash 'm' for later return mcint(a=0, b=ek.Pi / 2, f=lambda x: self.K(x, self.m_)) def forward(self): self.set_grad_out(self.grad_in('m_') * self.eval_grad()) def elliptic_k(m_): return ek.custom(EllipticK, m_) ek.enable_flag(ek.JitFlag.RecordLoops) x = m.Float(0.5) ek.enable_grad(x) y = elliptic_k(x) ek.forward(x) assert ek.allclose(y, 1.85407, rtol=5e-4) assert ek.allclose(ek.grad(y), 0.847213, rtol=5e-4) ek.disable_flag(ek.JitFlag.RecordLoops)
def test_ad_operations(package): Float, Array3f = package.Float, package.Array3f prepare(package) class MyStruct: ENOKI_STRUCT = {'a': Array3f, 'b': Float} def __init__(self): self.a = Array3f() self.b = Float() foo = ek.zero(MyStruct, 4) assert not ek.grad_enabled(foo.a) assert not ek.grad_enabled(foo.b) assert not ek.grad_enabled(foo) ek.enable_grad(foo) assert ek.grad_enabled(foo.a) assert ek.grad_enabled(foo.b) assert ek.grad_enabled(foo) foo_detached = ek.detach(foo) assert not ek.grad_enabled(foo_detached.a) assert not ek.grad_enabled(foo_detached.b) assert not ek.grad_enabled(foo_detached) x = Float(4.0) ek.enable_grad(x) foo.a += x foo.b += x * x ek.forward(x) foo_grad = ek.grad(foo) assert foo_grad.a == 1 assert foo_grad.b == 8 ek.set_grad(foo, 5.0) foo_grad = ek.grad(foo) assert foo_grad.a == 5.0 assert foo_grad.b == 5.0 ek.accum_grad(foo, 5.0) foo_grad = ek.grad(foo) assert foo_grad.a == 10.0 assert foo_grad.b == 10.0
def render_gradient(scene, passes, diff_params): """Render radiance and gradient image using forward autodiff""" from mitsuba.python.autodiff import render fsize = scene.sensors()[0].film().size() img = np.zeros((fsize[1], fsize[0], 3), dtype=np.float32) grad = np.zeros((fsize[1], fsize[0], 1), dtype=np.float32) for i in range(passes): img_i = render(scene) ek.forward(diff_params, i == passes - 1) grad_i = ek.gradient(img_i).numpy().reshape(fsize[1], fsize[0], -1)[:, :, [0]] img_i = img_i.numpy().reshape(fsize[1], fsize[0], -1) # Remove NaNs grad_i[grad_i != grad_i] = 0 img_i[img_i != img_i] = 0 grad += grad_i img += img_i return img / passes, grad / passes
def test22_scatter_fwd_permute(m): x = m.Float(4.0) ek.enable_grad(x) values_0 = x * ek.linspace(m.Float, 1, 9, 5) values_1 = x * ek.linspace(m.Float, 11, 19, 5) buf = ek.zero(m.Float, 10) idx_0 = ek.arange(m.UInt32, 5) idx_1 = ek.arange(m.UInt32, 5) + 5 ek.scatter(buf, values_0, idx_0, permute=False) ek.scatter(buf, values_1, idx_1, permute=False) ref = [4.0, 12.0, 20.0, 28.0, 36.0, 44.0, 52.0, 60.0, 68.0, 76.0] assert ek.allclose(buf, ref) ek.forward(x) grad = ek.grad(buf) ref_grad = [1.0, 3.0, 5.0, 7.0, 9.0, 11.0, 13.0, 15.0, 17.0, 19.0] assert ek.allclose(grad, ref_grad)
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 test47_nan_propagation(m): for i in range(2): x = ek.arange(m.Float, 10) ek.enable_grad(x) f0 = m.Float(0) y = ek.select(x < (20 if i == 0 else 0), x, x * (f0 / f0)) ek.backward(y) g = ek.grad(x) if i == 0: assert g == 1 else: assert ek.all(ek.isnan(g)) for i in range(2): x = ek.arange(m.Float, 10) ek.enable_grad(x) f0 = m.Float(0) y = ek.select(x < (20 if i == 0 else 0), x, x * (f0 / f0)) ek.forward(x) g = ek.grad(y) if i == 0: assert g == 1 else: assert ek.all(ek.isnan(g))
def test02_add_fwd(m): if True: a, b = m.Float(1), m.Float(2) ek.enable_grad(a, b) c = 2 * a + b ek.forward(a, retain_graph=True) assert ek.grad(c) == 2 ek.set_grad(c, 101) ek.forward(b) assert ek.grad(c) == 102 if True: a, b = m.Float(1), m.Float(2) ek.enable_grad(a, b) c = 2 * a + b ek.set_grad(a, 1.0) ek.enqueue(a) ek.traverse(m.Float, retain_graph=True, reverse=False) assert ek.grad(c) == 2 assert ek.grad(a) == 0 ek.set_grad(a, 1.0) ek.enqueue(a) ek.traverse(m.Float, retain_graph=False, reverse=False) assert ek.grad(c) == 4
def test13_differentiable_surface_interaction_ray_forward( 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) 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 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 test08_hsum_1_fwd(m): x = ek.linspace(m.Float, 0, 1, 10) ek.enable_grad(x) y = ek.hsum_async(ek.hsum_async(x) * x) ek.forward(x) assert ek.allclose(ek.grad(y), 100)
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 dK(self, x, m_): m_ = m.Float(m_) # Convert 'm' to differentiable type ek.enable_grad(m_) y = self.K(x, m_) ek.forward(m_) return ek.grad(y)
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 test10_hsum_2_fwd(m): x = ek.linspace(m.Float, 0, 1, 10) ek.enable_grad(x) y = ek.hsum_async(ek.hsum_async(x * x) * ek.hsum_async(x * x)) ek.forward(x) assert ek.allclose(ek.grad(y), 1900.0 / 27.0)
def run_ad(integrator, sc, fname, args): global time_threshold ad_config = args["AD"] if "spp" in ad_config: sc.opts.spp = ad_config["spp"] if "sppe" in ad_config: sc.opts.sppe = ad_config["sppe"] if "sppse" in ad_config: sc.opts.sppse = ad_config["sppse"] if "no_edge" in ad_config: for i in ad_config["no_edge"]: sc.param_map["Mesh[" + str(i) + "]"].enable_edges = False ro = sc.opts if ad_config["type"] == "mesh_transform": if len(ad_config["Mesh_ID"]) != len(ad_config["Mesh_dir"]): raise Exception("Mesh_ID and Mesh_dir have different sizes") elif ad_config["type"] == "mesh_rotate": if len(ad_config["Mesh_ID"]) != len(ad_config["axis"]): raise Exception("Mesh_ID and axis have different sizes") elif ad_config["type"] == "vertex_transform": if len(ad_config["Mesh_ID"]) != len(ad_config["Vertex_ID"]): raise Exception("Mesh_ID and Vertex_ID have different sizes") orig_vtx_pos = {} for j in ad_config["Mesh_ID"]: mesh_obj = sc.param_map["Mesh[" + str(j) + "]"] orig_vtx_pos[j] = ek.detach(mesh_obj.vertex_positions) elif ad_config["type"] == "material_roughness": base_roughness = {} for j in ad_config["BSDF_ID"]: bsdf_obj = sc.param_map["BSDF[" + str(j) + "]"] base_roughness[j] = (ek.detach(bsdf_obj.alpha_u.data), ek.detach(bsdf_obj.alpha_v.data)) elif ad_config["type"] == "envmap_rotate": if "Emitter_ID" not in ad_config: raise Exception("Missing Emitter_ID") else: raise Exception("Unknown transform") if "npass" in ad_config: npass = ad_config["npass"] elif "npass" in args: npass = args["npass"] else: npass = 1 num_sensors = sc.num_sensors img_ad = [None] * num_sensors t0 = time.process_time() t1 = t0 for i in range(npass): # AD config P = FloatD(0.) ek.set_requires_gradient(P) if ad_config["type"] == "mesh_transform": for j in range(len(ad_config["Mesh_ID"])): mesh_transform(sc, ad_config["Mesh_ID"][j], Vector3fD(ad_config["Mesh_dir"][j]) * P) elif ad_config["type"] == "mesh_rotate": for j in range(len(ad_config["Mesh_ID"])): mesh_rotate(sc, ad_config["Mesh_ID"][j], Vector3fD(ad_config["axis"][j]), P) elif ad_config["type"] == "vertex_transform": for j in range(len(ad_config["Mesh_ID"])): vertex_transform(sc, ad_config["Mesh_ID"][j], ad_config["Vertex_ID"][j], ad_config["Vertex_dir"][j], orig_vtx_pos[j], P) elif ad_config["type"] == "material_roughness": for j in ad_config["BSDF_ID"]: material_roughness(sc, j, base_roughness[j], P) elif ad_config["type"] == "envmap_rotate": envmap_rotate(sc, ad_config["Emitter_ID"], ad_config["axis"], P) # End AD config sc.configure() for sensor_id in range(num_sensors): if i == 0 and "guide" in ad_config: t2 = time.process_time() guide_info = ad_config["guide"] integrator.preprocess_secondary_edges( sc, sensor_id, np.array(guide_info["reso"]), guide_info["nround"]) print("guiding done in %.2f seconds." % (time.process_time() - t2)) img = integrator.renderD(sc, sensor_id) ek.forward(P, free_graph=True) grad_img = ek.gradient(img).numpy() grad_img[np.logical_not(np.isfinite(grad_img))] = 0. if i == 0: img_ad[sensor_id] = grad_img else: img_ad[sensor_id] += grad_img del img del P t2 = time.process_time() if t2 - t1 > time_threshold: print("(%d/%d) done in %.2f seconds." % (i + 1, npass, t2 - t0), end="\r") t1 = t2 print("(%d/%d) Total AD rendering time: %.2f seconds." % (npass, npass, t2 - t0)) for sensor_id in range(num_sensors): img = (img_ad[sensor_id] / float(npass)).reshape( (ro.height, ro.width, 3)) output = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite(fname[:-4] + "_" + str(sensor_id) + fname[-4:], output)
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)