def make_data(bs, nc, seed, nsteps, path='', z=0):
#initiate
pm = ParticleMesh(BoxSize=bs, Nmesh=(nc, nc, nc), dtype='f8')
ofolder = path + 'z%02d/L%04d_N%04d_S%04d_%02dstep/'%(z*10, bs, nc, seed, nsteps)
#
if pm.comm.rank == 0:
print('Make data for seed = %d'%seed)
print('Data to be saved at path - %s'%ofolder)
klin, plin = numpy.loadtxt('ics_matterpow_0.dat', unpack = True)
pk = interpolate(klin, plin)
cosmo = Planck15.clone(Omega_cdm = 0.2685, h = 0.6711, Omega_b = 0.049)
#pk = EHPower(Planck15, redshift=0)
#cosmo = Planck15
s_truth = pm.generate_whitenoise(seed, type='complex')\
.apply(lambda k, v: v * (pk(sum(ki **2 for ki in k) **0.5) / v.BoxSize.prod()) ** 0.5)\
.c2r()
#dynamics
aa = 1.0/(1+z)
if pm.comm.rank == 0:
print('Evolve to redshift = %0.1f, scale factor = %0.2f'%(z, aa))
stages = numpy.linspace(0.1, aa, nsteps, endpoint=True)
start = time()
dynamic_model = NBodyModel(cosmo, pm, B=2, steps=stages)
#dynamic_model = LPTModel(cosmo, pm, B=2, steps=stages)
#Save data
X, V, final = dynamic_model.get_code().compute(['X', 'V', 'final'], init={'parameters':s_truth})
end = time()
print('Time taken = ', end-start)
save_map(s_truth, ofolder + 'mesh', 's')
save_map(final, ofolder + 'mesh', 'd')
if pm.comm.rank == 0:
print('X, V computed')
cat = ArrayCatalog({'Position': X, 'Velocity' : V}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
kdd = KDDensity(cat).density
cat['KDDensity'] = kdd
cat['InitPosition'] = dynamic_model.get_code().engine.q
cat.save(ofolder + 'dynamic/1', ('InitPosition', 'Position', 'Velocity', 'KDDensity'))
if pm.comm.rank == 0:
print('dynamic model created')
#FOF
fof = FOF(cat, linking_length=0.2, nmin=12)
fofcat = fof.find_features(peakcolumn='KDDensity')
fofcat['Mass'] = fofcat['Length'] * cosmo.Om0 * 27.7455 * pm.BoxSize.prod() / pm.Nmesh.prod()
fofcat.save(ofolder+'FOF', ('CMPosition', 'CMVelocity', 'PeakPosition', 'PeakVelocity', 'Length', 'Mass'))
def test_matter_dominated():
C = Planck15.clone(T_cmb=0.1,
m_ncdm=[]) # suppress cmb and remove neutrino
a = numpy.logspace(-2, 0, 11, endpoint=True)
z = 1 / a - 1
pt = MatterDominated(C.Omega0_m, a=a)
# linear growth function
D1 = pt.D1(a)
D_CC = C.scale_independent_growth_factor(z)
assert_allclose(D1, D_CC, rtol=1e-3)
# linear growth rate
f1 = pt.f1(a)
f_CC = C.scale_independent_growth_rate(z)
assert_allclose(f1, f_CC, rtol=1e-3)
# second order quantities
D2 = pt.D2(a)
f2 = pt.f2(a)
def set_up():
params = get_params()
pm = ParticleMesh(Nmesh=params['Nmesh'],
BoxSize=params['BoxSize'],
comm=MPI.COMM_WORLD,
resampler='cic')
# generate initial conditions
cosmo = Planck15.clone(P_k_max=30)
x = pm.generate_uniform_particle_grid(shift=0.1)
BoxSize2D = [deg / 180. * np.pi for deg in params['BoxSize2D']]
logger = None
sim = lightcone.WLSimulation(stages=numpy.linspace(0.1,
1.0,
params['N_steps'],
endpoint=True),
cosmology=cosmo,
pm=pm,
boxsize2D=BoxSize2D,
params=params,
logger=None)
kmaps = [sim.mappm.create('real', value=0.) for ii in range(1)]
return pm, cosmo, x, kmaps, sim.DriftFactor, sim.mappm, sim
prefix = 'test'
fname = 's%d_%s' % (seed, prefix)
optfolder = ofolder + 'opt_%s/' % fname
if truth_pm.comm.rank == 0: print('Output Folder is %s' % optfolder)
for folder in [ofolder, optfolder]:
try:
os.makedirs(folder)
except:
pass
#initiate
klin, plin = numpy.loadtxt('../data/pklin_1.0000.txt', unpack=True)
ipk = interpolate(klin, plin)
cosmo = Planck15.clone(Omega_cdm=0.2685, h=0.6711, Omega_b=0.049)
##################################################################################
##setup the model
##
stages = numpy.linspace(0.1, aa, nsteps, endpoint=True)
dynamic_model = NBodyModel(cosmo, truth_pm, B=B, steps=stages)
#dynamic_model = ZAModel(cosmo, truth_pm, B=B, steps=stages)
#noise
#Artifically low noise since the data is constructed from the model
truth_noise_model = map.NoiseModel(
truth_pm, None,
noisevar * (truth_pm.BoxSize / truth_pm.Nmesh).prod(), 1234)
mock_model = map.MockModel(dynamic_model, params=[2, 1, 0.5])
from mpi4py import MPI
import numpy
from argparse import ArgumentParser
from nbodykit.cosmology import Planck15
from nbodykit.cosmology import EHPower
from nbodykit.cosmology.perturbation import PerturbationGrowth
from scipy.integrate import quad
PowerSpectrum = EHPower(Planck15, redshift=0.0)
pt = PerturbationGrowth(Planck15.clone(Tcmb0=0))
class FastPM:
def K(ai, af, ar):
return 1 / (ar**2 * pt.E(ar)) * (pt.Gf(af) - pt.Gf(ai)) / pt.gf(ar)
def D(ai, af, ar):
return 1 / (ar**3 * pt.E(ar)) * (pt.Gp(af) - pt.Gp(ai)) / pt.gp(ar)
class FastPM1:
def K(ai, af, ar):
def func(a):
return 1.0 / (a * a * pt.E(a))
return quad(func, ai, af)[0]
def D(ai, af, ar):
return 1 / (ar**3 * pt.E(ar)) * (pt.Gp(af) - pt.Gp(ai)) / pt.gp(ar)
from runtests.mpi import MPITest
from fastpm.core import leapfrog, autostages, Solver as CoreSolver
from fastpm.state import StateVector, Matter, Baryon, CDM, NCDM
from fastpm.multi import Solver
from pmesh.pm import ParticleMesh
from nbodykit.cosmology import Planck15, EHPower
import numpy
from numpy.testing import assert_allclose
from fastpm.multi import get_species_transfer_function_from_class
Planck15 = Planck15.clone(gauge='newtonian')
@MPITest([1, 4])
def test_solver(comm):
pm = ParticleMesh(BoxSize=512., Nmesh=[8, 8, 8], comm=comm)
solver = Solver(pm, Planck15, B=1)
P_prm = Planck15.Primordial.get_pkprim
tf = get_species_transfer_function_from_class(Planck15, 9)
Q = pm.generate_uniform_particle_grid(shift=0)
wn = solver.whitenoise(1234)
prm = solver.primordial(wn, P_prm)
ic = solver.lpt(prm, {
'0': (Baryon, tf['d_b'], tf['dd_b']),
'1': (CDM, tf['d_cdm'], tf['dd_cdm']),
'4': (NCDM, tf['d_ncdm[0]'], tf['dd_ncdm[0]']),
from __future__ import print_function
from numpy.testing import assert_raises, assert_array_equal, assert_allclose
from numpy.testing.decorators import skipif
import numpy
from pmesh.abopt import ParticleMesh, RealField, ComplexField, check_grad
from nbodykit.cosmology import Planck15, EHPower
cosmo = Planck15.clone(Tcmb0=0)
pm = ParticleMesh(BoxSize=1.0, Nmesh=(4, 4, 4), dtype='f8')
pk = EHPower(Planck15, redshift=0)
mask2d = pm.resize([4, 4, 1]).create(mode='real')
mask2d[...] = 1.0
from cosmo4d import map2d
from cosmo4d.nbody import NBodyModel
from cosmo4d.options import UseComplexSpaceOptimizer, UseRealSpaceOptimizer
def test_map2d():
dynamic_model = NBodyModel(cosmo, pm, 1, [1.0])
mock_model = map2d.MockModel(dynamic_model)
noise_model = map2d.NoiseModel(pm, mask2d, 1.0, 1234)
initial = pm.generate_whitenoise(1234, mode='real')
obs = mock_model.make_observable(initial)
assert_array_equal(obs.map2d.Nmesh, (4, 4, 1))
from mpi4py import MPI
import numpy
from argparse import ArgumentParser
from nbodykit.cosmology import Planck15
from nbodykit.cosmology import EHPower
from nbodykit.cosmology.perturbation import PerturbationGrowth
from scipy.integrate import quad
PowerSpectrum = EHPower(Planck15, redshift=0.0)
pt = PerturbationGrowth(Planck15.clone(Tcmb0=0))
class FastPM:
def K(ai, af, ar):
return 1 / (ar ** 2 * pt.E(ar)) * (pt.Gf(af) - pt.Gf(ai)) / pt.gf(ar)
def D(ai, af, ar):
return 1 / (ar ** 3 * pt.E(ar)) * (pt.Gp(af) - pt.Gp(ai)) / pt.gp(ar)
class FastPM1:
def K(ai, af, ar):
def func(a):
return 1.0 / (a * a * pt.E(a))
return quad(func, ai, af)[0]
def D(ai, af, ar):
return 1 / (ar ** 3 * pt.E(ar)) * (pt.Gp(af) - pt.Gp(ai)) / pt.gp(ar)
class FastPM2:
def K(ai, af, ar):
return 1 / (ar ** 2 * pt.E(ar)) * (pt.Gf(af) - pt.Gf(ai)) / pt.gf(ar)
def D(ai, af, ar):
def func(a):