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 backward(ctx, grad_out): ek.set_gradient(ctx.out, ek.FloatC(grad_out)) ek.FloatD.backward() result = (ek.gradient(ctx.in1).torch() if ek.requires_gradient(ctx.in1) else None, ek.gradient(ctx.in2).torch() if ek.requires_gradient(ctx.in2) else None) del ctx.out, ctx.in1, ctx.in2 ek.cuda_malloc_trim() return result
def step(self): """ Take a gradient step """ self.t += 1 from mitsuba.core import Float lr_t = ek.detach(Float(self.lr * ek.sqrt(1 - self.beta_2**self.t) / (1 - self.beta_1**self.t), literal=False)) for k, p in self.params.items(): g_p = ek.gradient(p) size = ek.slices(g_p) if size == 0: continue elif size != ek.slices(self.state[k][0]): # Reset state if data size has changed self._reset(k) m_tp, v_tp = self.state[k] m_t = self.beta_1 * m_tp + (1 - self.beta_1) * g_p v_t = self.beta_2 * v_tp + (1 - self.beta_2) * ek.sqr(g_p) self.state[k] = (m_t, v_t) u = ek.detach(p) - lr_t * m_t / (ek.sqrt(v_t) + self.epsilon) u = type(p)(u) ek.set_requires_gradient(u) self.params[k] = u
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 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 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 step(self): """ Take a gradient step """ self.t += 1 from enoki.cuda_autodiff import Float32 as Float lr_t = ek.detach( Float(self.lr * ek.sqrt(1 - self.beta_2**self.t) / (1 - self.beta_1**self.t), literal=False)) for k in self.bsdf_ad_keys: g_p = ek.gradient(self.bsdf_map[k].reflectance.data) size = ek.slices(g_p) assert (size == ek.slices(self.state[k][0])) m_tp, v_tp = self.state[k] m_t = self.beta_1 * m_tp + (1 - self.beta_1) * g_p v_t = self.beta_2 * v_tp + (1 - self.beta_2) * ek.sqr(g_p) self.state[k] = (m_t, v_t) u = ek.detach(self.bsdf_map[k].reflectance.data) - lr_t * m_t / ( ek.sqrt(v_t) + self.epsilon) u = type(self.bsdf_map[k].reflectance.data)(u) ek.set_requires_gradient(u) self.bsdf_map[k].reflectance.data = u for k in self.mesh_ad_keys: g_p = ek.gradient(self.mesh_map[k].vertex_positions) size = ek.slices(g_p) assert (size == ek.slices(self.state[k][0])) m_tp, v_tp = self.state[k] m_t = self.beta_1 * m_tp + (1 - self.beta_1) * g_p v_t = self.beta_2 * v_tp + (1 - self.beta_2) * ek.sqr(g_p) self.state[k] = (m_t, v_t) u = ek.detach(self.mesh_map[k].vertex_positions) - lr_t * m_t / ( ek.sqrt(v_t) + self.epsilon) u = type(self.mesh_map[k].vertex_positions)(u) ek.set_requires_gradient(u) self.mesh_map[k].vertex_positions = u
def backward(ctx, grad_output): try: ek.set_gradient(ctx.output, ek.detach(Float(grad_output))) Float.backward() result = tuple(ek.gradient(i).torch() if i is not None else None for i in ctx.inputs) del ctx.output del ctx.inputs ek.cuda_malloc_trim() return result except Exception as e: print("render_torch(): critical exception during " "backward pass: %s" % str(e)) raise e
def test01_set_gradient(): a = ek.FloatD(42, 10) ek.set_requires_gradient(a) with pytest.raises(TypeError): grad = ek.FloatD(-1, 10) ek.set_gradient(a, grad) grad = ek.FloatC(-1, 10) ek.set_gradient(a, grad) assert np.allclose(grad.numpy(), ek.gradient(a).numpy()) # Note: if `backward` is not called here, test03 segfaults later. # TODO: we should not need this, there's most likely some missing cleanup when `a` is destructed ek.FloatD.backward() del a, grad
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 step(self): """ Take a gradient step """ for k, p in self.params.items(): g_p = ek.gradient(p) size = ek.slices(g_p) if size == 0: continue if self.momentum != 0: if size != ek.slices(self.state[k]): # Reset state if data size has changed self._reset(k) self.state[k] = self.momentum * self.state[k] + g_p value = ek.detach(p) - self.lr_v * self.state[k] else: value = ek.detach(p) - self.lr_v * g_p value = type(p)(value) ek.set_requires_gradient(value) self.params[k] = value self.params.update()
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 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 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)
# Load the Cornell Box Thread.thread().file_resolver().append('cbox') scene = load_file( 'C:/MyFile/code/ray tracing/misuba2/test/gpu_autodiff/cbox/cbox.xml') # Find differentiable scene parameters params = traverse(scene) # Keep track of derivatives with respect to one parameter param_0 = params['red.reflectance.value'] ek.set_requires_gradient(param_0) # Differentiable simulation image = render(scene, spp=4) # Assign the gradient [1, 1, 1] to the 'red.reflectance.value' input ek.set_gradient(param_0, [1, 1, 1], backward=False) # Forward-propagate previously assigned gradients Float.forward() # The gradients have been propagated to the output image image_grad = ek.gradient(image) # .. write them to a PNG file crop_size = scene.sensors()[0].film().crop_size() fname = 'C:/MyFile/code/ray tracing/misuba2/test/gpu_autodiff/output/out.png' write_bitmap(fname, image_grad, crop_size) print('Wrote forward differentiation image to: {}'.format(fname))
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)
diff_init = ek.hsum(ek.sqr(params_opt['displacements'] - disp_tex_data_init * amp)) / ek.slices(params_opt['displacements']) diff_ref = ek.hsum(ek.sqr(params_opt['displacements'] - disp_tex_data_ref * amp)) / ek.slices(params_opt['displacements']) #print("diff_init:", diff_init[0]) #print("diff_ref:", diff_ref[0]) diff_vertex_init.append(diff_init[0]) diff_vertex_ref.append(diff_ref[0]) if(loss[0] != loss[0]): print("[WARNING] Skipping current iteration due to NaN loss.") continue ek.backward(loss) if ek.any(ek.isnan(ek.gradient(params_opt['displacements']))): print("[WARNING] NaNs in the displacement gradients. ({iteration:d})".format(iteration=i)) exit(-1) opt.step() if ek.any(ek.isnan(params_opt['displacements'])): print("[WARNING] NaNs in the vertex displacements. ({iteration:d})".format(iteration=i)) exit(-1) vertex_pos_optim = ravel(params[vertex_pos_key]) vertex_np = np.array(vertex_pos_optim) ind = np.linspace(0, vertex_np.shape[0]-1, num=vertex_np.shape[0]) ind = np.reshape(ind,(-1,1)) vertex_np = np.concatenate((vertex_np, ind), axis = 1) vertex_np = vertex_np[vertex_np[:,1].argsort()]