Coverage for estats/utils.py: 39%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright (C) 2019 ETH Zurich,
2# Institute for Particle Physics and Astrophysics
3# Author: Dominik Zuercher
5import numpy as np
6import healpy as hp
7import importlib
8import pathlib
9import glob
10import os
11from tqdm import tqdm
13from ekit import logger
15LOGGER = logger.init_logger(__name__)
18def _calc_alms(shear_map_1, shear_map_2, mode='A', method='KS',
19 shear_1_err=[], shear_2_err=[], lmax=None):
20 """
21 Calculates spherical harmonics given two shear maps.
22 Then converts to convergence spin-0 alms using indicated mass
23 mapping technique.
25 :param shear_map_1: First component Healpix shear map
26 :param shear_map_2: Second component Healpix shear map
27 :param mode: If E return E modes only, if B return B modes only,
28 if A return both
29 :param method: Mass mapping method.
30 Currently only Kaiser-Squires supported.
31 :param shear_1_err: Standard deviation of shear
32 component 1 pixel by pixel.
33 :param shear_1_err: Standard deviation of shear
34 component 2 pixel by pixel.
35 :param lmax: Maximum ell to consider for the alms.
36 :return: spherical harmonics E and B modes for spin-0 convergence field.
37 """
39 if lmax is None:
40 lmax = 3 * hp.npix2nside(shear_map_1.size) - 1
42 if (len(shear_map_1) == 0) | (len(shear_map_2) == 0):
43 if mode == 'A':
44 return [], []
45 elif mode == 'E':
46 return []
47 elif mode == 'B':
48 return []
50 if method == 'KS':
51 # Spin spherical harmonic decomposition
53 # numerical problems can lead to some pixels being inf / nan
54 # remove manually
56 is_inf = np.logical_or(np.isinf(shear_map_1), np.isinf(shear_map_2))
57 is_nan = np.logical_or(np.isnan(shear_map_1), np.isnan(shear_map_2))
58 is_bad = np.logical_or(is_inf, is_nan)
59 shear_map_1[is_bad] = 0.0
60 shear_map_2[is_bad] = 0.0
62 # manually set UNSEENs to 0 since healpy messes this up sometimes
63 shear_map_1[np.logical_not(shear_map_1 > hp.UNSEEN)] = 0.0
64 shear_map_2[np.logical_not(shear_map_2 > hp.UNSEEN)] = 0.0
66 alm_T, alm_E, alm_B = hp.map2alm(
67 [np.zeros_like(shear_map_1), shear_map_1, shear_map_2], lmax=lmax)
69 ell = hp.Alm.getlm(lmax)[0]
70 ell[ell == 1] = 0
72 # Multiply by the appropriate factor for spin convertion
73 fac = np.where(np.logical_and(ell != 1, ell != 0),
74 - np.sqrt(((ell + 1.0) * ell) / ((ell + 2.0)
75 * (ell - 1.0))), 0)
76 if mode == 'A':
77 alm_E *= fac
78 alm_B *= fac
79 return alm_E, alm_B
80 elif mode == 'E':
81 alm_E *= fac
82 return alm_E
83 elif mode == 'B':
84 alm_B *= fac
85 return alm_B
86 elif method == 'N0oB':
87 LOGGER.error("N0oB is not supported anymore!")
88 # NOTE: Not sure about healpix sign convention yet
89 if (len(shear_1_err) == 0) | (len(shear_2_err) == 0):
90 raise Exception(
91 "Not all maps given that are needed to use Null B mode "
92 "prior method")
94 # Need std(shear_i) and galaxy count in each pixel (n1, n2, n_gals)
95 noise_map_1 = shear_1_err
96 noise_map_2 = shear_2_err
97 mask = shear_map_1 > hp.UNSEEN
98 shear_map_1[~mask] = 0.0
99 shear_map_2[~mask] = 0.0
100 noise_map_1[~mask] = 0.0
101 noise_map_2[~mask] = 0.0
103 assert np.all(noise_map_1 >= 0.0)
104 assert np.all(noise_map_2 >= 0.0)
105 # upgrade to higher resolution
106 NSIDE_high = 2 * hp.npix2nside(shear_map_1)
107 mask = hp.pixelfunc.ud_grade(mask, nside_out=NSIDE_high)
108 alm_mute = hp.map2alm(noise_map_1)
109 noise_1 = hp.alm2map(alm_mute, nside=NSIDE_high)
110 alm_mute = hp.map2alm(noise_map_2)
111 noise_2 = hp.alm2map(alm_mute, nside=NSIDE_high)
113 noise_1 = np.abs(noise_1)
114 noise_2 = np.abs(noise_2)
116 # checks noise matrix (sometimes when you upgrade your noise matrix to
117 # higher nside, some pixels have a very small value and this might
118 # screw up things)
119 mask_std = noise_2 < np.mean(noise_2[mask]) - 3 * np.std(noise_2[mask])
120 noise_2[mask_std] = np.mean(noise_2[mask]) - 3 * np.std(noise_2[mask])
122 mask_std = noise_1 < np.mean(noise_1[mask]) - 3 * np.std(noise_1[mask])
123 noise_1[mask_std] = np.mean(noise_1[mask]) - 3 * np.std(noise_1[mask])
125 # calc alms of shear
126 alm_mute = hp.map2alm(shear_map_1)
127 e1 = hp.alm2map(alm_mute, nside=NSIDE_high)
128 alm_mute = hp.map2alm(shear_map_2)
129 e2 = hp.alm2map(alm_mute, nside=NSIDE_high)
131 En, Bn = e1 * 0., e1 * 0.
132 power = 1
133 iterations = 15
134 xx = 2
135 ft = np.abs((np.mean(
136 (0.5 * ((noise_1**power)[mask] + (noise_2**power)[mask])))))
137 for i in tqdm(range(iterations)):
138 e1n, e2n = _gk_inv(En, Bn, nside=NSIDE_high,
139 lmax=NSIDE_high * xx - 1)
140 d1 = ft * (e1 - e1n) / (noise_1**power)
141 d2 = ft * (e2 - e2n) / (noise_2**power)
142 REn, RBn, _ = _g2k_sphere_2(
143 d1, d2, mask, nside=NSIDE_high, lmax=NSIDE_high * xx,
144 nosh=True)
145 En = REn + En
146 Bn = RBn + Bn
147 alm_B = hp.map2alm(Bn, lmax=lmax)
148 alm_E = hp.map2alm(En, lmax=lmax)
149 if mode == 'A':
150 return alm_E, alm_B
151 elif mode == 'E':
152 return alm_E
153 elif mode == 'B':
154 return alm_B
155 else:
156 raise Exception(
157 "Mass mapping technique named {} not found.".format(method))
160def import_executable(stat, func):
161 """
162 Import function func from stat plugin.
163 :param verbose: Verbose mode
164 :param logger: Logging instance
165 """
166 executable = importlib.import_module(
167 '.stats.{}'.format(stat), package='estats')
168 executable = getattr(executable, func)
169 return executable
172def _get_plugin_contexts(alloweds=[], typess=[], defaultss=[]):
173 """
174 Gets the contexts for the different statistics plugins
175 """
177 # get all context parameters from the different statistics plugins
178 dir = pathlib.Path(__file__).parent.absolute()
179 plugin_paths = glob.glob('{}/stats/*.py'.format(dir))
180 plugins = []
181 for plugin in plugin_paths:
182 if ('__init__' in plugin):
183 continue
184 else:
185 plugins.append(os.path.basename(plugin.split('.py')[0]))
187 stat_types = {}
189 for plugin in plugins:
190 allowed, defaults, types, stat_type \
191 = import_executable(plugin, 'context')()
192 for ii, a in enumerate(allowed):
193 if a in alloweds:
194 continue
195 else:
196 alloweds.append(a)
197 typess.append(types[ii])
198 defaultss.append(defaults[ii])
200 stat_types[plugin] = stat_type
202 # add extra keywords
203 alloweds.append('stat_types')
204 defaultss.append(stat_types)
205 types.append('dict')
206 return alloweds, typess, defaultss
209def _gk_inv(K, KB, nside, lmax):
210 alms = hp.map2alm(K, lmax=lmax, pol=False) # Spin transform!
211 ell, emm = hp.Alm.getlm(lmax=lmax)
212 kalmsE = alms / (1. * ((ell * (ell + 1.))
213 / ((ell + 2.) * (ell - 1))) ** 0.5)
214 kalmsE[ell == 0] = 0.0
215 alms = hp.map2alm(KB, lmax=lmax, pol=False) # Spin transform!
216 ell, emm = hp.Alm.getlm(lmax=lmax)
217 kalmsB = alms / (1. * ((ell * (ell + 1.))
218 / ((ell + 2.) * (ell - 1))) ** 0.5)
219 kalmsB[ell == 0] = 0.0
220 _, e1t, e2t = hp.alm2map([kalmsE, kalmsE, kalmsB],
221 nside=nside, lmax=lmax, pol=True)
222 return e1t, e2t
225def _g2k_sphere_2(gamma1, gamma2, mask, nside=1024, lmax=2048, nosh=True):
226 """
227 Convert shear to convergence on a sphere. In put are all healpix maps.
228 """
229 gamma1_mask = gamma1 * mask
230 gamma2_mask = gamma2 * mask
231 KQU_masked_maps = [gamma1_mask, gamma1_mask, gamma2_mask]
232 # this fails for some reason
233 alms = hp.map2alm(KQU_masked_maps, lmax=lmax, pol=True) # Spin transform!
234 ell, emm = hp.Alm.getlm(lmax=lmax)
235 if nosh:
236 almsE = alms[1] * 1. * ((ell * (ell + 1.))
237 / ((ell + 2.) * (ell - 1))) ** 0.5
238 almsB = alms[2] * 1. * ((ell * (ell + 1.))
239 / ((ell + 2.) * (ell - 1))) ** 0.5
240 else:
241 almsE = alms[1] * 1.
242 almsB = alms[2] * 1.
243 almsE[ell == 0] = 0.0
244 almsB[ell == 0] = 0.0
245 almsE[ell == 1] = 0.0
246 almsB[ell == 1] = 0.0
247 almssm = [alms[0], almsE, almsB]
248 E_map = hp.alm2map(almssm[1], nside=nside, lmax=lmax, pol=False)
249 B_map = hp.alm2map(almssm[2], nside=nside, lmax=lmax, pol=False)
250 return E_map, B_map, almsE