Documentation for the tasks

def compute_classification(
    adata: AnnData,
    classes: List[str],
    label_decoders: Dict[str, Any],
    labels_hierarchy: Dict[str, Any],
    metric_type: List[str] = ["macro", "micro", "weighted"],
    use_unknown: bool = False,
) -> Dict[str, Dict[str, float]]:
    """
    Compute classification metrics for the given annotated data.

    Args:
        adata (AnnData): The annotated data matrix of shape n_obs x n_vars. Rows correspond to cells and columns to genes.
        classes (List[str]): List of class labels to be used for classification.
        label_decoders (Dict[str, Any]): Dictionary of label decoders for each class.
        labels_hierarchy (Dict[str, Any]): Dictionary representing the hierarchy of labels.
        metric_type (List[str], optional): List of metric types to compute. Defaults to ["macro", "micro", "weighted"].

    Returns:
        Dict[str, Dict[str, float]]: A dictionary containing classification metrics for each class.
    """
    metrics = {}
    for clss in classes:
        res = []
        if clss not in adata.obs.columns:
            print("not in columns")
            continue
        labels_topred = label_decoders[clss].values()
        if clss in labels_hierarchy:
            parentdf = (
                bt.CellType.filter()
                .df(include=["parents__ontology_id", "ontology_id"])
                .set_index("ontology_id")[["parents__ontology_id"]]
            )
            parentdf.parents__ontology_id = parentdf.parents__ontology_id.astype(str)
            class_groupings = {
                k: get_descendants(k, parentdf) for k in set(adata.obs[clss].unique())
            }
        tokeep = np.array([True] * adata.shape[0])
        for i, (pred, true) in enumerate(adata.obs[["pred_" + clss, clss]].values):
            if pred == true:
                res.append(true)
                continue
            if true == "unknown":
                tokeep[i] = False
            if clss in labels_hierarchy:
                if true in class_groupings:
                    if pred == "unknown" and not use_unknown:
                        tokeep[i] = False
                    res.append(true if pred in class_groupings[true] else "")
                    continue
                elif true not in labels_topred:
                    raise ValueError(f"true label {true} not in available classes")
            elif true not in labels_topred:
                raise ValueError(f"true label {true} not in available classes")
            res.append("")
        metrics[clss] = {}
        metrics[clss]["accuracy"] = np.mean(
            np.array(res)[tokeep] == adata.obs[clss].values[tokeep]
        )
        for x in metric_type:
            metrics[clss][x] = f1_score(
                np.array(res)[tokeep], adata.obs[clss].values[tokeep], average=x
            )
    return metrics

def compute_corr(
    out: np.ndarray,
    to: np.ndarray,
    doplot: bool = True,
    compute_mean_regress: bool = False,
    plot_corr_size: int = 64,
) -> dict:
    """
    Compute the correlation between the output and target matrices.

    Args:
        out (np.ndarray): The output matrix.
        to (np.ndarray): The target matrix.
        doplot (bool, optional): Whether to generate a plot of the correlation coefficients. Defaults to True.
        compute_mean_regress (bool, optional): Whether to compute mean regression. Defaults to False.
        plot_corr_size (int, optional): The size of the plot for correlation. Defaults to 64.

    Returns:
        dict: A dictionary containing the computed metrics.
    """
    metrics = {}
    corr_coef, p_value = spearmanr(
        out,
        to.T,
    )
    corr_coef[p_value > 0.05] = 0
    # corr_coef[]
    # only on non zero values,
    # compare a1-b1 corr with a1-b(n) corr. should be higher

    # Plot correlation coefficient
    val = plot_corr_size + 2 if compute_mean_regress else plot_corr_size
    metrics.update(
        {"recons_corr": np.mean(corr_coef[val:, :plot_corr_size].diagonal())}
    )
    if compute_mean_regress:
        metrics.update(
            {
                "mean_regress": np.mean(
                    corr_coef[
                        plot_corr_size : plot_corr_size + 2,
                        :plot_corr_size,
                    ].flatten()
                )
            }
        )
    if doplot:
        plt.figure(figsize=(10, 5))
        plt.imshow(corr_coef, cmap="coolwarm", interpolation="none", vmin=-1, vmax=1)
        plt.colorbar()
        plt.title('Correlation Coefficient of expr and i["x"]')
        plt.show()
    return metrics

