compute mixing matrices from field definitions

ntessore · ntessore · commit 24b960d8ae0e · 2024-01-02T11:52:35.000Z
diff --git a/examples/example.ipynb b/examples/example.ipynb
@@ -452,7 +452,9 @@
     "* A position field (`\"P\"`) for angular clustering and galaxy-galaxy lensing;\n",
     "* A shear field (`\"G\"`) for cosmic shear and galaxy-galaxy lensing.\n",
     "\n",
-    "We also specify that the shear field should flip the sign of the \"G2\" column by adding a minus sign."
+    "We also specify that the shear field should flip the sign of the \"G2\" column by adding a minus sign.\n",
+    "\n",
+    "Finally, we define the optional names of the weight functions (`\"V\"`, `\"W\"`) of the fields. These will be used further down below to compute mixing matrices for the fields."
    ]
   },
   {
@@ -465,8 +467,8 @@
     "lonlat = ('RIGHT_ASCENSION', 'DECLINATION')\n",
     "\n",
     "fields = {\n",
-    "    'P': Positions(*lonlat),\n",
-    "    'G': Shears(*lonlat, 'G1', '-G2', 'WEIGHT'),\n",
+    "    'P': Positions(*lonlat, weight=\"V\"),\n",
+    "    'G': Shears(*lonlat, 'G1', '-G2', 'WEIGHT', weight=\"W\"),\n",
     "}"
    ]
   },
@@ -520,7 +522,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2a0ea2d6fc02466fb00059925c44032e",
+       "model_id": "63d6558897214c0fbad71ed81e70f7db",
        "version_major": 2,
        "version_minor": 0
       },
@@ -534,13 +536,13 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">/Users/ntessore/code/heracles-ec/heracles/heracles/fields.py:247: UserWarning: position and visibility maps have \n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">/Users/ntessore/code/heracles-ec/heracles/heracles/fields.py:254: UserWarning: position and visibility maps have \n",
        "different NSIDE\n",
        "  warnings.warn(\"position and visibility maps have different NSIDE\")\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "/Users/ntessore/code/heracles-ec/heracles/heracles/fields.py:247: UserWarning: position and visibility maps have \n",
+       "/Users/ntessore/code/heracles-ec/heracles/heracles/fields.py:254: UserWarning: position and visibility maps have \n",
        "different NSIDE\n",
        "  warnings.warn(\"position and visibility maps have different NSIDE\")\n"
       ]
@@ -690,7 +692,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "73b9269235fb47b9b79d26f456f8962e",
+       "model_id": "ed443fdc7af948e1b3dfeb1e6740b8a9",
        "version_major": 2,
        "version_minor": 0
       },
