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 test03_shapes_parameters_grad_enabled(variant_gpu_autodiff_rgb):
from mitsuba.core.xml import load_string
from mitsuba.python.util import traverse
scene = load_string("""
<scene version="2.0.0">
<shape type="obj" id="box">
<string name="filename" value="resources/data/tests/obj/cbox_smallbox.obj"/>
</shape>
<shape type="sphere"/>
</scene>
""")
# Initial scene should always return False
assert scene.shapes_grad_enabled() == False
# Get scene parameters
params = traverse(scene)
# Only parameters of the shape should affect the result of that method
bsdf_param_key = 'box.bsdf.reflectance.value'
ek.set_requires_gradient(params[bsdf_param_key])
params.set_dirty(bsdf_param_key)
params.update()
assert scene.shapes_grad_enabled() == False
# When setting one of the shape's param to require gradient, method should return True
shape_param_key = 'box.vertex_positions_buf'
ek.set_requires_gradient(params[shape_param_key])
params.set_dirty(shape_param_key)
params.update()
assert scene.shapes_grad_enabled() == True
def test11_parameters_grad_enabled(variant_gpu_autodiff_rgb):
from mitsuba.core.xml import load_string
shape = load_string('''
<shape type="obj" version="2.0.0">
<string name="filename" value="resources/data/common/meshes/rectangle.obj"/>
</shape>
''')
assert shape.parameters_grad_enabled() == False
# Get the shape's parameters
params = traverse(shape)
# Only parameters of the shape should affect the result of that method
bsdf_param_key = 'bsdf.reflectance.value'
ek.set_requires_gradient(params[bsdf_param_key])
params.set_dirty(bsdf_param_key)
params.update()
assert shape.parameters_grad_enabled() == False
# When setting one of the shape's param to require gradient, method should return True
shape_param_key = 'vertex_positions_buf'
ek.set_requires_gradient(params[shape_param_key])
params.set_dirty(shape_param_key)
params.update()
assert shape.parameters_grad_enabled() == True
def __init__(self, bsdf_map, mesh_map, bsdf_ad_keys, mesh_ad_keys,\
lr, beta_1=0.9, beta_2=0.999, epsilon=1e-8):
from enoki.cuda_autodiff import Float32 as Float
# Ensure that the JIT compiler does merge 'lr' into the PTX code
# (this would trigger a recompile every time it is changed)
self.lr = lr
self.lr_v = ek.detach(Float(lr, literal=False))
self.bsdf_map = bsdf_map
self.mesh_map = mesh_map
self.bsdf_ad_keys = bsdf_ad_keys
self.mesh_ad_keys = mesh_ad_keys
self.beta_1 = beta_1
self.beta_2 = beta_2
self.epsilon = epsilon
self.t = 0
self.state = {}
for k in bsdf_ad_keys:
ek.set_requires_gradient(bsdf_map[k].reflectance.data)
size = ek.slices(bsdf_map[k].reflectance.data)
self.state[k] = (ek.detach(
type(bsdf_map[k].reflectance.data).zero(size)),
ek.detach(
type(
bsdf_map[k].reflectance.data).zero(size)))
for k in mesh_ad_keys:
ek.set_requires_gradient(mesh_map[k].vertex_positions)
size = ek.slices(mesh_map[k].vertex_positions)
self.state[k] = (ek.detach(
type(mesh_map[k].vertex_positions).zero(size)),
ek.detach(
type(
mesh_map[k].vertex_positions).zero(size)))
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 forward(ctx, arg1, arg2):
ctx.in1 = ek.FloatD(arg1)
ctx.in2 = ek.FloatD(arg2)
ek.set_requires_gradient(ctx.in1, arg1.requires_grad)
ek.set_requires_gradient(ctx.in2, arg2.requires_grad)
ctx.out = ek.atan2(ctx.in1, ctx.in2)
out_torch = ctx.out.torch()
ek.cuda_malloc_trim()
return out_torch
def disable_gradients(self):
"""Temporarily disable the generation of gradients."""
for _, p in self.params.items():
ek.set_requires_gradient(p, False)
try:
yield
finally:
for _, p in self.params.items():
ek.set_requires_gradient(p, True)
def forward(ctx, scene, params, *args):
try:
assert len(args) % 2 == 0
args = dict(zip(args[0::2], args[1::2]))
spp = None
sensor_index = 0
unbiased = True
ctx.inputs = [None, None]
for k, v in args.items():
if k == 'spp':
spp = v
elif k == 'sensor_index':
sensor_index = v
elif k == 'unbiased':
unbiased = v
elif params is not None:
params[k] = type(params[k])(v)
ctx.inputs.append(None)
ctx.inputs.append(
params[k] if v.requires_grad else None)
ctx.inputs.append(None)
ctx.inputs.append(None)
if type(spp) is not tuple:
spp = (spp, spp)
result = None
if params is not None:
params.update()
if unbiased:
result = render(scene,
spp=spp[0],
sensor_index=sensor_index).torch()
for v in ctx.inputs:
if v is not None:
ek.set_requires_gradient(v)
ctx.output = render(scene,
spp=spp[1],
sensor_index=sensor_index)
if result is None:
result = ctx.output.torch()
ek.cuda_malloc_trim()
return result
except Exception as e:
print("render_torch(): critical exception during "
"forward pass: %s" % str(e))
raise e
def __init__(self, params, lr):
"""
Parameter ``params``:
dictionary ``(name: variable)`` of differentiable parameters to be
optimized.
Parameter ``lr``:
learning rate
"""
self.set_learning_rate(lr)
self.params = params
self.state = {}
for k, p in self.params.items():
ek.set_requires_gradient(p)
self._reset(k)
def get_diff_param(scene):
# Create a differentiable hyperparameter
diff_param = Float(0.0)
ek.set_requires_gradient(diff_param)
# Update vertices so that they depend on diff_param
params = traverse(scene)
t = Transform4f.translate(Vector3f(1.0) * diff_param)
vertex_positions = params['light_shape.vertex_positions']
vertex_positions_t = t.transform_point(vertex_positions)
params['light_shape.vertex_positions'] = vertex_positions_t
# Update the scene
params.update()
return diff_param
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 get_diff_param(scene):
# Create a differentiable hyperparameter
diff_param = mitsuba.core.Float(0.0)
ek.set_requires_gradient(diff_param)
# Update vertices so that they depend on diff_param
properties = traverse(scene)
t = mitsuba.core.Transform4f.translate(mitsuba.core.Vector3f(1.0, 0.0, 0.0) * diff_param)
vertex_positions = properties['object.vertex_positions']
vertex_positions_t = t.transform_point(vertex_positions)
properties['object.vertex_positions'] = vertex_positions_t
# Update the scene
properties.update()
return diff_param
def __init__(self, params, lr):
"""
Parameter ``params``:
dictionary ``(name: variable)`` of differentiable parameters to be
optimized.
Parameter ``lr``:
learning rate
"""
self.set_learning_rate(lr)
self.params = params
if not params.all_differentiable():
raise Exception('Optimizer.__init__(): all parameters should '
'be differentiable!')
self.state = {}
for k, p in self.params.items():
ek.set_requires_gradient(p)
self._reset(k)
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 test04_differentiable_surface_interaction_ray_forward(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)
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 cylinder section, si.uv.x move by 1 / 2pi
si = pi.compute_surface_interaction(ray)
ek.forward(ray.o.x)
assert ek.allclose(ek.gradient(si.uv), [1 / (2 * ek.pi), 0])
# If the ray origin is shifted along the cylinder length, si.uv.y move by 1
si = pi.compute_surface_interaction(ray)
ek.forward(ray.o.z)
assert ek.allclose(ek.gradient(si.uv), [0, 1])
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 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 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 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 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)
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)
from mitsuba.core import Thread, Float
from mitsuba.core.xml import load_file
from mitsuba.python.util import traverse
from mitsuba.python.autodiff import render, write_bitmap
# 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()
