Coverage for tests/test_actions.py: 98%

1# This file is part of analysis_tools. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22import unittest 

23 

24import astropy.units as u 

25import numpy as np 

26import pandas as pd 

27 

28import lsst.utils.tests 

29from lsst.analysis.tools.actions.keyedData.calcDistances import CalcRelativeDistances 

30from lsst.analysis.tools.actions.scalar import ( 

31 ApproxFloor, 

32 CountAction, 

33 MeanAction, 

34 MedianAction, 

35 SigmaMadAction, 

36 StdevAction, 

37 WeightedMeanAction, 

38) 

39from lsst.analysis.tools.actions.vector import ( 

40 CalcBinnedStatsAction, 

41 CalcMomentSize, 

42 CalcRhoStatistics, 

43 ConvertFluxToMag, 

44 DownselectVector, 

45 ExtinctionCorrectedMagDiff, 

46 LoadVector, 

47 MagDiff, 

48) 

49from lsst.analysis.tools.actions.vector.mathActions import ( 

50 AddVector, 

51 ConstantValue, 

52 DivideVector, 

53 FractionalDifference, 

54 Log10Vector, 

55 MultiplyVector, 

56 RaiseFromBaseVector, 

57 RaiseToPowerVector, 

58 SqrtVector, 

59 SquareVector, 

60 SubtractVector, 

61) 

62from lsst.analysis.tools.actions.vector.selectors import ( 

63 CoaddPlotFlagSelector, 

64 FlagSelector, 

65 GalaxySelector, 

66 RangeSelector, 

67 SetSelector, 

68 SkyObjectSelector, 

69 SnSelector, 

70 StarSelector, 

71 VectorSelector, 

72) 

73from lsst.analysis.tools.math import sqrt 

74from lsst.pex.config import FieldValidationError 

75 

76 

77class TestScalarActions(unittest.TestCase): 

78 """ "Test ScalarActions""" 

79 

80 def setUp(self): 

81 x = np.arange(100, dtype=float) 

82 x[31] = np.nan 

83 x[41] = np.nan 

84 x[59] = np.nan 

85 y = x**2 

86 self.data = { 

87 "r_y": x, 

88 "i_y": y, 

89 "z_y": y.reshape((10, 10)), 

90 } 

91 mask = np.zeros(100, dtype=bool) 

92 mask[1:50] = 1 

93 self.mask = { 

94 "r": mask, 

95 "i": mask, 

96 "z": mask.reshape((10, 10)), 

97 } 

98 

99 def _testScalarActionAlmostEqual(self, cls, truth, maskedTruth, **kwargs): 

100 action = cls(vectorKey="{band}_y", **kwargs) 

101 schema = [col for col, colType in action.getInputSchema()] 

102 if isinstance(action, WeightedMeanAction): 

103 self.assertEqual(schema, ["{band}_y", kwargs["weightsKey"]]) 

104 else: 

105 self.assertEqual(schema, ["{band}_y"]) 

106 

107 result = action(self.data, band="i") 

108 self.assertAlmostEqual(result, truth) 

109 

110 result = action(self.data, band="z") 

111 self.assertAlmostEqual(result, truth) 

112 

113 result = action(self.data, mask=self.mask["i"], band="i") 

114 self.assertAlmostEqual(result, maskedTruth) 

115 

116 result = action(self.data, mask=self.mask["z"], band="z") 

117 self.assertAlmostEqual(result, maskedTruth) 

118 

119 def testMedianAction(self): 

120 self._testScalarActionAlmostEqual(MedianAction, 2500, 576) 

121 

122 def testMeanAction(self): 

123 self._testScalarActionAlmostEqual(MeanAction, 3321.9278350515465, 803.8936170212766) 

124 

125 def testStdevAction(self): 

126 self._testScalarActionAlmostEqual(StdevAction, 2984.5855649976297, 733.5989754407842) 

127 

128 def testSigmaMadAction(self): 

129 self._testScalarActionAlmostEqual(SigmaMadAction, 3278.033505115886, 759.0923358748682) 

130 

131 def testCountAction(self): 

132 self._testScalarActionAlmostEqual(CountAction, 97, 47) 

133 

