store input catalogue label in map metadata

heracles-ec · Nov 12, 2023 · 0154ef5 · 0154ef5
1 parent f4a097d
commit 0154ef5
Show file tree

Hide file tree

Showing 4 changed files with 127 additions and 84 deletions.
diff --git a/heracles/io.py b/heracles/io.py
@@ -36,14 +36,20 @@
 
 
 _METADATA_COMMENTS = {
+    "catalog": "catalog of map",
     "spin": "spin weight of map",
     "kernel": "mapping kernel of map",
+    "nside": "NSIDE parameter of HEALPix map",
     "power": "area power of map",
+    "catalog_1": "catalog of first map",
     "spin_1": "spin weight of first map",
     "kernel_1": "mapping kernel of first map",
+    "nside_1": "NSIDE parameter of first HEALPix map",
     "power_1": "area power of first map",
+    "catalog_2": "catalog of second map",
     "spin_2": "spin weight of second map",
     "kernel_2": "mapping kernel of second map",
+    "nside_2": "NSIDE parameter of second HEALPix map",
     "power_2": "area power of second map",
     "noisbias": "noise bias of spectrum",
 }

diff --git a/heracles/maps.py b/heracles/maps.py
@@ -331,9 +331,11 @@ def mapper(page: "CatalogPage") -> None:
         # set metadata of array
         update_metadata(
             pos,
+            catalog=catalog.label,
             spin=0,
             nbar=nbar,
             kernel="healpix",
+            nside=self.nside,
             power=power,
             bias=bias,
         )
@@ -434,9 +436,11 @@ def mapper(page: "CatalogPage") -> None:
         # set metadata of array
         update_metadata(
             val,
+            catalog=catalog.label,
             spin=0,
             wbar=wbar,
             kernel="healpix",
+            nside=self.nside,
             power=power,
             bias=bias,
         )
@@ -588,9 +592,11 @@ def mapper(page: "CatalogPage") -> None:
         # set metadata of array
         update_metadata(
             val,
+            catalog=catalog.label,
             spin=self.spin,
             wbar=wbar,
             kernel="healpix",
+            nside=self.nside,
             power=power,
             bias=bias,
         )
@@ -627,7 +633,14 @@ def __call__(self, catalog: "Catalog") -> MapData:
             # make a copy for updates to metadata
             vmap = np.copy(vmap)
 
-        update_metadata(vmap, spin=0, kernel="healpix", power=0)
+        update_metadata(
+            vmap,
+            catalog=catalog.label,
+            spin=0,
+            kernel="healpix",
+            nside=self.nside,
+            power=0,
+        )
 
         return vmap
 
@@ -691,7 +704,15 @@ def mapper(page: "CatalogPage") -> None:
             power = 1
 
         # set metadata of arrays
-        update_metadata(wht, spin=0, wbar=wbar, kernel="healpix", power=power)
+        update_metadata(
+            wht,
+            catalog=catalog.label,
+            spin=0,
+            wbar=wbar,
+            kernel="healpix",
+            nside=self.nside,
+            power=power,
+        )
 
         # return the weight map
         return wht
