data-beans 0.10.13

//! Train/test and K-fold splitting of a backend, by CELL.
//!
//! # Why this lives in data-beans
//!
//! Splitting is a data operation, not a model one: the folds have to be
//! identical no matter which engine consumes them, or two methods are compared
//! on different data and the comparison means nothing. Putting it beside
//! `subsample` and `subset-columns` gives every downstream crate — `senna`,
//! `pinto`, anything later — one splitter and one seed convention.
//!
//! # The grouped split is the point on spatial data
//!
//! A random cell split leaves nearly every test cell ringed by training cells.
//! On spatial data adjacent cells are near-duplicates, so the model has
//! effectively seen each test cell already and the score comes out optimistic —
//! it measures interpolation between neighbours, not generalization.
//!
//! So the split is by *region*, not by cell. Give it coordinates and a grid
//! resolution — `--coord positions.csv --grid 8` tiles the bounding box 8×8 and
//! assigns whole tiles — and a test region gets a training boundary rather than
//! a training interior. The gap between a random and a tiled split is worth
//! reporting rather than choosing between.
//!
//! `--groups` is the escape hatch for a grouping that is not geometric: donor,
//! sample, slide, or any precomputed labelling.
//!
//! # Ablating features
//!
//! Deliberately not a flag here. Dropping genes from the training half composes
//! out of what already exists:
//!
//! ```text
//! data-beans split data.zarr -o cv --test-frac 0.2
//! data-beans subsample cv.train.zarr.zip --gene-frac 0.8 -o cv.train.ablated
//! ```
//!
//! Keeping the two operations separate means the split stays reproducible from
//! its own seed regardless of what the ablation does.

use crate::hdf5_io::*;
use crate::sparse_io::*;
use crate::zarr_io::{finalize_output, prepare_output};

use clap::Args;
use log::info;
use rand::rngs::SmallRng;
use rand::seq::SliceRandom;
use rand::SeedableRng;
use rustc_hash::FxHashMap as HashMap;

use crate::handlers::explore::{off_screen, PlanArgs};
use crate::interactive::ui::Tui;

#[path = "split_plan.rs"]
mod plan;

#[derive(Args, Debug)]
pub struct SplitArgs {
    /// data file -- `.zarr`, `.zarr.zip`, or `.h5`
    pub data_file: Box<str>,

    /// fraction of cells (columns) held out as the test half, in (0, 1)
    #[arg(long, conflicts_with = "folds")]
    pub test_frac: Option<f64>,

    /// number of cross-validation folds; writes {out}.fold{k}.{train,test}
    #[arg(long, conflicts_with = "test_frac")]
    pub folds: Option<usize>,

    /// cell coordinates -- split by REGION instead of by cell (CSV/TSV/parquet)
    ///
    /// First column is the cell name; the coordinate columns are chosen with
    /// --coord-columns. Pair with --grid.
    #[arg(long, value_name = "FILE", conflicts_with = "groups")]
    pub coord: Option<Box<str>>,

    /// 0-based coordinate column indices, comma separated [default: 1,2]
    #[arg(long, value_name = "I,J", value_delimiter = ',', requires = "coord")]
    pub coord_columns: Option<Vec<usize>>,

    /// tiles per axis over the coordinate bounding box; 8 gives an 8x8 grid
    #[arg(long, default_value_t = 8, requires = "coord")]
    pub grid: usize,

    /// two-column table (cell name, group label), or one label per line in
    /// column order (qc's {out}.batch.gz) -- keep cells sharing a label on
    /// the same side: donor, sample, slide. Use --coord for geometry instead.
    #[arg(long, value_name = "TABLE")]
    pub groups: Option<Box<str>>,

    /// RNG seed for a reproducible split
    #[arg(long, default_value_t = 42)]
    pub seed: u64,

    /// backend for the output files
    #[arg(long, value_enum, default_value = "zarr")]
    pub backend: SparseIoBackend,

    /// output prefix: writes {output}.train.* and {output}.test.*
    #[arg(short, long)]
    pub output: Box<str>,

    /// keep `.zarr` directories instead of producing `.zarr.zip` archives
    #[arg(long = "no-zip", default_value_t = true, action = clap::ArgAction::SetFalse)]
    pub zip: bool,

    #[command(flatten)]
    pub view: PlanArgs,
}

/// How the cells are split: what `--test-frac`, `--folds`, `--grid` and
/// `--seed` say, or what the plan ended on.
#[derive(Debug, Clone, Copy, PartialEq)]
struct Choice {
    test_frac: Option<f64>,
    folds: Option<usize>,
    grid: usize,
    seed: u64,
}

