Skip to content

Commit a40dabc

Browse files
sebastianpinedasebastianpineda
committed
Backed io cleanup
1 parent 50b1daa commit a40dabc

4 files changed

Lines changed: 413 additions & 23 deletions

File tree

‎src/actionet/bindings/wp_io.cpp‎

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,28 @@ namespace {
3636
const auto* ptr = static_cast<const double*>(buf.ptr);
3737
return std::vector<double>(ptr, ptr + static_cast<size_t>(buf.shape[0]));
3838
}
39+
40+
// Copy a 1-D int64 numpy array into an arma::uvec. Consolidates the
41+
// ``col_indices``/``row_indices`` copy loops that were previously
42+
// duplicated inline in ``backed_take_columns``.
43+
arma::uvec int64_array_to_uvec(const py::array_t<int64_t>& arr) {
44+
auto buf = arr.request();
45+
if (buf.ndim != 1) {
46+
throw std::runtime_error(
47+
"expected a 1D int64 array for backed index conversion");
48+
}
49+
const auto n = static_cast<arma::uword>(buf.size);
50+
arma::uvec out(n);
51+
const auto* ptr = static_cast<const int64_t*>(buf.ptr);
52+
for (size_t i = 0; i < static_cast<size_t>(n); ++i) {
53+
if (ptr[i] < 0) {
54+
throw std::runtime_error(
55+
"backed index arrays must be non-negative");
56+
}
57+
out(i) = static_cast<arma::uword>(ptr[i]);
58+
}
59+
return out;
60+
}
3961
} // namespace
4062

