def gravity(self, dx):
q = self.q
# use Force PM resolution including B factor.
pm = self.fpm
x = q + dx
layout = decompose(x, pm)
xl = exchange(x, layout)
rho = paint(xl, 1.0, None, pm)
# convert to 1 + delta
fac = 1.0 * pm.Nmesh.prod() / pm.comm.allreduce(len(q))
rho = rho * fac
rhok = r2c(rho)
p = apply_transfer(rhok, fourier_space_laplace)
r1 = []
for d in range(pm.ndim):
dx1_c = apply_transfer(p, fourier_space_neg_gradient(d,
pm,
order=1))
dx1_r = c2r(dx1_c)
dx1l = readout(dx1_r, xl, None)
dx1 = gather(dx1l, layout)
r1.append(dx1)
f = linalg.stack(r1, axis=-1)
return dict(f=f, potk=p)
def gravity(dx, q, pm):
x = q + dx
#def w(q): print('q', q)
#watchpoint(x, w)
#watchpoint(x, lambda x: print('x', q))
#watchpoint(dx, lambda dx: print('dx', dx))
layout = decompose(x, pm)
xl = exchange(x, layout)
rho = paint(xl, 1.0, None, pm)
# convert to 1 + delta
fac = 1.0 * pm.Nmesh.prod() / pm.comm.allreduce(len(q))
rho = rho * fac
rhok = r2c(rho)
p = apply_transfer(rhok, fourier_space_laplace)
r1 = []
for d in range(pm.ndim):
dx1_c = apply_transfer(p, fourier_space_neg_gradient(d, pm, order=1))
dx1_r = c2r(dx1_c)
dx1l = readout(dx1_r, xl, None)
dx1 = gather(dx1l, layout)
r1.append(dx1)
f = linalg.stack(r1, axis=-1)
return dict(f=f, potk=p)
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):
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 lpt1(rhok, q, pm):
p = apply_transfer(rhok, fourier_space_laplace)
layout = decompose(q, pm)
r1 = []
for d in range(pm.ndim):
dx1_c = apply_transfer(p, fourier_space_neg_gradient(d, pm, order=1))
dx1_r = c2r(dx1_c)
dx1 = readout(dx1_r, q, layout)
r1.append(dx1)
dx1 = linalg.stack(r1, axis=-1)
return dict(dx1=dx1)
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 model(self, x):
return linalg.stack([x, x], axis=0)