134 def testApproxFloorAction(self): 

135 self._testScalarActionAlmostEqual(ApproxFloor, 9216.0, 2352.5) 

136 

137 def testWeightedMeanAction(self): 

138 self._testScalarActionAlmostEqual( 

139 WeightedMeanAction, 6003.42828813228, 1473.3447052907393, weightsKey="{band}_y" 

140 ) 

141 

142 

143class TestVectorActions(unittest.TestCase): 

144 """Test VectorActions""" 

145 

146 def setUp(self): 

147 self.size = 5 

148 r = np.arange(float(self.size)) + 1 

149 i = r**2 

150 rFlag = np.ones(self.size) 

151 rFlag[1] = 0 

152 iFlag = np.ones(self.size) 

153 iFlag[3] = 0 

154 data = { 

155 "r_vector": r, 

156 "i_vector": i, 

157 "r_flag": rFlag, 

158 "i_flag": iFlag, 

159 "g_vector": 3 * r, 

160 "E(B-V)": r[::-1], 

161 "r_ixx": r**2, 

162 "r_iyy": r[::-1] ** 2, 

163 "r_ixy": r * r[::-1], 

164 "g_iixx": r**2, 

165 "g_iiyy": r[::-1] ** 2, 

166 "g_iixy": r * r[::-1], 

167 } 

168 

169 self.data = pd.DataFrame.from_dict(data) 

170 

171 def _checkSchema(self, action, truth): 

172 schema = sorted([col for col, colType in action.getInputSchema()]) 

173 self.assertEqual(schema, truth) 

174 

175 # VectorActions with their own files 

176 

177 def testCalcBinnedStats(self): 

178 selector = RangeSelector(vectorKey="r_vector", minimum=0, maximum=self.size) 

179 prefix = "a_" 

180 stats = CalcBinnedStatsAction(name_prefix=prefix, selector_range=selector, key_vector="r_vector") 

181 result = stats(self.data) 

182 mask = selector(self.data) 

183 values = self.data["r_vector"][mask] 

184 median = np.median(values) 

185 truth = { 

186 stats.name_mask: mask, 

187 stats.name_median: median, 

188 stats.name_sigmaMad: 1.482602218505602 * np.median(np.abs(values - median)), 

189 stats.name_count: np.sum(mask), 

190 stats.name_select_maximum: np.max(values), 

191 stats.name_select_median: median, 

192 stats.name_select_minimum: np.min(values), 

193 "range_maximum": self.size, 

194 "range_minimum": 0, 

195 } 

196 self.assertEqual(list(result.keys()), list(truth.keys())) 

197 

198 np.testing.assert_array_almost_equal(result[stats.name_sigmaMad], truth[stats.name_sigmaMad]) 

199 del truth[stats.name_sigmaMad] 

200 

201 np.testing.assert_array_equal(result[stats.name_mask], truth[stats.name_mask]) 

202 del truth[stats.name_mask] 

203 

204 for key, value in truth.items(): 

205 self.assertEqual(result[key], value, key) 

206 

207 def testCalcMomentSize(self): 

208 xx = self.data["r_ixx"] 

209 yy = self.data["r_iyy"] 

210 xy = self.data["r_ixy"] 

211 

212 # Test determinant with defaults 

213 action = CalcMomentSize() 

214 result = action(self.data, band="r") 

215 schema = [col for col, colType in action.getInputSchema()] 

216 self.assertEqual(sorted(schema), ["{band}_ixx", "{band}_ixy", "{band}_iyy"]) 

217 truth = 0.25 * (xx * yy - xy**2) 

218 np.testing.assert_array_almost_equal(result, truth) 

219 

220 # Test trace with columns specified 

221 action = CalcMomentSize( 

222 colXx="{band}_iixx", 

223 colYy="{band}_iiyy", 

224 colXy="{band}_iixy", 

225 sizeType="trace", 

226 ) 

227 result = action(self.data, band="g") 

228 schema = [col for col, colType in action.getInputSchema()] 

229 self.assertEqual(sorted(schema), ["{band}_iixx", "{band}_iiyy"]) 

230 truth = sqrt(0.5 * (xx + yy)) 

231 np.testing.assert_array_almost_equal(result, truth) 

232 

233 CalcMomentSize(sizeType="trace", colXy=None).validate() 

234 with self.assertRaises(FieldValidationError): 

235 CalcMomentSize(sizeType="determinant", colXy=None).validate() 

236 

237 # def testCalcE(self): TODO: implement 

238 

239 # def testCalcEDiff(self): TODO: implement 

240 

241 # def testCalcE1(self): TODO: implement 

242 

243 # def testCalcE2(self): TODO: implement 

244 

245 # MathActions 

246 

247 def _testMath(self, ActionType, truth, num_vectors: int = 2, compare_exact: bool = False, **kwargs): 

248 letters = ("A", "B") 

249 bands = ("r", "i") 

250 actions = { 

251 f"action{letters[num]}": LoadVector(vectorKey=f"{{band{num + 1}}}_vector") 

252 for num in range(num_vectors) 

253 } 

254 action = ActionType(**actions, **kwargs) 

255 kwargs_bands = {f"band{num + 1}": bands[num] for num in range(num_vectors)} 

256 result = action(self.data, **kwargs_bands) 

257 self._checkSchema(action, [action.vectorKey for action in actions.values()]) 

258 if compare_exact: 

259 np.testing.assert_array_equal(result, truth) 

260 else: 

261 np.testing.assert_array_almost_equal(result, truth) 

262 

263 def testConstant(self): 

264 truth = [42.0] 

265 action = ConstantValue(value=truth[0]) 

266 self._checkSchema(action, []) 

267 result = action({}) 

268 np.testing.assert_array_equal(result, truth) 

269 

270 def testAdd(self): 

271 truth = [2.0, 6.0, 12.0, 20.0, 30.0] 

272 self._testMath(AddVector, truth, compare_exact=True) 

273 

274 def testSubtract(self): 

275 truth = [0.0, -2.0, -6.0, -12.0, -20.0] 

276 self._testMath(SubtractVector, truth, compare_exact=True) 

277 

278 def testMultiply(self): 

279 truth = [1.0, 8.0, 27.0, 64.0, 125.0] 

280 self._testMath(MultiplyVector, truth, compare_exact=False) 

281 

282 def testDivide(self): 

283 truth = 1 / np.arange(1, 6) 

284 self._testMath(DivideVector, truth, compare_exact=False) 

285 

286 def testSqrt(self): 

287 truth = sqrt(np.arange(1, 6)) 

288 self._testMath(SqrtVector, truth, compare_exact=False, num_vectors=1) 

289 

290 def testSquare(self): 

291 truth = np.arange(1, 6) ** 2 

292 self._testMath(SquareVector, truth, compare_exact=True, num_vectors=1) 

293 

294 def testRaiseFromBase(self): 

295 power = np.arange(1, 6) 

296 for base in (-2.3, 0.6): 

297 truth = base**power 

298 self._testMath(RaiseFromBaseVector, truth, compare_exact=False, base=base, num_vectors=1) 

299 

300 def testRaiseToPower(self): 

301 base = np.arange(1, 6) 

302 for power in (-2.3, 0.6): 

303 truth = base**power 

304 self._testMath(RaiseToPowerVector, truth, compare_exact=False, power=power, num_vectors=1) 

305 

306 def testLog10(self): 

307 truth = np.log10(np.arange(1, 6)) 

308 self._testMath(Log10Vector, truth, compare_exact=False, num_vectors=1) 

309 

310 def testFractionalDifference(self): 

311 truth = [0.0, -0.5, -0.6666666666666666, -0.75, -0.8] 

312 self._testMath(FractionalDifference, truth, compare_exact=False) 

313 

314 # Basic vectorActions 

315 

316 # def testLoadVector(self): TODO: implement 

317 

318 def testDownselectVector(self): 

319 selector = FlagSelector(selectWhenTrue=["{band}_flag"]) 

320 action = DownselectVector(vectorKey="{band}_vector", selector=selector) 

321 result = action(self.data, band="r") 

322 self._checkSchema(action, ["{band}_flag", "{band}_vector"]) 

323 np.testing.assert_array_equal(result, np.array([1, 3, 4, 5])) 

