-
Notifications
You must be signed in to change notification settings - Fork 7
/
polar_sim.py
76 lines (68 loc) · 2.61 KB
/
polar_sim.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import numpy as np, argparse, time, sys
from enlib import enmap, powspec, curvedsky, sharp, wcsutils, memory
parser = argparse.ArgumentParser()
parser.add_argument("powspec")
parser.add_argument("omap")
parser.add_argument("-R", "--radius", type=float, default=4)
parser.add_argument("-r", "--res", type=float, default=0.5)
parser.add_argument("-s", "--supersample", type=float, default=2.0)
parser.add_argument("-m", "--mmax", type=int, default=None)
parser.add_argument("--ncomp", type=int, default=3)
args = parser.parse_args()
deg2rad = np.pi/180
min2rad = np.pi/180/60
ncomp = args.ncomp
class dprint:
def __init__(self, desc):
self.desc = desc
def __enter__(self):
self.t1 = time.time()
def __exit__(self, type, value, traceback):
sys.stderr.write("%6.2f %6.3f %6.3f %s\n" % (time.time()-self.t1,memory.current()/1024.**3, memory.max()/1024.**3, self.desc))
with dprint("spec"):
ps = powspec.read_spectrum(args.powspec)[:ncomp,:ncomp]
# Construct our output coordinates, a zea system. My standard
# constructor doesn't handle pole crossing, so do it manually.
with dprint("construct omap"):
R = args.radius*deg2rad
res = args.res*min2rad
wo = wcsutils.WCS(naxis=2)
wo.wcs.ctype = ["RA---ZEA","DEC--ZEA"]
wo.wcs.crval = [0,90]
wo.wcs.cdelt = [res/deg2rad, res/deg2rad]
wo.wcs.crpix = [1,1]
x, y = wo.wcs_world2pix(0,90-R/deg2rad,1)
y = int(np.ceil(y))
n = 2*y-1
wo.wcs.crpix = [y,y]
omap = enmap.zeros((n,n),wo)
# Construct our projection coordinates this is a CAR system in order
# to make interpolation easy.
with dprint("construct imap"):
ires = np.array([1,1./np.sin(R)])*res/args.supersample
shape, wi = enmap.geometry(pos=[[np.pi/2-R,-np.pi],[np.pi/2,np.pi]], res=ires, proj="car")
imap = enmap.zeros((ncomp,)+shape, wi)
# Define SHT for interpolation pixels
with dprint("construct sht"):
minfo = curvedsky.map2minfo(imap)
lmax_ideal = np.pi/res
ps = ps[:,:,:lmax_ideal]
lmax = ps.shape[-1]
# We do not need all ms when centered on the pole. To reach 1e-10 relative
# error up to R, we need mmax approx 9560*R in radians
mmax = args.mmax or int(R*9560)
ainfo = sharp.alm_info(lmax, mmax)
sht = sharp.sht(minfo, ainfo)
with dprint("curvedsky tot"):
with dprint("rand alm"):
alm = curvedsky.rand_alm(ps, ainfo=ainfo, seed=1, m_major=False)
with dprint("alm2map"):
sht.alm2map(alm[:1], imap[:1].reshape(1,-1))
if ncomp == 3:
sht.alm2map(alm[1:3], imap[1:3,:].reshape(2,-1), spin=2)
del alm
# Make a test map to see if we can project between these
with dprint("project"):
omap = enmap.project(imap, omap.shape, omap.wcs, mode="constant", cval=np.nan)
del imap
enmap.write_map(args.omap, omap)