sparse_ir_core/

freq.rs

//! Matsubara frequency implementation for SparseIR
//!
//! This module provides Matsubara frequency types for both fermionic and
//! bosonic statistics, matching the C++ implementation in freq.hpp.

use num_complex::Complex64;
use std::cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd};
use std::fmt;
use std::ops::{Add, Neg, Sub};

use crate::error::Error;
use crate::traits::{Bosonic, Fermionic, Statistics, StatisticsType};

/// Matsubara frequency for a specific statistics type
///
/// This represents the Matsubara frequency ν = nπ/β, where:
/// - n is the stored integer (the *reduced* Matsubara frequency, returned by
///   [`n`](Self::n)): odd for fermionic statistics (n = ±1, ±3, …) and even
///   for bosonic statistics (n = 0, ±2, ±4, …)
/// - β is the inverse temperature, passed to [`value`](Self::value) and
///   [`value_imaginary`](Self::value_imaginary)
///
/// In terms of the conventional Matsubara index m of ν = (2m + ζ)π/β, with
/// ζ = 1 for fermionic and ζ = 0 for bosonic statistics, the stored integer is
/// n = 2m + ζ; it is *not* m itself. [`new`](Self::new) rejects an n whose
/// parity does not match the statistics.
///
/// The statistics type S is checked at compile time to ensure type safety.
///
/// # Examples
/// ```
/// use sparse_ir::{BosonicFreq, FermionicFreq};
/// use std::f64::consts::PI;
///
/// let beta = 10.0;
///
/// // n = 1 is the lowest positive fermionic frequency, ω = π/β (k = 0).
/// let w = FermionicFreq::new(1).unwrap();
/// assert_eq!(w.n(), 1);
/// assert!((w.value(beta) - PI / beta).abs() < 1e-15);
///
/// // n = 2 is the lowest positive bosonic frequency, ω = 2π/β (k = 1).
/// let nu = BosonicFreq::new(2).unwrap();
/// assert!((nu.value(beta) - 2.0 * PI / beta).abs() < 1e-15);
///
/// // The parity of n must match the statistics.
/// assert!(FermionicFreq::new(2).is_err());
/// assert!(BosonicFreq::new(1).is_err());
/// ```
///
/// Adding or subtracting frequencies gives the statistics of the result:
/// `FermionicFreq ± FermionicFreq` is a [`BosonicFreq`], and a fermionic and
/// a bosonic frequency (in either order) give a [`FermionicFreq`].
#[derive(Debug, Clone, Copy)]
pub struct MatsubaraFreq<S: StatisticsType> {
    n: i64,
    _phantom: std::marker::PhantomData<S>,
}

// Type aliases for convenience
pub type FermionicFreq = MatsubaraFreq<Fermionic>;
pub type BosonicFreq = MatsubaraFreq<Bosonic>;

impl<S: StatisticsType> MatsubaraFreq<S> {
    /// Get the reduced Matsubara frequency n, where ν = nπ/β
    ///
    /// n is odd for fermionic and even for bosonic statistics.
    pub fn n(&self) -> i64 {
        self.n
    }

    /// Create a new Matsubara frequency
    ///
    /// # Arguments
    /// * `n` - The reduced Matsubara frequency n of ν = nπ/β: odd for
    ///   fermionic statistics, even for bosonic statistics
    ///
    /// # Errors
    /// [`Error::InvalidMatsubaraIndex`] if `n` is even for fermionic or odd
    /// for bosonic statistics
    ///
    /// # Examples
    /// ```
    /// use sparse_ir::freq::{FermionicFreq, BosonicFreq};
    ///
    /// let fermionic = FermionicFreq::new(1).unwrap();  // OK: odd n for fermionic
    /// let bosonic = BosonicFreq::new(0).unwrap();      // OK: even n for bosonic
    /// ```
    pub fn new(n: i64) -> Result<Self, Error> {
        // Check if the frequency is allowed for this statistics type
        let allowed = match S::STATISTICS {
            Statistics::Fermionic => n % 2 != 0, // Fermionic: odd n only
            Statistics::Bosonic => n % 2 == 0,   // Bosonic: even n only
        };

        if !allowed {
            return Err(Error::InvalidMatsubaraIndex {
                n,
                statistics: S::STATISTICS,
            });
        }

        Ok(Self {
            n,
            _phantom: std::marker::PhantomData,
        })
    }