324 

325 # def testMultiCriteriaDownselectVector(self): TODO: implement 

326 

327 # Astronomical vectorActions 

328 

329 # def testCalcSn(self): TODO: implement 

330 

331 def testConvertFluxToMag(self): 

332 truth = [ 

333 31.4, 

334 29.89485002168, 

335 29.0143937264, 

336 28.38970004336, 

337 27.90514997832, 

338 ] 

339 action = ConvertFluxToMag(vectorKey="{band}_vector") 

340 result = action(self.data, band="i") 

341 self._checkSchema(action, ["{band}_vector"]) 

342 np.testing.assert_array_almost_equal(result, truth) 

343 

344 # def testConvertUnits(self): TODO: implement 

345 

346 def testMagDiff(self): 

347 # Use the same units as the data so that we know the difference exactly 

348 # without conversions. 

349 # testExtinctionCorrectedMagDiff will test the conversions 

350 truth = np.array(2 * self.data["r_vector"]) * u.ABmag 

351 action = MagDiff( 

352 col1="{band1}_vector", 

353 col2="{band2}_vector", 

354 fluxUnits1="mag(AB)", 

355 fluxUnits2="mag(AB)", 

356 returnMillimags=False, 

357 ) 

358 result = action(self.data, band1="g", band2="r") 

359 self._checkSchema(action, ["{band1}_vector", "{band2}_vector"]) 

360 np.testing.assert_array_almost_equal(result, truth.value) 

361 

362 def testExtinctionCorrectedMagDiff(self): 

363 for returnMillimags in (True, False): 

364 # Check that conversions are working properly 

365 magDiff = MagDiff( 

366 col1="{band1}_vector", 

367 col2="{band2}_vector", 

368 fluxUnits1="jansky", 

369 fluxUnits2="jansky", 

370 returnMillimags=returnMillimags, 

371 ) 

372 action = ExtinctionCorrectedMagDiff( 

373 magDiff=magDiff, 

374 band1="g", 

375 band2="r", 

376 ebvCol="E(B-V)", 

377 extinctionCoeffs={"g": 0.2, "r": 1.5}, 

378 ) 

379 

380 result = action(self.data, band1="g", band2="r") 

381 lhs = (np.array(self.data["g_vector"]) * u.jansky).to(u.ABmag) 

382 rhs = (np.array(self.data["r_vector"]) * u.jansky).to(u.ABmag) 

383 diff = lhs - rhs 

384 correction = np.array((0.2 - 1.5) * self.data["E(B-V)"]) * u.mag 

385 if returnMillimags: 

386 diff = diff.to(u.mmag) 

387 correction = correction.to(u.mmag) 

388 truth = diff - correction 

389 self._checkSchema(action, ["E(B-V)", "{band1}_vector", "{band2}_vector"]) 

390 np.testing.assert_array_almost_equal(result, truth.value) 

391 

392 # Test with hard coded bands 

393 magDiff = MagDiff( 

394 col1="g_vector", 

395 col2="r_vector", 

396 fluxUnits1="jansky", 

397 fluxUnits2="jansky", 

398 returnMillimags=False, 

399 ) 

400 action = ExtinctionCorrectedMagDiff( 

401 magDiff=magDiff, 

402 ebvCol="E(B-V)", 

403 extinctionCoeffs={"g": 0.2, "r": 1.5}, 

404 ) 

405 result = action(self.data) 

406 lhs = (np.array(self.data["g_vector"]) * u.jansky).to(u.ABmag) 

407 rhs = (np.array(self.data["r_vector"]) * u.jansky).to(u.ABmag) 

408 diff = lhs - rhs 

409 correction = np.array((0.2 - 1.5) * self.data["E(B-V)"]) * u.mag 

410 truth = diff - correction 

411 self._checkSchema(action, ["E(B-V)", "g_vector", "r_vector"]) 

412 np.testing.assert_array_almost_equal(result, truth.value) 

413 

414 # def testRAcosDec(self): TODO: implement 

415 

416 # Statistical vectorActions 

417 

418 # def testPerGroupStatistic(self): TODO: implement 

419 

420 # def testResidualWithPerGroupStatistic(self): TODO: implement 

