Add streaming CPU CD refits and memory-bounded QC

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)

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,

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