gwseq-io 0.2.0

//! The HiC reader.
//!
//! Three memoised caches — expected-value vectors, normalization vectors and
//! matrix metadata. Each is a `Mutex<HashMap<Key, Arc<_>>>`: the `Arc` is what
//! lets a lookup drop the lock before the vector is used, and it is carried
//! into the read rather than cloned out of, since a vector is one `f32` per bin
//! of a chromosome and copying two per read would dominate a small request.

use std::collections::HashMap;
use std::sync::Arc;

use ndarray::Array2;
use parking_lot::Mutex;

use crate::arrays::CooMatrix;
use crate::error::{Error, Result};
use crate::genomic::{ChrMap, Locs};
use crate::parallel::Executor;
use crate::source::ByteSource;

use super::block::{block_numbers, read_block, ContactRecord, RecordContext};
use super::header::{vector_key, HiCFooter, HiCHeader};
use super::matrix::{matrix_key, parse_loc2d, Loc2D, MatrixMetadata};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum HiCMode {
    #[default]
    Observed,
    /// Observed / expected.
    Oe,
    Expected,
}

impl std::str::FromStr for HiCMode {
    type Err = Error;
    fn from_str(s: &str) -> Result<Self> {
        match s.to_ascii_lowercase().as_str() {
            "observed" => Ok(HiCMode::Observed),
            "oe" => Ok(HiCMode::Oe),
            "expected" => Ok(HiCMode::Expected),
            o => Err(Error::invalid(format!(
                "mode {o} invalid (observed, oe or expected)"
            ))),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Unit {
    #[default]
    Bp,
    Frag,
}

impl Unit {
    pub fn as_str(self) -> &'static str {
        match self {
            Unit::Bp => "bp",
            Unit::Frag => "frag",
        }
    }
}

impl std::str::FromStr for Unit {
    type Err = Error;
    fn from_str(s: &str) -> Result<Self> {
        match s.to_ascii_lowercase().as_str() {
            "bp" => Ok(Unit::Bp),
            "frag" => Ok(Unit::Frag),
            o => Err(Error::invalid(format!("unit {o} invalid (bp or frag)"))),
        }
    }
}

/// The vectors a read values its contacts through.
#[derive(Debug, Clone, Default)]
pub struct Normalizations {
    pub x: Arc<Vec<f32>>,
    pub y: Arc<Vec<f32>>,
    pub expected: Arc<Vec<f32>>,
}

#[derive(Debug)]
struct Inner {
    source: Arc<dyn ByteSource>,
    executor: Executor,
}

pub struct HiCReader {
    inner: Option<Inner>,
    path: String,
    header: HiCHeader,
    footer: HiCFooter,

    expected_cache: Mutex<HashMap<String, Arc<Vec<f32>>>>,
    norm_cache: Mutex<HashMap<String, Arc<Vec<f32>>>>,
    matrix_cache: Mutex<HashMap<String, Arc<MatrixMetadata>>>,
}

impl std::fmt::Debug for HiCReader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("HiCReader")
            .field("path", &self.path)
            .field("version", &self.header.version)
            .field("chromosomes", &self.header.chr_map.len())
            .field("closed", &self.is_closed())
            .finish()
    }
}

impl HiCReader {
    pub fn open(
        path: &str,
        parallel: i64,
        block_size: Option<u64>,
        max_blocks: Option<usize>,
    ) -> Result<Self> {
        let source = crate::source::open(path, block_size, max_blocks)?;
        Self::from_source(source, path, parallel)
    }

    pub(crate) fn from_source(
        source: Arc<dyn ByteSource>,
        path: &str,
        parallel: i64,
    ) -> Result<Self> {
        let header = super::header::read_header(source.as_ref())?;
        let footer = super::header::read_footer(source.as_ref(), &header)?;
        Ok(Self {
            inner: Some(Inner {
                source,
                executor: Executor::new(parallel)?,
            }),
            path: path.to_string(),
            header,
            footer,
            expected_cache: Mutex::new(HashMap::new()),
            norm_cache: Mutex::new(HashMap::new()),
            matrix_cache: Mutex::new(HashMap::new()),
        })
    }

