diff --git a/pertpy/tools/_augur.py b/pertpy/tools/_augur.py index 9f1573fd..4f9760b8 100644 --- a/pertpy/tools/_augur.py +++ b/pertpy/tools/_augur.py @@ -518,7 +518,13 @@ def run_cross_validation( categorical=True, random_state=42) >>> results = ag_rfc.run_cross_validation(subsample=subsample, folds=3, subsample_idx=0, random_state=42, zero_division=0) """ - x = subsample.to_df() + # Pass the dense matrix instead of subsample.to_df(); its arrow-backed string columns make scikit-learn re-validate dtypes on every fold and scorer, while the values (and results) stay identical. + x = subsample.X + if sparse.issparse(x): + x = x.toarray() + x = np.asarray(x) + genes = subsample.var_names.tolist() + n_genes = len(genes) y = subsample.obs["y_"] scorer = self.set_scorer(multiclass=len(y.unique()) > 2, zero_division=zero_division) folds = StratifiedKFold(n_splits=folds, random_state=random_state, shuffle=True) @@ -540,21 +546,21 @@ def run_cross_validation( feature_importances = defaultdict(list) if isinstance(self.estimator, RandomForestClassifier | RandomForestRegressor): for fold, estimator in list(zip(range(len(results["estimator"])), results["estimator"], strict=False)): - feature_importances["genes"].extend(x.columns.tolist()) + feature_importances["genes"].extend(genes) feature_importances["feature_importances"].extend(estimator.feature_importances_.tolist()) - feature_importances["subsample_idx"].extend(len(x.columns) * [subsample_idx]) - feature_importances["fold"].extend(len(x.columns) * [fold]) + feature_importances["subsample_idx"].extend(n_genes * [subsample_idx]) + feature_importances["fold"].extend(n_genes * [fold]) # standardized coefficients with Agresti method # cf. https://think-lab.github.io/d/205/#3 if isinstance(self.estimator, LogisticRegression): for fold, estimator in list(zip(range(len(results["estimator"])), results["estimator"], strict=False)): - feature_importances["genes"].extend(x.columns.tolist()) + feature_importances["genes"].extend(genes) feature_importances["feature_importances"].extend( (estimator.coef_ * estimator.coef_.std()).flatten().tolist() ) - feature_importances["subsample_idx"].extend(len(x.columns) * [subsample_idx]) - feature_importances["fold"].extend(len(x.columns) * [fold]) + feature_importances["subsample_idx"].extend(n_genes * [subsample_idx]) + feature_importances["fold"].extend(n_genes * [fold]) results["feature_importances"] = feature_importances @@ -795,7 +801,9 @@ def predict( logger.warning("Set smaller span value in the case of a `segmentation fault` error.") logger.warning("Set larger span in case of svddc or other near singularities error.") adata.obs["augur_score"] = nan - for cell_type in track(adata.obs["cell_type"].unique(), description="Processing data..."): + + eligible_cell_types = [] + for cell_type in adata.obs["cell_type"].unique(): cell_type_subsample = adata[adata.obs["cell_type"] == cell_type].copy() if augur_mode in ("default", "permute") and len(cell_type_subsample) >= min_cells: @@ -825,42 +833,51 @@ def predict( f"subsample size {subsample_size}." ) else: - results[cell_type] = Parallel(n_jobs=n_threads)( - delayed(self.cross_validate_subsample)( - adata=cell_type_subsample, - augur_mode=augur_mode, - subsample_size=subsample_size, - folds=folds, - feature_perc=feature_perc, - subsample_idx=i, - random_state=random_state, - zero_division=zero_division, - ) - for i in range(n_subsamples) - ) - # summarize scores for cell type - results["summary_metrics"][cell_type] = self.average_metrics(results[cell_type]) - - # add scores as observation to anndata - mask = adata.obs["cell_type"].str.startswith(cell_type) - adata.obs.loc[mask, "augur_score"] = results["summary_metrics"][cell_type]["mean_augur_score"] - - # concatenate feature importances for each subsample cv - subsample_feature_importances_dicts = [cv["feature_importances"] for cv in results[cell_type]] - - for dictionary in subsample_feature_importances_dicts: - for key, value in dictionary.items(): - results["feature_importances"][key].extend(value) - results["feature_importances"]["cell_type"].extend( - [cell_type] - * (len(results["feature_importances"]["genes"]) - len(results["feature_importances"]["cell_type"])) - ) + eligible_cell_types.append((cell_type, cell_type_subsample)) + + # Fan subsamples out across all eligible cell types in one pool so workers don't idle at each cell-type boundary; submission order is preserved for regrouping. + flat_tasks = [(cell_type, sub, i) for cell_type, sub in eligible_cell_types for i in range(n_subsamples)] + cv_results = Parallel(n_jobs=n_threads, return_as="generator")( + delayed(self.cross_validate_subsample)( + adata=sub, + augur_mode=augur_mode, + subsample_size=subsample_size, + folds=folds, + feature_perc=feature_perc, + subsample_idx=i, + random_state=random_state, + zero_division=zero_division, + ) + for _, sub, i in flat_tasks + ) + cv_results = list(track(cv_results, total=len(flat_tasks), description="Processing data...")) + + for offset, (cell_type, _) in enumerate(eligible_cell_types): + results[cell_type] = cv_results[offset * n_subsamples : (offset + 1) * n_subsamples] + + # summarize scores for cell type + results["summary_metrics"][cell_type] = self.average_metrics(results[cell_type]) + + # add scores as observation to anndata + mask = adata.obs["cell_type"].str.startswith(cell_type) + adata.obs.loc[mask, "augur_score"] = results["summary_metrics"][cell_type]["mean_augur_score"] + + # concatenate feature importances for each subsample cv + subsample_feature_importances_dicts = [cv["feature_importances"] for cv in results[cell_type]] + + for dictionary in subsample_feature_importances_dicts: + for key, value in dictionary.items(): + results["feature_importances"][key].extend(value) + results["feature_importances"]["cell_type"].extend( + [cell_type] + * (len(results["feature_importances"]["genes"]) - len(results["feature_importances"]["cell_type"])) + ) - for idx, cv in zip(range(n_subsamples), results[cell_type], strict=False): - results["full_results"]["idx"].extend([idx] * folds) - results["full_results"]["augur_score"].extend(cv["test_augur_score"]) - results["full_results"]["folds"].extend(range(folds)) - results["full_results"]["cell_type"].extend([cell_type] * folds * n_subsamples) + for idx, cv in zip(range(n_subsamples), results[cell_type], strict=False): + results["full_results"]["idx"].extend([idx] * folds) + results["full_results"]["augur_score"].extend(cv["test_augur_score"]) + results["full_results"]["folds"].extend(range(folds)) + results["full_results"]["cell_type"].extend([cell_type] * folds * n_subsamples) # make sure one cell type worked if len(results) <= 2: logger.warning("No cells types had more than min_cells needed. Please adjust data or min_cells parameter.")