def test_compute(): def transfer(k): return 2.0 engine = ParticleMeshEngine(pm) code = CodeSegment(engine) code.r2c(real='r', complex='c') code.transfer(complex='c', tf=transfer) code.c2r(complex='c', real='r') code.to_scalar(x='r', y='sum') field = pm.generate_whitenoise(seed=1234).c2r() norm = code.compute('sum', init={'r': field}) assert_allclose(norm, field.cnorm() * 4)
def test_transfer_real(): def transfer(k): return k[0] field = pm.generate_whitenoise(seed=1234, mode='real') engine = ParticleMeshEngine(pm) code = CodeSegment(engine) code.r2c(complex='c', real='r') code.transfer(complex='c', tf=transfer) code.c2r(complex='c', real='r') check_grad(code, 'r', 'r', init={'r': field}, eps=1e-4, rtol=1e-8)
def test_transfer_real(): def transfer(k): return k[0] field = pm.generate_whitenoise(seed=1234, type='real') engine = ParticleMeshEngine(pm) code = CodeSegment(engine) code.r2c(complex='c', real='r') code.transfer(complex='c', tf=transfer) code.c2r(complex='c', real='r') check_grad(code, 'r', 'r', init={'r': field}, eps=1e-4, rtol=1e-8)
def test_vjp(): def transfer(k): k2 = sum(ki **2 for ki in k) k2[k2 == 0] = 1.0 # return 1 / k2 return 2.0 engine = ParticleMeshEngine(pm) code = CodeSegment(engine) code.r2c(real='r', complex='c') code.transfer(complex='c', tf=transfer) code.c2r(complex='c', real='r') code.multiply(x1='r', x2=Literal(0.1), y='r') code.to_scalar(x='r', y='sum') field = pm.generate_whitenoise(seed=1234).c2r() norm, tape = code.compute('sum', init={'r': field}, return_tape=True) assert_allclose(norm, field.cnorm() * 4 * 0.1 ** 2) vjp = tape.get_vjp() _r = vjp.compute('_r', init={'_sum': 1.0}) assert_allclose(_r, field * 4 * 2 * 0.1 * 0.1)