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: