use kiddo::SquaredEuclidean;
use crate::spherical_trig_funcs::{chord_distance, convert_equitorial_to_cartesian};
use crate::tree::build_kd_tree;
#[derive(Clone, Debug)]
pub struct PositionCatalog {
pub ra_deg: Vec<f64>,
pub dec_deg: Vec<f64>
}
pub fn count_galaxies_around_point(
catalog: PositionCatalog,
ra_point_deg: f64,
dec_point_deg: f64,
angular_difference_deg: f64,
) -> usize {
let chord_length = chord_distance(angular_difference_deg);
let tree = build_kd_tree(catalog.ra_deg, catalog.dec_deg);
let point_cartesian = convert_equitorial_to_cartesian(&ra_point_deg, &dec_point_deg);
let within = tree.within::<SquaredEuclidean>(&point_cartesian, chord_length.powi(2));
within.len()
}
pub fn count_galaxies_around_points(
catalog: PositionCatalog,
evaluate_catalog: PositionCatalog,
angular_differences_deg: Vec<f64>,
) -> Vec<usize> {
let chord_lengths = angular_differences_deg
.iter()
.map(|&ang| chord_distance(ang));
let tree = build_kd_tree(catalog.ra_deg, catalog.dec_deg);
let points_cartesian: Vec<[f64; 3]> = evaluate_catalog.ra_deg
.iter()
.zip(evaluate_catalog.dec_deg)
.map(|(ra, dec)| convert_equitorial_to_cartesian(ra, &dec))
.collect();
points_cartesian
.iter()
.zip(chord_lengths)
.map(|(point, chord)| tree.within::<SquaredEuclidean>(point, chord.powi(2)).len())
.collect()
}
pub fn calculate_completeness(observed_catalog: PositionCatalog, target_catalog: PositionCatalog, angular_radius_deg: Vec<f64>) -> Vec<f64> {
let number_observed = count_galaxies_around_points(observed_catalog.clone(), observed_catalog.clone(), angular_radius_deg.clone());
let number_target = count_galaxies_around_points(target_catalog, observed_catalog.clone(), angular_radius_deg);
number_observed.iter()
.zip(number_target.iter())
.map(|(obs, tar)| (*obs as f64)/(*tar as f64))
.collect::<Vec<f64>>()
}
#[cfg(test)]
mod tests {
use std::iter::zip;
use super::*;
#[test]
fn test_counting_around_a_single_point() {
let ras = vec![120., 120.9999, 120., 180.];
let decs = vec![0., 0., 0., -45.];
let cat = PositionCatalog {ra_deg: ras, dec_deg: decs};
let result_3 = count_galaxies_around_point(cat.clone(), 120., 0., 1.);
let result_2 = count_galaxies_around_point(cat.clone(), 120., 0., 0.1);
let result_1 = count_galaxies_around_point(cat.clone(), 179., -46., 5.);
assert_eq!(result_3, 3);
assert_eq!(result_2, 2);
assert_eq!(result_1, 1);
}
#[test]
fn test_counting_around_multiple_points() {
let ras = vec![120., 120.9999, 120., 180.];
let decs = vec![0., 0., 0., -45.];
let ra_points = vec![120., 120., 179.];
let dec_points = vec![0., 0., -46.];
let catalog = PositionCatalog {ra_deg: ras, dec_deg: decs};
let eval_catalog = PositionCatalog {ra_deg: ra_points, dec_deg: dec_points};
let ang_seps = vec![1., 0.1, 5.];
let results = count_galaxies_around_points(catalog, eval_catalog, ang_seps);
let answers = [3, 2, 1];
for (res, ans) in zip(results, answers) {
assert_eq!(res, ans)
}
}
#[test]
fn test_completeness() {
let target = PositionCatalog {ra_deg: vec![20.;4], dec_deg: vec![-20.;4]};
let observed = PositionCatalog {ra_deg: vec![20.;3], dec_deg: vec![-20.;3]};
let ang_dist = vec![0.000001;3];
let result = calculate_completeness(observed, target, ang_dist);
for (res, ans) in zip(result, vec![0.75; 3]) {
assert_eq!(res, ans)
}
let target = PositionCatalog {ra_deg: vec![20., 30., 40., 50.], dec_deg: vec![-20., -18., 20., 36.]};
let observed = PositionCatalog {ra_deg: vec![20., 30., 40., 50.], dec_deg: vec![-20., -18., 20., 36.]};
let bad_observed = PositionCatalog {ra_deg: vec![20., 30., 40.], dec_deg: vec![-20., -18., 20.]};
let ang_dist = vec![0.1;4];
let ang_dist_bad = vec![0.001; 3];
let result = calculate_completeness(observed, target.clone(), ang_dist);
for (res, ans) in zip(result, vec![1.0; 4]) {
assert_eq!(res, ans)
}
let bad_result = calculate_completeness(bad_observed, target, ang_dist_bad);
for (res, ans) in zip(bad_result, vec![1.0; 3]) {
assert_eq!(res, ans)
}
}
}