Skip to content

Commit 4b4be9b

Browse files
GerenjieGerenjie
authored
Merge pull request #266 from lsst/tickets/DM-48611
DM-48611 : Add additional SSSource columns
2 parents 4010046 + f92575a commit 4b4be9b

File tree

2 files changed

+80
-26
lines changed

2 files changed

+80
-26
lines changed

‎python/lsst/ap/association/ssoAssociation.py‎

Lines changed: 79 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,10 @@
2525

2626
from astropy import units as u
2727
from astropy.table import Table
28+
from astropy.coordinates import SkyCoord
2829
import healpy as hp
2930
import numpy as np
31+
from numpy.polynomial.chebyshev import Chebyshev, chebval
3032
import pandas as pd
3133
from scipy.spatial import cKDTree
3234

@@ -103,24 +105,63 @@ def run(self, diaSourceCatalog, solarSystemObjects, exposure):
103105
ref_time = mjd_midpoint - solarSystemObjects["tmin"].values[0]
104106

105107
solarSystemObjects['obs_position'] = solarSystemObjects.apply(lambda row: np.array([
106-
np.polynomial.chebyshev.chebval(ref_time, row['obs_x_poly']),
107-
np.polynomial.chebyshev.chebval(ref_time, row['obs_y_poly']),
108-
np.polynomial.chebyshev.chebval(ref_time, row['obs_z_poly'])
108+
chebval(ref_time, row['obs_x_poly']),
109+
chebval(ref_time, row['obs_y_poly']),
110+
chebval(ref_time, row['obs_z_poly'])
111+
]), axis=1)
112+
solarSystemObjects['obs_velocity'] = solarSystemObjects.apply(lambda row: np.array([
113+
chebval(ref_time, Chebyshev(row['obs_x_poly']).deriv().coef),
114+
chebval(ref_time, Chebyshev(row['obs_y_poly']).deriv().coef),
115+
chebval(ref_time, Chebyshev(row['obs_z_poly']).deriv().coef),
109116
]), axis=1)
110-
111117
solarSystemObjects['obj_position'] = solarSystemObjects.apply(lambda row: np.array([
112-
np.polynomial.chebyshev.chebval(ref_time, row['obj_x_poly']),
113-
np.polynomial.chebyshev.chebval(ref_time, row['obj_y_poly']),
114-
np.polynomial.chebyshev.chebval(ref_time, row['obj_z_poly'])
118+
chebval(ref_time, row['obj_x_poly']),
119+
chebval(ref_time, row['obj_y_poly']),
120+
chebval(ref_time, row['obj_z_poly'])
121+
]), axis=1)
122+
solarSystemObjects['obj_velocity'] = solarSystemObjects.apply(lambda row: np.array([
123+
chebval(ref_time, Chebyshev(row['obj_x_poly']).deriv().coef),
124+
chebval(ref_time, Chebyshev(row['obj_y_poly']).deriv().coef),
125+
chebval(ref_time, Chebyshev(row['obj_z_poly']).deriv().coef),
115126
]), axis=1)
116-
117127
vector = np.vstack(solarSystemObjects['obj_position'].values
118128
- solarSystemObjects['obs_position'].values)
119129
solarSystemObjects[['ra', 'dec']] = np.vstack(hp.vec2ang(vector, lonlat=True)).T
120130
solarSystemObjects['obs_position_x'], solarSystemObjects['obs_position_y'], \
121131
solarSystemObjects['obs_position_z'] = np.array(list(solarSystemObjects['obs_position'].values)).T
122-
solarSystemObjects['obj_position_x'], solarSystemObjects['obj_position_y'], \
123-
solarSystemObjects['obj_position_z'] = np.array(list(solarSystemObjects['obj_position'].values)).T
132+
solarSystemObjects['heliocentricX'], solarSystemObjects['heliocentricY'], \
133+
solarSystemObjects['heliocentricZ'] = np.array(list(solarSystemObjects['obj_position'].values)).T
134+
solarSystemObjects['obs_velocity_x'], solarSystemObjects['obs_velocity_y'], \
135+
solarSystemObjects['obs_velocity_z'] = np.array(list(solarSystemObjects['obs_velocity'].values)).T
136+
solarSystemObjects['heliocentricVX'], solarSystemObjects['heliocentricVY'], \
137+
solarSystemObjects['heliocentricVZ'] = np.array(list(solarSystemObjects['obj_velocity'].values)).T
138+
solarSystemObjects['topocentric_position'], solarSystemObjects['topocentric_velocity'] = (
139+
solarSystemObjects['obj_position'] - solarSystemObjects['obs_position'],
140+
solarSystemObjects['obj_velocity'] - solarSystemObjects['obs_velocity'],
141+
)
142+
solarSystemObjects['topocentricX'], solarSystemObjects['topocentricY'], \
143+
solarSystemObjects['topocentricZ'] = (
144+
np.array(list(solarSystemObjects['topocentric_position'].values)).T
145+
)
146+
solarSystemObjects['topocentricVX'], solarSystemObjects['topocentricVY'], \
147+
solarSystemObjects['topocentricVZ'] = (
148+
np.array(list(solarSystemObjects['topocentric_velocity'].values)).T
149+
)
150+
solarSystemObjects['heliocentricVX'], solarSystemObjects['heliocentricVY'], \
151+
solarSystemObjects['heliocentricVZ'] = np.array(list(solarSystemObjects['obj_velocity'].values)).T
152+
solarSystemObjects['heliocentricDist'], solarSystemObjects['topocentricDist'] = (
153+
solarSystemObjects['obj_position'].apply(np.linalg.norm),
154+
solarSystemObjects['topocentric_position'].apply(np.linalg.norm)
155+
)
156+
solarSystemObjects['phaseAngle'] = (
157+
(solarSystemObjects['obj_position'] * solarSystemObjects['topocentric_position']
158+
/ solarSystemObjects['heliocentricDist'] / solarSystemObjects['topocentricDist'])
159+
.apply(np.sum).apply(np.arccos).apply(np.degrees)
160+
)
161+
162+
stateVectorColumns = ['heliocentricX', 'heliocentricY', 'heliocentricZ', 'heliocentricVX',
163+
'heliocentricVY', 'heliocentricVZ', 'topocentricX', 'topocentricY',
164+
'topocentricZ', 'topocentricVX', 'topocentricVY', 'topocentricVZ']
124165

