Merge pull request #95 from LSSTDESC/u/jrbogart/gaia_slice

U/jrbogart/gaia slice

skycatalogs/objects/base_object.py

@@ -600,17 +600,10 @@ def __getitem__(self, key):
         elif type(key) == slice:
             if key.start is None:
                 key.start = 0
-            ixdata = [i for i in range(min(key.stop, len(self._ra)))]
-            ixes = itertools.islice(ixdata, key.start, key.stop, key.step)
-            return [self._object_class(self._ra[i], self._dec[i], self._id[i],
-                                       object_type, self, i)
-                    for i in ixes]
+            return [self.__getitem__(i) for i in range(self.__len__())[key]]
 
         elif type(key) == tuple and isinstance(key[0], Iterable):
-            #  check it's a list of int-like?
-            return [self._object_class(self._ra[i], self._dec[i], self._id[i],
-                                       object_type, self, i)
-                    for i in key[0]]
+            return[self.__getitem__(i) for i in key[0]]
 
     def get_partition_id(self):
         return self._partition_id
@@ -725,8 +718,8 @@ def __getitem__(self, key):
         If key is a slice return a list of object
         '''
         one_only = isinstance(key, int) or isinstance(key, np.int64)
-        is_slice = type(key) == slice
-        is_list = type(key) == tuple and isinstance(key[0], Iterable)
+        is_slice = isinstance(key, slice)
+        is_list = isinstance(key, tuple) and isinstance(key[0], Iterable)
 
         if one_only:
             start = key

skycatalogs/objects/gaia_object.py

@@ -85,9 +85,9 @@ def __init__(self, obj_pars, parent_collection, index):
         # dec = np.degrees(obj_pars['coord_dec'])
         ra = obj_pars['ra_deg']
         dec = obj_pars['dec_deg']
-        # Form the object id from the GAIA catalog id with the string
-        # 'gaia_dr2_' prepended.
-        id_prefix = GaiaCollection.get_config()['id_prefix']
+        # Form the object id from the GAIA catalog id with a string
+        # like 'gaia_dr2_' prepended.
+        id_prefix = GaiaCollection._id_prefix
         obj_id = f"{id_prefix}{obj_pars['id']}"
         super().__init__(ra, dec, obj_id, 'gaia_star',
                          belongs_to=parent_collection, belongs_index=index)
@@ -235,7 +235,7 @@ class GaiaCollection(ObjectCollection):
     @classmethod
     def set_config(cls, config):
         GaiaCollection._gaia_config = config
-
+        GaiaCollection._id_prefix = config['id_prefix']
     @classmethod
     def get_config(cls):
         return GaiaCollection._gaia_config
@@ -350,7 +350,6 @@ def __init__(self, df, sky_catalog, source_type, use_lut, mjd):
         self.df = df
         self._sky_catalog = sky_catalog
         self._partition_id = None
-        self._id = np.array([f"gaia_dr2_{df.iloc[key]['id']}" for key in range(len(df))])
         self._mask = None
         self._object_type_unique = source_type
         self._use_lut = use_lut
@@ -372,22 +371,20 @@ def mjd(self):
         return self._mjd
 
     def __getitem__(self, key):
+        cols = ('id', 'ra_deg', 'dec_deg', 'phot_bp_mean_flux',
+                'phot_rp_mean_flux')
         if isinstance(key, int) or isinstance(key, np.int64):
-            row = {col: self.df[col][key] for col in ('id', 'ra_deg',
-                                                      'dec_deg',
-                                                      'phot_bp_mean_flux',
-                                                      'phot_rp_mean_flux')}
+            row = {col: self.df[col][key] for col in cols}
             return GaiaObject(row, self, key)
 
-        elif type(key) == slice:
-            ixdata = [i for i in range(min(key.stop, len(self._id)))]
-            ixes = itertools.islice(ixdata, key.start, key.stop, key.step)
-            return [self._object_class(self.df.iloc[i], self, i) for i in ixes]
+        elif isinstance(key, slice):
+            # ixdata = [i for i in range(min(key.stop, len(self.df['id'])))]
+            # ixes = itertools.islice(ixdata, key.start, key.stop, key.step)
+            # return [self.__getitem__(i) for i in ixes]
+            return [self.__getitem__(i) for i in range(len(self.df))[key]]
 
         elif type(key) == tuple and isinstance(key[0], Iterable):
-            #  check it's a list of int-like?
-            return [self._object_class(self.df.iloc[i], self,
-                                       i) for i in key[0]]
+            return [self.__getitem__(i) for i in key[0]]
 
     def __len__(self):
         return len(self.df)

skycatalogs/objects/sso_object.py

@@ -163,14 +163,7 @@ def __getitem__(self, key):
         elif type(key) == slice:
             if key.start is None:
                 key.start = 0
-            ixdata = [i for i in range(min(key.stop, len(self._ra)))]
-            ixes = itertools.islice(ixdata, key.start, key.stop, key.step)
-            return [self._object_class(self._ra[i], self._dec[i], self._id[i],
-                                       object_type, self, i, self._mjds[i])
-                    for i in ixes]
+            return [self.__getitem__(i) for i in range(self.__len__())[key]]
 
         elif type(key) == tuple and isinstance(key[0], Iterable):
-            #  check it's a list of int-like?
-            return [self._object_class(self._ra[i], self._dec[i], self._id[i],
-                                       object_type, self, i, self._mjds[i])
-                    for i in key[0]]
+            return [self.__getitem__(i) for i in key[0]]