    pub fn header(&self) -> &HiCHeader {
        &self.header
    }
    pub fn footer(&self) -> &HiCFooter {
        &self.footer
    }
    pub fn chr_sizes(&self) -> &ChrMap {
        &self.header.chr_map
    }
    pub fn normalizations(&self) -> &[String] {
        &self.footer.normalizations
    }
    pub fn units(&self) -> &[String] {
        &self.footer.units
    }
    /// The resolutions the file was built at, for one unit, in file order.
    ///
    /// Left in the order the header lists them — largest first, as a hic file
    /// writes them — rather than sorted, because a caller indexing this by
    /// position would otherwise get a different resolution than the file's own
    /// order gives.
    pub fn bin_sizes(&self, unit: Unit) -> &[i64] {
        self.header.resolutions(unit)
    }
    pub fn path(&self) -> &str {
        &self.path
    }
    pub fn is_closed(&self) -> bool {
        self.inner.is_none()
    }
    pub fn parallel(&self) -> usize {
        self.inner.as_ref().map_or(0, |i| i.executor.parallel())
    }

    pub fn close(&mut self) {
        if let Some(inner) = self.inner.take() {
            inner.source.close();
        }
    }

    fn inner(&self) -> Result<&Inner> {
        self.inner.as_ref().ok_or_else(|| Error::Closed {
            path: self.path.clone(),
        })
    }

    /// Resolve a request's two windows against the file's own resolutions.
    ///
    /// The coordinates go through `Locs` before `parse_loc2d`. The two checks
    /// are not the same one: `Locs` demands one start and one end *per chromosome id*,
    /// while `parse_loc2d` accepts one or two of each and duplicates a single
    /// value across both axes. Without the first, two chromosomes and one start
    /// read the same window twice instead of being refused.
    pub fn parse_loc(&self, req: &HiCRequest) -> Result<Loc2D> {
        let locs = Locs::spans(&req.chr_ids, &req.starts, &req.ends)?;
        parse_loc2d(
            &self.header.chr_map,
            self.header.resolutions(req.unit),
            &locs.chr_ids,
            &locs.starts,
            &locs.ends,
            req.bin_size,
            req.bin_count.map(|n| n as i64),
            req.full_bin,
        )
    }

    fn expected_values(
        &self,
        chr: i64,
        normalization: &str,
        bin_size: i64,
        unit: Unit,
    ) -> Result<Arc<Vec<f32>>> {
        let key = vector_key(normalization, bin_size, unit.as_str(), Some(chr));
        if let Some(hit) = self.expected_cache.lock().get(&key) {
            return Ok(hit.clone());
        }
        let values = Arc::new(super::header::compute_expected_values(
            &self.footer,
            chr,
            unit.as_str(),
            bin_size,
            normalization,
        )?);
        self.expected_cache.lock().insert(key, values.clone());
        Ok(values)
    }

    fn normalization_vector(
        &self,
        inner: &Inner,
        chr: i64,
        normalization: &str,
        bin_size: i64,
        unit: Unit,
    ) -> Result<Arc<Vec<f32>>> {
        let key = vector_key(normalization, bin_size, unit.as_str(), Some(chr));
        if let Some(hit) = self.norm_cache.lock().get(&key) {
            return Ok(hit.clone());
        }
        let values = Arc::new(super::header::read_normalization_vector(
            inner.source.as_ref(),
            &self.footer,
            self.header.version,
            chr,
            unit.as_str(),
            bin_size,
            normalization,
        )?);
        self.norm_cache.lock().insert(key, values.clone());
        Ok(values)
    }