    /// Create a Matsubara frequency without validation
    ///
    /// # Safety
    /// This function bypasses the validation in `new()`. Only use when you're
    /// certain the frequency is valid for the statistics type.
    pub unsafe fn new_unchecked(n: i64) -> Self {
        Self {
            n,
            _phantom: std::marker::PhantomData,
        }
    }

    /// Get the reduced Matsubara frequency n, where ν = nπ/β (same as [`n`](Self::n))
    pub fn get_n(&self) -> i64 {
        self.n
    }

    /// Compute the real frequency value: n * π / β
    ///
    /// # Arguments
    /// * `beta` - Inverse temperature
    ///
    /// # Returns
    /// The real frequency value
    pub fn value(&self, beta: f64) -> f64 {
        self.n as f64 * std::f64::consts::PI / beta
    }

    /// Compute the imaginary frequency value: i * n * π / β
    ///
    /// # Arguments
    /// * `beta` - Inverse temperature
    ///
    /// # Returns
    /// The imaginary frequency value as a complex number
    pub fn value_imaginary(&self, beta: f64) -> Complex64 {
        Complex64::new(0.0, self.value(beta))
    }

    /// Get the statistics type
    pub fn statistics(&self) -> Statistics {
        S::STATISTICS
    }

    /// Convert to i64 (for compatibility with C++ operator long long())
    pub fn into_i64(self) -> i64 {
        self.n
    }
}

// Default implementations
impl Default for FermionicFreq {
    fn default() -> Self {
        // Default fermionic frequency is n=1 (smallest positive odd frequency)
        unsafe { Self::new_unchecked(1) }
    }
}

impl Default for BosonicFreq {
    fn default() -> Self {
        // Default bosonic frequency is n=0 (zero frequency)
        unsafe { Self::new_unchecked(0) }
    }
}

// Conversion to i64
impl<S: StatisticsType> From<MatsubaraFreq<S>> for i64 {
    fn from(freq: MatsubaraFreq<S>) -> Self {
        freq.n
    }
}

// Operator overloading: addition and subtraction
//
// The parity of n + m and n - m follows from the parities of n and m, so the
// output type carries the statistics of the result: fermionic ± fermionic is
// bosonic, fermionic ± bosonic (either order) is fermionic, bosonic ±
// bosonic is bosonic. The i64 arithmetic is the usual one (a panic on
// overflow in debug builds, wrapping in release builds); wrapping keeps the
// parity, so the output is always a valid frequency.
macro_rules! impl_freq_arithmetic {
    ($lhs:ty, $rhs:ty => $out:ty) => {
        impl Add<$rhs> for $lhs {
            type Output = $out;

            fn add(self, other: $rhs) -> $out {
                // SAFETY: not a memory-safety requirement; the parity of the
                // sum is that of $out (see above).
                unsafe { <$out>::new_unchecked(self.n + other.n) }
            }
        }

        impl Sub<$rhs> for $lhs {
            type Output = $out;

            fn sub(self, other: $rhs) -> $out {
                // SAFETY: as for `add`
                unsafe { <$out>::new_unchecked(self.n - other.n) }
            }
        }
    };
}

impl_freq_arithmetic!(FermionicFreq, FermionicFreq => BosonicFreq);
impl_freq_arithmetic!(FermionicFreq, BosonicFreq => FermionicFreq);
impl_freq_arithmetic!(BosonicFreq, FermionicFreq => FermionicFreq);
impl_freq_arithmetic!(BosonicFreq, BosonicFreq => BosonicFreq);

// Operator overloading: Negation
impl<S: StatisticsType> Neg for MatsubaraFreq<S> {
    type Output = Self;

    fn neg(self) -> Self {
        // Negation preserves parity, so this is always valid
        unsafe { Self::new_unchecked(-self.n) }
    }
}

// Comparison operators
impl<S: StatisticsType> PartialEq for MatsubaraFreq<S> {
    fn eq(&self, other: &Self) -> bool {
        self.n == other.n
    }
}

impl<S: StatisticsType> Eq for MatsubaraFreq<S> {}

impl<S: StatisticsType> PartialOrd for MatsubaraFreq<S> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<S: StatisticsType> Ord for MatsubaraFreq<S> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.n.cmp(&other.n)
    }
}

// Hash implementation (needed for some collections)
impl<S: StatisticsType> std::hash::Hash for MatsubaraFreq<S> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.n.hash(state);
    }
}

