WIP. Don't use Marlu to read from mwalib.

WIP: Re-implement progress bars.
MWATelescope · Jul 21, 2023 · dc0ba36 · dc0ba36
1 parent 8427004
commit dc0ba36
Show file tree

Hide file tree

Showing 8 changed files with 367 additions and 146 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "birli"
 description = "A preprocessing pipeline for the Murchison Widefield Array"
-version = "0.9.2"
+version = "0.10.0"
 readme = "README.md"
 homepage = "https://github.com/MWATelescope/Birli"
 repository = "https://github.com/MWATelescope/Birli"
@@ -12,7 +12,7 @@ authors = [
     "Luke A. Williams <[email protected]>",
 ]
 edition = "2021"
-rust-version = "1.63"
+rust-version = "1.64"
 license = "MPL-2.0"
 keywords = ["radioastronomy", "mwa", "astronomy", "aoflagger", "cotter"]
 categories = ["science", "parsing"]
@@ -43,7 +43,7 @@ indicatif = { version = "0.17.0", features = ["rayon"] }
 itertools = "0.10.0"
 lazy_static = "1.4.0"
 log = "0.4.0"
-marlu = "0.9.2"
+marlu = "0.10.0"
 regex = "1.4.0"
 thiserror = "1.0.0"
 
@@ -64,7 +64,7 @@ csv = "1.1"
 float-cmp = "0.9"
 glob = "0.3"
 lexical = "6.0"
-marlu = { version = "0.9.2", features = ["approx"] }
+marlu = { version = "0.10.0", features = ["approx"] }
 ndarray = { version = "0.15.4", features = ["approx-0_5"] }
 tempfile = "3.3"
 

diff --git a/src/cli.rs b/src/cli.rs
@@ -1,13 +1,30 @@
 //! Command Line Interface helpers for Birli
 