421 

422 

423class TestVectorRhoStats(unittest.TestCase): 

424 """Test Rho stats""" 

425 

426 def setUp(self): 

427 # generate data just for testCalcRhoStatistics. 

428 np.random.seed(42) 

429 sizeRho = 1000 

430 size_src = np.random.normal(scale=1e-3, size=sizeRho) 

431 e1_src = np.random.normal(scale=1e-3, size=sizeRho) 

432 e2_src = np.random.normal(scale=1e-3, size=sizeRho) 

433 

434 size_psf = np.random.normal(scale=1e-3, size=sizeRho) 

435 e1_psf = np.random.normal(scale=1e-3, size=sizeRho) 

436 e2_psf = np.random.normal(scale=1e-3, size=sizeRho) 

437 

438 src_data = np.array( 

439 [self.getMatrixElements(size, e1, e2) for size, e1, e2 in zip(size_src, e1_src, e2_src)] 

440 ) 

441 psf_data = np.array( 

442 [self.getMatrixElements(size, e1, e2) for size, e1, e2 in zip(size_psf, e1_psf, e2_psf)] 

443 ) 

444 

445 dataRhoStats = { 

446 "coord_ra": np.random.uniform(-120, 120, sizeRho), 

447 "coord_dec": np.random.uniform(-120, 120, sizeRho), 

448 "r_ixx": src_data[:, 0], 

449 "r_iyy": src_data[:, 1], 

450 "r_ixy": src_data[:, 2], 

451 "r_ixxPSF": psf_data[:, 0], 

452 "r_iyyPSF": psf_data[:, 1], 

453 "r_ixyPSF": psf_data[:, 2], 

454 } 

455 

456 self.dataRhoStats = pd.DataFrame.from_dict(dataRhoStats) 

457 

458 # Needed for testCalcRhoStatistics. 

459 @staticmethod 

460 def getMatrixElements(size, e1, e2): 

461 # putting guards just in case e1 or e2 are superior to 1, but unlikely. 

462 if abs(e1) >= 1: 462 ↛ 463line 462 didn't jump to line 463 because the condition on line 462 was never true

463 e1 = 0 

464 if abs(e2) >= 1: 464 ↛ 465line 464 didn't jump to line 465 because the condition on line 464 was never true

465 e2 = 0 

466 e = sqrt(e1**2 + e2**2) 

467 q = (1 - e) / (1 + e) 

468 phi = 0.5 * np.arctan2(e2, e1) 

469 rot = np.array([[np.cos(phi), np.sin(phi)], [-np.sin(phi), np.cos(phi)]]) 

470 ell = np.array([[size**2, 0], [0, (size * q) ** 2]]) 

471 L = np.dot(rot.T, ell.dot(rot)) 

472 return [L[0, 0], L[1, 1], L[0, 1]] 

473 

474 def testCalcRhoStatistics(self): 

475 # just check if runs 

476 rho = CalcRhoStatistics() 

477 rho.treecorr.nbins = 21 

478 rho.treecorr.min_sep = 0.01 

479 rho.treecorr.max_sep = 100.0 

480 rho.treecorr.sep_units = "arcmin" 

481 result = rho(self.dataRhoStats, band="r") 

482 for rho in result: 

483 if rho != "rho3alt": 

484 self.assertEqual(np.sum(np.isfinite(result[rho].xip)), len(result[rho].xip)) 

485 self.assertEqual(np.sum(np.isfinite(result[rho].xim)), len(result[rho].xim)) 

486 else: 

487 self.assertEqual(np.sum(np.isfinite(result[rho].xi)), len(result[rho].xi)) 

488 

489 

490class TestVectorSelectors(unittest.TestCase): 

491 def setUp(self): 

492 self.size = 20 

493 falseFlags = { 

494 "{band}_psfFlux_flag": [1], 

495 "{band}_pixelFlags_saturatedCenter": [3], 

496 "{band}_extendedness_flag": [5], 

497 "coord_flag": [7], 

498 "i_pixelFlags_edge": [13], 

499 "r_pixelFlags_edge": [15], 

500 "i_pixelFlags_nodata": [14], 

501 "r_pixelFlags_nodata": [16], 

502 "sky_object": [13, 15, 17], 

503 } 

