data-beans 0.11.127

//! Sampler primitives for paired ATAC + RNA simulation.
//!
//! Ported verbatim from `chickpea/src/simulation/sample.rs` (the original
//! `sim-link` simulator) with `Mat` substituted for `nalgebra::DMatrix<f32>`
//! so this module pulls in no chickpea-internal types.

use legume_numeric::matrix::traits::MatOps;
use nalgebra::DMatrix;
use rand::prelude::*;
use rand_distr::Poisson;
use rayon::prelude::*;

pub type Mat = DMatrix<f32>;

/// Sample dictionary β[D × K] via column-wise softmax of random logits, so
/// each topic (column) is a probability distribution over features and
/// `Σ_d β[d, k] = 1`. With column-stochastic θ this gives the rate model
/// `λ = depth · β · θ` an expected library size of `depth` per cell.
///
/// The seeded N(0,1) fill stays scalar (preserves seed reproducibility);
/// the softmax delegates to `normalize_exp_logits_columns_inplace`
/// (numerically-stable log-sum-exp from `legume_numeric::matrix`).
pub fn sample_dictionary(d: usize, k: usize, rng: &mut impl Rng) -> Mat {
    let normal = rand_distr::Normal::new(0.0f32, 1.0).unwrap();
    let buf: Vec<f32> = (0..d * k).map(|_| normal.sample(rng)).collect();
    let mut beta = Mat::from_vec(d, k, buf);
    beta.normalize_exp_logits_columns_inplace();
    beta
}

/// Generate per-cell depth multipliers `depth · exp(σ·N(0,1) − σ²/2)`.
/// The `−σ²/2` centering makes `E[depth_j] = depth` so the supplied `depth`
/// is the expected library size, not biased by the log-normal mean drift.
pub fn sample_cell_depths(
    n: usize,
    depth: usize,
    cell_sd_log_depth: f32,
    rng: &mut impl Rng,
) -> Vec<f32> {
    let sigma = cell_sd_log_depth;
    let center = 0.5 * sigma * sigma;
    let normal = rand_distr::Normal::new(0.0f32, sigma).unwrap();
    let depth_f = depth as f32;
    (0..n)
        .map(|_| depth_f * (normal.sample(rng) - center).exp())
        .collect()
}

/// Sample nested topic proportions [K_total, N] with hierarchical assignments.
///
/// Each cell gets a coarse topic (PVE-controlled), then a subtype within it.
/// The dominant coarse topic `k*` is mixed with a uniform background,
/// `θ_coarse[k*] = pve_coarse + (1 − pve_coarse)/K` and `(1 − pve_coarse)/K`
/// elsewhere, and likewise for the dominant subtype within `k*` with
/// `pve_sub` and `K_sub`, so every column of both matrices sums to 1.
/// Returns (θ_full[K_total, N], θ_coarse[K, N]) where K_total = K × K_sub.
/// θ_coarse is the marginalization: θ_coarse[k,n] = Σ_s θ_full[k*K_sub+s, n].
///
/// When K_sub = 1, θ_full == θ_coarse.
pub fn sample_nested_topic_proportions(
    k: usize,
    k_sub: usize,
    n: usize,
    pve_coarse: f32,
    pve_sub: f32,
    rseed: u64,
) -> (Mat, Mat) {
    let k_total = k * k_sub;
    let pve_c = pve_coarse.clamp(0.0, 1.0);
    let pve_s = pve_sub.clamp(0.0, 1.0);
    let bg_coarse = (1.0 - pve_c) / k.max(1) as f32;
    let bg_sub = (1.0 - pve_s) / k_sub.max(1) as f32;

    let cells: Vec<(Vec<f32>, Vec<f32>)> = (0..n)
        .into_par_iter()
        .map(|j| {
            let mut rng = StdRng::seed_from_u64(rseed.wrapping_add(j as u64));
            let coarse_dist = rand_distr::Uniform::new(0, k).unwrap();
            let sub_dist = rand_distr::Uniform::new(0, k_sub).unwrap();

            let dom_k = coarse_dist.sample(&mut rng);
            let dom_s = sub_dist.sample(&mut rng);

            let mut full_col = vec![0.0f32; k_total];
            let mut coarse_col = vec![0.0f32; k];

            for kk in 0..k {
                let coarse_weight = if kk == dom_k {
                    pve_c + bg_coarse
                } else {
                    bg_coarse
                };
                coarse_col[kk] = coarse_weight;

                for s in 0..k_sub {
                    let sub_weight = if kk == dom_k {
                        if s == dom_s {
                            pve_s + bg_sub
                        } else {
                            bg_sub
                        }
                    } else {
                        1.0 / k_sub as f32
                    };
                    full_col[kk * k_sub + s] = coarse_weight * sub_weight;
                }
            }

            (full_col, coarse_col)
        })
        .collect();

    let mut theta_full = Mat::zeros(k_total, n);
    let mut theta_coarse = Mat::zeros(k, n);
    for (j, (full_col, coarse_col)) in cells.into_iter().enumerate() {
        for (i, v) in full_col.into_iter().enumerate() {
            theta_full[(i, j)] = v;
        }
        for (i, v) in coarse_col.into_iter().enumerate() {
            theta_coarse[(i, j)] = v;
        }
    }

    (theta_full, theta_coarse)
}

