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'))
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])
##################################################################################
### Create and save data if not found
### Comment the following lines since we do not have any pregenerated data yet
##
s_truth = truth_pm.generate_whitenoise(seed, mode='complex')\
def solve(Nmesh, x0, rtol, run, Nsm):
pm = truth_pm.resize(Nmesh=(Nmesh, Nmesh, Nmesh))
atol = pm.Nmesh.prod() * rtol
x0 = pm.upsample(x0, keep_mean=True)
#data = data_n.downsample(pm)
#IDEAL no noise limit
data = data_p.downsample(pm)
# smooth the data. This breaks the noise model but we don't need it
# for lower resolution anyways.
sml = pm.BoxSize[0] / Nmesh * Nsm
dynamic_model = NBodyModel(cosmo, truth_pm, B=B, steps=stages)
mock_model = map.MockModel(dynamic_model)
# an approximate noise model, due to smoothing this is correct only at large scale.
noise_model = truth_noise_model.downsample(pm)
obj = map.SmoothedObjective(mock_model, noise_model, data, prior_ps=pk, sml=sml)#, noised=noised)
prior, chi2 = obj.get_code().compute(['prior', 'chi2'], init={'parameters': data.s})
if pm.comm.rank == 0:
print(prior, chi2) # for 2d chi2 is close to total pixels.
fit_p = mock_model.make_observable(data.s)
r = obj.evaluate(fit_p, data)
try:
os.makedirs(optfolder + '%s' % run)
except:
pass
try:
os.makedirs(optfolder + '%s/2pt' % run)
except:
pass
obj.save_report(r, optfolder + "%s/truth.png" % run)
dg.save_2ptreport(r, optfolder + "%s/2pt/truth.png" % run, pm)
optimizer = LBFGS(m=10, diag_update=scalar_diag)
prob = obj.get_problem(atol=atol, precond=UseComplexSpaceOptimizer)
def monitor(state):
if pm.comm.rank == 0:
print(state)
if state.nit % 5 == 0:
fit_p = mock_model.make_observable(state['x'])
if state.nit % 20 == 0:
fit_p.save(optfolder + '%s/%04d/fit_p' % (run, state['nit']))
r = obj.evaluate(fit_p, data)
#obj.save_report(r, optfolder + "%s/%s%02d-%04d.png"% (run, prefix, int(Nsm*10), state['nit']))
dg.save_report(r, optfolder + "%s/%s_N%02d-%04d.png"% (run, prefix, int(Nsm*10), state['nit']), pm)
dg.save_2ptreport(r, optfolder + "%s/2pt/%s_N%02d-%04d.png"% (run, prefix, int(Nsm*10), state['nit']), pm)
if pm.comm.rank == 0:
print('saved')
state = optimizer.minimize(prob, x0=x0, monitor=monitor)
fit_p = mock_model.make_observable(state['x'])
fit_p.save(optfolder + '%s/best-fit' % run)
r = obj.evaluate(fit_p, data)
obj.save_report(r, optfolder + "%s/%s%02d-best-fit.png" % (run, prefix, int(Nsm*10)))
dg.save_2ptreport(r, optfolder + "%s/2pt/%s_N%02d-best-fit.png" % (run, prefix, int(Nsm*10)), pm)
return state.x
else: dnewfolder = '/global/cscratch1/sd/chmodi/m3127/cm_lowres/%dstepT-B%d/%d-%d-9100-fixed/'%(nsteps, B, bs, nc*2) #dnewfolder = '/global/cscratch1/sd/chmodi/m3127/cm_lowres/%dstepT-B%d/%d-%d-9100-fixed/'%(nsteps, B, bs, nc*2) s_truth = BigFileMesh(dnewfolder + 'linear', 'LinearDensityK').paint() dyn = BigFileCatalog(dnewfolder + 'fastpm_%0.4f/1'%aa) dlayout = new_pm.decompose(dyn['Position']) d_truth = new_pm.paint(dyn['Position'], layout=dlayout) ## #Model params = numpy.loadtxt(optfolder + '/params.txt') stages = numpy.linspace(0.01, aa, nsteps, endpoint=True) if pmdisp: dynamic_model = NBodyModel(cosmo, new_pm, B=B, steps=stages) else: dynamic_model = ZAModel(cosmo, new_pm, B=B, steps=stages) if rank == 0: print(dynamic_model) #noise if stage2 is not None: truth_noise_model = mapnoise.ThermalNoise(new_pm, seed=100, aa=aa, att=stage2,spread=spread, hex=hex, limk=2, Ns=Ndish) else: truth_noise_model = mapnoise.ThermalNoise(new_pm, seed=None, aa=aa, att=stage2,spread=spread, hex=hex, Ns=Ndish) wedge_noise_model = mapnoise.WedgeNoiseModel(pm=new_pm, power=1, seed=100, kmin=kmin, angle=angle) #Create and save data if not found ################# mock_model = map.MockModel(dynamic_model, params=params, rsdpos=rsdpos, rsdfac=rsdfac) try: data_p = map.Observable.load(optfolder+'/datap_up') except: