Merge pull request #52 from LSSTDESC/u/jchiang/add_proper_motion_for_…

…Gaia_stars

add proper motion for gaia stars

data/ci_sample/repo/butler.yaml

@@ -0,0 +1,14 @@
+datastore:
+  cls: lsst.daf.butler.datastores.fileDatastore.FileDatastore
+  records:
+    table: file_datastore_records
+  root: <butlerRoot>
+registry:
+  db: sqlite:///<butlerRoot>/gen3.sqlite3
+  managers:
+    attributes: lsst.daf.butler.registry.attributes.DefaultButlerAttributeManager
+    collections: lsst.daf.butler.registry.collections.synthIntKey.SynthIntKeyCollectionManager
+    datasets: lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID
+    datastores: lsst.daf.butler.registry.bridge.monolithic.MonolithicDatastoreRegistryBridgeManager
+    dimensions: lsst.daf.butler.registry.dimensions.static.StaticDimensionRecordStorageManager
+    opaque: lsst.daf.butler.registry.opaque.ByNameOpaqueTableStorageManager

python/desc/skycatalogs/objects/base_object.py

@@ -355,7 +355,7 @@ class ObjectCollection(Sequence):
     rather than a single number.  There are some additional methods
     '''
     def __init__(self, ra, dec, id, object_type, partition_id, sky_catalog,
-                 region=None, mask=None, readers=None, row_group=0):
+                 region=None, mjd=None, mask=None, readers=None, row_group=0):
         '''
         Parameters
         ra, dec      float, array-like of same length
