Merge pull request #174 from litebird/noise_weighting

NET weight in mapping.py

docs/source/mapmaking.rst

@@ -82,7 +82,17 @@ detectors::
 Binner
 ------
 
-To be written.
+Once you have generated a set of observations, either on a single
+process or distributed over several mpi processes, you can create a 
+simple binned map with the function :func:`.make_bin_map`. This function
+takes: a single (or a list) of :class:`.Observations`, the Healpix
+resolution of the output map (``nside``) and produces a coadded map.
+It assumes white noise and each detector gets weighted by 
+:math:`1 / NET^2`. Optionally, if the parameter do_covariance is True,
+it can output the covariance matrix in an array of shape 
+`(12 * nside * nside, 3, 3)`. This is how should be called::
+
+    map, cov = lbs.make_bin_map(obs, 128, do_covariance=True)
 
 
 Destriper
@@ -145,6 +155,11 @@ Here is the complete source code of the example and the result:
 API reference
 -------------
 
+.. automodule:: litebird_sim.mapping
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 .. automodule:: litebird_sim.destriper
     :members:
     :undoc-members:

litebird_sim/mapping.py

@@ -15,26 +15,28 @@
 
 
 @njit
-def _accumulate_map_and_info(tod, pix, psi, info):
+def _accumulate_map_and_info(tod, pix, psi, weights, info):
     # Fill the upper triangle of the information matrix and use the lower
     # triangle for the RHS of the map-making equation
     assert tod.shape == pix.shape == psi.shape
-    tod = tod.ravel()
-    pix = pix.ravel()
-    psi = psi.ravel()
-    for d, p, a in zip(tod, pix, psi):
-        cos = np.cos(2 * a)
-        sin = np.sin(2 * a)
-        info_pix = info[p]
-        info_pix[0, 0] += 1.0
-        info_pix[0, 1] += cos
-        info_pix[0, 2] += sin
-        info_pix[1, 1] += cos * cos
-        info_pix[1, 2] += sin * cos
-        info_pix[2, 2] += sin * sin
-        info_pix[1, 0] += d
-        info_pix[2, 0] += d * cos
-        info_pix[2, 1] += d * sin
+
+    ndets = tod.shape[0]
+
+    for idet in range(ndets):
+        for d, p, a in zip(tod[idet], pix[idet], psi[idet]):
+            one = 1.0 / np.sqrt(weights[idet])
+            cos = np.cos(2 * a) / np.sqrt(weights[idet])
+            sin = np.sin(2 * a) / np.sqrt(weights[idet])
+            info_pix = info[p]
+            info_pix[0, 0] += one * one
+            info_pix[0, 1] += one * cos
+            info_pix[0, 2] += one * sin
+            info_pix[1, 1] += cos * cos
+            info_pix[1, 2] += sin * cos
+            info_pix[2, 2] += sin * sin
+            info_pix[1, 0] += d * one * one
+            info_pix[2, 0] += d * cos * one
+            info_pix[2, 1] += d * sin * one
 
 
 def _extract_map_and_fill_info(info):
@@ -47,7 +49,9 @@ def _extract_map_and_fill_info(info):
     return rhs
 
 
-def make_bin_map(obss: Union[Observation, List[Observation]], nside):
+def make_bin_map(
+    obss: Union[Observation, List[Observation]], nside, do_covariance=False
+):
     """Bin Map-maker
 
     Map a list of observations
@@ -65,11 +69,15 @@ def make_bin_map(obss: Union[Observation, List[Observation]], nside):
             If the observations are distributed over some communicator(s), they
             must share the same group processes.
         nside (int): HEALPix nside of the output map
-    Returs:
+        do_covariance (bool): optional, if true it returns also covariance
+
+    Returns:
         array: T, Q, U maps (stacked). The shape is `(3, 12 * nside * nside)`.
             All the detectors of all the observations contribute to the map.
             If the observations are distributed over some communicator(s), all
-            the processes (contribute and) hold a copy of the map
+            the processes (contribute and) hold a copy of the map.
+            Optionally can return the covariance matrix in an array of shape
+            `(12 * nside * nside, 3, 3)`
     """
     n_pix = hp.nside2npix(nside)
     info = np.zeros((n_pix, 3, 3))
@@ -80,7 +88,12 @@ def make_bin_map(obss: Union[Observation, List[Observation]], nside):
         obs_list = obss
 
     for obs in obs_list:
