gwseq-io 0.2.0

//! Matrix metadata and the two-dimensional locus.

use indexmap::IndexMap;

use crate::bytes::LeCursor;
use crate::error::{Error, Result};
use crate::genomic::{ChrEntry, ChrMap};
use crate::source::ByteSource;

use super::header::HiCIndexItem;

/// One side of a two-dimensional window, resolved against the bin grid.
#[derive(Debug, Clone)]
pub struct Side {
    pub chr: ChrEntry,
    pub start: i64,
    pub end: i64,
    pub binned_start: i64,
    pub binned_end: i64,
    pub bin_start: i64,
    pub bin_end: i64,
}

/// A request's two windows.
#[derive(Debug, Clone)]
pub struct Loc2D {
    pub x: Side,
    pub y: Side,
    pub bin_size: i64,
    /// The request named its chromosomes the other way round, so the axes were
    /// swapped to read them — a hic file stores only one of the two — and the
    /// result is transposed back before it is handed over.
    pub reversed: bool,
}

impl Loc2D {
    #[inline]
    pub fn is_intra(&self) -> bool {
        self.x.chr.index == self.y.chr.index
    }

    /// Distance from the diagonal in base pairs.
    ///
    /// On one chromosome that is simply `|y - x|`. Across two it is the distance
    /// to the line the window's own corners define, which is what makes
    /// `min_distance`/`max_distance` mean something on an inter-chromosomal
    /// matrix at all.
    pub fn distance_from_diagonal(&self, x: i64, y: i64) -> i64 {
        if self.is_intra() {
            return (y - x).abs();
        }
        let x_span = self.x.binned_end - self.x.binned_start;
        let y_span = self.y.binned_end - self.y.binned_start;
        if x_span <= 0 || y_span <= 0 {
            return 0;
        }
        let a = y_span as f64 / x_span as f64;
        let b = self.y.binned_start as f64 - a * self.x.binned_start as f64;
        let vertical = (y as f64 - (a * x as f64 + b)).abs();
        let horizontal = (x as f64 - (y as f64 - b) / a).abs();
        vertical.min(horizontal).round() as i64
    }
}

/// Turn a pair of genomic intervals into the two-dimensional window a hic
/// matrix is read through.
///
/// `bin_count` cannot be honoured exactly: a hic file holds a fixed set of
/// resolutions, so the nearest one is chosen and the window ends up with however
/// many bins that gives — unlike the bbi readers, where a locus is rescaled to
/// the bin count asked for. `exact_bin_count` is what resizes it afterwards.
// Eight: the map and the file's resolutions, then the request's ids, starts,
// ends and three binning parameters. A struct would hide which of the three a
// caller actually set, and that is what this branches on.
#[allow(clippy::too_many_arguments)]
pub fn parse_loc2d(
    map: &ChrMap,
    available_bin_sizes: &[i64],
    chr_ids: &[String],
    starts: &[i64],
    ends: &[i64],
    bin_size: Option<i64>,
    bin_count: Option<i64>,
    full_bin: bool,
) -> Result<Loc2D> {
    // Every path below reads a resolution out of this, and a file carrying none
    // for the unit asked for — which is every file without fragment-delimited
    // maps, for unit "frag" — would otherwise be indexed into empty.
    if available_bin_sizes.is_empty() {
        return Err(Error::invalid("file has no resolution for this unit"));
    }
    if bin_count == Some(0) {
        return Err(Error::invalid(
            "bin count must be positive, or negative to disregard it",
        ));
    }
    // A negative bin size means "the finest the file has". Zero means nothing,
    // and would reach the binning arithmetic below as a division by it.
    if bin_size == Some(0) {
        return Err(Error::invalid(
            "bin size must be positive, or negative to use the finest available",
        ));
    }

    let pair = |values: &[i64], what: &str| -> Result<(i64, i64)> {
        match values {
            [only] => Ok((*only, *only)),
            [a, b] => Ok((*a, *b)),
            _ => Err(Error::invalid(format!("1 or 2 {what} must be specified"))),
        }
    };
    let ids = match chr_ids {
        [only] => (only.clone(), only.clone()),
        [a, b] => (a.clone(), b.clone()),
        _ => return Err(Error::invalid("1 or 2 chromosomes must be specified")),
    };
    let (x_start, y_start) = pair(starts, "start positions")?;
    let (x_end, y_end) = pair(ends, "end positions")?;

    let make = |id: &str, start: i64, end: i64| -> Result<Side> {
        let chr = map.resolve(id)?.clone();
        if start > end {
            return Err(Error::invalid(format!(
                "window {}:{start}-{end} ends before it starts",
                chr.id
            )));
        }
        Ok(Side {
            chr,
            start,
            end,
            binned_start: 0,
            binned_end: 0,
            bin_start: 0,
            bin_end: 0,
        })
    };
    let mut x = make(&ids.0, x_start, x_end)?;
    let mut y = make(&ids.1, y_start, y_end)?;

    // A hic file stores one side of the diagonal, indexed by the lower
    // chromosome first, so a request naming them the other way round is read
    // swapped and transposed back on the way out.
    let mut reversed = false;
    if x.chr.index > y.chr.index {
        std::mem::swap(&mut x, &mut y);
        reversed = true;
    }

    let bin_size = match bin_count {
        Some(count) if count > 0 => {
            let span = ((x.end - x.start) + (y.end - y.start)) / 2;
            let wanted = (span + count - 1) / count;
            *available_bin_sizes
                .iter()
                .min_by_key(|available| (*available - wanted).abs())
                .expect("checked non-empty")
        }
        _ => match bin_size {
            Some(size) if size > 0 => size,
            _ => *available_bin_sizes.iter().min().expect("checked non-empty"),
        },
    };

    for side in [&mut x, &mut y] {
        side.binned_start = side.start / bin_size * bin_size;
        side.binned_end = if full_bin {
            (side.end + bin_size - 1) / bin_size * bin_size
        } else {
            side.end / bin_size * bin_size
        };
        side.bin_start = side.binned_start / bin_size;
        side.bin_end = side.binned_end / bin_size;
    }

    Ok(Loc2D {
        x,
        y,
        bin_size,
        reversed,
    })
}