504 

505 trueFlags = { 

506 "detect_isPatchInner": [9], 

507 "detect_isDeblendedSource": [11], 

508 } 

509 

510 flux = np.arange(self.size) * 10 

511 fluxErr = np.ones(self.size) * 0.1 

512 extendedness = np.arange(20) / 20 - 0.1 

513 

514 self.data = { 

515 "r_psfFlux": flux, 

516 "r_psfFluxErr": fluxErr, 

517 "i_cmodelFlux": flux[::-1], 

518 "i_cmodelFluxError": fluxErr[::-1], 

519 "r_cmodelFlux": flux, 

520 "r_cmodelFluxError": fluxErr, 

521 "i_extendedness": extendedness, 

522 "i_extended": extendedness, 

523 } 

524 bands = ("r", "i") 

525 for band in bands: 

526 for flag, bits in falseFlags.items(): 

527 vector = np.zeros(self.size, dtype=bool) 

528 for bit in bits: 

529 vector[bit] = 1 

530 self.data[flag.format(band=band)] = vector 

531 

532 for flag, bits in trueFlags.items(): 

533 vector = np.ones(self.size, dtype=bool) 

534 for bit in bits: 

535 vector[bit] = 0 

536 self.data[flag.format(band=band)] = vector 

537 

538 def _checkSchema(self, action, truth): 

539 schema = [col for col, colType in action.getInputSchema()] 

540 self.assertEqual(sorted(schema), sorted(truth)) 

541 

542 def testFlagSelector(self): 

543 selector = FlagSelector( 

544 selectWhenFalse=["{band}_psfFlux_flag"], selectWhenTrue=["detect_isPatchInner"] 

545 ) 

546 self._checkSchema(selector, ["detect_isPatchInner", "{band}_psfFlux_flag"]) 

547 result = selector(self.data, band="r") 

548 truth = np.ones(self.size, dtype=bool) 

549 truth[1] = False 

550 truth[9] = False 

551 np.testing.assert_array_equal(result, truth) 

552 

553 def testCoaddPlotFlagSelector(self): 

554 # Test defaults 

555 selector = CoaddPlotFlagSelector() 

556 keys = [ 

557 "{band}_psfFlux_flag", 

558 "{band}_pixelFlags_saturatedCenter", 

559 "{band}_extendedness_flag", 

560 "sky_object", 

561 "coord_flag", 

562 "detect_isPatchInner", 

563 "detect_isDeblendedSource", 

564 ] 

565 self._checkSchema(selector, keys) 

566 # Specifying a band will format all keys containing band 

567 selector.bands = ["i"] 

568 self._checkSchema(selector, [key.format(band=selector.bands[0]) for key in keys]) 

569 

570 result = selector(self.data) 

571 truth = np.ones(self.size, dtype=bool) 

572 for bit in (1, 3, 5, 7, 9, 11, 13, 15, 17): 

573 truth[bit] = 0 

574 np.testing.assert_array_equal(result, truth) 

575 

576 # Test bands override 

577 selector = CoaddPlotFlagSelector( 

578 selectWhenFalse=["{band}_psfFlux_flag"], 

579 selectWhenTrue=["detect_isDeblendedSource"], 

580 ) 

581 self._checkSchema(selector, ["{band}_psfFlux_flag", "detect_isDeblendedSource"]) 

582 selector.bands = ["i", "r"] 

583 self._checkSchema(selector, ["i_psfFlux_flag", "r_psfFlux_flag", "detect_isDeblendedSource"]) 

584 result = selector(self.data) 

585 truth = np.ones(self.size, dtype=bool) 

586 for bit in (1, 11): 

587 truth[bit] = 0 

588 np.testing.assert_array_equal(result, truth) 

589 

590 def testRangeSelector(self): 

591 selector = RangeSelector(vectorKey="r_psfFlux", minimum=np.nextafter(20, 30), maximum=50) 

592 self._checkSchema(selector, ["r_psfFlux"]) 

593 result = self.data["r_psfFlux"][selector(self.data)] 

594 truth = [30, 40] 

595 np.testing.assert_array_equal(result, truth) 