+use std::{
+    convert::Into,
+    ffi::OsString,
+    fmt::{Debug, Display},
+    time::Duration,
+};
+
+use cfg_if::cfg_if;
+use clap::{arg, command, ErrorKind::ArgumentNotFound, PossibleValue, ValueHint::FilePath};
+use itertools::{izip, Itertools};
+use log::{debug, info, trace};
+use mwalib::{
+    built_info::PKG_VERSION as MWALIB_PKG_VERSION, fitsio_sys::CFITSIO_VERSION, CableDelaysApplied,
+    CorrelatorContext, GeometricDelaysApplied,
+};
+use prettytable::{format as prettyformat, row, table};
+
 use crate::{
     error::{
         BirliError,
         BirliError::DryRun,
         CLIError::{InvalidCommandLineArgument, InvalidRangeSpecifier},
     },
     flags::FlagContext,
-    io::{aocal::AOCalSols, IOContext},
+    io::{aocal::AOCalSols, read_mwalib, IOContext},
     marlu::{
         built_info::PKG_VERSION as MARLU_PKG_VERSION,
         constants::{
@@ -23,22 +40,6 @@ use crate::{
     passband_gains::{PFB_COTTER_2014_10KHZ, PFB_JAKE_2022_200HZ},
     with_increment_duration, Axis, Complex, FlagFileSet, PreprocessContext, VisSelection,
 };
-use cfg_if::cfg_if;
-use clap::{arg, command, ErrorKind::ArgumentNotFound, PossibleValue, ValueHint::FilePath};
-use itertools::{izip, Itertools};
-use log::{debug, info, trace};
-
-use mwalib::{
-    built_info::PKG_VERSION as MWALIB_PKG_VERSION, fitsio_sys::CFITSIO_VERSION, CableDelaysApplied,
-    CorrelatorContext, GeometricDelaysApplied,
-};
-use prettytable::{format as prettyformat, row, table};
-use std::{
-    convert::Into,
-    ffi::OsString,
-    fmt::{Debug, Display},
-    time::Duration,
-};
 
 cfg_if! {
     if #[cfg(feature = "aoflagger")] {
@@ -1542,6 +1543,7 @@ impl<'a> BirliContext<'a> {
                     obs_ctx.name.as_deref(),
                     antenna_names,
                     antenna_positions.clone(),
+                    true,
                     Some(&history),
                 )
                 .expect("unable to initialize uvfits writer")
@@ -1554,6 +1556,7 @@ impl<'a> BirliContext<'a> {
                 obs_ctx.array_pos,
                 antenna_positions,
                 dut1,
+                true,
             );
             println!(
                 "Writing to MS: {} with {} chans selected",
@@ -1673,11 +1676,12 @@ impl<'a> BirliContext<'a> {
             // populate visibilities
             with_increment_duration!(
                 "read",
-                chunk_vis_sel.read_mwalib(
+                read_mwalib(
+                    &chunk_vis_sel,
                     corr_ctx,
                     jones_array.view_mut(),
                     flag_array.view_mut(),
-                    prep_ctx.draw_progress,
+                    prep_ctx.draw_progress
                 )?
             );
 
@@ -1725,12 +1729,7 @@ impl<'a> BirliContext<'a> {
                 with_increment_duration!(
                     "write",
                     uvfits_writer
-                        .write_vis(
-                            jones_array.view(),
-                            weight_array.view(),
-                            &chunk_vis_ctx,
-                            prep_ctx.draw_progress,
-                        )
+                        .write_vis(jones_array.view(), weight_array.view(), &chunk_vis_ctx)
                         .expect("unable to write uvfits")
                 );
             }
@@ -1740,12 +1739,7 @@ impl<'a> BirliContext<'a> {
                 with_increment_duration!(
                     "write",
                     ms_writer
-                        .write_vis(
-                            jones_array.view(),
-                            weight_array.view(),
-                            &chunk_vis_ctx,
-                            prep_ctx.draw_progress,
-                        )
+                        .write_vis(jones_array.view(), weight_array.view(), &chunk_vis_ctx)
                         .expect("unable to write ms")
                 );
             }

diff --git a/src/corrections.rs b/src/corrections.rs
@@ -29,7 +29,7 @@ use thiserror::Error;
 /// # Examples
 ///
 /// ```rust
-/// use birli::{correct_cable_lengths, mwalib::CorrelatorContext, VisSelection};
+/// use birli::{correct_cable_lengths, mwalib::CorrelatorContext, VisSelection, io::read_mwalib};
 ///
 /// // define our input files
 /// let metafits_path = "tests/data/1297526432_mwax/1297526432.metafits";
@@ -52,8 +52,7 @@ use thiserror::Error;
 /// let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap();
 ///
 /// // read visibilities out of the gpubox files
-/// vis_sel
-///     .read_mwalib(&corr_ctx, jones_array.view_mut(), flag_array.view_mut(), false)
+/// read_mwalib(&vis_sel, &corr_ctx, jones_array.view_mut(), flag_array.view_mut(), false)
 ///     .unwrap();
 ///
 /// correct_cable_lengths(&corr_ctx, jones_array.view_mut(), &vis_sel.coarse_chan_range, false);
@@ -158,7 +157,7 @@ pub fn correct_cable_lengths(
 /// ```rust
 /// use birli::{
 ///     FlagContext, correct_geometry, mwalib::CorrelatorContext, VisSelection,
-///     correct_cable_lengths
+///     correct_cable_lengths, io::read_mwalib
 /// };
 ///
 /// // define our input files
@@ -186,8 +185,7 @@ pub fn correct_cable_lengths(
 /// let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap();
 ///
 /// // read visibilities out of the gpubox files
-/// vis_sel
-///     .read_mwalib(&corr_ctx, jones_array.view_mut(), flag_array.view_mut(), false)
+/// read_mwalib(&vis_sel, &corr_ctx, jones_array.view_mut(), flag_array.view_mut(), false)
 ///     .unwrap();
 ///
 /// correct_cable_lengths(&corr_ctx, jones_array.view_mut(), &vis_sel.coarse_chan_range, false);
@@ -706,6 +704,7 @@ mod tests {
         approx::assert_abs_diff_eq,
         compare_jones,
         corrections::{DigitalGainCorrection, PassbandCorrection, ScrunchType},
+        io::read_mwalib,
         test_common::{get_mwa_ord_context, get_mwax_context},
         VisSelection,
     };
@@ -720,14 +719,14 @@ mod tests {
         let mut flag_array = vis_sel.allocate_flags(fine_chans_per_coarse).unwrap();
         let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap();
         // read visibilities out of the gpubox files
-        vis_sel
-            .read_mwalib(
-                &corr_ctx,
-                jones_array.view_mut(),
-                flag_array.view_mut(),
-                false,
-            )
-            .unwrap();
+        read_mwalib(
+            &vis_sel,
+            &corr_ctx,
+            jones_array.view_mut(),
+            flag_array.view_mut(),
+            false,
+        )
+        .unwrap();
 
         let coarse_chan_indices: Vec<_> = vis_sel.coarse_chan_range.clone().collect();
         let all_freqs_hz = corr_ctx.get_fine_chan_freqs_hz_array(&coarse_chan_indices);
@@ -860,14 +859,14 @@ mod tests {
         let mut flag_array = vis_sel.allocate_flags(fine_chans_per_coarse).unwrap();
         let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap();
         // read visibilities out of the gpubox files
-        vis_sel
-            .read_mwalib(
-                &corr_ctx,
-                jones_array.view_mut(),
-                flag_array.view_mut(),
-                false,
-            )
-            .unwrap();
+        read_mwalib(
+            &vis_sel,
+            &corr_ctx,
+            jones_array.view_mut(),
+            flag_array.view_mut(),
+            false,
+        )
+        .unwrap();
 
         let coarse_chan_indices: Vec<_> = vis_sel.coarse_chan_range.clone().collect();
         let all_freqs_hz = corr_ctx.get_fine_chan_freqs_hz_array(&coarse_chan_indices);
@@ -1009,14 +1008,14 @@ mod tests {
         let mut flag_array = vis_sel.allocate_flags(fine_chans_per_coarse).unwrap();
         let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap();
         // read visibilities out of the gpubox files
-        vis_sel
-            .read_mwalib(
-                &corr_ctx,
-                jones_array.view_mut(),
-                flag_array.view_mut(),
-                false,
-            )
-            .unwrap();
+        read_mwalib(
+            &vis_sel,
+            &corr_ctx,
+            jones_array.view_mut(),
+            flag_array.view_mut(),
+            false,
+        )
+        .unwrap();
 
         // ts 0, chan 0 (cc 0, fc 0), baseline 0
         let viz_0_0_0 = jones_array[(0, 0, 0)];
@@ -1147,14 +1146,14 @@ mod tests {
         let mut flag_array = vis_sel.allocate_flags(fine_chans_per_coarse).unwrap();
         let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap();
         // read visibilities out of the gpubox files
-        vis_sel
-            .read_mwalib(
-                &corr_ctx,
-                jones_array.view_mut(),
-                flag_array.view_mut(),
-                false,
-            )
-            .unwrap();
+        read_mwalib(
+            &vis_sel,
+            &corr_ctx,
+            jones_array.view_mut(),
+            flag_array.view_mut(),
+            false,
+        )
+        .unwrap();
 
         let coarse_chan_indices: Vec<_> = vis_sel.coarse_chan_range.clone().collect();
         let all_freqs_hz = corr_ctx.get_fine_chan_freqs_hz_array(&coarse_chan_indices);
@@ -1296,14 +1295,14 @@ mod tests {
         let mut flag_array = vis_sel.allocate_flags(fine_chans_per_coarse).unwrap();
         let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap();
         // read visibilities out of the gpubox files
-        vis_sel
-            .read_mwalib(
-                &corr_ctx,
-                jones_array.view_mut(),
-                flag_array.view_mut(),
-                false,
-            )
-            .unwrap();
+        read_mwalib(
+            &vis_sel,
+            &corr_ctx,
+            jones_array.view_mut(),
+            flag_array.view_mut(),
+            false,
+        )
+        .unwrap();
 
         let ant_pairs = vis_sel.get_ant_pairs(&corr_ctx.metafits_context);