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 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 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 test17_opaque():
Array3f = get_class('enoki.llvm.Array3f')
v = ek.opaque(Array3f, 4.0)
assert ek.width(v) == 1
assert ek.allclose(v, 4.0)
for i in range(len(v)):
assert not v[i].is_literal_()
assert v[i].is_evaluated_()
v = ek.opaque(Array3f, 4.0, 1)
assert ek.width(v) == 1
assert ek.allclose(v, 4.0)
v = ek.opaque(Array3f, 4.0, 10)
assert ek.width(v) == 10
assert ek.allclose(v, 4.0)
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