Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import warnings 

2import numpy as np 

3import pandas as pd 

4 

5__all__ = ['Orbits'] 

6 

7 

8class Orbits(object): 

9 """Orbits reads, checks for required values, and stores orbit parameters for moving objects. 

10 

11 Instantiate the class and then use readOrbits or setOrbits to set the orbit values. 

12 

13 self.orbits stores the orbital parameters, as a pandas dataframe. 

14 self.dataCols defines the columns required, although objId, H, g, and sed_filename are optional. 

15 """ 

16 def __init__(self): 

17 self.orbits = None 

18 self.orb_format = None 

19 

20 # Specify the required columns/values in the self.orbits dataframe. 

21 # Which columns are required depends on self.orb_format. 

22 self.dataCols = {} 

23 self.dataCols['COM'] = ['objId', 'q', 'e', 'inc', 'Omega', 'argPeri', 

24 'tPeri', 'epoch', 'H', 'g', 'sed_filename'] 

25 self.dataCols['KEP'] = ['objId', 'a', 'e', 'inc', 'Omega', 'argPeri', 

26 'meanAnomaly', 'epoch', 'H', 'g', 'sed_filename'] 

27 self.dataCols['CAR'] = ['objId', 'x', 'y', 'z', 'xdot', 'ydot', 'zdot', 

28 'epoch', 'H', 'g', 'sed_filename'] 

29 

30 def __len__(self): 

31 return len(self.orbits) 

32 

33 def __getitem__(self, i): 

34 orb = Orbits() 

35 orb.setOrbits(self.orbits.iloc[i]) 

36 return orb 

37 

38 def __iter__(self): 

39 for i, orbit in self.orbits.iterrows(): 

40 orb = Orbits() 

41 orb.setOrbits(orbit) 

42 yield orb 

43 

44 def __eq__(self, otherOrbits): 

45 if isinstance(otherOrbits, Orbits): 

46 if self.orb_format != otherOrbits.orb_format: 

47 return False 

48 for col in self.dataCols[self.orb_format]: 

49 if not np.all(self.orbits[col].values == otherOrbits.orbits[col].values): 

50 return False 

51 else: 

52 return True 

53 else: 

54 return False 

55 

56 def __neq__(self, otherOrbits): 

57 if self == otherOrbits: 

58 return False 

59 else: 

60 return True 

61 

62 def setOrbits(self, orbits): 

63 """Set and validate orbital parameters contain all required values. 

64 

65 Sets self.orbits and self.orb_format. 

66 If objid is not present in orbits, a sequential series of integers will be used. 

67 If H is not present in orbits, a default value of 20 will be used. 

68 If g is not present in orbits, a default value of 0.15 will be used. 

69 If sed_filename is not present in orbits, either C or S type will be assigned, 

70 according to the semi-major axis value. 

71 

72 Parameters 

73 ---------- 

74 orbits : pandas.DataFrame, pandas.Series or numpy.ndarray 

75 Array-like object containing orbital parameter information. 

76 """ 

77 # Do we have a single item or multiples? 

78 if isinstance(orbits, pd.Series): 

79 # Passed a single SSO in Series, convert to a DataFrame. 

80 orbits = pd.DataFrame([orbits]) 

81 elif isinstance(orbits, np.ndarray): 

82 # Passed a numpy array, convert to DataFrame. 

83 orbits = pd.DataFrame.from_records(orbits) 

84 elif isinstance(orbits, np.record): 

85 # This was a single object in a numpy array and we should be a bit fancy. 

86 orbits = pd.DataFrame.from_records([orbits], columns=orbits.dtype.names) 

87 elif isinstance(orbits, pd.DataFrame): 

88 # This was a pandas dataframe .. but we probably want to drop the index and recount. 

89 orbits.reset_index(drop=True, inplace=True) 

90 

91 if 'index' in orbits: 

92 del orbits['index'] 

93 

94 nSso = len(orbits) 

95 

96 # Error if orbits is empty (this avoids hard-to-interpret error messages from pyoorb). 

97 if nSso == 0: 

98 raise ValueError('Length of the orbits dataframe was 0.') 

99 

100 # Discover which type of orbital parameters we have on disk. 

101 self.orb_format = None 

102 if 'FORMAT' in orbits: 

103 if ~(orbits['FORMAT'] == orbits['FORMAT'].iloc[0]).all(): 

104 raise ValueError('All orbital elements in the set should have the same FORMAT.') 