// Display implementation
impl<S: StatisticsType> fmt::Display for MatsubaraFreq<S> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.n {
            0 => write!(f, "0"),
            1 => write!(f, "π/β"),
            -1 => write!(f, "-π/β"),
            n => write!(f, "{}π/β", n),
        }
    }
}

// Utility functions
/// Get the sign of a Matsubara frequency
#[allow(dead_code)]
pub(crate) fn sign<S: StatisticsType>(freq: &MatsubaraFreq<S>) -> i32 {
    if freq.n > 0 {
        1
    } else if freq.n < 0 {
        -1
    } else {
        0
    }
}

/// Get the fermionic sign based on statistics type
#[allow(dead_code)]
pub(crate) fn fermionic_sign<S: StatisticsType>() -> i32 {
    match S::STATISTICS {
        Statistics::Fermionic => -1,
        Statistics::Bosonic => 1,
    }
}

/// Create a zero frequency (bosonic only)
#[allow(dead_code)]
pub(crate) fn zero() -> BosonicFreq {
    unsafe { BosonicFreq::new_unchecked(0) }
}

/// Check if a frequency is zero
#[allow(dead_code)]
#[doc(hidden)]
pub fn is_zero<S: StatisticsType>(freq: &MatsubaraFreq<S>) -> bool {
    match S::STATISTICS {
        Statistics::Fermionic => false, // Fermionic frequencies are never zero
        Statistics::Bosonic => freq.n == 0,
    }
}

/// Compare two Matsubara frequencies of potentially different statistics types
#[allow(dead_code)]
pub(crate) fn is_less<S1: StatisticsType, S2: StatisticsType>(
    a: &MatsubaraFreq<S1>,
    b: &MatsubaraFreq<S2>,
) -> bool {
    a.get_n() < b.get_n()
}