125166
maskedObjects = self._maskToCcdRegion(
126167
solarSystemObjects,
@@ -144,24 +185,28 @@ def run(self, diaSourceCatalog, solarSystemObjects, exposure):
144185
# picks the DiaSource with the shortest distance. We can do something
145186
# fancier later.
146187
ssSourceData = []
147-
ras, decs, expected_ras, expected_decs = [], [], [], []
188+
ras, decs, residual_ras, residual_decs, dia_ids = [], [], [], [], []
148189
diaSourceCatalog["ssObjectId"] = 0
190+
source_column = 'id'
191+
if 'diaSourceId' in diaSourceCatalog.columns:
192+
source_column = 'diaSourceId'
149193
for index, ssObject in maskedObjects.iterrows():
150194
ssoVect = self._radec_to_xyz(ssObject["ra"], ssObject["dec"])
151195
# Which DIA Sources fall within r?
152196
dist, idx = tree.query(ssoVect, distance_upper_bound=maxRadius)
153197
if len(idx) == 1 and np.isfinite(dist[0]):
154198
nFound += 1
155199
diaSourceCatalog.loc[diaSourceCatalog.index[idx[0]], "ssObjectId"] = ssObject["ssObjectId"]
156-
ssSourceData.append(ssObject[["ssObjectId", "obs_position_x", "obs_position_y",
157-
"obs_position_z", "obj_position_x", "obj_position_y",
158-
"obj_position_z"]].values)
200+
ssSourceData.append(ssObject[["ssObjectId", "phaseAngle", "heliocentricDist",
201+
"topocentricDist"] + stateVectorColumns].values)
159202
dia_ra = diaSourceCatalog.loc[diaSourceCatalog.index[idx[0]], "ra"]
160203
dia_dec = diaSourceCatalog.loc[diaSourceCatalog.index[idx[0]], "dec"]
204+
dia_id = diaSourceCatalog.loc[diaSourceCatalog.index[idx[0]], source_column]
161205
ras.append(dia_ra)
162206
decs.append(dia_dec)
163-
expected_ras.append(ssObject["ra"])
164-
expected_decs.append(ssObject["dec"])
207+
dia_ids.append(dia_id)
208+
residual_ras.append(dia_ra - ssObject["ra"])
209+
residual_decs.append(dia_dec - ssObject["dec"])
165210
maskedObjects.loc[index, 'associated'] = True
166211
else:
167212
maskedObjects.loc[index, 'associated'] = False
@@ -171,18 +216,26 @@ def run(self, diaSourceCatalog, solarSystemObjects, exposure):
171216
self.metadata['nExpectedSsObjects'] = nSolarSystemObjects
172217
assocSourceMask = diaSourceCatalog["ssObjectId"] != 0
173218
unAssocObjectMask = np.logical_not(maskedObjects['associated'].values)
174-
ssSourceData = pd.DataFrame(ssSourceData, columns=["ssObjectId", "obs_position_x", "obs_position_y",
175-
"obs_position_z", "obj_position_x",
176-
"obj_position_y", "obj_position_z"])
219+
ssSourceData = pd.DataFrame(ssSourceData,
220+
columns=[
221+
"ssObjectId", "phaseAngle", "heliocentricDist", "topocentricDist"
222+
] + stateVectorColumns)
177223
ssSourceData["ra"] = ras
178224
ssSourceData["dec"] = decs
179-
ssSourceData["expected_ra"] = expected_ras
180-
ssSourceData["expected_dec"] = expected_decs
181-
unassociatedObjects = maskedObjects[unAssocObjectMask].drop(columns=['obs_x_poly', 'obs_y_poly',
182-
'obs_z_poly', 'obj_x_poly',
183-
'obj_y_poly', 'obj_z_poly',
184-
'obs_position', 'obj_position',
185-
'associated'])
225+
ssSourceData["residualRa"] = residual_ras
226+
ssSourceData["residualDec"] = residual_decs
227+
ssSourceData[source_column] = dia_ids
228+
coords = SkyCoord(ra=ssSourceData['ra'].values * u.deg, dec=ssSourceData['dec'].values * u.deg)
229+
ssSourceData['galacticL'] = coords.galactic.l.deg
230+
ssSourceData['galacticB'] = coords.galactic.b.deg
231+
ssSourceData['eclipticLambda'] = coords.barycentrictrueecliptic.lon.deg
232+
ssSourceData['eclipticBeta'] = coords.barycentrictrueecliptic.lat.deg
233+
columns_to_drop = [
234+
"obs_position", "obs_velocity", "obj_position", "obj_velocity", "topocentric_position",
235+
"topocentric_velocity", "obs_x_poly", "obs_y_poly", "obs_z_poly", "obj_x_poly", "obj_y_poly",
236+
"obj_z_poly", "associated"
237+
]
238+
unassociatedObjects = maskedObjects[unAssocObjectMask].drop(columns=columns_to_drop)
186239
return pipeBase.Struct(
187240
ssoAssocDiaSources=diaSourceCatalog[assocSourceMask].reset_index(drop=True),
188241
unAssocDiaSources=diaSourceCatalog[~assocSourceMask].reset_index(drop=True),

‎tests/test_ssoAssociationTask.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ def setUp(self):
7474
pd.DataFrame([[45, 45]], columns=["ra", "dec"])],
7575
ignore_index=True)
7676
self.testDiaSources["ssObjectId"] = 0
77+
self.testDiaSources["diaSourceId"] = [i for i in range(len(self.testDiaSources))]
7778

7879
# Grab a subset to treat as solar system objects
7980
self.testSsObjects = pd.DataFrame()

0 commit comments

Comments
 (0)