105 self.orb_format = orbits['FORMAT'].iloc[0] 

106 # Backwards compatibility .. a bit. CART is deprecated, so swap it to CAR. 

107 if self.orb_format == 'CART': 

108 self.orb_format = 'CAR' 

109 del orbits['FORMAT'] 

110 # Check that the orbit format is approximately right. 

111 if self.orb_format == 'COM': 

112 if 'q' not in orbits: 

113 raise ValueError('The stated format was COM, but "q" not present in orbital elements?') 

114 if self.orb_format == 'KEP': 

115 if 'a' not in orbits: 

116 raise ValueError('The stated format was KEP, but "a" not present in orbital elements?') 

117 if self.orb_format == 'CAR': 

118 if 'x' not in orbits: 

119 raise ValueError('The stated format was CAR but "x" not present in orbital elements?') 

120 if self.orb_format is None: 

121 # Try to figure out the format, if it wasn't provided. 

122 if 'q' in orbits: 

123 self.orb_format = 'COM' 

124 elif 'a' in orbits: 

125 self.orb_format = 'KEP' 

126 elif 'x' in orbits: 

127 self.orb_format = 'CAR' 

128 else: 

129 raise ValueError("Can't determine orbital type, as neither q, a or x in input orbital elements.\n" 

130 "Was attempting to base orbital element quantities on header row, " 

131 "with columns: \n%s" % orbits.columns) 

132 

133 # Check that the orbit epoch is within a 'reasonable' range, to detect possible column mismatches. 

134 general_epoch = orbits['epoch'].head(1).values[0] 

135 # Look for epochs between 1800 and 2200 - this is primarily to check if people used MJD (and not JD). 

136 expect_min_epoch = -21503. 

137 expect_max_epoch = 124594. 

138 if general_epoch < expect_min_epoch or general_epoch > expect_max_epoch: 

139 raise ValueError("The epoch detected for this orbit is odd - %f. " 

140 "Expecting a value between %.1f and %.1f (MJD!)" % (general_epoch, 

141 expect_min_epoch, 

142 expect_max_epoch)) 

143 

144 # If these columns are not available in the input data, auto-generate them. 

145 if 'objId' not in orbits: 

146 objId = np.arange(0, nSso, 1) 

147 orbits = orbits.assign(objId = objId) 

148 if 'H' not in orbits: 

149 orbits = orbits.assign(H = 20.0) 

150 if 'g' not in orbits: 

151 orbits = orbits.assign(g = 0.15) 

152 if 'sed_filename' not in orbits: 

153 orbits = orbits.assign(sed_filename = self.assignSed(orbits)) 

154 

155 # Make sure we gave all the columns we need. 

156 for col in self.dataCols[self.orb_format]: 

157 if col not in orbits: 

158 raise ValueError('Missing required orbital element %s for orbital format type %s' 

159 % (col, self.orb_format)) 

160 

161 # Check to see if we have duplicates. 

162 if len(orbits['objId'].unique()) != nSso: 

163 warnings.warn('There are duplicates in the orbit objId values' + 

164 ' - was this intended? (continuing).') 

165 # All is good. 

166 self.orbits = orbits 

167 

168 def assignSed(self, orbits, randomSeed=None): 

169 """Assign either a C or S type SED, depending on the semi-major axis of the object. 

170 P(C type) = 0 (a<2); 0.5*a - 1 (2<a<4); 1 (a > 4), 

171 based on figure 23 from Ivezic et al 2001 (AJ, 122, 2749). 

172 

173 Parameters 

174 ---------- 

175 orbits : pandas.DataFrame, pandas.Series or numpy.ndarray 

176 Array-like object containing orbital parameter information. 

177 

178 Returns 

179 ------- 

180 numpy.ndarray 

181 Array containing the SED type for each object in 'orbits'. 

182 """ 

183 # using fig. 23 from Ivezic et al. 2001 (AJ, 122, 2749), 

184 # we can approximate the sed types with a simple linear form: 

185 # p(C) = 0 for a<2 

186 # p(C) = 0.5*a-1 for 2<a<4 

187 # p(C) = 1 for a>4 

188 # where a is semi-major axis, and p(C) is the probability that 

189 # an asteroid is C type, with p(S)=1-p(C) for S types. 

190 if 'a' in orbits: 

191 a = orbits['a'] 

192 elif 'q' in orbits: 

193 a = orbits['q'] / (1 - orbits['e']) 