@@ -849,7 +870,6 @@ def transform_maps(
 
     # convert maps to alms, taking care of complex and spin-weighted maps
     for (k, i), m in maps.items():
-        nside = hp.get_nside(m)
         if isinstance(lmax, Mapping):
             _lmax = lmax.get((k, i)) or lmax.get(k)
         else:
@@ -877,8 +897,7 @@ def transform_maps(
             alms = {(f"{k}_E", i): alms[1], (f"{k}_B", i): alms[2]}
 
         for ki, alm in alms.items():
-            if md:
-                update_metadata(alm, nside=nside, **md)
+            update_metadata(alm, **md)
             out[ki] = alm
 
         del m, alms, alm

diff --git a/tests/test_io.py b/tests/test_io.py
@@ -33,7 +33,15 @@ def mock_cls(rng):
     import numpy as np
 
     cl = rng.random(101)
-    cl.dtype = np.dtype(cl.dtype, metadata={"nside_1": 32, "nside_2": 64})
+    cl.dtype = np.dtype(
+        cl.dtype,
+        metadata={
+            "catalog_1": "cat-a.fits",
+            "nside_1": 32,
+            "catalog_2": "cat-b.fits",
+            "nside_2": 64,
+        },
+    )
 
     return {
         ("P", "P", 0, 0): cl,
@@ -191,9 +199,9 @@ def test_write_read_maps(rng, tmp_path):
     v = rng.random(npix)
     g = rng.random((2, npix))
 
-    p.dtype = np.dtype(p.dtype, metadata={"spin": 0})
-    v.dtype = np.dtype(v.dtype, metadata={"spin": 0})
-    g.dtype = np.dtype(g.dtype, metadata={"spin": 0})
+    p.dtype = np.dtype(p.dtype, metadata={"catalog": "cat.fits", "spin": 0})
+    v.dtype = np.dtype(v.dtype, metadata={"catalog": "cat.fits", "spin": 0})
+    g.dtype = np.dtype(g.dtype, metadata={"catalog": "cat.fits", "spin": 2})
 
     maps = {
         ("P", 1): p,

diff --git a/tests/test_maps.py b/tests/test_maps.py
@@ -100,7 +100,13 @@ def test_visibility_map(nside, vmap):
         assert result is not vmap
 
         assert result.shape == (12 * nside_out**2,)
-        assert result.dtype.metadata == {"spin": 0, "kernel": "healpix", "power": 0}
+        assert result.dtype.metadata == {
+            "catalog": catalog.label,
+            "spin": 0,
+            "kernel": "healpix",
+            "nside": nside_out,
+            "power": 0,
+        }
         assert np.isclose(result.mean(), fsky)
 
     # test missing visibility map
@@ -124,31 +130,31 @@ def test_position_map(nside, catalog, vmap):
 
     nbar = 4.0
     assert m.shape == (npix,)
-    assert m.dtype.metadata == pytest.approx(
-        {
-            "spin": 0,
-            "nbar": nbar,
-            "kernel": "healpix",
-            "power": 0,
-            "bias": bias / nbar**2,
-        },
-    )
+    assert m.dtype.metadata == {
+        "catalog": catalog.label,
+        "spin": 0,
+        "nbar": nbar,
+        "kernel": "healpix",
+        "nside": nside,
+        "power": 0,
+        "bias": pytest.approx(bias / nbar**2),
+    }
     np.testing.assert_array_equal(m, 0)
 
     # compute number count map
 
     m = map_catalog(PositionMap(nside, "ra", "dec", overdensity=False), catalog)
 
     assert m.shape == (npix,)
-    assert m.dtype.metadata == pytest.approx(
-        {
-            "spin": 0,
-            "nbar": 4.0,
-            "kernel": "healpix",
-            "power": 1,
-            "bias": bias,
-        },
-    )
+    assert m.dtype.metadata == {
+        "catalog": catalog.label,
+        "spin": 0,
+        "nbar": 4.0,
+        "kernel": "healpix",
+        "nside": nside,
+        "power": 1,
+        "bias": pytest.approx(bias),
+    }
     np.testing.assert_array_equal(m, 4)
 
     # compute overdensity maps with visibility map
@@ -159,30 +165,30 @@ def test_position_map(nside, catalog, vmap):
     m = map_catalog(PositionMap(nside, "ra", "dec"), catalog)
 
     assert m.shape == (12 * nside**2,)
-    assert m.dtype.metadata == pytest.approx(
-        {
-            "spin": 0,
-            "nbar": nbar,
-            "kernel": "healpix",
-            "power": 0,
-            "bias": bias / nbar**2,
-        },
-    )
+    assert m.dtype.metadata == {
+        "catalog": catalog.label,
+        "spin": 0,
+        "nbar": pytest.approx(nbar),
+        "kernel": "healpix",
+        "nside": nside,
+        "power": 0,
+        "bias": pytest.approx(bias / nbar**2),
+    }
 
     # compute number count map with visibility map
 
     m = map_catalog(PositionMap(nside, "ra", "dec", overdensity=False), catalog)
 
     assert m.shape == (12 * nside**2,)
-    assert m.dtype.metadata == pytest.approx(
-        {
-            "spin": 0,
-            "nbar": nbar,
-            "kernel": "healpix",
-            "power": 1,
-            "bias": bias,
-        },
-    )
+    assert m.dtype.metadata == {
+        "catalog": catalog.label,
+        "spin": 0,
+        "nbar": pytest.approx(nbar),
+        "kernel": "healpix",
+        "nside": nside,
+        "power": 1,
+        "bias": pytest.approx(bias),
+    }
 
 
 def test_scalar_map(nside, catalog):
@@ -200,29 +206,29 @@ def test_scalar_map(nside, catalog):
     bias = (4 * np.pi / npix / npix) * v2
 
     assert m.shape == (npix,)
-    assert m.dtype.metadata == pytest.approx(
-        {
-            "spin": 0,
-            "wbar": wbar,
-            "kernel": "healpix",
-            "power": 0,
-            "bias": bias / wbar**2,
-        },
-    )
+    assert m.dtype.metadata == {
+        "catalog": catalog.label,
+        "spin": 0,
+        "wbar": pytest.approx(wbar),
+        "kernel": "healpix",
+        "nside": nside,
+        "power": 0,
+        "bias": pytest.approx(bias / wbar**2),
+    }
     np.testing.assert_array_almost_equal(m, 0)
 
     m = map_catalog(ScalarMap(nside, "ra", "dec", "g1", "w", normalize=False), catalog)
 
     assert m.shape == (npix,)
-    assert m.dtype.metadata == pytest.approx(
-        {
-            "spin": 0,
-            "wbar": wbar,
-            "kernel": "healpix",
-            "power": 1,
-            "bias": bias,
-        },
-    )
+    assert m.dtype.metadata == {
+        "catalog": catalog.label,
+        "spin": 0,
+        "wbar": pytest.approx(wbar),
+        "kernel": "healpix",
+        "nside": nside,
+        "power": 1,
+        "bias": pytest.approx(bias),
+    }
     np.testing.assert_array_almost_equal(m, 0)
 
 
@@ -242,15 +248,15 @@ def test_complex_map(nside, catalog):
     bias = (4 * np.pi / npix / npix) * v2 / 2
 
     assert m.shape == (2, npix)
-    assert m.dtype.metadata == pytest.approx(
-        {
-            "spin": 2,
-            "wbar": wbar,
-            "kernel": "healpix",
-            "power": 0,
-            "bias": bias / wbar**2,
-        },
-    )
+    assert m.dtype.metadata == {
+        "catalog": catalog.label,
+        "spin": 2,
+        "wbar": pytest.approx(wbar),
+        "kernel": "healpix",
+        "nside": nside,
+        "power": 0,
+        "bias": pytest.approx(bias / wbar**2),
+    }
     np.testing.assert_array_almost_equal(m, 0)
 
     m = map_catalog(
@@ -259,15 +265,15 @@ def test_complex_map(nside, catalog):
     )
 
     assert m.shape == (2, npix)
-    assert m.dtype.metadata == pytest.approx(
-        {
-            "spin": 1,
-            "wbar": wbar,
-            "kernel": "healpix",
-            "power": 1,
-            "bias": bias,
-        },
-    )
+    assert m.dtype.metadata == {
+        "catalog": catalog.label,
+        "spin": 1,
+        "wbar": pytest.approx(wbar),
+        "kernel": "healpix",
+        "nside": nside,
+        "power": 1,
+        "bias": pytest.approx(bias),
+    }
     np.testing.assert_array_almost_equal(m, 0)
 
 
@@ -282,9 +288,11 @@ def test_weight_map(nside, catalog):
 
     assert m.shape == (12 * nside**2,)
     assert m.dtype.metadata == {
+        "catalog": catalog.label,
         "spin": 0,
         "wbar": wbar,
         "kernel": "healpix",
+        "nside": nside,
         "power": 0,
     }
     np.testing.assert_array_almost_equal(m, w / wbar)
@@ -293,9 +301,11 @@ def test_weight_map(nside, catalog):
 
     assert m.shape == (12 * nside**2,)
     assert m.dtype.metadata == {
+        "catalog": catalog.label,
         "spin": 0,
         "wbar": wbar,
         "kernel": "healpix",
+        "nside": nside,
         "power": 1,
     }
     np.testing.assert_array_almost_equal(m, w)
@@ -308,9 +318,9 @@ def test_transform_maps(rng):
     npix = 12 * nside**2
 
     t = rng.standard_normal(npix)
-    update_metadata(t, spin=0, a=1)
+    update_metadata(t, spin=0, nside=nside, a=1)
     p = rng.standard_normal((2, npix))
-    update_metadata(p, spin=2, b=2)
+    update_metadata(p, spin=2, nside=nside, b=2)
 
     # single scalar map
     maps = {("T", 0): t}