def LDL(param, X, pm, Nstep, baryon=True):
fac = 1.0 * pm.Nmesh.prod() / pm.comm.allreduce(len(X), op=MPI.SUM)
X = Displacement(param, X, pm, Nstep)
#paint particle overdensity field
layout = fastpm.decompose(X, pm)
Xl = fastpm.exchange(X, layout)
delta = fac * fastpm.paint(Xl, 1., None, pm)
if baryon:
#take parameters
mu = linalg.take(param, 5*Nstep, axis=0)
b1 = linalg.take(param, 5*Nstep+1, axis=0)
b0 = linalg.take(param, 5*Nstep+2, axis=0)
mu = mpi.allbcast(mu, comm=pm.comm)
mu = linalg.broadcast_to(mu, eval(delta, lambda x : x.shape))
b1 = mpi.allbcast(b1, comm=pm.comm)
b1 = linalg.broadcast_to(b1, eval(delta, lambda x : x.shape))
b0 = mpi.allbcast(b0, comm=pm.comm)
b0 = linalg.broadcast_to(b0, eval(delta, lambda x : x.shape))
#Field transformation
F = ReLU(b1 * (delta+1e-8) ** mu + b0) #definition of b0 is different from the paper
else:
F = delta
return F
def LDL(param, X, pm, Nstep, baryon=True):
fac = 1.0 * pm.Nmesh.prod() / pm.comm.allreduce(len(X), op=MPI.SUM)
X = Displacement(param, X, pm, Nstep)
#paint particle overdensity field
layout = fastpm.decompose(X, pm)
Xl = fastpm.exchange(X, layout)
delta = fac * fastpm.paint(Xl, 1., None, pm)
if baryon:
#take parameters
gamma = linalg.take(param, 5*Nstep, axis=0)
bias0 = linalg.take(param, 5*Nstep+1, axis=0)
bias1 = linalg.take(param, 5*Nstep+2, axis=0)
gamma = mpi.allbcast(gamma, comm=pm.comm)
gamma = linalg.broadcast_to(gamma, eval(delta, lambda x : x.shape))
bias0 = mpi.allbcast(bias0, comm=pm.comm)
bias0 = linalg.broadcast_to(bias0, eval(delta, lambda x : x.shape))
bias1 = mpi.allbcast(bias1, comm=pm.comm)
bias1 = linalg.broadcast_to(bias1, eval(delta, lambda x : x.shape))
#Field transformation
F = ReLU(bias0 * (delta+1e-8) ** gamma + bias1)
else:
F = delta
return F
def test(x):
size = stdlib.eval(x, lambda x: numpy.prod(x.shape))
mean = linalg.sum(x, axis=None) / size
mean = linalg.broadcast_to(mean, stdlib.eval(x, lambda x: x.shape))
x = x - mean
y = linalg.sum(x**2, axis=None) / size
return y
def interp(self, dx, p, dx_PGD, ax, ap, ai, af):
di, df = self.cosmo.comoving_distance(
1. / numpy.array([ai, af], dtype=float) - 1.)
zero_map = Literal(self.mappm.create('real', value=0.))
kmaps = [zero_map for ii in range(len(self.ds))]
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_ = self.makemap(xy, w * mask) * self.factor
kmaps[ii] = kmaps[ii] + kmap_
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 no_interp(self, dx, p, dx_PGD, ai, af, jj):
dx = dx + dx_PGD
di, df = self.cosmo.comoving_distance(
1. / numpy.array([ai, af], dtype=object) - 1.)
zero_map = Literal(self.mappm.create('real', value=0.))
kmaps = [zero_map for ii in range(len(self.ds))]
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((dx + self.q) % self.pm.BoxSize, M,
boxshift)
d_approx = self.rotate.build(M=M, boxshift=boxshift).compute(
'd', init=dict(x=self.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
kmaps[ii] = kmaps[ii] + kmap_
return kmaps
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 Displacement(param, X, pm, Nstep):
#Lagrangian displacement
#normalization constant for overdensity
fac = 1.0 * pm.Nmesh.prod() / pm.comm.allreduce(len(X), op=MPI.SUM)
for i in range(Nstep):
#move the particles across different MPI ranks
layout = fastpm.decompose(X, pm)
xl = fastpm.exchange(X, layout)
delta = fac * fastpm.paint(xl, 1.0, None, pm)
#take parameters
alpha = linalg.take(param, 5*i, axis=0)
gamma = linalg.take(param, 5*i+1, axis=0)
kh = linalg.take(param, 5*i+2, axis=0)
kl = linalg.take(param, 5*i+3, axis=0)
n = linalg.take(param, 5*i+4, axis=0)
#delta**gamma
gamma = mpi.allbcast(gamma, comm=pm.comm)
gamma = linalg.broadcast_to(gamma, eval(delta, lambda x : x.shape))
delta = (delta+1e-8) ** gamma
#Fourier transform
deltak = fastpm.r2c(delta)
#Green's operator in Fourier space
Filter = Literal(pm.create(type='complex', value=1).apply(lambda k, v: k.normp(2, zeromode=1e-8) ** 0.5))
kh = mpi.allbcast(kh, comm=pm.comm)
kh = linalg.broadcast_to(kh, eval(Filter, lambda x : x.shape))
kl = mpi.allbcast(kl, comm=pm.comm)
kl = linalg.broadcast_to(kl, eval(Filter, lambda x : x.shape))
n = mpi.allbcast(n, comm=pm.comm)
n = linalg.broadcast_to(n, eval(Filter, lambda x : x.shape))
Filter = - unary.exp(-Filter**2/kl**2) * unary.exp(-kh**2/Filter**2) * Filter**n
Filter = compensate2factor(Filter)
p = complex_mul(deltak, Filter)
#gradient of potential
r1 = []
for d in range(pm.ndim):
dx1_c = fastpm.apply_transfer(p, fastpm.fourier_space_neg_gradient(d, pm, order=1))
dx1_r = fastpm.c2r(dx1_c)
dx1l = fastpm.readout(dx1_r, xl, None)
dx1 = fastpm.gather(dx1l, layout)
r1.append(dx1)
#displacement
S = linalg.stack(r1, axis=-1)
alpha = mpi.allbcast(alpha, comm=pm.comm)
alpha = linalg.broadcast_to(alpha, eval(S, lambda x : x.shape))
S = S * alpha
X = X+S
return X
def func(x, n):
a = linalg.take(n, 0, axis=0)
a = linalg.broadcast_to(a, eval(x, lambda x: x.shape))
return x + a
