From d595d8f1ede0ccfa82b69426bcb1f3d6103ca2af Mon Sep 17 00:00:00 2001 From: d3v-null Date: Tue, 10 Feb 2026 15:43:19 +0800 Subject: [PATCH] analytic beams - Added support for "mwa_pb" and "RTS" analytic beams. - Updated beam subcommand to accept beam type via -b or --beam-type flag. - Fixed issue where beam subcommand failed to set up FEE beam if MWA_BEAM_FILE was not set. - Updated documentation to reflect new analytic beam options and usage instructions. --- CHANGELOG.md | 12 ++ examples/plot_beam_responses.py | 8 +- mdbook/src/defs/beam.md | 9 +- mdbook/src/installation/post.md | 3 + src/beam/analytic.rs | 337 ++++++++++++++++++++++++++++++++ src/beam/error.rs | 11 +- src/beam/fee.rs | 32 +-- src/beam/mod.rs | 36 +++- src/cli/beam.rs | 41 ++-- src/cli/common/beam/mod.rs | 117 ++++++++++- src/cli/error.rs | 5 +- src/gpu/mod.rs | 5 + src/model/cpu.rs | 30 +-- 13 files changed, 550 insertions(+), 96 deletions(-) create mode 100644 src/beam/analytic.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index b3b3e777..4e034bb6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,18 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## Unreleased + +### Added +- Support for the "mwa_pb" and "RTS" analytic beams. + +### Changed +- The beam subcommand now takes a beam type with the -b or --beam-type flag. + +### Fixed +- Using the beam subcommand would fail to set up an FEE beam if MWA_BEAM_FILE + was not set. + ## [0.7.0] - 2026-02-09 ### Fixed diff --git a/examples/plot_beam_responses.py b/examples/plot_beam_responses.py index 84f954a6..75fc3f1f 100755 --- a/examples/plot_beam_responses.py +++ b/examples/plot_beam_responses.py @@ -1,9 +1,15 @@ #!/usr/bin/env python3 +import sys + import numpy as np import matplotlib.pyplot as plt -data = np.genfromtxt(fname="beam_responses.tsv", delimiter="\t", skip_header=0) +if len(sys.argv) == 1: + file = "beam_responses.tsv" +else: + file = sys.argv[1] +data = np.genfromtxt(fname=file, delimiter="\t", skip_header=0) fig, ax = plt.subplots(1, 2, subplot_kw=dict(projection="polar")) p = ax[0].scatter(data[:, 0], data[:, 1], c=data[:, 2]) diff --git a/mdbook/src/defs/beam.md b/mdbook/src/defs/beam.md index a4221300..23df50d8 100644 --- a/mdbook/src/defs/beam.md +++ b/mdbook/src/defs/beam.md @@ -15,6 +15,13 @@ In addition, the FEE beam code needs an HDF5 file to function. See the [post-installation instructions](../installation/post.md) for information on getting that set up. +It is possible to use the "analytic" MWA beam, and this does not require an +additional file (but the FEE beam is the default selection). There are two +"flavours": `mwa_pb` and `RTS`, and these can be used by specifying `--beam-type +analytic-mwa_pb` and `--beam-type analytic-rts`, respectively. The differences +between the flavours is not huge, but I (CHJ) suggest the `RTS` flavour if in +doubt, as it seems to look a little better. + ## Errors Beam code usually does not error, but if it does it's likely because: @@ -22,4 +29,4 @@ Beam code usually does not error, but if it does it's likely because: 1. There aren't exactly 16 dipole delays; 2. There aren't exactly 16 or 32 dipole gains per tile; or 3. There's something wrong with the FEE HDF5 file. The official file is well - tested. \ No newline at end of file + tested. diff --git a/mdbook/src/installation/post.md b/mdbook/src/installation/post.md index e37169c6..7f7847d2 100644 --- a/mdbook/src/installation/post.md +++ b/mdbook/src/installation/post.md @@ -16,6 +16,9 @@ Move the `h5` file anywhere you like, and put the file path in the export MWA_BEAM_FILE=/path/to/mwa_full_embedded_element_pattern.h5 ``` +It is possible to use the analytic beam instead of the FEE beam, meaning you +don't need the HDF5 file, but the FEE beam is probably better. + See the README for [`hyperbeam`](https://github.com/MWATelescope/mwa_hyperbeam) for more info. ~~~ diff --git a/src/beam/analytic.rs b/src/beam/analytic.rs new file mode 100644 index 00000000..3ab5a9c3 --- /dev/null +++ b/src/beam/analytic.rs @@ -0,0 +1,337 @@ +// 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/. + +//! Code for analytic beam calculations. + +use log::debug; +use marlu::{AzEl, Jones}; +use mwa_hyperbeam::analytic::AnalyticType; +use ndarray::prelude::*; + +use super::{partial_to_full, validate_delays, Beam, BeamError, BeamType, Delays}; + +#[cfg(any(feature = "cuda", feature = "hip"))] +use super::{BeamGpu, DevicePointer, GpuFloat}; + +/// A wrapper of the `AnalyticBeam` struct in hyperbeam that implements the +/// [`Beam`] trait. +pub(crate) struct AnalyticBeam { + hyperbeam_object: mwa_hyperbeam::analytic::AnalyticBeam, + analytic_type: AnalyticType, + delays: Array2, + gains: Array2, + ideal_delays: [u32; 16], +} + +impl AnalyticBeam { + pub(crate) fn new_mwa_pb( + num_tiles: usize, + delays: Delays, + gains: Option>, + ) -> Result { + Self::new_inner(AnalyticType::MwaPb, num_tiles, delays, gains) + } + + pub(crate) fn new_rts( + num_tiles: usize, + delays: Delays, + gains: Option>, + ) -> Result { + Self::new_inner(AnalyticType::Rts, num_tiles, delays, gains) + } + + fn new_inner( + at: AnalyticType, + num_tiles: usize, + delays: Delays, + gains: Option>, + ) -> Result { + // Check that the delays are sensible. + validate_delays(&delays, num_tiles)?; + + let ideal_delays = delays.get_ideal_delays(); + debug!("Ideal dipole delays: {:?}", ideal_delays); + + let delays = match delays { + Delays::Full(d) => d, + Delays::Partial(d) => partial_to_full(d, num_tiles), + }; + + // If no gains were provided, assume all are alive. + let gains = match gains { + Some(g) => { + debug!("Using supplied dipole gains"); + g + } + None => { + debug!("No dipole gains supplied; setting all to 1"); + Array2::ones((delays.len_of(Axis(0)), 32)) + } + }; + + // Complain if the dimensions of delays and gains don't match. + if delays.dim().0 != gains.dim().0 { + return Err(BeamError::DelayGainsDimensionMismatch { + delays: delays.dim().0, + gains: gains.dim().0, + }); + } + + // Wrap the `AnalyticBeam` out of hyperbeam with our own `AnalyticBeam`. + let hyperbeam_object = mwa_hyperbeam::analytic::AnalyticBeam::new_custom( + at, + at.get_default_dipole_height(), + 4, + ); + Ok(AnalyticBeam { + hyperbeam_object, + analytic_type: at, + delays, + gains, + ideal_delays, + }) + } + + fn calc_jones_inner( + &self, + azel: AzEl, + freq_hz: f64, + delays: &[u32], + amps: &[f64], + latitude_rad: f64, + ) -> Result, mwa_hyperbeam::analytic::AnalyticBeamError> { + self.hyperbeam_object.calc_jones_pair( + azel.az, + azel.za(), + freq_hz as _, + delays, + amps, + latitude_rad, + true, + ) + } + + fn _calc_jones_array( + &self, + azels: &[AzEl], + freq_hz: f64, + delays: &[u32], + amps: &[f64], + latitude_rad: f64, + ) -> Result>, mwa_hyperbeam::analytic::AnalyticBeamError> { + self.hyperbeam_object.calc_jones_array( + azels, + freq_hz as _, + delays, + amps, + latitude_rad, + true, + ) + } + + fn calc_jones_array_inner( + &self, + azels: &[AzEl], + freq_hz: f64, + delays: &[u32], + amps: &[f64], + latitude_rad: f64, + results: &mut [Jones], + ) -> Result<(), mwa_hyperbeam::analytic::AnalyticBeamError> { + self.hyperbeam_object.calc_jones_array_inner( + azels, + freq_hz as _, + delays, + amps, + latitude_rad, + true, + results, + ) + } +} + +impl Beam for AnalyticBeam { + fn get_beam_type(&self) -> BeamType { + match self.analytic_type { + AnalyticType::MwaPb => BeamType::AnalyticMwaPb, + AnalyticType::Rts => BeamType::AnalyticRts, + } + } + + fn get_num_tiles(&self) -> usize { + self.delays.len_of(Axis(0)) + } + + fn get_ideal_dipole_delays(&self) -> Option<[u32; 16]> { + Some(self.ideal_delays) + } + + fn get_dipole_delays(&self) -> Option>> { + Some(self.delays.to_shared()) + } + + fn get_dipole_gains(&self) -> Option>> { + Some(self.gains.to_shared()) + } + + fn get_beam_file(&self) -> Option<&std::path::Path> { + None + } + + fn calc_jones( + &self, + azel: marlu::AzEl, + freq_hz: f64, + tile_index: Option, + latitude_rad: f64, + ) -> Result, BeamError> { + if let Some(tile_index) = tile_index { + if tile_index > self.delays.len_of(Axis(0)) { + return Err(BeamError::BadTileIndex { + got: tile_index, + max: self.delays.len_of(Axis(0)), + }); + } + let delays = self.delays.slice(s![tile_index, ..]); + let amps = self.gains.slice(s![tile_index, ..]); + let j = self.calc_jones_inner( + azel, + freq_hz, + delays.as_slice().unwrap(), + amps.as_slice().unwrap(), + latitude_rad, + )?; + Ok(j) + } else { + let delays = &self.ideal_delays; + let amps = [1.0; 32]; + let j = self.calc_jones_inner(azel, freq_hz, delays, &s, latitude_rad)?; + Ok(j) + } + } + + fn calc_jones_array( + &self, + azels: &[AzEl], + freq_hz: f64, + tile_index: Option, + latitude_rad: f64, + ) -> Result>, BeamError> { + let mut jones = vec![Jones::default(); azels.len()]; + Beam::calc_jones_array_inner(self, azels, freq_hz, tile_index, latitude_rad, &mut jones)?; + Ok(jones) + } + + fn calc_jones_array_inner( + &self, + azels: &[marlu::AzEl], + freq_hz: f64, + tile_index: Option, + latitude_rad: f64, + results: &mut [marlu::Jones], + ) -> Result<(), BeamError> { + if let Some(tile_index) = tile_index { + if tile_index > self.delays.len_of(Axis(0)) { + return Err(BeamError::BadTileIndex { + got: tile_index, + max: self.delays.len_of(Axis(0)), + }); + } + let delays = self.delays.slice(s![tile_index, ..]); + let amps = self.gains.slice(s![tile_index, ..]); + self.calc_jones_array_inner( + azels, + freq_hz, + delays.as_slice().unwrap(), + amps.as_slice().unwrap(), + latitude_rad, + results, + )?; + } else { + let delays = &self.ideal_delays; + let amps = [1.0; 32]; + self.calc_jones_array_inner(azels, freq_hz, delays, &s, latitude_rad, results)?; + } + Ok(()) + } + + fn find_closest_freq(&self, desired_freq_hz: f64) -> f64 { + desired_freq_hz + } + + fn empty_coeff_cache(&self) {} + + #[cfg(any(feature = "cuda", feature = "hip"))] + fn prepare_gpu_beam(&self, freqs_hz: &[u32]) -> Result, BeamError> { + let gpu_beam = unsafe { + self.hyperbeam_object + .gpu_prepare(self.delays.view(), self.gains.view())? + }; + let freq_map = (0..freqs_hz.len()).map(|i| i as i32).collect::>(); + let d_freq_map = DevicePointer::copy_to_device(&freq_map)?; + Ok(Box::new(AnalyticBeamGpu { + hyperbeam_object: gpu_beam, + d_freqs_hz: DevicePointer::copy_to_device(freqs_hz)?, + d_freq_map, + })) + } +} + +#[cfg(any(feature = "cuda", feature = "hip"))] +struct AnalyticBeamGpu { + hyperbeam_object: mwa_hyperbeam::analytic::AnalyticBeamGpu, + d_freqs_hz: DevicePointer, + d_freq_map: DevicePointer, +} + +#[cfg(any(feature = "cuda", feature = "hip"))] +impl BeamGpu for AnalyticBeamGpu { + unsafe fn calc_jones_pair( + &self, + az_rad: &[GpuFloat], + za_rad: &[GpuFloat], + latitude_rad: f64, + d_jones: *mut std::ffi::c_void, + ) -> Result<(), BeamError> { + let d_az_rad = DevicePointer::copy_to_device(az_rad)?; + let d_za_rad = DevicePointer::copy_to_device(za_rad)?; + self.hyperbeam_object.calc_jones_device_pair_inner( + d_az_rad.get(), + d_za_rad.get(), + az_rad.len().try_into().expect("not bigger than i32::MAX"), + self.d_freqs_hz.get(), + self.d_freqs_hz + .get_num_elements() + .try_into() + .expect("not bigger than i32::MAX"), + latitude_rad as GpuFloat, + true, + d_jones, + )?; + Ok(()) + } + + fn get_beam_type(&self) -> BeamType { + BeamType::FEE + } + + fn get_tile_map(&self) -> *const i32 { + self.hyperbeam_object.get_device_tile_map() + } + + fn get_freq_map(&self) -> *const i32 { + self.d_freq_map.get() + } + + fn get_num_unique_tiles(&self) -> i32 { + self.hyperbeam_object.get_num_unique_tiles() + } + + fn get_num_unique_freqs(&self) -> i32 { + self.d_freqs_hz + .get_num_elements() + .try_into() + .expect("not bigger than i32::MAX") + } +} diff --git a/src/beam/error.rs b/src/beam/error.rs index aeba9a45..6402f712 100644 --- a/src/beam/error.rs +++ b/src/beam/error.rs @@ -35,11 +35,14 @@ pub enum BeamError { #[error("Got tile index {got}, but the biggest tile index is {max}")] BadTileIndex { got: usize, max: usize }, - #[error("hyperbeam error: {0}")] - Hyperbeam(#[from] mwa_hyperbeam::fee::FEEBeamError), + #[error("hyperbeam FEE error: {0}")] + HyperbeamFee(#[from] mwa_hyperbeam::fee::FEEBeamError), - #[error("hyperbeam init error: {0}")] - HyperbeamInit(#[from] mwa_hyperbeam::fee::InitFEEBeamError), + #[error("hyperbeam init FEE error: {0}")] + HyperbeamInitFee(#[from] mwa_hyperbeam::fee::InitFEEBeamError), + + #[error("hyperbeam analytic error: {0}")] + HyperbeamAnalytic(#[from] mwa_hyperbeam::analytic::AnalyticBeamError), #[cfg(any(feature = "cuda", feature = "hip"))] #[error(transparent)] diff --git a/src/beam/fee.rs b/src/beam/fee.rs index c1e525b6..c49ffcff 100644 --- a/src/beam/fee.rs +++ b/src/beam/fee.rs @@ -33,6 +33,7 @@ impl FEEBeam { gains: Option>, file: Option<&Path>, ) -> Result { + // Check that the delays are sensible. validate_delays(&delays, num_tiles)?; let ideal_delays = delays.get_ideal_delays(); @@ -115,7 +116,7 @@ impl FEEBeam { ) } - fn calc_jones_array( + fn _calc_jones_array( &self, azels: &[AzEl], freq_hz: f64, @@ -225,33 +226,8 @@ impl Beam for FEEBeam { tile_index: Option, latitude_rad: f64, ) -> Result>, BeamError> { - // The FEE beam is defined only at specific frequencies. For this - // reason, rather than making a unique hash for every single different - // frequency, round specified frequency (`freq_hz`) to the nearest beam - // frequency and use that for the hash. - let beam_freq = self.find_closest_freq(freq_hz); - - let jones = if let Some(tile_index) = tile_index { - if tile_index > self.delays.len_of(Axis(0)) { - return Err(BeamError::BadTileIndex { - got: tile_index, - max: self.delays.len_of(Axis(0)), - }); - } - let delays = self.delays.slice(s![tile_index, ..]); - let amps = self.gains.slice(s![tile_index, ..]); - self.calc_jones_array( - azels, - beam_freq, - delays.as_slice().unwrap(), - amps.as_slice().unwrap(), - latitude_rad, - )? - } else { - let delays = &self.ideal_delays; - let amps = [1.0; 32]; - self.calc_jones_array(azels, beam_freq, delays, &s, latitude_rad)? - }; + let mut jones = vec![Jones::default(); azels.len()]; + Beam::calc_jones_array_inner(self, azels, freq_hz, tile_index, latitude_rad, &mut jones)?; Ok(jones) } diff --git a/src/beam/mod.rs b/src/beam/mod.rs index ff19c958..ee1bf383 100644 --- a/src/beam/mod.rs +++ b/src/beam/mod.rs @@ -12,11 +12,13 @@ //! implication being that a sky-model source's brightness is always assumed to //! be correct when at zenith. +mod analytic; mod error; mod fee; #[cfg(test)] mod tests; +pub(crate) use analytic::AnalyticBeam; pub(crate) use error::BeamError; pub(crate) use fee::FEEBeam; @@ -50,6 +52,14 @@ pub enum BeamType { #[default] FEE, + /// The mwa_pb flavour of the analytic beam. + #[strum(serialize = "analytic-mwa_pb")] + AnalyticMwaPb, + + /// The RTS flavour of the analytic beam. + #[strum(serialize = "analytic-rts")] + AnalyticRts, + /// a.k.a. [`NoBeam`]. Only returns identity matrices. #[strum(serialize = "none")] None, @@ -387,6 +397,11 @@ impl BeamGpu for NoBeamGpu { } } +/// Create a beam object given the beam type (if not supplied, fall back to a +/// default). This function _should not_ be used generally; setting up beam +/// objects with this function cannot use dipole amps, and the FEE beam will +/// always be set up with the `MWA_BEAM_FILE` environment variable, which the +/// user may not have set. pub fn create_beam_object( beam_type: Option<&str>, num_tiles: usize, @@ -409,7 +424,6 @@ pub fn create_beam_object( BeamType::FEE => { debug!("Setting up a FEE beam object"); - // Check that the delays are sensible. validate_delays(&dipole_delays, num_tiles)?; @@ -421,6 +435,26 @@ pub fn create_beam_object( None, )?)) } + + BeamType::AnalyticMwaPb => { + debug!("Setting up an \"mwa_pb\" analytic beam object"); + validate_delays(&dipole_delays, num_tiles)?; + Ok(Box::new(AnalyticBeam::new_mwa_pb( + num_tiles, + dipole_delays, + None, + )?)) + } + + BeamType::AnalyticRts => { + debug!("Setting up an \"RTS\" analytic beam object"); + validate_delays(&dipole_delays, num_tiles)?; + Ok(Box::new(AnalyticBeam::new_rts( + num_tiles, + dipole_delays, + None, + )?)) + } } } diff --git a/src/cli/beam.rs b/src/cli/beam.rs index db12b288..846c8b35 100644 --- a/src/cli/beam.rs +++ b/src/cli/beam.rs @@ -13,31 +13,20 @@ use clap::Parser; use marlu::AzEl; use num_traits::{Float, FromPrimitive}; -use crate::{ - beam::{create_beam_object, Delays, BEAM_TYPES_COMMA_SEPARATED}, - HyperdriveError, -}; - -lazy_static::lazy_static! { - static ref BEAM_TYPE_HELP: String = format!("The type of beam to use. Supported types: {}", *BEAM_TYPES_COMMA_SEPARATED); -} +use crate::{beam::Delays, HyperdriveError}; /// Generate beam response values. #[derive(Parser, Debug)] pub struct BeamArgs { - #[clap(help = BEAM_TYPE_HELP.as_str())] - beam_type: String, + #[clap(flatten)] + beam_args: super::common::BeamArgs, /// The frequency to use for the beam model [MHz]. #[clap(short, long, default_value = "150")] freq_mhz: f64, - /// If specified, use these dipole delays for the MWA pointing. e.g. 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 - #[clap(short, long, multiple_values(true))] - delays: Option>, - /// The array latitude to use. This only affects the parallactic-angle - /// correction. + /// correction for the FEE beam, but is required for the analytic beam. #[clap(short, long, allow_hyphen_values = true, default_value = "-27.0")] latitude_deg: f64, @@ -102,8 +91,7 @@ fn gen_azzas( fn calc_cpu(args: &BeamArgs) -> Result<(), HyperdriveError> { let BeamArgs { - beam_type, - delays, + beam_args, freq_mhz, latitude_deg, max_za, @@ -113,11 +101,9 @@ fn calc_cpu(args: &BeamArgs) -> Result<(), HyperdriveError> { gpu: _, } = args; - let beam = create_beam_object( - Some(beam_type.as_str()), - 1, - Delays::Partial(delays.clone().unwrap_or(vec![0; 16])), - )?; + let beam = beam_args + .clone() + .parse(1, Some(Delays::Partial(vec![0; 16])), None, None)?; let mut out = BufWriter::new(File::create(output)?); let azels: Vec<_> = gen_azzas(max_za.to_radians(), step.to_radians()) @@ -145,8 +131,7 @@ fn calc_gpu(args: &BeamArgs) -> Result<(), HyperdriveError> { use crate::gpu::{DevicePointer, GpuFloat, GpuJones}; let BeamArgs { - beam_type, - delays, + beam_args, freq_mhz, latitude_deg, max_za, @@ -155,11 +140,9 @@ fn calc_gpu(args: &BeamArgs) -> Result<(), HyperdriveError> { gpu: _, } = args; - let beam = create_beam_object( - Some(beam_type.as_str()), - 1, - Delays::Partial(delays.clone().unwrap_or(vec![0; 16])), - )?; + let beam = beam_args + .clone() + .parse(1, Some(Delays::Partial(vec![0; 16])), None, None)?; let gpu_beam = beam.prepare_gpu_beam(&[(freq_mhz * 1e6) as u32])?; let mut out = BufWriter::new(File::create(output)?); diff --git a/src/cli/common/beam/mod.rs b/src/cli/common/beam/mod.rs index 1acf7444..4dc6d968 100644 --- a/src/cli/common/beam/mod.rs +++ b/src/cli/common/beam/mod.rs @@ -14,7 +14,10 @@ use serde::{Deserialize, Serialize}; use super::{InfoPrinter, Warn}; use crate::{ - beam::{Beam, BeamError, BeamType, Delays, FEEBeam, NoBeam, BEAM_TYPES_COMMA_SEPARATED}, + beam::{ + AnalyticBeam, Beam, BeamError, BeamType, Delays, FEEBeam, NoBeam, + BEAM_TYPES_COMMA_SEPARATED, + }, io::read::VisInputType, }; @@ -242,6 +245,118 @@ impl BeamArgs { }; Box::new(beam) } + + BeamType::AnalyticMwaPb | BeamType::AnalyticRts => { + match beam_type { + BeamType::AnalyticMwaPb => { + debug!("Setting up an mwa_pb-flavoured analytic beam object"); + printer.push_line("Type: Analytic (mwa_pb)".into()); + } + BeamType::AnalyticRts => { + debug!("Setting up an RTS-flavoured analytic beam object"); + printer.push_line("Type: Analytic (RTS)".into()); + } + BeamType::FEE => unreachable!(), + BeamType::None => unreachable!(), + } + + let mut dipole_delays = match user_dipole_delays { + Some(d) => Some(Delays::parse(d)?), + None => data_dipole_delays, + } + .ok_or(BeamError::NoDelays("Analytic"))?; + trace!("Attempting to use delays:"); + match &dipole_delays { + Delays::Full(d) => { + for row in d.outer_iter() { + trace!("{row}"); + } + } + Delays::Partial(d) => trace!("{d:?}"), + } + let dipole_gains = if unity_dipole_gains { + printer.push_line("Assuming all dipoles are \"alive\"".into()); + None + } else { + // If we don't have dipole gains from the input data, then + // we issue a warning that we must assume no dead dipoles. + if dipole_gains.is_none() { + match input_data_type { + Some(VisInputType::MeasurementSet) => [ + "Measurement sets cannot supply dead dipole information.".into(), + "Without a metafits file, we must assume all dipoles are alive.".into(), + "This will make beam Jones matrices inaccurate in sky-model generation." + .into(), + ] + .warn(), + Some(VisInputType::Uvfits) => [ + "uvfits files cannot supply dead dipole information.".into(), + "Without a metafits file, we must assume all dipoles are alive.".into(), + "This will make beam Jones matrices inaccurate in sky-model generation." + .into(), + ] + .warn(), + Some(VisInputType::Raw) => { + unreachable!("Raw data inputs always specify dipole gains") + } + None => (), + } + } + dipole_gains + }; + if let Some(dipole_gains) = dipole_gains.as_ref() { + trace!("Attempting to use dipole gains:"); + for row in dipole_gains.outer_iter() { + trace!("{row}"); + } + + // Currently, the only way to have dipole gains other than + // zero or one is by using Aman's "DipAmps" metafits column. + if dipole_gains.iter().any(|&g| g != 0.0 && g != 1.0) { + printer.push_line( + "Using Aman's 'DipAmps' dipole gains from the metafits".into(), + ); + } else { + let num_tiles_with_dead_dipoles = dipole_gains + .outer_iter() + .filter(|tile_dipole_gains| { + tile_dipole_gains.iter().any(|g| g.abs() < f64::EPSILON) + }) + .count(); + printer.push_line( + format!( + "Using dead dipole information ({num_tiles_with_dead_dipoles} tiles affected)" + ) + .into(), + ); + } + } else { + // If we don't have dipole gains, we must assume all dipoles + // are "alive". But, if any dipole delays are 32, then the + // beam code will still ignore those dipoles. So use ideal + // dipole delays for all tiles. + dipole_delays.set_to_ideal_delays(); + let ideal_delays = dipole_delays.get_ideal_delays(); + + // Warn the user if they wanted unity dipole gains but the + // ideal dipole delays contain 32. + if unity_dipole_gains && ideal_delays.iter().any(|&v| v == 32) { + "Some ideal dipole delays are 32; these dipoles will not have unity gains" + .warn() + } + } + + let beam = match beam_type { + BeamType::AnalyticMwaPb => { + AnalyticBeam::new_mwa_pb(total_num_tiles, dipole_delays, dipole_gains)? + } + BeamType::AnalyticRts => { + AnalyticBeam::new_rts(total_num_tiles, dipole_delays, dipole_gains)? + } + _ => unreachable!("only analytic beams should be here"), + }; + Box::new(beam) + } }; if let Some(d) = beam.get_ideal_dipole_delays() { diff --git a/src/cli/error.rs b/src/cli/error.rs index 3737553d..5091bb75 100644 --- a/src/cli/error.rs +++ b/src/cli/error.rs @@ -468,8 +468,9 @@ impl From for HyperdriveError { | BeamError::DelayGainsDimensionMismatch { .. } => Self::Delays(s), BeamError::Unrecognised(_) | BeamError::BadTileIndex { .. } - | BeamError::Hyperbeam(_) - | BeamError::HyperbeamInit(_) => Self::Beam(s), + | BeamError::HyperbeamFee(_) + | BeamError::HyperbeamInitFee(_) + | BeamError::HyperbeamAnalytic(_) => Self::Beam(s), #[cfg(any(feature = "cuda", feature = "hip"))] BeamError::Gpu(_) => Self::Beam(s), } diff --git a/src/gpu/mod.rs b/src/gpu/mod.rs index 432a65ff..0b75c511 100644 --- a/src/gpu/mod.rs +++ b/src/gpu/mod.rs @@ -214,6 +214,11 @@ impl DevicePointer { self.size } + /// Get the number of elements allocated against the buffer. + pub(crate) fn get_num_elements(&self) -> usize { + self.size / std::mem::size_of::() + } + /// Allocate a number of bytes on the device. #[track_caller] pub(crate) fn malloc(size: usize) -> Result, GpuError> { diff --git a/src/model/cpu.rs b/src/model/cpu.rs index 5660e0d9..6c11303b 100644 --- a/src/model/cpu.rs +++ b/src/model/cpu.rs @@ -119,7 +119,7 @@ impl<'a> SkyModellerCpu<'a> { continue; } - let (gains, delays) = fix_amps_ndarray(gains, delays); + let (gains, delays) = mwa_hyperbeam::fix_amps_ndarray(gains, delays); let mut unique_tile_hasher = DefaultHasher::new(); delays.hash(&mut unique_tile_hasher); @@ -934,31 +934,3 @@ impl<'a> super::SkyModeller<'a> for SkyModellerCpu<'a> { Ok(()) } } - -/// Ensure that any delays of 32 have an amplitude (dipole gain) of 0. The -/// results are bad otherwise! Also ensure that we have 32 dipole gains (amps) -/// here. Also return a Rust array of delays for convenience. -/// -/// TODO: This is copy+pasted from `hyperbeam`; make that function public and -/// use it instead. -fn fix_amps_ndarray(amps: ArrayView1, delays: ArrayView1) -> ([f64; 32], [u32; 16]) { - let mut full_amps: [f64; 32] = [1.0; 32]; - full_amps - .iter_mut() - .zip(amps.iter().cycle()) - .zip(delays.iter().cycle()) - .for_each(|((out_amp, &in_amp), &delay)| { - if delay == 32 { - *out_amp = 0.0; - } else { - *out_amp = in_amp; - } - }); - - // So that we don't have to do .as_slice().unwrap() on our ndarrays outside - // of this function, return a Rust array of delays here. - let mut delays_a: [u32; 16] = [0; 16]; - delays_a.iter_mut().zip(delays).for_each(|(da, d)| *da = *d); - - (full_amps, delays_a) -}