add plotting resources for kernel fit and E/B mode (#24)

* add plotting resources for kernel fit and E/B mode

* add test for ploting (just check if run)

* ran black

* remove try / expect to catch error

* black

* add matplotlib in requirements

* add back try/expect

* remove 3.7 tests and replace by 3.11 (#25)

* add test to see if e/b mode is running

* add comments from Clare

requirements.txt

@@ -1,6 +1,7 @@
 setuptools>=65.5.1
 numpy>=1.16.5
 scipy
+matplotlib
 iminuit>1.3,<1.4; python_version < '3.0'  # iminuit 1.4 fails on python 2.7.
 iminuit>2;        python_version >= '3.0' # iminuit changed API in v2.0.  We use the 2.x API.
 treecorr>=4.2

tests/test_hyp_search.py

@@ -124,6 +124,18 @@ def test_hyperparameter_search_2d():
         )
         gp.initialize(x, y, y_err=y_err)
         gp.solve()
+        # Test if the plot is running (not if correct).
+        if opt == "anisotropic":
+            try:
+                fig = gp.plot_fitted_kernel()
+            except:
+                raise ValueError("Failed to plot fitted kernel")
+        # Test if E/B decomposition is running (not if correct).
+        if opt == "anisotropic":
+            try:
+                e_mode, b_mode, logr = treegp.comp_eb(x[:, 0], x[:, 1], y, y)
+            except:
+                raise ValueError("Failed to compute E/B decomposition")
         # test if found hyperparameters are close the true hyperparameters.
         np.testing.assert_allclose(kernel_skl.theta, gp.kernel.theta, atol=5e-1)
 

treegp/__init__.py

@@ -18,6 +18,8 @@
 
 from .meanify import meanify
 
+from .utils import comp_eb
+
 __all__ = [
     "__version__",
     "__version_info__",
@@ -29,4 +31,5 @@
     "AnisotropicVonKarman",
     "eval_kernel",
     "meanify",
+    "comp_eb",
 ]

treegp/gp_interp.py

@@ -279,3 +279,89 @@ def return_log_likelihood(self, theta=None):
         )
         log_likelihood = logl.log_likelihood(kernel)
         return log_likelihood
+
+    def plot_fitted_kernel(self):
+
+        if self.optimizer in ["none", "log-likehood", "two-pcf"]:
+            raise NotImplementedError(
+                "This method is only available for anisotropic optimizer"
+            )
+        import os
+
+        if os.getenv("GITHUB_ACTIONS") == "true":
+            # Code is running under GitHub Actions
+            import matplotlib
+
+            matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+
+        EXT = [
+            np.min(self._optimizer._2pcf_dist[:, 0]),
+            np.max(self._optimizer._2pcf_dist[:, 0]),
+            np.min(self._optimizer._2pcf_dist[:, 1]),
+            np.max(self._optimizer._2pcf_dist[:, 1]),
+        ]
+
+        CM = plt.cm.seismic
+
+        MAX = np.max(self._optimizer._2pcf)
+        N = int(np.sqrt(len(self._optimizer._2pcf)))
+        fig = plt.figure(figsize=(14, 4))
+        plt.subplots_adjust(wspace=0.5, left=0.07, right=0.95, bottom=0.1, top=0.92)
+
+        plt.subplot(1, 3, 1)
+        plt.imshow(
+            self._optimizer._2pcf.reshape(N, N),
+            extent=EXT,
+            interpolation="nearest",
+            origin="lower",
+            vmin=-MAX,
+            vmax=MAX,
+            cmap=CM,
+        )
+        cbar = plt.colorbar()
+        cbar.formatter.set_powerlimits((0, 0))
+        cbar.update_ticks()
+        cbar.set_label("$\\xi$", fontsize=16)
+        plt.xlabel(r"$\Delta x$", fontsize=16)
+        plt.ylabel(r"$\Delta y$", fontsize=16)
+        plt.title("Measured 2-PCF", fontsize=16)
+
+        plt.subplot(1, 3, 2)
+        plt.imshow(
+            self._optimizer._2pcf_fit.reshape(N, N),
+            extent=EXT,
+            interpolation="nearest",
+            origin="lower",
+            vmin=-MAX,
+            vmax=MAX,
+            cmap=CM,
+        )
+        cbar = plt.colorbar()
+        cbar.formatter.set_powerlimits((0, 0))
+        cbar.update_ticks()
+        cbar.set_label("$\\xi'$", fontsize=16)
+        plt.xlabel(r"$\Delta x$", fontsize=16)
+        plt.title("Fitted 2-PCF", fontsize=16)
+
+        plt.subplot(1, 3, 3)
+        residuals = self._optimizer._2pcf.reshape(
+            N, N
+        ) - self._optimizer._2pcf_fit.reshape(N, N)
+        plt.imshow(
+            residuals,
+            extent=EXT,
+            interpolation="nearest",
+            origin="lower",
+            vmin=-MAX,
+            vmax=MAX,
+            cmap=CM,
+        )
+        cbar = plt.colorbar()
+        cbar.formatter.set_powerlimits((0, 0))
+        cbar.update_ticks()
+        cbar.set_label("$\\xi - \\xi'$", fontsize=16)
+        plt.xlabel(r"$\Delta x$", fontsize=16)
+        plt.title("Difference", fontsize=16)
+
+        return fig

