def make_kappa_maps(cat, nside, zs_list, ds_list, localsize, nbar): """ Make kappa maps at a list of ds Return kappa, Nm in shape of (n_ds, localsize), kappabar in shape of (n_ds,) The maps are distributed in memory, and localsize is the size of map on this rank. """ dl = (abs(cat['Position'] **2).sum(axis=-1)) ** 0.5 chunks = dl.chunks ra = cat['RA'] dec = cat['DEC'] zl = (1 / cat['Aemit'] - 1) ipix = da.apply_gufunc(lambda ra, dec, nside: healpix.ang2pix(nside, numpy.radians(90-dec), numpy.radians(ra)), '(),()->()', ra, dec, nside=nside) npix = healpix.nside2npix(nside) ipix = ipix.compute() dl = dl.persist() cat.comm.barrier() if cat.comm.rank == 0: cat.logger.info("ipix and dl are persisted") area = (4 * numpy.pi / npix) * dl**2 Om = cat.attrs['OmegaM'][0] kappa_list = [] kappabar_list = [] Nm_list = [] for zs, ds in zip(zs_list, ds_list): LensKernel = da.apply_gufunc(lambda dl, zl, Om, ds: wlen(Om, dl, zl, ds), "(), ()-> ()", dl, zl, Om=Om, ds=ds) weights = (LensKernel / (area * nbar)) weights = weights.compute() cat.comm.barrier() if cat.comm.rank == 0: cat.logger.info("source plane %g weights are persisted" % zs) Wmap, Nmap = weighted_map(ipix, npix, weights, localsize, cat.comm) cat.comm.barrier() if cat.comm.rank == 0: cat.logger.info("source plane %g maps generated" % zs) # compute kappa bar # this is a simple integral, but we do not know dl, dz relation # so do it with values from a subsample of particles every = (cat.csize // 100000) kappa1 = Wmap if every == 0: every = 1 # use GatherArray, because it is faster than comm.gather at this scale # (> 4000 ranks on CrayMPI) ssdl = GatherArray(dl[::every].compute(), cat.comm) ssLensKernel = GatherArray(LensKernel[::every].compute(), cat.comm) if cat.comm.rank == 0: arg = ssdl.argsort() ssdl = ssdl[arg] ssLensKernel = ssLensKernel[arg] kappa1bar = numpy.trapz(ssLensKernel, ssdl) else: kappa1bar = None kappa1bar = cat.comm.bcast(kappa1bar) cat.comm.barrier() if cat.comm.rank == 0: cat.logger.info("source plane %g bar computed " % zs) kappa_list.append(kappa1) kappabar_list.append(kappa1bar) Nm_list.append(Nmap) """ # estimate nbar dlmin = dl.min() dlmax = dl.max() volume = (Nmap > 0).sum() / len(Nmap) * 4 / 3 * numpy.pi * (dlmax**3 - dlmin ** 3) """ # returns number rather than delta, since we do not know fsky here. #Nmap = Nmap / cat.csize * cat.comm.allreduce((Nmap > 0).sum()) # to overdensity. return numpy.array(kappa_list), numpy.array(kappabar_list), numpy.array(Nm_list)
def main(ns): if ns.zlmax is None: ns.zlmax = max(ns.zs) zs_list = ns.zs ###### JL hardcode zs_list #zs_list = numpy.arange(ns.zs, 2.21, 0.1) zs_list = ns.zs zlmin = ns.zlmin zlmax = zs_list[-1]#ns.zlmax # no need to be accurate here ds_list = Planck15.comoving_distance(zs_list) path = ns.source cat = BigFileCatalog(path, dataset=ns.dataset) kappa = 0 Nm = 0 kappabar = 0 npix = healpix.nside2npix(ns.nside) localsize = npix * (cat.comm.rank + 1) // cat.comm.size - npix * (cat.comm.rank) // cat.comm.size nbar = (cat.attrs['NC'] ** 3 / cat.attrs['BoxSize'] ** 3 * cat.attrs['ParticleFraction'])[0] # print('DEBUG BoxSize', cat.attrs['BoxSize']) Nsteps = int(numpy.round((zlmax - zlmin) / ns.zstep)) if Nsteps < 2 : Nsteps = 2 z = numpy.linspace(zlmax, zlmin, Nsteps+1, endpoint=True) if cat.comm.rank == 0: cat.logger.info("Splitting data redshift bins %s" % str(z)) kappa_all = numpy.zeros((Nsteps, len(zs_list), localsize)) for i, (z1, z2) in enumerate(zip(z[:-1], z[1:])): import gc gc.collect() if cat.comm.rank == 0: cat.logger.info("nbar = %g, zlmin = %g, zlmax = %g zs = %s" % (nbar, z2, z1, zs_list)) slice = read_range(cat, 1/(1 + z1), 1 / (1 + z2)) if slice.csize == 0: continue if cat.comm.rank == 0: cat.logger.info("read %d particles" % slice.csize) kappa1, kappa1bar, Nm1 = make_kappa_maps(slice, ns.nside, zs_list, ds_list, localsize, nbar) kappa = kappa + kappa1 kappa_all[i] = kappa1 Nm = Nm + Nm1 kappabar = kappabar + kappa1bar cat.comm.barrier() if cat.comm.rank == 0: # use bigfile because it allows concurrent write to different datasets. cat.logger.info("writing to %s", ns.output) # array to get all map slices if cat.comm.rank == 0: kappa1_all = numpy.zeros((Nsteps, int(12*ns.nside**2))) for i, (zs, ds) in enumerate(zip(zs_list, ds_list)): std = numpy.std(cat.comm.allgather(len(kappa[i]))) mean = numpy.mean(cat.comm.allgather(len(kappa[i]))) if cat.comm.rank == 0: cat.logger.info("started gathering source plane %s, size-var = %g, size-bar = %g" % (zs, std, mean)) kappa1 = GatherArray(kappa[i], cat.comm) Nm1 = GatherArray(Nm[i], cat.comm) # get slices of kappa map for j in range(Nsteps): kappa1_allj = GatherArray(kappa_all[j,i], cat.comm) if cat.comm.rank == 0: kappa1_all[j] = kappa1_allj if cat.comm.rank == 0: cat.logger.info("done gathering source plane %s" % zs) if cat.comm.rank == 0: fname = ns.output + "/WL-%02.2f-N%04d" % (zs, ns.nside) cat.logger.info("started writing source plane %s" % zs) with bigfile.File(fname, create=True) as ff: print('DEBUG', kappa1_all.shape, len(kappa1_all), numpy.dtype((kappa1_all.dtype, kappa1_all.shape[1:]))) ds1 = ff.create_from_array("kappa", kappa1, Nfile=1) ds2 = ff.create_from_array("Nm", Nm1, Nfile=1) #ds3 = ff.create_from_array("kappa_all", kappa1_all.T, Nfile=1)#, memorylimit=1024*1024*1024) for d in ds1, ds2:#, ds3: d.attrs['kappabar'] = kappabar[i] d.attrs['nside'] = ns.nside d.attrs['zlmin'] = zlmin d.attrs['zlmax'] = zlmax d.attrs['zstep'] = ns.zstep d.attrs['zs'] = zs d.attrs['ds'] = ds d.attrs['nbar'] = nbar cat.comm.barrier() if cat.comm.rank == 0: # use bigfile because it allows concurrent write to different datasets. cat.logger.info("source plane at %g written. " % zs)
def main(ns): if ns.zlmax is None: ns.zlmax = max(ns.zs) zs_list = ns.zs zlmin = ns.zlmin zlmax = ns.zlmax # no need to be accurate here ds_list = Planck15.comoving_distance(zs_list) path = ns.source #'/global/cscratch1/sd/yfeng1/m3127/desi/1536-9201-40eae2464/lightcone/usmesh/' cat = BigFileCatalog(path, dataset=ns.dataset) kappa = 0 Nm = 0 kappabar = 0 npix = healpix.nside2npix(ns.nside) localsize = npix * (cat.comm.rank + 1) // cat.comm.size - npix * ( cat.comm.rank) // cat.comm.size nbar = (cat.attrs['NC']**3 / cat.attrs['BoxSize']**3 * cat.attrs['ParticleFraction'])[0] Nsteps = int(numpy.round((zlmax - zlmin) / ns.zstep)) if Nsteps < 2: Nsteps = 2 z = numpy.linspace(zlmax, zlmin, Nsteps, endpoint=True) if cat.comm.rank == 0: cat.logger.info("Splitting data redshift bins %s" % str(z)) for z1, z2 in zip(z[:-1], z[1:]): import gc gc.collect() if cat.comm.rank == 0: cat.logger.info("nbar = %g, zlmin = %g, zlmax = %g zs = %s" % (nbar, z2, z1, zs_list)) slice = read_range(cat, 1 / (1 + z1), 1 / (1 + z2)) if slice.csize == 0: continue if cat.comm.rank == 0: cat.logger.info("read %d particles" % slice.csize) kappa1, kappa1bar, Nm1 = make_kappa_maps(slice, ns.nside, zs_list, ds_list, localsize, nbar) kappa = kappa + kappa1 Nm = Nm + Nm1 kappabar = kappabar + kappa1bar cat.comm.barrier() if cat.comm.rank == 0: # use bigfile because it allows concurrent write to different datasets. cat.logger.info("writing to %s", ns.output) for i, (zs, ds) in enumerate(zip(zs_list, ds_list)): std = numpy.std(cat.comm.allgather(len(kappa[i]))) mean = numpy.mean(cat.comm.allgather(len(kappa[i]))) if cat.comm.rank == 0: cat.logger.info( "started gathering source plane %s, size-var = %g, size-bar = %g" % (zs, std, mean)) kappa1 = GatherArray(kappa[i], cat.comm) Nm1 = GatherArray(Nm[i], cat.comm) if cat.comm.rank == 0: cat.logger.info("done gathering source plane %s" % zs) if cat.comm.rank == 0: fname = ns.output + "/WL-%02.2f-N%04d" % (zs, ns.nside) cat.logger.info("started writing source plane %s" % zs) with bigfile.File(fname, create=True) as ff: ds1 = ff.create_from_array("kappa", kappa1, Nfile=1) ds2 = ff.create_from_array("Nm", Nm1, Nfile=1) for d in ds1, ds2: d.attrs['kappabar'] = kappabar[i] d.attrs['nside'] = ns.nside d.attrs['zlmin'] = zlmin d.attrs['zlmax'] = zlmax d.attrs['zs'] = zs d.attrs['ds'] = ds d.attrs['nbar'] = nbar cat.comm.barrier() if cat.comm.rank == 0: # use bigfile because it allows concurrent write to different datasets. cat.logger.info("source plane at %g written. " % zs)