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