/// Split a backend into train and test halves by cell, or into K folds.
///
/// Every cell lands in exactly one side (or exactly one fold's test half), so
/// the halves are disjoint and together cover the input — checked, not assumed.
pub fn run_split(args: &SplitArgs) -> anyhow::Result<()> {
    let session = args.view.session("split")?;
    let (backend_in, file_in) = resolve_backend_file(&args.data_file, None)?;
    let data = open_sparse_matrix(&file_in, &backend_in)?;
    let n_columns = data
        .num_columns()
        .ok_or_else(|| anyhow::anyhow!("backend has no `ncol`"))?;
    let column_names = data.column_names()?;
    let mut choice = Choice {
        test_frac: args.test_frac,
        folds: args.folds,
        grid: args.grid,
        seed: args.seed,
    };
    let mut tui = session.as_ref().map(Tui::new);
    if let Some(tui) = tui.as_mut() {
        match plan::choose(tui, args, &column_names, choice)? {
            Some(c) => choice = c,
            None => {
                info!("split: cancelled; nothing written");
                return Ok(());
            }
        }
    }
    let splits = split_columns(args, &column_names, choice)?;
    off_screen(tui.as_mut(), || {
        write_splits(args, &*data, n_columns, &splits, choice)
    })
}

/// The cells of each train and test half: their columns, by fold.
type Splits = Vec<(Vec<usize>, Vec<usize>)>;

/// The units shuffled and assigned whole, under `choice`: one column each
/// by default, one grid tile under `--coord`, one label under `--groups`.
/// Everything after is the same in all three cases, which is why grouped
/// splitting costs nothing extra to support.
fn column_groups(
    args: &SplitArgs,
    column_names: &[Box<str>],
    grid: usize,
) -> anyhow::Result<Vec<Vec<usize>>> {
    Ok(match (args.coord.as_deref(), args.groups.as_deref()) {
        (Some(path), _) => {
            let coord_columns = args.coord_columns.clone().unwrap_or_else(|| vec![1, 2]);
            column_groups_from_grid(path, grid, &coord_columns, column_names)?
        }
        (_, Some(path)) => column_groups_from_table(path, column_names)?,
        _ => (0..column_names.len()).map(|i| vec![i]).collect(),
    })
}

/// The halves `choice` makes of the columns of `column_names`.
fn split_columns(
    args: &SplitArgs,
    column_names: &[Box<str>],
    choice: Choice,
) -> anyhow::Result<Splits> {
    let column_groups = column_groups(args, column_names, choice.grid)?;
    // Two, not one: a single group cannot be divided, and `--test-frac` would
    // reach `clamp(1, 0)` and panic. One group is a real outcome of a coarse
    // `--grid` or a `--groups` table with a single label, so it gets a message
    // that names the cause rather than a backtrace.
    anyhow::ensure!(
        column_groups.len() >= 2,
        "only {} group(s) to split: every cell falls in the same tile or group, \
         so there is no way to hold any of them out. Raise --grid, or check that \
         --groups / --coord really vary across cells.",
        column_groups.len()
    );

    // Group ids first, then the columns they cover: the partition never sees a
    // column, so a grouped split cannot leak one across the boundary by
    // construction rather than by a later check.
    Ok(partition_groups(
        column_groups.len(),
        choice.test_frac,
        choice.folds,
        choice.seed,
    )?
    .into_iter()
    .map(|(train_groups, test_groups)| {
        (
            columns_in_groups(&column_groups, &train_groups),
            columns_in_groups(&column_groups, &test_groups),
        )
    })
    .collect())
}

/// `{output}.train` and `{output}.test`, or `{output}.fold{k}.*` for folds.
fn half_name(args: &SplitArgs, choice: Choice, fold: usize, half: &str) -> String {
    match choice.folds {
        Some(_) => format!("{}.fold{fold}.{half}", args.output),
        None => format!("{}.{half}", args.output),
    }
}

fn write_splits(
    args: &SplitArgs,
    data: &dyn SparseIo<IndexIter = Vec<usize>>,
    n_columns: usize,
    splits: &Splits,
    choice: Choice,
) -> anyhow::Result<()> {
    for (fold, (train_columns, test_columns)) in splits.iter().enumerate() {
        anyhow::ensure!(
            train_columns.len() + test_columns.len() == n_columns,
            "split lost cells: {} train + {} test != {n_columns}",
            train_columns.len(),
            test_columns.len()
        );
        let tag = |half: &str| half_name(args, choice, fold, half);
        write_half(data, train_columns, &tag("train"), args)?;
        write_half(data, test_columns, &tag("test"), args)?;
        info!(
            "fold {fold}: {} train cells, {} test cells",
            train_columns.len(),
            test_columns.len()
        );
    }
    Ok(())
}

