diff --git a/src/too/database.py b/src/too/database.py index 1479f1d..d2d46be 100644 --- a/src/too/database.py +++ b/src/too/database.py @@ -135,6 +135,19 @@ def validate_too_targets(targets: polars.DataFrame): "At least one magnitude value is required." ) + # Check that if one of the Sloan magnitudes is set, all are set. + sloan_mags = targets.select(polars.col("^[ugriz]_mag$")) + # Rows where at least one value is not null. + sloan_mags_not_null = sloan_mags.filter( + polars.fold(False, lambda a, b: a | b.is_not_null(), polars.all()) + ) + # Rows where at least one value is null. + sloan_mags_null = sloan_mags_not_null.filter( + polars.fold(False, lambda a, b: a | b.is_null(), polars.all()) + ) + if sloan_mags_null.height > 0: + raise ValidationError("Found rows with incomplete Sloan magnitudes.") + if set(targets["fiber_type"].unique()) != set(["APOGEE", "BOSS"]): raise ValidationError( "Invalid fiber_type values. Valid values are 'APOGEE' and 'BOSS'." diff --git a/src/too/mock.py b/src/too/mock.py index 0026762..738ff11 100644 --- a/src/too/mock.py +++ b/src/too/mock.py @@ -217,7 +217,6 @@ def create_mock_too_catalogue( int(n_unknown * fraction_unknown_sdss), ) sdss_unknown_targets = sdss_unknown_targets.with_columns( - polars.col(r"^[ugiz]_mag$").replace({}, default=None), catalogid=None, sdss_id=None, ) @@ -248,7 +247,11 @@ def create_mock_too_catalogue( "ra": numpy.random.uniform(0, 360, n_random), "dec": numpy.random.uniform(-90, 90, n_random), "gaia_g_mag": polars.Series(random_gaia, dtype=polars.Float32), + "u_mag": polars.Series(random_sdss, dtype=polars.Float32), + "g_mag": polars.Series(random_sdss, dtype=polars.Float32), "r_mag": polars.Series(random_sdss, dtype=polars.Float32), + "i_mag": polars.Series(random_sdss, dtype=polars.Float32), + "z_mag": polars.Series(random_sdss, dtype=polars.Float32), }, ) diff --git a/tests/test_database.py b/tests/test_database.py index 7f6dd2b..1ab4e47 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -105,6 +105,7 @@ def test_validate_too_target_passes(too_mock: polars.DataFrame): ("active", "Null 'active' column values found"), ("mag_columns", "ToOs found with missing magnitudes"), ("fiber_type", "Invalid fiber_type values."), + ("sloan_mags", "Found rows with incomplete Sloan magnitudes."), ], ) def test_validate_too_target_fails( @@ -133,6 +134,10 @@ def test_validate_too_target_fails( ) elif test_mode == "fiber_type": too_mock_test[0, "fiber_type"] = "INVALID" + elif test_mode == "sloan_mags": + too_mock_test = too_mock_test.with_columns( + u_mag=polars.lit(None, dtype=polars.Float32) + ) with pytest.raises(ValidationError, match=error_message): validate_too_targets(too_mock_test)