diff --git a/src/hr_selection_function/gaia.py b/src/hr_selection_function/gaia.py index 64348a7..9041543 100644 --- a/src/hr_selection_function/gaia.py +++ b/src/hr_selection_function/gaia.py @@ -17,14 +17,12 @@ class GaiaDensityEstimator: @requires_data def __init__(self): - self._map = pd.read_parquet( - _CONFIG["data_dir"] / "m10_hp7.parquet" - ) # Todo make this + self._map = pd.read_parquet(_CONFIG["data_dir"] / "density_hp7.parquet") def __call__( self, - l: ArrayLike, - b: ArrayLike, + l: ArrayLike, # noqa: E741 + b: ArrayLike, # noqa: E741 pmra: ArrayLike, pmdec: ArrayLike, parallax: ArrayLike, @@ -61,26 +59,41 @@ class GaiaDensityEstimator: When parallax values are outside of allowed range, or when shapes of arguments mismatch. """ - # Todo: check shapes too - l, b, pmra, pmdec, parallax = ( + l, b, pmra, pmdec, parallax = ( # noqa: E741 np.atleast_1d(l), np.atleast_1d(b), np.atleast_1d(pmra), np.atleast_1d(pmdec), np.atleast_1d(parallax), ) + self._check_call_arguments(l, b, pmra, pmdec, parallax) + df_indices = self._generate_dataframe_indices(l, b) + map_values = self._interpolate_values(df_indices, parallax) + pdf_values = self._query_multivariate_normal(map_values, pmra, pmdec) + density_per_field = self._convert_pdf_to_field_dens(map_values, pdf_values) + return density_per_field + + def _check_call_arguments(self, l, b, pmra, pmdec, parallax): # noqa: E741 if not l.shape == b.shape == pmra.shape == pmdec.shape == parallax.shape: raise ValueError("Shapes of arguments must match.") if np.any(parallax > 2) or np.any(parallax < 0): raise ValueError( - "Parallax must be positive and less than 2 mas (i.e. distance > 500 pc.)" + "Parallax must be positive and less than 2 mas (i.e. distance > 500 pc)" ) - # Get indices of everything to fetch + def _generate_dataframe_indices(self, l, b): # noqa: E741 + """The dataframe the map is stored in is 2D, but should really be 3D - this + quick indexing function grabs all required info at a given pixel. + """ id_pix = hp.ang2pix(2**7, l, b, nest=True, lonlat=True) - df_indices = self._generate_dataframe_indices(id_pix) + id_bin = self._N_BINS * id_pix + id_all_bins = (id_bin.reshape(-1, 1) + np.arange(self._N_BINS)).flatten() + return id_all_bins - # Perform interpolations + def _interpolate_values(self, df_indices, parallax): + """Interpolates all relevant values on the map between neighboring parallax + bins of our query point at a given parallax. + """ parallax_map = self._map.loc[df_indices, "parallax"].to_numpy().reshape(-1, 21) values = dict() for col in ("mean_0", "mean_1", "cov_0", "cov_1", "cov_3", "n_stars_per_mas"): @@ -92,31 +105,24 @@ class GaiaDensityEstimator: # Since the matrix is symmetric, we skip this one values["cov_2"] = values["cov_1"] + return values - # Query the multivar normal + def _query_multivariate_normal(self, map_values, pmra, pmdec): proper_motions = np.asarray([pmra, pmdec]).T - means = np.asarray([values["mean_0"], values["mean_1"]]).T - covariances = np.asarray([values[f"cov_{i}"] for i in range(4)]).T.reshape( - -1, 2, 2 - ) + means = np.asarray([map_values["mean_0"], map_values["mean_1"]]).T + covariances = np.asarray([map_values[f"cov_{i}"] for i in range(4)]).T - normal_values = vectorized_multivariate_normal( - proper_motions, means, covariances + return vectorized_multivariate_normal( + proper_motions, means, covariances.reshape(-1, 2, 2) ) + def _convert_pdf_to_field_dens(self, map_values, pdf_values): density_per_degree = ( - values["n_stars_per_mas"] * normal_values / self._HP7_PIXEL_AREA + map_values["n_stars_per_mas"] * pdf_values / self._HP7_PIXEL_AREA ) - return density_per_degree * self._HP5_FIELD_AREA - - def _generate_dataframe_indices(self, id_pix): - """The dataframe the map is stored in is 2D, but should really be 3D - this - quick indexing function grabs all required info at a given pixel. - """ - return ( - self._N_BINS * id_pix.reshape(-1, 1) + np.arange(self._N_BINS) - ).flatten() + density_per_field = density_per_degree * self._HP5_FIELD_AREA + return density_per_field class M10Estimator: @@ -148,7 +154,7 @@ class MSubsampleEstimator: def __init__(self): self._map = pd.read_parquet(_CONFIG["data_dir"] / "subsample_cuts_hp7.parquet") - def __call__(self, l: ArrayLike, b: ArrayLike) -> ArrayLike: + def __call__(self, l: ArrayLike, b: ArrayLike) -> ArrayLike: # noqa: E741 """Estimates the median magnitude of stars removed from Gaia data by HR23's subsample cuts at the given l and b values. diff --git a/tests/test_gaia.py b/tests/test_gaia.py new file mode 100644 index 0000000..e00ecba --- /dev/null +++ b/tests/test_gaia.py @@ -0,0 +1,144 @@ +from hr_selection_function.gaia import ( + M10Estimator, + MSubsampleEstimator, + GaiaDensityEstimator, +) +import numpy as np +import healpy as hp + + +pixels_for_testing = [0, 57, 5893, 58932, 3, 382, 382, 159532, 9694, 196607] + + +def test_M10Estimator(): + model = M10Estimator() + + # Some values looked up manually from the GaiaUnlimited package + test_values = [ + 21.0278, + 20.9648, + 21.0118, + 20.9611, + 20.9353, + 20.9351, + 20.9351, + 21.0907, + 20.9455, + 21.0904, + ] + + # Convert into ra/dec + ra, dec = hp.pix2ang(2**7, pixels_for_testing, nest=True, lonlat=True) + + result = model(ra, dec) + np.testing.assert_allclose(result, test_values) + + +def test_MSubsampleEstimator(): + model = MSubsampleEstimator() + + # Some values looked up manually from the dataframe + test_values = [ + 20.473399, + 20.460272, + 20.783859, + 20.615404, + 20.396479, + 20.648832, + 20.648832, + 20.727514, + 20.797851, + 20.266360, + ] + + # Convert into ra/dec + l, b = hp.pix2ang(2**7, pixels_for_testing, nest=True, lonlat=True) # noqa: E741 + + result = model(l, b) + np.testing.assert_allclose(result, test_values) + + +def test_GaiaDensityEstimator(): + l, b, pmra, pmdec, parallax, result_expected = _get_test_density_data() # noqa: E741 + + model = GaiaDensityEstimator() + result_actual = model(l, b, pmra, pmdec, parallax) + + np.testing.assert_allclose(result_actual, result_expected) + + +def _get_test_density_data(): + return [ + [ + 195.44967469843573, + 172.28930309346387, + 104.75874782403532, + 173.70832020346367, + 79.20595430423523, + 153.9832103408225, + 76.36865570717138, + 75.91799679118454, + 269.8536005362042, + 191.86293977733843, + ], + [ + -0.3109413445191218, + -0.09090826566516344, + -1.970488037640793, + -2.534745624741465, + 1.975546326136809, + -9.477611780892138, + -0.36600912016968506, + 2.33643211648706, + -2.839079251269452, + -8.251482667198461, + ], + [ + 0.41893829079702133, + 0.45112952571743464, + -0.8864160229048276, + -0.00944812985700761, + -1.945501466102864, + -0.2501085125841758, + -1.1535825632274828, + -1.9011820856534454, + -5.555678886414437, + 1.0248139396694183, + ], + [ + 0.5954508838720914, + -0.5945587927019814, + -2.2983623730571368, + -1.2255169778814685, + -3.1292798157719655, + -0.4026029905131187, + -2.047954786200685, + -3.2077251711446126, + 5.1387683882097965, + 0.23107733843403835, + ], + [ + 0.11226267600669465, + 0.09336113589903527, + 0.1087746540178971, + 0.11248267295854439, + 0.07306373461755672, + 0.14866912811593022, + 0.09964444948414934, + 0.09891577538495654, + 0.16250218203179784, + 0.343948431873036, + ], + [ + 26670.688142002928, + 32035.043148849825, + 77662.2672329491, + 32962.30258016591, + 5788.030329350739, + 22917.85569373121, + 3136.150415190452, + 127494.21366708462, + 82293.64350935556, + 23629.392112275982, + ], + ]