def test51_scatter_reduce_fwd_eager(m):
with EagerMode():
for i in range(3):
idx1 = ek.arange(m.UInt, 5)
idx2 = ek.arange(m.UInt, 4) + 3
x = ek.linspace(m.Float, 0, 1, 5)
y = ek.linspace(m.Float, 1, 2, 4)
buf = ek.zero(m.Float, 10)
if i % 2 == 0:
ek.enable_grad(buf)
ek.set_grad(buf, 1)
if i // 2 == 0:
ek.enable_grad(x, y)
ek.set_grad(x, 1)
ek.set_grad(y, 1)
x.label = "x"
y.label = "y"
buf.label = "buf"
buf2 = m.Float(buf)
ek.scatter_reduce(ek.ReduceOp.Add, buf2, x, idx1)
ek.scatter_reduce(ek.ReduceOp.Add, buf2, y, idx2)
s = ek.dot_async(buf2, buf2)
# Verified against Mathematica
assert ek.allclose(ek.detach(s), 15.5972)
assert ek.allclose(ek.grad(s), (25.1667 if i // 2 == 0 else 0) +
(17 if i % 2 == 0 else 0))
def test16_custom(cname):
t = get_class(cname)
v1 = ek.zero(t, 100)
v2 = ek.empty(t, 100)
assert len(v1.state) == 100
assert len(v2.inc) == 100
v2.state = v1.state
v1.state = ek.arange(type(v1.state), 100)
v3 = ek.select(v1.state < 10, v1, v2)
assert v3.state[3] == 3
assert v3.state[11] == 0
assert ek.width(v3) == 100
v4 = ek.zero(t, 1)
ek.schedule(v4)
ek.resize(v4, 200)
assert ek.width(v4) == 200
assert ek.width(v3) == 100
v4 = ek.zero(t, 1)
ek.resize(v4, 200)
assert ek.width(v4) == 200
index = ek.arange(type(v1.state), 100)
ek.scatter(v4, v1, index)
v5 = ek.gather(t, v4, index)
ek.eval(v5)
assert v5.state == v1.state and v5.inc == v1.inc
def test_roundtrip_dlpack_all(package):
Float, Array3f = package.Float, package.Array3f
prepare(package)
a1 = Array3f(
ek.arange(Float, 10),
ek.arange(Float, 10) + 100,
ek.arange(Float, 10) + 1000,
)
a2 = a1.dlpack()
a3 = Array3f(a2)
assert a1 == a3
assert a1.x == Float(a1.x.dlpack())
def test_roundtrip_numpy_all(package):
pytest.importorskip("numpy")
Float, Array3f = package.Float, package.Array3f
prepare(package)
a1 = Array3f(
ek.arange(Float, 10),
ek.arange(Float, 10) + 100,
ek.arange(Float, 10) + 1000,
)
a2 = a1.numpy()
a3 = Array3f(a2)
assert a1 == a3
assert a1.x == Float(a1.x.numpy())
def test_roundtrip_pytorch_llvm():
pytest.importorskip("torch")
prepare(enoki.llvm.ad)
from enoki.llvm.ad import Float, Array3f
a1 = Array3f(
ek.arange(Float, 10),
ek.arange(Float, 10) + 100,
ek.arange(Float, 10) + 1000,
)
a2 = a1.torch()
a3 = Array3f(a2)
assert a1 == a3
assert a1.x == Float(a1.x.torch())
def test_roundtrip_pytorch_jax():
pytest.importorskip("jax")
prepare(enoki.cuda.ad)
from enoki.cuda.ad import Float, Array3f
a1 = Array3f(
ek.arange(Float, 10),
ek.arange(Float, 10) + 100,
ek.arange(Float, 10) + 1000,
)
a2 = a1.jax()
a3 = Array3f(a2)
assert a1 == a3
assert a1.x == Float(a1.x.jax())
def test20_scatter_add_rev(m):
for i in range(3):
idx1 = ek.arange(m.UInt, 5)
idx2 = ek.arange(m.UInt, 4) + 3
x = ek.linspace(m.Float, 0, 1, 5)
y = ek.linspace(m.Float, 1, 2, 4)
buf = ek.zero(m.Float, 10)
if i % 2 == 0:
ek.enable_grad(buf)
if i // 2 == 0:
ek.enable_grad(x, y)
x.label = "x"
y.label = "y"
buf.label = "buf"
buf2 = m.Float(buf)
ek.scatter_add(buf2, x, idx1)
ek.scatter_add(buf2, y, idx2)
ref_buf = m.Float(0.0000, 0.2500, 0.5000, 1.7500, 2.3333, 1.6667,
2.0000, 0.0000, 0.0000, 0.0000)
assert ek.allclose(ref_buf, buf2, atol=1e-4)
assert ek.allclose(ref_buf, buf, atol=1e-4)
s = ek.dot_async(buf2, buf2)
print(ek.graphviz_str(s))
ek.backward(s)
ref_x = m.Float(0.0000, 0.5000, 1.0000, 3.5000, 4.6667)
ref_y = m.Float(3.5000, 4.6667, 3.3333, 4.0000)
if i // 2 == 0:
assert ek.allclose(ek.grad(y), ek.detach(ref_y), atol=1e-4)
assert ek.allclose(ek.grad(x), ek.detach(ref_x), atol=1e-4)
else:
assert ek.grad(x) == 0
assert ek.grad(y) == 0
if i % 2 == 0:
assert ek.allclose(ek.grad(buf), ek.detach(ref_buf) * 2, atol=1e-4)
else:
assert ek.grad(buf) == 0
def test22_scatter_fwd(m):
x = m.Float(4.0)
ek.enable_grad(x)
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)
ek.forward(x, retain_graph=True)
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 eval(self, pos, vel):
pos, vel = m.Array2f(pos), m.Array2f(vel)
# Run for 100 iterations
it, max_it = m.UInt32(0), 100
# Allocate scratch space
n = max(ek.width(pos), ek.width(vel))
self.temp_pos = ek.empty(m.Array2f, n * max_it)
self.temp_vel = ek.empty(m.Array2f, n * max_it)
loop = m.Loop(pos, vel, it)
while loop.cond(it < max_it):
# Store current loop variables
index = it * n + ek.arange(m.UInt32, n)
ek.scatter(self.temp_pos, pos, index)
ek.scatter(self.temp_vel, vel, index)
# Update loop variables
pos_out, vel_out = self.timestep(pos, vel)
pos.assign(pos_out)
vel.assign(vel_out)
it += 1
# Ensure output and temp. arrays are evaluated at this point
ek.eval(pos, vel)
return pos, vel
def test02_multiple_values(pkg, variant):
p = get_class(pkg)
i = ek.arange(p.Int, 0, 10)
v = ek.zero(p.Array3f, 10)
if variant == 1:
v.y = p.Float(0)
loop = p.Loop(i, v)
while loop.cond(i < 5):
i.assign(i + 1)
f = p.Float(i)
v.x += f
v.y += 2 * f
v.z += 4 * f
if variant == 0:
ek.eval(i, v)
else:
ek.eval(i)
ek.eval(v.x)
ek.eval(v.y)
ek.eval(v.z)
assert i == p.Int(5, 5, 5, 5, 5, 5, 6, 7, 8, 9)
assert v.y == p.Int(30, 28, 24, 18, 10, 0, 0, 0, 0, 0)
def backward(self):
grad_pos, grad_vel = self.grad_out()
# Run for 100 iterations
it = m.UInt32(100)
loop = m.Loop(it, grad_pos, grad_vel)
n = ek.width(grad_pos)
while loop.cond(it > 0):
# Retrieve loop variables, reverse chronological order
it -= 1
index = it * n + ek.arange(m.UInt32, n)
pos = ek.gather(m.Array2f, self.temp_pos, index)
vel = ek.gather(m.Array2f, self.temp_vel, index)
# Differentiate loop body in reverse mode
ek.enable_grad(pos, vel)
pos_out, vel_out = self.timestep(pos, vel)
ek.set_grad(pos_out, grad_pos)
ek.set_grad(vel_out, grad_vel)
ek.enqueue(pos_out, vel_out)
ek.traverse(m.Float, reverse=True)
# Update loop variables
grad_pos.assign(ek.grad(pos))
grad_vel.assign(ek.grad(vel))
self.set_grad_in('pos', grad_pos)
self.set_grad_in('vel', grad_vel)
def test01_ctr(pkg):
p = get_class(pkg)
i = ek.arange(p.Int, 0, 10)
loop = p.Loop(i)
while loop.cond(i < 5):
i += 1
assert i == p.Int(5, 5, 5, 5, 5, 5, 6, 7, 8, 9)
def test14_roundtrip(package, dest):
pytest.importorskip(dest)
Float, Array3f, Array4f = package.Float, package.Array3f, package.Array4f
Matrix3f, Matrix4f, Matrix44f = package.Matrix3f, package.Matrix4f, package.Matrix44f
prepare(package)
def to_dest(a):
if dest == 'numpy':
return a.numpy()
if dest == 'torch':
return a.torch()
if dest == 'jax':
return a.jax()
v = Array3f(
(1.0 + ek.arange(Float, 5)),
(1.0 + ek.arange(Float, 5)) * 2,
(1.0 + ek.arange(Float, 5)) * 3,
)
assert(v == Array3f(to_dest(v)))
m = Matrix3f([
Array3f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]),
Array3f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]) * 10,
Array3f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]) * 100
])
assert(m ==Matrix3f(to_dest(m)))
m = Matrix4f([
Array4f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [4.0, 4.5]]),
Array4f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [4.0, 4.5]]) * 10,
Array4f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [4.0, 4.5]]) * 100,
Array4f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [4.0, 4.5]]) * 1000
])
assert(m == Matrix4f(to_dest(m)))
m = Matrix44f([
[1.0, 2.0, 3.0, 4.0],
[1.0, 2.0, 3.0, 4.0],
[1.0, 2.0, 3.0, 4.0],
[1.0, 2.0, 3.0, 4.0]
])
assert(m == Matrix44f(to_dest(m)))
def test02_radical_inverse_vectorized(variant_scalar_rgb):
from mitsuba.core import RadicalInverse
v = RadicalInverse()
for index, prime in enumerate(gen_primes()):
if index >= 1024:
break
result = v.eval(index, ek.arange(10, dtype=ek.uint64))
for i in range(len(result)):
assert ek.abs(r_inv(prime, i) - result[i]) < 1e-7
def test01_ctr(pkg):
p = get_class(pkg)
i = ek.arange(p.Int, 0, 10)
loop = p.Loop("MyLoop", lambda: i)
while loop(i < 5):
i = i + 1
assert i == p.Int(5, 5, 5, 5, 5, 5, 6, 7, 8, 9)
def test09_loop_simplification(pkg):
# Test a regression where most loop variables are optimized away
p = get_class(pkg)
res = p.Float(0.0)
active = p.Bool(True)
depth = p.UInt32(1000)
loop = p.Loop("TestLoop")
loop.put(lambda: (active, depth, res))
loop.init()
while loop(active):
res += ek.arange(p.Float, 10)
depth -= 1
active &= (depth > 0)
del loop
del active
del depth
gc.collect()
gc.collect()
assert res == ek.arange(p.Float, 10) * 1000
def test07_gather_ravel_unravel(pkg):
get_class(pkg.__name__)
str_1 = '[[0.0, 1.0, 2.0],\n [3.0, 4.0, 5.0],\n [6.0, 7.0, 8.0],\n' \
' [9.0, 10.0, 11.0]]'
str_2 = '[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0]'
a = ek.arange(pkg.Float, 100)
b = ek.gather(pkg.Array3f, a, pkg.UInt(0, 1, 2, 3))
assert repr(b) == str_1
c = ek.ravel(b)
assert repr(c) == str_2
d = ek.unravel(pkg.Array3f, c)
assert repr(d) == str_1
def test03_sample_eval_pdf(variant_scalar_rgb, interaction):
from mitsuba.core.math import InvPi
from mitsuba.core.warp import square_to_uniform_sphere
from mitsuba.render import BSDFContext
from mitsuba.core.xml import load_string
bsdf = load_string("""<bsdf version="2.0.0" type="twosided">
<bsdf type="diffuse">
<rgb name="reflectance" value="0.1, 0.1, 0.1"/>
</bsdf>
<bsdf type="diffuse">
<rgb name="reflectance" value="0.9, 0.9, 0.9"/>
</bsdf>
</bsdf>""")
n = 5
ctx = BSDFContext()
for u in ek.arange(UInt32, n):
for v in ek.arange(UInt32, n):
interaction.wi = square_to_uniform_sphere([u / float(n-1),
v / float(n-1)])
up = ek.dot(interaction.wi, [0, 0, 1]) > 0
for x in ek.arange(UInt32, n):
for y in ek.arange(UInt32, n):
sample = [x / float(n-1), y / float(n-1)]
(bs, s_value) = bsdf.sample(ctx, interaction, 0.5, sample)
if ek.any(s_value > 0):
# Multiply by square_to_cosine_hemisphere_theta
s_value *= bs.wo[2] * InvPi
if not up:
s_value *= -1
e_value = bsdf.eval(ctx, interaction, bs.wo)
p_pdf = bsdf.pdf(ctx, interaction, bs.wo)
assert ek.allclose(s_value, e_value, atol=1e-2)
assert ek.allclose(bs.pdf, p_pdf)
assert not ek.any(ek.isnan(e_value) | ek.isnan(s_value))
def test21_scatter_add_fwd(m):
for i in range(3):
idx1 = ek.arange(m.UInt, 5)
idx2 = ek.arange(m.UInt, 4) + 3
x = ek.linspace(m.Float, 0, 1, 5)
y = ek.linspace(m.Float, 1, 2, 4)
buf = ek.zero(m.Float, 10)
if i % 2 == 0:
ek.enable_grad(buf)
ek.set_grad(buf, 1)
if i // 2 == 0:
ek.enable_grad(x, y)
ek.set_grad(x, 1)
ek.set_grad(y, 1)
x.label = "x"
y.label = "y"
buf.label = "buf"
buf2 = m.Float(buf)
ek.scatter_add(buf2, x, idx1)
ek.scatter_add(buf2, y, idx2)
s = ek.dot_async(buf2, buf2)
if i % 2 == 0:
ek.enqueue(buf)
if i // 2 == 0:
ek.enqueue(x, y)
ek.traverse(m.Float, reverse=False)
# Verified against Mathematica
assert ek.allclose(ek.detach(s), 15.5972)
assert ek.allclose(ek.grad(s), (25.1667 if i // 2 == 0 else 0) +
(17 if i % 2 == 0 else 0))
def test02_radical_inverse_vectorized(variant_scalar_rgb):
from mitsuba.core import RadicalInverse
try:
from mitsuba.packet_rgb.core.qmc import RadicalInverseP
except ImportError:
pytest.skip("packet_rgb mode not enabled")
v = RadicalInverse()
v_p = RadicalInverseP()
for index in range(1024):
result = v_p.eval_scrambled(index, ek.arange(10, dtype=ek.uint64))
for i in range(len(result)):
assert ek.abs(v.eval_scrambled(index, i) - result[i]) < 1e-7
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 generate_masks(owners, scene):
#Generate rays that are used for both the binary mask and the distance mask
sensor = scene.sensors()[0]
film = sensor.film()
sampler = sensor.sampler()
film_size = film.crop_size()
spp = 3
# Seed the sampler
total_sample_count = ek.hprod(film_size) * spp
if sampler.wavefront_size() != total_sample_count:
sampler.seed(UInt64.arange(total_sample_count))
# Enumerate discrete sample & pixel indices, and uniformly sample
# positions within each pixel.
pos = ek.arange(UInt32, total_sample_count)
pos //= spp
scale = Vector2f(1.0 / film_size[0], 1.0 / film_size[1])
pos = Vector2f(Float(pos % int(film_size[0])),
Float(pos // int(film_size[0])))
pos += sampler.next_2d()
# Sample rays starting from the camera sensor
rays, weights = sensor.sample_ray_differential(time=0,
sample1=sampler.next_1d(),
sample2=pos * scale,
sample3=0)
#Masks dictionary containas a mask for each params
masks = dict()
#For each parameter owner, we generate a mask
for key in owners:
scene_dict = dict()
scene_dict["type"] = "scene"
scene_dict["camera"] = scene.sensors()[0]
for item in owners[key]:
scene_dict[item.id()] = item
#dummy_scene contains only camera and objects for associated with this owner
dummy_scene = load_dict(scene_dict)
surface_interaction = dummy_scene.ray_intersect(rays)
result = surface_interaction.t + 1.0
result[~(surface_interaction.t is Infinity)] = 1
result[~surface_interaction.is_valid()] = 0
masks[key] = result
return masks
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 load(t, i, offset):
size = shape[-1 - i]
stride = strides[-1 - i]
if i == ndim - 1:
if type(offset) is int and stride == 1:
return data
else:
i = enoki.arange(enoki.int32_array_t(t), size)
return t.gather_(data, offset + stride * i, True, False)
else:
result = t()
for j in range(size):
result[j] = load(t.Value, i + 1, offset + stride * j)
return result
def test08_divmod(cname):
t = get_class(cname)
index = ek.arange(t, 10000000)
index[index < len(index) // 2] = -index
index *= 256203161
for i in range(1, 100):
assert index // i == index // ek.opaque(t, i, 1)
assert index % i == index % ek.opaque(t, i, 1)
if t.IsSigned:
for i in range(1, 100):
assert index // -i == index // ek.opaque(t, -i, 1)
assert index % -i == index % ek.opaque(t, -i, 1)
def sample_ray(
self,
time,
sample1, # wavelength
sample2, # pos
sample3, # dir
active):
wavelengths, spec_weight = self.m_intensity.sample(
SurfaceInteraction3f(), ek.arange(sample1), active)
trafo = self.m_world_transform.eval(ref.time)
ray = Ray3f(trafo * Point3f(0), warp.square_to_uniform_sphere(sample3),
time, wavelengths)
# print(spec_weight.class_().name())
return (ray, spec_weight * 4.0 * Pi)
def test18_slice():
Float = get_class('enoki.llvm.Float')
Array2f = get_class('enoki.llvm.Array2f')
a = ek.arange(Float, 10)
for i in range(10):
assert ek.slice(a, i) == i
a2 = Array2f(a, a * 10)
for i in range(10):
assert ek.slice(a2, i)[0] == i
assert ek.slice(a2, i)[1] == i * 10
a3 = (a, a * 10)
for i in range(10):
assert ek.slice(a3, i)[0] == i
assert ek.slice(a3, i)[1] == i * 10
def sample_ray(
self,
time,
sample1, # wavelength
sample2, # pos
sample3, # dir
active):
ps = self.m_shape.sample_position(time, sample2, active)
local = warp.square_to_cosine_hemisphere(sample3)
si = SurfaceInteraction3f(ps, 0)
wavelengths, spec_weight = self.m_radiance.sample(
si, ek.arange(sample1), active)
ray = Ray3f(ps.p, Frame3f(ps.n).to_world(local), time, wavelengths)
return (ray, spec_weight * self.m_area_times_pi)
def test32_atan2(m):
x = ek.linspace(m.Float, -.8, .8, 10)
y = m.Float(ek.arange(m.Int, 10) & 1) * 1 - .5
ek.enable_grad(x, y)
z = ek.atan2(y, x)
ek.backward(z)
assert ek.allclose(
z,
m.Float(-2.58299, 2.46468, -2.29744, 2.06075, -1.74674, 1.39486,
-1.08084, 0.844154, -0.676915, 0.558599))
assert ek.allclose(
ek.grad(x),
m.Float(0.561798, -0.784732, 1.11724, -1.55709, 1.93873, -1.93873,
1.55709, -1.11724, 0.784732, -0.561798))
assert ek.allclose(
ek.grad(y),
m.Float(-0.898876, -0.976555, -0.993103, -0.83045, -0.344663, 0.344663,
0.83045, 0.993103, 0.976555, 0.898876))
def test06_test_collatz(pkg, variant):
p = get_class(pkg)
def collatz(value: p.Int):
counter = p.Int(0)
loop = p.Loop(value, counter)
while (loop.cond(ek.neq(value, 1))):
is_even = ek.eq(value & 1, 0)
value.assign(ek.select(is_even, value // 2, 3 * value + 1))
counter += 1
return value, counter
value, ctr = collatz(ek.arange(p.Int, 1, 11))
if variant == 0:
ek.eval(value, ctr)
elif variant == 1:
ek.eval(value)
ek.eval(ctr)
elif variant == 2:
ek.eval(ctr)
ek.eval(value)
assert value == p.Int([1] * 10)
assert ctr == p.Int([0, 1, 7, 2, 5, 8, 16, 3, 19, 6])