194 elif 'x' in orbits: 

195 # This definitely isn't right, but it's a placeholder to make it work for now. 

196 a = np.sqrt(orbits['x']**2 + orbits['y']**2 + orbits['z']**2) 

197 else: 

198 raise ValueError('Need either a or q (plus e) in orbit data frame.') 

199 

200 if not hasattr(self, "_rng"): 

201 if randomSeed is not None: 

202 self._rng = np.random.RandomState(randomSeed) 

203 else: 

204 self._rng = np.random.RandomState(42) 

205 

206 chance = self._rng.random_sample(len(orbits)) 

207 prob_c = 0.5 * a - 1.0 

208 # if chance <= prob_c: 

209 sedvals = np.where(chance <= prob_c, 'C.dat', 'S.dat') 

210 return sedvals 

211 

212 def readOrbits(self, orbitfile, delim=None, skiprows=None): 

213 """Read orbits from a file, generating a pandas dataframe containing columns matching dataCols, 

214 for the appropriate orbital parameter format (currently accepts COM, KEP or CAR formats). 

215 

216 After reading and standardizing the column names, calls self.setOrbits to validate the 

217 orbital parameters. Expects angles in orbital element formats to be in degrees. 

218 

219 Note that readOrbits uses pandas.read_csv to read the data file with the orbital parameters. 

220 Thus, it should have column headers specifying the column names .. 

221 unless skiprows = -1 or there is just no header line at all. 

222 in which case it is assumed to be a standard DES format file, with no header line. 

223 

224 Parameters 

225 ---------- 

226 orbitfile : str 

227 The name of the input file containing orbital parameter information. 

228 delim : str, optional 

229 The delimiter for the input orbit file. Default is None, will use delim_whitespace=True. 

230 skiprows : int, optional 

231 The number of rows to skip before reading the header information for pandas. 

232 Default is None, which will trigger a check of the file to look for the header columns. 

233 """ 

234 names = None 

235 

236 # If skiprows is set, then we will assume the user has handled this so that the 

237 # first line read has the header information. 

238 # But, if skiprows is not set, then we have to do some checking to see if there is 

239 # header information and which row it might start in. 

240 if skiprows is None: 

241 skiprows = -1 

242 # Figure out whether the header is in the first line, or if there are rows to skip. 

243 # We need to do a bit of juggling to do this before pandas reads the whole orbit file though. 

244 with open(orbitfile, 'r') as fp: 

245 headervalues = None 

246 for line in fp: 

247 values = line.split() 

248 try: 

249 # If it is a valid orbit line, we expect column 3 to be a number. 

250 float(values[3]) 

251 # And if it worked, we're done here (it's an orbit) - go on to parsing header values. 

252 break 

253 except (ValueError, IndexError): 

254 # This wasn't a valid number or there wasn't anything in the third value. 

255 # So this is either the header line or it's a comment line before the header columns. 

256 skiprows += 1 

257 headervalues = values 

258 

259 

260 if headervalues is not None: # (and skiprows > -1) 

261 # There is a header, but we also need to check if there is a comment key at the start 

262 # of the proper header line. 

263 # ... Because this varies as well, and is sometimes separated from header columns. 

264 linestart = headervalues[0] 

265 if linestart == '#' or linestart == '!!' or linestart == '##': 

266 names = headervalues[1:] 

267 else: 

268 names = headervalues 

269 # Add 1 to skiprows, so that we skip the header column line. 

270 skiprows += 1 

271 

272 # So now skiprows is a value. If it is -1, then there is no header information. 

273 if skiprows == -1: 

274 # No header; assume it's a typical DES file - 

275 # we'll assign the column names based on the FORMAT. 

276 names_COM = ('objId', 'FORMAT', 'q', 'e', 'i', 'node', 'argperi', 't_p', 

277 'H', 'epoch', 'INDEX', 'N_PAR', 'MOID', 'COMPCODE') 

278 names_KEP = ('objId', 'FORMAT', 'a', 'e', 'i', 'node', 'argperi', 'meanAnomaly', 

279 'H', 'epoch', 'INDEX', 'N_PAR', 'MOID', 'COMPCODE') 

280 names_CAR = ('objId', 'FORMAT', 'x', 'y', 'z', 'xdot', 'ydot', 'zdot', 

281 'H', 'epoch', 'INDEX', 'N_PAR', 'MOID', 'COMPCODE') 