-        _accumulate_map_and_info(obs.tod, obs.pixind, obs.psi, info)
+        try:
+            weights = obs.sampling_rate_shz * obs.net_ukrts ** 2
+        except AttributeError:
+            weights = np.ones(obs.n_detectors)
+
+            _accumulate_map_and_info(obs.tod, obs.pixind, obs.psi, weights, info)
 
     if all([obs.comm is None for obs in obs_list]) or not mpi.MPI_ENABLED:
         # Serial call
@@ -99,10 +112,19 @@ def make_bin_map(obss: Union[Observation, List[Observation]], nside):
 
     rhs = _extract_map_and_fill_info(info)
     try:
-        return np.linalg.solve(info, rhs)
+        res = np.linalg.solve(info, rhs)
     except np.linalg.LinAlgError:
         cond = np.linalg.cond(info)
         res = np.full_like(rhs, hp.UNSEEN)
         mask = cond < COND_THRESHOLD
         res[mask] = np.linalg.solve(info[mask], rhs[mask])
-        return res
+
+    if do_covariance:
+        try:
+            return res.T, np.linalg.inv(info)
+        except np.linalg.LinAlgError:
+            covmat = np.full_like(info, hp.UNSEEN)
+            covmat[mask] = np.linalg.inv(info[mask])
+            return res.T, covmat
+    else:
+        return res.T

test/test_mapping.py

@@ -30,7 +30,11 @@ def test_accumulate_map_and_info():
     res_info[:, 2, 1] = np.bincount(pix, tod * np.sin(2 * psi))
 
     info = np.zeros((2, 3, 3))
-    mapping._accumulate_map_and_info(tod, pix, psi, info)
+    weights = np.ones(1)
+    tod = np.expand_dims(tod, axis=0)
+    psi = np.expand_dims(psi, axis=0)
+    pix = np.expand_dims(pix, axis=0)
+    mapping._accumulate_map_and_info(tod, pix, psi, weights, info)
 
     assert np.allclose(res_info, info)
 
@@ -55,7 +59,7 @@ def test_make_bin_map_api_simulation(tmp_path):
         )
     )
     instr = lbs.InstrumentInfo(name="core", spin_boresight_angle_rad=np.radians(65))
-    det = lbs.DetectorInfo(name="foo", sampling_rate_hz=10)
+    det = lbs.DetectorInfo(name="foo", sampling_rate_hz=10, net_ukrts=1.0)
     obss = sim.create_observations(detectors=[det])
     pointings = lbs.get_pointings(
         obss[0],
@@ -104,9 +108,9 @@ def test_make_bin_map_basic_mpi():
         obs.psi = psi.reshape(2, 5)
 
     obs.set_n_blocks(n_blocks_time=obs.comm.size, n_blocks_det=1)
-    res = mapping.make_bin_map([obs], 1)[: len(res_map)]
+    res = mapping.make_bin_map([obs], 1).T[: len(res_map)]
     assert np.allclose(res, res_map)
 
     obs.set_n_blocks(n_blocks_time=1, n_blocks_det=obs.comm.size)
-    res = mapping.make_bin_map([obs], 1)[: len(res_map)]
+    res = mapping.make_bin_map([obs], 1).T[: len(res_map)]
     assert np.allclose(res, res_map)

test/test_scan_map.py

@@ -84,7 +84,7 @@ def test_scan_map():
         maps=in_map,
         input_map_in_galactic=False,
     )
-    out_map1 = lbs.make_bin_map(obs1, nside).T
+    out_map1 = lbs.make_bin_map(obs1, nside)
 
     times = obs2.get_times()
     obs2.pixind = np.empty(obs2.tod.shape, dtype=np.int32)
@@ -104,7 +104,7 @@ def test_scan_map():
             + np.sin(2 * obs2.psi[idet, :]) * maps[2, obs2.pixind[idet, :]]
         )
 
-    out_map2 = lbs.make_bin_map(obs2, nside).T
+    out_map2 = lbs.make_bin_map(obs2, nside)
 
     np.testing.assert_allclose(
         out_map1, in_map["Boresight_detector_T"], rtol=1e-6, atol=1e-6
@@ -135,7 +135,7 @@ def test_scan_map():
         input_map_in_galactic=True,
         fill_psi_and_pixind_in_obs=True,
     )
-    out_map1 = lbs.make_bin_map(obs1, nside).T
+    out_map1 = lbs.make_bin_map(obs1, nside)
 
     np.testing.assert_allclose(
         out_map1, in_map_G["Boresight_detector_T"], rtol=1e-6, atol=1e-6