diff --git a/benches/bench.rs b/benches/bench.rs index 2cf5d07..4160f4f 100644 --- a/benches/bench.rs +++ b/benches/bench.rs @@ -35,9 +35,8 @@ fn fee(c: &mut Criterion) { let iau_order = true; b.iter(|| { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - beam.calc_jones_pair( - az, - za, + beam.calc_jones( + (az, za), freq, &delays, &s, @@ -60,9 +59,8 @@ fn fee(c: &mut Criterion) { let iau_order = true; let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); b.iter(|| { - beam.calc_jones_pair( - az, - za, + beam.calc_jones( + (az, za), freq, &delays, &s, @@ -86,9 +84,8 @@ fn fee(c: &mut Criterion) { let iau_order = true; let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); // Prime the cache. - beam.calc_jones_pair( - az, - za, + beam.calc_jones( + (az, za), freq, &delays, &s, @@ -98,9 +95,8 @@ fn fee(c: &mut Criterion) { ) .unwrap(); b.iter(|| { - beam.calc_jones_pair( - az, - za, + beam.calc_jones( + (az, za), freq, &delays, &s, @@ -128,9 +124,8 @@ fn fee(c: &mut Criterion) { let iau_order = true; let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); // Prime the cache. - beam.calc_jones_pair( - az[0], - za[0], + beam.calc_jones( + (az[0], za[0]), freq, &delays, &s, @@ -140,9 +135,8 @@ fn fee(c: &mut Criterion) { ) .unwrap(); b.iter(|| { - beam.calc_jones_array_pair( - &az, - &za, + beam.calc_jones_array( + (&az, &za), freq, &delays, &s, @@ -172,9 +166,8 @@ fn fee(c: &mut Criterion) { let iau_order = true; let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); // Prime the cache. - beam.calc_jones_pair( - az[0], - za[0], + beam.calc_jones( + (az[0], za[0]), freq, &delays, &s, @@ -187,9 +180,8 @@ fn fee(c: &mut Criterion) { az.par_iter() .zip(za.par_iter()) .map(|(&a, &z)| { - beam.calc_jones_pair( - a, - z, + beam.calc_jones( + (a, z), freq, &delays, &s, @@ -253,9 +245,8 @@ fn fee(c: &mut Criterion) { c.bench_function("calc_jones_array 100000 dirs", |b| { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); // Prime the cache. - beam.calc_jones_pair( - az_double[0], - za_double[0], + beam.calc_jones( + (az_double[0], za_double[0]), freqs[0], delays.as_slice().unwrap(), amps.as_slice().unwrap(), @@ -265,9 +256,8 @@ fn fee(c: &mut Criterion) { ) .unwrap(); b.iter(|| { - beam.calc_jones_array_pair( - &az_double, - &za_double, + beam.calc_jones_array( + (&az_double, &za_double), freqs[0], delays.as_slice().unwrap(), amps.as_slice().unwrap(), diff --git a/build.rs b/build.rs index 16074c3..6176aa1 100644 --- a/build.rs +++ b/build.rs @@ -314,7 +314,7 @@ mod gpu { }; for arch in arches { - hip_target.flag(&format!("--offload-arch={arch}")); + hip_target.flag(format!("--offload-arch={arch}")); } match env::var("DEBUG").as_deref() { diff --git a/examples/analytic_cuda.rs b/examples/analytic_cuda.rs index ba3aef8..4e08e54 100644 --- a/examples/analytic_cuda.rs +++ b/examples/analytic_cuda.rs @@ -55,8 +55,8 @@ fn main() -> Result<(), Box> { // Set up the directions to test. The type depends on the GPU precision. let mut azels = Vec::with_capacity(num_directions); for i in 0..num_directions { - let az = 0.4 + 0.3 * PI * (i / num_directions) as f64; - let za = 0.3 + 0.4 * FRAC_PI_2 * (i / num_directions) as f64; + let az = 0.4 + 0.3 * PI * i as f64 / num_directions as f64; + let za = 0.3 + 0.4 * FRAC_PI_2 * i as f64 / num_directions as f64; azels.push(AzEl::from_radians(az, FRAC_PI_2 - za)); } diff --git a/examples/analytic_hip.rs b/examples/analytic_hip.rs index e073fdf..4b304e0 100644 --- a/examples/analytic_hip.rs +++ b/examples/analytic_hip.rs @@ -55,8 +55,8 @@ fn main() -> Result<(), Box> { // Set up the directions to test. The type depends on the GPU precision. let mut azels = Vec::with_capacity(num_directions); for i in 0..num_directions { - let az = 0.4 + 0.3 * PI * (i / num_directions) as f64; - let za = 0.3 + 0.4 * FRAC_PI_2 * (i / num_directions) as f64; + let az = 0.4 + 0.3 * PI * i as f64 / num_directions as f64; + let za = 0.3 + 0.4 * FRAC_PI_2 * i as f64 / num_directions as f64; azels.push(AzEl::from_radians(az, FRAC_PI_2 - za)); } @@ -89,6 +89,7 @@ fn main() -> Result<(), Box> { norm_to_zenith, )?; + #[allow(clippy::useless_conversion)] let diff = jones[(0, 0, 0)] - Jones::::from(jones_cpu); println!("Difference between first GPU and CPU Jones matrices"); diff --git a/examples/fee_cuda.rs b/examples/fee_cuda.rs index bc5c206..ab9f505 100644 --- a/examples/fee_cuda.rs +++ b/examples/fee_cuda.rs @@ -59,8 +59,8 @@ fn main() -> Result<(), Box> { // Set up the directions to test. The type depends on the GPU precision. let mut azels = Vec::with_capacity(num_directions); for i in 0..num_directions { - let az = 0.4 + 0.3 * PI * (i / num_directions) as f64; - let za = 0.3 + 0.4 * FRAC_PI_2 * (i / num_directions) as f64; + let az = 0.4 + 0.3 * PI * i as f64 / num_directions as f64; + let za = 0.3 + 0.4 * FRAC_PI_2 * i as f64 / num_directions as f64; azels.push(AzEl::from_radians(az, FRAC_PI_2 - za)); } diff --git a/examples/fee_hip.rs b/examples/fee_hip.rs index 4006a59..4d81b4f 100644 --- a/examples/fee_hip.rs +++ b/examples/fee_hip.rs @@ -59,8 +59,8 @@ fn main() -> Result<(), Box> { // Set up the directions to test. The type depends on the GPU precision. let mut azels = Vec::with_capacity(num_directions); for i in 0..num_directions { - let az = 0.4 + 0.3 * PI * (i / num_directions) as f64; - let za = 0.3 + 0.4 * FRAC_PI_2 * (i / num_directions) as f64; + let az = 0.4 + 0.3 * PI * i as f64 / num_directions as f64; + let za = 0.3 + 0.4 * FRAC_PI_2 * i as f64 / num_directions as f64; azels.push(AzEl::from_radians(az, FRAC_PI_2 - za)); } @@ -95,6 +95,7 @@ fn main() -> Result<(), Box> { iau_order, )?; + #[allow(clippy::useless_conversion)] let diff = jones[(0, 0, 0)] - Jones::::from(jones_cpu); println!("Difference between first GPU and CPU Jones matrices"); diff --git a/src/analytic/ffi/mod.rs b/src/analytic/ffi/mod.rs index 9be3e72..20587b4 100644 --- a/src/analytic/ffi/mod.rs +++ b/src/analytic/ffi/mod.rs @@ -146,9 +146,8 @@ pub unsafe extern "C" fn analytic_calc_jones( let amps_s = slice::from_raw_parts(amps, num_amps as usize); // Using the passed-in beam, get the beam response (Jones matrix). - match beam.calc_jones_pair( - az_rad, - za_rad, + match beam.calc_jones( + (az_rad, za_rad), freq_hz, delays_s, amps_s, @@ -243,9 +242,8 @@ pub unsafe extern "C" fn analytic_calc_jones_array( let amps_s = slice::from_raw_parts(amps, num_amps as usize); let results_s = slice::from_raw_parts_mut(jones.cast(), num_azza as usize); - ffi_error!(beam.calc_jones_array_pair_inner( - az, - za, + ffi_error!(beam.calc_jones_array_inner( + (az, za), freq_hz, delays_s, amps_s, diff --git a/src/analytic/ffi/tests.rs b/src/analytic/ffi/tests.rs index 600ba2f..1f84a53 100644 --- a/src/analytic/ffi/tests.rs +++ b/src/analytic/ffi/tests.rs @@ -301,9 +301,8 @@ fn test_calc_jones_gpu_via_ffi() { for (mut out, freq) in out.outer_iter_mut().zip(freqs) { unsafe { let cpu_results = (*beam) - .calc_jones_array_pair( - &az, - &za, + .calc_jones_array( + (&az, &za), freq, delays.as_slice().unwrap(), amps.as_slice().unwrap(), diff --git a/src/analytic/gpu/tests.rs b/src/analytic/gpu/tests.rs index f71e350..8f296bc 100644 --- a/src/analytic/gpu/tests.rs +++ b/src/analytic/gpu/tests.rs @@ -47,9 +47,8 @@ fn test_analytic( { for (mut out, &freq) in out.outer_iter_mut().zip(freqs) { let cpu_results = beam - .calc_jones_array_pair( - &az, - &za, + .calc_jones_array( + (&az, &za), freq, delays.as_slice().unwrap(), amps.as_slice().unwrap(), @@ -172,9 +171,8 @@ fn test_cram() { // Compare with CPU. let cpu_results = beam - .calc_jones_pair( - az_rad[0], - za_rad[0], + .calc_jones( + (az_rad[0], za_rad[0]), freq_hz[0], delays.as_slice().unwrap(), amps.as_slice().unwrap(), diff --git a/src/analytic/mod.rs b/src/analytic/mod.rs index b77bb26..6fef637 100644 --- a/src/analytic/mod.rs +++ b/src/analytic/mod.rs @@ -21,7 +21,10 @@ use std::f64::consts::{FRAC_PI_2, TAU}; use marlu::{c64, constants::VEL_C, rayon, AzEl, Jones}; use rayon::prelude::*; -use crate::constants::{DELAY_STEP, MWA_DPL_SEP}; +use crate::{ + constants::{DELAY_STEP, MWA_DPL_HGT, MWA_DPL_HGT_RTS, MWA_DPL_SEP}, + direction::HorizCoord, +}; #[cfg(any(feature = "cuda", feature = "hip"))] use ndarray::prelude::*; @@ -42,13 +45,14 @@ impl AnalyticType { /// type. pub fn get_default_dipole_height(self) -> f64 { match self { - AnalyticType::MwaPb => 0.278, - AnalyticType::Rts => 0.30, + AnalyticType::MwaPb => MWA_DPL_HGT, + AnalyticType::Rts => MWA_DPL_HGT_RTS, } } } -/// The main struct to be used for calculating analytic pointings. +/// The struct used to calculate beam-response Jones matrices for the analytic +/// beam implementation. pub struct AnalyticBeam { /// The height of the MWA dipoles we're simulating \[metres\]. /// @@ -134,53 +138,40 @@ impl AnalyticBeam { /// 16 elements, and `amps` can have 16 or 32 elements. A CRAM tile has 8 /// bowties per row; `delays` must have 64 elements, and `amps` can have 64 /// or 128 elements. - pub fn calc_jones( - &self, - azel: AzEl, - freq_hz: u32, - delays: &[u32], - amps: &[f64], - latitude_rad: f64, - norm_to_zenith: bool, - ) -> Result, AnalyticBeamError> { - self.calc_jones_pair( - azel.az, - azel.za(), - freq_hz, - delays, - amps, - latitude_rad, - norm_to_zenith, - ) - } - - /// Calculate the beam-response Jones matrix for a given direction and - /// pointing. /// - /// `delays` and `amps` apply to each dipole in an MWA tile in the M&C - /// order; see - /// . - /// `delays` *must* have `bowties_per_row * bowties_per_row` elements (which - /// was declared when `AnalyticBeam` was created), whereas `amps` can have - /// this number or double elements; if the former is given, then these map - /// 1:1 with bowties. If double are given, then the *smallest* of the two - /// amps corresponding to a bowtie's dipoles is used. + /// # Examples /// - /// e.g. A normal MWA tile has 4 bowties per row. `delays` must then have - /// 16 elements, and `amps` can have 16 or 32 elements. A CRAM tile has 8 - /// bowties per row; `delays` must have 64 elements, and `amps` can have 64 - /// or 128 elements. - #[allow(clippy::too_many_arguments)] - pub fn calc_jones_pair( + /// ``` + /// use std::f64::consts::FRAC_PI_2; + /// + /// use marlu::{AzEl, Jones, constants::MWA_LAT_RAD}; + /// use mwa_hyperbeam::analytic::AnalyticBeam; + /// + /// let direction = AzEl::from_radians(0.4, 0.7); + /// let freq_hz = 150e6 as u32; + /// let delays = vec![0; 16]; + /// let amps = vec![1.0; 16]; + /// let latitude_rad = MWA_LAT_RAD; + /// let norm_to_zenith = true; + /// let beam = AnalyticBeam::new_rts(); + /// let result = beam.calc_jones(direction, freq_hz, &delays, &s, latitude_rad, norm_to_zenith).unwrap(); + /// + /// // Floats can be used too. + /// let direction = (0.4, FRAC_PI_2 - 0.7); + /// let result2 = beam.calc_jones(direction, freq_hz, &delays, &s, latitude_rad, norm_to_zenith).unwrap(); + /// assert_eq!(result, result2); + /// ``` + pub fn calc_jones( &self, - az_rad: f64, - za_rad: f64, + direction: C, freq_hz: u32, delays: &[u32], amps: &[f64], latitude_rad: f64, norm_to_zenith: bool, ) -> Result, AnalyticBeamError> { + let az_rad = direction.get_az(); + let za_rad = direction.get_za(); if za_rad > FRAC_PI_2 { return Err(AnalyticBeamError::BelowHorizon { za: za_rad }); } @@ -222,11 +213,11 @@ impl AnalyticBeam { Ok(jones) } - /// Calculate the beam-response Jones matrices for many directions - /// given a pointing and latitude. This is basically a wrapper around - /// `calc_jones` that efficiently calculates the Jones matrices in - /// parallel. The number of parallel threads used can be controlled by - /// setting `RAYON_NUM_THREADS`. + /// Calculate the beam-response Jones matrices for many directions given a + /// pointing and latitude. This is basically a wrapper around `calc_jones` + /// that efficiently calculates the Jones matrices in parallel. The number + /// of parallel threads used can be controlled by setting + /// `RAYON_NUM_THREADS`. /// /// `delays` and `amps` apply to each dipole in an MWA tile in the M&C /// order; see @@ -241,18 +232,48 @@ impl AnalyticBeam { /// 16 elements, and `amps` can have 16 or 32 elements. A CRAM tile has 8 /// bowties per row; `delays` must have 64 elements, and `amps` can have 64 /// or 128 elements. - pub fn calc_jones_array( + /// + /// # Examples + /// + /// ``` + /// use std::f64::consts::FRAC_PI_2; + /// + /// use marlu::{AzEl, Jones, constants::MWA_LAT_RAD}; + /// use mwa_hyperbeam::analytic::AnalyticBeam; + /// + /// let directions = vec![AzEl::from_radians(0.4, 0.7), AzEl::from_radians(0.5, 0.8)]; + /// let freq_hz = 150e6 as u32; + /// let delays = vec![0; 16]; + /// let amps = vec![1.0; 16]; + /// let latitude_rad = MWA_LAT_RAD; + /// let norm_to_zenith = true; + /// let beam = AnalyticBeam::new_rts(); + /// let results = beam.calc_jones_array(directions, freq_hz, &delays, &s, latitude_rad, norm_to_zenith).unwrap(); + /// + /// // Floats can be used, but these need to be grouped by azimuth and ZA. + /// let azimuths = vec![0.4, 0.5]; + /// let zenith_angles = vec![FRAC_PI_2 - 0.7, FRAC_PI_2 - 0.8]; + /// let results2 = beam.calc_jones_array((&azimuths, &zenith_angles), freq_hz, &delays, &s, latitude_rad, norm_to_zenith).unwrap(); + /// assert_eq!(results, results2); + /// ``` + pub fn calc_jones_array( &self, - azels: &[AzEl], + directions: I, freq_hz: u32, delays: &[u32], amps: &[f64], latitude_rad: f64, norm_to_zenith: bool, - ) -> Result>, AnalyticBeamError> { - let mut results = vec![Jones::default(); azels.len()]; + ) -> Result>, AnalyticBeamError> + where + C: HorizCoord, + I: IntoParallelIterator, + I2: IndexedParallelIterator, + { + let directions = directions.into_par_iter(); + let mut results = vec![Jones::default(); directions.len()]; self.calc_jones_array_inner( - azels, + directions, freq_hz, delays, amps, @@ -263,9 +284,10 @@ impl AnalyticBeam { Ok(results) } - /// Calculate the Jones matrices for many directions given a pointing and - /// latitude. This is the same as `calc_jones_array` but uses pre-allocated - /// memory. + /// Calculate the beam-response Jones matrices for many directions given a + /// pointing and latitude. This is the same as `calc_jones_array` but uses + /// pre-allocated memory. `results` should have a length equal to or greater + /// than `directions`. /// /// `delays` and `amps` apply to each dipole in an MWA tile in the M&C /// order; see @@ -280,183 +302,51 @@ impl AnalyticBeam { /// 16 elements, and `amps` can have 16 or 32 elements. A CRAM tile has 8 /// bowties per row; `delays` must have 64 elements, and `amps` can have 64 /// or 128 elements. - #[allow(clippy::too_many_arguments)] - pub fn calc_jones_array_inner( - &self, - azels: &[AzEl], - freq_hz: u32, - delays: &[u32], - amps: &[f64], - latitude_rad: f64, - norm_to_zenith: bool, - results: &mut [Jones], - ) -> Result<(), AnalyticBeamError> { - for azel in azels { - let za = azel.za(); - if za > FRAC_PI_2 { - return Err(AnalyticBeamError::BelowHorizon { za }); - } - } - let num_bowties = usize::from(self.bowties_per_row * self.bowties_per_row); - if delays.len() != num_bowties { - return Err(AnalyticBeamError::IncorrectDelaysLength { - got: delays.len(), - expected: num_bowties, - }); - } - if amps.len() != num_bowties && amps.len() != num_bowties * 2 { - return Err(AnalyticBeamError::IncorrectAmpsLength { - got: amps.len(), - expected1: num_bowties, - expected2: num_bowties * 2, - }); - } - - let amps = fix_amps(amps, delays); - let (amps, delays) = if matches!(self.beam_type, AnalyticType::Rts) { - reorder_to_rts(&s, delays) - } else { - (amps.to_vec(), delay_ints_to_floats(delays)) - }; - - let lambda_m = VEL_C / freq_hz as f64; - let (s_lat, c_lat) = latitude_rad.sin_cos(); - azels - .par_iter() - .zip(results.par_iter_mut()) - .try_for_each(|(&azel, result)| { - if azel.za() > FRAC_PI_2 { - return Err(AnalyticBeamError::BelowHorizon { za: azel.za() }); - } - - let j = self.calc_jones_inner( - azel.az, - azel.za(), - lambda_m, - latitude_rad, - s_lat, - c_lat, - &delays, - &s, - norm_to_zenith, - ); - *result = j; - - Ok(()) - }) - } - - /// Calculate the beam-response Jones matrices for many directions given a - /// pointing. This is basically a wrapper around `calc_jones` that - /// efficiently calculates the Jones matrices in parallel. The number of - /// parallel threads used can be controlled by setting `RAYON_NUM_THREADS`. /// - /// `delays` and `amps` apply to each dipole in an MWA tile in the M&C - /// order; see - /// . - /// `delays` *must* have `bowties_per_row * bowties_per_row` elements (which - /// was declared when `AnalyticBeam` was created), whereas `amps` can have - /// this number or double elements; if the former is given, then these map - /// 1:1 with bowties. If double are given, then the *smallest* of the two - /// amps corresponding to a bowtie's dipoles is used. + /// # Examples /// - /// e.g. A normal MWA tile has 4 bowties per row. `delays` must then have - /// 16 elements, and `amps` can have 16 or 32 elements. A CRAM tile has 8 - /// bowties per row; `delays` must have 64 elements, and `amps` can have 64 - /// or 128 elements. - #[allow(clippy::too_many_arguments)] - pub fn calc_jones_array_pair( - &self, - az_rad: &[f64], - za_rad: &[f64], - freq_hz: u32, - delays: &[u32], - amps: &[f64], - latitude_rad: f64, - norm_to_zenith: bool, - ) -> Result>, AnalyticBeamError> { - for &za in za_rad { - if za > FRAC_PI_2 { - return Err(AnalyticBeamError::BelowHorizon { za }); - } - } - let num_bowties = usize::from(self.bowties_per_row * self.bowties_per_row); - if delays.len() != num_bowties { - return Err(AnalyticBeamError::IncorrectDelaysLength { - got: delays.len(), - expected: num_bowties, - }); - } - if amps.len() != num_bowties && amps.len() != num_bowties * 2 { - return Err(AnalyticBeamError::IncorrectAmpsLength { - got: amps.len(), - expected1: num_bowties, - expected2: num_bowties * 2, - }); - } - - let amps = fix_amps(amps, delays); - let (amps, delays) = if matches!(self.beam_type, AnalyticType::Rts) { - reorder_to_rts(&s, delays) - } else { - (amps.to_vec(), delay_ints_to_floats(delays)) - }; - - let lambda_m = VEL_C / freq_hz as f64; - let (s_lat, c_lat) = latitude_rad.sin_cos(); - let out = az_rad - .par_iter() - .zip(za_rad.par_iter()) - .map(|(&az, &za)| { - self.calc_jones_inner( - az, - za, - lambda_m, - latitude_rad, - s_lat, - c_lat, - &delays, - &s, - norm_to_zenith, - ) - }) - .collect(); - Ok(out) - } - - /// Calculate the Jones matrices for many directions given a pointing. This - /// is the same as `calc_jones_array_pair` but uses pre-allocated memory. + /// ``` + /// use std::f64::consts::FRAC_PI_2; /// - /// `delays` and `amps` apply to each dipole in an MWA tile in the M&C - /// order; see - /// . - /// `delays` *must* have `bowties_per_row * bowties_per_row` elements (which - /// was declared when `AnalyticBeam` was created), whereas `amps` can have - /// this number or double elements; if the former is given, then these map - /// 1:1 with bowties. If double are given, then the *smallest* of the two - /// amps corresponding to a bowtie's dipoles is used. + /// use marlu::{AzEl, Jones, constants::MWA_LAT_RAD}; + /// use mwa_hyperbeam::analytic::AnalyticBeam; /// - /// e.g. A normal MWA tile has 4 bowties per row. `delays` must then have - /// 16 elements, and `amps` can have 16 or 32 elements. A CRAM tile has 8 - /// bowties per row; `delays` must have 64 elements, and `amps` can have 64 - /// or 128 elements. + /// let directions = vec![AzEl::from_radians(0.4, 0.7), AzEl::from_radians(0.5, 0.8)]; + /// let freq_hz = 150e6 as u32; + /// let delays = vec![0; 16]; + /// let amps = vec![1.0; 16]; + /// let latitude_rad = MWA_LAT_RAD; + /// let norm_to_zenith = true; + /// // Make the results buffer the right size, fill with default values which will be overwritten + /// let mut results = vec![Jones::default(); directions.len()]; + /// assert_eq!(results[0][0].re, 0.0); + /// let beam = AnalyticBeam::new_rts(); + /// beam.calc_jones_array_inner(&directions, freq_hz, &delays, &s, latitude_rad, norm_to_zenith, &mut results).unwrap(); + /// assert_ne!(results[0][0].re, 0.0); + /// + /// // Floats can be used, but these need to be grouped by azimuth and ZA. + /// let azimuths = vec![0.4, 0.5]; + /// let zenith_angles = vec![FRAC_PI_2 - 0.7, FRAC_PI_2 - 0.8]; + /// let mut results2 = vec![Jones::default(); directions.len()]; + /// beam.calc_jones_array_inner((&azimuths, &zenith_angles), freq_hz, &delays, &s, latitude_rad, norm_to_zenith, &mut results2).unwrap(); + /// assert_eq!(results, results2); + /// ``` #[allow(clippy::too_many_arguments)] - pub fn calc_jones_array_pair_inner( + pub fn calc_jones_array_inner( &self, - az_rad: &[f64], - za_rad: &[f64], + directions: I, freq_hz: u32, delays: &[u32], amps: &[f64], latitude_rad: f64, norm_to_zenith: bool, results: &mut [Jones], - ) -> Result<(), AnalyticBeamError> { - for &za in za_rad { - if za > FRAC_PI_2 { - return Err(AnalyticBeamError::BelowHorizon { za }); - } - } + ) -> Result<(), AnalyticBeamError> + where + C: HorizCoord, + I: IntoParallelIterator, + I2: IndexedParallelIterator, + { let num_bowties = usize::from(self.bowties_per_row * self.bowties_per_row); if delays.len() != num_bowties { return Err(AnalyticBeamError::IncorrectDelaysLength { @@ -481,18 +371,19 @@ impl AnalyticBeam { let lambda_m = VEL_C / freq_hz as f64; let (s_lat, c_lat) = latitude_rad.sin_cos(); - az_rad - .par_iter() - .zip(za_rad.par_iter()) + directions + .into_par_iter() .zip(results.par_iter_mut()) - .try_for_each(|((&az, &za), result)| { - if za > FRAC_PI_2 { - return Err(AnalyticBeamError::BelowHorizon { za }); + .try_for_each(|(dir, result)| { + let az_rad = dir.get_az(); + let za_rad = dir.get_za(); + if za_rad > FRAC_PI_2 { + return Err(AnalyticBeamError::BelowHorizon { za: za_rad }); } let j = self.calc_jones_inner( - az, - za, + az_rad, + za_rad, lambda_m, latitude_rad, s_lat, diff --git a/src/analytic/tests.rs b/src/analytic/tests.rs index 0ef97eb..ce9676f 100644 --- a/src/analytic/tests.rs +++ b/src/analytic/tests.rs @@ -227,9 +227,8 @@ macro_rules! test_analytic { norm_to_zenith, expected, } = $args; - let result = $beam.calc_jones_pair( - az_rad, - za_rad, + let result = $beam.calc_jones( + (az_rad, za_rad), freq_hz, &delays, &s, @@ -278,9 +277,8 @@ fn mwa_pb_5() { #[test] fn mwa_pb_single_matches_array() { let beam = AnalyticBeam::new(); - let result = beam.calc_jones_pair( - 4.724779027649792, - 0.3230075857, + let result = beam.calc_jones( + (4.724779027649792, 0.3230075857), 200e6 as _, &[0, 2, 4, 6, 0, 1, 2, 3, 10, 12, 14, 16, 0, 4, 8, 12], &[1.0; 16], @@ -290,9 +288,8 @@ fn mwa_pb_single_matches_array() { assert!(result.is_ok()); let result = result.unwrap(); - let result_a = beam.calc_jones_array_pair( - &[4.724779027649792], - &[0.3230075857], + let result_a = beam.calc_jones_array( + (&[4.724779027649792], &[0.3230075857]), 200e6 as _, &[0, 2, 4, 6, 0, 1, 2, 3, 10, 12, 14, 16, 0, 4, 8, 12], &[1.0; 16], @@ -339,9 +336,8 @@ fn rts_5() { #[test] fn rts_single_matches_array() { let beam = AnalyticBeam::new_rts(); - let result = beam.calc_jones_pair( - 4.724779027649792, - 0.3230075857, + let result = beam.calc_jones( + (4.724779027649792, 0.3230075857), 200e6 as _, &[0, 2, 4, 6, 0, 1, 2, 3, 10, 12, 14, 16, 0, 4, 8, 12], &[1.0; 16], @@ -351,9 +347,8 @@ fn rts_single_matches_array() { assert!(result.is_ok()); let result = result.unwrap(); - let result_a = beam.calc_jones_array_pair( - &[4.724779027649792], - &[0.3230075857], + let result_a = beam.calc_jones_array( + (&[4.724779027649792], &[0.3230075857]), 200e6 as _, &[0, 2, 4, 6, 0, 1, 2, 3, 10, 12, 14, 16, 0, 4, 8, 12], &[1.0; 16], @@ -430,9 +425,8 @@ fn test_cram() { let norm_to_zenith = true; let beam = AnalyticBeam::new_custom(AnalyticType::MwaPb, 0.3, 8); - let result = beam.calc_jones_pair( - az_rad, - za_rad, + let result = beam.calc_jones( + (az_rad, za_rad), freq_hz, &delays, &s, diff --git a/src/bin/verify-beam-file.rs b/src/bin/verify-beam-file.rs index ceed769..20df833 100644 --- a/src/bin/verify-beam-file.rs +++ b/src/bin/verify-beam-file.rs @@ -29,8 +29,14 @@ fn test_file(beam_file: &str) -> Result<(), InitFEEBeamError> { for &file_freq in beam.get_freqs() { println!("Testing freq {file_freq}"); // If this blows up, we know there's a problem... - beam.calc_jones_pair( - 0.0, 0.0, file_freq, &[0; 16], &[1.0; 16], false, None, false, + beam.calc_jones( + (0.0, 0.0), + file_freq, + &[0; 16], + &[1.0; 16], + false, + None, + false, ) .unwrap(); } diff --git a/src/constants.rs b/src/constants.rs index 65b8110..a77c87e 100644 --- a/src/constants.rs +++ b/src/constants.rs @@ -17,4 +17,10 @@ pub(crate) const J_POWER_TABLE: [c64; 4] = [ ]; /// MWA dipole separation \[metres\] -pub(crate) const MWA_DPL_SEP: f64 = 1.100; +pub const MWA_DPL_SEP: f64 = 1.100; + +/// MWA dipole height (according to mwa_pb) \[metres\] +pub const MWA_DPL_HGT: f64 = 0.278; + +/// MWA dipole height (according to the RTS) \[metres\] +pub const MWA_DPL_HGT_RTS: f64 = 0.30; diff --git a/src/direction.rs b/src/direction.rs new file mode 100644 index 0000000..b05ecce --- /dev/null +++ b/src/direction.rs @@ -0,0 +1,74 @@ +use std::f64::consts::FRAC_PI_2; + +use marlu::AzEl; + +/// A trait that describes a coordinate pair in a horizonal coordinate system. +/// MWA beam codes historially specified (azimuth, zenith angle) rather than +/// the perhaps-more-familiar "alt az" (altitude, azimuth), so we conform with +/// history here. The MWA also prefers "elevation" instead of "altitude". +pub trait HorizCoord: Copy { + /// Get the azimuth of this coordinate. + fn get_az(&self) -> f64; + /// Get the elevation of this coordinate. + fn get_el(&self) -> f64; + /// Get the zenith angle of this coordinate. + fn get_za(&self) -> f64; +} + +impl HorizCoord for &C { + fn get_az(&self) -> f64 { + (*self).get_az() + } + + fn get_el(&self) -> f64 { + (*self).get_el() + } + + fn get_za(&self) -> f64 { + (*self).get_za() + } +} + +impl HorizCoord for AzEl { + fn get_az(&self) -> f64 { + self.az + } + + fn get_el(&self) -> f64 { + self.el + } + + fn get_za(&self) -> f64 { + self.za() + } +} + +/// We assume that a tuple of floats is (azimuth, zenith angle), both in +/// radians. +impl HorizCoord for (f64, f64) { + fn get_az(&self) -> f64 { + self.0 + } + + fn get_el(&self) -> f64 { + FRAC_PI_2 - self.1 + } + + fn get_za(&self) -> f64 { + self.1 + } +} + +impl HorizCoord for (&f64, &f64) { + fn get_az(&self) -> f64 { + *self.0 + } + + fn get_el(&self) -> f64 { + FRAC_PI_2 - self.1 + } + + fn get_za(&self) -> f64 { + *self.1 + } +} diff --git a/src/fee/ffi/mod.rs b/src/fee/ffi/mod.rs index 3f2d1fa..a50e6a2 100644 --- a/src/fee/ffi/mod.rs +++ b/src/fee/ffi/mod.rs @@ -215,9 +215,8 @@ pub unsafe extern "C" fn fee_calc_jones( let amps_s = slice::from_raw_parts(amps, num_amps as usize); // Using the passed-in beam, get the beam response (Jones matrix). - match beam.calc_jones_pair( - az_rad, - za_rad, + match beam.calc_jones( + (az_rad, za_rad), freq_hz, delays_s, amps_s, @@ -337,9 +336,8 @@ pub unsafe extern "C" fn fee_calc_jones_array( let amps_s = slice::from_raw_parts(amps, num_amps as usize); let results_s = slice::from_raw_parts_mut(jones.cast(), num_azza as usize); - ffi_error!(beam.calc_jones_array_pair_inner( - az, - za, + ffi_error!(beam.calc_jones_array_inner( + (az, za), freq_hz, delays_s, amps_s, diff --git a/src/fee/ffi/tests.rs b/src/fee/ffi/tests.rs index 733df4c..10134d0 100644 --- a/src/fee/ffi/tests.rs +++ b/src/fee/ffi/tests.rs @@ -361,9 +361,8 @@ fn test_calc_jones_gpu_via_ffi() { for (mut out, freq) in out.outer_iter_mut().zip(freqs) { unsafe { let cpu_results = (*beam) - .calc_jones_array_pair( - &az, - &za, + .calc_jones_array( + (&az, &za), freq, delays.as_slice().unwrap(), amps.as_slice().unwrap(), diff --git a/src/fee/gpu/tests.rs b/src/fee/gpu/tests.rs index 95d90db..496736e 100644 --- a/src/fee/gpu/tests.rs +++ b/src/fee/gpu/tests.rs @@ -7,6 +7,7 @@ use approx::{assert_abs_diff_eq, assert_abs_diff_ne}; use marlu::constants::MWA_LAT_RAD; use ndarray::prelude::*; +use rayon::prelude::*; use serial_test::serial; use super::*; @@ -55,9 +56,8 @@ fn test_gpu_calc_jones_no_norm() { { for (mut out, freq) in out.outer_iter_mut().zip(freqs) { let cpu_results = beam - .calc_jones_array_pair( - &az, - &za, + .calc_jones_array( + (&az, &za), freq, delays.as_slice().unwrap(), amps.as_slice().unwrap(), @@ -127,9 +127,8 @@ fn test_gpu_calc_jones_w_norm() { { for (mut out, freq) in out.outer_iter_mut().zip(freqs) { let cpu_results = beam - .calc_jones_array_pair( - &az, - &za, + .calc_jones_array( + (&az, &za), freq, delays.as_slice().unwrap(), amps.as_slice().unwrap(), @@ -199,9 +198,8 @@ fn test_gpu_calc_jones_w_norm_and_parallactic() { { for (mut out, freq) in out.outer_iter_mut().zip(freqs) { let cpu_results = beam - .calc_jones_array_pair( - &az, - &za, + .calc_jones_array( + (&az, &za), freq, delays.as_slice().unwrap(), amps.as_slice().unwrap(), @@ -325,9 +323,8 @@ fn test_gpu_calc_jones_deduplication() { { for (mut out, freq) in out.outer_iter_mut().zip(freqs) { let cpu_results = beam - .calc_jones_array_pair( - &az, - &za, + .calc_jones_array( + (&az, &za), freq, delays.as_slice().unwrap(), amps.as_slice().unwrap(), @@ -415,9 +412,8 @@ fn test_gpu_calc_jones_deduplication_w_norm() { { for (mut out, freq) in out.outer_iter_mut().zip(freqs) { let cpu_results = beam - .calc_jones_array_pair( - &az, - &za, + .calc_jones_array( + (&az, &za), freq, delays.as_slice().unwrap(), amps.as_slice().unwrap(), @@ -495,9 +491,8 @@ fn test_gpu_calc_jones_no_amps() { { for (mut out, freq) in out.outer_iter_mut().zip(freqs.iter()) { let cpu_results = beam - .calc_jones_array_pair( - &az, - &za, + .calc_jones_array( + (&az, &za), *freq, delays.as_slice().unwrap(), amps.as_slice().unwrap(), diff --git a/src/fee/mod.rs b/src/fee/mod.rs index 7d78d7a..48ee274 100644 --- a/src/fee/mod.rs +++ b/src/fee/mod.rs @@ -2,8 +2,7 @@ // 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/. -//! Code to implement the MWA Fully Embedded Element (FEE) beam, a.k.a. "the -//! 2016 beam". +//! Code for the MWA Fully Embedded Element (FEE) beam, a.k.a. "the 2016 beam". mod error; mod ffi; @@ -32,12 +31,14 @@ use rayon::prelude::*; use crate::{ constants::*, + direction::HorizCoord, factorial::FACTORIAL, legendre::p1sin, types::{CacheKey, Pol}, }; -/// The main struct to be used for calculating Jones matrices. +/// The struct used to calculate beam-response Jones matrices for the Fully +/// Embedded Element (FEE) beam, a.k.a. "the 2016 beam". #[allow(clippy::upper_case_acronyms)] pub struct FEEBeam { /// The [`hdf5_metno::File`] struct associated with the opened HDF5 file. It is @@ -463,52 +464,34 @@ impl FEEBeam { /// `delays` *must* have 16 elements, whereas `amps` can have 16 or 32 /// elements; if 16 are given, then these map 1:1 with dipoles, otherwise /// the first 16 are for X dipole elements, and the next 16 are for Y. - #[allow(clippy::too_many_arguments)] - pub fn calc_jones( - &self, - azel: AzEl, - freq_hz: u32, - delays: &[u32], - amps: &[f64], - norm_to_zenith: bool, - latitude_rad: Option, - iau_order: bool, - ) -> Result, FEEBeamError> { - self.calc_jones_pair( - azel.az, - azel.za(), - freq_hz, - delays, - amps, - norm_to_zenith, - latitude_rad, - iau_order, - ) - } - - /// Calculate the beam-response Jones matrix for a given direction and - /// pointing. If `latitude_rad` is *not* supplied, the result will match - /// the original specification of the FEE beam code (possibly more useful - /// for engineers). /// - /// Astronomers are more likely to want to specify `latitude_rad` (which - /// will apply the parallactic-angle correction using the Earth latitude - /// provided for the telescope) and `iau_order`. If `latitude_rad` is - /// `None`, then `iau_reorder` does nothing. See this document for more - /// information: - /// + /// # Examples /// - /// `delays` and `amps` apply to each dipole in an MWA tile in the M&C - /// order; see - /// . - /// `delays` *must* have 16 elements, whereas `amps` can have 16 or 32 - /// elements; if 16 are given, then these map 1:1 with dipoles, otherwise - /// the first 16 are for X dipole elements, and the next 16 are for Y. + /// ``` + /// use std::f64::consts::FRAC_PI_2; + /// + /// use marlu::{AzEl, Jones, constants::MWA_LAT_RAD}; + /// use mwa_hyperbeam::fee::FEEBeam; + /// + /// let direction = AzEl::from_radians(0.4, 0.7); + /// let freq_hz = 150e6 as u32; + /// let delays = vec![0; 16]; + /// let amps = vec![1.0; 16]; + /// let norm_to_zenith = true; + /// let latitude_rad = Some(MWA_LAT_RAD); + /// let iau_order = false; + /// let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); + /// let result = beam.calc_jones(direction, freq_hz, &delays, &s, norm_to_zenith, latitude_rad, iau_order).unwrap(); + /// + /// // Floats can be used too. + /// let direction = (0.4, FRAC_PI_2 - 0.7); + /// let result2 = beam.calc_jones(direction, freq_hz, &delays, &s, norm_to_zenith, latitude_rad, iau_order).unwrap(); + /// assert_eq!(result, result2); + /// ``` #[allow(clippy::too_many_arguments)] - pub fn calc_jones_pair( + pub fn calc_jones( &self, - az_rad: f64, - za_rad: f64, + direction: C, freq_hz: u32, delays: &[u32], amps: &[f64], @@ -534,6 +517,8 @@ impl FEEBeam { // Populate the coefficients cache if it isn't already populated. let coeffs = self.get_modes(freq_hz, delays, &full_amps)?; + let az_rad = direction.get_az(); + let za_rad = direction.get_za(); let mut jones = calc_jones_direct(az_rad, za_rad, &coeffs, norm_jones); if let Some(latitude_rad) = latitude_rad { apply_parallactic_correction(az_rad, za_rad, latitude_rad, iau_order, &mut jones); @@ -553,20 +538,51 @@ impl FEEBeam { /// `delays` *must* have 16 elements, whereas `amps` can have 16 or 32 /// elements; if 16 are given, then these map 1:1 with dipoles, otherwise /// the first 16 are for X dipole elements, and the next 16 are for Y. + /// + /// # Examples + /// + /// ``` + /// use std::f64::consts::FRAC_PI_2; + /// + /// use marlu::{AzEl, Jones, constants::MWA_LAT_RAD}; + /// use mwa_hyperbeam::fee::FEEBeam; + /// + /// let directions = vec![AzEl::from_radians(0.4, 0.7), AzEl::from_radians(0.5, 0.8)]; + /// let freq_hz = 150e6 as u32; + /// let delays = vec![0; 16]; + /// let amps = vec![1.0; 16]; + /// let norm_to_zenith = true; + /// let latitude_rad = Some(MWA_LAT_RAD); + /// let iau_order = false; + /// let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); + /// let results = beam.calc_jones_array(&directions, freq_hz, &delays, &s, norm_to_zenith, latitude_rad, iau_order).unwrap(); + /// + /// // Floats can be used, but these need to be grouped by azimuth and ZA. + /// let azimuths = vec![0.4, 0.5]; + /// let zenith_angles = vec![FRAC_PI_2 - 0.7, FRAC_PI_2 - 0.8]; + /// let results2 = beam.calc_jones_array((&azimuths, &zenith_angles), freq_hz, &delays, &s, norm_to_zenith, latitude_rad, iau_order).unwrap(); + /// assert_eq!(results, results2); + /// ``` #[allow(clippy::too_many_arguments)] - pub fn calc_jones_array( + pub fn calc_jones_array( &self, - azels: &[AzEl], + directions: I, freq_hz: u32, delays: &[u32], amps: &[f64], norm_to_zenith: bool, latitude_rad: Option, iau_order: bool, - ) -> Result>, FEEBeamError> { - let mut results = vec![Jones::default(); azels.len()]; + ) -> Result>, FEEBeamError> + where + C: HorizCoord, + I: IntoParallelIterator, + I2: IndexedParallelIterator, + { + let directions = directions.into_par_iter(); + let mut results = vec![Jones::default(); directions.len()]; self.calc_jones_array_inner( - azels, + directions, freq_hz, delays, amps, @@ -578,8 +594,10 @@ impl FEEBeam { Ok(results) } - /// Calculate the Jones matrices for many directions given a pointing. This - /// is the same as `calc_jones_array` but uses pre-allocated memory. + /// Calculate the beam-response Jones matrices for many directions given a + /// pointing. This is the same as `calc_jones_array` but uses pre-allocated + /// memory. `results` should have a length equal to or greater than + /// `directions`. /// /// `delays` and `amps` apply to each dipole in an MWA tile in the M&C /// order; see @@ -587,103 +605,40 @@ impl FEEBeam { /// `delays` *must* have 16 elements, whereas `amps` can have 16 or 32 /// elements; if 16 are given, then these map 1:1 with dipoles, otherwise /// the first 16 are for X dipole elements, and the next 16 are for Y. - #[allow(clippy::too_many_arguments)] - pub fn calc_jones_array_inner( - &self, - azels: &[AzEl], - freq_hz: u32, - delays: &[u32], - amps: &[f64], - norm_to_zenith: bool, - latitude_rad: Option, - iau_order: bool, - results: &mut [Jones], - ) -> Result<(), FEEBeamError> { - if delays.len() != 16 { - return Err(FEEBeamError::IncorrectDelaysLength(delays.len())); - } - let delays: &[u32; 16] = delays.try_into().unwrap(); - if !(amps.len() == 16 || amps.len() == 32) { - return Err(FEEBeamError::IncorrectAmpsLength(amps.len())); - } - let full_amps = fix_amps(amps, delays); - - // If we're normalising the beam, get the normalisation Jones matrix here. - let norm_jones = match norm_to_zenith { - true => Some(self.get_norm_jones(freq_hz)?), - false => None, - }; - - // Populate the coefficients cache if it isn't already populated. - let coeffs = self.get_modes(freq_hz, delays, &full_amps)?; - - azels - .par_iter() - .zip(results.par_iter_mut()) - .for_each(|(&azel, result)| { - let az = azel.az; - let za = azel.za(); - let mut jones = calc_jones_direct(az, za, &coeffs, norm_jones); - if let Some(latitude_rad) = latitude_rad { - apply_parallactic_correction(az, za, latitude_rad, iau_order, &mut jones); - } - *result = jones; - }); - Ok(()) - } - - /// Calculate the Jones matrices for many directions given a pointing. This - /// is basically a wrapper around `calc_jones` that efficiently calculates - /// the Jones matrices in parallel. The number of parallel threads used can - /// be controlled by setting `RAYON_NUM_THREADS`. /// - /// `delays` and `amps` apply to each dipole in an MWA tile in the M&C - /// order; see - /// . - /// `delays` *must* have 16 elements, whereas `amps` can have 16 or 32 - /// elements; if 16 are given, then these map 1:1 with dipoles, otherwise - /// the first 16 are for X dipole elements, and the next 16 are for Y. - #[allow(clippy::too_many_arguments)] - pub fn calc_jones_array_pair( - &self, - az_rad: &[f64], - za_rad: &[f64], - freq_hz: u32, - delays: &[u32], - amps: &[f64], - norm_to_zenith: bool, - latitude_rad: Option, - iau_order: bool, - ) -> Result>, FEEBeamError> { - let mut results = vec![Jones::default(); az_rad.len()]; - self.calc_jones_array_pair_inner( - az_rad, - za_rad, - freq_hz, - delays, - amps, - norm_to_zenith, - latitude_rad, - iau_order, - &mut results, - )?; - Ok(results) - } - - /// Calculate the Jones matrices for many directions given a pointing. This - /// is the same as `calc_jones_array_pair` but uses pre-allocated memory. + /// # Examples /// - /// `delays` and `amps` apply to each dipole in an MWA tile in the M&C - /// order; see - /// . - /// `delays` *must* have 16 elements, whereas `amps` can have 16 or 32 - /// elements; if 16 are given, then these map 1:1 with dipoles, otherwise - /// the first 16 are for X dipole elements, and the next 16 are for Y. + /// ``` + /// use std::f64::consts::FRAC_PI_2; + /// + /// use marlu::{AzEl, Jones, constants::MWA_LAT_RAD}; + /// use mwa_hyperbeam::fee::FEEBeam; + /// + /// let directions = vec![AzEl::from_radians(0.4, 0.7), AzEl::from_radians(0.5, 0.8)]; + /// let freq_hz = 150e6 as u32; + /// let delays = vec![0; 16]; + /// let amps = vec![1.0; 16]; + /// let norm_to_zenith = true; + /// let latitude_rad = Some(MWA_LAT_RAD); + /// let iau_order = false; + /// // Make the results buffer the right size, fill with default values which will be overwritten + /// let mut results = vec![Jones::default(); directions.len()]; + /// assert_eq!(results[0][0].re, 0.0); + /// let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); + /// beam.calc_jones_array_inner(&directions, freq_hz, &delays, &s, norm_to_zenith, latitude_rad, iau_order, &mut results).unwrap(); + /// assert_ne!(results[0][0].re, 0.0); + /// + /// // Floats can be used, but these need to be grouped by azimuth and ZA. + /// let azimuths = vec![0.4, 0.5]; + /// let zenith_angles = vec![FRAC_PI_2 - 0.7, FRAC_PI_2 - 0.8]; + /// let mut results2 = vec![Jones::default(); directions.len()]; + /// beam.calc_jones_array_inner((&azimuths, &zenith_angles), freq_hz, &delays, &s, norm_to_zenith, latitude_rad, iau_order, &mut results2).unwrap(); + /// assert_eq!(results, results2); + /// ``` #[allow(clippy::too_many_arguments)] - pub fn calc_jones_array_pair_inner( + pub fn calc_jones_array_inner( &self, - az_rad: &[f64], - za_rad: &[f64], + directions: I, freq_hz: u32, delays: &[u32], amps: &[f64], @@ -691,7 +646,12 @@ impl FEEBeam { latitude_rad: Option, iau_order: bool, results: &mut [Jones], - ) -> Result<(), FEEBeamError> { + ) -> Result<(), FEEBeamError> + where + C: HorizCoord, + I: IntoParallelIterator, + I2: IndexedParallelIterator, + { if delays.len() != 16 { return Err(FEEBeamError::IncorrectDelaysLength(delays.len())); } @@ -710,14 +670,21 @@ impl FEEBeam { // Populate the coefficients cache if it isn't already populated. let coeffs = self.get_modes(freq_hz, delays, &full_amps)?; - az_rad - .par_iter() - .zip(za_rad.par_iter()) + directions + .into_par_iter() .zip(results.par_iter_mut()) - .for_each(|((&az, &za), result)| { - let mut jones = calc_jones_direct(az, za, &coeffs, norm_jones); + .for_each(|(dir, result)| { + let az_rad = dir.get_az(); + let za_rad = dir.get_za(); + let mut jones = calc_jones_direct(az_rad, za_rad, &coeffs, norm_jones); if let Some(latitude_rad) = latitude_rad { - apply_parallactic_correction(az, za, latitude_rad, iau_order, &mut jones); + apply_parallactic_correction( + az_rad, + za_rad, + latitude_rad, + iau_order, + &mut jones, + ); } *result = jones; }); diff --git a/src/fee/tests.rs b/src/fee/tests.rs index ed8d3c2..3e02666 100644 --- a/src/fee/tests.rs +++ b/src/fee/tests.rs @@ -575,9 +575,8 @@ fn test_get_modes2() { #[serial] fn test_calc_jones_eng() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 45.0_f64.to_radians(), - 10.0_f64.to_radians(), + let result = beam.calc_jones( + (45.0_f64.to_radians(), 10.0_f64.to_radians()), 51200000, &[0; 16], &[1.0; 16], @@ -601,9 +600,8 @@ fn test_calc_jones_eng() { #[serial] fn test_calc_jones_eng_2() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 70.0_f64.to_radians(), - 10.0_f64.to_radians(), + let result = beam.calc_jones( + (70.0_f64.to_radians(), 10.0_f64.to_radians()), 51200000, &[3, 2, 1, 0, 3, 2, 1, 0, 3, 2, 1, 0, 3, 2, 1, 0], &[ @@ -629,8 +627,14 @@ fn test_calc_jones_eng_2() { #[serial] fn test_calc_jones_eng_norm() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 0.1_f64, 0.1_f64, 150000000, &[0; 16], &[1.0; 16], true, None, false, + let result = beam.calc_jones( + (0.1_f64, 0.1_f64), + 150000000, + &[0; 16], + &[1.0; 16], + true, + None, + false, ); assert!(result.is_ok()); let jones = result.unwrap(); @@ -648,9 +652,8 @@ fn test_calc_jones_eng_norm() { #[serial] fn test_calc_jones_eng_norm_2() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 0.1_f64, - 0.1_f64, + let result = beam.calc_jones( + (0.1_f64, 0.1_f64), 150000000, &[3, 2, 1, 0, 3, 2, 1, 0, 3, 2, 1, 0, 3, 2, 1, 0], &[ @@ -676,9 +679,8 @@ fn test_calc_jones_eng_norm_2() { #[serial] fn test_calc_jones() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 45.0_f64.to_radians(), - 10.0_f64.to_radians(), + let result = beam.calc_jones( + (45.0_f64.to_radians(), 10.0_f64.to_radians()), 51200000, &[0; 16], &[1.0; 16], @@ -702,9 +704,8 @@ fn test_calc_jones() { #[serial] fn test_calc_jones_norm() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 0.1_f64, - 0.1_f64, + let result = beam.calc_jones( + (0.1_f64, 0.1_f64), 150000000, &[0; 16], &[1.0; 16], @@ -742,9 +743,8 @@ fn test_calc_jones_norm() { #[serial] fn test_calc_jones_pa() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 0.1_f64, - 0.1_f64, + let result = beam.calc_jones( + (0.1_f64, 0.1_f64), 150000000, &[0; 16], &[1.0; 16], @@ -755,8 +755,14 @@ fn test_calc_jones_pa() { assert!(result.is_ok()); let pa = result.unwrap(); - let result = beam.calc_jones_pair( - 0.1_f64, 0.1_f64, 150000000, &[0; 16], &[1.0; 16], true, None, false, + let result = beam.calc_jones( + (0.1_f64, 0.1_f64), + 150000000, + &[0; 16], + &[1.0; 16], + true, + None, + false, ); assert!(result.is_ok()); let not_pa = result.unwrap(); @@ -768,9 +774,8 @@ fn test_calc_jones_pa() { #[serial] fn test_calc_jones_iau() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 0.1_f64, - 0.1_f64, + let result = beam.calc_jones( + (0.1_f64, 0.1_f64), 150000000, &[0; 16], &[1.0; 16], @@ -781,9 +786,8 @@ fn test_calc_jones_iau() { assert!(result.is_ok()); let j_iau = result.unwrap(); - let result = beam.calc_jones_pair( - 0.1_f64, - 0.1_f64, + let result = beam.calc_jones( + (0.1_f64, 0.1_f64), 150000000, &[0; 16], &[1.0; 16], @@ -808,9 +812,10 @@ fn test_calc_jones_iau() { #[serial] fn test_calc_jones_array() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 45.0_f64.to_radians(), - 10.0_f64.to_radians(), + let az = 45.0_f64.to_radians(); + let za = 10.0_f64.to_radians(); + let result = beam.calc_jones( + (az, za), 51200000, &[0; 16], &[1.0; 16], @@ -821,9 +826,8 @@ fn test_calc_jones_array() { assert!(result.is_ok()); let jones = result.unwrap(); - let result = beam.calc_jones_array_pair( - &[45.0_f64.to_radians()], - &[10.0_f64.to_radians()], + let result = beam.calc_jones_array( + (&[az], &[za]), 51200000, &[0; 16], &[1.0; 16], @@ -837,10 +841,9 @@ fn test_calc_jones_array() { assert_eq!(jones_array.len(), 1); assert_eq!(jones, jones_array[0]); - // Ensure that FEEBeam::calc_jones_array is the same as - // FEEBeam::calc_jones_array_pair. + // Ensure that using `AzEl` gives the same results as above. let result = beam.calc_jones_array( - &[AzEl::from_degrees(45.0, 80.0)], + [AzEl::from_radians(az, FRAC_PI_2 - za)], 51200000, &[0; 16], &[1.0; 16], @@ -858,9 +861,8 @@ fn test_calc_jones_array() { #[serial] fn test_empty_cache() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 45.0_f64.to_radians(), - 10.0_f64.to_radians(), + let result = beam.calc_jones( + (45.0_f64.to_radians(), 10.0_f64.to_radians()), 51200000, &[0; 16], &[1.0; 16], @@ -893,9 +895,8 @@ fn test_get_freqs() { #[serial] fn test_cache_is_used() { let beam = FEEBeam::new("mwa_full_embedded_element_pattern.h5").unwrap(); - let result = beam.calc_jones_pair( - 45.0_f64.to_radians(), - 10.0_f64.to_radians(), + let result = beam.calc_jones( + (45.0_f64.to_radians(), 10.0_f64.to_radians()), 51200000, &[0; 16], &[1.0; 16], @@ -906,9 +907,8 @@ fn test_cache_is_used() { assert!(result.is_ok()); result.unwrap(); - let result = beam.calc_jones_pair( - 45.0_f64.to_radians(), - 10.0_f64.to_radians(), + let result = beam.calc_jones( + (45.0_f64.to_radians(), 10.0_f64.to_radians()), 51200000, &[0; 16], &[1.0; 16], diff --git a/src/ffi.rs b/src/ffi.rs index feaf2fb..91b3f04 100644 --- a/src/ffi.rs +++ b/src/ffi.rs @@ -5,7 +5,7 @@ //! General FFI code for error handling. //! //! Most of this is derived from -//! use std::{ cell::RefCell, diff --git a/src/lib.rs b/src/lib.rs index 0dba742..31e8f36 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,7 +5,8 @@ //! Primary beam code for the Murchison Widefield Array. pub mod analytic; -mod constants; +pub mod constants; +mod direction; mod factorial; pub mod fee; mod ffi; @@ -16,12 +17,16 @@ mod types; mod python; // Re-exports. +pub use analytic::AnalyticBeam; +pub use fee::FEEBeam; cfg_if::cfg_if! { if #[cfg(any(feature = "cuda", feature = "hip"))] { mod gpu; - /// The float type use in GPU code. This depends on how `hyperbeam` was + /// The float type used in GPU code. This depends on how `hyperbeam` was /// compiled (used cargo feature "gpu-single" or not). pub use gpu::{GpuFloat, GpuComplex}; + pub use analytic::AnalyticBeamGpu; + pub use fee::FEEBeamGpu; } } diff --git a/src/python/analytic.rs b/src/python/analytic.rs index c00ab34..5e2aeb0 100644 --- a/src/python/analytic.rs +++ b/src/python/analytic.rs @@ -68,9 +68,8 @@ impl AnalyticBeam { latitude_rad: f64, norm_to_zenith: Option, ) -> PyResult>> { - let jones = self.beam.calc_jones_pair( - az_rad, - za_rad, + let jones = self.beam.calc_jones( + (az_rad, za_rad), // hyperbeam expects an int for the frequency. By specifying that // Python should pass in a float, it also allows an int to be passed // in (!?). Convert the float here in Rust for usage in hyperbeam. @@ -106,9 +105,8 @@ impl AnalyticBeam { latitude_rad: f64, norm_to_zenith: Option, ) -> PyResult>> { - let jones = self.beam.calc_jones_array_pair( - &az_rad, - &za_rad, + let jones = self.beam.calc_jones_array( + (&az_rad, &za_rad), freq_hz.round() as _, &delays, &s, diff --git a/src/python/fee.rs b/src/python/fee.rs index 2e3499c..cfc1691 100644 --- a/src/python/fee.rs +++ b/src/python/fee.rs @@ -10,6 +10,7 @@ use self::ndarray::prelude::*; use num_complex::Complex64 as c64; use numpy::*; use pyo3::prelude::*; +use rayon::prelude::*; use crate::fee::FEEBeam as FEEBeamRust; #[cfg(any(feature = "cuda", feature = "hip"))] @@ -79,9 +80,8 @@ impl FEEBeam { latitude_rad: Option, iau_order: Option, ) -> PyResult>> { - let jones = self.beam.calc_jones_pair( - az_rad, - za_rad, + let jones = self.beam.calc_jones( + (az_rad, za_rad), // hyperbeam expects an int for the frequency. By specifying that // Python should pass in a float, it also allows an int to be passed // in (!?). Convert the float here in Rust for usage in hyperbeam. @@ -124,9 +124,8 @@ impl FEEBeam { latitude_rad: Option, iau_order: Option, ) -> PyResult>> { - let jones = self.beam.calc_jones_array_pair( - &az_rad, - &za_rad, + let jones = self.beam.calc_jones_array( + (&az_rad, &za_rad), freq_hz.round() as _, &delays, &s,