@@ -923,7 +925,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7eea7c123c304cfeb5f2150ab6db9012",
+       "model_id": "2d17b05fc35f4e95aadb24b0e1fe989c",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1027,7 +1029,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "32abfabf091a441b9c0b716b3d476aae",
+       "model_id": "1cccc3fe4f304b4fa8b25ecc5da7972f",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1114,22 +1116,22 @@
     }
    ],
    "source": [
-    "mms = mixing_matrices(cls_mm, l1max=lmax, l2max=lmax, l3max=lmax_mm)"
+    "mms = mixing_matrices(fields, cls_mm, l1max=lmax, l2max=lmax, l3max=lmax_mm)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The mixing matrices are returned in a toc dict with slightly different names: They use a two letter code that indicates which angular power spectra are mixed by a given matrix:\n",
+    "The mixing matrices are returned in a toc dict with familiar names: They correspond to the _output_ angular power spectra after mixing by a given matrix:\n",
     "\n",
-    "* `00` for `(P, P) → (P, P)`;\n",
-    "* `0+` for `(P, E) → (P, E)` and `(P, B) → (P, B)`;\n",
-    "* `++` for `(E, E) → (E, E)` and `(B, B) → (B, B)`;\n",
-    "* `--` for `(E, E) → (B, B)` and `(B, B) → (E, E)`;\n",
-    "* `+-` for `(E, B) → (E, B)`.\n",
+    "* `P, P` for `(P, P) → (P, P)`;\n",
+    "* `P, G_E` for `(P, G_E) → (P, G_E)`, as well as `(P, G_B) → (P, G_B)`;\n",
+    "* `G_E, G_E` for `(G_E, G_E) → (G_E, G_E)`, as well as `(G_B, G_B) → (G_B, G_B)`;\n",
+    "* `G_B, G_B` for `(G_E, G_E) → (G_B, G_B)`, as well as `(G_B, G_B) → (G_E, G_E)`;\n",
+    "* `G_E, G_B` for `(G_E, G_B) → (G_E, G_B)`.\n",
     "\n",
-    "For more details, see the paper by Brown, Castro & Taylor (2005). Note that the `+-` mixing matrix here is not the $W^{+-}$ matrix from said paper."
+    "For more details, see e.g. the paper by Brown, Castro & Taylor (2005)."
    ]
   },
   {
@@ -1140,21 +1142,19 @@
     {
      "data": {
       "text/plain": [
-       "[('00', 0, 0),\n",
-       " ('0+', 0, 0),\n",
-       " ('00', 0, 1),\n",
-       " ('0+', 0, 1),\n",
-       " ('00', 0, 2),\n",
-       " ('0+', 0, 2),\n",
-       " ('00', 0, 3),\n",
-       " ('0+', 0, 3),\n",
-       " ('00', 0, 4),\n",
-       " ('0+', 0, 4),\n",
-       " ('00', 0, 5),\n",
-       " ('0+', 0, 5),\n",
-       " ('00', 0, 6),\n",
-       " ('0+', 0, 6),\n",
-       " ('00', 0, 7)]"
+       "[('P', 'P', 0, 0),\n",
+       " ('P', 'G_E', 0, 0),\n",
+       " ('P', 'P', 0, 1),\n",
+       " ('P', 'G_E', 0, 1),\n",
+       " ('P', 'P', 0, 2),\n",
+       " ('P', 'G_E', 0, 2),\n",
+       " '...',\n",
+       " ('G_E', 'G_B', 8, 9),\n",
+       " ('P', 'P', 9, 9),\n",
+       " ('P', 'G_E', 9, 9),\n",
+       " ('G_E', 'G_E', 9, 9),\n",
+       " ('G_B', 'G_B', 9, 9),\n",
+       " ('G_E', 'G_B', 9, 9)]"
       ]
      },
      "execution_count": 37,
@@ -1163,7 +1163,7 @@
     }
    ],
    "source": [
-    "list(mms.keys())[:15]"
+    "list(mms.keys())[:6] + ['...'] + list(mms.keys())[-6:]"
    ]
   },
   {
@@ -1234,7 +1234,7 @@
     }
    ],
    "source": [
-    "plt.imshow(mms['++', i, i], cmap='binary',\n",
+    "plt.imshow(mms['G_E', 'G_E', i, i], cmap='binary',\n",
     "            norm=mpl.colors.LogNorm(vmin=1e-7))\n",
     "plt.colorbar(pad=0.025, fraction=0.0465)\n",
     "plt.show()"
@@ -1569,15 +1569,15 @@
     "        cl = camb_cls[f'W{i+1}xW{i+j+1}']\n",
     "\n",
     "        if s1.source_type == 'counts' and s2.source_type == 'counts':\n",
-    "            theory_cls['P', 'P', i1, i2] = mms['00', i1, i2] @ cl\n",
+    "            theory_cls['P', 'P', i1, i2] = mms['P', 'P', i1, i2] @ cl\n",
     "        elif s1.source_type == 'lensing' and s2.source_type == 'lensing':\n",
-    "            theory_cls['G_E', 'G_E', i1, i2] = mms['++', i1, i2] @ (cl*fl**2)\n",
-    "            theory_cls['G_B', 'G_B', i1, i2] = mms['--', i1, i2] @ (cl*fl**2)\n",
+    "            theory_cls['G_E', 'G_E', i1, i2] = mms['G_E', 'G_E', i1, i2] @ (cl*fl**2)\n",
+    "            theory_cls['G_B', 'G_B', i1, i2] = mms['G_B', 'G_B', i1, i2] @ (cl*fl**2)\n",
     "        elif s1.source_type == 'counts' and s2.source_type == 'lensing':\n",
-    "            theory_cls['P', 'G_E', i1, i2] = mms['0+', i1, i2] @ (cl*fl)\n",
+    "            theory_cls['P', 'G_E', i1, i2] = mms['P', 'G_E', i1, i2] @ (cl*fl)\n",
     "            theory_cls['P', 'G_B', i1, i2] = np.zeros_like(cl)\n",
     "        elif s1.source_type == 'lensing' and s2.source_type == 'counts':\n",
-    "            theory_cls['P', 'G_E', i2, i1] = mms['0+', i2, i1] @ (cl*fl)\n",
+    "            theory_cls['P', 'G_E', i2, i1] = mms['P', 'G_E', i2, i1] @ (cl*fl)\n",
     "            theory_cls['P', 'G_B', i2, i1] = np.zeros_like(cl)"
    ]
   },