def default_benchmark(
    model: torch.nn.Module,
    folder_dir: str = FILE_LOC + "/../../data/",
    dataset: str = FILE_LOC + "/../../data/gNNpgpo6gATjuxTE7CCp.h5ad",
    do_class: bool = True,
    coarse: bool = False,
) -> dict:
    """
    Run the default benchmark for embedding and annotation using the scPRINT model.

    Args:
        model (torch.nn.Module): The scPRINT model to be used for embedding and annotation.
        folder_dir (str, optional): The directory containing data files.
        dataset (str, optional): The dataset to use for benchmarking. Can be a path or URL.
        do_class (bool, optional): Whether to perform classification. Defaults to True.
        coarse (bool, optional): Whether to use coarse cell type annotations. Defaults to False.

    Returns:
        dict: A dictionary containing the benchmark metrics.
    """
    if dataset.startswith("https://"):
        adata = sc.read(
            folder_dir
            + dataset.split("/")[-1]
            + (".h5ad" if not dataset.endswith(".h5ad") else ""),
            backup_url=dataset,
        )
    else:
        adata = sc.read_h5ad(dataset)
    if adata.shape[0] > 100_000:
        adata = adata[
            adata.obs_names[np.random.choice(adata.shape[0], 100_000, replace=False)]
        ]
    max_len = 4000 if adata.X.sum(1).mean() < 50_000 else 8000
    batch_size = 64 if adata.X.sum(1).mean() < 50_000 else 32
    log_every = 10_000
    if dataset.split("/")[-1] in ["24539942", "24539828"]:  # lung and pancreas
        adata.obs["organism_ontology_term_id"] = "NCBITaxon:9606"
        use_layer = "counts"
        is_symbol = True
        batch_key = "tech" if dataset.split("/")[-1] == "24539828" else "batch"
        label_key = "celltype" if dataset.split("/")[-1] == "24539828" else "cell_type"
        adata.obs["cell_type_ontology_term_id"] = adata.obs[label_key].replace(
            COARSE if coarse else FINE
        )
        adata.obs["assay_ontology_term_id"] = adata.obs[batch_key].replace(
            COARSE if coarse else FINE
        )
    else:
        use_layer = None
        is_symbol = False
        batch_key = (
            "batch"
            if dataset.split("/")[-1] == "661d5ec2-ca57-413c-8374-f49b0054ddba.h5ad"
            else "assay_ontology_term_id"
        )
        label_key = "cell_type_ontology_term_id"
    preprocessor = Preprocessor(
        use_layer=use_layer,
        is_symbol=is_symbol,
        force_preprocess=True,
        skip_validate=True,
        do_postp=model.expr_emb_style == "metacell",
        drop_non_primary=False,
    )
    adata = preprocessor(adata.copy())
    if model.expr_emb_style == "metacell":
        sc.pp.neighbors(adata, use_rep="X_pca")
    embedder = Embedder(
        pred_embedding=(
            model.pred_embedding if model.pred_embedding is not None else ["all"]
        ),
        doclass=do_class,
        max_len=max_len,
        doplot=False,
        keep_all_labels_pred=False,
        save_every=log_every,
        batch_size=batch_size,
        how="random expr",
    )
    adata, metrics = embedder(model, adata)

    bm = Benchmarker(
        adata,
        batch_key=batch_key,
        label_key=label_key,
        embedding_obsm_keys=["scprint_emb"],
    )
    bm.benchmark()
    metrics.update(
        {"scib": bm.get_results(min_max_scale=False).T.to_dict()["scprint_emb"]}
    )
    if model.class_scale > 0:
        metrics["classif"] = compute_classification(
            adata, model.classes, model.label_decoders, model.labels_hierarchy
        )
    return metrics

