use np testing to assert equal in gaia tools tests

henrysky · Jul 20, 2024 · f5cb1ab · f5cb1ab
1 parent 445ac40
commit f5cb1ab
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 39 deletions.
diff --git a/src/astroNN/data/__init__.py b/src/astroNN/data/__init__.py
@@ -28,13 +28,11 @@ def data_description():
     | 2019-July-02 - Updated - Henry Leung (University of Toronto)
     """
     items = {
-        "anderson_2017_dr14_parallax.npz": "Anderson 2017 improved Gaia TGAS parallax from Data-Driven Stellar Model",
+        "aspcap_l31c_masks.npy": "ASPCAP l31c (DR14) elements windows mask represented by bits on the 7514px spectrum",
         "dr13_contmask.npz": "APOGEE DR13 continuum mask from Bovy's APOGEE tools",
         "dr14_contmask.npz": "APOGEE DR14 continuum mask from Bovy's APOGEE tools",
         "dr16_contmask.npz": "APOGEE DR16 continuum mask",
-        "gaiadr2_apogeedr14_parallax.npz": "Gaia DR2 - APOGEE DR14 matches, indices corresponds "
-        "to APOGEE allstar DR14 file",
-        "aspcap_l31c_masks.npy": "ASPCAP l31c (DR14) elements windows mask represented by bits on the 7514px spectrum",
+        "dr17_contmask.npz": "APOGEE DR17 continuum mask",
     }
 
     for item in items:

diff --git a/tests/test_gaia_tools.py b/tests/test_gaia_tools.py
@@ -87,7 +87,7 @@ def test_logsol(self):
         npt.assert_array_equal(logsol_to_fakemag(fakemag_to_logsol(np.array([100, 100, 100]))), [100., 100., 100.])
         npt.assert_equal(fakemag_to_logsol(MAGIC_NUMBER), MAGIC_NUMBER)
         npt.assert_equal(logsol_to_fakemag(fakemag_to_logsol(MAGIC_NUMBER)), MAGIC_NUMBER)
-        npt.assert_equal(np.any(fakemag_to_logsol([MAGIC_NUMBER, 1000]) == MAGIC_NUMBER), True)
+        npt.assert_equal(fakemag_to_logsol([MAGIC_NUMBER, 1000])[1], MAGIC_NUMBER)
 
         npt.assert_equal(logsol_to_absmag(absmag_to_logsol(99.)), 99.)
         self.assertAlmostEqual(logsol_to_absmag(absmag_to_logsol(-99.)), -99.)
@@ -96,23 +96,23 @@ def test_logsol(self):
         npt.assert_array_almost_equal(logsol_to_absmag(absmag_to_logsol(np.array([99., 99., 99.]))), [99., 99., 99.])
         npt.assert_equal(absmag_to_logsol(MAGIC_NUMBER), MAGIC_NUMBER)
         npt.assert_equal(logsol_to_absmag(absmag_to_logsol(MAGIC_NUMBER)), MAGIC_NUMBER)
-        npt.assert_equal(np.any(absmag_to_logsol([MAGIC_NUMBER, 1000]) == MAGIC_NUMBER), True)
+        npt.assert_equal(absmag_to_logsol([MAGIC_NUMBER, 1000])[0], MAGIC_NUMBER)
 
     def test_extinction(self):
         from astroNN.gaia import extinction_correction
 
-        npt.assert_equal(np.any([extinction_correction(10., -90.) == -9999.]), False)
+        npt.assert_raises(AssertionError, npt.assert_array_equal, extinction_correction(10., -90.)[1], -9999)
         npt.assert_equal(extinction_correction(10., -90.), 10.)
-        npt.assert_equal(np.any([extinction_correction(-99.99, -90.) == -9999.]), True)
+        npt.assert_equal(np.any([extinction_correction(-99.99, -90.) == -9999.]))
 
     def test_known_regression(self):
         # prevent regression of known bug
         npt.assert_equal(mag_to_absmag(1., MAGIC_NUMBER), MAGIC_NUMBER)
         npt.assert_equal(mag_to_absmag(MAGIC_NUMBER, MAGIC_NUMBER), MAGIC_NUMBER)
         npt.assert_equal(
-            np.all(mag_to_absmag(MAGIC_NUMBER, MAGIC_NUMBER, 1.) == (MAGIC_NUMBER, MAGIC_NUMBER)), True)
+            np.all(mag_to_absmag(MAGIC_NUMBER, MAGIC_NUMBER, 1.), (MAGIC_NUMBER, MAGIC_NUMBER)))
         npt.assert_equal(
-            np.all(mag_to_fakemag(MAGIC_NUMBER, MAGIC_NUMBER, 1.) == (MAGIC_NUMBER, MAGIC_NUMBER)), True)
+            np.all(mag_to_fakemag(MAGIC_NUMBER, MAGIC_NUMBER, 1.), (MAGIC_NUMBER, MAGIC_NUMBER)))
 
         npt.assert_equal(mag_to_fakemag(1., MAGIC_NUMBER), MAGIC_NUMBER)
         npt.assert_equal(mag_to_fakemag(MAGIC_NUMBER, MAGIC_NUMBER), MAGIC_NUMBER)
@@ -121,35 +121,6 @@ def test_known_regression(self):
         npt.assert_equal(fakemag_to_pc(-1., 2.).value, MAGIC_NUMBER)
         npt.assert_equal(absmag_to_pc(1., MAGIC_NUMBER).value, MAGIC_NUMBER)
 
-    def test_anderson(self):
-        from astroNN.gaia import anderson_2017_parallax
-        # To load the improved parallax
-        # Both parallax and para_var is in mas
-        # cuts=True to cut bad data (negative parallax and percentage error more than 20%)
-        ra, dec, parallax, para_err = anderson_2017_parallax(cuts=True)
-        npt.assert_equal(np.any([parallax == -9999.]), False)  # assert no -9999
-
-    def test_dr2_parallax(self):
-        from astroNN.gaia import gaiadr2_parallax
-        # To load the improved parallax
-        # Both parallax and para_var is in mas
-        # cuts=True to cut bad data (negative parallax and percentage error more than 20%)
-        ra, dec, parallax, para_err = gaiadr2_parallax(cuts=True)
-        ra02, dec02, parallax02, para_err02 = gaiadr2_parallax(cuts=0.2)
-        ra01, dec01, parallax01, para_err01 = gaiadr2_parallax(cuts=0.1)
-        # assert no -9999.
-        npt.assert_equal(np.any([parallax == -9999.]), False)
-        # assert cuts = True equals 0.2
-        npt.assert_equal(np.any([ra == ra02]), True)
-        npt.assert_equal((ra01.shape[0] < ra02.shape[0]), True)
-        # assert no ridiculous parallax if do cut
-        npt.assert_equal(np.any([((para_err / parallax) > 0.2) & (parallax < 0.)]), False)
-
-        ra, dec, parallax, para_err = gaiadr2_parallax(cuts=False)
-        # assert some -9999. due to no cuts
-        npt.assert_equal(np.any([parallax == -9999.]), True)
-        ra, dec, parallax, para_err = gaiadr2_parallax(cuts=True, keepdims=True)
-
 
 if __name__ == '__main__':
     unittest.main()