WIP implementation of search

stellarphot/core.py

@@ -6,6 +6,8 @@
 from astropy.time import Time
 from astropy.units import Quantity
 
+from astroquery.vizier import Vizier
+
 from pydantic import BaseModel, root_validator
 
 import numpy as np
@@ -773,6 +775,95 @@ def catalog_name(self):
     def catalog_source(self):
         return self.meta['catalog_source']
 
+    @classmethod
+    def search_vizier_catalog(cls,
+                              frame_wcs_or_center,
+                              shape,
+                              desired_catalog,
+                              ra_column='RAJ2000',
+                              dec_column='DEJ2000',
+                              radius=0.5 * u.degree,
+                              clip_by_frame=True,
+                              padding=100,
+                              colname_map=None,
+                              prepare_catalog=None):
+        """
+        Return the items from catalog that are within the search radius and
+        (optionally) within the field of view of a frame.
+
+        Parameters
+        ----------
+
+        frame_wcs : `astropy.wcs.WCS`
+            WCS of the image of interest.
+
+        shape : tuple of int
+            Shape of the image of interest.
+
+        desired_catalog : str
+            Name of the catalog to be searched.
+
+        ra_column : str, optional
+            Name of the column in the catalog that contains the RA values.
+            Default is 'RAJ2000'.
+
+        dec_column : str, optional
+            Name of the column in the catalog that contains the Dec values.
+            Default is 'DEJ2000'.
+
+        radius : float, optional
+            Radius, in degrees, around which to search. Default is 0.5.
+
+        clip_by_frame : bool, optional
+            If ``True``, only return items that are within the field of view
+            of the frame. Default is ``True``.
+
+        padding : int, optional
+            Coordinates need to be at least this many pixels in from the edge
+            of the frame to be considered in the field of view. Default value
+            is 100.
+
+        Returns
+        -------
+
+        `astropy.table.Table`
+            Table of catalog information for stars in the field of view.
+
+        """
+
+        if isinstance(frame_wcs_or_center, SkyCoord):
+            # Center was passed in, just use it.
+            center = frame_wcs_or_center
+            if clip_by_frame:
+                raise ValueError('To clip entries by frame you must use '
+                                'a WCS as the first argument.')
+        else:
+            # Find the center of the frame
+            center = frame_wcs_or_center.pixel_to_world(shape[1] / 2,
+                                                        shape[0] / 2)
+
+        # Get catalog via cone search
+        Vizier.ROW_LIMIT = -1  # Set row_limit to have no limit
+        cat = Vizier.query_region(center, radius=radius, catalog=desired_catalog)
+
+        # Vizier always returns list even if there is only one element. Grab that
+        # element.
+        cat = cat[0]
+        cat_coords = SkyCoord(ra=cat[ra_column], dec=cat[dec_column])
+        if clip_by_frame:
+            x, y = frame_wcs_or_center.all_world2pix(cat_coords.ra, cat_coords.dec, 0)
+            in_x = (x >= padding) & (x <= frame_wcs_or_center.pixel_shape[0] - padding)
+            in_y = (y >= padding) & (y <= frame_wcs_or_center.pixel_shape[1] - padding)
+            in_fov = in_x & in_y  # in_frame(frame_wcs_or_center, cat_coords, padding=padding)
+        else:
+            in_fov = np.ones([len(cat_coords)], dtype=bool)
+
+        final_cat = cat[in_fov]
+        if prepare_catalog is not None:
+            final_cat = prepare_catalog(cat[in_fov])
+
+        return cls(input_data=final_cat, colname_map=colname_map)
+
 
 class SourceListData(BaseEnhancedTable):
     """

stellarphot/tests/test_core.py

@@ -4,7 +4,7 @@
 from astropy.table import Table, Column
 from astropy.time import Time
 from astropy.io import ascii
-from astropy.coordinates import EarthLocation
+from astropy.coordinates import EarthLocation, SkyCoord
 from astropy.utils.data import get_pkg_data_filename
 from pydantic import ValidationError
 from stellarphot.core import (Camera, BaseEnhancedTable, PhotometryData,
@@ -487,6 +487,31 @@ def test_catalog_recursive():
                                    catalog_source="Vizier", colname_map=vsx_colname_map)
 
 
+def test_catalog_vizier_search():
+    # Do a cone search with a small enough radius to return exaclty one star,
+    # DQ Psc, which happens to already be in the test data.
+    coordinate = SkyCoord(ra=359.94371 * u.deg, dec=-0.2801 * u.deg)
+    vsx_map = dict(
+        Name='id',
+        RAJ2000='ra',
+        DEJ2000='dec',
+        V='mag',
+        passband='passband'
+    )
+    def prep_vsx(cat):
+        cat['passband'] = 'V'
+        return cat
+
+    my_cat = CatalogData.search_vizier_catalog(coordinate,
+                                               None,
+                                               'B/vsx/vsx',
+                                               radius=0.1*u.arcmin,
+                                               clip_by_frame=False,
+                                               colname_map=vsx_map,
+                                               prepare_catalog=prep_vsx)
+    assert my_cat['id'][0] == 'DQ Psc'
+
+
 # Load test apertures
 test_sl_data = ascii.read(get_pkg_data_filename('data/test_sourcelist.ecsv'),
                              format='ecsv',