def display_confusion_matrix(
    nadata, pred="conv_pred_cell_type_ontology_term_id", true="cell_type"
):
    """
    Display the confusion matrix for true vs predicted cell types.

    Args:
        nadata (AnnData): Annotated data object containing predictions and ground truth.
        pred (str): Column name for predictions. Defaults to "conv_pred_cell_type_ontology_term_id".
        true (str): Column name for ground truth. Defaults to "cell_type".
    """
    counts = None
    for k, v in nadata.obs[true].value_counts().items():
        name = k + " - " + str(v)
        if counts is None:
            counts = pd.DataFrame(
                nadata.obs.loc[
                    nadata.obs[true] == k,
                    pred,
                ].value_counts()
            ).rename(columns={"count": name})
        else:
            counts = pd.concat(
                [
                    counts,
                    pd.DataFrame(
                        nadata.obs.loc[
                            nadata.obs[true] == k,
                            pred,
                        ].value_counts(),
                    ).rename(columns={"count": name}),
                ],
                axis=1,
            )
    counts = counts.T
    # Fill NaN values with 0 for visualization
    counts_filled = counts.fillna(0)

    # Create the heatmap
    plt.figure(figsize=(12, 10))

    # Convert to percentages (row-wise normalization)
    counts_percentage = counts_filled.div(counts_filled.sum(axis=1), axis=0) * 100
    counts_percentage = counts_percentage.iloc[:, counts_percentage.values.max(0) > 5]

    ax = sns.heatmap(
        counts_percentage,
        cmap="Blues",
        cbar_kws={"label": "Percentage (%)"},
        linewidths=0.5,
        square=True,
    )
    # place the x-label on top
    ax.xaxis.set_label_position("top")
    ax.xaxis.tick_top()

    plt.title(
        "Confusion Matrix: " + true + " vs " + pred + " (Percentage)",
        fontsize=16,
        pad=20,
    )
    ax.set_xlabel(pred, fontsize=12)
    ax.set_ylabel(true + " (with counts)", fontsize=12)
    ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="left", fontsize=12)
    ax.set_yticklabels(ax.get_yticklabels(), rotation=0, fontsize=14)
    plt.tight_layout()
    plt.show()

