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 model(self, x):
a = linalg.take(x, [[2, 2], [2, 2]], axis=0)
b = linalg.take(a, 0, axis=0)
c = linalg.take(b, 0, axis=0)
#d = linalg.take(b, 0, axis=0)
#watchpoint(b, lambda b:print('b=', b))
#watchpoint(c, lambda c:print('c=', c))
# watchpoint(d, lambda d:print('d=', d))
return c + c
def rotate(self, x, M, boxshift):
"""
rotates, shift, and separates particle coordinates into distance and xy position
x: particle positions
M: rotation matrix
boxshift: shift vector
"""
y = linalg.einsum('ij,kj->ki', (M, x))
y = y + self.pm.BoxSize * boxshift
d = linalg.take(y, 2, axis=1)
xy = linalg.take(y, (0, 1), axis=1)
return dict(xy=xy, d=d)
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):
b = 0
for i in range(3):
b = b + linalg.take(x, i, axis=0)
return b
def model(self, x):
return linalg.sum(linalg.take(x, [2, 2], axis=0), axis=0)
def model(self, x):
x_ = linalg.take(x, 0, axis=0)
elems = linalg.take(x, 1, axis=0)
elem = linalg.take(elems, 0, axis=0)
res = lightcone.list_put(x_, elem, self.i)
return res
def model(self, x):
y = linalg.einsum('ij,kj->ki', (self.M, x))
y = y + self.pm.BoxSize * self.boxshift
d = linalg.take(y, 2, axis=1)
xy = linalg.take(y, (0, 1), axis=1)
return xy #d
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
def model(self, x):
x1 = linalg.take(x, 0, axis=0)
x2 = linalg.take(x, 1, axis=0)
return type(self).ufunc(x1, x2)