From 62cefb86ba1a5745e2838e2709738d1c23e1f4be Mon Sep 17 00:00:00 2001
From: Sean McGuire <123987820+smcguire-cmu@users.noreply.github.com>
Date: Fri, 1 Nov 2024 17:41:59 -0400
Subject: [PATCH] Add plot moc method (#491)

* add plot moc method

* fix return

* fix return type

* fix hipscat

* remove plot_moc
---
 src/lsdb/catalog/dataset/healpix_dataset.py | 21 ++++++++++++++-------
 tests/lsdb/catalog/test_catalog.py          | 12 ++++++++++++
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/lsdb/catalog/dataset/healpix_dataset.py b/src/lsdb/catalog/dataset/healpix_dataset.py
index 57e9a87c..9aa5f15e 100644
--- a/src/lsdb/catalog/dataset/healpix_dataset.py
+++ b/src/lsdb/catalog/dataset/healpix_dataset.py
@@ -11,12 +11,14 @@
 import nested_pandas as npd
 import numpy as np
 import pandas as pd
+from astropy.visualization.wcsaxes import WCSAxes
 from dask.delayed import Delayed, delayed
 from hats.catalog.healpix_dataset.healpix_dataset import HealpixDataset as HCHealpixDataset
 from hats.inspection.visualize_catalog import plot_healpix_map
 from hats.pixel_math import HealpixPixel
 from hats.pixel_math.healpix_pixel_function import get_pixel_argsort
 from hats.pixel_math.spatial_index import SPATIAL_INDEX_COLUMN
+from matplotlib.figure import Figure
 from pandas._libs import lib
 from pandas._typing import AnyAll, Axis, IndexLabel
 from pandas.api.extensions import no_default
@@ -399,7 +401,7 @@ def skymap(
         projection="MOL",
         plotting_args: Dict | None = None,
         **kwargs,
-    ):
+    ) -> tuple[Figure, WCSAxes]:
         """Plot a skymap of an aggregate function applied over each partition
 
         Args:
@@ -424,19 +426,24 @@ def skymap(
             plotting_args = {}
         return plot_healpix_map(img, projection=projection, **plotting_args)
 
-    def plot_pixels(self, projection: str = "MOL", **kwargs):
+    def plot_pixels(self, projection: str = "MOL", **kwargs) -> tuple[Figure, WCSAxes]:
         """Create a visual map of the pixel density of the catalog.
 
         Args:
-            projection (str) The map projection to use. Valid values include:
-                - moll - Molleweide projection (default)
-                - gnom - Gnomonic projection
-                - cart - Cartesian projection
-                - orth - Orthographic projection
+            projection (str) The map projection to use. Available projections listed at
+            https://docs.astropy.org/en/stable/wcs/supported_projections.html
             kwargs (dict): additional keyword arguments to pass to plotting call.
         """
         return self.hc_structure.plot_pixels(projection=projection, **kwargs)
 
+    def plot_coverage(self, **kwargs) -> tuple[Figure, WCSAxes]:
+        """Create a visual map of the coverage of the catalog.
+
+        Args:
+            kwargs: additional keyword arguments to pass to hats.Catalog.plot_moc
+        """
+        return self.hc_structure.plot_moc(**kwargs)
+
     def to_hats(
         self,
         base_catalog_path: str | Path | UPath,
diff --git a/tests/lsdb/catalog/test_catalog.py b/tests/lsdb/catalog/test_catalog.py
index 1abd4aea..c2648a2e 100644
--- a/tests/lsdb/catalog/test_catalog.py
+++ b/tests/lsdb/catalog/test_catalog.py
@@ -531,6 +531,18 @@ def test_plot_pixels(small_sky_order1_catalog, mocker):
     )
 
 
+# pylint: disable=no-member
+def test_plot_coverage(small_sky_order1_catalog, mocker):
+    mocker.patch("hats.catalog.healpix_dataset.healpix_dataset.plot_moc")
+    small_sky_order1_catalog.plot_coverage()
+
+    hc.catalog.healpix_dataset.healpix_dataset.plot_moc.assert_called_once()
+    assert (
+        hc.catalog.healpix_dataset.healpix_dataset.plot_moc.call_args[0][0]
+        == small_sky_order1_catalog.hc_structure.moc
+    )
+
+
 def test_square_bracket_columns(small_sky_order1_catalog):
     columns = ["ra", "dec", "id"]
     column_subset = small_sky_order1_catalog[columns]