fix test fail

prob-ml · Aug 6, 2024 · 4c28796 · 4c28796
1 parent c4b9c27
commit 4c28796
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/tests/test_catalogs.py b/tests/test_catalogs.py
@@ -97,11 +97,12 @@ def test_restrict_tile_cat_to_brightest(self, multi_source_tilecat):
         assert cat.get_brightest_sources_per_tile(band=2).max_sources == 1
 
     def test_filter_tile_cat_by_flux(self, multi_source_tilecat):
-        cat = multi_source_tilecat.filter_by_flux(300, 2000)
+        cat = multi_source_tilecat.filter_by_flux(300)
         assert cat.max_sources == 2
-        assert cat["n_sources"].sum() == 3
-        assert cat["galaxy_fluxes"].sum() == 2100.0
-        assert cat["galaxy_fluxes"].max() == 1000.0
+        assert cat["n_sources"].sum() == 4
+        r_band_flux = cat["galaxy_fluxes"][..., 2:3]
+        r_band_flux = torch.where(cat.galaxy_bools, r_band_flux, torch.inf)
+        assert r_band_flux.min().item() == 500
 
     def test_bin_full_cat_by_flux(self):
         d = {