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 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 sample_functor(sample, *args):
n = ek.slices(sample)
plugin = instantiate(args)
mi, ctx = make_context(n)
wo, pdf = plugin.sample(ctx, mi, [sample[0], sample[1]])
w = Float.full(1.0, ek.slices(pdf))
w[ek.eq(pdf, 0)] = 0
return wo, w
def sample_functor(sample, *args):
n = ek.slices(sample)
plugin = instantiate(args)
(si, ctx) = make_context(n)
bs, weight = plugin.sample(ctx, si, sample[0], [sample[1], sample[2]])
w = Float.full(1.0, ek.slices(weight))
w[ek.all(ek.eq(weight, 0))] = 0
return bs.wo, w
def _reset(self, key):
""" Zero-initializes the internal state associated with a parameter """
if self.momentum == 0:
return
p = self.params[key]
size = ek.slices(p)
self.state[key] = ek.detach(type(p).zero(size))
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))
def test05_load_simple_mesh(variant_scalar_rgb):
from mitsuba.core.xml import load_string
"""Tests the OBJ and PLY loaders on a simple example."""
for mesh_format in ["obj", "ply"]:
shape = load_string("""
<shape type="{0}" version="2.0.0">
<string name="filename" value="resources/data/tests/{0}/cbox_smallbox.{0}"/>
</shape>
""".format(mesh_format))
positions = shape.vertex_positions_buffer()
faces = shape.faces_buffer()
assert shape.has_vertex_normals()
assert ek.slices(positions) == 72
assert ek.slices(faces) == 36
assert ek.allclose(faces[6:9], [4, 5, 6])
assert ek.allclose(positions[:5], [130, 165, 65, 82, 165])
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 test02_ply_triangle(variant_scalar_rgb):
from mitsuba.core import UInt32
from mitsuba.core.xml import load_string
m = load_string("""
<shape type="ply" version="0.5.0">
<string name="filename" value="data/triangle.ply"/>
<boolean name="face_normals" value="true"/>
</shape>
""")
positions = m.vertex_positions_buffer()
faces = m.faces_buffer()
assert not m.has_vertex_normals()
assert ek.slices(positions) == 9
assert ek.allclose(positions[0:3], [0, 0, 0])
assert ek.allclose(positions[3:6], [0, 0, 1])
assert ek.allclose(positions[6:9], [0, 1, 0])
assert ek.slices(faces) == 3
assert faces[0] == UInt32(0)
assert faces[1] == UInt32(1)
assert faces[2] == UInt32(2)
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 unravel(source, target, dim=3):
idx = UInt32.arange(ek.slices(source))
for i in range(dim):
ek.scatter(target, source[i], dim * idx + i)
if not os.path.isdir(output_path):
os.makedirs(output_path)
# Generate the scene and image with a plain glass panel
scene = make_scene(path_str, spp_ref, width, height)
image_plain = render(scene)
write_bitmap(output_path + "out_plain.png", image_plain, (width, height))
print("Writing " + "out_plain.png")
params = traverse(scene)
print(params)
positions_buf = params['grid_mesh.vertex_positions_buf']
positions_initial = ravel(positions_buf)
normals_initial = ravel(params['grid_mesh.vertex_normals_buf'])
vertex_count = ek.slices(positions_initial)
filename = 'scene/diffuser_surface_1.jpg'
Thread.thread().file_resolver().append(os.path.dirname(filename))
# Create a texture with the reference displacement map
disp_tex_1 = xml.load_dict({
"type" : "bitmap",
"filename": "diffuser_surface_1.jpg",
"to_uv" : ScalarTransform4f.scale([1, -1, 1]) # texture is upside-down
}).expand()[0]
# Create a texture with another displacement map
disp_tex_2 = xml.load_dict({
"type" : "bitmap",
"filename": "diffuser_surface_2.jpg",
def pdf_functor(wo, *args):
n = ek.slices(wo)
plugin = instantiate(args)
(si, ctx) = make_context(n)
return plugin.pdf(ctx, si, wo)
def pdf_functor(w, *args):
plugin = instantiate(args)
si = SurfaceInteraction3f.zero(ek.slices(w))
si.wavelengths = w
return plugin.pdf(si)[0]
def sample_functor(sample, *args):
plugin = instantiate(args)
si = SurfaceInteraction3f.zero(ek.slices(sample))
wavelength, weight = plugin.sample(si, sample_shifted(sample[0]))
return Vector1f(wavelength[0])
def pdf_functor(wo, *args):
n = ek.slices(wo)
plugin = instantiate(args)
mi, ctx = make_context(n)
return plugin.eval(ctx, mi, wo)
def _reset(self, key):
""" Zero-initializes the internal state associated with a parameter """
p = self.params[key]
size = ek.slices(p)
self.state[key] = (ek.detach(type(p).zero(size)),
ek.detach(type(p).zero(size)))
if not os.path.isdir(output_path):
os.makedirs(output_path)
# Generate the scene and image with a plain glass panel
scene = make_scene(path_str, spp_ref, width, height)
image_plain = render(scene)
write_bitmap(output_path + "out_plain.png", image_plain, (width, height))
print("Writing " + "out_plain.png")
params = traverse(scene)
print(params)
positions_buf = params['grid_mesh.vertex_positions_buf']
positions_initial = ravel(positions_buf)
normals_initial = ravel(params['grid_mesh.vertex_normals_buf'])
vertex_count = ek.slices(positions_initial)
filename = 'scene/diffuser_surface_1.jpg'
Thread.thread().file_resolver().append(os.path.dirname(filename))
# Create a texture with the reference displacement map
disp_tex_1 = xml.load_dict({
"type": "bitmap",
"filename": "diffuser_surface_1.jpg",
"to_uv": ScalarTransform4f.scale([1, -1, 1]) # texture is upside-down
}).expand()[0]
# Create a texture with another displacement map
disp_tex_2 = xml.load_dict({
"type": "bitmap",
"filename": "diffuser_surface_2.jpg",
