From dc0ba367197e6d9db423ff4e9aae8885a604a843 Mon Sep 17 00:00:00 2001 From: "Christopher H. Jordan" Date: Thu, 20 Jul 2023 11:15:02 +0800 Subject: [PATCH] WIP. Don't use Marlu to read from mwalib. WIP: Re-implement progress bars. --- Cargo.lock | 22 ++-- Cargo.toml | 8 +- src/cli.rs | 56 ++++----- src/corrections.rs | 91 ++++++++------- src/flags.rs | 27 +++-- src/io/mod.rs | 269 +++++++++++++++++++++++++++++++++++++++---- src/lib.rs | 5 +- src/preprocessing.rs | 35 +++--- 8 files changed, 367 insertions(+), 146 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4378c43..8eaaf03 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -110,7 +110,7 @@ dependencies = [ [[package]] name = "birli" -version = "0.9.2" +version = "0.10.0" dependencies = [ "aoflagger_sys", "approx", @@ -891,7 +891,7 @@ checksum = "cef509aa9bc73864d6756f0d34d35504af3cf0844373afe9b8669a5b8005a729" dependencies = [ "console", "number_prefix", - "portable-atomic", + "portable-atomic 0.3.20", "rayon", "unicode-width", ] @@ -1101,9 +1101,9 @@ dependencies = [ [[package]] name = "marlu" -version = "0.9.2" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f80c9394e86d76322203f66758e51dd54c1696c6eeb8bae3bd59e7a15f6f3d6" +checksum = "04412543ad9d5db8c4f8c322950fbdbf218c4e1635cb67840433df13416f5366" dependencies = [ "approx", "built", @@ -1113,7 +1113,6 @@ dependencies = [ "fitsio-sys", "flate2", "hifitime", - "indicatif", "itertools", "lazy_static", "log", @@ -1321,9 +1320,18 @@ dependencies = [ [[package]] name = "portable-atomic" -version = "0.3.19" +version = "0.3.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e30165d31df606f5726b090ec7592c308a0eaf61721ff64c9a3018e344a8753e" +dependencies = [ + "portable-atomic 1.4.1", +] + +[[package]] +name = "portable-atomic" +version = "1.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26f6a7b87c2e435a3241addceeeff740ff8b7e76b74c13bf9acb17fa454ea00b" +checksum = "edc55135a600d700580e406b4de0d59cb9ad25e344a3a091a97ded2622ec4ec6" [[package]] name = "prettytable-rs" diff --git a/Cargo.toml b/Cargo.toml index 6d74025..ebfcca7 100644 --- 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 ", ] 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 index 49c5353..75332a8 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,5 +1,22 @@ //! 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, @@ -7,7 +24,7 @@ use crate::{ 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 index 8cb63c5..73e16ef 100644 --- 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); diff --git a/src/flags.rs b/src/flags.rs index 60f31d3..41eef75 100644 --- a/src/flags.rs +++ b/src/flags.rs @@ -312,7 +312,7 @@ pub fn flag_baseline_view_to_flagmask( /// /// ``` /// use birli::{FlagContext, flag_jones_array_existing, write_flags, -/// mwalib::CorrelatorContext, cxx_aoflagger_new, VisSelection}; +/// mwalib::CorrelatorContext, cxx_aoflagger_new, VisSelection, io::read_mwalib}; /// use tempfile::tempdir; /// /// // define our input files @@ -339,8 +339,7 @@ pub fn flag_baseline_view_to_flagmask( /// 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(); /// /// // use the default strategy file location for MWA @@ -460,7 +459,7 @@ pub fn flag_jones_array( /// Here's an example of how to flag some visibility files /// /// ```rust -/// use birli::{FlagContext, write_flags, mwalib::CorrelatorContext, VisSelection}; +/// use birli::{FlagContext, write_flags, mwalib::CorrelatorContext, VisSelection, io::read_mwalib}; /// use tempfile::tempdir; /// /// // define our input files @@ -502,8 +501,7 @@ pub fn flag_jones_array( /// 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(); /// /// // write the flags to disk as .mwaf @@ -734,6 +732,7 @@ mod tests_aoflagger { use crate::{ flags::{flag_jones_array, flag_jones_array_existing, FlagContext}, + io::read_mwalib, BirliError, VisSelection, }; use aoflagger_sys::cxx_aoflagger_new; @@ -872,14 +871,14 @@ mod tests_aoflagger { ) .unwrap(); let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap(); - 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 strategy_filename = &aoflagger.FindStrategyFileMWA(); diff --git a/src/io/mod.rs b/src/io/mod.rs index eaed138..d612776 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -11,7 +11,10 @@ use std::{ path::{Path, PathBuf}, }; -use log::trace; +use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle}; +use itertools::izip; +use log::{trace, warn}; +use marlu::{mwalib, SelectionError, VisSelection}; use crate::{ marlu::{ @@ -19,9 +22,10 @@ use crate::{ hifitime::Duration, io::{ms::MeasurementSetWriter, uvfits::UvfitsWriter, VisWrite}, mwalib::{CorrelatorContext, MwalibError}, + rayon::prelude::*, Jones, LatLngHeight, MwaObsContext, ObsContext, RADec, VisContext, ENH, }, - ndarray::{ArrayView3, ArrayViewMut3}, + ndarray::prelude::*, }; use self::error::IOError; @@ -82,6 +86,227 @@ pub trait ReadableVis: Sync + Send { ) -> Result<(), IOError>; } +/// Read the visibilities for this selection into the jones array using mwalib, +/// flag visiblities if they are not provided. +/// +/// # Errors +/// +/// Can raise [`SelectionError::BadArrayShape`] if `jones_array` or `flag_array` does not match the +/// expected shape of this selection. +/// +/// # Examples +/// +/// ```rust +/// use marlu::{mwalib::CorrelatorContext, VisSelection}; +/// use birli::io::read_mwalib; +/// +/// // define our input files +/// let metafits_path = "tests/data/1297526432_mwax/1297526432.metafits"; +/// let gpufits_paths = vec![ +/// "tests/data/1297526432_mwax/1297526432_20210216160014_ch117_000.fits", +/// "tests/data/1297526432_mwax/1297526432_20210216160014_ch117_001.fits", +/// "tests/data/1297526432_mwax/1297526432_20210216160014_ch118_000.fits", +/// "tests/data/1297526432_mwax/1297526432_20210216160014_ch118_001.fits", +/// ]; +/// +/// // Create an mwalib::CorrelatorContext for accessing visibilities. +/// let corr_ctx = CorrelatorContext::new(metafits_path, &gpufits_paths).unwrap(); +/// +/// // Determine which timesteps and coarse channels we want to use +/// let img_timestep_idxs = &corr_ctx.common_timestep_indices; +/// let good_timestep_idxs = &corr_ctx.common_good_timestep_indices; +/// +/// let mut vis_sel = VisSelection::from_mwalib(&corr_ctx).unwrap(); +/// vis_sel.timestep_range = +/// *img_timestep_idxs.first().unwrap()..(*img_timestep_idxs.last().unwrap() + 1); +/// +/// // Create a blank array to store flags and visibilities +/// let fine_chans_per_coarse = corr_ctx.metafits_context.num_corr_fine_chans_per_coarse; +/// 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 +/// read_mwalib(&vis_sel, &corr_ctx, jones_array.view_mut(), flag_array.view_mut(), false) +/// .unwrap(); +/// +/// let dims_common = jones_array.dim(); +/// +/// // now try only with good timesteps +/// vis_sel.timestep_range = +/// *good_timestep_idxs.first().unwrap()..(*good_timestep_idxs.last().unwrap() + 1); +/// +/// // read visibilities out of the gpubox files +/// 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_mwalib(&vis_sel, &corr_ctx, jones_array.view_mut(), flag_array.view_mut(), false) +/// .unwrap(); +/// +/// let dims_good = jones_array.dim(); +/// +/// // different selections have different sized arrays. +/// assert_ne!(dims_common, dims_good); +/// ``` +pub fn read_mwalib( + vis_sel: &VisSelection, + corr_ctx: &CorrelatorContext, + mut jones_array: ArrayViewMut3>, + mut flag_array: ArrayViewMut3, + draw_progress: bool, +) -> Result<(), SelectionError> { + let fine_chans_per_coarse = corr_ctx.metafits_context.num_corr_fine_chans_per_coarse; + let shape = vis_sel.get_shape(fine_chans_per_coarse); + let (num_timesteps, _, _) = shape; + let num_coarse_chans = vis_sel.coarse_chan_range.len(); + + if jones_array.dim() != shape { + return Err(SelectionError::BadArrayShape { + argument: "jones_array".to_string(), + function: "VisSelection::read_mwalib".to_string(), + expected: format!("{shape:?}"), + received: format!("{:?}", jones_array.dim()), + }); + }; + + if flag_array.dim() != shape { + return Err(SelectionError::BadArrayShape { + argument: "flag_array".to_string(), + function: "VisSelection::read_mwalib".to_string(), + expected: format!("{shape:?}"), + received: format!("{:?}", flag_array.dim()), + }); + }; + + // since we are using read_by_baseline_into_buffer, the visibilities are read in order: + // baseline,frequency,pol,r,i + + // compiler optimization + let floats_per_chan = 8; + assert_eq!( + corr_ctx.metafits_context.num_visibility_pols * 2, + floats_per_chan + ); + + let floats_per_baseline = floats_per_chan * fine_chans_per_coarse; + let floats_per_hdu = floats_per_baseline * corr_ctx.metafits_context.num_baselines; + + // Progress bar draw target + let draw_target = if draw_progress { + ProgressDrawTarget::stderr() + } else { + ProgressDrawTarget::hidden() + }; + // a progress bar containing the progress bars associated with this method + let multi_progress = MultiProgress::with_draw_target(draw_target); + // a vector of progress bars for the visibility reading progress of each channel. + let read_progress: Vec = vis_sel + .coarse_chan_range + .clone() + .map(|mwalib_coarse_chan_idx| { + let channel_progress = multi_progress.add( + ProgressBar::new(num_timesteps as _) + .with_style( + ProgressStyle::default_bar() + .template("{msg:16}: [{wide_bar:.blue}] {pos:4}/{len:4}") + .unwrap() + .progress_chars("=> "), + ) + .with_position(0) + .with_message(format!("coarse_chan {mwalib_coarse_chan_idx:03}")), + ); + channel_progress.set_position(0); + channel_progress + }) + .collect(); + // The total reading progress bar. + let total_progress = multi_progress.add( + ProgressBar::new((num_timesteps * num_coarse_chans) as _) + .with_style( + ProgressStyle::default_bar() + .template( + "{msg:16}: [{elapsed_precise}] [{wide_bar:.cyan/blue}] {percent:3}% ({eta:5})", + ) + .unwrap() + .progress_chars("=> "), + ) + .with_position(0) + .with_message("loading hdus"), + ); + + // Load HDUs from each coarse channel. arrays: [timestep][chan][baseline] + jones_array + .axis_chunks_iter_mut(Axis(1), fine_chans_per_coarse) + .into_par_iter() + .zip(flag_array.axis_chunks_iter_mut(Axis(1), fine_chans_per_coarse)) + .zip(vis_sel.coarse_chan_range.clone()) + .zip(read_progress) + .try_for_each( + |(((mut jones_array, mut flag_array), coarse_chan_idx), progress)| { + progress.set_position(0); + + // buffer: [baseline][chan][pol][complex] + let mut hdu_buffer: Vec = vec![0.0; floats_per_hdu]; + + // arrays: [chan][baseline] + for (mut jones_array, mut flag_array, timestep_idx) in izip!( + jones_array.outer_iter_mut(), + flag_array.outer_iter_mut(), + vis_sel.timestep_range.clone(), + ) { + match corr_ctx.read_by_baseline_into_buffer( + timestep_idx, + coarse_chan_idx, + hdu_buffer.as_mut_slice(), + ) { + Ok(()) => { + // arrays: [chan] + for (mut jones_array, baseline_idx) in izip!( + jones_array.axis_iter_mut(Axis(1)), + vis_sel.baseline_idxs.iter() + ) { + // buffer: [chan][pol][complex] + let hdu_baseline_chunk = &hdu_buffer + [baseline_idx * floats_per_baseline..][..floats_per_baseline]; + for (jones, hdu_chan_chunk) in izip!( + jones_array.iter_mut(), + hdu_baseline_chunk.chunks_exact(floats_per_chan) + ) { + *jones = Jones::from([ + hdu_chan_chunk[0], + hdu_chan_chunk[1], + hdu_chan_chunk[2], + hdu_chan_chunk[3], + hdu_chan_chunk[4], + hdu_chan_chunk[5], + hdu_chan_chunk[6], + hdu_chan_chunk[7], + ]); + } + } + } + Err(mwalib::GpuboxError::NoDataForTimeStepCoarseChannel { .. }) => { + warn!( + "Flagging missing HDU @ ts={}, cc={}", + timestep_idx, coarse_chan_idx + ); + flag_array.fill(true); + } + Err(e) => return Err(e), + } + + progress.inc(1); + total_progress.inc(1); + } + progress.finish(); + Ok(()) + }, + )?; + + // We're done! + total_progress.finish(); + + Ok(()) +} + /// Write the given ndarrays of flags and [`Jones`] matrix visibilities to a /// uvfits file. /// @@ -98,7 +323,7 @@ pub trait ReadableVis: Sync + Send { /// marlu::mwalib::CorrelatorContext, /// get_weight_factor, /// flag_to_weight_array, -/// VisSelection +/// VisSelection, io::read_mwalib /// }; /// /// // define our input files @@ -125,8 +350,7 @@ pub trait ReadableVis: Sync + Send { /// 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(); /// /// // write the visibilities to disk as .uvfits @@ -221,10 +445,11 @@ pub fn write_uvfits>( obs_ctx.name.as_deref(), antenna_names, antenna_positions, + true, None, )?; - uvfits_writer.write_vis(jones_array, weight_array, &vis_ctx, draw_progress)?; + uvfits_writer.write_vis(jones_array, weight_array, &vis_ctx)?; uvfits_writer.finalise()?; @@ -249,7 +474,7 @@ pub fn write_uvfits>( /// marlu::mwalib::CorrelatorContext, /// get_weight_factor, /// flag_to_weight_array, -/// FlagContext, +/// FlagContext, io::read_mwalib /// }; /// /// // define our input files @@ -276,8 +501,7 @@ pub fn write_uvfits>( /// 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(); /// /// // write the visibilities to disk as .ms @@ -349,6 +573,7 @@ pub fn write_ms>( obs_ctx.array_pos, obs_ctx.ant_positions_geodetic().collect(), Duration::from_seconds(corr_ctx.metafits_context.dut1.unwrap_or(0.0)), + true, ); ms_writer @@ -356,12 +581,7 @@ pub fn write_ms>( .unwrap(); ms_writer - .write_vis( - jones_array.view(), - weight_array.view(), - &vis_ctx, - draw_progress, - ) + .write_vis(jones_array.view(), weight_array.view(), &vis_ctx) .unwrap(); trace!("end write_ms"); @@ -375,7 +595,8 @@ pub fn write_ms>( mod tests_aoflagger { use crate::{ flags::{flag_jones_array_existing, flag_to_weight_array, get_weight_factor}, - write_uvfits, FlagContext, VisSelection, + io::{read_mwalib, write_uvfits}, + FlagContext, VisSelection, }; use aoflagger_sys::cxx_aoflagger_new; use fitsio::errors::check_status as fits_check_status; @@ -430,14 +651,14 @@ mod tests_aoflagger { ) .unwrap(); let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap(); - 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(); // use the default strategy file location for MWA let strategy_filename = &aoflagger.FindStrategyFileMWA(); diff --git a/src/lib.rs b/src/lib.rs index be1d3c5..e76b512 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,7 +10,7 @@ //! write_flags, //! mwalib::CorrelatorContext, write_uvfits, //! get_weight_factor, flag_to_weight_array, -//! FlagContext, VisSelection +//! FlagContext, VisSelection, io::read_mwalib //! }; //! use tempfile::tempdir; //! @@ -51,8 +51,7 @@ //! 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(); //! //! // write the flags to disk as .mwaf diff --git a/src/preprocessing.rs b/src/preprocessing.rs index dd9d74e..9b593ae 100644 --- a/src/preprocessing.rs +++ b/src/preprocessing.rs @@ -277,9 +277,10 @@ mod tests { use crate::{ flag_to_weight_array, flags::get_weight_factor, + io::{read_mwalib, write_uvfits}, passband_gains::PFB_JAKE_2022_200HZ, test_common::{compare_uvfits_with_csv, get_1254670392_avg_paths}, - write_uvfits, FlagContext, VisSelection, + FlagContext, VisSelection, }; use super::*; @@ -324,14 +325,14 @@ mod tests { ) .unwrap(); let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap(); - 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(); // generate weights let weight_factor = get_weight_factor(&corr_ctx); @@ -400,14 +401,14 @@ mod tests { let fine_chans_per_coarse = corr_ctx.metafits_context.num_corr_fine_chans_per_coarse; 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(); - 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 mut weight_array = vis_sel.allocate_weights(fine_chans_per_coarse).unwrap(); weight_array.fill(get_weight_factor(&corr_ctx) as _);