flare 0.2.0

/// Generate a frequency grid for a given time array `t`, and optional minimum and maximum frequencies.
///
/// # Arguments
///
/// * `t` - Time array (slice of f64)
/// * `fmin` - Optional minimum frequency (f64)
/// * `fmax` - Optional maximum frequency (f64)
/// * `oversample` - Oversampling factor (usize), typically between 1 and 10
///
/// # Returns
///
/// * `Vec<f64>` - A vector of frequencies for the periodogram.
///
///
/// # Examples
///
/// ```
/// use rand::Rng;
/// use flare::period::freq_grid;
///
/// // randomize a list of 1000 data points over 6 years
/// let n_points = 1000;
/// let mut rng = rand::rng();
/// let mut t: Vec<f64> = Vec::with_capacity(n_points);
/// for _ in 0..n_points {
///     // generate a random time between 0 and 6*365 days
///     let random_time = rng.random_range(0.0..(6.0 * 365.0));
///     t.push(random_time);
/// }
/// // Sort the time array to ensure it's in ascending order
/// t.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
///
/// // Now generate the frequency grid, from 10 days to 5 minutes
/// let f_min = 1.0 / (10.0 * 24.0 * 60.0);
/// let f_max = 1.0 / (5.0 / 60.0);
/// let fgrid = freq_grid(&t, Some(f_min), Some(f_max), 3); // Generate frequency grid
/// assert!(!fgrid.is_empty(), "Frequency grid should not be empty");
/// assert!(fgrid.len() > 70000, "Frequency grid should have at least 7000 values for the given parameters");
/// ```
pub fn freq_grid(t: &[f64], fmin: Option<f64>, fmax: Option<f64>, oversample: usize) -> Vec<f64> {
    let trange =
        t.iter().cloned().fold(f64::NAN, f64::max) - t.iter().cloned().fold(f64::NAN, f64::min);
    let texp = t.windows(2).map(|w| w[1] - w[0]).fold(f64::NAN, f64::min);
    let fres = 1.0 / trange / oversample as f64;

    let fmax = fmax.unwrap_or(0.5 / texp);
    let fmin = fmin.unwrap_or(fres);

    let mut fgrid = Vec::new();
    let mut f = fmin;
    while f < fmax {
        fgrid.push(f);
        f += fres;
    }
    fgrid
}

#[inline(always)]
fn fpw_statistic(times: &[f64], ivar: &[f64], ivar_y: &[f64], f: f64, n_bins: usize) -> f64 {
    let mut vt_cinv_v = [0.0; 20];
    let mut yt_cinv_v = [0.0; 20];

    let scale = f * n_bins as f64;

    // Accumulate inverse variance and ivar * y products for each bin
    for i in 0..times.len() {
        let index = ((times[i] * scale) as usize) % n_bins;
        vt_cinv_v[index] += ivar[i];
        yt_cinv_v[index] += ivar_y[i];
    }

    // Compute the final statistic
    let mut sum = 0.0;
    for i in 0..n_bins {
        sum += yt_cinv_v[i].powi(2) / vt_cinv_v[i];
    }
    sum * 0.5
}

