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()]