/// Marginalize β_ext[P, K×K_sub] → β_atac[P, K] by averaging over subtypes
/// (uniform `P(s | k) = 1/K_sub`). With column-stochastic β_ext, β_atac's
/// columns also sum to 1.
///
/// Vectorized via nalgebra column-views — each output column is a sum of
/// `K_sub` contiguous source columns, scaled by `1/K_sub`.
pub fn marginalize_dictionary(beta_ext: &Mat, k: usize, k_sub: usize) -> Mat {
    let p = beta_ext.nrows();
    let inv_k_sub = 1.0_f32 / k_sub.max(1) as f32;
    let mut beta = Mat::zeros(p, k);
    for kk in 0..k {
        let mut acc = beta.column_mut(kk);
        for s in 0..k_sub {
            acc += beta_ext.column(kk * k_sub + s);
        }
        acc.scale_mut(inv_k_sub);
    }
    beta
}

/// Sample gene-topic effects γ[G × K] from LogNormal(0, σ²).
pub fn sample_gene_topic_effects(
    g: usize,
    k: usize,
    gene_topic_sd: f32,
    rng: &mut impl Rng,
) -> Mat {
    let normal = rand_distr::Normal::new(0.0f32, gene_topic_sd).unwrap();
    let buf: Vec<f32> = (0..g * k).map(|_| normal.sample(rng).exp()).collect();
    Mat::from_vec(g, k, buf)
}

/// Sample sparse indicator matrix M[G × P] with nearby causal peaks.
///
/// Causal peaks are selected from the same chromosome as the gene,
/// within a local window. This ensures distance-based cis-windows
/// can recover the true links.
///
/// Layout: gene i → chr (i % n_chr) + 1, position ~ i / n_chr.
///         peak j → chr (j % n_chr) + 1, position ~ j / n_chr.
///
/// Returns (gene_indices, peak_indices) where each pair represents
/// one entry M[gene, peak] = 1.
pub fn sample_indicator_matrix(
    n_genes: usize,
    n_peaks: usize,
    n_linked_genes: usize,
    n_causal_per_gene: usize,
    n_chromosomes: usize,
    rng: &mut impl Rng,
) -> (Vec<usize>, Vec<usize>) {
    let mut genes = Vec::new();
    let mut peaks = Vec::new();

    let mut chr_peaks: Vec<Vec<usize>> = vec![Vec::new(); n_chromosomes];
    for p in 0..n_peaks {
        chr_peaks[p % n_chromosomes].push(p);
    }

    let mut gene_indices: Vec<usize> = (0..n_genes).collect();
    gene_indices.shuffle(rng);
    let linked_genes = &gene_indices[..n_linked_genes.min(n_genes)];

    for &gi in linked_genes {
        let chr = gi % n_chromosomes;
        let peaks_on_chr = &chr_peaks[chr];
        if peaks_on_chr.is_empty() {
            continue;
        }

        let gene_pos = gi / n_chromosomes;

        let center = gene_pos.min(peaks_on_chr.len() - 1);

        let half_window = (n_causal_per_gene * 3).max(10).min(peaks_on_chr.len() / 2);
        let lo = center.saturating_sub(half_window);
        let hi = (center + half_window + 1).min(peaks_on_chr.len());
        let window_size = hi - lo;

        let n_causal = n_causal_per_gene.min(window_size);
        let sampled = rand::seq::index::sample(rng, window_size, n_causal);
        for idx in sampled.into_iter() {
            genes.push(gi);
            peaks.push(peaks_on_chr[lo + idx]);
        }
    }

    (genes, peaks)
}

