gwseq-io 0.2.0

//! Array shaping shared by more than one format.
//!
//! The two bilinear resizes behind `exact_bin_count`, and
//! `compress_sparse_by_color`, which the module exposes as a free function.
//!
//! Hand-written rather than reached for from an image crate, because the
//! sample positions, the edge behaviour *and the arithmetic the four corners
//! are combined with* decide the output, and a matrix of contacts is not a
//! photograph — a resized cell is a number someone will do statistics on.
//!
//! Two things about that arithmetic are deliberate and neither is obvious:
//!
//! - The interpolation is grouped as
//!   `(tl*(1-cf) + tr*cf)*(1-rf) + (bl*(1-cf) + br*cf)*rf`, which is
//!   algebraically the same as summing four weighted corners and is not the
//!   same `f64`. Nesting keeps each partial sum between two corners that are
//!   neighbours, and so of similar magnitude.
//! - Each of those three `a*b + c*d` pairs is a `mul_add` — one rounding
//!   instead of two, and a single `fmadd` instruction on every target this
//!   builds for. Written as a separate multiply and add it is a *different*
//!   answer, by about an ulp; `mul_add` is the more accurate one.

use crate::error::{Error, Result};

/// A COO sparse matrix, as `HiCReader::read_sparse_values` returns.
///
/// Repeated coordinates accumulate, per the COO convention: an element's value
/// is the sum of every entry naming it.
#[derive(Debug, Clone, Default)]
pub struct CooMatrix {
    pub values: Vec<f32>,
    pub row: Vec<u32>,
    pub col: Vec<u32>,
    pub shape: (usize, usize),
}

/// Resize a flat row-major matrix by bilinear interpolation.
///
/// The scale maps an output coordinate onto a source one so that the two grids
/// share their *corners* — output 0 samples source 0, and the last output
/// samples the last source — which is what keeps a resize of a matrix to its
/// own shape an identity.
pub fn bilinear(
    data: &[f32],
    shape: (usize, usize),
    new_shape: (usize, usize),
) -> Result<Vec<f32>> {
    let (old_rows, old_cols) = shape;
    let (new_rows, new_cols) = new_shape;
    if data.len() != old_rows * old_cols {
        return Err(Error::invalid(
            "bilinear resize: data size does not match shape",
        ));
    }
    if data.is_empty() && new_rows > 0 && new_cols > 0 {
        return Err(Error::invalid(
            "bilinear resize: cannot resize an empty array to a non-empty shape",
        ));
    }
    if new_rows == 0 || new_cols == 0 {
        return Ok(Vec::new());
    }

    let (row_scale, col_scale) = scales((old_rows, old_cols), (new_rows, new_cols));

    let mut out = vec![0.0f32; new_rows * new_cols];
    for i in 0..new_rows {
        // Everything about the source rows is fixed by `i`, so it is worked out
        // once per output row rather than once per output cell.
        let sr = i as f64 * row_scale;
        let r0 = sr as usize;
        let r1 = (r0 + 1).min(old_rows - 1);
        let rf = sr - r0 as f64;
        let (top, bottom) = (r0 * old_cols, r1 * old_cols);

        for j in 0..new_cols {
            let sc = j as f64 * col_scale;
            let c0 = sc as usize;
            let c1 = (c0 + 1).min(old_cols - 1);
            let cf = sc - c0 as f64;

            let tl = data[top + c0] as f64;
            let tr = data[top + c1] as f64;
            let bl = data[bottom + c0] as f64;
            let br = data[bottom + c1] as f64;
            out[i * new_cols + j] = interpolate(tl, tr, bl, br, rf, cf);
        }
    }
    Ok(out)
}

