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)