WIP. Peeling!

WIP: Write thorough tests that intended visibilities make their way to the peel
functions. Currently, averaging to 1.28 MHz isn't happening
WIP: Add UVW cutoffs
WIP: Fix error handling in src/cli/error.rs

This is a squashed commit of months of work by CHJ and Dev. Big props to Dev for
their help.

.github/workflows/bench.yml

@@ -0,0 +1,44 @@
+---
+name: Benchmarks
+
+on:
+  push:
+    pull_request:
+    branches:
+      - "**"
+
+env:
+  CARGO_TERM_COLOR: always
+  CARGO_INCREMENTAL: 0
+  MWA_BEAM_FILE: /usr/local/mwa_full_embedded_element_pattern.h5
+
+jobs:
+  test:
+    name: Benchmarks
+    runs-on: self-hosted
+    steps:
+      - name: Checkout sources
+        uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+
+      - name: Install stable toolchain
+        uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: stable
+          override: true
+
+      - name: Cargo Bench
+        run: |
+          cargo bench
+        env:
+          HYPERDRIVE_TEST_DIR: /home/runner/data
+
+      - name: Zip benchmark results
+        run: zip -r criterion.zip target/criterion/*
+
+      - uses: actions/upload-artifact@v2
+        with:
+          name: criterion.zip
+          path: criterion.zip

src/averaging/mod.rs

@@ -478,7 +478,8 @@ pub(super) fn parse_freq_average_factor(
             // Scale the quantity by the unit, if required.
             let quantity = match time_format {
                 FreqFormat::Hz => quantity,
-                FreqFormat::kHz => 1000.0 * quantity,
+                FreqFormat::kHz => 1e3 * quantity,
+                FreqFormat::MHz => 1e6 * quantity,
             };
             let factor = quantity / freq_res;
             // Reject non-integer floats.
@@ -615,3 +616,95 @@ fn vis_average_weights_non_zero(
     }
     *weight_to = weight_sum as f32;
 }
+
+/// This function is the same as `vis_average`, except it assumes that any
+/// flagged visibilities are indicated by weights equal to 0.0 (or -0.0). This
+/// allows the code to be more efficient.
+pub(crate) fn vis_average_no_negative_weights(
+    jones_from_tfb: ArrayView3<Jones<f32>>,
+    mut jones_to_fb: ArrayViewMut2<Jones<f32>>,
+    weight_from_tfb: ArrayView3<f32>,
+    mut weight_to_fb: ArrayViewMut2<f32>,
+    flagged_chan_indices: &HashSet<u16>,
+) {
+    let avg_time = jones_from_tfb.len_of(Axis(0));
+    let avg_freq = (jones_from_tfb.len_of(Axis(1)) as f64 / jones_to_fb.len_of(Axis(0)) as f64)
+        .ceil() as usize;
+
+    // iterate along time axis in chunks of avg_time
+    jones_from_tfb
+        .axis_chunks_iter(Axis(0), avg_time)
+        .zip(weight_from_tfb.axis_chunks_iter(Axis(0), avg_time))
+        .for_each(|(jones_chunk_tfb, weight_chunk_tfb)| {
+            jones_chunk_tfb
+                .axis_iter(Axis(2))
+                .zip(weight_chunk_tfb.axis_iter(Axis(2)))
+                .zip(jones_to_fb.axis_iter_mut(Axis(1)))
+                .zip(weight_to_fb.axis_iter_mut(Axis(1)))
+                .for_each(
+                    |(((jones_chunk_tf, weight_chunk_tf), mut jones_to_f), mut weight_to_f)| {
+                        jones_chunk_tf
+                            .axis_chunks_iter(Axis(1), avg_freq)
+                            .zip(weight_chunk_tf.axis_chunks_iter(Axis(1), avg_freq))
+                            .enumerate()
+                            .filter(|(i, _)| !flagged_chan_indices.contains(&(*i as u16)))
+                            .map(|(_, d)| d)
+                            .zip(jones_to_f.iter_mut())
+                            .zip(weight_to_f.iter_mut())
+                            .for_each(
+                                |(((jones_chunk_tf, weight_chunk_tf), jones_to), weight_to)| {
+                                    vis_average_weights_are_zero(
+                                        jones_chunk_tf,
+                                        weight_chunk_tf,
+                                        jones_to,
+                                        weight_to,
+                                    );
+                                },
+                            );
+                    },
+                );
+        });
+}
+
+/// Average a chunk of visibilities and weights (both must have the same
+/// dimensions) into an output vis and weight. This function needs the input
+/// weights to be 0 or greater; this allows the averaging algorithm to be
+/// simpler (as well as hopefully being faster), while also allowing further
+/// averaging of the output visibilities without complicated logic.
+#[inline]
+fn vis_average_weights_are_zero(
+    jones_chunk_tf: ArrayView2<Jones<f32>>,
+    weight_chunk_tf: ArrayView2<f32>,
+    jones_to: &mut Jones<f32>,
+    weight_to: &mut f32,
+) {
+    let mut jones_weighted_sum = Jones::default();
+    let mut weight_sum = 0.0;
+
+    // iterate through time chunks
+    jones_chunk_tf
+        .outer_iter()
+        .zip_eq(weight_chunk_tf.outer_iter())
+        .for_each(|(jones_chunk_f, weights_chunk_f)| {
+            jones_chunk_f
+                .iter()
+                .zip_eq(weights_chunk_f.iter())
+                .for_each(|(jones, weight)| {
+                    // Any flagged visibilities would have a weight <= 0, but
+                    // we've already capped them to 0. This means we don't need
+                    // to check the value of the weight when accumulating
+                    // unflagged visibilities; the flagged ones contribute
+                    // nothing.
+
+                    let jones = Jones::<f64>::from(*jones);
+                    let weight = *weight as f64;
+                    jones_weighted_sum += jones * weight;
+                    weight_sum += weight;
+                });
+        });
+
+    if weight_sum > 0.0 {
+        *jones_to = Jones::from(jones_weighted_sum / weight_sum);
+        *weight_to = weight_sum as f32;
+    }
+}

src/cli/di_calibrate/tests.rs

@@ -27,7 +27,6 @@ use crate::{
         MsReader, VisRead,
     },
     math::TileBaselineFlags,
-    params::CalVis,
     tests::{
         get_reduced_1090008640_ms, get_reduced_1090008640_raw, get_reduced_1090008640_uvfits,
         DataAsStrings,
@@ -918,14 +917,11 @@ fn test_1090008640_calibration_quality_raw() {
     };
 
     let params = args.parse().unwrap();
-    let CalVis {
-        vis_data,
-        vis_model,
-        ..
-    } = params
+    let mut cal_vis = params
         .get_cal_vis()
         .expect("Couldn't read data and generate a model");
-    test_1090008640_quality(params, vis_data.view(), vis_model.view());
+    cal_vis.scale_by_weights(Some(&params.baseline_weights));
+    test_1090008640_quality(params, cal_vis.vis_data.view(), cal_vis.vis_model.view());
 }
 
 #[test]
@@ -962,14 +958,11 @@ fn test_1090008640_calibration_quality_ms() {
     };
 
     let params = args.parse().unwrap();
-    let CalVis {
-        vis_data,
-        vis_model,
-        ..
-    } = params
+    let mut cal_vis = params
         .get_cal_vis()
         .expect("Couldn't read data and generate a model");
-    test_1090008640_quality(params, vis_data.view(), vis_model.view());
+    cal_vis.scale_by_weights(Some(&params.baseline_weights));
+    test_1090008640_quality(params, cal_vis.vis_data.view(), cal_vis.vis_model.view());
 }
 
 #[test]
@@ -1005,12 +998,9 @@ fn test_1090008640_calibration_quality_uvfits() {
     };
 
     let params = args.parse().unwrap();
-    let CalVis {
-        vis_data,
-        vis_model,
-        ..
-    } = params
+    let mut cal_vis = params
         .get_cal_vis()
         .expect("Couldn't read data and generate a model");
-    test_1090008640_quality(params, vis_data.view(), vis_model.view());
+    cal_vis.scale_by_weights(Some(&params.baseline_weights));
+    test_1090008640_quality(params, cal_vis.vis_data.view(), cal_vis.vis_model.view());
 }

src/cli/error.rs

@@ -10,6 +10,7 @@ use thiserror::Error;
 use super::{
     common::InputVisArgsError,
     di_calibrate::DiCalArgsError,
+    peel::PeelArgsError,
     solutions::{SolutionsApplyArgsError, SolutionsPlotError},
     srclist::SrclistByBeamError,
     vis_convert::VisConvertArgsError,
@@ -24,7 +25,7 @@ use crate::{
         GlobError,
     },
     model::ModelError,
-    params::{DiCalibrateError, VisConvertError, VisSimulateError, VisSubtractError},
+    params::{DiCalibrateError, PeelError, VisConvertError, VisSimulateError, VisSubtractError},
     solutions::{SolutionsReadError, SolutionsWriteError},
     srclist::{ReadSourceListError, SrclistError, WriteSourceListError},
 };
@@ -39,6 +40,10 @@ pub enum HyperdriveError {
     #[error("{0}\n\nSee for more info: {URL}/user/di_cal/intro.html")]
     DiCalibrate(String),
 
+    /// An error related to peeling.
+    #[error("{0}\n\nSee for more info: {URL}/*****.html")]
+    Peel(String),
+
     /// An error related to solutions-apply.
     #[error("{0}\n\nSee for more info: {URL}/user/solutions_apply/intro.html")]
     SolutionsApply(String),
@@ -164,6 +169,49 @@ impl From<DiCalibrateError> for HyperdriveError {
     }
 }
 
+impl From<PeelArgsError> for HyperdriveError {
+    fn from(e: PeelArgsError) -> Self {
+        match e {
+            PeelArgsError::NoOutput
+            | PeelArgsError::NoChannels
+            | PeelArgsError::ZeroPasses
+            | PeelArgsError::ParseIonoTimeAverageFactor(_)
+            | PeelArgsError::ParseIonoFreqAverageFactor(_)
+            | PeelArgsError::IonoTimeFactorNotInteger
+            | PeelArgsError::IonoFreqFactorNotInteger
+            | PeelArgsError::IonoTimeResNotMultiple { .. }
+            | PeelArgsError::IonoFreqResNotMultiple { .. }
+            | PeelArgsError::IonoTimeFactorZero
+            | PeelArgsError::IonoFreqFactorZero
+            | PeelArgsError::ParseUvwMin(_)
+            | PeelArgsError::ParseUvwMax(_) => Self::Generic(e.to_string()),
+            PeelArgsError::Glob(e) => Self::from(e),
+            PeelArgsError::VisRead(e) => Self::from(e),
+            PeelArgsError::FileWrite(e) => Self::from(e),
+            PeelArgsError::SourceList(e) => Self::from(e),
+            PeelArgsError::Beam(e) => Self::from(e),
+            PeelArgsError::Model(e) => Self::from(e),
+            PeelArgsError::IO(e) => Self::from(e),
+            #[cfg(any(feature = "cuda", feature = "hip"))]
+            PeelArgsError::Gpu(e) => Self::from(e),
+        }
+    }
+}
+
+impl From<PeelError> for HyperdriveError {
+    fn from(e: PeelError) -> Self {
+        match e {
+            PeelError::VisRead(e) => Self::from(e),
+            PeelError::FileWrite(e) => Self::from(e),
+            PeelError::Beam(e) => Self::from(e),
+            PeelError::Model(e) => Self::from(e),
+            PeelError::IO(e) => Self::from(e),
+            #[cfg(any(feature = "cuda", feature = "hip"))]
+            PeelError::Gpu(e) => Self::from(e),
+        }
+    }
+}
+
 impl From<SolutionsApplyArgsError> for HyperdriveError {
     fn from(e: SolutionsApplyArgsError) -> Self {
         let s = e.to_string();

src/cli/mod.rs

@@ -18,6 +18,7 @@ mod beam;
 mod di_calibrate;
 mod dipole_gains;
 mod error;
+mod peel;
 mod solutions;
 mod srclist;
 mod vis_convert;
@@ -94,6 +95,9 @@ https://mwatelescope.github.io/mwa_hyperdrive/user/di_cal/intro.html"#
     )]
     DiCalibrate(di_calibrate::DiCalArgs),
 
+    #[clap(about = r#"Peeling!"#)]
+    Peel(peel::PeelArgs),
+
     #[clap(alias = "convert-vis")]
     #[clap(about = r#"Convert visibilities from one type to another.
     https://mwatelescope.github.io/mwa_hyperdrive/user/vis_convert/intro.html"#)]
@@ -171,6 +175,7 @@ impl Hyperdrive {
         // Print the version of hyperdrive and its build-time information.
         let sub_command = match &self.command {
             Command::DiCalibrate(_) => "di-calibrate",
+            Command::Peel(_) => "peel",
             Command::VisConvert(_) => "vis-convert",
             Command::VisSimulate(_) => "vis-simulate",
             Command::VisSubtract(_) => "vis-subtract",
@@ -209,6 +214,10 @@ impl Hyperdrive {
                 merge_save_run!(args)
             }
 
+            Command::Peel(args) => {
+                merge_save_run!(args)
+            }
+
             Command::VisConvert(args) => {
                 merge_save_run!(args)
             }

src/cli/peel/error.rs

@@ -0,0 +1,72 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#[derive(thiserror::Error, Debug)]
+pub(crate) enum PeelArgsError {
+    #[error("No calibration output was specified. There must be at least one calibration solution file.")]
+    NoOutput,
+
+    #[error("The data either contains no frequency channels or all channels are flagged")]
+    NoChannels,
+
+    #[error("The number of iono sub passes cannot be 0")]
+    ZeroPasses,
+
+    #[error("Error when parsing iono time average factor: {0}")]
+    ParseIonoTimeAverageFactor(crate::unit_parsing::UnitParseError),
+
+    #[error("Error when parsing iono freq. average factor: {0}")]
+    ParseIonoFreqAverageFactor(crate::unit_parsing::UnitParseError),
+
+    #[error("Iono time average factor isn't an integer")]
+    IonoTimeFactorNotInteger,
+
+    #[error("Iono freq. average factor isn't an integer")]
+    IonoFreqFactorNotInteger,
+
+    #[error(
+        "Iono time resolution isn't a multiple of input data's: {out} seconds vs {inp} seconds"
+    )]
+    IonoTimeResNotMultiple { out: f64, inp: f64 },
+
+    #[error("Iono freq. resolution isn't a multiple of input data's: {out} Hz vs {inp} Hz")]
+    IonoFreqResNotMultiple { out: f64, inp: f64 },
+
+    #[error("Iono time average factor cannot be 0")]
+    IonoTimeFactorZero,
+
+    #[error("Iono freq. average factor cannot be 0")]
+    IonoFreqFactorZero,
+
+    #[error("Error when parsing minimum UVW cutoff: {0}")]
+    ParseUvwMin(crate::unit_parsing::UnitParseError),
+
+    #[error("Error when parsing maximum UVW cutoff: {0}")]
+    ParseUvwMax(crate::unit_parsing::UnitParseError),
+
+    #[error(transparent)]
+    Glob(#[from] crate::io::GlobError),
+
+    #[error(transparent)]
+    VisRead(#[from] crate::io::read::VisReadError),
+
+    #[error(transparent)]
+    FileWrite(#[from] crate::io::write::FileWriteError),
+
+    #[error("Error when trying to read source list: {0}")]
+    SourceList(#[from] crate::srclist::ReadSourceListError),
+
+    #[error(transparent)]
+    Beam(#[from] crate::beam::BeamError),
+
+    #[error(transparent)]
+    Model(#[from] crate::model::ModelError),
+
+    #[error(transparent)]
+    IO(#[from] std::io::Error),
+
+    #[cfg(any(feature = "cuda", feature = "hip"))]
+    #[error(transparent)]
+    Gpu(#[from] crate::gpu::GpuError),
+}