From 1429e7a95c3562f9440a68807adbbad72edf3abe Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Thu, 15 Oct 2026 06:37:15 +0000 Subject: [PATCH] added disk-streaming consensus refit for large matrix --- src/cnmf/cnmf.py | 36 ++++-- src/cnmf/qc.py | 12 +- src/cnmf/streaming.py | 254 +++++++++++++++++++++++++++++++++++++ tests/test_consensus_qc.py | 17 ++- tests/test_streaming.py | 208 ++++++++++++++++++++++++++++++ 5 files changed, 514 insertions(+), 13 deletions(-) create mode 100644 src/cnmf/streaming.py create mode 100644 tests/test_streaming.py diff --git a/src/cnmf/cnmf.py b/src/cnmf/cnmf.py index b7b4347..66c6376 100755 --- a/src/cnmf/cnmf.py +++ b/src/cnmf/cnmf.py @@ -774,7 +774,7 @@ def combine_nmf(self, k, skip_missing_files=False, remove_individual_iterations= return combined_spectra - def refit_usage(self, X, spectra): + def refit_usage(self, X, spectra, streaming=False): """ Takes an input data matrix and a fixed spectra and uses NNLS to find the optimal usage matrix. Generic kwargs for NMF are loaded from self.paths['nmf_run_parameters']. @@ -796,14 +796,19 @@ def refit_usage(self, X, spectra): else: refit_nmf_kwargs.update(dict(n_components = spectra.shape[0], H = spectra, update_H = False)) - _, rf_usages = self._nmf(X, nmf_kwargs=refit_nmf_kwargs) + if not streaming: + _, rf_usages = self._nmf(X, nmf_kwargs=refit_nmf_kwargs) + else: + from cnmf.streaming import _refit_from_cnmf + values = X.values if isinstance(X, pd.DataFrame) else X + rf_usages = _refit_from_cnmf(self, values, spectra) if (type(X) is pd.DataFrame) and (type(spectra) is pd.DataFrame): rf_usages = pd.DataFrame(rf_usages, index=X.index, columns=spectra.index) return(rf_usages) - def refit_spectra(self, X, usage): + def refit_spectra(self, X, usage, streaming=False): """ Takes an input data matrix and a fixed usage matrix and uses NNLS to find the optimal spectra matrix. Generic kwargs for NMF are loaded from self.paths['nmf_run_parameters']. @@ -818,12 +823,20 @@ def refit_spectra(self, X, usage): usage : pandas.DataFrame or numpy.ndarray, cells X programs Non-negative spectra of expression programs """ + if streaming: + from cnmf.streaming import _refit_from_cnmf + values = X.values if isinstance(X, pd.DataFrame) else X + result = _refit_from_cnmf(self, values, usage, spectra=True) + if isinstance(X, pd.DataFrame) and isinstance(usage, pd.DataFrame): + result = pd.DataFrame(result, index=usage.columns, columns=X.columns) + return result return(self.refit_usage(X.T, usage.T).T) def consensus(self, k, density_threshold=0.5, local_neighborhood_size=0.30, show_clustering=True, build_ref=True, skip_density_and_return_after_stats=False, close_clustergram_fig=False, - refit_usage=True, normalize_tpm_spectra=False, norm_counts=None): + refit_usage=True, normalize_tpm_spectra=False, norm_counts=None, + streaming=False): """ Obtain consensus estimates of spectra and usages from a cNMF run and output a clustergram of the consensus matrix. Assumes prepare, factorize, and combine steps have already been run. @@ -917,7 +930,8 @@ def consensus(self, k, density_threshold=0.5, local_neighborhood_size=0.30, show median_spectra = (median_spectra.T/median_spectra.sum(1)).T # Obtain reconstructed count matrix by re-fitting usage and computing dot product: usage.dot(spectra) - rf_usages = self.refit_usage(norm_counts.X, median_spectra) + refit_options = dict(streaming=True) if streaming else {} + rf_usages = self.refit_usage(norm_counts.X, median_spectra, **refit_options) rf_usages = pd.DataFrame(rf_usages, index=norm_counts.obs.index, columns=median_spectra.index) if skip_density_and_return_after_stats: @@ -952,7 +966,9 @@ def consensus(self, k, density_threshold=0.5, local_neighborhood_size=0.30, show # with usages fixed and TPM as the input matrix tpm = sc.read(self.paths['tpm']) tpm_stats = load_df_from_npz(self.paths['tpm_stats']) - spectra_tpm = self.refit_spectra(tpm.X, norm_usages.astype(tpm.X.dtype)) + # Streaming keeps fitted factors in FP64; do not round them to X's storage dtype. + refit_dtype = np.float64 if streaming else tpm.X.dtype + spectra_tpm = self.refit_spectra(tpm.X, norm_usages.astype(refit_dtype), **refit_options) spectra_tpm = pd.DataFrame(spectra_tpm, index=rf_usages.columns, columns=tpm.var.index) if normalize_tpm_spectra: spectra_tpm = spectra_tpm.div(spectra_tpm.sum(axis=1), axis=0) * 1e6 @@ -974,7 +990,8 @@ def consensus(self, k, density_threshold=0.5, local_neighborhood_size=0.30, show spectra_tpm_rf = spectra_tpm.loc[:,hvgs] spectra_tpm_rf = spectra_tpm_rf.div(tpm_stats.loc[hvgs, '__std'], axis=1) - rf_usages = self.refit_usage(norm_tpm.X, spectra_tpm_rf.astype(norm_tpm.X.dtype)) + refit_dtype = np.float64 if streaming else norm_tpm.X.dtype + rf_usages = self.refit_usage(norm_tpm.X, spectra_tpm_rf.astype(refit_dtype), **refit_options) rf_usages = pd.DataFrame(rf_usages, index=norm_counts.obs.index, columns=spectra_tpm_rf.index) save_df_to_npz(median_spectra, self.paths['consensus_spectra']%(k, density_threshold_repl)) @@ -1099,6 +1116,7 @@ def consensus(self, k, density_threshold=0.5, local_neighborhood_size=0.30, show kmeans_cluster_labels=kmeans_cluster_labels, normalized_counts=norm_counts.X, refit_usages=density_filtered_refit_usages.values, + streaming=streaming, ) @@ -1280,6 +1298,7 @@ def main(): parser.add_argument('--local-neighborhood-size', type=float, help='[consensus] Fraction of the number of replicates to use as nearest neighbors for local density filtering', default=0.30) parser.add_argument('--show-clustering', dest='show_clustering', help='[consensus] Produce a clustergram figure summarizing the spectra clustering', action='store_true') parser.add_argument('--build-reference', dest='build_reference', help='[consensus] Generates a reference spectra for use in starCAT', action='store_true', default=True) + parser.add_argument('--streaming', action='store_true', help='[consensus, CPU CD] Stream refit calculations; input loading and QC are not yet streaming') parser.add_argument("--engine", type=str.lower, choices=["cpu", "gpu"], help="[factorize,consensus] NMF engine to use (default cpu)", default="cpu") @@ -1325,7 +1344,8 @@ def main(): for k in ks: merged_spectra = load_df_from_npz(cnmf_obj.paths['merged_spectra']%k) cnmf_obj.consensus(k, args.local_density_threshold, args.local_neighborhood_size, args.show_clustering, - args.build_reference, close_clustergram_fig=True) + args.build_reference, close_clustergram_fig=True, + streaming=args.streaming) elif args.command == 'k_selection_plot': cnmf_obj.k_selection_plot(close_fig=True) diff --git a/src/cnmf/qc.py b/src/cnmf/qc.py index 31ed4e5..02a3219 100644 --- a/src/cnmf/qc.py +++ b/src/cnmf/qc.py @@ -213,14 +213,20 @@ def run_qc( kmeans_cluster_labels, normalized_counts, refit_usages, + streaming=False, ): """Calculate and save QC using state already produced by ``consensus``.""" # local_density contains all spectra, while l2_spectra has already been density-filtered. nruns = len(local_density) nruns_per_gep = nruns/k - if sp.issparse(normalized_counts): - normalized_counts = np.asarray(normalized_counts.todense()) - run_r2, run_sse, run_tss = get_r2( + if streaming: + from cnmf.streaming import reconstruction_stats_streaming + reconstruction_stats = reconstruction_stats_streaming + else: + if sp.issparse(normalized_counts): + normalized_counts = np.asarray(normalized_counts.todense()) + reconstruction_stats = get_r2 + run_r2, run_sse, run_tss = reconstruction_stats( normalized_counts, refit_usages, median_spectra.values, diff --git a/src/cnmf/streaming.py b/src/cnmf/streaming.py new file mode 100644 index 0000000..df0429d --- /dev/null +++ b/src/cnmf/streaming.py @@ -0,0 +1,254 @@ +"""FP64 CPU CD streaming refits and an opt-in cNMF adapter. + +Inputs must already have the intended expression normalization, gene order and +scaling. Usages must have the same cell order as expression. No implicit TPM, +usage normalization, gene selection, or scaling is performed here. +""" +from contextlib import contextmanager +from concurrent.futures import ThreadPoolExecutor +from itertools import islice +from pathlib import Path +from tempfile import TemporaryDirectory +import warnings + +import anndata as ad +import h5py +import numpy as np +from scipy import sparse +from sklearn.decomposition import non_negative_factorization +from sklearn.exceptions import ConvergenceWarning +from threadpoolctl import threadpool_limits + +# Working-buffer budget, not a whole-process RAM limit. Leave 25% headroom. +MEMORY_BUDGET = 16 * 1024**3 +DEFAULT_THREADS = 4 +DEFAULT_WORKERS = 4 + + +def _chunk_size(shape, k, fixed_bytes=0): + """Budget FP64 blocks, residuals, solver copies and fixed factors.""" + available = int(MEMORY_BUDGET * .75) - fixed_bytes + bytes_per_row = DEFAULT_WORKERS * 8 * (6 * shape[1] + 4 * k) + if available < bytes_per_row: + raise ValueError('Streaming working-memory budget cannot fit one row') + return min(max(1, (shape[0] + DEFAULT_WORKERS - 1) // DEFAULT_WORKERS), + available // bytes_per_row) + + +@contextmanager +def _expression(path): + """Open only X, not AnnData metadata, layers or embeddings.""" + if not isinstance(path, (str, Path)): + yield path + return + path = Path(path) + if path.suffix == '.npy': + x = np.load(path, mmap_mode='r') + yield x + else: + with h5py.File(path, 'r') as handle: + x = handle['X'] + if isinstance(x, h5py.Group): + x = ad.io.sparse_dataset(x) + if x.format != 'csr': + raise ValueError('Streaming currently requires CSR or dense X, not CSC') + yield x + + +def _checked(value): + values = value.data if sparse.issparse(value) else value + if not np.isfinite(values).all() or (values < 0).any(): + raise ValueError('Refit inputs must be finite and nonnegative') + return value + + +def _blocks(x, chunk_size): + """Promote only the selected expression block, before products or reductions.""" + for start in range(0, x.shape[0], chunk_size): + stop = min(start + chunk_size, x.shape[0]) + yield start, stop, _checked(x[start:stop].astype(np.float64, copy=False)) + + +def _validate(chunk_size, max_iter, tol, alpha, l1_ratio): + if not isinstance(chunk_size, int) or chunk_size < 1: + raise ValueError('chunk_size must be a positive integer') + if max_iter < 1 or tol < 0 or alpha < 0 or not 0 <= l1_ratio <= 1: + raise ValueError('Invalid iteration, tolerance or regularization settings') + + +def _map_blocks(function, x, chunk_size): + """Read on the caller thread; bound loaded chunks and preserve reduction order.""" + blocks = _blocks(x, chunk_size) + with ThreadPoolExecutor(max_workers=DEFAULT_WORKERS) as pool: + while True: + batch = list(islice(blocks, DEFAULT_WORKERS)) + if not batch: + break + futures = [pool.submit(function, *block) for block in batch] + for future in futures: + yield future.result() + # Release inputs/results before reading the next batch (no eager prefetch). + del batch, futures, future + + +def _refit_from_cnmf(obj, expression, fixed_factor, spectra=False): + """Adapter for existing cNMF refits; no preprocessing or engine mutation.""" + import yaml + with open(obj.paths['nmf_run_parameters']) as stream: + options = yaml.safe_load(stream) + if options.get('solver', 'mu') != 'cd' or options.get('beta_loss', 'frobenius') != 'frobenius': + raise ValueError('Streaming refits require CD with Frobenius loss') + if options.get('shuffle', False): + raise ValueError('Streaming refits currently require shuffle=False') + fixed_factor = np.asarray(fixed_factor) + function = refit_spectra_streaming if spectra else refit_usage_streaming + with TemporaryDirectory(prefix='cnmf-streaming-', dir=Path(obj.paths['nmf_run_parameters']).parent) as folder: + output = Path(folder) / 'factor.npy' + function(expression, fixed_factor, output, + max_iter=options.get('max_iter', 1000), tol=options.get('tol', 1e-4), + alpha=options.get('alpha_W', 0.0), l1_ratio=options.get('l1_ratio', 0.0)) + return np.load(output) + + +@threadpool_limits.wrap(limits=DEFAULT_THREADS, user_api='blas') +def refit_usage_streaming(expression_path, spectra, output_path, chunk_size=None, + max_iter=1000, tol=1e-4, alpha=0.0, l1_ratio=0.0): + """ + Fit full-data usages with fixed spectra using CPU CD and save an NPY file. + + spectra : array + Fixed K-by-genes matrix, in the input expression's gene order. + alpha : float + sklearn alpha_W: penalty scales with the number of expression genes. + + Returns a receipt with chunk iteration counts. Output is published only + after every chunk succeeds; existing output is never overwritten. + """ + h = _checked(np.array(spectra, dtype=np.float64, order='C', copy=True)) + output_path = Path(output_path) + if output_path.exists(): + raise FileExistsError(output_path) + with _expression(expression_path) as x: + if len(x.shape) != 2 or min(x.shape) < 1 or h.ndim != 2 or h.shape[0] < 1 or h.shape[1] != x.shape[1]: + raise ValueError('Expression and spectra shapes do not match') + if chunk_size is None: + chunk_size = _chunk_size(x.shape, h.shape[0], + DEFAULT_WORKERS * (4*h.nbytes + 8*h.shape[0]**2)) + _validate(chunk_size, max_iter, tol, alpha, l1_ratio) + output_path.parent.mkdir(parents=True, exist_ok=True) + iterations = [] + with TemporaryDirectory(prefix='.refit-', dir=output_path.parent) as folder: + temporary = Path(folder) / 'usages.npy' + w = np.lib.format.open_memmap(temporary, mode='w+', dtype=np.float64, + shape=(x.shape[0], h.shape[0])) + def fit_block(start, stop, block): + fitted, _, n_iter = non_negative_factorization( + block, H=h, update_H=False, n_components=h.shape[0], + solver='cd', beta_loss='frobenius', max_iter=max_iter, tol=tol, + alpha_W=alpha, alpha_H=0.0, l1_ratio=l1_ratio, + shuffle=False, random_state=0) + w[start:stop] = fitted + return int(n_iter) + iterations.extend(_map_blocks(fit_block, x, chunk_size)) + w.flush() + del w + temporary.replace(output_path) + return dict(output=str(output_path), iterations=iterations, dtype='float64') + + +@threadpool_limits.wrap(limits=DEFAULT_THREADS, user_api='blas') +def refit_spectra_streaming(expression_path, usages, output_path, chunk_size=None, + max_iter=1000, tol=1e-4, alpha=0.0, l1_ratio=0.0): + """ + Fit full-gene spectra with W fixed, keeping W in RAM and streaming X. + + usages : array + Cells-by-K matrix in expression row order. Apply cNMF's usage + normalization before calling when reproducing its TPM-spectrum refit. + alpha : float + Equivalent to sklearn alpha_W on the transposed fixed-H problem + refit_usage(X.T, usages.T); penalty scales with the full cell count. + + Uses sklearn's private CD sweep kernel; parity tests guard compatibility. + Saves K-by-genes FP64 NPY, returning convergence and direct streamed SSE. + """ + from sklearn.decomposition._cdnmf_fast import _update_cdnmf_fast + + w = _checked(np.asarray(usages, dtype=np.float64, order='C')) + output_path = Path(output_path) + if output_path.exists(): + raise FileExistsError(output_path) + with _expression(expression_path) as x: + if len(x.shape) != 2 or min(x.shape) < 1 or w.ndim != 2 or w.shape[1] < 1 or w.shape[0] != x.shape[0]: + raise ValueError('Expression and usages shapes do not match') + if chunk_size is None: + chunk_size = _chunk_size(x.shape, w.shape[1], + w.nbytes + (32 + 8*DEFAULT_WORKERS)*x.shape[1]*w.shape[1] + + 8*w.shape[1]**2) + _validate(chunk_size, max_iter, tol, alpha, l1_ratio) + a = np.ascontiguousarray(w.T @ w) + # Transposed layout matches sklearn's fixed-H refit on X.T. + b = np.zeros((x.shape[1], w.shape[1]), dtype=np.float64) + def cross_product(start, stop, block): + return block.T @ w[start:stop] + for partial in _map_blocks(cross_product, x, chunk_size): + b += partial + a.flat[::w.shape[1]+1] += x.shape[0] * alpha * (1-l1_ratio) + b -= x.shape[0] * alpha * l1_ratio + ht = np.zeros_like(b) + order = np.arange(w.shape[1], dtype=np.intp) + converged = False + for n_iter in range(1, max_iter+1): + violation = _update_cdnmf_fast(ht, a, b, order) + if n_iter == 1: + initial_violation = violation + if initial_violation == 0 or violation / initial_violation <= tol: + converged = True + break + if not converged and tol > 0: + warnings.warn('Streaming spectra refit reached max_iter', ConvergenceWarning) + h = ht.T + # Direct residual evaluation avoids cancellation in the statistics formula. + def squared_error(start, stop, block): + residual = w[start:stop] @ h + if sparse.issparse(block): + block = block.toarray() + residual -= block + return float(np.square(residual).sum(dtype=np.float64)) + sse = sum(_map_blocks(squared_error, x, chunk_size)) + output_path.parent.mkdir(parents=True, exist_ok=True) + with TemporaryDirectory(prefix='.refit-', dir=output_path.parent) as folder: + temporary = Path(folder) / 'spectra.npy' + np.save(temporary, h) + temporary.replace(output_path) + return dict(output=str(output_path), iterations=n_iter, converged=converged, + SSE=sse, dtype='float64') + + +@threadpool_limits.wrap(limits=DEFAULT_THREADS, user_api='blas') +def reconstruction_stats_streaming(expression, usages, spectra): + """Original global-mean R²/SSE/TSS, with two FP64 row-streamed passes.""" + h = np.asarray(spectra, dtype=np.float64) + with _expression(expression) as x: + rows = _chunk_size(x.shape, h.shape[0], h.nbytes + usages.nbytes) + def block_sum(start, stop, block): + return float(block.sum(dtype=np.float64)) + total = sum(_map_blocks(block_sum, x, rows)) + mean = total / (x.shape[0] * x.shape[1]) + def block_errors(start, stop, block): + if sparse.issparse(block): + block = block.toarray() + residual = np.asarray(usages[start:stop], dtype=np.float64) @ h + residual -= block + np.square(residual, out=residual) + block_sse = float(residual.sum(dtype=np.float64)) + centered = block - mean + np.square(centered, out=centered) + return block_sse, float(centered.sum(dtype=np.float64)) + sse = tss = 0.0 + for block_sse, block_tss in _map_blocks(block_errors, x, rows): + sse += block_sse + tss += block_tss + with np.errstate(divide='ignore', invalid='ignore'): + r2 = 1 - np.divide(sse, tss) + return r2, sse, tss diff --git a/tests/test_consensus_qc.py b/tests/test_consensus_qc.py index b4c2692..3bcba60 100644 --- a/tests/test_consensus_qc.py +++ b/tests/test_consensus_qc.py @@ -1,5 +1,7 @@ import numpy as np import pandas as pd +import pytest +from scipy import sparse from cnmf.qc import edist, run_qc @@ -25,7 +27,9 @@ def test_edistance_matches_squared_distance_definition(): assert actual.columns.tolist() == [1, 2] -def test_run_qc_matches_run_qc_cnmf_reference_output(tmp_path): +@pytest.mark.parametrize('streaming', [False, True]) +@pytest.mark.parametrize('sparse_input', [False, True]) +def test_run_qc_matches_run_qc_cnmf_reference_output(tmp_path, monkeypatch, streaming, sparse_input): """Match the non-LOO QC output produced by sc-blipper's run_qc_cnmf.py.""" # Reference source: ../sc-blipper/bin/run_qc_cnmf.py. This synthetic fixture # reproduces the state immediately after that script's clustering and fixed-H @@ -74,6 +78,14 @@ def test_run_qc_matches_run_qc_cnmf_reference_output(tmp_path): output = str(tmp_path / "reference") prefix = "k_3.dt_0_5" + import cnmf.qc as qc + import cnmf.streaming as streaming_module + def wrong_path(*args, **kwargs): + raise AssertionError('Wrong QC reconstruction path') + if streaming: + monkeypatch.setattr(qc, 'get_r2', wrong_path) + else: + monkeypatch.setattr(streaming_module, 'reconstruction_stats_streaming', wrong_path) run_qc( output=output, prefix=prefix, @@ -83,8 +95,9 @@ def test_run_qc_matches_run_qc_cnmf_reference_output(tmp_path): l2_spectra=l2_spectra, local_density=local_density, kmeans_cluster_labels=kmeans_cluster_labels, - normalized_counts=normalized_counts, + normalized_counts=sparse.csr_matrix(normalized_counts) if sparse_input else normalized_counts, refit_usages=refit_usages, + streaming=streaming, ) annotation_path = tmp_path / f"reference.{prefix}.annotation.tsv" diff --git a/tests/test_streaming.py b/tests/test_streaming.py new file mode 100644 index 0000000..4990ce8 --- /dev/null +++ b/tests/test_streaming.py @@ -0,0 +1,208 @@ +"""Deterministic synthetic parity against sklearn fixed-factor CD, no downloads.""" +import anndata as ad +import numpy as np +import pytest +from scipy import sparse +from sklearn.decomposition import non_negative_factorization +from cnmf.streaming import refit_usage_streaming, refit_spectra_streaming + + +@pytest.mark.parametrize('storage', ['dense', 'csr', 'backed']) +@pytest.mark.parametrize('dtype', [np.float32, np.float64]) +def test_streamed_qc(tmp_path, monkeypatch, storage, dtype): + import cnmf.streaming as streaming + from cnmf.qc import get_r2 + rng = np.random.default_rng(9) + x = rng.random((37, 19)).astype(dtype) + w = rng.random((37, 3)).astype(dtype) + h = rng.random((3, 19)).astype(dtype) + expected = get_r2(x.astype(float), w.astype(float), h.astype(float)) + source = sparse.csr_matrix(x) if storage == 'csr' else x + if storage == 'backed': + source = tmp_path / 'qc.h5ad' + ad.AnnData(sparse.csr_matrix(x)).write_h5ad(source) + monkeypatch.setattr(streaming, '_chunk_size', lambda *args: 5) + actual = streaming.reconstruction_stats_streaming(source, w, h) + np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12) + + +def test_auto_chunk_budget(): + from cnmf.streaming import _chunk_size, MEMORY_BUDGET, DEFAULT_WORKERS + rows = _chunk_size((8_000_000, 20_000), 20, 1_280_000_000) + assert 1 < rows < 8_000_000 + assert DEFAULT_WORKERS * rows * 8 * (6*20_000 + 4*20) + 1_280_000_000 <= MEMORY_BUDGET * .75 + assert _chunk_size((10, 20), 3) == 3 + with pytest.raises(ValueError, match='one row'): + _chunk_size((10, 20), 3, MEMORY_BUDGET) + + +def test_default_threads(tmp_path, monkeypatch): + import cnmf.streaming as streaming + from threadpoolctl import threadpool_info + original = streaming.non_negative_factorization + observed = [] + def record(*args, **kwargs): + observed.extend(p['num_threads'] for p in threadpool_info() if p['user_api'] == 'blas') + return original(*args, **kwargs) + monkeypatch.setattr(streaming, 'non_negative_factorization', record) + refit_usage_streaming(np.ones((10, 5)), np.ones((2, 5)), tmp_path / 'w.npy') + assert observed and set(observed) == {4} + + +def test_four_chunk_workers_bounded_and_ordered(monkeypatch): + import threading + import cnmf.streaming as streaming + barrier = threading.Barrier(4) + thread_ids = set() + reads = [] + reader_thread = threading.get_ident() + original = streaming._blocks + def blocks(x, rows): + for item in original(x, rows): + assert threading.get_ident() == reader_thread + reads.append(item[0]) + yield item + def work(start, stop, block): + thread_ids.add(threading.get_ident()) + barrier.wait(timeout=10) + return start + monkeypatch.setattr(streaming, '_blocks', blocks) + results = streaming._map_blocks(work, np.ones((8, 2)), 1) + assert next(results) == 0 + assert reads == [0, 1, 2, 3] + assert list(results) == list(range(1, 8)) + assert len(thread_ids) == 4 + + +def test_worker_count_parity(tmp_path, monkeypatch): + import cnmf.streaming as streaming + rng = np.random.default_rng(18) + x, w, h = rng.random((41, 17)), rng.random((41, 3)), rng.random((3, 17)) + results = [] + for workers in [1, 4]: + monkeypatch.setattr(streaming, 'DEFAULT_WORKERS', workers) + wp, hp = tmp_path / ('w%d.npy' % workers), tmp_path / ('h%d.npy' % workers) + refit_usage_streaming(x, h, wp, chunk_size=5) + refit_spectra_streaming(x, w, hp, chunk_size=5) + monkeypatch.setattr(streaming, '_chunk_size', lambda *args: 5) + results.append((np.load(wp), np.load(hp), streaming.reconstruction_stats_streaming(x, w, h))) + for single, parallel in zip(*results): + np.testing.assert_array_equal(single, parallel) + + +@pytest.mark.parametrize('storage', ['npy', 'dense', 'csr']) +@pytest.mark.parametrize('chunk_size', [7, 25, 200]) +@pytest.mark.parametrize('alpha', [0., .01]) +@pytest.mark.parametrize('input_dtype', [np.float32, np.float64]) +def test_refits(tmp_path, storage, chunk_size, alpha, input_dtype): + rng = np.random.default_rng(42) + x = rng.gamma(2, 1, (83, 31)).astype(input_dtype) + w = rng.random((83, 3)).astype(input_dtype) + h = rng.random((3, 31)).astype(input_dtype) + x64, w64, h64 = [a.astype(np.float64) for a in (x, w, h)] + path = tmp_path / ('x.npy' if storage == 'npy' else 'x.h5ad') + if storage == 'npy': + np.save(path, x) + else: + ad.AnnData(sparse.csr_matrix(x) if storage == 'csr' else x).write_h5ad(path) + kwargs = dict(solver='cd', update_H=False, max_iter=1000, tol=1e-10, + alpha_W=alpha, alpha_H=0., l1_ratio=.4, shuffle=False) + expected_w, _, _ = non_negative_factorization(x64, H=h64, n_components=3, **kwargs) + out = tmp_path / 'w.npy' + refit_usage_streaming(path, h, out, chunk_size, tol=1e-10, alpha=alpha, l1_ratio=.4) + np.testing.assert_allclose(np.load(out), expected_w, atol=1e-8, rtol=1e-7) + assert np.load(out).dtype == np.float64 + expected_ht, _, _ = non_negative_factorization(x64.T, H=w64.T, n_components=3, **kwargs) + out_h = tmp_path / 'h.npy' + receipt = refit_spectra_streaming(path, w, out_h, chunk_size, tol=1e-10, + alpha=alpha, l1_ratio=.4) + actual = np.load(out_h) + np.testing.assert_allclose(actual, expected_ht.T, atol=1e-10, rtol=1e-8) + assert actual.dtype == np.float64 + np.testing.assert_allclose(receipt['SSE'], np.square(x64-w64@actual).sum(), rtol=1e-12) + if storage == 'npy': + np.testing.assert_array_equal(np.load(path), x) + assert np.load(path, mmap_mode='r').dtype == input_dtype + with pytest.raises(FileExistsError): + refit_spectra_streaming(path, w, out_h) + + +def test_zero_program_and_invalid_input(tmp_path): + x = np.ones((20, 10)) + w = np.ones((20, 2)) + w[:, 1] = 0 + path = tmp_path / 'x.npy' + np.save(path, x) + out = tmp_path / 'h.npy' + refit_spectra_streaming(path, w, out, chunk_size=3) + assert np.all(np.load(out)[1] == 0) + x[-1, -1] = np.nan + np.save(path, x) + with pytest.raises(ValueError, match='finite'): + refit_usage_streaming(path, np.ones((2, 10)), tmp_path / 'bad.npy', chunk_size=3) + assert not (tmp_path / 'bad.npy').exists() + assert not list(tmp_path.glob('.refit-*')) + + +def test_csc_rejected(tmp_path): + path = tmp_path / 'csc.h5ad' + ad.AnnData(sparse.csc_matrix(np.ones((10, 5)))).write_h5ad(path) + with pytest.raises(ValueError, match='CSC'): + refit_usage_streaming(path, np.ones((2, 5)), tmp_path / 'out.npy') + + +def test_cnmf_consensus_streaming_parity(tmp_path, monkeypatch): + import sys + import pandas as pd + import yaml + from cnmf.cnmf import cNMF, main, load_df_from_npz, save_df_to_npz + rng = np.random.default_rng(73) + counts = tmp_path / 'counts.h5ad' + ad.AnnData(rng.poisson(8, (60, 12)).astype(float)).write_h5ad(counts) + prototypes = rng.random((2, 6)) + .1 + objects = [] + for name in ['original', 'streamed']: + obj = cNMF(str(tmp_path), name) + obj.prepare(str(counts), components=[2], n_iter=3, densify=True, + seed=14, num_highvar_genes=6, max_NMF_iter=1000, + beta_loss='frobenius') + with open(obj.paths['nmf_run_parameters']) as f: + options = yaml.safe_load(f) + options.update(solver='cd', tol=1e-10) + with open(obj.paths['nmf_run_parameters'], 'w') as f: + yaml.safe_dump(options, f) + genes = ad.read_h5ad(obj.paths['normalized_counts']).var_names + merged = pd.DataFrame(np.vstack([prototypes*.98, prototypes, prototypes*1.02]), + index=['iter%d_topic%d' % (i,j) for i in range(3) for j in (1,2)], + columns=genes) + save_df_to_npz(merged, obj.paths['merged_spectra'] % 2) + objects.append(obj) + import cnmf.cnmf as cnmf_module + original_qc = cnmf_module.run_qc + qc_flags = [] + def record_qc(*args, **kwargs): + qc_flags.append(kwargs.get('streaming', False)) + return original_qc(*args, **kwargs) + monkeypatch.setattr(cnmf_module, 'run_qc', record_qc) + objects[0].consensus(2, density_threshold=2., local_neighborhood_size=.5, + show_clustering=False) + import cnmf.streaming as streaming + monkeypatch.setattr(streaming, '_chunk_size', lambda *args: 7) + original = streaming._refit_from_cnmf + calls = [] + def record(*args, **kwargs): + calls.append(kwargs.get('spectra', False)) + return original(*args, **kwargs) + monkeypatch.setattr(streaming, '_refit_from_cnmf', record) + monkeypatch.setattr(sys, 'argv', [ + 'cnmf', 'consensus', '--output-dir', str(tmp_path), '--name', 'streamed', + '-k', '2', '--local-density-threshold', '2.0', + '--local-neighborhood-size', '.5', '--streaming']) + main() + assert calls == [False, True, False] + assert qc_flags == [False, True] + for key in ['consensus_spectra', 'consensus_usages', 'gene_spectra_tpm', + 'gene_spectra_score', 'starcat_spectra']: + a, b = [load_df_from_npz(o.paths[key] % (2, '2_0')) for o in objects] + assert a.index.equals(b.index) and a.columns.equals(b.columns) + np.testing.assert_allclose(a, b, rtol=1e-6, atol=1e-7)