    /// The matrix for one chromosome pair at one resolution, or `None` when the
    /// file holds no matrix for that **pair** at all.
    ///
    /// A miss reads the whole record — every resolution of that pair — and
    /// caches all of them, since the file stores them together and a second
    /// resolution of the same pair is the common next request.
    ///
    /// The two absences are different answers. A pair the file simply does not
    /// carry is ordinary — anything against chrM, or a pair below the coverage
    /// threshold whatever wrote the file used — and comes back as no contacts,
    /// which is what straw and hictk report and what a caller sweeping every
    /// pair of a genome needs. A *resolution* the pair does not have is the
    /// caller asking for something the file cannot serve, and still raises,
    /// listing what it has. The chromosome names were resolved through `ChrMap`
    /// long before here, so a typo is already an `UnknownChromosome` and cannot
    /// arrive as an empty matrix.
    fn matrix(
        &self,
        inner: &Inner,
        loc: &Loc2D,
        unit: Unit,
    ) -> Result<Option<Arc<MatrixMetadata>>> {
        let (chr1, chr2) = (loc.x.chr.index as i64, loc.y.chr.index as i64);
        let key = matrix_key(chr1, chr2, loc.bin_size, unit.as_str());
        if let Some(hit) = self.matrix_cache.lock().get(&key) {
            return Ok(Some(hit.clone()));
        }

        let index_key = format!("{chr1}_{chr2}");
        let Some(item) = self.footer.master_index.get(&index_key).copied() else {
            return Ok(None);
        };
        let matrices =
            super::matrix::read_matrix_metadata(inner.source.as_ref(), item, chr1, chr2)?;

        let mut available = Vec::new();
        {
            let mut cache = self.matrix_cache.lock();
            for matrix in matrices {
                available.push(format!("{}{}", matrix.bin_size, matrix.unit));
                let entry_key = matrix_key(chr1, chr2, matrix.bin_size, &matrix.unit);
                cache.insert(entry_key, Arc::new(matrix));
            }
            if let Some(hit) = cache.get(&key) {
                return Ok(Some(hit.clone()));
            }
        }
        Err(Error::invalid(format!(
            "no matrix for {key} (available for this pair: {})",
            available.join(", ")
        )))
    }

    /// Every contact of a window, decoded on the reader's threads.
    fn records(&self, req: &HiCRequest, loc: &Loc2D) -> Result<Vec<ContactRecord>> {
        let inner = self.inner()?;
        let normalization = req.normalization.to_ascii_lowercase();
        let Some(matrix) = self.matrix(inner, loc, req.unit)? else {
            // No matrix for this pair: no contacts, and every bin reads as
            // `def_value`. See `matrix`.
            return Ok(Vec::new());
        };

        // Held as `Arc`s rather than cloned out of them, which is the whole
        // reason the caches hand `Arc`s back: a normalization vector is one
        // `f32` per bin of a chromosome, and copying two of them per read was
        // the cost this was built to avoid.
        let mut vectors = Normalizations::default();
        if normalization != "none" {
            vectors.x = self.normalization_vector(
                inner,
                loc.x.chr.index as i64,
                &normalization,
                loc.bin_size,
                req.unit,
            )?;
            vectors.y = self.normalization_vector(
                inner,
                loc.y.chr.index as i64,
                &normalization,
                loc.bin_size,
                req.unit,
            )?;
        }
        if req.mode != HiCMode::Observed && loc.is_intra() {
            vectors.expected = self.expected_values(
                loc.x.chr.index as i64,
                &normalization,
                loc.bin_size,
                req.unit,
            )?;
        }

        // The mean of an inter-chromosomal matrix, which stands in for the
        // expected value there. Counted in bins, so a chromosome shorter than
        // one bin holds none of them and there is no mean to take.
        let mut average_value = f32::NAN;
        if !loc.is_intra() {
            let x_bins = loc.x.chr.size / loc.bin_size;
            let y_bins = loc.y.chr.size / loc.bin_size;
            if x_bins > 0 && y_bins > 0 {
                average_value = matrix.sum_counts / x_bins as f32 / y_bins as f32;
            }
        }

        let numbers = block_numbers(
            loc,
            &matrix,
            req.max_distance,
            self.header.version,
            req.triangle,
        )?;
        let blocks: Vec<_> = numbers
            .iter()
            .filter_map(|n| matrix.blocks.get(n).copied())
            .collect();

        let ctx = RecordContext {
            loc,
            normalization: &normalization,
            mode: req.mode,
            vectors: &vectors,
            average_value,
            min_distance: req.min_distance,
            max_distance: req.max_distance,
        };
        let per_block = inner.executor.map_batches(&blocks, |_, block| {
            let raw = inner
                .source
                .read_exact_at(block.position, block.size.max(0) as usize)?;
            read_block(raw, self.header.version, *block, &ctx, &self.path)
        })?;
        Ok(per_block.into_iter().flatten().collect())
    }