/// Flatten group ids to the column indices they cover, sorted so the backend
/// read stays sequential.
fn columns_in_groups(column_groups: &[Vec<usize>], picked: &[usize]) -> Vec<usize> {
    let mut columns: Vec<usize> = picked
        .iter()
        .flat_map(|&g| column_groups[g].iter().copied())
        .collect();
    columns.sort_unstable();
    columns
}

/// Cells grouped into a regular grid over their coordinates.
///
/// A grid rather than a clustering on purpose: the tiles are then a property of
/// the tissue's geometry alone, reproducible from the coordinate file without a
/// model run, and a reader can say exactly how large a held-out region was.
/// Empty tiles simply do not appear.
fn column_groups_from_grid(
    path: &str,
    tiles_per_axis: usize,
    coord_columns: &[usize],
    column_names: &[Box<str>],
) -> anyhow::Result<Vec<Vec<usize>>> {
    anyhow::ensure!(tiles_per_axis >= 2, "--grid must be at least 2");
    anyhow::ensure!(
        coord_columns.len() == 2,
        "--coord-columns needs exactly two indices; got {}",
        coord_columns.len()
    );
    let coords = read_coords(path, coord_columns)?;

    let mut coord_of_column: Vec<(f64, f64)> = Vec::with_capacity(column_names.len());
    let mut n_missing = 0usize;
    for name in column_names {
        match coords.get(name.as_ref()) {
            Some(&xy) => coord_of_column.push(xy),
            None => {
                n_missing += 1;
                coord_of_column.push((f64::NAN, f64::NAN));
            }
        }
    }
    anyhow::ensure!(
        n_missing == 0,
        "{path}: {n_missing} of {} cells have no coordinate",
        column_names.len()
    );

    let (mut x0, mut y0, mut x1, mut y1) = (f64::MAX, f64::MAX, f64::MIN, f64::MIN);
    for &(x, y) in &coord_of_column {
        x0 = x0.min(x);
        y0 = y0.min(y);
        x1 = x1.max(x);
        y1 = y1.max(y);
    }
    // A degenerate span would put every cell in tile 0; widen it so the grid is
    // still well defined rather than silently collapsing to one group.
    let span = |lo: f64, hi: f64| if hi > lo { hi - lo } else { 1.0 };
    let (width, height) = (span(x0, x1), span(y0, y1));

    let mut group_of_tile: HashMap<usize, usize> = HashMap::default();
    let mut column_groups: Vec<Vec<usize>> = Vec::new();
    for (column, &(x, y)) in coord_of_column.iter().enumerate() {
        let tx = (((x - x0) / width) * tiles_per_axis as f64)
            .floor()
            .clamp(0.0, (tiles_per_axis - 1) as f64) as usize;
        let ty = (((y - y0) / height) * tiles_per_axis as f64)
            .floor()
            .clamp(0.0, (tiles_per_axis - 1) as f64) as usize;
        let tile = ty * tiles_per_axis + tx;
        let next = column_groups.len();
        let group = *group_of_tile.entry(tile).or_insert(next);
        if group == next {
            column_groups.push(Vec::new());
        }
        column_groups[group].push(column);
    }
    info!(
        "spatial split: {} cells in {} non-empty tiles of a \
         {tiles_per_axis}x{tiles_per_axis} grid ({path})",
        column_names.len(),
        column_groups.len()
    );
    Ok(column_groups)
}

/// `cell -> (x, y)` from a CSV/TSV/parquet whose first column is the cell name.
fn read_coords(path: &str, cols: &[usize]) -> anyhow::Result<HashMap<Box<str>, (f64, f64)>> {
    let (i, j) = (cols[0], cols[1]);
    let mut out: HashMap<Box<str>, (f64, f64)> = HashMap::default();
    if path.ends_with(".parquet") {
        let names = legume_numeric::matrix::parquet::read_parquet_string_column(path, 0)?;
        let mat =
            <nalgebra::DMatrix<f32> as legume_numeric::matrix::traits::IoOps>::from_parquet(path)?;
        for (r, name) in names.iter().enumerate() {
            out.insert(
                name.clone(),
                (
                    f64::from(mat.mat[(r, i - 1)]),
                    f64::from(mat.mat[(r, j - 1)]),
                ),
            );
        }
        return Ok(out);
    }
    let text = std::fs::read_to_string(path)
        .map_err(|e| anyhow::anyhow!("reading coordinates {path}: {e}"))?;
    for (lineno, line) in text.lines().enumerate() {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        let f: Vec<&str> = if line.contains('\t') {
            line.split('\t').collect()
        } else if line.contains(',') {
            line.split(',').collect()
        } else {
            line.split_whitespace().collect()
        };
        let max = i.max(j);
        anyhow::ensure!(
            f.len() > max,
            "{path}:{}: {} columns, need index {max}",
            lineno + 1,
            f.len()
        );
        // A header line has non-numeric coordinates; skip it rather than failing.
        let (Ok(x), Ok(y)) = (f[i].trim().parse::<f64>(), f[j].trim().parse::<f64>()) else {
            continue;
        };
        out.insert(f[0].trim().into(), (x, y));
    }
    anyhow::ensure!(!out.is_empty(), "{path}: no coordinate rows parsed");
    Ok(out)
}