/// The same, for a sparse matrix.
///
/// Only the small neighbourhood of output cells each non-zero input cell can
/// influence is visited, so the cost follows the number of non-zeros rather
/// than the full output area. Entries that interpolate to exactly zero are left
/// out, and the result comes back in row-major order — which is what
/// [`compress_sparse_by_color`] needs, since it merges a run only against the
/// last span it opened.
pub fn bilinear_sparse(data: &CooMatrix, new_shape: (usize, usize)) -> Result<CooMatrix> {
    let (old_rows, old_cols) = data.shape;
    let (new_rows, new_cols) = new_shape;

    if data.row.len() != data.values.len() || data.col.len() != data.values.len() {
        return Err(Error::invalid(
            "bilinear resize: values, row and col must have equal length",
        ));
    }
    for i in 0..data.values.len() {
        if data.row[i] as usize >= old_rows || data.col[i] as usize >= old_cols {
            return Err(Error::invalid(
                "bilinear resize: coordinate outside declared shape",
            ));
        }
    }
    if (old_rows == 0 || old_cols == 0) && new_rows > 0 && new_cols > 0 {
        return Err(Error::invalid(
            "bilinear resize: cannot resize an empty array to a non-empty shape",
        ));
    }
    // Empty output: nothing to interpolate into. Returning early also keeps the
    // `new_rows - 1` extents below from wrapping around.
    if new_rows == 0 || new_cols == 0 {
        return Ok(CooMatrix {
            shape: new_shape,
            ..Default::default()
        });
    }

    let (row_scale, col_scale) = scales((old_rows, old_cols), (new_rows, new_cols));

    // Dense lookup for O(1) sparse reads, accumulating repeated coordinates.
    let mut sparse_map = std::collections::HashMap::with_capacity(data.values.len());
    for i in 0..data.values.len() {
        let key = ((data.row[i] as u64) << 32) | data.col[i] as u64;
        *sparse_map.entry(key).or_insert(0.0f32) += data.values[i];
    }
    let get = |r: usize, c: usize| -> f32 {
        sparse_map
            .get(&(((r as u64) << 32) | c as u64))
            .copied()
            .unwrap_or(0.0)
    };

    // Pass 1: which output cells any non-zero input cell can influence.
    // Collecting the set first means a cell reached from several inputs is
    // computed once rather than once per input, and it bounds the degenerate
    // case: at a scale of zero every input influences the whole output.
    let mut targets = std::collections::HashSet::new();
    let full_output = new_rows * new_cols;
    for idx in 0..data.values.len() {
        // Skipping zeroes cannot skip a coordinate that matters: an accumulated
        // value is only non-zero if some entry naming it was, and that entry
        // visits the same neighbourhood. Interpolation reads the accumulated
        // total through `get`, so duplicates are never counted twice.
        if data.values[idx] == 0.0 {
            continue;
        }
        // Output cells this input influences, by inverse mapping out = in /
        // scale. Output `o` samples source rows floor(o * scale) and the next,
        // so `src_r` is read when o * scale lies in [src_r - 1, src_r + 1) — a
        // half-width of 1/scale output cells, not 1. The two are equal only at
        // scale == 1; upsampling with a fixed +/-1 window silently drops
        // contributions.
        //
        // A scale of zero means the output or the input has a single row, and
        // every output row then draws on source row 0.
        let (r_min, r_max) = influence(data.row[idx], row_scale, new_rows);
        let (c_min, c_max) = influence(data.col[idx], col_scale, new_cols);
        for r in r_min..=r_max {
            for c in c_min..=c_max {
                targets.insert(((r as u64) << 32) | c as u64);
            }
        }
        // Every output cell is already spoken for, so no further input can add
        // one. Without this the zero-scale case keeps re-inserting the whole
        // output grid, once per non-zero.
        if targets.len() >= full_output {
            break;
        }
    }

    // Pass 2: interpolate each influenced output cell once. Sorted, which for a
    // key of (row << 32 | col) is row-major order.
    let mut ordered: Vec<u64> = targets.into_iter().collect();
    ordered.sort_unstable();

    let mut out = CooMatrix {
        shape: new_shape,
        values: Vec::with_capacity(ordered.len()),
        row: Vec::with_capacity(ordered.len()),
        col: Vec::with_capacity(ordered.len()),
    };
    for key in ordered {
        let (or_, oc) = ((key >> 32) as u32, key as u32);
        let sr = or_ as f64 * row_scale;
        let sc = oc as f64 * col_scale;
        let r0 = sr as usize;
        let c0 = sc as usize;
        let r1 = (r0 + 1).min(old_rows - 1);
        let c1 = (c0 + 1).min(old_cols - 1);
        let (rf, cf) = (sr - r0 as f64, sc - c0 as f64);

        let interp = interpolate(
            get(r0, c0) as f64,
            get(r0, c1) as f64,
            get(r1, c0) as f64,
            get(r1, c1) as f64,
            rf,
            cf,
        );
        if interp != 0.0 {
            out.row.push(or_);
            out.col.push(oc);
            out.values.push(interp);
        }
    }
    Ok(out)
}

/// One output cell from its four source corners.
///
/// Three `mul_add`s: fused, so each pair rounds once rather than twice. See
/// the module docs on why the expression is nested rather than flattened.
fn interpolate(tl: f64, tr: f64, bl: f64, br: f64, rf: f64, cf: f64) -> f32 {
    let top = tl.mul_add(1.0 - cf, tr * cf);
    let bottom = bl.mul_add(1.0 - cf, br * cf);
    top.mul_add(1.0 - rf, bottom * rf) as f32
}