/// Build derived RNA dictionary W[G × K] = M × β.
pub fn build_derived_dictionary(
    indicator_genes: &[usize],
    indicator_peaks: &[usize],
    beta_dk: &Mat,
    n_genes: usize,
) -> Mat {
    let k = beta_dk.ncols();
    let mut w = Mat::zeros(n_genes, k);

    for (&gi, &pi) in indicator_genes.iter().zip(indicator_peaks.iter()) {
        for kk in 0..k {
            w[(gi, kk)] += beta_dk[(pi, kk)];
        }
    }

    w
}

/// Sample Poisson counts from a per-(feature, cell) log-rate `logits_dn`.
///
/// Each cell follows a multinomial-Poisson: softmax the cell's column over
/// features (numerically stable), scale by the cell's library size `depth_j`,
/// and draw `Poisson(depth_j · p_fj)` per feature. The softmax fixes the
/// expected library at `depth_j`, so the variance budget encoded in `logits`
/// (see `combine_budget_logits`) governs only the relative per-feature rates.
pub fn sample_poisson_from_logits(
    logits_dn: &Mat,
    depths: &[f32],
    rseed: u64,
) -> Vec<(u64, u64, f32)> {
    let d = logits_dn.nrows();
    depths
        .par_iter()
        .enumerate()
        .flat_map_iter(|(j, &depth_j)| {
            let mut rng = StdRng::seed_from_u64(rseed.wrapping_add(j as u64));
            let col = logits_dn.column(j);
            let maxv = col.max();
            let denom: f32 = (0..d)
                .map(|f| (col[f] - maxv).exp())
                .sum::<f32>()
                .max(1e-30);
            let mut out = Vec::new();
            for f in 0..d {
                let rate = depth_j * (col[f] - maxv).exp() / denom;
                if rate > 1e-6 {
                    let count = sample_poisson_safe(rate, &mut rng);
                    if count > 0.0 {
                        out.push((f as u64, j as u64, count));
                    }
                }
            }
            out
        })
        .collect()
}

/// One `Poisson(rate)` draw; 0 for a rate the distribution rejects
/// (non-positive or non-finite).
fn sample_poisson_safe(rate: f32, rng: &mut impl Rng) -> f32 {
    Poisson::new(rate as f64)
        .map(|d| d.sample(rng) as f32)
        .unwrap_or(0.0)
}

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

    /// Both θ matrices are column-stochastic, θ_full marginalizes to
    /// θ_coarse, and the dominant topic carries `pve + (1 − pve)/K`.
    #[test]
    fn nested_proportions_are_proportions() {
        for &(k, k_sub, pve_c, pve_s) in &[
            (10usize, 1usize, 0.8f32, 1.0f32),
            (8, 1, 0.5, 1.0),
            (4, 2, 0.6, 0.5),
            (3, 3, 0.0, 0.0),
            (1, 1, 0.3, 0.3),
        ] {
            let n = 25;
            let (full, coarse) = sample_nested_topic_proportions(k, k_sub, n, pve_c, pve_s, 5);
            for j in 0..n {
                let sc: f32 = coarse.column(j).sum();
                let sf: f32 = full.column(j).sum();
                assert!((sc - 1.0).abs() < 1e-5, "K={k}: coarse column sums to {sc}");
                assert!((sf - 1.0).abs() < 1e-5, "K={k}: full column sums to {sf}");
                for kk in 0..k {
                    let block: f32 = (0..k_sub).map(|s| full[(kk * k_sub + s, j)]).sum();
                    assert!((block - coarse[(kk, j)]).abs() < 1e-5);
                }
                let top = coarse.column(j).max();
                let want = pve_c + (1.0 - pve_c) / k as f32;
                assert!((top - want).abs() < 1e-5, "dominant {top} vs {want}");
            }
        }
    }

    /// A high rate is still a random draw: it varies between seeds and its
    /// spread is Poisson (variance ≈ mean), not a rounded constant.
    #[test]
    fn high_rates_keep_poisson_noise() {
        let rate = 2000.0_f32;
        let mut rng = StdRng::seed_from_u64(1);
        let draws: Vec<f64> = (0..2000)
            .map(|_| sample_poisson_safe(rate, &mut rng) as f64)
            .collect();
        let mean = draws.iter().sum::<f64>() / draws.len() as f64;
        let var = draws.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / draws.len() as f64;
        assert!((mean - rate as f64).abs() < 10.0, "mean {mean}");
        assert!((0.7..1.3).contains(&(var / rate as f64)), "variance {var}");
    }
}