general mcms implemented (#237)

Author: damonge

pymaster/__init__.py

@@ -99,6 +99,7 @@
     compute_coupled_cell_flat,
     compute_full_master_flat,
     uncorr_noise_deprojection_bias,
+    get_general_coupling_matrix,
 )
 from pymaster.covariance import (  # noqa
     NmtCovarianceWorkspace,

pymaster/tests/test_master.py

@@ -767,3 +767,36 @@ def test_fkp_normalization_errors():
     # Two potentially different catalog fields
     with pytest.raises(ValueError):
         nmt.NmtWorkspace.from_fields(fc, fcb, b, normalization='FKP')
+
+
+def test_general_mcmc():
+    # Create disc mask
+    nside = 256
+    nell = 3*nside
+    npix = hp.nside2npix(nside)
+    mask = np.zeros(npix)
+    mask[hp.query_disc(nside, [1, 0, 0], np.radians(30))] = 1
+
+    f0 = nmt.NmtField(mask, None, spin=0)
+    f2 = nmt.NmtField(mask, None, spin=2)
+    b = nmt.NmtBin.from_nside_linear(nside, nlb=10)
+    w00 = nmt.NmtWorkspace.from_fields(f0, f0, b)
+    w02 = nmt.NmtWorkspace.from_fields(f0, f2, b)
+    w22 = nmt.NmtWorkspace.from_fields(f2, f2, b)
+
+    mcm00 = w00.get_coupling_matrix()
+    mcm02 = w02.get_coupling_matrix().reshape([nell, 2, nell, 2])
+    mcm02 = mcm02[:, 0, :, 0]
+    mcm22 = w22.get_coupling_matrix().reshape([nell, 4, nell, 4])
+    mcmee_ee = mcm22[:, 0, :, 0]
+    mcmee_bb = mcm22[:, 0, :, 3]
+    mcm22 = mcmee_ee + mcmee_bb
+
+    pclm = hp.anafast(mask)
+    m00 = nmt.get_general_coupling_matrix(pclm, 0, 0, 0, 0)
+    m02 = nmt.get_general_coupling_matrix(pclm, 0, 2, 0, 2)
+    m22 = nmt.get_general_coupling_matrix(pclm, 2, 2, 2, 2)
+
+    assert np.amax(np.fabs(m00-mcm00)/np.amax(mcm00)) < 1E-10
+    assert np.amax(np.fabs(m02-mcm02)/np.amax(mcm02)) < 1E-10
+    assert np.amax(np.fabs(m22-mcm22)/np.amax(mcm22)) < 1E-10

pymaster/workspaces.py

@@ -1172,8 +1172,33 @@ def compute_full_master_flat(f1, f2, b, cl_noise=None, cl_guess=None,
 
 
 def get_general_coupling_matrix(pcl_mask, s1, s2, n1, n2):
-    """ Returns a general mode-coupling matrix
-    TODO
+    """ Returns a general mode-coupling matrix of the form
+
+    .. math::
+      M_{\\ell \\ell'}=\\sum_{\\ell''}
+      \\frac{(2\\ell'+1)(2\\ell''+1)}{4\\pi}
+      \\tilde{C}^{uv}_\\ell
+      \\left(\\begin{array}{ccc}
+      \\ell & \\ell' & \\ell'' \\\\
+      n_1 & -s_1 & s_1-n_1
+      \\end{array}\\right)
+      \\left(\\begin{array}{ccc}
+      \\ell & \\ell' & \\ell'' \\\\
+      n_2 & -s_2 & s_2-n_2
+      \\end{array}\\right)
+
+    Args:
+        pcl_mask (`array`): 1D array containing the power spectrum
+          of the masks :math:`\\tilde{C}_\\ell^{uw}`.
+        s1 (:obj:`int`): spin index :math:`s_1` above.
+        s2 (:obj:`int`): spin index :math:`s_2` above.
+        n1 (:obj:`int`): spin index :math:`n_1` above.
+        n2 (:obj:`int`): spin index :math:`n_2` above.
+
+    Returns:
+        (`array`): 2D array of shape ``[nl, nl]``, where ``nl`` is
+        the size of ``pcl_mask``, containing the mode-coupling
+        matrix for multipoles from 0 to ``nl-1``.
     """
 
     lmax = len(pcl_mask)-1