/// A single output row or column samples only index 0, so a scale of 0 is right
/// rather than a division by zero.
fn scales(old: (usize, usize), new: (usize, usize)) -> (f64, f64) {
    let row = if new.0 > 1 {
        (old.0 - 1) as f64 / (new.0 - 1) as f64
    } else {
        0.0
    };
    let col = if new.1 > 1 {
        (old.1 - 1) as f64 / (new.1 - 1) as f64
    } else {
        0.0
    };
    (row, col)
}

/// The inclusive range of output indices one source index can reach.
fn influence(src: u32, scale: f64, new_len: usize) -> (u32, u32) {
    if scale == 0.0 {
        return (0, new_len as u32 - 1);
    }
    let out = src as f64 / scale;
    let radius = 1.0 / scale;
    let min = (out - radius).floor().max(0.0) as u32;
    let max = (out + radius).ceil().min((new_len - 1) as f64) as u32;
    (min, max)
}

/// Group the entries of a sparse matrix by value bin, as runs of adjacent
/// cells.
///
/// `color_count` equal-width bins are laid over `[0, max value]`, each
/// `max value / color_count` wide. The last is closed at both ends, so the
/// maximum falls in it rather than in one of its own.
///
/// Returns one flat list of `{row, column, length}` triples per bin, holding
/// the entries whose value falls in it. Consecutive cells of a row extend the
/// last triple instead of opening one, so a run of one colour is a single span.
///
/// What a renderer draws, worked out once instead of per frame. The merge only
/// looks at the last triple of a bin, so entries out of row-major order still
/// come back, as spans of length 1.
///
/// Values outside the scale are clamped to its ends rather than dropped, as a
/// colour scale does. NaN and infinite values belong to no bin and are left out
/// of both the scale and the output. A matrix with no scale to spread comes
/// back as empty bins.
///
/// `row` and `col` must be as long as `values`; the binding layer checks that
/// and reports a mismatch as its own kind of failure. A short one is treated
/// here as no entry at all.
pub fn compress_sparse_by_color(
    values: &[f32],
    row: &[u32],
    col: &[u32],
    color_count: u32,
) -> Result<Vec<Vec<u32>>> {
    if color_count == 0 {
        return Err(Error::invalid(
            "compress_sparse_by_color: color_count must be positive",
        ));
    }
    let mut result = vec![Vec::<u32>::new(); color_count as usize];
    if values.is_empty() {
        return Ok(result);
    }

    // The scale starts at 0 rather than at the smallest value: a sparse matrix
    // is mostly the zeros it does not store, and a scale that left them off its
    // low end would colour them as something else.
    let min_value = 0.0f32;
    let mut max_value = min_value;
    for &value in values {
        // Non-finite entries take no part in the scale: an infinity would
        // stretch it until every other value fell in the first bin, and a NaN
        // would leave the whole scale NaN and every output bin empty.
        if value.is_finite() && value > max_value {
            max_value = value;
        }
    }
    // No scale to spread the values over.
    if max_value <= min_value {
        return Ok(result);
    }
    let range = max_value - min_value;
    // Scaled by the bin count, not the last bin index, so the bins are the
    // equal ones the contract promises. `color_count - 1` left the top bin
    // holding the maximum alone — for a renderer, a colour almost nothing is
    // drawn in.
    let scale = color_count as f32;
    let last_bin = color_count - 1;

    for i in 0..values.len().min(row.len()).min(col.len()) {
        if !values[i].is_finite() {
            continue;
        }
        let bin = ((values[i] - min_value) / range * scale).floor();
        // Clamped before the cast: the maximum itself lands on `color_count`,
        // and clamping it here is what closes the top bin at both ends.
        let value_bin = if bin <= 0.0 {
            0
        } else if bin >= last_bin as f32 {
            last_bin
        } else {
            bin as u32
        };

        let spans = &mut result[value_bin as usize];
        let n = spans.len();
        if n >= 3 && spans[n - 3] == row[i] && spans[n - 2] + spans[n - 1] == col[i] {
            spans[n - 1] += 1;
            continue;
        }
        spans.extend_from_slice(&[row[i], col[i], 1]);
    }
    Ok(result)
}

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

    fn coo(entries: &[(u32, u32, f32)], shape: (usize, usize)) -> CooMatrix {
        CooMatrix {
            values: entries.iter().map(|e| e.2).collect(),
            row: entries.iter().map(|e| e.0).collect(),
            col: entries.iter().map(|e| e.1).collect(),
            shape,
        }
    }

    #[test]
    fn resizing_to_its_own_shape_is_the_identity() {
        let data: Vec<f32> = (0..12).map(|v| v as f32).collect();
        assert_eq!(bilinear(&data, (3, 4), (3, 4)).unwrap(), data);
    }

    #[test]
    fn the_corners_are_kept_whatever_the_new_shape() {
        let data = vec![1.0f32, 2.0, 3.0, 4.0];
        let out = bilinear(&data, (2, 2), (5, 5)).unwrap();
        assert_eq!(out[0], 1.0);
        assert_eq!(out[4], 2.0);
        assert_eq!(out[20], 3.0);
        assert_eq!(out[24], 4.0);
    }

    #[test]
    fn a_midpoint_is_the_mean_of_its_four_neighbours() {
        let data = vec![0.0f32, 10.0, 20.0, 30.0];
        let out = bilinear(&data, (2, 2), (3, 3)).unwrap();
        // The centre of a 3x3 samples exactly halfway in both directions.
        assert_eq!(out[4], 15.0);
        // And the edge midpoints are the means of their two neighbours.
        assert_eq!(out[1], 5.0);
        assert_eq!(out[3], 10.0);
    }

    #[test]
    fn shrinking_samples_rather_than_averages() {
        let data: Vec<f32> = (0..16).map(|v| v as f32).collect();
        let out = bilinear(&data, (4, 4), (2, 2)).unwrap();
        // Corners of the source, since the grids share them.
        assert_eq!(out, [0.0, 3.0, 12.0, 15.0]);
    }

    #[test]
    fn a_single_output_cell_samples_the_first_source_one() {
        let data = vec![7.0f32, 8.0, 9.0, 10.0];
        assert_eq!(bilinear(&data, (2, 2), (1, 1)).unwrap(), [7.0]);
    }

    #[test]
    fn a_mismatched_shape_is_refused_rather_than_read_past() {
        // These messages reach Python unchanged and callers match on
        // fragments of them, so the prefix is contract. It names the operation
        // — `bilinear` and `bilinear_sparse` share it — rather than either
        // function, which is what keeps the two saying the same thing about the
        // same mistake.
        let err = bilinear(&[1.0, 2.0], (3, 4), (2, 2))
            .unwrap_err()
            .to_string();
        assert_eq!(err, "bilinear resize: data size does not match shape");
        let err = bilinear(&[], (0, 0), (2, 2)).unwrap_err().to_string();
        assert_eq!(
            err,
            "bilinear resize: cannot resize an empty array to a non-empty shape"
        );
    }

    #[test]
    fn an_empty_target_is_empty_not_an_error() {
        assert!(bilinear(&[1.0, 2.0, 3.0, 4.0], (2, 2), (0, 5))
            .unwrap()
            .is_empty());
    }

    // -- sparse ------------------------------------------------------------

    #[test]
    fn a_sparse_resize_agrees_with_the_dense_one_cell_for_cell() {
        // The two are separate implementations of the same interpolation, and
        // the sparse one visits a neighbourhood rather than the whole grid.
        let entries = [(0u32, 0u32, 1.0f32), (1, 2, 5.0), (3, 3, -2.0), (2, 1, 4.5)];
        let sparse = coo(&entries, (4, 4));
        let mut dense = vec![0.0f32; 16];
        for (r, c, v) in entries {
            dense[r as usize * 4 + c as usize] = v;
        }
        for new_shape in [(2, 2), (4, 4), (7, 7), (3, 5)] {
            let want = bilinear(&dense, (4, 4), new_shape).unwrap();
            let got = bilinear_sparse(&sparse, new_shape).unwrap();
            assert_eq!(got.shape, new_shape);
            for i in 0..got.values.len() {
                let flat = got.row[i] as usize * new_shape.1 + got.col[i] as usize;
                assert_eq!(got.values[i], want[flat], "{new_shape:?} entry {i}");
            }
            // And every cell it left out is one the dense resize made zero.
            let listed: std::collections::HashSet<usize> = (0..got.values.len())
                .map(|i| got.row[i] as usize * new_shape.1 + got.col[i] as usize)
                .collect();
            for (flat, value) in want.iter().enumerate() {
                assert!(
                    *value == 0.0 || listed.contains(&flat),
                    "{new_shape:?} {flat}"
                );
            }
        }
    }

    #[test]
    fn a_sparse_resize_comes_back_in_row_major_order() {
        // The hash-set pass 1 is what made this worth asserting: unsorted
        // output degrades compress_sparse_by_color to one span per cell.
        let sparse = coo(
            &[(5, 5, 1.0), (0, 9, 2.0), (9, 0, 3.0), (2, 2, 4.0)],
            (10, 10),
        );
        let out = bilinear_sparse(&sparse, (6, 6)).unwrap();
        let keys: Vec<u64> = (0..out.values.len())
            .map(|i| ((out.row[i] as u64) << 32) | out.col[i] as u64)
            .collect();
        assert!(keys.windows(2).all(|w| w[0] < w[1]), "{keys:?}");
    }

    #[test]
    fn repeated_sparse_coordinates_accumulate() {
        let sparse = coo(&[(0, 0, 1.0), (0, 0, 2.0)], (2, 2));
        let out = bilinear_sparse(&sparse, (2, 2)).unwrap();
        assert_eq!(out.values[0], 3.0);
    }

    #[test]
    fn a_sparse_coordinate_outside_the_shape_is_refused() {
        let sparse = coo(&[(4, 0, 1.0)], (2, 2));
        let err = bilinear_sparse(&sparse, (2, 2)).unwrap_err().to_string();
        assert_eq!(err, "bilinear resize: coordinate outside declared shape");
    }

    // -- compress_sparse_by_color ------------------------------------------

    fn cells_of(spans: &[u32]) -> Vec<u32> {
        spans.chunks(3).flat_map(|s| s[1]..s[1] + s[2]).collect()
    }

    #[test]
    fn the_colour_scale_is_split_into_equal_bins() {
        // It scaled by color_count - 1, so the top bin held the exact maximum
        // alone where the contract promises N equal bins.
        let values: Vec<f32> = (0..=100).map(|v| v as f32).collect();
        let row = vec![0u32; 101];
        let col: Vec<u32> = (0..101).collect();
        let bins = compress_sparse_by_color(&values, &row, &col, 4).unwrap();

        let counts: Vec<usize> = bins.iter().map(|b| cells_of(b).len()).collect();
        assert!(
            counts.iter().max().unwrap() - counts.iter().min().unwrap() <= 1,
            "{counts:?}"
        );
        assert!(
            counts[3] > 1 && cells_of(&bins[3]).contains(&100),
            "{counts:?}"
        );
        let mut all: Vec<u32> = bins.iter().flat_map(|b| cells_of(b)).collect();
        all.sort_unstable();
        assert_eq!(all, (0..101).collect::<Vec<u32>>());
    }

    #[test]
    fn a_run_of_one_colour_is_a_single_span() {
        let values = vec![5.0f32; 200];
        let row = vec![0u32; 200];
        let col: Vec<u32> = (0..200).collect();
        let bins = compress_sparse_by_color(&values, &row, &col, 4).unwrap();
        let spans: usize = bins.iter().map(|b| b.len() / 3).sum();
        assert_eq!(spans, 1);
    }

    #[test]
    fn a_new_row_opens_a_span_rather_than_extending_the_last() {
        let values = vec![5.0f32; 4];
        let bins = compress_sparse_by_color(&values, &[0, 0, 1, 1], &[0, 1, 0, 1], 1).unwrap();
        assert_eq!(bins[0], [0, 0, 2, 1, 0, 2]);
    }

    #[test]
    fn non_finite_values_take_no_part_in_the_scale_or_the_output() {
        let values = [0.5f32, f32::INFINITY, 2.0, f32::NAN];
        let bins = compress_sparse_by_color(&values, &[0; 4], &[0, 1, 2, 3], 2).unwrap();
        let listed: Vec<u32> = bins.iter().flat_map(|b| cells_of(b)).collect();
        assert_eq!(listed.len(), 2, "{bins:?}");
        // The infinity is what makes this worth asserting: had it counted, the
        // scale would run to it and both values would fall in the first bin.
        // Ignored, the scale is [0, 2] and they land one per bin -- 2.0 by the
        // clamp that closes the top one, since floor(2/2*2) is the bin count.
        assert_eq!(cells_of(&bins[0]), [0]);
        assert_eq!(cells_of(&bins[1]), [2]);
    }

    #[test]
    fn a_matrix_with_no_scale_comes_back_as_empty_bins() {
        let bins = compress_sparse_by_color(&[0.0, -3.0], &[0, 0], &[0, 1], 3).unwrap();
        assert_eq!(bins.len(), 3);
        assert!(bins.iter().all(|b| b.is_empty()));
        assert_eq!(compress_sparse_by_color(&[], &[], &[], 2).unwrap().len(), 2);
    }

    #[test]
    fn a_zero_color_count_is_refused() {
        let err = compress_sparse_by_color(&[1.0], &[0], &[0], 0)
            .unwrap_err()
            .to_string();
        assert!(err.contains("must be positive"), "{err}");
        // The ragged case belongs to the caller; see the doc comment.
    }
}