mwa_hyperbeam 0.5.0

// 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/.

//! A generic interface for MWA primary beam code.

pub(crate) mod error;

pub(crate) use error::BeamError;

use marlu::{ndarray, AzEl, Jones};
use ndarray::Array2;

cfg_if::cfg_if! {
    if #[cfg(feature = "cuda")] {
        use marlu::cuda::*;
        use super::CudaFloat;
    }
}

/// Supported beam types.
#[allow(clippy::upper_case_acronyms)]
pub enum BeamType {
    /// Fully-embedded element beam.
    FEE,

    /// Analytic beam (mwa_pb implementation).
    Analytic,

    /// Analytic beam (RTS implementation).
    AnalyticRTS,

    /// a.k.a. `NoBeam`. Only returns identity matrices.
    None,
}

/// A trait abstracting beam code functions.
pub trait Beam: Sync + Send {
    /// Get the type of beam.
    fn get_beam_type(&self) -> BeamType;

    /// Get the number of tiles associated with this beam. This is determined by
    /// how many delays have been provided.
    fn get_num_tiles(&self) -> usize;

    /// Calculate the beam-response Jones matrix for an [AzEl] direction. The
    /// pointing information is not needed because it was provided when `self`
    /// was created.
    fn calc_jones(
        &self,
        azel: AzEl,
        freq_hz: f64,
        tile_index: usize,
    ) -> Result<Jones<f64>, BeamError>;

    /// Calculate the beam-response Jones matrices for multiple [AzEl]
    /// directions. The pointing information is not needed because it was
    /// provided when `self` was created.
    fn calc_jones_array(
        &self,
        azels: &[AzEl],
        freq_hz: f64,
        tile_index: usize,
    ) -> Result<Vec<Jones<f64>>, BeamError>;

    /// Given a frequency in Hz, find the closest frequency that the beam code
    /// is defined for. An example of when this is important is with the FEE
    /// beam code, which can only give beam responses at specific frequencies.
    /// On the other hand, the analytic beam can be used at any frequency.
    fn find_closest_freq(&self, desired_freq_hz: f64) -> f64;

    /// If this [Beam] supports it, empty the coefficient cache.
    fn empty_coeff_cache(&self);

    #[cfg(feature = "cuda")]
    /// Using the tile information from this [Beam] and frequencies to be used,
    /// return a [BeamCUDA]. This object only needs pointings to calculate beam
    /// response [Jones] matrices.
    ///
    /// # Safety
    ///
    /// This function interfaces directly with the CUDA API. Rust errors attempt
    /// to catch problems but there are no guarantees.
    unsafe fn prepare_cuda_beam(&self, freqs_hz: &[u32]) -> Result<Box<dyn BeamCUDA>, BeamError>;
}

/// A trait abstracting beam code functions on a CUDA-capable device.
#[cfg(feature = "cuda")]
pub trait BeamCUDA {
    /// Calculate the Jones matrices for an [AzEl] direction. The pointing
    /// information is not needed because it was provided when `self` was
    /// created.
    ///
    /// # Safety
    ///
    /// This function interfaces directly with the CUDA API. Rust errors attempt
    /// to catch problems but there are no guarantees.
    unsafe fn calc_jones(
        &self,
        azels: &[AzEl],
    ) -> Result<DevicePointer<Jones<CudaFloat>>, BeamError>;

    /// Get the type of beam.
    fn get_beam_type(&self) -> BeamType;

    /// Get a pointer to the device tile map. This is necessary to access
    /// de-duplicated beam Jones matrices on the device.
    fn get_tile_map(&self) -> *const i32;

    /// Get a pointer to the device freq map. This is necessary to access
    /// de-duplicated beam Jones matrices on the device.
    fn get_freq_map(&self) -> *const i32;

    /// Get the number of de-duplicated frequencies associated with this
    /// [BeamCUDA].
    fn get_num_unique_freqs(&self) -> i32;
}

/// An enum to track whether MWA dipole delays are provided and/or necessary.
#[derive(Debug)]
pub enum Delays {
    /// Delays are fully specified.
    Full(Array2<u32>),

    /// Delays are specified for a single tile. If this can't be refined, then
    /// we must assume that these dipoles apply to all tiles.
    Partial(Vec<u32>),

    /// Delays have not been provided, but are necessary for calibration.
    None,

    /// Delays are not necessary, probably because no beam code is being used.
    NotNecessary,
}