282 # First use names_COM, and then change if required. 

283 orbits = pd.read_csv(orbitfile, delim_whitespace=True, header=None, names=names_COM) 

284 

285 if orbits['FORMAT'][0] == 'KEP': 

286 orbits.columns = names_KEP 

287 elif orbits['FORMAT'][0] == 'CAR': 

288 orbits.columns = names_CAR 

289 

290 else: 

291 if delim is None: 

292 orbits = pd.read_csv(orbitfile, delim_whitespace=True, skiprows=skiprows, 

293 names=names) 

294 else: 

295 orbits = pd.read_csv(orbitfile, sep=delim, skiprows=skiprows, names=names) 

296 

297 # Drop some columns that are typically present in DES files but that we don't need. 

298 if 'INDEX' in orbits: 

299 del orbits['INDEX'] 

300 if 'N_PAR' in orbits: 

301 del orbits['N_PAR'] 

302 if 'MOID' in orbits: 

303 del orbits['MOID'] 

304 if 'COMPCODE' in orbits: 

305 del orbits['COMPCODE'] 

306 if 'tmp' in orbits: 

307 del orbits['tmp'] 

308 

309 # Normalize the column names to standard values and identify the orbital element types. 

310 ssoCols = orbits.columns.values.tolist() 

311 

312 # These are the alternative possibilities for various column headers 

313 # (depending on file version, origin, etc.) 

314 # that might need remapping from the on-file values to our standardized values. 

315 altNames = {} 

316 altNames['objId'] = ['objId', 'objid', '!!ObjID', '!!OID', '!!S3MID', 'OID', 'S3MID' 

317 'objid(int)', 'full_name', '#name'] 

318 altNames['q'] = ['q'] 

319 altNames['a'] = ['a'] 

320 altNames['e'] = ['e', 'ecc'] 

321 altNames['inc'] = ['inc', 'i', 'i(deg)', 'incl'] 

322 altNames['Omega'] = ['Omega', 'omega', 'node', 'om', 'node(deg)', 

323 'BigOmega', 'Omega/node', 'longNode'] 

324 altNames['argPeri'] = ['argPeri', 'argperi', 'omega/argperi', 'w', 'argperi(deg)', 'peri'] 

325 altNames['tPeri'] = ['tPeri', 't_p', 'timeperi', 't_peri', 'T_peri'] 

326 altNames['epoch'] = ['epoch', 't_0', 'Epoch', 'epoch_mjd'] 

327 altNames['H'] = ['H', 'magH', 'magHv', 'Hv', 'H_v'] 

328 altNames['g'] = ['g', 'phaseV', 'phase', 'gV', 'phase_g', 'G'] 

329 altNames['meanAnomaly'] = ['meanAnomaly', 'meanAnom', 'M', 'ma'] 

330 altNames['sed_filename'] = ['sed_filename', 'sed'] 

331 altNames['xdot'] = ['xdot', 'xDot'] 

332 altNames['ydot'] = ['ydot', 'yDot'] 

333 altNames['zdot'] = ['zdot', 'zDot'] 

334 

335 # Update column names that match any of the alternatives above. 

336 for name, alternatives in altNames.items(): 

337 intersection = list(set(alternatives) & set(ssoCols)) 

338 if len(intersection) > 1: 

339 raise ValueError('Received too many possible matches to %s in orbit file %s' 

340 % (name, orbitfile)) 

341 if len(intersection) == 1: 

342 idx = ssoCols.index(intersection[0]) 

343 ssoCols[idx] = name 

344 # Assign the new column names back to the orbits dataframe. 

345 orbits.columns = ssoCols 

346 # Validate and assign orbits to self. 

347 self.setOrbits(orbits) 

348 

349 def updateOrbits(self, neworb): 

350 """Update existing orbits with new values, leaving OrbitIds, H, g, and sed_filenames in place. 

351 

352 Example use: transform orbital parameters (using PyOrbEphemerides) and then replace original values. 

353 Example use: propagate orbital parameters (using PyOrbEphemerides) and then replace original values. 

354 

355 Parameters 

356 ---------- 

357 neworb: pandas.DataFrame 

358 """ 

359 col_orig = ['objId', 'sed_filename'] 

360 new_order = ['objId'] + [n for n in neworb.columns] + ['sed_filename'] 

361 updated_orbits = neworb.join(self.orbits[col_orig])[new_order] 

362 self.setOrbits(updated_orbits) 

363