Function compute_seurat_hvg

src/memory/processing/hvg/mod.rs:493–572

Compute highly variable genes using the Seurat method. The Seurat method models the mean-variance relationship by: 1. Calculating dispersion (variance/mean) for each gene 2. Binning genes by expression level 3. Computing expected dispersion within each bin 4. Normalizing dispersions to z-scores 5. Selecting genes with high normalized dispersion ## Method Details - Uses log1p transformation of me

(
    adata: &IMAnnData,
    x: &IMArrayElement,
    params: HVGParams,
)

/// - `dispersions_norm`: Z-score normalized dispersions
/// - `highly_variable`: Boolean selection mask
fn compute_seurat_hvg(
    adata: &IMAnnData,
    x: &IMArrayElement,
    params: HVGParams,
) -> anyhow::Result<()> {
    let n_obs = adata.n_obs();

    // Calculate means from raw counts
    let raw_means: Vec<f64> = x
        .sum_whole(&Direction::COLUMN)?
        .iter()
        .map(|sum: &f64| sum / n_obs as f64)
        .collect();

    let variances: Vec<f64> = x.variance_whole::<u32, f64>(&Direction::COLUMN)?;

    // Calculate dispersions with proper handling of zero means
    let dispersions: Vec<f64> = raw_means
        .iter()
        .zip(variances.iter())
        .map(|(&mean, &var)| {
            let safe_mean = if mean > 1e-12 { mean } else { 1e-12 };
            var / safe_mean
        })
        .collect();

    // For Seurat flavor, use log1p of means for binning and storage
    // This matches what Python does AFTER reverting log normalization
    let log1p_means: Vec<f64> = raw_means.iter().map(|&x| (x + 1.0).ln()).collect();

    // Log dispersions with NaN for zero dispersions (matching Python)
    let log_dispersions: Vec<f64> = dispersions
        .iter()
        .map(|&x| {
            if x > 0.0 {
                x.ln()
            } else {
                f64::NAN // Python sets dispersion[dispersion == 0] = np.nan
            }
        })
        .collect();

    let n_bins = params.n_bins;

    // Use equal-width binning on log1p_means (like Python's pd.cut)
    let (bin_indices, _) = equal_width_binning(&log1p_means, n_bins)?;

    // Calculate mean and std for each bin
    let (mut bin_means, mut bin_stds) =
        calculate_bin_stats(&log_dispersions, &bin_indices, n_bins)?;

    // Handle single-gene bins (like Python's _postprocess_dispersions_seurat)
    postprocess_seurat_dispersions(&mut bin_means, &mut bin_stds)?;

    // Normalize dispersions
    let normalized_dispersions =
        normalize_dispersions(&log_dispersions, &bin_indices, &bin_means, &bin_stds)?;