/// Cells grouped by their group label, from a group table (see
/// [`super::group_table`]).
fn column_groups_from_table(
    path: &str,
    column_names: &[Box<str>],
) -> anyhow::Result<Vec<Vec<usize>>> {
    let labels = super::group_table::column_labels(path, column_names)?;
    let (_, column_groups, _) = super::group_table::groups_of(&labels);
    info!(
        "grouped split: {} cells in {} groups from {path}",
        column_names.len(),
        column_groups.len()
    );
    Ok(column_groups)
}

/// Write one half as a new backend, streamed.
///
/// The selected columns' exact nnz comes off the resident indptr and the
/// survivors flow through the validated streaming writer in byte-budgeted
/// slabs, so peak residency is bounded by the slab budget rather than scaling
/// with the half — the wall the old triplet-materialising version hit at
/// imaging scale is gone, and so is the doc comment that described it.
fn write_half(
    data: &dyn SparseIo<IndexIter = Vec<usize>>,
    selected_columns: &[usize],
    output: &str,
    args: &SplitArgs,
) -> anyhow::Result<()> {
    let row_names = data.row_names()?;
    let all_column_names = data.column_names()?;
    let kept_column_names: Vec<Box<str>> = selected_columns
        .iter()
        .map(|&c| all_column_names[c].clone())
        .collect();

    let (effective_output, backend_out, file_out) =
        prepare_output(output, args.backend.clone(), args.zip)?;
    info!("Streaming {} columns to {output}", selected_columns.len());
    let (n_rows, n_cols, nnz) = crate::column_subset::stream_column_selection(
        data,
        selected_columns,
        None,
        &row_names,
        &kept_column_names,
        file_out.as_ref(),
        &backend_out,
    )?;

    let final_path = finalize_output(&file_out, &effective_output)?;
    info!("wrote {final_path} ({n_rows} genes x {n_cols} cells, {nnz} non-zeros)");
    Ok(())
}

/// Partition `n_groups` shuffled groups into `(train, test)` group-index pairs.
///
/// Separated from the column bookkeeping so the part that decides who is held
/// out can be tested without a backend on disk — this is the piece where an
/// off-by-one silently produces an overlapping or incomplete fold, which no
/// downstream number would reveal.
fn partition_groups(
    n_groups: usize,
    test_frac: Option<f64>,
    folds: Option<usize>,
    seed: u64,
) -> anyhow::Result<Vec<(Vec<usize>, Vec<usize>)>> {
    let mut shuffled_groups: Vec<usize> = (0..n_groups).collect();
    let mut rng = SmallRng::seed_from_u64(seed);
    shuffled_groups.shuffle(&mut rng);

    match (test_frac, folds) {
        (Some(f), _) => {
            anyhow::ensure!(
                f > 0.0 && f < 1.0,
                "--test-frac must be in (0, 1); got {f}. A percentage goes in as e.g. 0.2"
            );
            // Clamped so neither side is ever empty: a fraction that rounds to
            // zero, or to everything, is a silently useless split rather than an
            // error the user would see.
            let n_test = ((f * n_groups as f64).round() as usize).clamp(1, n_groups - 1);
            let (test_groups, train_groups) = shuffled_groups.split_at(n_test);
            Ok(vec![(train_groups.to_vec(), test_groups.to_vec())])
        }
        (_, Some(k)) => {
            anyhow::ensure!(k >= 2, "--folds must be at least 2");
            anyhow::ensure!(
                k <= n_groups,
                "--folds {k} exceeds the {n_groups} groups available to split"
            );
            Ok((0..k)
                .map(|f| {
                    let take = |in_test: bool| -> Vec<usize> {
                        shuffled_groups
                            .iter()
                            .enumerate()
                            .filter(|(i, _)| (i % k == f) == in_test)
                            .map(|(_, &g)| g)
                            .collect()
                    };
                    (take(false), take(true))
                })
                .collect())
        }
        (None, None) => anyhow::bail!("specify one of --test-frac or --folds"),
    }
}

#[cfg(all(test, feature = "ndarray"))]
mod tests;