treegp/utils.py

@@ -0,0 +1,111 @@
+import numpy as np
+
+# TO DO: E/B mode computation using numpy should use
+# treecorr instead.
+
+
+def vcorr(x, y, dx, dy, rmin=5.0 / 3600.0, rmax=1.5, dlogr=0.05, maxpts=30000):
+    """
+    Produce angle-averaged 2-point correlation functions of astrometric error
+    for the supplied sample of data, using brute-force pair counting.
+    Output are the following functions:
+    logr - mean log of radius in each bin
+    xi_+ - <vr1 vr2 + vt1 vt2> = <vx1 vx2 + vy1 vy2>
+    xi_- - <vr1 vr2 - vt1 vt2>
+    xi_x - <vr1 vt2 + vt1 vr2>
+    xi_z2 - <vx1 vx2 - vy1 vy2 + 2 i vx1 vy2>
+
+    Parameters
+    ----------
+    x, y : array_like. positions of objects.
+    dx, dy : array_like. astrometric shift.
+    rmin : float. minimum separation in degrees. (default: 5.0 / 3600.0)
+    rmax : float. maximum separation in degrees. (default: 1.5)
+    dlogr : float. bin size in log(r). (default: 0.05)
+    maxpts : int. maximum number of points to use. (default: 30000)
+
+    Returns
+    -------
+    logr, xiplus, ximinus, xicross, xiz2 : array_like.
+    """
+    if len(x) > maxpts:
+        # Subsample array to get desired number of points
+        rate = float(maxpts) / len(x)
+        print("Subsampling rate {:5.3f}%".format(rate * 100.0))
+        use = np.random.random(len(x)) <= rate
+        x = x[use]
+        y = y[use]
+        dx = dx[use]
+        dy = dy[use]
+    print("Length ", len(x))
+    # Get index arrays that make all unique pairs
+    i1, i2 = np.triu_indices(len(x))
+    # Omit self-pairs
+    use = i1 != i2
+    i1 = i1[use]
+    i2 = i2[use]
+    del use
+
+    # Make complex separation vector
+    dr = 1j * (y[i2] - y[i1])
+    dr += x[i2] - x[i1]
+
+    # log radius vector used to bin data
+    logdr = np.log(np.absolute(dr))
+    logrmin = np.log(rmin)
+    bins = int(np.ceil(np.log(rmax / rmin) / dlogr))
+    hrange = (logrmin, logrmin + bins * dlogr)
+    counts = np.histogram(logdr, bins=bins, range=hrange)[0]
+    logr = np.histogram(logdr, bins=bins, range=hrange, weights=logdr)[0] / counts
+
+    # First accumulate un-rotated stats
+    v = dx + 1j * dy
+    vv = dx[i1] * dx[i2] + dy[i1] * dy[i2]
+    xiplus = np.histogram(logdr, bins=bins, range=hrange, weights=vv)[0] / counts
+    vv = v[i1] * v[i2]
+    xiz2 = np.histogram(logdr, bins=bins, range=hrange, weights=vv)[0] / counts
+
+    # Now rotate into radial / perp components
+    vv *= np.conj(dr)
+    vv *= np.conj(dr)
+    dr = dr.real * dr.real + dr.imag * dr.imag
+    vv /= dr
+    del dr
+    ximinus = np.histogram(logdr, bins=bins, range=hrange, weights=vv)[0] / counts
+    xicross = np.imag(ximinus)
+    ximinus = np.real(ximinus)
+
+    return logr, xiplus, ximinus, xicross, xiz2
+
+
+def xiB(logr, xiplus, ximinus):
+    """
+    Return estimate of pure B-mode correlation function
+    """
+    # Integral of d(log r) ximinus(r) from r to infty:
+    dlogr = np.zeros_like(logr)
+    dlogr[1:-1] = 0.5 * (logr[2:] - logr[:-2])
+    tmp = np.array(ximinus) * dlogr
+    integral = np.cumsum(tmp[::-1])[::-1]
+    return 0.5 * (xiplus - ximinus) + integral
+
+
+def comp_eb(u, v, du, dv, **kwargs):
+    """
+    Compute E/B decomposition of astrometric error correlation function
+
+    Parameters
+    ----------
+    u, v : array_like. positions of objects.
+    du, dv : array_like. astrometric shift.
+
+    returns
+    -------
+    xie, xib, logr : array_like. E-mode, B-mode,
+    and log of binned distance separation in 2-point correlation function.
+    """
+
+    logr, xiplus, ximinus, xicross, xiz2 = vcorr(u, v, du, dv, **kwargs)
+    xib = xiB(logr, xiplus, ximinus)
+    xie = xiplus - xib
+    return xie, xib, logr