def test_symbol_add():
with Builder() as m:
a, b = m.input('a', 'b')
m.output(c=(a + b) * b / b**2)
c = m.compute('c', init=dict(a=2, b=3))
assert round(c, 4) == round((2 + 3) * 3 / 3**2, 4)
def setup(self):
with Builder() as m:
x = m.input('x')
x = self.model(x)
# avoid using the bound function.
y = type(self).to_scalar(x)
m.output(y=y)
self.m = m
y_ = []
for x_ in self.x_:
# run a step along x_
xl = self.x - x_ * (self.epsilon * 0.5)
xr = self.x + x_ * (self.epsilon * 0.5)
# numerical
yl = self.m.compute(init=dict(x=xl), vout='y', return_tape=False)
yr = self.m.compute(init=dict(x=xr), vout='y', return_tape=False)
y_1 = (yr - yl) / self.epsilon
y_.append(y_1)
y, tape = self.m.compute(init=dict(x=self.x),
vout='y',
return_tape=True)
self.tape = tape
self.y_ = y_
def test_symbol_eval():
with Builder() as m:
a = m.input('a')
m.output(
b=a.eval(lambda a: len(a)),
c=a.eval("len(x)"),
)
m.compute(['b', 'c'], init=dict(a=[1, 2, 3]))
m.compute_with_vjp(init=dict(a=[1, 2, 3]), v=dict(_b=1.0, _c=1.0))
m.compute_with_jvp(['b', 'c'], init=dict(a=[1, 2, 3]), v=dict(a_=1.0))
def __init__(self, comm, forward_operator, residuals, vectorspace=None):
self.residuals = residuals
self.forward_operator = forward_operator
self.comm = comm
if vectorspace is None:
vectorspace = MPIVectorSpace(comm)
with Builder() as m:
x = m.input('x')
y = 0
fx = forward_operator(x)
# fixme: need a way to directly include a subgraphs
# rather than building it again.
for operator in self.residuals:
r = operator(*fx)
chi2 = MPIChiSquareOperator(r, comm)
y = linalg.add(y, chi2)
m.output(y=y)
def objective(x):
return m.compute(vout='y', init=dict(x=x))
def gradient(x):
y, [vjp] = m.compute_with_vjp(init=dict(x=x), v=dict(_y=1.0))
return vjp
def hessian_vector_product(x, v):
Dv = 0
replay = forward_operator.precompute(x=x)
for operator in self.residuals:
with Builder() as m:
xx = m.input('x')
fx = replay(xx)
r = operator(*fx)
m.output(y=r)
y, [Dv1] = m.compute_with_gnDp(vout='y',
init=dict(x=x),
v=dict(x_=v))
Dv = Dv + Dv1
# H is 2 JtJ, see wikipedia on Gauss Newton.
return Dv * 2
BaseProblem.__init__(self,
vs=vectorspace,
objective=objective,
gradient=gradient,
hessian_vector_product=hessian_vector_product)
def hessian_vector_product(x, v):
Dv = 0
replay = forward_operator.precompute(x=x)
for operator in self.residuals:
with Builder() as m:
xx = m.input('x')
fx = replay(xx)
r = operator(*fx)
m.output(y=r)
y, [Dv1] = m.compute_with_gnDp(vout='y',
init=dict(x=x),
v=dict(x_=v))
Dv = Dv + Dv1
# H is 2 JtJ, see wikipedia on Gauss Newton.
return Dv * 2
def test_vmad3_functional():
""" this test demonstrates building a model directly """
with Builder() as m:
a, b = m.input('a', 'b')
t1 = add(a, a)
t2 = add(b, 0)
c = add(t1, t2)
m.output(c=c)
print("----- model -----")
pprint(m)
pprint(m[:])
print("----- compute -----")
init = dict(a=3, b=4)
c = m.compute(init=init, vout='c')
print(init, c)
print("----- tape -----")
init = dict(a=3, b=4)
c, tape = m.compute(init=init, vout='c', return_tape=True)
print(init, c)
pprint(tape)
print("----- vjp -----")
vjp = tape.get_vjp()
pprint(vjp)
pprint(vjp[:])
init = dict(_c=1.0)
_a, _b = vjp.compute(init=init, vout=['_a', '_b'], monitor=print)
print('_a, _b = ', _a, _b)
print("----- jvp -----")
jvp = tape.get_jvp()
pprint(jvp)
pprint(jvp[:])
init = dict(a_=1.0, b_=1.0)
c_, = jvp.compute(init=init, vout=['c_'], monitor=print)
print('c_ = ', c_)
def test_subclass_symbol():
class MySymbol(Symbol):
@operator
class __add__:
ain = 'x', 'y'
aout = 'z'
def apl(self, x, y):
return dict(z=x + y)
def vjp(self, _z):
return dict(_x=_z, _y=_z)
def jvp(self, x_, y_):
return dict(z_=x_ + y_)
with Builder() as m:
a, b = m.input(MySymbol('a'), MySymbol('b'))
c = a + b
m.output(c=c)
m.compute('c', init=dict(a=1, b=2))
m.compute_with_vjp(init=dict(a=1, b=2), v=dict(_c=1.0))
m.compute_with_jvp(['c'], init=dict(a=1, b=2), v=dict(a_=1.0, b_=1.0))
def _make_model(self):
with Builder() as m:
x = m.input('x')
y = self.model(x)
m.output(y=y)
return m
from vmad import Builder
from fastpm.force.lpt import lpt1, lpt2source
from vmad.lib import fastpm
import numpy
from nbodykit.cosmology import Planck15, LinearPower
pm = fastpm.ParticleMesh([32, 32, 32], BoxSize=128.)
powerspectrum = LinearPower(Planck15, 0)
q = pm.generate_uniform_particle_grid()
with Builder() as model:
x = model.input('x')
wnk = fastpm.as_complex_field(x, pm)
rhok = fastpm.induce_correlation(wnk, powerspectrum, pm)
dx1, dx2 = fastpm.lpt(rhok, q, pm)
dx, p, f = fastpm.nbody(rhok, q, [0.1, 0.6, 1.0], Planck15, pm)
dx0, p0, f0 = fastpm.nbody(rhok, q, [0.1], Planck15, pm)
model.output(dx1=dx1, dx2=dx2, dx=dx, p=p, dx0=dx0, p0=p0, f0=f0, f=f)
wn = pm.generate_whitenoise(555, unitary=True)
x = wn[...]
x = numpy.stack([x.real, x.imag], -1)
from fastpm.core import Solver, leapfrog