596 

597 def testSetSelector(self): 

598 n_values = 3 

599 values = self.data["r_psfFlux"][:n_values] 

600 selector = SetSelector(vectorKeys=("r_psfFlux", "i_cmodelFlux"), values=values) 

601 self._checkSchema(selector, ("r_psfFlux", "i_cmodelFlux")) 

602 result = selector(self.data) 

603 truth = np.zeros_like(result) 

604 truth[:n_values] = True 

605 # i_cModelFlux is just r_psfFlux reversed 

606 truth[-n_values:] = True 

607 np.testing.assert_array_equal(result, truth) 

608 

609 def testSnSelector(self): 

610 # test defaults 

611 selector = SnSelector() 

612 keys = [ 

613 "{band}_psfFlux", 

614 "{band}_psfFluxErr", 

615 ] 

616 self._checkSchema(selector, keys) 

617 result = selector(self.data, bands=["r"]) 

618 truth = np.ones(self.size, dtype=bool) 

619 truth[:6] = 0 

620 np.testing.assert_array_equal(result, truth) 

621 

622 # test overrides 

623 selector = SnSelector( 

624 fluxType="{band}_cmodelFlux", 

625 threshold=200.0, 

626 uncertaintySuffix="Error", 

627 bands=["r", "i"], 

628 ) 

629 keys = [ 

630 "{band}_cmodelFlux", 

631 "{band}_cmodelFluxError", 

632 ] 

633 self._checkSchema(selector, keys) 

634 result = selector(self.data) 

635 truth = np.ones(self.size, dtype=bool) 

636 truth[:3] = 0 

637 truth[-3:] = 0 

638 np.testing.assert_array_equal(result, truth) 

639 

640 def testSkyObjectSelector(self): 

641 # Test with kwargs 

642 selector = SkyObjectSelector() 

643 keys = ["{band}_pixelFlags_edge", "{band}_pixelFlags_nodata", "sky_object"] 

644 self._checkSchema(selector, keys) 

645 result = selector(self.data, bands=["i"]) 

646 truth = np.zeros(self.size, dtype=bool) 

647 truth[15] = 1 

648 truth[17] = 1 

649 np.testing.assert_array_equal(result, truth) 

650 

651 # Test overrides 

652 selector = SkyObjectSelector(bands=["i", "r"]) 

653 self._checkSchema(selector, keys) 

654 result = selector(self.data) 

655 truth = np.zeros(self.size, dtype=bool) 

656 truth[17] = 1 

657 np.testing.assert_array_equal(result, truth) 

658 

659 def testStarSelector(self): 

660 # test default 

661 selector = StarSelector() 

662 self._checkSchema(selector, ["{band}_extendedness"]) 

663 result = selector(self.data, band="i") 

664 truth = (self.data["i_extendedness"] >= 0) & (self.data["i_extendedness"] < 0.5) 

665 np.testing.assert_array_almost_equal(result, truth) 

666 

667 # Test overrides 

668 selector = StarSelector(vectorKey="i_extended", extendedness_maximum=0.3) 

669 result = selector(self.data, band="i") 

670 truth = (self.data["i_extendedness"] >= 0) & (self.data["i_extendedness"] < 0.3) 

671 np.testing.assert_array_almost_equal(result, truth) 

672 

673 def testGalaxySelector(self): 

674 # test default 

675 selector = GalaxySelector() 

676 self._checkSchema(selector, ["{band}_extendedness"]) 

677 result = selector(self.data, band="i") 

678 truth = self.data["i_extendedness"] > 0.5 

679 np.testing.assert_array_almost_equal(result, truth) 

680 

681 # Test overrides 

682 selector = GalaxySelector(vectorKey="i_extended", extendedness_minimum=0.3) 

683 result = selector(self.data, band="i") 

684 truth = self.data["i_extendedness"] > 0.3 

685 np.testing.assert_array_almost_equal(result, truth) 

686 

687 def testVectorSelector(self): 

688 selector = VectorSelector(vectorKey="{band}_psfFlux_flag") 

689 self._checkSchema(selector, ["{band}_psfFlux_flag"]) 

690 result = selector(self.data, band="i")