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 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 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 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 scatter_add_(self, target, index, mask):
assert target.Depth == 1
sr = max(len(self), len(index), len(mask))
for i in range(sr):
_ek.scatter_add(target, self[i], index[i], mask[i])
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