    /// Dense `(loc1 bins, loc2 bins)`.
    ///
    /// A bin holding a contact the file **cannot value** — one the chosen
    /// normalization has no factor for, or an expected value of zero — comes
    /// back `NaN`, which is a different answer from `def_value`, meaning "not
    /// observed at all". This matches straw and hictk.
    ///
    /// An out-of-range window pads with `def_value` like bbi; only the block
    /// enumeration is clamped.
    pub fn read_values(&self, req: &HiCRequest) -> Result<Array2<f32>> {
        let loc = self.parse_loc(req)?;
        let records = self.records(req, &loc)?;

        let rows = (loc.x.bin_end - loc.x.bin_start).max(0) as usize;
        let cols = (loc.y.bin_end - loc.y.bin_start).max(0) as usize;
        let mut flat = vec![req.def_value; rows * cols];
        for record in &records {
            let r = record.x_bin - loc.x.bin_start;
            let c = record.y_bin - loc.y.bin_start;
            if r >= 0 && (r as usize) < rows && c >= 0 && (c as usize) < cols {
                flat[r as usize * cols + c as usize] = record.value;
            }
            // The file stores one side of the diagonal, so the other is filled
            // from it — unless `triangle`, which asks for the stored side alone.
            if loc.is_intra() && !req.triangle {
                let r = record.y_bin - loc.x.bin_start;
                let c = record.x_bin - loc.y.bin_start;
                if r >= 0 && (r as usize) < rows && c >= 0 && (c as usize) < cols {
                    flat[r as usize * cols + c as usize] = record.value;
                }
            }
        }

        let (mut rows, mut cols, mut flat) = if loc.reversed {
            (cols, rows, transpose(&flat, rows, cols))
        } else {
            (rows, cols, flat)
        };

        if req.exact_bin_count {
            if let Some(count) = req.bin_count {
                flat = crate::arrays::bilinear(&flat, (rows, cols), (count, count))?;
                rows = count;
                cols = count;
            }
        }
        Array2::from_shape_vec((rows, cols), flat)
            .map_err(|e| Error::invalid(format!("output shape {rows}x{cols}: {e}")))
    }

    /// The same, sparse.
    ///
    /// No `def_value`: a cell this does not list is one no contact reached,
    /// which is what a sparse matrix says by leaving it out.
    pub fn read_sparse_values(&self, req: &HiCRequest) -> Result<CooMatrix> {
        let loc = self.parse_loc(req)?;
        let records = self.records(req, &loc)?;

        let rows = (loc.x.bin_end - loc.x.bin_start).max(0) as usize;
        let cols = (loc.y.bin_end - loc.y.bin_start).max(0) as usize;
        let mut out = CooMatrix {
            shape: (rows, cols),
            ..Default::default()
        };
        let push = |r: i64, c: i64, value: f32, out: &mut CooMatrix| {
            if r >= 0 && (r as usize) < rows && c >= 0 && (c as usize) < cols {
                out.values.push(value);
                out.row.push(r as u32);
                out.col.push(c as u32);
            }
        };
        for record in &records {
            push(
                record.x_bin - loc.x.bin_start,
                record.y_bin - loc.y.bin_start,
                record.value,
                &mut out,
            );
            // Except on the diagonal itself, where the two are the same cell.
            // Writing it twice is invisible in the dense matrix, which assigns,
            // and doubles it in this one, which is a list of entries a reader is
            // expected to sum.
            if loc.is_intra() && !req.triangle && record.x_bin != record.y_bin {
                push(
                    record.y_bin - loc.x.bin_start,
                    record.x_bin - loc.y.bin_start,
                    record.value,
                    &mut out,
                );
            }
        }

        // By row then column, so the entries come back in a stable order
        // whatever order the blocks were decoded in.
        let mut out = sort_row_major(out);

        if loc.reversed {
            std::mem::swap(&mut out.row, &mut out.col);
            out.shape = (out.shape.1, out.shape.0);
            // Swapping the two coordinate arrays leaves them ordered by the
            // *old* row, which is the new column. Sorting again is what keeps
            // the row-major promise on a reversed request, and what
            // `compress_sparse_by_color` needs downstream.
            out = sort_row_major(out);
        }

        if req.exact_bin_count {
            if let Some(count) = req.bin_count {
                out = crate::arrays::bilinear_sparse(&out, (count, count))?;
            }
        }
        Ok(out)
    }
}