/// One resolution of one chromosome pair: where its blocks are and how they are
/// laid out.
#[derive(Debug, Clone)]
pub struct MatrixMetadata {
    pub chr1_index: i64,
    pub chr2_index: i64,
    pub unit: String,
    pub bin_size: i64,
    pub sum_counts: f32,
    pub block_bin_count: i64,
    pub block_column_count: i64,
    pub blocks: IndexMap<i64, HiCIndexItem>,
}

/// The key a matrix is cached under.
pub fn matrix_key(chr1: i64, chr2: i64, bin_size: i64, unit: &str) -> String {
    format!("chr_index={chr1}_{chr2}|bin_size={bin_size}|unit={unit}")
}

pub fn read_matrix_metadata(
    source: &dyn ByteSource,
    item: HiCIndexItem,
    chr1_index: i64,
    chr2_index: i64,
) -> Result<Vec<MatrixMetadata>> {
    let path = source.path();
    // The master index says how long the record is, so it is read in one go
    // rather than streamed.
    let buf = source.read_at(item.position, item.size.max(0) as usize)?;
    let mut c = LeCursor::new(&buf, item.position, path);

    let file_chr1 = c.read_i32()? as i64;
    let file_chr2 = c.read_i32()? as i64;
    let bin_size_count = c.read_i32()? as i64;
    if file_chr1 != chr1_index || file_chr2 != chr2_index {
        return Err(Error::corrupt(
            path,
            item.position,
            "matrix metadata chr indices mismatch",
        ));
    }

    let mut matrices = Vec::with_capacity(bin_size_count.clamp(0, 64) as usize);
    for _ in 0..bin_size_count.max(0) {
        let unit = c.take_cstr()?.to_ascii_lowercase();
        c.skip(4)?; // bin size index in the header
        let sum_counts = c.read_f32()?;
        c.skip(4)?; // occupied cell count
        c.skip(4)?; // 5th percentile estimate
        c.skip(4)?; // 95th percentile estimate
        let bin_size = c.read_i32()? as i64;
        let block_bin_count = c.read_i32()? as i64;
        let block_column_count = c.read_i32()? as i64;
        let block_count = c.read_i32()? as i64;

        let mut blocks = IndexMap::with_capacity(block_count.clamp(0, 1 << 16) as usize);
        for _ in 0..block_count.max(0) {
            let number = c.read_i32()? as i64;
            let position = c.read_u64()?;
            let size = c.read_i32()? as i64;
            blocks.insert(number, HiCIndexItem { position, size });
        }
        matrices.push(MatrixMetadata {
            chr1_index: file_chr1,
            chr2_index: file_chr2,
            unit,
            bin_size,
            sum_counts,
            block_bin_count,
            block_column_count,
            blocks,
        });
    }
    Ok(matrices)
}

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

    fn map() -> ChrMap {
        ChrMap::from_indexed_entries([
            ("chr1".to_string(), 1_000_000, 0),
            ("chr2".to_string(), 500_000, 1),
        ])
    }

    fn loc(chr_ids: &[&str], starts: &[i64], ends: &[i64], bin: Option<i64>) -> Result<Loc2D> {
        let ids: Vec<String> = chr_ids.iter().map(|s| s.to_string()).collect();
        parse_loc2d(
            &map(),
            &[5000, 10000, 25000],
            &ids,
            starts,
            ends,
            bin,
            None,
            false,
        )
    }

    #[test]
    fn one_chromosome_is_used_for_both_axes() {
        let l = loc(&["chr1"], &[0], &[100_000], Some(5000)).unwrap();
        assert_eq!(l.x.chr.id, "chr1");
        assert_eq!(l.y.chr.id, "chr1");
        assert!(l.is_intra());
        assert_eq!((l.x.bin_start, l.x.bin_end), (0, 20));
    }

    #[test]
    fn the_axes_swap_when_the_second_chromosome_sorts_first() {
        let l = loc(&["chr2", "chr1"], &[0, 0], &[100_000, 100_000], Some(5000)).unwrap();
        assert!(l.reversed, "the request named them the other way round");
        assert_eq!(l.x.chr.id, "chr1", "the lower chromosome is read first");
        assert_eq!(l.y.chr.id, "chr2");

        // Named in file order, nothing is swapped.
        let l = loc(&["chr1", "chr2"], &[0, 0], &[100_000, 100_000], Some(5000)).unwrap();
        assert!(!l.reversed);
    }

    #[test]
    fn a_negative_bin_size_takes_the_finest_resolution() {
        assert_eq!(
            loc(&["chr1"], &[0], &[100_000], Some(-1)).unwrap().bin_size,
            5000
        );
        assert_eq!(
            loc(&["chr1"], &[0], &[100_000], None).unwrap().bin_size,
            5000
        );
    }

    #[test]
    fn a_bin_count_picks_the_nearest_available_resolution() {
        let ids = ["chr1".to_string()];
        // 100 kb over 10 bins wants 10 000, which the file has exactly.
        let l = parse_loc2d(
            &map(),
            &[5000, 10000, 25000],
            &ids,
            &[0],
            &[100_000],
            None,
            Some(10),
            false,
        )
        .unwrap();
        assert_eq!(l.bin_size, 10000);
        // 100 kb over 4 bins wants 25 000.
        let l = parse_loc2d(
            &map(),
            &[5000, 10000, 25000],
            &ids,
            &[0],
            &[100_000],
            None,
            Some(4),
            false,
        )
        .unwrap();
        assert_eq!(l.bin_size, 25000);
        // Nothing exact: 100 kb over 7 bins wants ~14 286, nearest is 10 000.
        let l = parse_loc2d(
            &map(),
            &[5000, 10000, 25000],
            &ids,
            &[0],
            &[100_000],
            None,
            Some(7),
            false,
        )
        .unwrap();
        assert_eq!(l.bin_size, 10000);
    }

    #[test]
    fn full_bin_rounds_the_end_up_instead_of_down() {
        let down = loc(&["chr1"], &[0], &[12_000], Some(5000)).unwrap();
        assert_eq!(down.x.bin_end, 2); // 12 000 floors to bin 2
        let ids = ["chr1".to_string()];
        let up = parse_loc2d(
            &map(),
            &[5000],
            &ids,
            &[0],
            &[12_000],
            Some(5000),
            None,
            true,
        )
        .unwrap();
        assert_eq!(up.x.bin_end, 3);
    }

    #[test]
    fn the_refusals_say_what_was_wrong() {
        let ids = ["chr1".to_string()];
        let err = parse_loc2d(&map(), &[], &ids, &[0], &[10], None, None, false)
            .unwrap_err()
            .to_string();
        assert!(err.contains("no resolution for this unit"), "{err}");

        let err = parse_loc2d(&map(), &[5000], &ids, &[0], &[10], Some(0), None, false)
            .unwrap_err()
            .to_string();
        assert!(err.contains("bin size must be positive"), "{err}");

        let err = parse_loc2d(&map(), &[5000], &ids, &[0], &[10], None, Some(0), false)
            .unwrap_err()
            .to_string();
        assert!(err.contains("bin count must be positive"), "{err}");

        let err = loc(&["chr1"], &[100], &[10], Some(5000))
            .unwrap_err()
            .to_string();
        assert!(err.contains("ends before it starts"), "{err}");

        let three = ["chr1".to_string(), "chr2".to_string(), "chr1".to_string()];
        let err = parse_loc2d(&map(), &[5000], &three, &[0], &[10], None, None, false)
            .unwrap_err()
            .to_string();
        assert!(err.contains("1 or 2 chromosomes"), "{err}");
    }

    #[test]
    fn distance_from_the_diagonal_is_the_offset_on_one_chromosome() {
        let l = loc(&["chr1"], &[0], &[100_000], Some(5000)).unwrap();
        assert_eq!(l.distance_from_diagonal(10_000, 10_000), 0);
        assert_eq!(l.distance_from_diagonal(10_000, 35_000), 25_000);
        assert_eq!(l.distance_from_diagonal(35_000, 10_000), 25_000);
    }

    #[test]
    fn across_two_chromosomes_it_is_measured_against_the_windows_own_line() {
        let l = loc(&["chr1", "chr2"], &[0, 0], &[100_000, 100_000], Some(5000)).unwrap();
        assert!(!l.is_intra());
        // The window is square, so its "diagonal" is y = x.
        assert_eq!(l.distance_from_diagonal(20_000, 20_000), 0);
        assert!(l.distance_from_diagonal(20_000, 60_000) > 0);
    }

    #[test]
    fn the_matrix_key_is_built_the_same_way_both_ways() {
        assert_eq!(
            matrix_key(0, 1, 5000, "bp"),
            "chr_index=0_1|bin_size=5000|unit=bp"
        );
    }
}