-
Notifications
You must be signed in to change notification settings - Fork 12
/
mapsim.py
111 lines (107 loc) · 4.44 KB
/
mapsim.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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
from __future__ import division, print_function
""""This module pulls together all the different effects that can go into
CMB map simulations, such as curved sky, lensing, aberration, etc."""
import numpy as np
from . import enmap, curvedsky, lensing, aberration, utils
def rand_map(shape, wcs, ps, lmax=None, curved=True, lens=True, aberrate=True,
beta=None, dir=None, seed=None, dtype=None, verbose=False, recenter=False,
pad=0):
if curved:
return rand_map_curved(shape, wcs, ps, lmax=lmax, lens=lens, aberrate=aberrate,
beta=beta, dir=dir, seed=seed, dtype=dtype, verbose=verbose, recenter=recenter)
else:
return rand_map_flat(shape, wcs, ps, lmax=lmax, lens=lens, aberrate=aberrate,
beta=beta, dir=dir, seed=seed, dtype=dtype, verbose=verbose, recenter=recenter,
pad=pad)
def rand_map_curved(shape, wcs, ps, lmax=None, lens=True, aberrate=True,
beta=None, dir=None, seed=None, dtype=None, verbose=False, recenter=False):
"""Simulate a random curved-sky map. The input spectrum should be
[{phi,T,E,B},{phi,T,E,b},nl] of lens is True, and just [{T,E,B},{T,E,B},nl]
otherwise."""
if dtype is None: dtype = np.float64
if dir is None: dir = aberration.dir_equ
if beta is None: beta = aberration.beta
ctype = np.result_type(dtype,0j)
if verbose: print("Generating alms")
alm = curvedsky.rand_alm(ps, lmax=lmax, seed=seed, dtype=ctype)
# Before any corrections are applied, the pixel positions map
# directly to the raw cmb positions, and there is no induced polarization
# rotation or amplitude modulation.
if verbose: print("Computing observed coordinates")
pos = enmap.posmap(shape, wcs)
ang = enmap.zeros(shape[-2:], wcs)
amp = enmap.zeros(shape[-2:], wcs)+1
# Lensing remaps positions
if lens:
phi_alm, alm = alm[0], alm[1:]
if verbose: print("Computing lensing gradient")
grad = curvedsky.alm2map(phi_alm, enmap.zeros((2,)+shape[-2:], wcs, dtype=dtype), deriv=True)
del phi_alm
if verbose: print("Applying lensing gradient")
pos = enmap.samewcs(lensing.offset_by_grad(pos, grad, pol=True, geodesic=True), pos)
ang += pos[2]
# Aberration remaps positions and modulates amplitudes
if aberrate:
if verbose: print("Computing aberration")
pos = enmap.samewcs(aberration.remap(pos[1::-1], dir=dir, beta=beta, recenter=recenter), pos)
ang += pos[2]
amp *= pos[3]
pos = pos[1::-1]
# Simulate the sky at the observed locations
if verbose: print("Simulating sky signal")
map = curvedsky.alm2map_pos(alm, pos)
del alm, pos
# Apply polarization rotation
if verbose: print("Applying polarization rotation")
map = enmap.rotate_pol(map, ang)
# and modulation
if verbose: print("Applying mouldation")
map *= amp
return map
def rand_map_flat(shape, wcs, ps, lmax=None, lens=True, aberrate=True,
beta=None, dir=None, seed=None, dtype=None, verbose=False, recenter=False,
pad=0):
"""Simulate a random flat-sky map. The input spectrum should be
[{phi,T,E,B},{phi,T,E,b},nl] of lens is True, and just [{T,E,B},{T,E,B},nl]
otherwise."""
if dtype is None: dtype = np.float64
if dir is None: dir = aberration.dir_equ
if beta is None: beta = aberration.beta
ctype = np.result_type(dtype,0j)
if verbose: print("Generating unlensed cmb")
# No position calculation necessary if we're not lensing or aberrating.
if not lens and not aberrate:
return enmap.rand_map(shape, wcs, ps, seed=seed)
# Otherwise we must deal with various displacements
if aberrate: pad += np.pi*beta*1.2
pad_pix = int(pad / enmap.pixsize(shape, wcs)**0.5)
if pad_pix > 0:
if verbose: print("Padding")
template = enmap.zeros(shape, wcs, np.int16)
template, pslice = enmap.pad(template, pad_pix, return_slice=True)
pshape, pwcs = template.shape, template.wcs
else:
pshape, pwcs = shape, wcs
# Simulate (padded) lensing map
if lens:
maps = enmap.rand_map((ps.shape[0],)+pshape[-2:], pwcs, ps)
phi, unlensed = maps[0], maps[1:]
if verbose: print("Lensing")
m = lensing.lens_map_flat(unlensed, phi)
else:
m = enmap.rand_map((ps.shape[0],)+pshape[-2:], pwcs, ps)
# Then handle aberration if necessary
if aberrate:
if verbose: print("Computing aberration displacement")
pos = m.posmap()
pos = enmap.samewcs(aberration.remap(pos[1::-1], dir=dir, beta=beta, recenter=recenter), pos)
amp = pos[3]
pos = pos[1::-1]
if verbose: print("Interpolating aberration")
m = enmap.samewcs(m.at(pos, mask_nan=False), m)
if verbose: print("Applying modulation")
m *= amp
if pad_pix > 0:
if verbose: print("Unpadding")
m = m[pslice]
return m