def add_noise(self, obs):
pm = obs.mapp.pm
field = obs.mapp
def tf(k):
kmesh = sum(ki**2 for ki in k)**0.5
mask = [numpy.ones_like(ki) for ki in k]
mask[2] *= abs(k[2]) >= self.kmin
mask = numpy.prod(mask)
mask2 = numpy.ones_like(mask)
if self.angle > 0:
kperp = (k[0]**2 + k[1]**2)**0.5
kangle = abs(k[2]) / (kperp + 1e-10)
angles = (numpy.arctan(kangle) * 180 / numpy.pi)
mask2[angles < self.angle] = 0
fgmask = mask * mask2
return fgmask
fieldc = field.r2c()
fieldc.apply(lambda k, v: nyquist_mask(tf(k), v) * v, out=Ellipsis)
field = fieldc.c2r()
obsret = Observable(mapp=field, s=obs.s, d=obs.d)
return obsret
def save2dphoto(Nmu=4, numd=1e-2, aa=None, scatter=False):
if numd > 1e-2:
print('Too high number density')
sys.exit()
num = int(numd * bs**3)
if aa is None: aas = [0.3333, 0.2000, 0.1429]
else: aas = [aa]
for ia, aa in enumerate(aas):
zz = 1 / aa - 1
sigz = lambda z: 120 * ((1 + z) / 5)**-0.5
##
cat = BigFileCatalog(
'/global/cscratch1/sd/chmodi/m3127/H1mass/highres/10240-9100/fastpm_%0.4f/Hcat-Numd-%04d/'
% (aa, 1e-2 * 1e4))
if scatter:
pos = cat['Position'][:num].compute()
dz = np.random.normal(0, sigz(zz), size=pos[:, -1].size)
pos[:, -1] += dz
layout = pm.decompose(pos)
hmesh = pm.paint(pos, layout=layout)
else:
pos = cat['Position'][:num].compute()
layout = pm.decompose(pos)
hmesh = pm.paint(pos, layout=layout)
def tf(k): #Photoz smoothing
kmesh = sum(ki**2 for ki in k)**0.5
kmesh[kmesh == 0] = 1
mumesh = k[2] / kmesh
weights = np.exp(-kmesh**2 * mumesh**2 * sigz(zz)**2 / 2.)
return weights
hmesh /= hmesh.cmean()
if not scatter:
hmeshc = hmesh.r2c()
hmeshc.apply(lambda k, v: nyquist_mask(tf(k), v) * v, out=Ellipsis)
hmesh = hmeshc.c2r()
ph = FFTPower(hmesh, mode='2d', Nmu=Nmu).power
#
for iw, wopt in enumerate(['opt', 'pess']):
#for iw, wopt in enumerate(['opt']):
for it, thopt in enumerate(['opt', 'pess', 'reas']):
#for it, thopt in enumerate([ 'reas']):
if rank == 0: print(aa, wopt, thopt)
angle = np.round(mapn.wedge(zz, att=wopt, angle=True), 0)
#dpath = '/global/cscratch1/sd/chmodi/m3127/21cm_cleaning/recon/fastpm_%0.4f/wedge_kmin%0.2f_ang%0.1f/'%(aa, 0.03, angle)
dpath = '/global/cscratch1/sd/chmodi/m3127/21cm_cleaning/recon/fastpm_%0.4f/wedge_kmin%0.2f_%s/' % (
aa, 0.03, wopt)
dpath += 'L%04d-N%04d-R//thermal-%s-hex/ZA/opt_s999_h1massA_fourier_rsdpos/' % (
bs, nc, thopt)
if scatter:
ofolder = '../../data/ZArecon-rsd/kmin-003_wedge-%s/thermal-%s-hex/photog-Numd%04d-Nmu%d/' % (
wopt, thopt, numd * 1e4, Nmu)
else:
ofolder = '../../data/ZArecon-rsd/kmin-003_wedge-%s/thermal-%s-hex/photo-Numd%04d-Nmu%d/' % (
wopt, thopt, numd * 1e4, Nmu)
try:
os.makedirs(ofolder)
except:
pass
if rank == 0: print(ofolder)
if args.pp:
datapp = BigFileMesh(dpath + '/dataw_up/', 'mapp').paint()
bpaths = [dpath + 'upsample2/%d-0.00//best-fit' % nc2] + [
dpath + 'upsample2/%d-0.00//%04d/fit_p/' % (nc2, i)
for i in range(100, 50, -20)
]
else:
datapp = BigFileMesh(dpath + '/dataw/', 'mapp').paint()
bpaths = [dpath + '%d-0.00//best-fit' % nc] + [
dpath + '%d-0.00//%04d/fit_p/' % (nc, i)
for i in range(100, 50, -20)
]
for path in bpaths:
if os.path.isdir(path):
break
if rank == 0: print(path)
bfit = BigFileMesh(path, 'mapp').paint()
pxrh = FFTPower(hmesh, second=bfit, mode='2d', Nmu=Nmu).power
pxwh = FFTPower(hmesh, second=datapp, mode='2d', Nmu=Nmu).power
fname = ofolder + 'photo-L%04d_%0.4f.txt' % (bs, aa)
if args.pp: fname = fname[:-4] + '-up.txt'
np.savetxt(fname, ph['power'].real)
fname = ofolder + 'xdataw-L%04d_%0.4f.txt' % (bs, aa)
if args.pp: fname = fname[:-4] + '-up.txt'
np.savetxt(fname, pxwh['power'].real)
fname = ofolder + 'xrecon-L%04d_%0.4f.txt' % (bs, aa)
if args.pp: fname = fname[:-4] + '-up.txt'
np.savetxt(fname, pxrh['power'].real)