diff --git a/benches/expensive_benches.rs b/benches/expensive_benches.rs index ffc8f63..62e6081 100644 --- a/benches/expensive_benches.rs +++ b/benches/expensive_benches.rs @@ -175,9 +175,7 @@ fn bench_correct_geometry_mwax_half_1247842824(crt: &mut Criterion) { correct_geometry( black_box(&corr_ctx), black_box(jones_array.view_mut()), - black_box(&vis_sel.timestep_range), - black_box(&vis_sel.coarse_chan_range), - black_box(&vis_sel.baseline_idxs), + black_box(&vis_sel), None, None, false, @@ -206,9 +204,7 @@ fn bench_correct_geometry_ord_half_1196175296(crt: &mut Criterion) { correct_geometry( black_box(&corr_ctx), black_box(jones_array.view_mut()), - black_box(&vis_sel.timestep_range), - black_box(&vis_sel.coarse_chan_range), - black_box(&vis_sel.baseline_idxs), + black_box(&vis_sel), None, None, false, diff --git a/src/corrections.rs b/src/corrections.rs index 7cd9f72..9770dbb 100644 --- a/src/corrections.rs +++ b/src/corrections.rs @@ -12,7 +12,7 @@ use marlu::{ io::error::BadArrayShape, mwalib::{CorrelatorContext, MWAVersion}, precession::precess_time, - Complex, LatLngHeight, RADec, XyzGeodetic, UVW, + Complex, LatLngHeight, RADec, VisSelection, XyzGeodetic, UVW, }; use std::{f64::consts::TAU, ops::Range}; use thiserror::Error; @@ -198,23 +198,16 @@ pub fn correct_cable_lengths( /// correct_geometry( /// &corr_ctx, /// jones_array.view_mut(), -/// &vis_sel.timestep_range, -/// &vis_sel.coarse_chan_range, -/// &vis_sel.baseline_idxs, +/// &vis_sel, /// None, /// None, /// false, /// ); /// ``` -#[allow(clippy::too_many_arguments)] pub fn correct_geometry( corr_ctx: &CorrelatorContext, mut jones_array: ArrayViewMut3>, - // TODO: next 3 args could just be - // vis_sel: &VisSelection, - timestep_range: &Range, - coarse_chan_range: &Range, - baseline_idxs: &[usize], + vis_sel: &VisSelection, array_pos: Option, phase_centre: Option, draw_progress: bool, @@ -229,10 +222,10 @@ pub fn correct_geometry( LatLngHeight::mwa() }); - let timesteps = &corr_ctx.timesteps[timestep_range.clone()]; + let timesteps = &corr_ctx.timesteps[vis_sel.timestep_range.clone()]; - let all_freqs_hz = - corr_ctx.get_fine_chan_freqs_hz_array(&coarse_chan_range.clone().collect::>()); + let all_freqs_hz = corr_ctx + .get_fine_chan_freqs_hz_array(&vis_sel.coarse_chan_range.clone().collect::>()); let jones_dims = jones_array.dim(); let integration_time_s = corr_ctx.metafits_context.corr_int_time_ms as f64 / 1000.0; @@ -242,15 +235,7 @@ pub fn correct_geometry( let tiles_xyz_geod = XyzGeodetic::get_tiles(&corr_ctx.metafits_context, array_pos.latitude_rad); // use baseline_idxs to select antpairs out of corr_ctx.metafits_context.baselines - let ant_pairs = baseline_idxs - .iter() - .map(|b| { - ( - corr_ctx.metafits_context.baselines[*b].ant1_index, - corr_ctx.metafits_context.baselines[*b].ant2_index, - ) - }) - .collect::>(); + let ant_pairs = vis_sel.get_ant_pairs(&corr_ctx.metafits_context); let centroid_timestamps = timesteps .iter() .map(|t| Epoch::from_gpst_seconds(t.gps_time_ms as f64 / 1000.0 + integration_time_s / 2.0)) @@ -1122,9 +1107,7 @@ mod tests { correct_geometry( &corr_ctx, jones_array.view_mut(), - &vis_sel.timestep_range, - &vis_sel.coarse_chan_range, - &vis_sel.baseline_idxs, + &vis_sel, None, None, false, @@ -1273,9 +1256,7 @@ mod tests { correct_geometry( &corr_ctx, jones_array.view_mut(), - &vis_sel.timestep_range, - &vis_sel.coarse_chan_range, - &vis_sel.baseline_idxs, + &vis_sel, None, None, false, diff --git a/src/preprocessing.rs b/src/preprocessing.rs index d3a6d25..9d2d988 100644 --- a/src/preprocessing.rs +++ b/src/preprocessing.rs @@ -235,9 +235,7 @@ impl<'a> PreprocessContext<'a> { correct_geometry( corr_ctx, jones_array.view_mut(), - &vis_sel.timestep_range, - &vis_sel.coarse_chan_range, - &vis_sel.baseline_idxs, + vis_sel, Some(self.array_pos), Some(self.phase_centre), self.draw_progress,