def default_benchmark(
    model: Any,
    default_dataset: str = "sroy",
    cell_types: List[str] = [],
    maxlayers: int = 16,
    maxgenes: int = 5000,
    batch_size: int = 32,
    maxcells: int = 1024,
) -> dict:
    """
    default_benchmark function to run the default scPRINT GRN benchmark

    Args:
        model (Any): The scPRINT model to be used for the benchmark.
        default_dataset (str, optional): The default dataset to use for benchmarking. Defaults to "sroy".
        cell_types (List[str], optional): List of cell types to include in the benchmark. Defaults to [].
        maxlayers (int, optional): Maximum number of layers to use from the model. Defaults to 16.
        maxgenes (int, optional): Maximum number of genes to consider. Defaults to 5000.
        batch_size (int, optional): Batch size for processing. Defaults to 32.
        maxcells (int, optional): Maximum number of cells to consider. Defaults to 1024.

    Returns:
        dict: A dictionary containing the benchmark metrics.
    """
    metrics = {}
    layers = list(range(model.nlayers))[max(0, model.nlayers - maxlayers) :]
    clf_omni = None
    if default_dataset == "sroy":
        preprocessor = Preprocessor(
            is_symbol=True,
            force_preprocess=True,
            skip_validate=True,
            do_postp=model.expr_emb_style == "metacell",
            min_valid_genes_id=5000,
            min_dataset_size=64,
            keepdata=True,
        )
        clf_self = None
        todo = [
            ("han", "human", "full"),
            ("mine", "human", "full"),
            ("han", "human", "chip"),
            ("han", "human", "ko"),
            ("tran", "mouse", "full"),
            ("zhao", "mouse", "full"),
            ("tran", "mouse", "chip"),
            ("tran", "mouse", "ko"),
        ]
        for da, spe, gt in todo:
            if gt != "full":
                continue
            if "NCBITaxon:10090" not in model.organisms and spe == "mouse":
                continue
            print(da + "_" + gt)
            preadata = get_sroy_gt(get=da, species=spe, gt=gt)
            adata = preprocessor(preadata.copy())
            if model.expr_emb_style == "metacell":
                sc.pp.neighbors(adata, use_rep="X_pca")
            grn_inferer = GNInfer(
                layer=layers,
                how="most var within",
                preprocess=(
                    "softpick"
                    if model.attention in ["softpick", "softpick-flash"]
                    else "softmax"
                ),
                head_agg="none",
                filtration="none",
                num_genes=maxgenes,
                num_workers=8,
                max_cells=maxcells,
                batch_size=batch_size,
            )
            grn = grn_inferer(model, adata)
            grn.varp["all"] = grn.varp["GRN"]
            grn.var["ensembl_id"] = grn.var.index
            grn.var["symbol"] = make_index_unique(grn.var["symbol"].astype(str))
            grn.var.index = grn.var["symbol"]
            grn.varp["GRN"] = grn.varp["all"].mean(-1).T
            metrics["mean_" + da + "_" + gt] = BenGRN(
                grn, do_auc=True, doplot=False
            ).compare_to(other=preadata)
            grn.varp["GRN"] = grn.varp["GRN"].T
            if spe == "human":
                metrics["mean_" + da + "_" + gt + "_base"] = BenGRN(
                    grn, do_auc=True, doplot=False
                ).scprint_benchmark()

            ## OMNI
            if clf_omni is None:
                grn.varp["GRN"] = grn.varp["all"]
                _, m, clf_omni = train_classifier(
                    grn,
                    C=1,
                    train_size=0.9,
                    class_weight={1: 800, 0: 1},
                    shuffle=True,
                    return_full=False,
                )
                joblib.dump(clf_omni, "clf_omni.pkl")
                metrics["omni_classifier"] = m
            coef = clf_omni.coef_[0] if clf_omni.coef_.shape[0] == 1 else clf_omni.coef_
            grn.varp["GRN"] = grn.varp["all"][:, :, coef > 0].mean(-1)
            if spe == "human":
                metrics["omni_" + da + "_" + gt + "_base"] = BenGRN(
                    grn, do_auc=True, doplot=True
                ).scprint_benchmark()
            grn.varp["GRN"] = grn.varp["GRN"].T
            metrics["omni_" + da + "_" + gt] = BenGRN(
                grn, do_auc=True, doplot=False
            ).compare_to(other=preadata)

            ## SELF
            if clf_self is None:
                grn.varp["GRN"] = np.transpose(grn.varp["all"], (1, 0, 2))
                _, m, clf_self = train_classifier(
                    grn,
                    other=preadata,
                    C=1,
                    train_size=0.5,
                    class_weight={1: 40, 0: 1},
                    shuffle=False,
                    return_full=False,
                )
                metrics["self_classifier"] = m
            coef = clf_self.coef_[0] if clf_self.coef_.shape[0] == 1 else clf_self.coef_
            grn.varp["GRN"] = grn.varp["all"][:, :, coef > 0].mean(-1).T
            metrics["self_" + da + "_" + gt] = BenGRN(
                grn, do_auc=True, doplot=False
            ).compare_to(other=preadata)
            if spe == "human":
                grn.varp["GRN"] = grn.varp["GRN"].T
                metrics["self_" + da + "_" + gt + "_base"] = BenGRN(
                    grn, do_auc=True, doplot=True
                ).scprint_benchmark()

            ## chip / ko
            if (da, spe, "chip") in todo:
                preadata = get_sroy_gt(get=da, species=spe, gt="chip")
                grn.varp["GRN"] = grn.varp["all"].mean(-1).T
                metrics["mean_" + da + "_" + "chip"] = BenGRN(
                    grn, do_auc=True, doplot=False
                ).compare_to(other=preadata)
                grn.varp["GRN"] = grn.varp["all"][:, :, coef > 0].mean(-1).T
                metrics["omni_" + da + "_" + "chip"] = BenGRN(
                    grn, do_auc=True, doplot=False
                ).compare_to(other=preadata)

                grn.varp["GRN"] = grn.varp["all"][:, :, coef > 0].mean(-1).T
                metrics["self_" + da + "_" + "chip"] = BenGRN(
                    grn, do_auc=True, doplot=False
                ).compare_to(other=preadata)
            if (da, spe, "ko") in todo:
                preadata = get_sroy_gt(get=da, species=spe, gt="ko")
                grn.varp["GRN"] = grn.varp["all"].mean(-1).T
                metrics["mean_" + da + "_" + "ko"] = BenGRN(
                    grn, do_auc=True, doplot=False
                ).compare_to(other=preadata)
                grn.varp["GRN"] = grn.varp["all"][:, :, coef > 0].mean(-1).T
                metrics["omni_" + da + "_" + "ko"] = BenGRN(
                    grn, do_auc=True, doplot=False
                ).compare_to(other=preadata)
                grn.varp["GRN"] = grn.varp["all"][:, :, coef > 0].mean(-1).T
                metrics["self_" + da + "_" + "ko"] = BenGRN(
                    grn, do_auc=True, doplot=False
                ).compare_to(other=preadata)
            del grn
    elif default_dataset == "gwps":
        adata = get_perturb_gt()
        preprocessor = Preprocessor(
            force_preprocess=True,
            keepdata=True,
            skip_validate=True,
            do_postp=model.expr_emb_style == "metacell",
            min_valid_genes_id=maxgenes,
            min_dataset_size=64,
        )
        nadata = preprocessor(adata.copy())
        if model.expr_emb_style == "metacell":
            sc.pp.neighbors(nadata, use_rep="X_pca")
        nadata.var["isTF"] = False
        nadata.var.loc[nadata.var.gene_name.isin(grnutils.TF), "isTF"] = True
        nadata.var["isTF"].sum()
        grn_inferer = GNInfer(
            layer=layers,
            how="most var within",
            preprocess=(
                "softpick"
                if model.attention in ["softpick", "softpick-flash"]
                else "softmax"
            ),
            head_agg="none",
            filtration="none",
            num_genes=maxgenes,
            max_cells=maxcells,
            num_workers=8,
            batch_size=batch_size,
        )
        grn = grn_inferer(model, nadata)
        del nadata
        grn.varp["all"] = grn.varp["GRN"]

        grn.varp["GRN"] = grn.varp["all"].mean(-1).T
        metrics["mean"] = BenGRN(grn, do_auc=True, doplot=False).compare_to(other=adata)
        grn.var["ensembl_id"] = grn.var.index
        grn.var.index = grn.var["symbol"]
        grn.varp["GRN"] = grn.varp["all"].mean(-1)
        metrics["mean_base"] = BenGRN(
            grn, do_auc=True, doplot=False
        ).scprint_benchmark()

        grn.varp["GRN"] = grn.varp["all"]
        grn.var.index = grn.var["ensembl_id"]
        _, m, clf_omni = train_classifier(
            grn,
            C=1,
            train_size=0.9,
            class_weight={1: 800, 0: 1},
            shuffle=True,
            doplot=False,
            return_full=False,
            use_col="gene_name",
        )
        coef = clf_omni.coef_[0] if clf_omni.coef_.shape[0] == 1 else clf_omni.coef_
        grn.varp["GRN"] = grn.varp["all"][:, :, coef > 0].mean(-1).T
        metrics["omni"] = BenGRN(grn, do_auc=True, doplot=False).compare_to(other=adata)
        metrics["omni_classifier"] = m
        grn.var.index = grn.var["symbol"]
        grn.varp["GRN"] = grn.varp["GRN"].T
        metrics["omni_base"] = BenGRN(
            grn, do_auc=True, doplot=False
        ).scprint_benchmark()
        grn.varp["GRN"] = np.transpose(grn.varp["all"], (1, 0, 2))
        grn.var.index = grn.var["ensembl_id"]
        _, m, clf_self = train_classifier(
            grn,
            other=adata,
            C=1,
            train_size=0.5,
            class_weight={1: 40, 0: 1},
            doplot=False,
            shuffle=False,
            return_full=False,
            use_col="ensembl_id",
        )
        coef = clf_self.coef_[0] if clf_self.coef_.shape[0] == 1 else clf_self.coef_
        grn.varp["GRN"] = grn.varp["all"][:, :, coef > 0].mean(-1).T
        metrics["self"] = BenGRN(grn, do_auc=True, doplot=False).compare_to(other=adata)
        metrics["self_classifier"] = m
        grn.var.index = grn.var["symbol"]
        grn.varp["GRN"] = grn.varp["GRN"].T
        metrics["self_base"] = BenGRN(
            grn, do_auc=True, doplot=False
        ).scprint_benchmark()
    elif default_dataset == "genernib":
        raise ValueError("Not implemented")
        # for adata in [NORMAN, OP, ADAMSON]:
        #   adata = sc.read_h5ad(adata)
        #   adata.obs["organism_ontology_term_id"] = "NCBITaxon:9606"
        #   preprocessor = Preprocessor(
        #       force_preprocess=False,
        #       skip_validate=True,
        #       drop_non_primary=False,
        #       do_postp=False,
        #       min_valid_genes_id=1000,
        #       min_dataset_size=64,
        #       keepdata=True,
        #       is_symbol=True,
        #       use_raw=False,
        #   )
        #   adata = preprocessor(adata.copy())
        #   run_gene_rnib(
        #      adata=adata,
        #      model=model,
        #      layer=layers,
        #      how="most var within",
        #      preprocess="softmax",
        #   )
        #   grn_inferer = GNInfer(
        #      how="most var across",
        #      preprocess="softmax",
        #      head_agg="mean",
        #      filtration="none",
        #      forward_mode="none",
        #      num_genes=3_000,
        #      max_cells=3000,
        #      batch_size=10,
        #      cell_type_col="perturbation",
        #      layer=list(range(model.nlayers))[:],
        #   )
        # grn = grn_inferer(model, adata, cell_type="ctrl")
        # grn.var.index = make_index_unique(grn.var["symbol"].astype(str))

    else:
        # max_genes=4000
        if default_dataset.startswith("https://"):
            adata = sc.read(
                FILEDIR + "/../../data/" + default_dataset.split("/")[-1],
                backup_url=default_dataset,
            )
        else:
            adata = sc.read_h5ad(default_dataset)
        if default_dataset.split("/")[-1] in ["yBCKp6HmXuHa0cZptMo7.h5ad"]:
            use_layer = "counts"
            is_symbol = True
        else:
            use_layer = None
            is_symbol = False

        preprocessor = Preprocessor(
            use_layer=use_layer,
            is_symbol=is_symbol,
            force_preprocess=True,
            skip_validate=True,
            do_postp=model.expr_emb_style == "metacell",
            drop_non_primary=False,
        )
        adata = preprocessor(adata.copy())

        adata.var["isTF"] = False
        adata.var.loc[adata.var.symbol.isin(grnutils.TF), "isTF"] = True
        if model.expr_emb_style == "metacell":
            if "X_pca" not in adata.obsm:
                sc.pp.pca(adata, n_comps=50)
            sc.pp.neighbors(adata, use_rep="X_pca")
        for celltype in list(adata.obs["cell_type"].unique())[:14]:
            # print(celltype)
            # grn_inferer = GNInfer(
            #    layer=layers,
            #    how="random expr",
            #    preprocess="softmax",
            #    head_agg="max",
            #    filtration="none",
            #    num_workers=8,
            #    num_genes=2200,
            #    max_cells=maxcells,
            #    batch_size=batch_size,
            # )
            #
            # grn = grn_inferer(model, adata[adata.X.sum(1) > 500], cell_type=celltype)
            # grn.var.index = make_index_unique(grn.var["symbol"].astype(str))
            # metrics[celltype + "_scprint"] = BenGRN(
            #    grn, doplot=False
            # ).scprint_benchmark()
            # del grn
            # gc.collect()
            grn_inferer = GNInfer(
                layer=layers,
                how="most var across",
                preprocess=(
                    "softpick"
                    if model.attention in ["softpick", "softpick-flash"]
                    else "softmax"
                ),
                head_agg="none",
                filtration="none",
                num_workers=8,
                num_genes=maxgenes,
                max_cells=maxcells,
                batch_size=batch_size,
            )
            grn = grn_inferer(model, adata[adata.X.sum(1) > 500], cell_type=celltype)
            grn.var.index = make_index_unique(grn.var["symbol"].astype(str))
            grn.varp["all"] = grn.varp["GRN"]
            grn.varp["GRN"] = grn.varp["GRN"].mean(-1)
            metrics[celltype + "_scprint_mean"] = BenGRN(
                grn, doplot=False
            ).scprint_benchmark()
            if clf_omni is None:
                grn.varp["GRN"] = grn.varp["all"]
                _, m, clf_omni = train_classifier(
                    grn,
                    C=1,
                    train_size=0.6,
                    max_iter=300,
                    class_weight={1: 800, 0: 1},
                    return_full=False,
                    shuffle=True,
                    doplot=False,
                )
                joblib.dump(clf_omni, "clf_omni.pkl")
                metrics["classifier"] = m
            coef = clf_omni.coef_[0] if clf_omni.coef_.shape[0] == 1 else clf_omni.coef_
            grn.varp["GRN"] = grn.varp["all"][:, :, coef > 0].mean(-1)
            metrics[celltype + "_scprint_class"] = BenGRN(
                grn, doplot=False
            ).scprint_benchmark()
            del grn
            gc.collect()
    return metrics