def upsample2(Sl, Ql): #pm has the same resolution as Ql
layout = add(Ql, 1)
print(id(Model.get_model(layout)), len(Model.get_model(layout)))
Ql1 = add(Ql, layout)
print(id(Model.get_model(layout)), id(Model.get_model(Ql1)), len(Model.get_model(layout)))
Sl1 = add(Sl, layout)
print(id(Model.get_model(layout)), id(Model.get_model(Ql1)), len(Model.get_model(layout)))
dis_d = add(Ql1, Sl1)
print(id(Model.get_model(layout)), id(Model.get_model(Ql1)), len(Model.get_model(layout)))
return dis_d
def model(self, x):
layout = fastpm.decompose(x, pm=self.pm)
x1 = fastpm.exchange(x, layout)
y1 = fastpm.paint(x1, mass=1.0, layout=None, pm=self.pm)
y = linalg.add(y1,
self.mesh) # biasing a bit to get non-zero derivatives.
return y
def test_operator_watchpoint():
from vmad.core.stdlib import watchpoint
foo = [0]
def monitor(x):
foo[0] = 1
with Builder() as m:
a = m.input('a')
b = linalg.add(a, a)
watchpoint(a, monitor=monitor)
m.output(c=b)
init = [('a', 1)]
# for node in m:
# print('m', node)
c, tape = m.compute(init=init, vout='c', return_tape=True)
# vjp = tape.get_vjp()
# for node in vjp:
# print('vjp', node)
[c], [_a] = m.compute_with_vjp(init=init, v=[('_c', 1.0)])
assert foo[0] == 1
assert c == 2
assert _a == 2.0
foo[0] = 0
c, c_ = m.compute_with_jvp(vout='c', init=init, v=[('a_', 1.0)])
assert foo[0] == 1
assert c == 2
assert c_ == 2.0
def test_model_partial_out():
with Builder() as m:
a = m.input('a')
t1 = add(x1=a, x2=a)
m.output(c=t1)
m.output(a=a)
init = dict(a=3)
(a, c), tape = m.compute(init=init, vout=['a', 'c'], return_tape=True)
assert c == 6
assert a == 3
vjp = tape.get_vjp()
# test two outputs individually
init = dict(_c=1.0, _a=0.0)
_a = vjp.compute(init=init, vout='_a', monitor=print)
assert _a == 2.0
init = dict(_c=0.0, _a=1.0)
_a = vjp.compute(init=init, vout='_a', monitor=print)
assert _a == 1.0
jvp = tape.get_jvp()
init = dict(a_=1.0)
a_, c_ = jvp.compute(init=init, vout=['a_', 'c_'], monitor=print)
assert c_ == 2.0
assert a_ == 1.0
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_model_compute_with_jvp():
with Builder() as m:
a = m.input('a')
t1 = add(x1=a, x2=a)
m.output(b=t1)
init = [('a', 1)]
b, b_ = m.compute_with_jvp(vout='b', init=init, v=[('a_', 1.0)])
assert b == 2.0
assert b_ == 2.0
def test_model_compute_with_vjp():
with Builder() as m:
a = m.input('a')
t1 = add(x1=a, x2=a)
m.output(b=t1)
init = [('a', 1)]
[b], [_a] = m.compute_with_vjp(init=init, v=[('_b', 1.0)])
assert b == 2.0
assert _a == 2.0
def test_model_attr():
import numpy
with Builder() as m:
a, b = m.input('a', 'b')
d = add(x1=b, x2=1)
t1 = add(x1=a, x2=b.eval(lambda b: b.size))
m.output(c=t1)
init = dict(a=2, b=numpy.array([2,]))
c, tape = m.compute(init=init, vout='c', return_tape=True)
assert c == 3
vjp = tape.get_vjp()
init = dict(_c=1.0)
_a = vjp.compute(init=init, vout='_a', monitor=print)
assert _a == 1.0
jvp = tape.get_jvp()
init = dict(a_=1.0)
c_ = jvp.compute(init=init, vout='c_', monitor=print)
assert c_ == 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 interp(self, dx, p, kmaps, dx_PGD, ax, ap, ai, af):
di, df = self.cosmo.comoving_distance(
1. / numpy.array([ai, af], dtype=float) - 1.)
for M in self.imgen.generate(di, df):
# if lower end of box further away than source -> do nothing
if df > self.max_ds:
continue
else:
M, boxshift = M
# positions of unevolved particles after rotation
d_approx = self.rotate.build(M=M, boxshift=boxshift).compute(
'd', init=dict(x=self.q))
z_approx = z_chi.apl.impl(node=None,
cosmo=self.cosmo,
z_chi_int=self.z_chi_int,
chi=d_approx)['z']
a_approx = 1. / (z_approx + 1.)
# move particles to a_approx, then add PGD correction
dx1 = dx + p * self.DriftFactor(a_approx, ax,
ap)[:, None] + dx_PGD
# rotate their positions
xy, d = self.rotate((dx1 + self.q) % self.pm.BoxSize, M,
boxshift)
# projection
xy = ((xy - self.pm.BoxSize[:2] * 0.5) /
linalg.broadcast_to(linalg.reshape(d, (len(self.q), 1)),
(len(self.q), 2)) +
self.mappm.BoxSize * 0.5)
for ii, ds in enumerate(self.ds):
w = self.wlen(d, ds)
mask = stdlib.eval(
d,
lambda d, di=di, df=df, ds=ds, d_approx=d_approx: 1.0 *
(d_approx < di) * (d_approx >= df) * (d <= ds))
kmap = list_elem(kmaps, ii)
kmap_ = self.makemap(xy, w * mask) * self.factor
kmap = linalg.add(kmap_, kmap)
kmaps = list_put(kmaps, kmap, ii)
return kmaps
def model(self, x):
di, df = self.cosmo.comoving_distance(1. /
numpy.array([self.ai, self.af]) -
1.)
kmap = lightcone.list_elem(self.kmaps, 0)
for M in self.sim.imgen.generate(di, df):
#if lower end of box further away than source -> do nothing
if df > self.sim.max_ds:
continue
else:
M, boxshift = M
#positions of unevolved particles after rotation
d_approx = self.sim.rotate.build(
M=M, boxshift=boxshift).compute('d', init=dict(x=self.q))
z_approx = lightcone.z_chi.apl.impl(
node=None,
cosmo=self.cosmo,
z_chi_int=self.sim.z_chi_int,
chi=d_approx)['z']
a_approx = 1. / (z_approx + 1.)
#move particles to a_approx, then add PGD correction
dx1 = x + self.p * self.DriftFactor(
a_approx, self.ac, self.ac)[:, None] + self.dx_PGD
#rotate their positions
xy, d = self.sim.rotate((dx1 + self.q) % self.pm.BoxSize, M,
boxshift)
#projection
xy = (xy - self.pm.BoxSize[:2] * 0.5) / linalg.broadcast_to(
linalg.reshape(d, (np.prod(self.pm.Nmesh), 1)),
(np.prod(self.pm.Nmesh), 2)) + self.sim.mappm.BoxSize * 0.5
for ii, ds in enumerate(self.sim.ds):
w = self.sim.wlen(d, ds)
mask = stdlib.eval(
d,
lambda d, di=di, df=df, ds=ds, d_approx=d_approx: 1.0 *
(d_approx < di) * (d_approx >= df) * (d <= ds))
#kmap = lightcone.list_elem(self.kmaps,ii)
kmap_ = self.sim.makemap(xy, w * mask)
kmap = linalg.add(kmap_, kmap)
#self.kmaps = lightcone.list_put(self.kmaps,kmap,ii)
#kmap_ = RealField(lightcone.list_elem(self.kmaps,0))
return kmap
def test_model_many_rewrites():
# this is a nasty model with many variable rewrites.
n = 2
with Builder() as m:
x = m.input('x')
for i in range(2):
x = add(x1=x, x2=x)
m.output(y=x)
init = dict(x=1.0)
y, tape = m.compute(init=init, vout='y', return_tape=True)
assert y == 4.0
vjp = tape.get_vjp()
init = dict(_y = 1.0)
_x = vjp.compute(init=init, vout='_x', monitor=print)
assert _x == 4.0
def test_model_partial():
with Builder() as m:
a = m.input('a')
t1 = add(x1=a, x2=a)
m.output(c=t1)
init = dict(a=3)
c, tape = m.compute(init=init, vout='c', return_tape=True)
assert c == 6
vjp = tape.get_vjp()
init = dict(_c=1.0)
_a = vjp.compute(init=init, vout='_a', monitor=print)
assert _a == 2.0
jvp = tape.get_jvp()
init = dict(a_=1.0)
c_ = jvp.compute(init=init, vout='c_', monitor=print)
assert c_ == 2.0
def no_interp(self, kmaps, q, ai, af, jj):
di, df = self.cosmo.comoving_distance(
1. / numpy.array([ai, af], dtype=object) - 1.)
#q = np.random.random(self.q.shape)*self.pm.BoxSize[0]
for M in self.imgen.generate(di, df):
# if lower end of box further away than source -> do nothing
if df > self.max_ds:
if self.params['logging']:
self.logger.info('imgen passed, %d' % jj)
continue
else:
if self.params['logging']:
self.logger.info('imgen with projection, %d' % jj)
M, boxshift = M
xy, d = self.rotate(q, M, boxshift)
d_approx = self.rotate.build(M=M, boxshift=boxshift).compute(
'd', init=dict(x=q))
xy = ((xy - self.pm.BoxSize[:2] * 0.5) /
linalg.broadcast_to(linalg.reshape(d, (len(self.q), 1)),
(len(self.q), 2)) +
self.mappm.BoxSize * 0.5)
for ii, ds in enumerate(self.ds):
if self.params['logging']:
self.logger.info('projection, %d' % jj)
w = self.wlen(d, ds)
mask = stdlib.eval(
d,
lambda d, di=di, df=df, ds=ds, d_approx=d_approx: 1.0 *
(d_approx < di) * (d_approx >= df) * (d <= ds))
kmap_ = self.makemap(xy, w * mask) * self.factor
kmap = list_elem(kmaps, ii)
kmap = linalg.add(kmap_, kmap)
kmaps = list_put(kmaps, kmap, ii)
return kmaps
def model(self, x):
c = linalg.pack_complex(x, x)
r, i = linalg.unpack_complex(c)
return linalg.add(r, i)
def model(self, x):
return linalg.add(x, 5.0)
def model(self, x):
return linalg.add(linalg.abs(x), x)
def main(self, x, n):
for i in range(n):
x = add(x1=x, x2=x)
return dict(y=x)
def example_func(x, n):
for i in range(n):
x = add(x1=x, x2=x)
return dict(y=x)
def main(model, a, b, n):
for i in range(n):
a = add(a, a)
t2 = add(b, 0)
return dict(c=add(a, t2))
def model(self, x):
layout = fastpm.decompose(x, pm=self.pm)
y = fastpm.paint(x, layout=layout, mass=1.0, pm=self.pm)
y = linalg.add(y,
self.mesh) # biasing a bit to get non-zero derivatives.
return y
def model(self, x):
dx1, dx2 = fastpm.lpt(fastpm.r2c(x), q=self.pos, pm=self.pm)
return linalg.add(dx1, dx2)