@@ -368,6 +368,8 @@ def __init__(self, ra, dec, id, object_type, partition_id, sky_catalog,
                      calling open_catalog
         region       maybe be used to determine which objects are in the
                      collection
+        mjd          MJD of the observation epoch.  This is used by
+                     time-varying objects, e.g., SNe, stars.
         mask         indices to be masked off, e.g. in case not all objects
                      in the partition (such as healpixel) are in the region
         readers      may be used to recover properties for objects in this
@@ -399,6 +401,11 @@ def __init__(self, ra, dec, id, object_type, partition_id, sky_catalog,
             self._uniform_object_type = True
 
         self._region = region
+        self._mjd = mjd
+
+    @property
+    def mjd(self):
+        return self._mjd
 
     @property
     def partition_id(self):

python/desc/skycatalogs/objects/gaia_object.py

@@ -3,6 +3,7 @@
 import warnings
 import itertools
 import numpy as np
+import erfa
 import astropy.modeling
 from astropy import units as u
 import galsim
@@ -115,7 +116,7 @@ def set_config(config):
     def get_config():
         return GaiaCollection._gaia_config
 
-    def load_collection(region, skycatalog):
+    def load_collection(region, skycatalog, mjd=None):
         if isinstance(region, Disk):
             ra = lsst.geom.Angle(region.ra, lsst.geom.degrees)
             dec = lsst.geom.Angle(region.dec, lsst.geom.degrees)
@@ -152,9 +153,33 @@ def load_collection(region, skycatalog):
         cat = ref_obj_loader.loadRegion(refcat_region, band).refCat
         df =  cat.asAstropy().to_pandas()
 
-        return GaiaCollection(df, skycatalog, source_type, use_lut)
-
-    def __init__(self, df, sky_catalog, source_type, use_lut):
+        if mjd is None:
+            raise RuntimeError("MJD needs to be provided for Gaia "
+                               "proper motion calculations.")
+        # The erfa.pmsafe code
+        # (https://pyerfa.readthedocs.io/en/latest/api/erfa.pmsafe.html)
+        # expects ra, dec in units of radians and pm_ra, pm_dec
+        # (proper motion) in units of radians/year.  These are the
+        # units used in the refcats files.  However, erfa.pmsafe
+        # expects parallax in arcseconds so convert from the refcats
+        # units of radians:
+        arcsec_per_radian = (1.0*u.radian).to(u.arcsec).value
+        df['parallax'] *= arcsec_per_radian
+        # radial velocity is missing from the Gaia DR2 refcats, so pass
+        # an array of zeros.
+        rv1 = np.zeros(len(df))
+        epNa = 2400000.5  # "part A" of starting and ending epochs for MJDs.
+        ep2b = mjd
+        pm_outputs = erfa.pmsafe(df['coord_ra'], df['coord_dec'],
+                                 df['pm_ra'], df['pm_dec'], df['parallax'],
+                                 rv1, epNa, df['epoch'], epNa, ep2b)
+        # Update ra, dec values.
+        df['coord_ra'] = pm_outputs[0]
+        df['coord_dec'] = pm_outputs[1]
+
+        return GaiaCollection(df, skycatalog, source_type, use_lut, mjd)
+
+    def __init__(self, df, sky_catalog, source_type, use_lut, mjd):
         self.df = df
         self._sky_catalog = sky_catalog
         self._partition_id = None
@@ -165,6 +190,7 @@ def __init__(self, df, sky_catalog, source_type, use_lut):
         self._rdrs = []
         self._object_class = GaiaObject
         self._use_lut = use_lut
+        self._mjd = mjd
 
     @property
     def native_columns(self):
@@ -174,6 +200,10 @@ def native_columns(self):
     def use_lut(self):
         return self._use_lut
 
+    @property
+    def mjd(self):
+        return self._mjd
+
     def __getitem__(self, key):
         if isinstance(key, int) or isinstance(key, np.int64):
             return GaiaObject(self.df.iloc[key], self, key)

python/desc/skycatalogs/skyCatalogs.py

@@ -371,14 +371,14 @@ def toplevel_only(self, object_types):
                 objs_copy.add(parent)
         return objs_copy
 
-    def get_objects_by_region(self, region, obj_type_set=None, datetime=None):
+    def get_objects_by_region(self, region, obj_type_set=None, mjd=None):
         '''
         Parameters
         ----------
         region         region is a named tuple(may be box or circle)
                        or object of type PolygonalRegion
         obj_type_set   Return only these objects. Defaults to all available
-        datetime       Python datetime object. Ignored except for SSO
+        mjd            MJD of observation epoch.
 
         Returns
         -------
@@ -405,20 +405,19 @@ def get_objects_by_region(self, region, obj_type_set=None, datetime=None):
         obj_types = self.toplevel_only(obj_types)
 
         for ot in obj_types:
-            new_list = self.get_object_type_by_region(region, ot,
-                                                      datetime=None)
+            new_list = self.get_object_type_by_region(region, ot, mjd=mjd)
             object_list.append_object_list(new_list)
 
         return object_list
 
-    def get_object_type_by_region(self, region, object_type, datetime=None):
+    def get_object_type_by_region(self, region, object_type, mjd=None):
         '''
         Parameters
         ----------
         region        box, circle or PolygonalRegion. Supported region
                       types made depend on object_type
         object_type   known object type without parent
-        datetime      Ignored for now
+        mjd           MJD of observation epoch.
 
         Returns all objects found
         '''
@@ -431,24 +430,22 @@ def get_object_type_by_region(self, region, object_type, datetime=None):
 
         coll_type = self.cat_cxt.lookup_collection_type(object_type)
         if coll_type is not None:
-            out_list.append_collection(coll_type.load_collection(region,
-                                                                        self))
+            out_list.append_collection(coll_type.load_collection(region, self,
+                                                                 mjd=mjd))
             return out_list
 
         if partition != 'None':
             if partition['type'] == 'healpix':
                 hps = self.get_hps_by_region(region, object_type)
                 for hp in hps:
-                    c = self.get_object_type_by_hp(hp, object_type, region,
-                                                    datetime)
+                    c = self.get_object_type_by_hp(hp, object_type, region, mjd)
                     if len(c) > 0:
                         out_list.append_object_list(c)
                 return out_list
         else:
             raise NotImplementedError(f'Unsupported object type {object_type}')
 
-    def get_object_type_by_hp(self, hp, object_type, region=None,
-                              datetime=None):
+    def get_object_type_by_hp(self, hp, object_type, region=None, mjd=None):
         object_list = ObjectList()
 
         #  Do we need to check more specifically by object type?
@@ -511,6 +508,7 @@ def get_object_type_by_hp(self, hp, object_type, region=None,
                         new_collection = ObjectCollection(ra_c, dec_c, id_c,
                                                           'galaxy', hp, self,
                                                           region=region,
+                                                          mjd=mjd,
                                                           mask=mask,
                                                           readers=the_readers,
                                                           row_group=rg)
@@ -524,6 +522,7 @@ def get_object_type_by_hp(self, hp, object_type, region=None,
                         new_collection = ObjectCollection(ra_c, dec_c, id_c,
                                                           object_type_c[0], hp,
                                                           self, region=region,
+                                                          mjd=mjd,
                                                           mask=mask,
                                                           readers=the_readers,
                                                           row_group=rg)
@@ -536,11 +535,11 @@ def get_object_type_by_hp(self, hp, object_type, region=None,
     #    but if region cut leaves too small a list, read more rowgroups
     #    to achieve a reasonable size list (or exhaust the file)
     def get_object_iterator_by_hp(self, hp, obj_type_set=None,
-                                  max_chunk=None, datetime=None):
+                                  max_chunk=None, mjd=None):
         '''
         Parameters
         ----------
-        datetime       Python datetime object.
+        mjd            MJD of observation epoch
         hp             A healpix id
         obj_type_set   Return only these objects. Defaults to all available
         max_chunk      If specified, iterator will return no more than this

tests/test_gaia_objects.py

@@ -0,0 +1,81 @@
+import os
+from pathlib import Path
+import unittest
+import numpy as np
+import pandas as pd
+import lsst.daf.butler as daf_butler
+from desc.skycatalogs import skyCatalogs
+from desc.skycatalogs.objects import GaiaCollection
+
+
+PACKAGE_DIR = os.path.dirname(os.path.abspath(str(Path(__file__).parent)))
+SKYCATALOG_ROOT = os.path.join(PACKAGE_DIR, "data")
+CATALOG_DIR = os.path.join(PACKAGE_DIR, "data", "ci_sample")
+
+
+CONFIG = {'area_partition': None,
+          'butler_parameters':
+          {'collections': 'refcats',
+           'dstype': 'gaia_dr2_20200414',
+           'repo': os.path.join(CATALOG_DIR, 'repo')},
+          'data_file_type': 'butler_refcat'}
+
+
+def get_gaia_data(butler_params):
+    butler = daf_butler.Butler(butler_params['repo'],
+                               collections=[butler_params['collections']])
+    refs = set(butler.registry.queryDatasets(butler_params['dstype']))
+    return pd.concat(butler.get(_).asAstropy().to_pandas() for _ in refs)
+
+
+class GaiaObjectTestCase(unittest.TestCase):
+    """TestCase class for GaiaObjects"""
+    def setUp(self):
+        self.df = get_gaia_data(CONFIG['butler_parameters'])
+        GaiaCollection.set_config(CONFIG)
+
+    def tearDown(self):
+        pass
+
+    def test_proper_motion(self):
+        ra, dec = 0, 0
+        radius = 1
+        region = skyCatalogs.Disk(ra, dec, radius*3600.0)
+
+        # Use start of proper motion epoch, so ra, dec values should
+        # be unchanged from refcat values.
+        epoch_a_values = set(self.df['epoch'])
+        assert len(epoch_a_values) == 1
+        mjd0 = epoch_a_values.pop()
+        object_list = GaiaCollection.load_collection(region, None, mjd=mjd0)
+        for index in np.random.choice(range(len(object_list)), size=20):
+            obj = object_list[index]
+            gaia_id = obj.id.split('_')[-1]
+            df = self.df.query(f"id=={gaia_id}")
+            row = df.iloc[0]
+            self.assertAlmostEqual(np.degrees(row.coord_ra), obj.ra, places=15)
+            self.assertAlmostEqual(np.degrees(row.coord_dec), obj.dec,
+                                   places=15)
+
+        self.assertEqual(mjd0, object_list.mjd)
+
+        # Use a plausible LSST observation date, and ensure that
+        # calculated proper motion offsets are at least 10% larger
+        # than the naive estimates.
+        mjd = 60584.  # 2024-10-01 00:00:00
+        object_list = GaiaCollection.load_collection(region, None, mjd=mjd)
+        for index in np.random.choice(range(len(object_list)), size=20):
+            obj = object_list[index]
+            gaia_id = obj.id.split('_')[-1]
+            df = self.df.query(f"id=={gaia_id}")
+            row = df.iloc[0]
+            self.assertGreaterEqual(abs(np.radians(obj.ra) - row.coord_ra),
+                                    abs(row['pm_ra']*(mjd - mjd0)/365.)/1.1)
+            self.assertGreaterEqual(abs(np.radians(obj.dec) - row.coord_dec),
+                                    abs(row['pm_dec']*(mjd - mjd0)/365.)/1.1)
+
+        self.assertEqual(mjd, object_list.mjd)
+
+
+if __name__ == '__main__':
+    unittest.main()