Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

U/jrbogart/gaia slice #95

Merged
merged 3 commits into from
May 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 4 additions & 11 deletions skycatalogs/objects/base_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
29 changes: 13 additions & 16 deletions skycatalogs/objects/gaia_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
11 changes: 2 additions & 9 deletions skycatalogs/objects/sso_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]