/// Factory function to create Statistics from zeta value
#[allow(dead_code)]
pub(crate) fn create_statistics(zeta: i64) -> Result<Statistics, String> {
    match zeta {
        1 => Ok(Statistics::Fermionic),
        0 => Ok(Statistics::Bosonic),
        _ => Err(format!("Unknown statistics type: zeta={}", zeta)),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_fermionic_frequency_creation() {
        // Valid fermionic frequencies (odd n)
        let freq1 = FermionicFreq::new(1).unwrap();
        assert_eq!(freq1.get_n(), 1);

        let freq3 = FermionicFreq::new(3).unwrap();
        assert_eq!(freq3.get_n(), 3);

        let freq_neg1 = FermionicFreq::new(-1).unwrap();
        assert_eq!(freq_neg1.get_n(), -1);

        // Invalid fermionic frequencies (even n)
        assert!(FermionicFreq::new(0).is_err());
        assert!(FermionicFreq::new(2).is_err());
        assert!(FermionicFreq::new(-2).is_err());
    }

    #[test]
    fn test_bosonic_frequency_creation() {
        // Valid bosonic frequencies (even n)
        let freq0 = BosonicFreq::new(0).unwrap();
        assert_eq!(freq0.get_n(), 0);

        let freq2 = BosonicFreq::new(2).unwrap();
        assert_eq!(freq2.get_n(), 2);

        let freq_neg2 = BosonicFreq::new(-2).unwrap();
        assert_eq!(freq_neg2.get_n(), -2);

        // Invalid bosonic frequencies (odd n)
        assert!(BosonicFreq::new(1).is_err());
        assert!(BosonicFreq::new(3).is_err());
        assert!(BosonicFreq::new(-1).is_err());
    }

    #[test]
    fn test_frequency_values() {
        let beta = 2.0;
        let pi = std::f64::consts::PI;

        let fermionic = FermionicFreq::new(1).unwrap();
        assert!((fermionic.value(beta) - pi / 2.0).abs() < 1e-14);

        let bosonic = BosonicFreq::new(0).unwrap();
        assert_eq!(bosonic.value(beta), 0.0);

        let bosonic2 = BosonicFreq::new(2).unwrap();
        assert!((bosonic2.value(beta) - pi).abs() < 1e-14);
    }

    #[test]
    fn test_imaginary_values() {
        let beta = 2.0;
        let pi = std::f64::consts::PI;

        let fermionic = FermionicFreq::new(1).unwrap();
        let imag = fermionic.value_imaginary(beta);
        assert!((imag.re - 0.0).abs() < 1e-14);
        assert!((imag.im - pi / 2.0).abs() < 1e-14);
    }

    #[test]
    fn test_operator_overloading() {
        // Bosonic: even + even = even (valid)
        let bfreq2 = BosonicFreq::new(2).unwrap();
        let bfreq4 = BosonicFreq::new(4).unwrap();

        let sum = bfreq2 + bfreq4;
        assert_eq!(sum.get_n(), 6);

        let diff = bfreq4 - bfreq2;
        assert_eq!(diff.get_n(), 2);

        // Negation preserves parity (valid for both statistics)
        let freq1 = FermionicFreq::new(1).unwrap();
        let neg = -freq1;
        assert_eq!(neg.get_n(), -1);

        let neg_b = -bfreq2;
        assert_eq!(neg_b.get_n(), -2);
    }

    /// Sums and differences have the statistics of their parity: fermionic ±
    /// fermionic is bosonic, fermionic ± bosonic is fermionic. Before the
    /// change the output had the type of the inputs, so FermionicFreq(1) +
    /// FermionicFreq(1) was a FermionicFreq with the even n = 2 in release
    /// builds (a debug_assert only).
    #[test]
    fn test_sum_and_difference_follow_the_statistics() {
        let f1 = FermionicFreq::new(1).unwrap();
        let f3 = FermionicFreq::new(3).unwrap();
        let b2 = BosonicFreq::new(2).unwrap();

        let sum: BosonicFreq = f1 + f3;
        assert_eq!(sum, BosonicFreq::new(4).unwrap());
        let diff: BosonicFreq = f1 - f3;
        assert_eq!(diff, BosonicFreq::new(-2).unwrap());

        let fb_sum: FermionicFreq = f1 + b2;
        assert_eq!(fb_sum, FermionicFreq::new(3).unwrap());
        let fb_diff: FermionicFreq = f3 - b2;
        assert_eq!(fb_diff, FermionicFreq::new(1).unwrap());
        let bf_sum: FermionicFreq = b2 + f1;
        assert_eq!(bf_sum, FermionicFreq::new(3).unwrap());
        let bf_diff: FermionicFreq = b2 - f3;
        assert_eq!(bf_diff, FermionicFreq::new(-1).unwrap());

        let bb: BosonicFreq = b2 + b2;
        assert_eq!(bb, BosonicFreq::new(4).unwrap());
    }

    #[test]
    fn test_comparison_operators() {
        let freq1 = FermionicFreq::new(1).unwrap();
        let freq3 = FermionicFreq::new(3).unwrap();
        let freq1_copy = FermionicFreq::new(1).unwrap();

        assert_eq!(freq1, freq1_copy);
        assert_ne!(freq1, freq3);
        assert!(freq1 < freq3);
        assert!(freq3 > freq1);
        assert!(freq1 <= freq1_copy);
        assert!(freq1 >= freq1_copy);
    }

    #[test]
    fn test_utility_functions() {
        let fermionic = FermionicFreq::new(1).unwrap();
        let bosonic = BosonicFreq::new(0).unwrap();
        let bosonic2 = BosonicFreq::new(2).unwrap();

        // Sign function
        assert_eq!(sign(&fermionic), 1);
        assert_eq!(sign(&bosonic), 0);
        assert_eq!(sign(&-fermionic), -1);

        // Fermionic sign
        assert_eq!(fermionic_sign::<Fermionic>(), -1);
        assert_eq!(fermionic_sign::<Bosonic>(), 1);

        // Zero frequency
        let zero_freq = zero();
        assert_eq!(zero_freq.get_n(), 0);

        // Is zero
        assert!(!is_zero(&fermionic));
        assert!(is_zero(&bosonic));
        assert!(!is_zero(&bosonic2));

        // Is less
        assert!(is_less(&fermionic, &bosonic2));
        assert!(!is_less(&bosonic2, &fermionic));
    }

    #[test]
    fn test_statistics_creation() {
        assert_eq!(create_statistics(1).unwrap(), Statistics::Fermionic);
        assert_eq!(create_statistics(0).unwrap(), Statistics::Bosonic);
        assert!(create_statistics(2).is_err());
    }

    #[test]
    fn test_display() {
        let freq1 = FermionicFreq::new(1).unwrap();
        let freq_neg1 = FermionicFreq::new(-1).unwrap();
        let freq0 = BosonicFreq::new(0).unwrap();
        let freq2 = BosonicFreq::new(2).unwrap();

        assert_eq!(format!("{}", freq1), "π/β");
        assert_eq!(format!("{}", freq_neg1), "-π/β");
        assert_eq!(format!("{}", freq0), "0");
        assert_eq!(format!("{}", freq2), "2π/β");
    }

    #[test]
    fn test_default_implementations() {
        let default_fermionic = FermionicFreq::default();
        assert_eq!(default_fermionic.get_n(), 1);

        let default_bosonic = BosonicFreq::default();
        assert_eq!(default_bosonic.get_n(), 0);
    }

    #[test]
    fn test_conversion_to_i64() {
        let freq = FermionicFreq::new(3).unwrap();
        let n: i64 = freq.into();
        assert_eq!(n, 3);

        let n_direct = freq.into_i64();
        assert_eq!(n_direct, 3);
    }

    #[test]
    fn test_sign() {
        let fermionic_pos = FermionicFreq::new(3).unwrap();
        assert_eq!(sign(&fermionic_pos), 1);

        let fermionic_neg = FermionicFreq::new(-5).unwrap();
        assert_eq!(sign(&fermionic_neg), -1);

        let bosonic_zero = BosonicFreq::new(0).unwrap();
        assert_eq!(sign(&bosonic_zero), 0);

        let bosonic_pos = BosonicFreq::new(4).unwrap();
        assert_eq!(sign(&bosonic_pos), 1);
    }

    #[test]
    fn test_fermionic_sign() {
        assert_eq!(fermionic_sign::<Fermionic>(), -1);
        assert_eq!(fermionic_sign::<Bosonic>(), 1);
    }

    #[test]
    fn test_zero() {
        let zero_freq = zero();
        assert_eq!(zero_freq.get_n(), 0);
        assert_eq!(zero_freq.value(1.0), 0.0);
    }

    #[test]
    fn test_is_zero() {
        // Fermionic frequencies are never zero
        let fermionic = FermionicFreq::new(1).unwrap();
        assert!(!is_zero(&fermionic));

        // Bosonic zero frequency
        let bosonic_zero = BosonicFreq::new(0).unwrap();
        assert!(is_zero(&bosonic_zero));

        // Non-zero bosonic frequency
        let bosonic_nonzero = BosonicFreq::new(2).unwrap();
        assert!(!is_zero(&bosonic_nonzero));
    }

    #[test]
    fn test_is_less() {
        let freq1 = FermionicFreq::new(1).unwrap();
        let freq3 = FermionicFreq::new(3).unwrap();
        assert!(is_less(&freq1, &freq3));
        assert!(!is_less(&freq3, &freq1));

        // Compare different statistics types
        let fermionic = FermionicFreq::new(1).unwrap();
        let bosonic = BosonicFreq::new(2).unwrap();
        assert!(is_less(&fermionic, &bosonic));

        // Same frequency
        assert!(!is_less(&freq1, &freq1));
    }

    #[test]
    fn test_create_statistics() {
        // Fermionic (zeta = 1)
        let fermionic = create_statistics(1).unwrap();
        assert_eq!(fermionic, Statistics::Fermionic);

        // Bosonic (zeta = 0)
        let bosonic = create_statistics(0).unwrap();
        assert_eq!(bosonic, Statistics::Bosonic);

        // Invalid zeta
        assert!(create_statistics(2).is_err());
        assert!(create_statistics(-1).is_err());
    }

    /// The parity check holds for negative n, whose remainder is negative.
    #[test]
    fn test_new_rejects_wrong_parity_with_typed_error() {
        use crate::error::Error;

        assert_eq!(
            FermionicFreq::new(2).unwrap_err(),
            Error::InvalidMatsubaraIndex {
                n: 2,
                statistics: Statistics::Fermionic,
            }
        );
        assert_eq!(
            BosonicFreq::new(-3).unwrap_err(),
            Error::InvalidMatsubaraIndex {
                n: -3,
                statistics: Statistics::Bosonic,
            }
        );
        assert_eq!(
            BosonicFreq::new(-3).unwrap_err().to_string(),
            "Matsubara frequency n = -3 is not allowed for bosonic statistics"
        );
        assert_eq!(FermionicFreq::new(-1).unwrap().n(), -1);
        assert_eq!(BosonicFreq::new(-2).unwrap().n(), -2);
    }
}