Something went wrong. Try again.
[READ-ONLY] Mirror of https://github.com/emilyhunt/hr_selection_function. Selection function for the Hunt & Reffert 2024 open cluster catalogue.
Something went wrong. Try again.
5.7 kB · 166 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167import numpy as npfrom hr_selection_function import HR24SelectionFunction, NStarsPredictorfrom .test_gaia import _get_test_density_datafrom astropy.coordinates import SkyCoordfrom astropy import units as uimport pytestimport pandas as pd
def test_HR24SelectionFunction(): ra, dec, pmra, pmdec, parallax, densities = _get_test_density_data(l_and_b=False) result_expected = _get_expected_detectability() result_expected_mean = _get_expected_detectability("mean")
# Basic setup of params n_clusters = len(ra) coords = SkyCoord( ra=np.asarray(ra) * u.deg, dec=np.asarray(dec) * u.deg, pm_ra_cosdec=np.asarray(pmra) * u.mas / u.yr, pm_dec=np.asarray(pmdec) * u.mas / u.yr, distance=(1000 / np.asarray(parallax)) * u.pc, ) n_stars = np.linspace(0, 100, num=n_clusters) errors = np.geomspace(1, 0.01, num=n_clusters) thresholds = np.linspace(1, 10, num=n_clusters) model = HR24SelectionFunction()
# Check against previous package values when using density directly result = model(densities, n_stars, errors, thresholds) np.testing.assert_allclose(result, result_expected, rtol=1e-2, atol=1e-8)
# Check against previous package values when using coordinates result = model(coords, n_stars, errors, thresholds) result_expected = _get_expected_detectability() np.testing.assert_allclose(result, result_expected, rtol=1e-2, atol=1e-8)
# Check one-example performance result_zero = model(densities[3], n_stars[3], errors[3], thresholds[3]) assert result[3] == result_zero
# Check input error raising. bad_density = densities.copy() bad_density[3] = np.nan with pytest.raises(ValueError, match="All input values must be finite"): model(bad_density, n_stars, errors, thresholds)
bad_density[3] = -1 with pytest.raises(ValueError, match="All input values must be >= 0"): model(bad_density, n_stars, errors, thresholds)
with pytest.raises(ValueError, match="Shape mismatch"): model(densities, n_stars, errors, 3.0)
bad_threshold = thresholds.copy() bad_threshold[3] = 10.1
with pytest.warns(UserWarning, match="Our model was only fit for CST threshold"): model(densities, n_stars, errors, bad_threshold)
# Integration test of other modes result = model(densities, n_stars, errors, thresholds, mode="mean") np.testing.assert_allclose(result, result_expected_mean, rtol=3e-2, atol=1e-8)
results = [] for i in range(100): results.append( model(densities, n_stars, errors, thresholds, mode="random_sample") )
result = np.mean(results, axis=0) np.testing.assert_allclose( result, result_expected, rtol=3e-1 ) # Big error cos samples
def test_n_stars_predictor(): # Setup test data coordinates, mass, extinction, log_age, metallicity, differential_extinction = ( get_test_cluster() ) test_df = pd.DataFrame.from_dict( dict( mass=mass, extinction=extinction, log_age=log_age, metallicity=metallicity, differential_extinction=differential_extinction, ) ) n_models = 5 n_clusters = len(mass)
# Check if it works at all tbh model = NStarsPredictor(models=n_models) result = model( coordinates, mass=mass, extinction=extinction, log_age=log_age, metallicity=metallicity, differential_extinction=differential_extinction, ) assert result[0].shape == result[1].shape == (n_clusters,) assert np.all(np.isfinite(result[0])) assert np.all(np.isfinite(result[1])) assert np.all(result[0] >= 0) assert np.all(result[1] > 0)
# Check the dataframe type works & gives near-identical answer result_df = model(coordinates, test_df) assert np.allclose(result[0], result_df[0]) assert np.allclose(result[1], result_df[1])
# Check some certain values assert np.allclose(result[0].mean(), 371.06874640000007) assert np.allclose(result[1].mean(), 0.05853619699999999)
# Check we get correct type in (mode == "full") result = model(coordinates, test_df, mode="full") assert result[0].shape == result[1].shape == (n_clusters, n_models)
def get_test_cluster(): coordinates = SkyCoord( l=np.linspace(0, 360, endpoint=False) * u.deg, b=np.linspace(-80, 80) * u.deg, distance=np.linspace(500, 5000) * u.pc, frame="galactic", ) mass = np.linspace(50, 5000) extinction = np.linspace(0, 5) log_age = np.linspace(10, 6) metallicity = np.linspace(-0.5, 0.5) differential_extinction = np.linspace(0, 1) return coordinates, mass, extinction, log_age, metallicity, differential_extinction
def _get_expected_detectability(mode="median"): if mode == "median": return [ 0.0, 0.007085723421269008, 0.010443618654938702, 0.042600789320291654, 0.18525819893633644, 0.9669179306367179, 0.8782480444661485, 0.965296187018378, 0.9699061206414834, 0.9728711248322095, ] if mode == "mean": return [ 0.0, 0.005972830700591403, 0.008844614639174742, 0.036354620277421756, 0.16404563065598354, 0.9647549051366275, 0.8688448242342212, 0.9626120540069898, 0.967396788562127, 0.9702594734699409, ] raise ValueError("Specified mode not recognized.")