/// Calculate the FPW statistic (https://arxiv.org/abs/2502.00243) for a given time series data.
///
/// # Arguments
///
/// * `t` - Time array (slice of f64)
/// * `y` - Measurement array (slice of f64)
/// * `dy` - Measurement errors (slice of f64), must be the same length as `y`
/// * `freqs` - Frequency grid (slice of f64), typically generated by `freq_grid`
/// * `n_bins` - Number of bins to use for the FPW statistic (must be <= 20)
///
/// # Returns
///
/// * `Vec<f64>` - A vector of FPW statistics for each frequency in the provided frequency grid.
///
/// # Examples
///
/// ```
/// use rand::Rng;
/// use flare::period::{fpw, freq_grid};
///
/// // Simulate some time series data
/// let n_points = 1000;
/// let period = 0.25; // in days
/// let mut rng = rand::rng();
/// let mut t: Vec<f64> = Vec::with_capacity(n_points);
/// for _ in 0..n_points {
///     // generate a random time between 0 and 6*365 days
///     let random_time = rng.random_range(0.0..(6.0 * 365.0));
///     t.push(random_time);
/// }
/// // Sort the time array to ensure it's in ascending order
/// t.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
///
/// // Now generate the sinusoidal light curve with some noise
/// let y: Vec<f64> = t
///     .iter()
///     .map(|&time| {
///         // Simulate a sinusoidal light curve with some noise
///         let true_value = (2.0 * std::f64::consts::PI * time / period).sin();
///         let noise = rng.random_range(-0.1..0.1); // Add some noise
///         true_value + noise
///     })
///     .collect();
///
/// // Define frequency grid, with a min freq of 10 days
/// let f_min = 1.0 / (10.0 * 24.0 * 60.0);
/// let f_max = 1.0 / (5.0 / 60.0);
/// let freqs = freq_grid(&t, Some(f_min), Some(f_max), 3); // Generate frequency grid
///
/// // the FPW algorithm requires a number of bins, typically between 5 and 20
/// let n_bins = 10; // choose a number of bins between 5 and 20
///
/// // Calculate FPW statistic for the given time series data
/// let fpw_stats = fpw(&t, &y, &vec![0.1; n_points], &freqs, n_bins);
///
/// // Get the index of the best frequency and its corresponding statistic
/// let (best_freq, best_stat) = {
///     let max_index = fpw_stats
///        .iter()
///         .enumerate()
///         .filter(|&(_, &x)| !x.is_nan())
///         .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
///         .map(|(i, _)| i)
///         .unwrap_or(0);
///
///     (freqs[max_index], fpw_stats[max_index])
/// };
///
/// let period = 1.0 / best_freq * 24.0; // convert from frequency in days to a period in hours
///
/// println!("Best period: {:.2} hours, statistic: {:.2}", period, best_stat);
///
/// assert!((period - 6.0).abs() < 1e-4, "The best period should be close to the true period of 6 hours");
/// ```
pub fn fpw(t: &[f64], y: &[f64], dy: &[f64], freqs: &[f64], n_bins: usize) -> Vec<f64> {
    if n_bins > 20 {
        panic!("n_bins must be less than or equal to 20");
    }

    let mut ivar = vec![0.0; t.len()];
    let mut ivar_y = vec![0.0; t.len()];
    for i in 0..t.len() {
        ivar[i] = 1.0 / (dy[i] * dy[i]);
        ivar_y[i] = ivar[i] * y[i];
    }

    // Calculate FPW statistic for each frequency
    freqs
        .iter()
        .map(|&freq| fpw_statistic(t, &ivar, &ivar_y, freq, n_bins))
        .collect()
}

/// Get the best frequency and its corresponding statistic from the FPW results.
///
/// # Arguments
///
/// * `freqs` - Frequency grid (slice of f64)
/// * `result` - FPW statistic results (slice of f64)
///
/// # Returns
///
/// * `(f64, f64)` - A tuple containing the best frequency and its corresponding statistic.
///
/// # Examples
///
/// ```
/// use flare::period::get_best_freq;
/// let freqs = vec![0.1, 0.2, 0.3, 0.4, 0.5]; // Example frequencies
/// let result = vec![0.5, 1.2, 0.8, 2.5, 1.0]; // Example FPW statistics
/// let (best_freq, best_stat) = get_best_freq(&freqs, &result);
/// assert_eq!(best_freq, 0.4);
/// assert_eq!(best_stat, 2.5);
/// ```
pub fn get_best_freq(freqs: &[f64], result: &[f64]) -> (f64, f64) {
    let max_index = result
        .iter()
        .enumerate()
        .filter(|&(_, &x)| !x.is_nan())
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
        .map(|(i, _)| i)
        .unwrap_or(0);

    (freqs[max_index], result[max_index])
}

/// Get the best frequencies and their statistics, returning the top `n` frequencies based on the FPW statistic.
///
/// # Arguments
///
/// * `freqs` - Frequency grid (slice of f64)
/// * `result` - FPW statistic results (slice of f64)
/// * `n` - Number of top frequencies to return
///
/// # Returns
///
/// * `Vec<(f64, f64)>` - A vector of tuples containing the top `n` frequencies and their corresponding statistics.
///
/// # Examples
///
/// ```
/// use flare::period::get_best_freqs;
/// let freqs = vec![0.1, 0.2, 0.3, 0.4, 0.5]; // Example frequencies
/// let result = vec![0.5, 1.2, 0.8, 2.5, 1.0]; // Example FPW statistics
/// let best_freqs = get_best_freqs(&freqs, &result, 3);
/// assert_eq!(best_freqs.len(), 3);
/// assert_eq!(best_freqs[0], (0.4, 2.5));
/// assert_eq!(best_freqs[1], (0.2, 1.2));
/// assert_eq!(best_freqs[2], (0.5, 1.0));
pub fn get_best_freqs(freqs: &[f64], result: &[f64], n: usize) -> Vec<(f64, f64)> {
    let mut freq_stat_pairs: Vec<(f64, f64)> = freqs
        .iter()
        .zip(result.iter())
        .filter(|&(_, &x)| !x.is_nan())
        .map(|(&f, &s)| (f, s))
        .collect();

    // Sort by statistic in descending order
    freq_stat_pairs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

    // Return the top n frequencies
    freq_stat_pairs.into_iter().take(n).collect()
}