def PGD_correction(X, alpha, kl, ks, pm, q):
layout = fastpm.decompose(X, pm)
xl = fastpm.exchange(X, layout)
rho = fastpm.paint(xl, 1.0, None, pm)
fac = 1.0 * pm.Nmesh.prod() / pm.comm.allreduce(len(q))
rho = rho * fac
rhok = fastpm.r2c(rho)
p = fastpm.apply_transfer(rhok, PGDkernel(kl, ks))
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)
S = linalg.stack(r1, axis=-1)
S = S * alpha
return S
def model(self, x):
from nbodykit.cosmology import Planck15
dx, p, f = fastpm.nbody(fastpm.r2c(x),
q=self.pos,
stages=[0.1, 0.5, 1.0],
pm=self.pm,
cosmology=Planck15)
return linalg.stack([dx, p, f], axis=-1)
def model(self, x):
c = fastpm.r2c(x)
digitizer = fastpm.apply_digitized.isotropic_wavenumber(self.kedges)
c = fastpm.apply_digitized(c,
tf=self.tf,
digitizer=digitizer,
kind='wavenumber')
r = fastpm.c2r(c)
return r
def model(self, x):
from nbodykit.cosmology import Planck15
sim = fastpm.FastPMSimulation(pm=self.pm,
cosmology=Planck15,
stages=[0.1, 0.5, 1.0],
q=self.pos)
dx, p, f = sim.run(fastpm.r2c(x))
return linalg.stack([dx, p, f], axis=-1)
def model(self, x):
compensation = self.pm.resampler.get_compensation()
layout = vmadfastpm.decompose(self.w, self.pm)
map = vmadfastpm.paint(self.w, x, layout, self.pm)
y = map + self.xx # bias needed to avoid zero derivative
# compensation for cic window
c = vmadfastpm.r2c(y)
c = vmadfastpm.apply_transfer(c,
lambda k: compensation(k, 1.0),
kind='circular')
map = vmadfastpm.c2r(c)
return map
def smoothed_residue(param, X, pm, Nstep, target, n, baryon=True):
F = LDL(param, X, pm, Nstep, baryon=baryon)
#residue field
residue = F - target
#smooth the field
Filter = pm.create(type='complex', value=1).apply(smoothing(n=n))
residuek = fastpm.r2c(residue)
residuek = residuek * Filter
residue = fastpm.c2r(residuek)
return residue
def smoothed_residue(param, X, pm, Nstep, target, n, baryon=True, index=1, field2=None):
F = LDL(param, X, pm, Nstep, baryon=baryon)
if index != 1:
F = F ** index
if field2 is not None:
F = F * field2
#residue field
residue = F - target
#smooth the field
Filter = pm.create(type='complex', value=1).apply(smoothing(n=n))
residuek = fastpm.r2c(residue)
residuek = residuek * Filter
residue = fastpm.c2r(residuek)
return residue
def makemap(self, xy, w):
"""
paint projected particles to 2D mesh
xy: particle positions in radians
w: weighting = projection kernel
"""
if (self.mappm.affine.period != 0).any():
raise RuntimeError("The ParticeMesh object must be non-periodic")
if self.mappm.ndim != 2:
raise RuntimeError(
"The ParticeMesh object must be 2 dimensional. ")
compensation = self.mappm.resampler.get_compensation()
layout = fastpm.decompose(xy, self.mappm)
map = fastpm.paint(xy, w, layout, self.mappm)
# compensation for cic window
c = fastpm.r2c(map)
c = fastpm.apply_transfer(c,
lambda k: compensation(k, 1.0),
kind='circular')
map = fastpm.c2r(c)
return map
def model(self, x):
c = fastpm.r2c(x)
c = fastpm.apply_transfer(c, tf=transfer)
r = fastpm.c2r(c)
return r
def model(self, x):
c = fastpm.r2c(x)
r = fastpm.c2r(c)
return r
def run(self, rho):
rhok = fastpm.r2c(rho)
dx, p = self.firststep(rhok)
pt = self.pt
stages = self.stages
q = self.q
Om0 = pt.Om0
zero_map = Literal(self.mappm.create('real', value=0.))
kmaps = [zero_map for ds in self.ds]
f, potk = self.gravity(dx)
jj = 0 #counting steps for saving snapshots
for ai, af in zip(stages[:-1], stages[1:]):
if self.params['logging']:
self.logger.info('fastpm step, %d' % jj)
# central scale factor
ac = (ai * af)**0.5
# kick (update momentum)
dp = f * (self.KickFactor(ai, ai, ac) * 1.5 * Om0)
p = p + dp
# drift (update positions)
ddx = p * self.DriftFactor(ai, ac, af)
dx = dx + ddx
if self.params['PGD']:
alpha = self.alpha0 * af**self.mu
dx_PGD = PGD.PGD_correction(q + dx, alpha, self.kl, self.ks,
self.fpm, q)
else:
dx_PGD = 0.
if self.params['save3D'] or self.params['save3Dpower']:
zf = 1. / af - 1.
zi = 1. / ai - 1.
if zi < self.params['zs_source']:
pos_raw = dx + q
pos = pos_raw + dx_PGD
stdlib.watchpoint(
pos_raw,
lambda pos, ii=jj, zi=zi, zf=zf, params=self.params:
save_snapshot(pos, ii, zi, zf, params, 'raw'))
stdlib.watchpoint(
pos,
lambda pos, ii=jj, zi=zi, zf=zf, params=self.params:
save_snapshot(pos, ii, zi, zf, params, 'PGD'))
jj += 1
kmaps_ = self.no_interp(
dx, p, dx_PGD, ai, af, jj, kmaps=ListPlaceholder(len(self.ds))
) #[Symbol('kmaps-%d-%d'%(ii,jj)) for ii in range(len(self.ds))])
for ii in range(len(self.ds)):
kmaps[ii] = kmaps[ii] + kmaps_[ii]
# force (compute force)
f, potk = self.gravity(dx)
# kick (update momentum)
dp = f * (self.KickFactor(ac, af, af) * 1.5 * Om0)
p = p + dp
return kmaps
def model(self, x):
return fastpm.lpt2src(fastpm.r2c(x), pm=self.pm)
def model(self, x):
dx1 = fastpm.lpt1(fastpm.r2c(x), q=self.pos, pm=self.pm)
return dx1
def model(self, x):
x1 = fastpm.r2c(self.x1)
x = fastpm.r2c(x)
return fastpm.cdot(x, x1)
def model(self, x):
dx1, dx2 = fastpm.lpt(fastpm.r2c(x), q=self.pos, pm=self.pm)
return linalg.add(dx1, dx2)
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 run_interpolated(self, rho):
rhok = fastpm.r2c(rho)
dx, p = self.firststep(rhok)
pt = self.pt
stages = self.stages
q = self.q
Om0 = pt.Om0
powers = []
zero_map = Literal(self.mappm.create('real', value=0.))
kmaps = [zero_map for ds in self.ds]
f, potk = self.gravity(dx)
for ai, af in zip(stages[:-1], stages[1:]):
# central scale factor
ac = (ai * af)**0.5
# kick
dp = f * (self.KickFactor(ai, ai, ac) * 1.5 * Om0)
p = p + dp
# drift
ddx = p * self.DriftFactor(ai, ac, ac)
dx = dx + ddx
if self.params['PGD']:
alpha = self.alpha0 * ac**self.mu
dx_PGD = PGD.PGD_correction(self.q + dx, alpha, self.kl,
self.ks, self.fpm, self.q)
else:
dx_PGD = 0.
#if interpolation is on, only take 'half' and then evolve according to their position
kmaps_ = self.interp(dx,
p,
dx_PGD,
ac,
ac,
ai,
af,
kmaps=ListPlaceholder(len(self.ds)))
for ii in range(len(self.ds)):
kmaps[ii] = kmaps[ii] + kmaps_[ii]
# drift
ddx = p * self.DriftFactor(ac, ac, af)
dx = dx + ddx
# force
f, potk = self.gravity(dx)
# kick
dp = f * (self.KickFactor(ac, af, af) * 1.5 * Om0)
p = p + dp
return dict(kmaps=kmaps)
