def check_uniform_wavefront_sampler(sampler, res=16, atol=0.5):
from mitsuba.core import Float, UInt32, UInt64, Vector2u
sample_count = sampler.sample_count()
sampler.set_samples_per_wavefront(sample_count)
sampler.seed(0, sample_count)
hist_1d = ek.zero(UInt32, res)
hist_2d = ek.zero(UInt32, res * res)
v_1d = ek.clamp(sampler.next_1d() * res, 0, res)
ek.scatter_add(
target=hist_1d,
index=UInt32(v_1d),
source=UInt32(1.0)
)
v_2d = Vector2u(ek.clamp(sampler.next_2d() * res, 0, res))
ek.scatter_add(
target=hist_2d,
index=UInt32(v_2d.x * res + v_2d.y),
source=UInt32(1.0)
)
assert ek.allclose(Float(hist_1d), float(sample_count) / res, atol=atol)
assert ek.allclose(Float(hist_2d), float(sample_count) / (res * res), atol=atol)
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 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 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 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 test05_side_effect_noloop(pkg):
p = get_class(pkg)
i = ek.zero(p.Int, 10)
j = ek.zero(p.Int, 10)
buf = ek.zero(p.Float, 10)
ek.disable_flag(ek.JitFlag.RecordLoops)
loop = p.Loop(i, j)
while loop.cond(i < 10):
j += i
i += 1
ek.scatter_add(target=buf, value=p.Float(i), index=0, mask=loop.mask())
assert i == p.Int([10] * 10)
assert buf == p.Float(550, *([0] * 9))
assert j == p.Int([45] * 10)
def test04_side_effect(pkg):
p = get_class(pkg)
i = ek.zero(p.Int, 10)
j = ek.zero(p.Int, 10)
buf = ek.zero(p.Float, 10)
loop = p.Loop(i, j)
while loop.cond(i < 10):
j += i
i += 1
ek.scatter_add(target=buf, value=p.Float(i), index=0)
ek.eval(i, j)
assert i == p.Int([10] * 10)
assert buf == p.Float(550, *([0] * 9))
assert j == p.Int([45] * 10)
def fallof(self, cosTheta):
# delta = (cosTheta - self.cosTotalWidth)/(self.cosFallOffStart - self.cosTotalWidth)
delta = (self.cutOffAngle -
ek.acos(cosTheta)) * self.m_invTransitionWidth
delta = ek.select(cosTheta > self.cosFallOffStart,
ek.full(type(delta), 1), delta)
delta = ek.select(cosTheta < self.cosTotalWidth, ek.zero(type(delta)),
delta)
return delta
def zero_(cls, size=1):
result = cls()
if cls.Size == Dynamic:
result.init_(size)
for i in range(size):
result.set_entry_(i, 0)
else:
for i in range(cls.Size):
result.set_entry_(i, _ek.zero(cls.Value, size))
return result
def test_58_diffloop_masking_rev(m, no_record):
fo = ek.zero(m.Float, 10)
fi = m.Float(1, 2)
i = m.UInt32(0, 5)
ek.enable_grad(fi)
loop = m.Loop("MyLoop", lambda: i)
while loop(i < 5):
ek.scatter_reduce(ek.ReduceOp.Add, fo, fi, i)
i += 1
ek.backward(fo)
assert ek.grad(fi) == m.Float(5, 0)
def test05_side_effect_noloop(pkg):
p = get_class(pkg)
i = ek.zero(p.Int, 10)
j = ek.zero(p.Int, 10)
buf = ek.zero(p.Float, 10)
ek.set_flag(ek.JitFlag.LoopRecord, False)
loop = p.Loop("MyLoop", lambda: (i, j))
while loop(i < 10):
j += i
i += 1
ek.scatter_reduce(op=ek.ReduceOp.Add,
target=buf,
value=p.Float(i),
index=0)
assert i == p.Int([10] * 10)
assert buf == p.Float(550, *([0] * 9))
assert j == p.Int([45] * 10)
def test03_failures(pkg):
p = get_class(pkg)
i = p.Int()
v = ek.zero(p.Array3f, 10)
if 'ad' in pkg:
i = p.Int(0)
ek.enable_grad(v)
with pytest.raises(ek.Exception) as e:
p.Loop("MyLoop", lambda: (i, v))
assert 'one of the supplied loop variables is attached to the AD graph' in str(
e.value)
def test_zero_initialization(package):
Float, Array3f = package.Float, package.Array3f
prepare(package)
class MyStruct:
ENOKI_STRUCT = {'a': Array3f, 'b': Float}
def __init__(self):
self.a = Array3f()
self.b = Float()
# Custom zero initialize callback
def zero_(self, size):
self.a += 1
foo = ek.zero(MyStruct, 4)
assert ek.width(foo) == 4
assert foo.a == 1
assert foo.b == 0
foo = ek.zero(MyStruct, 1)
ek.resize(foo, 8)
assert ek.width(foo) == 8
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 diag(a):
if _ek.is_matrix_v(a):
result = a.Value()
for i in range(a.Size):
result[i] = a[i, i]
return result
elif _ek.is_static_array_v(a):
name = _ek.detail.array_name('Matrix', a.Type, (a.Size, *a.Shape),
a.IsScalar)
module = _modules.get(a.__module__)
cls = getattr(module, name)
result = _ek.zero(cls)
for i in range(a.Size):
result[i, i] = a[i]
return result
else:
raise Exception('Unsupported type!')
def test_ad_operations(package):
Float, Array3f = package.Float, package.Array3f
prepare(package)
class MyStruct:
ENOKI_STRUCT = {'a': Array3f, 'b': Float}
def __init__(self):
self.a = Array3f()
self.b = Float()
foo = ek.zero(MyStruct, 4)
assert not ek.grad_enabled(foo.a)
assert not ek.grad_enabled(foo.b)
assert not ek.grad_enabled(foo)
ek.enable_grad(foo)
assert ek.grad_enabled(foo.a)
assert ek.grad_enabled(foo.b)
assert ek.grad_enabled(foo)
foo_detached = ek.detach(foo)
assert not ek.grad_enabled(foo_detached.a)
assert not ek.grad_enabled(foo_detached.b)
assert not ek.grad_enabled(foo_detached)
x = Float(4.0)
ek.enable_grad(x)
foo.a += x
foo.b += x * x
ek.forward(x)
foo_grad = ek.grad(foo)
assert foo_grad.a == 1
assert foo_grad.b == 8
ek.set_grad(foo, 5.0)
foo_grad = ek.grad(foo)
assert foo_grad.a == 5.0
assert foo_grad.b == 5.0
ek.accum_grad(foo, 5.0)
foo_grad = ek.grad(foo)
assert foo_grad.a == 10.0
assert foo_grad.b == 10.0
def test03_failures(pkg):
p = get_class(pkg)
i = p.Int()
v = ek.zero(p.Array3f, 10)
with pytest.raises(ek.Exception) as e:
p.Loop(i, v)
assert 'Variables provided to enoki::Loop() must be fully initialized' in str(
e.value)
if 'ad' in pkg:
i = p.Int(0)
ek.enable_grad(v)
with pytest.raises(ek.Exception) as e:
p.Loop(i, v)
assert 'Symbolic loop encountered a differentiable array with enabled gradients! This is not supported.' in str(
e.value)
def sample_direction(self, ref, sample, active):
# as the shape init happens after the emitter init
if (not (hasattr(self, "m_shape"))):
self.set_shape_area()
ds = self.m_shape.sample_direction(ref, sample, active)
active &= (ek.dot(ds.d, ds.n) < 0) & (ek.neq(ds.pdf, 0))
si = SurfaceInteraction3f(ds, ref.wavelengths)
# spatially varying
cosTheta = -ek.dot(ds.d, ds.n)
fall = self.fallof(cosTheta)
spec = self.m_radiance.eval(si, active) * fall / ds.pdf
ds.object = 0 # HACK
return (ds, ek.select(active, spec, ek.zero(type(spec))))
def broadcast_(self, value):
if not self.IsSpecial:
for i in range(len(self)):
self.set_entry_(i, value)
elif self.IsComplex:
self.set_entry_(0, value)
self.set_entry_(1, 0)
elif self.IsQuaternion:
for i in range(3):
self.set_entry_(i, 0)
self.set_entry_(3, value)
elif self.IsMatrix:
t = self.Value
for i in range(len(self)):
c = _ek.zero(t)
c.set_entry_(i, t.Value(value))
self.set_entry_(i, c)
else:
raise Exception("broadcast_(): don't know how to handle this type!")
def sincos(x):
Float = type(x)
Int = ek.int_array_t(Float)
xa = ek.abs(x)
j = Int(xa * 1.2732395447351626862)
j = (j + Int(1)) & ~Int(1)
y = Float(j)
Shift = Float.Type.Size * 8 - 3
sign_sin = ek.reinterpret_array(Float, j << Shift) ^ x
sign_cos = ek.reinterpret_array(Float, (~(j - Int(2)) << Shift))
y = xa - y * 0.78515625 \
- y * 2.4187564849853515625e-4 \
- y * 3.77489497744594108e-8
z = y * y
z |= ek.eq(xa, ek.Infinity)
s = poly2(z, -1.6666654611e-1,
8.3321608736e-3,
-1.9515295891e-4) * z
c = poly2(z, 4.166664568298827e-2,
-1.388731625493765e-3,
2.443315711809948e-5) * z
s = ek.fmadd(s, y, y)
c = ek.fmadd(c, z, ek.fmadd(z, -0.5, 1))
polymask = ek.eq(j & Int(2), ek.zero(Int))
return (
ek.mulsign(ek.select(polymask, s, c), sign_sin),
ek.mulsign(ek.select(polymask, c, s), sign_cos)
)
def test13_slice_setitem():
a = ek.zero(ek.scalar.ArrayXf, 5)
a[2] = 1.0
assert ek.allclose(a, [0, 0, 1, 0, 0])
a[2:] = [2.0, 1.0, 1.0]
assert ek.allclose(a, [0, 0, 2, 1, 1])
a[:] = 0.0
assert ek.allclose(a, [0, 0, 0, 0, 0])
v = ek.scalar.Array3f(0)
v[2] = 1.0
assert ek.allclose(v, [0, 0, 1])
v[1:] = 2.0
assert ek.allclose(v, [0, 2, 2])
m = ek.scalar.Matrix3f(0)
m[1:, 1] = 1.0
assert ek.allclose(m, [[0, 0, 0], [0, 1, 0], [0, 1, 0]])
m[0, 1:] = 2.0
assert ek.allclose(m, [[0, 2, 2], [0, 1, 0], [0, 1, 0]])
m[1:, 1:] = 3.0
assert ek.allclose(m, [[0, 2, 2], [0, 3, 3], [0, 3, 3]])
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 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 test01_record_loop(pkg):
p = get_class(pkg)
for i in range(3):
ek.set_flag(ek.JitFlag.LoopRecord, not i == 0)
ek.set_flag(ek.JitFlag.LoopOptimize, i == 2)
for j in range(2):
x = ek.arange(p.Int, 0, 10)
y = ek.zero(p.Float, 1)
z = p.Float(1)
loop = p.Loop("MyLoop", lambda: (x, y, z))
while loop(x < 5):
y += p.Float(x)
x += 1
z = z + 1
if j == 0:
ek.schedule(x, y, z)
assert z == p.Int(6, 5, 4, 3, 2, 1, 1, 1, 1, 1)
assert y == p.Int(10, 10, 9, 7, 4, 0, 0, 0, 0, 0)
assert x == p.Int(5, 5, 5, 5, 5, 5, 6, 7, 8, 9)
def map_backward(self, p):
from mitsuba.core import Vector2f, Float
return Vector2f(p.x, ek.zero(Float, len(p.x)))
def tabulate_histogram(self):
"""
Invoke the provided sampling strategy many times and generate a
histogram in the parameter domain. If ``sample_func`` returns a tuple
``(positions, weights)`` instead of just positions, the samples are
considered to be weighted.
"""
# Generate a table of uniform variates
from mitsuba.core import Float, Vector2f, Vector2u, Float32, \
UInt64, PCG32
rng = PCG32(initseq=ek.arange(UInt64, self.sample_count))
samples_in = getattr(mitsuba.core, 'Vector%if' % self.sample_dim)()
for i in range(self.sample_dim):
samples_in[i] = rng.next_float32() if Float is Float32 \
else rng.next_float64()
self.pdf_start = time.time()
# Invoke sampling strategy
samples_out = self.sample_func(samples_in)
if type(samples_out) is tuple:
weights_out = samples_out[1]
samples_out = samples_out[0]
else:
weights_out = Float(1.0)
# Map samples into the parameter domain
xy = self.domain.map_backward(samples_out)
# Sanity check
eps = self.bounds.extents() * 1e-4
in_domain = ek.all((xy >= self.bounds.min - eps)
& (xy <= self.bounds.max + eps))
if not ek.all(in_domain):
self._log('Encountered samples outside of the specified '
'domain: %s' % str(ek.compress(xy, ~in_domain)))
self.fail = True
# Normalize position values
xy = (xy - self.bounds.min) / self.bounds.extents()
xy = Vector2u(
ek.clamp(xy * Vector2f(self.res), 0, Vector2f(self.res - 1)))
# Compute a histogram of the positions in the parameter domain
self.histogram = ek.zero(Float, ek.hprod(self.res))
ek.scatter_add(target=self.histogram,
index=xy.x + xy.y * self.res.x,
source=weights_out)
self.pdf_end = time.time()
histogram_min = ek.hmin(self.histogram)
if not histogram_min >= 0:
self._log('Encountered a cell with negative sample '
'weights: %f' % histogram_min)
self.fail = True
self.histogram_sum = ek.hsum(self.histogram) / self.sample_count
if self.histogram_sum > 1.1:
self._log('Sample weights add up to a value greater '
'than 1.0: %f' % self.histogram_sum)
self.fail = True