4163
std::shared_ptr<actionet::MatrixOperator> create_backed_operator(
@@ -185,24 +207,12 @@ void init_io(py::module_ &m) {
185207
throw std::runtime_error("backed_take_columns: operator is null");
186208
}
187209

188-
// Convert col_indices.
189-
auto col_buf = col_indices_arr.request();
190-
arma::uvec col_indices(static_cast<arma::uword>(col_buf.size));
191-
auto* col_ptr = static_cast<int64_t*>(col_buf.ptr);
192-
for (size_t i = 0; i < static_cast<size_t>(col_buf.size); ++i) {
193-
col_indices(i) = static_cast<arma::uword>(col_ptr[i]);
194-
}
210+
arma::uvec col_indices = int64_array_to_uvec(col_indices_arr);
195211

196-
// Convert optional row_indices.
197212
arma::uvec row_indices;
198213
if (!row_indices_obj.is_none()) {
199214
py::array_t<int64_t> row_arr = row_indices_obj.cast<py::array_t<int64_t>>();
200-
auto row_buf = row_arr.request();
201-
row_indices.set_size(static_cast<arma::uword>(row_buf.size));
202-
auto* row_ptr = static_cast<int64_t*>(row_buf.ptr);
203-
for (size_t i = 0; i < static_cast<size_t>(row_buf.size); ++i) {
204-
row_indices(i) = static_cast<arma::uword>(row_ptr[i]);
205-
}
215+
row_indices = int64_array_to_uvec(row_arr);
206216
}
207217

208218
if (prefer_sparse) {

‎src/libactionet‎

‎tests/backed/test_backed_operator_parity.py‎

Lines changed: 336 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,336 @@
1+
"""Low-level parity tests for backed matrix operators.
2+
3+
Focused regressions for the paths the io/backed_h5ad refactor touches:
4+
- BackedSparseMatrixOperator::takeColumnsDense / takeColumnsSparse
5+
(CSR + CSC) — including duplicate column indices, arbitrary row_indices
6+
reorderings, and combined transforms (row_scale + log1p + log_scale).
7+
- BackedDenseMatrixOperator::takeColumnsDense / takeColumnsSparse
8+
(dense path via the same public binding) — same matrix.
9+
- Matvec/rmatvec/matmat/rmatmat semantics — exercised indirectly by
10+
``run_svd_backed_operator`` versus in-memory ``run_svd``.
11+
12+
The `apply_log1p` path uses the Paul Mineiro `fastlog()` approximation,
13+
so transforms involving log1p are compared with a looser (`rtol=5e-3`)
14+
tolerance per the operator header documentation.
15+
"""
16+
17+
from __future__ import annotations
18+
19+
import numpy as np
20+
import pytest
21+
import scipy.sparse as sp
22+
import pandas as pd
23+
import anndata as ad
24+
25+
import actionet._core as _core
26+
27+
28+
# ---------------------------------------------------------------------------
29+
# Small synthetic fixtures — kept intentionally lean so CI cost is minimal.
30+
# ---------------------------------------------------------------------------
31+
32+
_SHAPES = [(64, 32), (256, 128)]
33+
_DENSITIES = [0.15]
34+
_LOG_TOL = dict(rtol=5e-3, atol=5e-6) # fastlog approximation regime
35+
_EXACT_TOL = dict(rtol=1e-10, atol=1e-12) # no-log1p paths
36+
37+
38+
def _make_matrix(shape, density, fmt, seed):
39+
rng = np.random.default_rng(seed)
40+
if fmt == "dense":
41+
X = rng.random(shape) * 5.0
42+
X[rng.random(shape) > density] = 0.0 # sparsify but keep dense storage
43+
return X
44+
S = sp.random(*shape, density=density, random_state=rng,
45+
format=fmt, dtype=np.float64)
46+
S.data = np.abs(S.data) * 5.0 # counts-like, positive
47+
S.eliminate_zeros()
48+
if fmt == "csr":
49+
return sp.csr_matrix(S)
50+
return sp.csc_matrix(S)
51+
52+
53+
def _write_h5ad(path, X, fmt):
54+
n_obs, n_var = X.shape
55+
obs = pd.DataFrame(index=[f"c{i}" for i in range(n_obs)])
56+
var = pd.DataFrame(index=[f"g{i}" for i in range(n_var)])
57+
if fmt == "dense":
58+
adata = ad.AnnData(X=np.asarray(X, dtype=np.float64), obs=obs, var=var)
59+
else:
60+
adata = ad.AnnData(X=X, obs=obs, var=var)
61+
adata.write_h5ad(path)
62+
return str(path)
63+
64+
65+
def _apply_transform_reference(X_orig, row_scale, apply_log1p, log_scale):
66+
"""Reference transform mirroring the operator's lazy transform.
67+
68+
NOTE: The C++ path uses ``fastlog(1 + x)`` (float precision). Callers
69+
that pass ``apply_log1p=True`` must compare with the loose tolerance.
70+
"""
71+
if sp.issparse(X_orig):
72+
X = X_orig.copy().astype(np.float64)
73+
if row_scale is not None:
74+
X = sp.diags(row_scale, format="csc") @ X
75+
if apply_log1p:
76+
data = np.log1p(X.data.astype(np.float32)).astype(np.float64)
77+
if log_scale != 1.0:
78+
data = data * log_scale
79+
X = X.copy()
80+
X.data = data
81+
return X
82+
else:
83+
X = np.asarray(X_orig, dtype=np.float64).copy()
84+
if row_scale is not None:
85+
X = X * row_scale[:, None]
86+
if apply_log1p:
87+
X = np.log1p(X.astype(np.float32)).astype(np.float64)
88+
if log_scale != 1.0:
89+
X = X * log_scale
90+
return X
91+
92+
93+
def _open_op(path, *, chunk_size, row_scale, apply_log1p, log_scale, n_threads=1):
94+
kwargs = dict(
95+
file_path=path,
96+
group_path="/X",
97+
chunk_size=chunk_size,
98+
apply_log1p=apply_log1p,
99+
log_scale=log_scale,
100+
n_threads=n_threads,
101+
)
102+
if row_scale is not None:
103+
kwargs["row_scale_factors"] = np.asarray(row_scale, dtype=np.float64)
104+
return _core.create_backed_operator(**kwargs)
105+
106+
107+
# ---------------------------------------------------------------------------
108+
# Parametrised test matrix
109+
# ---------------------------------------------------------------------------
110+
111+
_FMTS = ["csr", "csc", "dense"]
112+
113+
_TRANSFORMS = [
114+
# (id, row_scale_fn, apply_log1p, log_scale, tol)
115+
("none", None, False, 1.0, _EXACT_TOL),
116+
("row_scale", lambda n: 1.0 + np.arange(n) / n, False, 1.0, _EXACT_TOL),
117+
("log1p", None, True, 1.0, _LOG_TOL),
118+
("row_scale+log1p", lambda n: 1.0 + np.arange(n) / n, True, 1.0, _LOG_TOL),
119+
("row_scale+log1p+logscale",
120+
lambda n: 1.0 + np.arange(n) / n, True, 1.0 / np.log(2.0), _LOG_TOL),
121+
]
122+
123+
124+
@pytest.fixture(scope="module")
125+
def _matrices():
126+
"""Materialised (X_orig, X_ref_by_transform) per (fmt, shape)."""
127+
cache = {}
128+
for fmt in _FMTS:
129+
for shape in _SHAPES:
130+
X_orig = _make_matrix(shape, _DENSITIES[0], fmt, seed=hash((fmt, shape)) & 0xFFFF)
131+
per_transform = {}
132+
for tid, rs_fn, apply_log1p, log_scale, _tol in _TRANSFORMS:
133+
rs = rs_fn(shape[0]) if rs_fn is not None else None
134+
per_transform[tid] = _apply_transform_reference(
135+
X_orig, rs, apply_log1p, log_scale,
136+
)
137+
cache[(fmt, shape)] = (X_orig, per_transform)
138+
return cache
139+
140+
141+
def _params_id(fmt, shape, tid):
142+
return f"{fmt}-{shape[0]}x{shape[1]}-{tid}"
143+
144+
145+
_ALL_PARAMS = [
146+
(fmt, shape, tid, rs_fn, apply_log1p, log_scale, tol)
147+
for fmt in _FMTS
148+
for shape in _SHAPES
149+
for (tid, rs_fn, apply_log1p, log_scale, tol) in _TRANSFORMS
150+
]
151+
152+
153+
# ---------------------------------------------------------------------------
154+
# takeColumnsDense parity (unique cols, all rows)
155+
# ---------------------------------------------------------------------------
156+
157+
@pytest.mark.parametrize(
158+
"fmt,shape,tid,rs_fn,apply_log1p,log_scale,tol",
159+
_ALL_PARAMS,
160+
ids=[_params_id(f, s, t) for (f, s, t, *_ ) in _ALL_PARAMS],
161+
)
162+
def test_take_columns_dense_unique(tmp_path, _matrices, fmt, shape, tid, rs_fn,
163+
apply_log1p, log_scale, tol):
164+
X_orig, refs = _matrices[(fmt, shape)]
165+
X_ref = refs[tid]
166+
if sp.issparse(X_ref):
167+
X_ref = X_ref.toarray()
168+
else:
169+
X_ref = np.asarray(X_ref)
170+
171+
path = _write_h5ad(tmp_path / "m.h5ad", X_orig, fmt=fmt)
172+
op = _open_op(
173+
path, chunk_size=17,
174+
row_scale=(rs_fn(shape[0]) if rs_fn is not None else None),
175+
apply_log1p=apply_log1p, log_scale=log_scale,
176+
)
177+
cols = np.array([0, 5, 3, shape[1] - 1, shape[1] // 2], dtype=np.int64)
178+
cols = np.unique(cols)
179+
out = _core.backed_take_columns(op, cols, prefer_sparse=False)
180+
181+
np.testing.assert_allclose(out, X_ref[:, cols], **tol)
182+
183+
184+
# ---------------------------------------------------------------------------
185+
# takeColumnsDense parity with duplicate columns
186+
# ---------------------------------------------------------------------------
187+
188+
@pytest.mark.parametrize(
189+
"fmt,shape,tid,rs_fn,apply_log1p,log_scale,tol",
190+
_ALL_PARAMS,
191+
ids=[_params_id(f, s, t) for (f, s, t, *_ ) in _ALL_PARAMS],
192+
)
193+
def test_take_columns_dense_duplicates(tmp_path, _matrices, fmt, shape, tid, rs_fn,
194+
apply_log1p, log_scale, tol):
195+
X_orig, refs = _matrices[(fmt, shape)]
196+
X_ref = refs[tid]
197+
if sp.issparse(X_ref):
198+
X_ref = X_ref.toarray()
199+
else:
200+
X_ref = np.asarray(X_ref)
201+
202+
path = _write_h5ad(tmp_path / "m.h5ad", X_orig, fmt=fmt)
203+
op = _open_op(
204+
path, chunk_size=17,
205+
row_scale=(rs_fn(shape[0]) if rs_fn is not None else None),
206+
apply_log1p=apply_log1p, log_scale=log_scale,
207+
)
208+
cols = np.array([3, 3, 3, 7, 3, 7, 1], dtype=np.int64)
209+
out = _core.backed_take_columns(op, cols, prefer_sparse=False)
210+
211+
np.testing.assert_allclose(out, X_ref[:, cols], **tol)
212+
213+
214+
# ---------------------------------------------------------------------------
215+
# takeColumnsDense parity with row_indices (subset and reorder)
216+
# ---------------------------------------------------------------------------
217+
218+
@pytest.mark.parametrize(
219+
"fmt,shape,tid,rs_fn,apply_log1p,log_scale,tol",
220+
_ALL_PARAMS,
221+
ids=[_params_id(f, s, t) for (f, s, t, *_ ) in _ALL_PARAMS],
222+
)
223+
def test_take_columns_dense_row_indices(tmp_path, _matrices, fmt, shape, tid, rs_fn,
224+
apply_log1p, log_scale, tol):
225+
X_orig, refs = _matrices[(fmt, shape)]
226+
X_ref = refs[tid]
227+
if sp.issparse(X_ref):
228+
X_ref = X_ref.toarray()
229+
else:
230+
X_ref = np.asarray(X_ref)
231+
232+
path = _write_h5ad(tmp_path / "m.h5ad", X_orig, fmt=fmt)
233+
op = _open_op(
234+
path, chunk_size=17,
235+
row_scale=(rs_fn(shape[0]) if rs_fn is not None else None),
236+
apply_log1p=apply_log1p, log_scale=log_scale,
237+
)
238+
cols = np.array([0, 5, 3], dtype=np.int64)
239+
rows = np.array([shape[0] - 1, 0, 4, shape[0] // 2, 2, 1], dtype=np.int64)
240+
out = _core.backed_take_columns(op, cols, row_indices=rows, prefer_sparse=False)
241+
242+
np.testing.assert_allclose(out, X_ref[np.ix_(rows, cols)], **tol)
243+
244+
245+
# ---------------------------------------------------------------------------
246+
# takeColumnsSparse parity (unique + duplicate)
247+
# ---------------------------------------------------------------------------
248+
249+
@pytest.mark.parametrize(
250+
"fmt,shape,tid,rs_fn,apply_log1p,log_scale,tol",
251+
# dense-storage input still returns a sparse output through the binding.
252+
_ALL_PARAMS,
253+
ids=[_params_id(f, s, t) for (f, s, t, *_ ) in _ALL_PARAMS],
254+
)
255+
def test_take_columns_sparse_duplicates(tmp_path, _matrices, fmt, shape, tid, rs_fn,
256+
apply_log1p, log_scale, tol):
257+
X_orig, refs = _matrices[(fmt, shape)]
258+
X_ref_dense = refs[tid].toarray() if sp.issparse(refs[tid]) else np.asarray(refs[tid])
259+
260+
path = _write_h5ad(tmp_path / "m.h5ad", X_orig, fmt=fmt)
261+
op = _open_op(
262+
path, chunk_size=17,
263+
row_scale=(rs_fn(shape[0]) if rs_fn is not None else None),
264+
apply_log1p=apply_log1p, log_scale=log_scale,
265+
)
266+
cols = np.array([4, 4, 9, 4, 9, 0], dtype=np.int64)
267+
out = _core.backed_take_columns(op, cols, prefer_sparse=True)
268+
assert sp.issparse(out)
269+
np.testing.assert_allclose(out.toarray(), X_ref_dense[:, cols], **tol)
270+
271+
272+
# ---------------------------------------------------------------------------
273+
# SVD parity — exercises matvec/rmatvec/matmat/rmatmat end-to-end via PRIMME.
274+
# Compares singular values (sign/rotation invariant, so only sigma is checked).
275+
# ---------------------------------------------------------------------------
276+
277+
@pytest.mark.parametrize("fmt", ["csr", "csc", "dense"])
278+
def test_backed_svd_sigma_parity(tmp_path, fmt):
279+
shape = (128, 64)
280+
X_orig = _make_matrix(shape, 0.20, fmt, seed=13)
281+
path = _write_h5ad(tmp_path / "svd.h5ad", X_orig, fmt=fmt)
282+
283+
op = _core.create_backed_operator(
284+
file_path=path, group_path="/X", chunk_size=32, n_threads=1,
285+
)
286+
k = 10
287+
res_backed = _core.run_svd_backed_operator(
288+
op, k=k, max_it=200, seed=1, algorithm=1, verbose=False, # 1 = Halko
289+
)
290+
291+
if fmt == "dense":
292+
X_ref = np.asarray(X_orig)
293+
else:
294+
X_ref = X_orig.toarray()
295+
# Reference sigma from full SVD (only sigma is compared).
296+
u_ref, s_ref, vt_ref = np.linalg.svd(X_ref, full_matrices=False)
297+
298+
s_backed = np.sort(np.asarray(res_backed["d"]).reshape(-1))[::-1]
299+
s_ref_top = np.sort(s_ref)[::-1][:k]
300+
np.testing.assert_allclose(s_backed[:k], s_ref_top, rtol=1e-3, atol=1e-6)
301+
302+
303+
# ---------------------------------------------------------------------------
304+
# SVD parity with lazy transform (row_scale + log1p) — matvec path exercised
305+
# with the operator-side transform applied at read time.
306+
# ---------------------------------------------------------------------------
307+
308+
@pytest.mark.parametrize("fmt", ["csr", "csc"])
309+
def test_backed_svd_sigma_parity_with_transform(tmp_path, fmt):
310+
shape = (128, 64)
311+
X_orig = _make_matrix(shape, 0.20, fmt, seed=17)
312+
path = _write_h5ad(tmp_path / "svd_t.h5ad", X_orig, fmt=fmt)
313+
314+
row_scale = 1.0 + np.arange(shape[0]) / shape[0]
315+
log_scale = 1.0 / np.log(2.0)
316+
317+
op = _core.create_backed_operator(
318+
file_path=path, group_path="/X", chunk_size=32,
319+
row_scale_factors=row_scale, apply_log1p=True, log_scale=log_scale,
320+
n_threads=1,
321+
)
322+
k = 10
323+
res_backed = _core.run_svd_backed_operator(
324+
op, k=k, max_it=200, seed=1, algorithm=1, verbose=False, # 1 = Halko
325+
)
326+
327+
X_ref = _apply_transform_reference(
328+
X_orig, row_scale, apply_log1p=True, log_scale=log_scale,
329+
)
330+
X_ref = X_ref.toarray() if sp.issparse(X_ref) else np.asarray(X_ref)
331+
_, s_ref, _ = np.linalg.svd(X_ref, full_matrices=False)
332+
333+
s_backed = np.sort(np.asarray(res_backed["d"]).reshape(-1))[::-1]
334+
s_ref_top = np.sort(s_ref)[::-1][:k]
335+
# fastlog approximation: loosen sigma tolerance
336+
np.testing.assert_allclose(s_backed[:k], s_ref_top, rtol=1e-2, atol=1e-4)

0 commit comments

Comments
 (0)