diff --git a/heracles/fields.py b/heracles/fields.py
@@ -82,10 +82,11 @@ def __init_subclass__(cls, *, spin: int | None = None) -> None:
                 break
         cls.__ncol = (ncol - nopt, ncol)
 
-    def __init__(self, *columns: str) -> None:
+    def __init__(self, *columns: str, weight: str | None = None) -> None:
         """Initialise the field."""
         super().__init__()
         self.__columns = self._init_columns(*columns) if columns else None
+        self.__weight = weight
         self._metadata: dict[str, Any] = {}
         if (spin := self.__spin) is not None:
             self._metadata["spin"] = spin
@@ -140,6 +141,11 @@ def spin(self) -> int:
             raise ValueError(msg)
         return spin
 
+    @property
+    def weight(self) -> str | None:
+        """Name of the weight function for this field."""
+        return self.__weight
+
     @abstractmethod
     async def __call__(
         self,
@@ -187,9 +193,10 @@ def __init__(
         *columns: str,
         overdensity: bool = True,
         nbar: float | None = None,
+        weight: str | None = None,
     ) -> None:
         """Create a position field."""
-        super().__init__(*columns)
+        super().__init__(*columns, weight=weight)
         self.__overdensity = overdensity
         self.__nbar = nbar
 
diff --git a/heracles/twopoint.py b/heracles/twopoint.py
@@ -25,7 +25,6 @@
 
 import healpy as hp
 import numpy as np
-from convolvecl import mixmat, mixmat_eb
 
 from .core import TocDict, toc_match, update_metadata
 
@@ -182,65 +181,80 @@ def debias_cls(cls, bias=None, *, inplace=False):
     return out
 
 
-def mixing_matrices(cls, *, l1max=None, l2max=None, l3max=None):
-    """compute mixing matrices from a set of cls"""
+def mixing_matrices(fields, cls, *, l1max=None, l2max=None, l3max=None, out=None):
+    """compute mixing matrices for fields from a set of cls"""
 
-    logger.info("computing two-point mixing matrices for %d cl(s)", len(cls))
-    t = time.monotonic()
+    from convolvecl import mixmat, mixmat_eb
+
+    # output dictionary if not provided
+    if out is None:
+        out = TocDict()
 
-    logger.info("using L1MAX = %s, L2MAX = %s, L3MAX = %s", l1max, l2max, l3max)
+    # inverse mapping of weights to fields
+    weights = {}
+    for key, field in fields.items():
+        if field.weight is not None:
+            if field.weight not in weights:
+                weights[field.weight] = {}
+            weights[field.weight][key] = field
 
-    # set of computed mixing matrices
-    mms = TocDict()
+    # keep track of combinations that have been done already
+    done = set()
 
     # go through the toc dict of cls and compute mixing matrices
-    # which mixing matrix is computed depends on the combination of V/W maps
+    # which mixing matrix is computed depends on the whts mapping
     for (k1, k2, i1, i2), cl in cls.items():
+        # if the weights are not named then skip this cl
+        try:
+            fields1 = weights[k1]
+            fields2 = weights[k2]
+        except KeyError:
+            continue
+
+        # deal with structured cl arrays
         if cl.dtype.names is not None:
             cl = cl["CL"]
-        if k1 == "V" and k2 == "V":
-            logger.info("computing 00 mixing matrix for bins %s, %s", i1, i2)
-            w00 = mixmat(cl, l1max=l1max, l2max=l2max, l3max=l3max)
-            mms["00", i1, i2] = w00
-        elif k1 == "V" and k2 == "W":
-            logger.info("computing 0+ mixing matrix for bins %s, %s", i1, i2)
-            w0p = mixmat(cl, l1max=l1max, l2max=l2max, l3max=l3max, spin=(0, 2))
-            mms["0+", i1, i2] = w0p
-        elif k1 == "W" and k2 == "V":
-            logger.info("computing 0+ mixing matrix for bins %s, %s", i2, i1)
-            w0p = mixmat(cl, l1max=l1max, l2max=l2max, l3max=l3max, spin=(2, 0))
-            mms["0+", i2, i1] = w0p
-        elif k1 == "W" and k2 == "W":
-            logger.info("computing ++, --, +- mixing matrices for bins %s, %s", i1, i2)
-            wpp, wmm, wpm = mixmat_eb(
-                cl,
-                l1max=l1max,
-                l2max=l2max,
-                l3max=l3max,
-                spin=(2, 2),
-            )
-            mms["++", i1, i2] = wpp
-            mms["--", i1, i2] = wmm
-            mms["+-", i1, i2] = wpm
-        else:
-            logger.warning(
-                "computing unknown %s x %s mixing matrix for bins %s, %s",
-                k1,
-                k2,
-                i1,
-                i2,
-            )
-            w = mixmat(cl, l1max=l1max, l2max=l2max, l3max=l3max)
-            mms[f"{k1}{k2}", i1, i2] = w
 
-    logger.info(
-        "computed %d mm(s) in %s",
-        len(mms),
-        timedelta(seconds=(time.monotonic() - t)),
-    )
+        # compute mixing matrices for all fields of this weight combination
+        for f1, f2 in product(fields1, fields2):
+            # check if this combination has been done already
+            if (f1, f2, i1, i2) in done or (f2, f1, i2, i1) in done:
+                continue
+            # otherwise, mark it as done
+            done.add((f1, f2, i1, i2))
+
+            # get spins of fields
+            spin1, spin2 = fields1[f1].spin, fields2[f2].spin
+
+            # if any spin is zero, then there is no need for E/B decomposition
+            if spin1 == 0 or spin2 == 0:
+                mm = mixmat(
+                    cl,
+                    l1max=l1max,
+                    l2max=l2max,
+                    l3max=l3max,
+                    spin=(spin1, spin2),
+                )
+                name1 = f1 if spin1 == 0 else f"{f1}_E"
+                name2 = f2 if spin2 == 0 else f"{f2}_E"
+                out[name1, name2, i1, i2] = mm
+                del mm
+            else:
+                # E/B decomposition for mixing matrix
+                mm_ee, mm_bb, mm_eb = mixmat_eb(
+                    cl,
+                    l1max=l1max,
+                    l2max=l2max,
+                    l3max=l3max,
+                    spin=(spin1, spin2),
+                )
+                out[f"{f1}_E", f"{f2}_E", i1, i2] = mm_ee
+                out[f"{f1}_B", f"{f2}_B", i1, i2] = mm_bb
+                out[f"{f1}_E", f"{f2}_B", i1, i2] = mm_eb
+                del mm_ee, mm_bb, mm_eb
 
     # return the toc dict of mixing matrices
-    return mms
+    return out
 
 
 def bin2pt(arr, bins, name, *, weights=None):
diff --git a/tests/test_fields.py b/tests/test_fields.py
@@ -115,9 +115,10 @@ async def __call__(self):
     with pytest.raises(ValueError, match="accepts 2 to 3 columns"):
         TestField("lon")
 
-    f = TestField("lon", "lat")
+    f = TestField("lon", "lat", weight="W")
 
     assert f.columns == ("lon", "lat", None)
+    assert f.weight == "W"
 
 
 def test_visibility(mapper, vmap):
diff --git a/tests/test_twopoint.py b/tests/test_twopoint.py
