sparse_npz 0.1.0

Reader and writer for SciPy sparse matrices saved in the NumPy .npz format (CSC and CSR).
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

use sparse_npz::{Compression, CscMatrix, Format, Values};

fn temp_path(tag: &str) -> std::path::PathBuf {
    let dir = std::env::temp_dir().join(format!("sparse_npz_test_{}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap();
    dir.join(format!("{tag}.npz"))
}

fn sample(values: Values) -> CscMatrix {
    // 3x4 matrix, column-major: col0 row0, col1 row2, col2 row1, col3 empty.
    CscMatrix {
        rows: 3,
        cols: 4,
        col_ptr: vec![0, 1, 2, 3, 3],
        row_indices: vec![0, 2, 1],
        values,
    }
}

#[test]
fn csc_write_then_read_roundtrips() {
    let path = temp_path("csc");
    let m = sample(Values::I64(vec![5, 3, 7]));
    m.write_npz(&path).unwrap();
    let back = CscMatrix::read_npz(&path).unwrap();
    assert_eq!(back, m);

    let mut got: Vec<(usize, usize, f64)> = back.entries().collect();
    got.sort_by_key(|&(r, c, _)| (c, r));
    assert_eq!(got, vec![(0, 0, 5.0), (2, 1, 3.0), (1, 2, 7.0)]);
}

#[test]
fn element_dtype_is_preserved_across_roundtrip() {
    let cases = [
        ("bool", Values::Bool(vec![true, true, true])),
        ("i8", Values::I8(vec![-5, 3, 7])),
        ("u8", Values::U8(vec![5, 3, 7])),
        ("i16", Values::I16(vec![-500, 300, 700])),
        ("u16", Values::U16(vec![500, 300, 700])),
        ("i32", Values::I32(vec![-5, 3, 7])),
        ("u32", Values::U32(vec![5, 3, 7])),
        ("i64", Values::I64(vec![-5, 3, 7])),
        ("u64", Values::U64(vec![5, 3, 7])),
        ("f32", Values::F32(vec![1.5, 2.5, -3.5])),
        ("f64", Values::F64(vec![1.5, 2.5, -3.5])),
    ];
    for (tag, values) in cases {
        let path = temp_path(tag);
        let m = sample(values);
        m.write_npz(&path).unwrap();
        let back = CscMatrix::read_npz(&path).unwrap();
        assert_eq!(back.values, m.values, "dtype {tag} not preserved");
    }
}

#[test]
fn csr_write_reports_csr_and_preserves_entries() {
    let path = temp_path("csr");
    let m = sample(Values::I64(vec![5, 3, 7]));
    m.write_npz_as(&path, Format::Csr).unwrap();

    let (back, format) = CscMatrix::read_npz_with_format(&path).unwrap();
    assert_eq!(format, Format::Csr);
    // Reading normalizes back to CSC, so the matrix must match the original.
    assert_eq!(back, m);
}

#[test]
fn csc_write_reports_csc() {
    let path = temp_path("csc_fmt");
    let m = sample(Values::I64(vec![5, 3, 7]));
    m.write_npz(&path).unwrap();
    let (_, format) = CscMatrix::read_npz_with_format(&path).unwrap();
    assert_eq!(format, Format::Csc);
}

#[test]
fn zstd_compressed_roundtrips_and_preserves_dtype() {
    let path = temp_path("zstd");
    let m = sample(Values::U32(vec![5, 3, 7]));
    m.write_npz_with(&path, Format::Csc, Compression::Zstd { level: None }).unwrap();
    // Reading auto-detects the Zstandard members.
    let back = CscMatrix::read_npz(&path).unwrap();
    assert_eq!(back, m);
}

#[test]
fn zstd_levels_roundtrip() {
    let m = sample(Values::I64(vec![5, 3, 7]));
    for level in [Some(-5), Some(1), Some(3), Some(19), Some(22), None] {
        let path = temp_path(&format!("zstd_lvl_{level:?}"));
        m.write_npz_with(&path, Format::Csc, Compression::Zstd { level }).unwrap();
        assert_eq!(CscMatrix::read_npz(&path).unwrap(), m);
    }
}

#[test]
fn deflate_levels_roundtrip() {
    let m = sample(Values::I64(vec![5, 3, 7]));
    for level in [Some(1), Some(6), Some(9), Some(10), Some(264), None] {
        let path = temp_path(&format!("deflate_lvl_{level:?}"));
        m.write_npz_with(&path, Format::Csc, Compression::Deflate { level }).unwrap();
        assert_eq!(CscMatrix::read_npz(&path).unwrap(), m);
    }
}

#[test]
fn zstd_and_deflate_read_identically() {
    let m = sample(Values::I64(vec![5, 3, 7]));
    let deflate = temp_path("cmp_deflate");
    let zstd = temp_path("cmp_zstd");
    m.write_npz_with(&deflate, Format::Csr, Compression::Deflate { level: Some(9) }).unwrap();
    m.write_npz_with(&zstd, Format::Csr, Compression::Zstd { level: Some(19) }).unwrap();
    let (a, fa) = CscMatrix::read_npz_with_format(&deflate).unwrap();
    let (b, fb) = CscMatrix::read_npz_with_format(&zstd).unwrap();
    assert_eq!(a, b);
    assert_eq!(fa, Format::Csr);
    assert_eq!(fb, Format::Csr);
}

#[test]
fn nnz_and_entries_agree() {
    let m = sample(Values::I64(vec![5, 3, 7]));
    assert_eq!(m.nnz(), 3);
    assert_eq!(m.entries().count(), 3);
}