/// The entries of a COO matrix ordered by row then column.
///
/// Not a stable sort of the arrays in place: the three run in parallel, so the
/// permutation is worked out once and applied to all three.
fn sort_row_major(coo: CooMatrix) -> CooMatrix {
    let mut order: Vec<usize> = (0..coo.values.len()).collect();
    order.sort_by_key(|i| (coo.row[*i], coo.col[*i]));
    CooMatrix {
        values: order.iter().map(|i| coo.values[*i]).collect(),
        row: order.iter().map(|i| coo.row[*i]).collect(),
        col: order.iter().map(|i| coo.col[*i]).collect(),
        shape: coo.shape,
    }
}

fn transpose(flat: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    let mut out = vec![0.0f32; flat.len()];
    for r in 0..rows {
        for c in 0..cols {
            out[c * rows + r] = flat[r * cols + c];
        }
    }
    out
}

#[derive(Debug, Clone)]
pub struct HiCRequest {
    pub chr_ids: Vec<String>,
    pub starts: Vec<i64>,
    pub ends: Vec<i64>,
    /// `None` takes the smallest available. Must be one the file has.
    pub bin_size: Option<i64>,
    /// Takes precedence over `bin_size`, selecting the closest bin size that
    /// gives about this many bins.
    pub bin_count: Option<usize>,
    /// Resize the output to match `bin_count` exactly, bilinearly.
    pub exact_bin_count: bool,
    pub full_bin: bool,
    pub def_value: f32,
    /// On one chromosome a hic file stores one side of the diagonal only.
    /// `triangle` reads that side alone rather than mirroring it — so a window
    /// lying entirely on the *other* side comes back empty.
    pub triangle: bool,
    pub min_distance: Option<i64>,
    pub max_distance: Option<i64>,
    pub normalization: String,
    pub mode: HiCMode,
    pub unit: Unit,
}

impl HiCRequest {
    pub fn new(chr_ids: Vec<String>, starts: Vec<i64>, ends: Vec<i64>) -> Self {
        Self {
            chr_ids,
            starts,
            ends,
            bin_size: None,
            bin_count: None,
            exact_bin_count: false,
            full_bin: false,
            def_value: 0.0,
            triangle: false,
            min_distance: None,
            max_distance: None,
            normalization: "none".into(),
            mode: HiCMode::Observed,
            unit: Unit::Bp,
        }
    }
}

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

    #[test]
    fn modes_and_units_parse_and_refuse() {
        assert_eq!(HiCMode::from_str("oe").unwrap(), HiCMode::Oe);
        assert_eq!(HiCMode::from_str("OBSERVED").unwrap(), HiCMode::Observed);
        let err = HiCMode::from_str("median").unwrap_err().to_string();
        assert!(err.contains("mode median invalid"), "{err}");

        assert_eq!(Unit::from_str("BP").unwrap(), Unit::Bp);
        assert_eq!(Unit::from_str("frag").unwrap(), Unit::Frag);
        let err = Unit::from_str("kb").unwrap_err().to_string();
        assert!(err.contains("unit kb invalid (bp or frag)"), "{err}");
    }

    #[test]
    fn transposing_swaps_the_axes() {
        // 2x3 row-major becomes 3x2.
        let flat = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
        assert_eq!(transpose(&flat, 2, 3), [1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
    }
}