diff --git a/docs/command_line.rst b/docs/command_line.rst index 2140786c..7612a5c2 100644 --- a/docs/command_line.rst +++ b/docs/command_line.rst @@ -133,28 +133,17 @@ Key Features: :path: pooled-sbs -scallops norm-features -======================= +scallops pert-map +=================== -The `scallops norm-features` command is used to normalize features. +The `scallops pert-map` command provides functionality for perturbation map building. .. argparse:: :module: scallops.__main__ :func: create_parsers :prog: scallops - :path: rank-features - - -scallops rank-features -======================= - -The `scallops rank-features` command is used to compute significance from the output of `scallops norm-features`. + :path: pert-map -.. argparse:: - :module: scallops.__main__ - :func: create_parsers - :prog: scallops - :path: rank-features scallops registration diff --git a/pyproject.toml b/pyproject.toml index 7db35fa0..216f1c7b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ classifiers = [ dependencies = [ "adjustText", - "anndata>=0.12.4", # https://github.com/scverse/anndata/issues/2166 + "anndata>=0.13.0", "bioio>=3", "bioio-nd2", "bioio-ome-tiff", diff --git a/scallops/__init__.py b/scallops/__init__.py index bbd5b3f6..8820c659 100644 --- a/scallops/__init__.py +++ b/scallops/__init__.py @@ -1,7 +1,13 @@ import warnings +import anndata + +from scallops.zarr_io import write_basic_dask_dask_dense # noqa: F401 + from .experiment.elements import Experiment # noqa: F401 warnings.filterwarnings( - "ignore", message="unclosed.*", category=ResourceWarning, module="aiohttp" + "ignore", + message="Unclosed client.* | client_session.*", ) +anndata.settings.auto_shard_zarr_v3 = False diff --git a/scallops/__main__.py b/scallops/__main__.py index 6147be4f..c45f34de 100644 --- a/scallops/__main__.py +++ b/scallops/__main__.py @@ -41,9 +41,8 @@ features_main, find_objects_main, illumination_correction_main, - norm_features_main, + pert_map_main, pooled_if_sbs_main, - rank_features_main, register_main, segment_main, stitch_main, @@ -67,8 +66,7 @@ def create_parsers(default_help: bool = False) -> argparse.ArgumentParser: segment_main._create_parser(subparsers, default_help) illumination_correction_main._create_parser(subparsers, default_help) dialout_main._create_parser(subparsers, default_help) - norm_features_main._create_parser(subparsers, default_help) - rank_features_main._create_parser(subparsers, default_help) + pert_map_main._create_parser(subparsers, default_help) register_main._create_parser(subparsers, default_help) stitch_main._create_stitch_parser(subparsers, default_help) stitch_main._create_stitch_preview_parser(subparsers, default_help) diff --git a/scallops/cli/features.py b/scallops/cli/features.py index c1f466db..791474ac 100644 --- a/scallops/cli/features.py +++ b/scallops/cli/features.py @@ -49,7 +49,7 @@ _to_parquet, is_parquet_file, pluralize, - read_anndata_zarr, + read_anndata, ) logger = _get_cli_logger() @@ -79,7 +79,7 @@ def _read_merged_or_objects( area_column = f"{_label_name_to_prefix[label_name]}_AreaShape_Area" if merge_path.lower().endswith(".zarr"): - data = read_anndata_zarr(merge_path, dask=True) + data = read_anndata(merge_path, dask=True) merged_df = data.obs columns = {area_column} assert area_column in data.var.index diff --git a/scallops/cli/norm_features.py b/scallops/cli/norm_features.py deleted file mode 100644 index 404e3df8..00000000 --- a/scallops/cli/norm_features.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Module for the Command-Line Interface (CLI) related to normalizing features. - -Authors: - - The SCALLOPS development team -""" - -import argparse -import json -import os - -import dask.array as da -import dask.dataframe as dd -import fsspec -import pyarrow as pa -import pyarrow.parquet as pq - -from scallops.cli.util import ( - _create_dask_client, - _create_default_dask_config, - _dask_workers_threads, - _get_cli_logger, - cli_metadata, - load_json, -) -from scallops.features.normalize import _convert_scale, normalize_features -from scallops.features.util import ( - _join_metadata, - _query_anndata, - _read_data, - _slice_anndata, -) -from scallops.io import is_parquet_file -from scallops.utils import _fix_json -from scallops.zarr_io import is_anndata_zarr - -logger = _get_cli_logger() - - -def run_pipeline_norm_features(arguments: argparse.Namespace): - paths = arguments.input - features = arguments.features - reference = arguments.reference - label_filter = arguments.label_filter - join_path = arguments.metadata - join_fields = arguments.join - if join_path is not None and join_fields is None: - raise ValueError("Please specify join fields") - - norm_output = arguments.output - force = arguments.force - no_version = arguments.no_version - by = arguments.by - normalize = arguments.method - n_neighbors = arguments.neighbors - mad_scale_factor = arguments.mad_scale_factor - centering = not arguments.no_centering - scaling = not arguments.no_scaling - if mad_scale_factor.lower() == "normal": - mad_scale_factor = _convert_scale(mad_scale_factor) - else: - mad_scale_factor = float(mad_scale_factor) - - robust = arguments.robust - dask_server_url = arguments.client - dask_cluster_parameters = ( - load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} - ) - if dask_server_url is None and arguments.dask_cluster is None: - dask_cluster_parameters = _dask_workers_threads() - suffix = os.path.splitext(norm_output.lower())[1] - if suffix not in {".zarr", ".parquet", ".pq"}: - norm_output = norm_output + ".zarr" - output_format = "zarr" if norm_output.lower().endswith("zarr") else "parquet" - if not force: - skip = False - if output_format == "zarr" and is_anndata_zarr(norm_output): - skip = True - - elif output_format == "parquet" and is_parquet_file(norm_output): - skip = True - if skip: - logger.info( - f"{norm_output} already exists, skipping. Use --force to overwrite." - ) - return - - metadata = {} - if not no_version: - metadata.update(cli_metadata()) - with ( - _create_default_dask_config(), - _create_dask_client(dask_server_url, **dask_cluster_parameters), - ): - data = _read_data(paths, features) - - if label_filter is not None: - data = _slice_anndata(data, _query_anndata(data, label_filter).index) - if join_path is not None: - _join_metadata( - data, - dd.read_csv(join_path) - if not join_path.lower().endswith(".parquet") - or join_path.lower().endswith(".pq") - else dd.read_parquet(join_path), - join_fields, - ) - logger.info(f"# labels: {data.shape[0]:,}, # features: {data.shape[1]:,}") - if centering or scaling: - chunks = list(data.X.chunksize) - feature_chunk_size = 10 - if chunks[1] != feature_chunk_size: - chunks[1] = feature_chunk_size - data.X = data.X.rechunk(tuple(chunks)) - data = normalize_features( - data, - reference, - normalize=normalize, - robust=robust, - by=by, - n_neighbors=n_neighbors, - mad_scale=mad_scale_factor, - centering=centering, - scaling=scaling, - ) - else: - logger.info("No normalization") - - if output_format == "zarr": - if not da.core._check_regular_chunks(data.X.chunks): - # need uniform chunks to save to zarr - chunks = list(data.X.chunksize) - chunks[0] = "auto" - data.X = data.X.rechunk(tuple(chunks)) - data.uns["scallops"] = _fix_json(metadata) - data.write_zarr(norm_output, convert_strings_to_categoricals=False) - - else: - data.X = data.X.compute() - df = data.to_df().join(data.obs) - table = pa.Table.from_pandas(df, preserve_index=True) - table = table.replace_schema_metadata( - { - "scallops".encode(): json.dumps(metadata).encode(), - **table.schema.metadata, - } - ) - fs, output_file = fsspec.url_to_fs(norm_output) - pq.write_table( - table, - norm_output, - filesystem=fs, - ) diff --git a/scallops/cli/norm_features_main.py b/scallops/cli/norm_features_main.py deleted file mode 100644 index 948833ae..00000000 --- a/scallops/cli/norm_features_main.py +++ /dev/null @@ -1,111 +0,0 @@ -import argparse - -from scallops.cli.arg_parser import _sort_groups -from scallops.cli.util import ( - dask_client_arg, - dask_cluster_arg, - force_arg, - no_version_arg, -) - - -def _run_norm_features(arguments: argparse.Namespace): - from scallops.cli.norm_features import run_pipeline_norm_features - - run_pipeline_norm_features(arguments) - - -def _create_parser(subparsers: argparse.ArgumentParser, default_help: bool) -> None: - parser = subparsers.add_parser( - "norm-features", - help="Normalize features from output of `merge` command", - formatter_class=( - argparse.ArgumentDefaultsHelpFormatter - if default_help - else argparse.HelpFormatter - ), - ) - required = parser.add_argument_group("required arguments") - required.add_argument( - "-i", "--input", help="Path to merged file(s)", required=True, nargs="+" - ) - - required.add_argument( - "--output", - help="Path to save normalized features in Zarr or Parquet format", - required=True, - ) - parser.add_argument( - "--features", - help="Features to include. If not specified, all features are used.", - nargs="*", - ) - - parser.add_argument( - "--label-filter", - help="Expression to filter labels (e.g. barcode_Q_mean_0/barcode_Q_mean > 0.5)", - ) - - parser.add_argument( - "--by", - help="Stratify by groups when normalizing.", - nargs="*", - ) - - parser.add_argument( - "--reference", - help="Reference expression to normalize to (e.g. gene_symbol=='NTC').", - ) - - parser.add_argument( - "--robust", - help="Use robust statistics for normalization.", - action="store_true", - ) - parser.add_argument( - "--method", - help="Normalization method", - choices=["zscore", "local-zscore"], - default="zscore", - ) - parser.add_argument( - "--neighbors", - help="Number of neighbors for local z-score", - default=100, - type=int, - ) - parser.add_argument( - "--no-centering", - help="Do not center the data before scaling.", - action="store_true", - ) - parser.add_argument( - "--no-scaling", - help="Do not scale the data by dividing by standard deviation.", - action="store_true", - ) - - parser.add_argument( - "--metadata", - help="Path to CSV or Parquet file containing metadata to join with merged data.", - ) - parser.add_argument( - "--join", - help="Field(s) to join on", - nargs="*", - ) - - parser.add_argument( - "--mad-scale-factor", - help="Numerical scale factor to divide median absolute deviation. " - "The string “normal” is also accepted, and results in scale being the" - " inverse of the standard normal quantile function at 0.75", - default="normal", - type=str, - ) - dask_client_arg(parser) - dask_cluster_arg(parser) - force_arg(parser) - no_version_arg(parser) - _sort_groups(parser) - parser.set_defaults(func=_run_norm_features) diff --git a/scallops/cli/pert_map.py b/scallops/cli/pert_map.py new file mode 100644 index 00000000..9f46466d --- /dev/null +++ b/scallops/cli/pert_map.py @@ -0,0 +1,778 @@ +"""Module for the Command-Line Interface (CLI) related to normalizing features. + +Authors: + - The SCALLOPS development team +""" + +import argparse +import json +import os + +import anndata +import dask.dataframe as dd +import fsspec +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq + +from scallops.cli.util import ( + _create_dask_client, + _create_default_dask_config, + _dask_workers_threads, + _get_cli_logger, + cli_metadata, + load_json, +) +from scallops.features.agg import agg_features +from scallops.features.decomposition import pca +from scallops.features.map_eval import pairwise_similarities, read_corum, recall +from scallops.features.normalize import ( + _convert_scale, + normalize_features, + typical_variation_normalization, +) +from scallops.features.preprocessing import filter_data +from scallops.features.rank import rank_features +from scallops.features.util import ( + _join_metadata, + _slice_anndata, + pandas_to_anndata, +) +from scallops.io import _to_parquet, is_anndata, is_parquet_file, read_anndata +from scallops.utils import _fix_json + +logger = _get_cli_logger() + + +def _read_data( + data_paths: list[str], + feature_filter: str | None = None, + label_filter: str | None = None, + use_dask: bool = True, +) -> anndata.AnnData: + if label_filter is not None: + if fsspec.url_to_fs(label_filter)[0].exists(label_filter): + label_filter = pd.read_parquet(label_filter).index + elif label_filter.endswith(".parquet"): + logger.warning(f"{label_filter} path not found.") + if feature_filter is not None: + if fsspec.url_to_fs(feature_filter)[0].exists(feature_filter): + feature_filter = pd.read_parquet(feature_filter).index + elif feature_filter.endswith(".parquet"): + logger.warning(f"{feature_filter} path not found.") + results = [] + for data_path in data_paths: + fs, data_path = fsspec.url_to_fs(data_path) + if "*" in data_path: + paths = fs.glob(data_path) + if len(paths) == 0: + raise ValueError(f"No files found at {data_path}.") + + else: + paths = [data_path] + for path in paths: + path = fs.unstrip_protocol(path) + if path.endswith(".parquet"): + d = dd.read_parquet(path) if use_dask else pd.read_parquet(path) + d = pandas_to_anndata(d) + elif path.endswith(".zarr") or path.endswith(".h5ad"): + d = read_anndata(path, dask=use_dask) + else: + raise ValueError(f"Unrecognized file type: {path}") + assert not d.obs.index.has_duplicates, "Duplicate obs index detected." + assert not d.var.index.has_duplicates, "Duplicate var index detected." + results.append(d) + + data = ( + results[0] if len(results) == 1 else anndata.concat(results, index_unique="-") + ) + + assert not data.obs.index.has_duplicates + if isinstance(label_filter, str): + label_filter = data.obs.query(label_filter).index + if isinstance(feature_filter, str): + feature_filter = data.var.query(feature_filter).index + if label_filter is not None or feature_filter is not None: + data = _slice_anndata(data, label_filter, feature_filter) + return data + + +def rechunk( + data: anndata.AnnData, + rechunk_label_size: str | None, + rechunk_feature_size: str | None, +) -> anndata.AnnData: + if rechunk_label_size == "": + rechunk_label_size = None + if rechunk_feature_size == "": + rechunk_feature_size = None + if rechunk_label_size is not None or rechunk_feature_size is not None: + if rechunk_label_size is not None and rechunk_label_size.isdigit(): + rechunk_label_size = int(rechunk_label_size) + if rechunk_feature_size is not None and rechunk_feature_size.isdigit(): + rechunk_feature_size = int(rechunk_feature_size) + if rechunk_label_size is None: + rechunk_label_size = data.X.chunksize[0] + if rechunk_feature_size is None: + rechunk_feature_size = data.X.chunksize[1] + data.X = data.X.rechunk((rechunk_label_size, rechunk_feature_size)) + return data + + +def run_recall(arguments: argparse.Namespace): + data_paths = arguments.input + force = arguments.force + no_version = arguments.no_version + dask_server_url = arguments.client + ground_truth_paths = arguments.ground_truth_corum + dask_cluster_parameters = ( + load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} + ) + output = arguments.output + recall_thresholds = arguments.threshold + if not force and is_parquet_file(output): + logger.info(f"{output} already exists, skipping. Use --force to overwrite.") + return + ground_truth = [] + for i in range(len(ground_truth_paths)): + corum_df = read_corum(ground_truth_paths[i]) + corum_df = corum_df.set_index(corum_df["a"] + "-" + corum_df["b"]) + corum_name = os.path.basename(ground_truth_paths[i]) + ground_truth.append((corum_name, corum_df)) + metadata = {} + if not no_version: + metadata.update(cli_metadata()) + with ( + _create_default_dask_config(), + _create_dask_client(dask_server_url, **dask_cluster_parameters), + ): + similarity_data = _read_data(data_paths) + similarity_data.X = similarity_data.X.compute() # load into memory + results = [] + gene_symbols = similarity_data.obs.index.values + for ground_truth_name, ground_truth_df in ground_truth: + indices_a = [] + indices_b = [] + + for i in range(len(gene_symbols)): + gene1 = gene_symbols[i] + for j in range(i): + key = gene1 + "-" + gene_symbols[j] + if key in ground_truth_df.index: + indices_a.append(i) + indices_b.append(j) + if len(indices_a) == 0: + raise ValueError("No genes found in ground truth.") + indices_a = np.array(indices_a) + indices_b = np.array(indices_b) + + query_distribution = similarity_data.X[indices_a, indices_b] + null_distribution = similarity_data.X[ + np.tril_indices(similarity_data.shape[0], k=-1) + ] + result = recall( + query_distribution=query_distribution, + null_distribution=null_distribution, + recall_thresholds=recall_thresholds, + ) + result["null"] = "all" + result["dataset"] = ground_truth_name + results.append(result) + + in_corum = similarity_data.obs.index.isin(ground_truth_df["a"].unique()) + similarity_data_corum = similarity_data.X[in_corum,][:, in_corum] + null_distribution_corum = similarity_data_corum[ + np.tril_indices(similarity_data_corum.shape[0], k=-1) + ] + result = recall( + query_distribution=query_distribution, + null_distribution=null_distribution_corum, + recall_thresholds=recall_thresholds, + ) + result["null"] = "CORUM" + result["dataset"] = ground_truth_name + results.append(result) + + df = pd.concat(results) + multi_threshold = False + for threshold in recall_thresholds: + if not np.isscalar(threshold): + multi_threshold = True + break + if multi_threshold: + df["threshold"] = df["threshold"].astype(str) + table = pa.Table.from_pandas(df, preserve_index=False) + table = table.replace_schema_metadata( + { + "scallops".encode(): json.dumps(metadata).encode(), + **table.schema.metadata, + } + ) + + fs, output = fsspec.url_to_fs(output) + pq.write_table( + table, + output, + filesystem=fs, + ) + + +def run_similarity_matrix(arguments: argparse.Namespace): + data_paths = arguments.input + force = arguments.force + no_version = arguments.no_version + + dask_server_url = arguments.client + dask_cluster_parameters = ( + load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} + ) + output = arguments.output + + if not force and is_anndata(output): + logger.info(f"{output} already exists, skipping. Use --force to overwrite.") + return + + metadata = {} + if not no_version: + metadata.update(cli_metadata()) + with ( + _create_default_dask_config(), + _create_dask_client(dask_server_url, **dask_cluster_parameters), + ): + data = _read_data(data_paths) + logger.info(f"# labels: {data.shape[0]:,}, # features: {data.shape[1]:,}") + data = anndata.AnnData( + X=pairwise_similarities(data), obs=data.obs, var=data.obs + ) + _write_anndata(data, output, metadata, None, None) + + +def run_aggregate(arguments: argparse.Namespace): + data_paths = arguments.input + label_filter = arguments.label_filter + feature_filter = arguments.feature_filter + join_path = arguments.metadata + join_fields = arguments.join + if join_path is not None and join_fields is None: + raise ValueError("Please specify join fields") + rechunk_label_size = arguments.rechunk_labels + rechunk_feature_size = arguments.rechunk_features + force = arguments.force + no_version = arguments.no_version + + dask_server_url = arguments.client + dask_cluster_parameters = ( + load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} + ) + if dask_server_url is None and arguments.dask_cluster is None: + dask_cluster_parameters = _dask_workers_threads(threads_per_worker=8) + by = arguments.by + output = arguments.output + + if not force and is_anndata(output): + logger.info(f"{output} already exists, skipping. Use --force to overwrite.") + return + center_reference_query = arguments.center_reference_query + metadata = {} + if not no_version: + metadata.update(cli_metadata()) + with ( + _create_default_dask_config(), + _create_dask_client(dask_server_url, **dask_cluster_parameters), + ): + data = _read_data(data_paths, feature_filter, label_filter) + data = rechunk(data, rechunk_label_size, rechunk_feature_size) + if join_path is not None: + _join_metadata( + data, + dd.read_csv(join_path) + if not join_path.lower().endswith(".parquet") + or join_path.lower().endswith(".pq") + else dd.read_parquet(join_path), + join_fields, + ) + logger.info(f"# labels: {data.shape[0]:,}, # features: {data.shape[1]:,}") + + if center_reference_query is not None: + data = normalize_features( + data=data, + normalize="zscore", + scaling=False, + robust=False, + reference_query=center_reference_query, + ) + + data = agg_features( + data=data, + by=by, + agg_func="mean", + ) + fs, output_dir = fsspec.url_to_fs(os.path.dirname(output)) + fs.makedirs(output_dir, exist_ok=True) + data.uns["scallops"] = _fix_json(metadata) + if output.lower().endswith(".zarr"): + data.write_zarr(output, convert_strings_to_categoricals=False) + else: + data.write_h5ad(output, convert_strings_to_categoricals=False) + + +def run_tvn(arguments: argparse.Namespace): + data_paths = arguments.input + label_filter = arguments.label_filter + feature_filter = arguments.feature_filter + join_path = arguments.metadata + join_fields = arguments.join + if join_path is not None and join_fields is None: + raise ValueError("Please specify join fields") + rechunk_label_size = arguments.rechunk_labels + rechunk_feature_size = arguments.rechunk_features + post_rechunk_label_size = arguments.post_rechunk_labels + post_rechunk_feature_size = arguments.post_rechunk_features + force = arguments.force + no_version = arguments.no_version + + dask_server_url = arguments.client + dask_cluster_parameters = ( + load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} + ) + reference_query = arguments.reference_query + by = arguments.by + output = arguments.output + + if not force and is_anndata(output): + logger.info(f"{output} already exists, skipping. Use --force to overwrite.") + return + + metadata = {} + if not no_version: + metadata.update(cli_metadata()) + with ( + _create_default_dask_config(), + _create_dask_client(dask_server_url, **dask_cluster_parameters), + ): + data = _read_data(data_paths, feature_filter, label_filter) + data = rechunk(data, rechunk_label_size, rechunk_feature_size) + if join_path is not None: + _join_metadata( + data, + dd.read_csv(join_path) + if not join_path.lower().endswith(".parquet") + or join_path.lower().endswith(".pq") + else dd.read_parquet(join_path), + join_fields, + ) + logger.info(f"# labels: {data.shape[0]:,}, # features: {data.shape[1]:,}") + data = typical_variation_normalization( + data=data, + reference_query=reference_query, + by=by, + ) + + _write_anndata( + data, output, metadata, post_rechunk_label_size, post_rechunk_feature_size + ) + + +def run_pca(arguments: argparse.Namespace): + data_paths = arguments.input + label_filter = arguments.label_filter + feature_filter = arguments.feature_filter + join_path = arguments.metadata + join_fields = arguments.join + batch_size = arguments.batch_size + if batch_size is not None and batch_size <= 0: + batch_size = None + if join_path is not None and join_fields is None: + raise ValueError("Please specify join fields") + rechunk_label_size = arguments.rechunk_labels + rechunk_feature_size = arguments.rechunk_features + post_rechunk_label_size = arguments.post_rechunk_labels + post_rechunk_feature_size = arguments.post_rechunk_features + force = arguments.force + no_version = arguments.no_version + + dask_server_url = arguments.client + dask_cluster_parameters = ( + load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} + ) + n_components = arguments.components + whiten = arguments.whiten + output = arguments.output + + if not force and is_anndata(output): + logger.info(f"{output} already exists, skipping. Use --force to overwrite.") + return + + metadata = {} + if not no_version: + metadata.update(cli_metadata()) + dask_config = {} + if ( + batch_size is None + and rechunk_label_size is not None + or rechunk_feature_size is not None + ): + dask_config = {"array.rechunk.method": "tasks"} # for dask PCA + with ( + _create_default_dask_config(dask_config), + _create_dask_client(dask_server_url, **dask_cluster_parameters), + ): + data = _read_data(data_paths, feature_filter, label_filter) + data = rechunk(data, rechunk_label_size, rechunk_feature_size) + if join_path is not None: + _join_metadata( + data, + dd.read_csv(join_path) + if not join_path.lower().endswith(".parquet") + or join_path.lower().endswith(".pq") + else dd.read_parquet(join_path), + join_fields, + ) + logger.info(f"# labels: {data.shape[0]:,}, # features: {data.shape[1]:,}") + data = pca( + data=data, n_components=n_components, whiten=whiten, batch_size=batch_size + ) + _write_anndata( + data, output, metadata, post_rechunk_label_size, post_rechunk_feature_size + ) + + +def _write_anndata(data, output, metadata, rechunk_label_size, rechunk_feature_size): + output = output.rstrip("/") + fs, output_dir = fsspec.url_to_fs(os.path.dirname(output)) + fs.makedirs(output_dir, exist_ok=True) + data.uns["scallops"] = _fix_json(metadata) + data = rechunk(data, rechunk_label_size, rechunk_feature_size) + if output.lower().endswith(".zarr"): + data.write_zarr(output, convert_strings_to_categoricals=False) + else: + data.write_h5ad(output, convert_strings_to_categoricals=False) + + +def run_rank_features(arguments: argparse.Namespace): + data_paths = arguments.input + label_filter = arguments.label_filter + feature_filter = arguments.feature_filter + join_path = arguments.metadata + join_fields = arguments.join + if join_path is not None and join_fields is None: + raise ValueError("Please specify join fields") + rechunk_label_size = arguments.rechunk_labels + rechunk_feature_size = arguments.rechunk_features + + force = arguments.force + no_version = arguments.no_version + by = arguments.by + + dask_server_url = arguments.client + dask_cluster_parameters = ( + load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} + ) + + method = arguments.rank_method + + rank_output = arguments.output + if rank_output is None: + rank_output = os.path.splitext(os.path.basename(data_paths[0]))[0] + ".parquet" + if len(data_paths) > 1: + logger.info(f"Saving results to {rank_output}") + + perturbation_column = arguments.perturbation + min_labels = arguments.min_labels + reference_value = arguments.reference + iqr_multiplier = arguments.iqr_multiplier + + if not rank_output.lower().endswith( + ".parquet" + ) and not rank_output.lower().endswith(".pq"): + rank_output = rank_output + ".parquet" + + if not force and is_parquet_file(rank_output): + logger.info( + f"{rank_output} already exists, skipping. Use --force to overwrite." + ) + return + + metadata = {} + if not no_version: + metadata.update(cli_metadata()) + with ( + _create_default_dask_config(), + _create_dask_client(dask_server_url, **dask_cluster_parameters), + ): + if label_filter is None: + label_filter = f"~`{perturbation_column}`.isna()" + data = _read_data(data_paths, feature_filter, label_filter) + data = rechunk(data, rechunk_label_size, rechunk_feature_size) + if join_path is not None: + _join_metadata( + data, + dd.read_csv(join_path) + if not join_path.lower().endswith(".parquet") + or join_path.lower().endswith(".pq") + else dd.read_parquet(join_path), + join_fields, + ) + logger.info(f"# labels: {data.shape[0]:,}, # features: {data.shape[1]:,}") + # columns_needed = set() + # columns_needed.add(perturbation_column) + # if by is not None: + # columns_needed.update(by) + # if label_filter is not None: + # columns_needed.update(_get_names_from_pd_query(label_filter)) + # if join_path is not None: + # columns_needed.update(join_fields) + # _load_coords(data, list(columns_needed)) + + rank_df = rank_features( + data=data, + by=by, + perturbation_column=perturbation_column, + reference_value=reference_value, + method=method, + min_labels=min_labels, + iqr_multiplier=iqr_multiplier, + ) + + fs, output_dir = fsspec.url_to_fs(os.path.dirname(rank_output)) + fs.makedirs(output_dir, exist_ok=True) + + if isinstance(rank_df, dd.DataFrame): + _to_parquet( + rank_df, + rank_output, + write_index=False, + custom_metadata=dict(scallops=json.dumps(metadata)), + ) + else: + table = pa.Table.from_pandas(rank_df, preserve_index=False) + table = table.replace_schema_metadata( + { + "scallops".encode(): json.dumps(metadata).encode(), + **table.schema.metadata, + } + ) + + fs, rank_output = fsspec.url_to_fs(rank_output) + pq.write_table( + table, + rank_output, + filesystem=fs, + ) + + +def run_norm_features(arguments: argparse.Namespace): + data_paths = arguments.input + label_filter = arguments.label_filter + feature_filter = arguments.feature_filter + join_path = arguments.metadata + join_fields = arguments.join + if join_path is not None and join_fields is None: + raise ValueError("Please specify join fields") + + rechunk_label_size = arguments.rechunk_labels + rechunk_feature_size = arguments.rechunk_features + post_rechunk_label_size = arguments.post_rechunk_labels + post_rechunk_feature_size = arguments.post_rechunk_features + force = arguments.force + no_version = arguments.no_version + by = arguments.by + + dask_server_url = arguments.client + dask_cluster_parameters = ( + load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} + ) + reference = arguments.reference + output = arguments.output + + normalize = arguments.method + n_neighbors = arguments.neighbors + mad_scale_factor = arguments.mad_scale_factor + centering = not arguments.no_centering + scaling = not arguments.no_scaling + if mad_scale_factor.lower() == "normal": + mad_scale_factor = _convert_scale(mad_scale_factor) + else: + mad_scale_factor = float(mad_scale_factor) + + robust = arguments.robust + max_value = arguments.max_value + batch_size = arguments.batch_size + if batch_size < 0: + batch_size = None + centroid_column_names = arguments.centroid_columns + if dask_server_url is None and arguments.dask_cluster is None: + dask_cluster_parameters = _dask_workers_threads(threads_per_worker=8) + + output_ext = os.path.splitext(os.path.basename(output.lower()))[1] + if output_ext == ".zarr": + output_format = "zarr" + elif output_ext == ".h5ad": + output_format = "h5ad" + else: + output_format = "parquet" + if not force: + skip = False + if output_format in ("zarr", "h5ad") and is_anndata(output): + skip = True + + elif output_format == "parquet" and is_parquet_file(output): + skip = True + if skip: + logger.info(f"{output} already exists, skipping. Use --force to overwrite.") + return + + metadata = {} + if not no_version: + metadata.update(cli_metadata()) + with ( + _create_default_dask_config( + { + "distributed.scheduler.worker-saturation": 1.0, + "optimization.fuse.active": False, + "distributed.admin.large-graph-warning-threshold": "100MB", + "distributed.worker.resources.process": 1, + } + ), + _create_dask_client(dask_server_url, **dask_cluster_parameters), + ): + data = _read_data(data_paths, feature_filter, label_filter) + data = rechunk(data, rechunk_label_size, rechunk_feature_size) + if join_path is not None: + _join_metadata( + data, + dd.read_csv(join_path) + if not join_path.lower().endswith(".parquet") + or join_path.lower().endswith(".pq") + else dd.read_parquet(join_path), + join_fields, + ) + _log_data_shape(data) + if centering or scaling: + data = normalize_features( + data, + reference, + normalize=normalize, + robust=robust, + by=by, + n_neighbors=n_neighbors, + mad_scale=mad_scale_factor, + centering=centering, + scaling=scaling, + max_value=max_value, + batch_size=batch_size, + centroid_column_names=centroid_column_names, + ) + _log_data_shape(data, "After normalization, ") + else: + logger.info("No normalization") + fs, output_dir = fsspec.url_to_fs(os.path.dirname(output)) + fs.makedirs(output_dir, exist_ok=True) + + if output_format in ("zarr", "h5ad"): + _write_anndata( + data, + output, + metadata, + post_rechunk_label_size, + post_rechunk_feature_size, + ) + else: + data.X = data.X.compute() + df = data.to_df().join(data.obs) + table = pa.Table.from_pandas(df, preserve_index=True) + table = table.replace_schema_metadata( + { + "scallops".encode(): json.dumps(metadata).encode(), + **table.schema.metadata, + } + ) + fs, output_file = fsspec.url_to_fs(output) + pq.write_table( + table, + output, + filesystem=fs, + ) + + +def run_filter_data(arguments: argparse.Namespace) -> None: + data_paths = arguments.input + label_filter = arguments.label_filter + feature_filter = arguments.feature_filter + join_path = arguments.metadata + join_fields = arguments.join + scale = not arguments.no_scale + if join_path is not None and join_fields is None: + raise ValueError("Please specify join fields") + rechunk_label_size = arguments.rechunk_labels + rechunk_feature_size = arguments.rechunk_features + post_rechunk_label_size = arguments.post_rechunk_labels + post_rechunk_feature_size = arguments.post_rechunk_features + force = arguments.force + no_version = arguments.no_version + by = arguments.by + + dask_server_url = arguments.client + dask_cluster_parameters = ( + load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} + ) + output = arguments.output + + if not force and is_anndata(output): + logger.info(f"Skipping {output}") + return + + n_features = arguments.n_features + min_feature_variance = arguments.min_feature_variance + max_feature_variance = arguments.max_feature_variance + max_cell_fraction_not_finite = arguments.max_cell_fraction_not_finite + if min_feature_variance is not None and min_feature_variance < 0: + min_feature_variance = None + if max_feature_variance is not None and max_feature_variance < 0: + max_feature_variance = None + if max_cell_fraction_not_finite is not None and max_cell_fraction_not_finite < 0: + max_cell_fraction_not_finite = None + + metadata = {} + if not no_version: + metadata.update(cli_metadata()) + with ( + _create_default_dask_config( + {"distributed.scheduler.locks.lease-timeout": "inf"} + ), + _create_dask_client(dask_server_url, **dask_cluster_parameters), + ): + data = _read_data(data_paths, feature_filter, label_filter) + data = rechunk(data, rechunk_label_size, rechunk_feature_size) + if join_path is not None: + _join_metadata( + data, + dd.read_csv(join_path) + if not join_path.lower().endswith(".parquet") + or join_path.lower().endswith(".pq") + else dd.read_parquet(join_path), + join_fields, + ) + _log_data_shape(data) + + data = filter_data( + data=data, + max_fraction_not_finite=max_cell_fraction_not_finite, + min_variance=min_feature_variance, + max_variance=max_feature_variance, + n_features=n_features, + by=by, + scale=scale, + ) + + _log_data_shape(data, "After filtering, ") + _write_anndata( + data, output, metadata, post_rechunk_label_size, post_rechunk_feature_size + ) + + +def _log_data_shape(data, prefix="", log_chunk_size=True): + logger.info(f"{prefix}# labels: {data.shape[0]:,}, # features: {data.shape[1]:,}") + if log_chunk_size: + logger.info(f"Chunk size: {data.X.chunksize[0]:,}, {data.X.chunksize[1]:,}") diff --git a/scallops/cli/pert_map_main.py b/scallops/cli/pert_map_main.py new file mode 100644 index 00000000..a7c3c265 --- /dev/null +++ b/scallops/cli/pert_map_main.py @@ -0,0 +1,592 @@ +import argparse + +from scallops.cli.arg_parser import _sort_groups +from scallops.cli.util import ( + dask_client_arg, + dask_cluster_arg, + force_arg, + no_version_arg, +) + + +def _run_filter_data(arguments: argparse.Namespace): + from scallops.cli.pert_map import run_filter_data + + run_filter_data(arguments) + + +def _run_pca(arguments: argparse.Namespace): + from scallops.cli.pert_map import run_pca + + run_pca(arguments) + + +def _run_tvn(arguments: argparse.Namespace): + from scallops.cli.pert_map import run_tvn + + run_tvn(arguments) + + +def _run_aggregate(arguments: argparse.Namespace): + from scallops.cli.pert_map import run_aggregate + + run_aggregate(arguments) + + +def input_arg(parser: argparse.ArgumentParser): + parser.add_argument( + "--input", + type=str, + nargs="+", + help="Path to one or more zarr, h5ad, or Parquet files or a " + "pattern to match files (e.g. s3://foo/*.zarr).", + ) + + +def common_args( + parser: argparse.ArgumentParser, + metadata: bool = True, + pre_rechunk: bool = True, + post_rechunk: bool = True, + dask_client_value: str | None = None, + rechunk_features: str | None = None, + rechunk_labels: str | None = None, +): + if metadata: + metadata_args(parser) + if pre_rechunk: + pre_rechunk_args( + parser, rechunk_features=rechunk_features, rechunk_labels=rechunk_labels + ) + if post_rechunk: + post_rechunk_args( + parser, rechunk_features=rechunk_features, rechunk_labels=rechunk_labels + ) + dask_client_arg(parser, dask_client_value) + dask_cluster_arg(parser) + force_arg(parser) + no_version_arg(parser) + _sort_groups(parser) + + +def pre_rechunk_args( + parser: argparse.ArgumentParser, + rechunk_features: str | None = None, + rechunk_labels: str | None = None, +): + parser.add_argument( + "--pre-rechunk-labels", + type=str, + dest="rechunk_labels", + default=rechunk_labels, + help="Rechunk dataset labels before processing.", + ) + + parser.add_argument( + "--pre-rechunk-features", + type=str, + dest="rechunk_features", + default=rechunk_features, + help="Rechunk dataset features before processing.", + ) + + +def post_rechunk_args( + parser: argparse.ArgumentParser, + rechunk_features: str | None = None, + rechunk_labels: str | None = None, +): + parser.add_argument( + "--post-rechunk-labels", + type=str, + default=rechunk_labels, + help="Rechunk dataset labels after processing.", + ) + + parser.add_argument( + "--post-rechunk-features", + type=str, + default=rechunk_features, + help="Rechunk dataset features after processing.", + ) + + +def _run_norm_features(arguments: argparse.Namespace): + from scallops.cli.pert_map import run_norm_features + + run_norm_features(arguments) + + +def _run_similarity_matrix(arguments: argparse.Namespace): + from scallops.cli.pert_map import run_similarity_matrix + + run_similarity_matrix(arguments) + + +def _run_recall(arguments: argparse.Namespace): + from scallops.cli.pert_map import run_recall + + run_recall(arguments) + + +def filter_args( + parser: argparse.ArgumentParser, + label_filter: bool = True, + feature_filter: bool = True, +): + if label_filter: + parser.add_argument( + "--label-filter", + type=str, + help="Query string to filter dataset before processing (e.g. gene_symbol!='foo') or path to " + "Parquet file containing label identifiers.", + ) + if feature_filter: + parser.add_argument( + "--feature-filter", + type=str, + help="Query string to filter dataset before processing (e.g. gene_symbol!='foo') or path to " + "Parquet file containing label identifiers.", + ) + + +def metadata_args(parser: argparse.ArgumentParser): + parser.add_argument( + "--metadata", + help="Path to CSV or Parquet file containing metadata to join with dataset.", + ) + parser.add_argument( + "--join", + help="Field(s) in metadata to join on", + nargs="*", + ) + + +def _create_similarity_matrix_parser( + subparsers: argparse.ArgumentParser, default_help: bool +) -> None: + parser = subparsers.add_parser( + "similarity-matrix", + help="Create pairwise similarity matrix", + formatter_class=( + argparse.ArgumentDefaultsHelpFormatter + if default_help + else argparse.HelpFormatter + ), + ) + required = parser.add_argument_group("required arguments") + input_arg(required) + + required.add_argument( + "--output", + help="Path to save result in zarr or h5ad format", + required=True, + ) + + required.add_argument( + "--by", + help="Perturbation column(s) in dataset observations to aggregate by.", + nargs="+", + ) + + common_args(parser=parser, metadata=False, pre_rechunk=False, post_rechunk=False) + parser.set_defaults(func=_run_similarity_matrix) + + +def _create_aggregate_parser( + subparsers: argparse.ArgumentParser, default_help: bool +) -> None: + parser = subparsers.add_parser( + "aggregate", + help="Run aggregatation", + formatter_class=( + argparse.ArgumentDefaultsHelpFormatter + if default_help + else argparse.HelpFormatter + ), + ) + required = parser.add_argument_group("required arguments") + input_arg(required) + + required.add_argument( + "--output", + help="Path to save result in zarr or h5ad format", + required=True, + ) + + required.add_argument( + "--by", + help="Perturbation column(s) in dataset observations to aggregate by.", + nargs="+", + ) + parser.add_argument( + "--center-reference-query", + help="Center the data to a reference before aggregating (e.g. gene_symbol=='NTC')", + ) + filter_args(parser) + + common_args(parser, dask_client_value="none", pre_rechunk=True, post_rechunk=True) + parser.set_defaults(func=_run_aggregate) + + +def _create_tvn_parser(subparsers: argparse.ArgumentParser, default_help: bool) -> None: + parser = subparsers.add_parser( + "tvn", + help="Run TNV", + formatter_class=( + argparse.ArgumentDefaultsHelpFormatter + if default_help + else argparse.HelpFormatter + ), + ) + required = parser.add_argument_group("required arguments") + input_arg(required) + + required.add_argument( + "--output", + help="Path to save result in zarr or h5ad format", + required=True, + ) + required.add_argument( + "--reference-query", + help="Query to extract reference observations (e.g. gene_symbol=='NTC')", + ) + parser.add_argument( + "--by", + help="Further align control and treatments in each group, using the covariance matrix of all negative " + "(reference) controls as the target and the covariance matrix of each group of negative controls " + "as the source.", + nargs="*", + ) + filter_args(parser) + + common_args(parser, pre_rechunk=True, post_rechunk=True, dask_client_value="none") + parser.set_defaults(func=_run_tvn) + + +def _create_recall_parser( + subparsers: argparse.ArgumentParser, default_help: bool +) -> None: + parser = subparsers.add_parser( + "recall", + help="Run recall", + formatter_class=( + argparse.ArgumentDefaultsHelpFormatter + if default_help + else argparse.HelpFormatter + ), + ) + required = parser.add_argument_group("required arguments") + input_arg(required) + + required.add_argument( + "--output", + help="Path to save result in Parquet format", + required=True, + ) + required.add_argument( + "--ground-truth-corum", + help="Path(s) to ground truth datasets from CORUM", + nargs="+", + ) + required.add_argument( + "--threshold", + help="Recall threshold", + nargs="+", + type=float, + default=[0.99, 0.95, 0.01, 0.05], + ) + + common_args( + parser, + metadata=False, + pre_rechunk=False, + post_rechunk=False, + dask_client_value="none", + ) + parser.set_defaults(func=_run_recall) + + +def _create_pca_parser(subparsers: argparse.ArgumentParser, default_help: bool) -> None: + parser = subparsers.add_parser( + "pca", + help="Run PCA", + formatter_class=( + argparse.ArgumentDefaultsHelpFormatter + if default_help + else argparse.HelpFormatter + ), + ) + required = parser.add_argument_group("required arguments") + input_arg(required) + + required.add_argument( + "--output", + help="Path to save result in zarr or h5ad format", + required=True, + ) + filter_args(parser) + + parser.add_argument( + "--whiten", + action="store_true", + help="When True the components vectors are multiplied by the " + "square root of n_samples and then divided by the singular " + "values to ensure uncorrelated outputs with unit " + "component-wise variances.", + ) + parser.add_argument( + "--components", type=int, default=128, help="Number of principal components" + ) + parser.add_argument( + "--batch-size", + type=int, + default=500_000, + help="Number of samples to use for each batch", + ) + common_args(parser, pre_rechunk=True, post_rechunk=True) + parser.set_defaults(func=_run_pca) + + +def _create_normalize_parser( + subparsers: argparse.ArgumentParser, default_help: bool +) -> None: + parser = subparsers.add_parser( + "normalize", + help="Normalize features", + formatter_class=( + argparse.ArgumentDefaultsHelpFormatter + if default_help + else argparse.HelpFormatter + ), + ) + required = parser.add_argument_group("required arguments") + input_arg(required) + + required.add_argument( + "--output", + help="Path to save normalized features in zarr, h5ad, or Parquet format", + required=True, + ) + filter_args(parser) + + parser.add_argument( + "--by", + help="Stratify by groups when normalizing.", + nargs="*", + ) + + parser.add_argument( + "--reference", + help="Reference expression to normalize to (e.g. gene_symbol=='NTC').", + ) + + parser.add_argument( + "--robust", + help="Use robust statistics for normalization.", + action="store_true", + ) + parser.add_argument( + "--method", + help="Normalization method", + choices=["zscore", "local-zscore"], + default="zscore", + ) + parser.add_argument( + "--neighbors", + help="Number of neighbors for local z-score", + default=100, + type=int, + ) + parser.add_argument( + "--no-centering", + help="Do not center the data before scaling.", + action="store_true", + ) + parser.add_argument( + "--no-scaling", + help="Do not scale the data by dividing by standard deviation.", + action="store_true", + ) + + parser.add_argument( + "--mad-scale-factor", + help="Numerical scale factor to divide median absolute deviation. " + "The string “normal” is also accepted, and results in scale being the" + " inverse of the standard normal quantile function at 0.75", + default="normal", + type=str, + ) + parser.add_argument( + "--max-value", help="Truncate to this value after scaling", type=float + ) + parser.add_argument( + "--batch-size", + help="Batch size to use for local z-score scaling to conserve memory", + default=25000, + type=int, + ) + parser.add_argument( + "--centroid-columns", + help="Columns for y and x centroids to use for local zscore.", + default=["Nuclei_AreaShape_Center_Y", "Nuclei_AreaShape_Center_X"], + nargs=2, + ) + + common_args(parser, pre_rechunk=True, post_rechunk=True) + parser.set_defaults(func=_run_norm_features) + + +def _create_filter_parser( + subparsers: argparse.ArgumentParser, default_help: bool +) -> None: + parser = subparsers.add_parser( + "filter", + help="Filter labels and features", + description="Filter labels and features.", + formatter_class=( + argparse.ArgumentDefaultsHelpFormatter + if default_help + else argparse.HelpFormatter + ), + ) + required = parser.add_argument_group("required arguments") + input_arg(required) + required.add_argument( + "--output", + help="Path to save result in zarr or h5ad format", + required=True, + ) + filter_args(parser) + parser.add_argument( + "--min-feature-variance", + type=float, + help="Minimum median feature variance across `by` to retain a feature.", + ) + parser.add_argument( + "--max-feature-variance", + type=float, + help="Maximum median feature variance across `by` to retain a feature.", + ) + parser.add_argument( + "--n-features", + type=int, + help="Select top n features by variance.", + ) + parser.add_argument( + "--no-scale", + action="store_true", + help="Do not min-max scale each feature before computing variance.", + ) + + parser.add_argument( + "--max-cell-fraction-not-finite", + default=0.25, + type=float, + help="Maximum fraction of non-finite values allowed per cell", + ) + + parser.add_argument( + "--by", + help="Metadata column(s) in dataset to stratify min-max scaling and variance computation (e.g. plate well).", + nargs="*", + ) + common_args( + parser, + pre_rechunk=True, + post_rechunk=True, + rechunk_features="auto", + rechunk_labels="auto", + ) + parser.set_defaults( + func=_run_filter_data, + ) + + +def _run_rank_features(arguments: argparse.Namespace): + from scallops.cli.pert_map import run_rank_features + + run_rank_features(arguments) + + +def _create_rank_parser( + subparsers: argparse.ArgumentParser, default_help: bool +) -> None: + parser = subparsers.add_parser( + "rank", + help="Rank features from output of `merge` command", + formatter_class=( + argparse.ArgumentDefaultsHelpFormatter + if default_help + else argparse.HelpFormatter + ), + ) + required = parser.add_argument_group("required arguments") + input_arg(required) + + required.add_argument( + "--output", + help="Path to Parquet file containing ranked features.", + ) + filter_args(parser) + parser.add_argument( + "--rank-method", + help="Method to rank features", + choices=["welch_t", "student_t", "mannwhitney"], + default="welch_t", + ) + + parser.add_argument( + "--iqr-multiplier", + help="Include values between Q25 - multiplier * IQR and Q75 - multiplier * IQR", + type=float, + ) + + parser.add_argument( + "--perturbation", + help="Field name to group perturbations", + default="gene_symbol", + ) + parser.add_argument( + "--reference", + help="Reference value in `perturbation` to compare against.", + required=True, + ) + + parser.add_argument( + "--by", + help="Stratify by groups when ranking.", + nargs="*", + ) + + parser.add_argument( + "--min-labels", + help="Require at least `min-labels` to include perturbation", + default=10, + type=int, + ) + + common_args(parser, pre_rechunk=True, post_rechunk=False) + parser.set_defaults(func=_run_rank_features) + + +def _create_parser(subparsers: argparse.ArgumentParser, default_help: bool): + parser = subparsers.add_parser( + "pert-map", + help="Perturbation map processing", + description="Perturbation map processing.", + formatter_class=( + argparse.ArgumentDefaultsHelpFormatter + if default_help + else argparse.HelpFormatter + ), + ) + subparsers = parser.add_subparsers(help="Sub-command help.") + _create_filter_parser(subparsers, default_help) + _create_normalize_parser(subparsers, default_help) + _create_rank_parser(subparsers, default_help) + _create_pca_parser(subparsers, default_help) + _create_tvn_parser(subparsers, default_help) + _create_aggregate_parser(subparsers, default_help) + _create_similarity_matrix_parser(subparsers, default_help) + _create_recall_parser(subparsers, default_help) diff --git a/scallops/cli/pooled_if_sbs.py b/scallops/cli/pooled_if_sbs.py index 9f30f60a..095ceaec 100644 --- a/scallops/cli/pooled_if_sbs.py +++ b/scallops/cli/pooled_if_sbs.py @@ -46,6 +46,7 @@ _images2fov, _set_up_experiment, _to_parquet, + is_anndata, is_parquet_file, ) from scallops.reads import ( @@ -77,7 +78,6 @@ _get_sep, _get_store_path, _write_zarr_image, - is_anndata_zarr, open_ome_zarr, read_ome_zarr_array, ) @@ -519,7 +519,7 @@ def merge_sbs_phenotype_pipeline( output_file = f"{output_dir}{image_key}.{output_format}" if not force and ( (output_format == "parquet" and is_parquet_file(output_file)) - or (output_format == "zarr" and is_anndata_zarr(output_file)) + or (output_format == "zarr" and is_anndata(output_file)) ): logger.info(f"Skipping merge for {image_key}") return [] diff --git a/scallops/cli/rank_features.py b/scallops/cli/rank_features.py deleted file mode 100644 index 149d4851..00000000 --- a/scallops/cli/rank_features.py +++ /dev/null @@ -1,139 +0,0 @@ -"""Module for the Command-Line Interface (CLI) related to ranking features. - -Authors: - - The SCALLOPS development team -""" - -import argparse -import json -import os - -import dask.dataframe as dd -import fsspec -import pyarrow as pa -import pyarrow.parquet as pq - -from scallops.cli.util import ( - _create_dask_client, - _create_default_dask_config, - _dask_workers_threads, - _get_cli_logger, - cli_metadata, - load_json, -) -from scallops.features.rank import rank_features -from scallops.features.util import ( - _join_metadata, - _query_anndata, - _read_data, - _slice_anndata, -) -from scallops.io import _to_parquet, is_parquet_file - -logger = _get_cli_logger() - - -def run_pipeline_rank_features(arguments: argparse.Namespace): - paths = arguments.input - features = arguments.features - method = arguments.rank_method - label_filter = arguments.label_filter - rank_output = arguments.output - if rank_output is None: - rank_output = os.path.splitext(os.path.basename(paths[0]))[0] + ".parquet" - if len(paths) > 1: - logger.info(f"Saving results to {rank_output}") - by = arguments.by - - join_path = arguments.metadata - join_fields = arguments.join - if join_path is not None and join_fields is None: - raise ValueError("Please specify join fields") - - force = arguments.force - no_version = arguments.no_version - perturbation_column = arguments.perturbation - min_labels = arguments.min_labels - reference_value = arguments.reference - iqr_multiplier = arguments.iqr_multiplier - - dask_server_url = arguments.client - dask_cluster_parameters = ( - load_json(arguments.dask_cluster) if arguments.dask_cluster is not None else {} - ) - if dask_server_url is None and arguments.dask_cluster is None: - dask_cluster_parameters = _dask_workers_threads() - - if not rank_output.lower().endswith( - ".parquet" - ) and not rank_output.lower().endswith(".pq"): - rank_output = rank_output + ".parquet" - - if not force and is_parquet_file(rank_output): - logger.info( - f"{rank_output} already exists, skipping. Use --force to overwrite." - ) - return - - metadata = {} - if not no_version: - metadata.update(cli_metadata()) - with ( - _create_default_dask_config(), - _create_dask_client(dask_server_url, **dask_cluster_parameters), - ): - data = _read_data(paths, features) - # columns_needed = set() - # columns_needed.add(perturbation_column) - # if by is not None: - # columns_needed.update(by) - # if label_filter is not None: - # columns_needed.update(_get_names_from_pd_query(label_filter)) - # if join_path is not None: - # columns_needed.update(join_fields) - # _load_coords(data, list(columns_needed)) - - if label_filter is None: - label_filter = f"~`{perturbation_column}`.isna()" - data = _slice_anndata(data, _query_anndata(data, label_filter).index) - - if join_path is not None: - _join_metadata( - data, - dd.read_csv(join_path) - if not join_path.lower().endswith(".parquet") - or join_path.lower().endswith(".pq") - else dd.read_parquet(join_path), - join_fields, - ) - rank_df = rank_features( - data=data, - by=by, - perturbation_column=perturbation_column, - reference_value=reference_value, - method=method, - min_labels=min_labels, - iqr_multiplier=iqr_multiplier, - ) - if isinstance(rank_df, dd.DataFrame): - _to_parquet( - rank_df, - rank_output, - write_index=False, - custom_metadata=dict(scallops=json.dumps(metadata)), - ) - else: - table = pa.Table.from_pandas(rank_df, preserve_index=False) - table = table.replace_schema_metadata( - { - "scallops".encode(): json.dumps(metadata).encode(), - **table.schema.metadata, - } - ) - - fs, rank_output = fsspec.url_to_fs(rank_output) - pq.write_table( - table, - rank_output, - filesystem=fs, - ) diff --git a/scallops/cli/rank_features_main.py b/scallops/cli/rank_features_main.py deleted file mode 100644 index 6ab4435c..00000000 --- a/scallops/cli/rank_features_main.py +++ /dev/null @@ -1,98 +0,0 @@ -import argparse - -from scallops.cli.arg_parser import _sort_groups -from scallops.cli.util import ( - dask_client_arg, - dask_cluster_arg, - force_arg, - no_version_arg, -) - - -def _run_rank_features(arguments: argparse.Namespace): - from scallops.cli.rank_features import run_pipeline_rank_features - - run_pipeline_rank_features(arguments) - - -def _create_parser(subparsers: argparse.ArgumentParser, default_help: bool) -> None: - parser = subparsers.add_parser( - "rank-features", - help="Rank features from output of `merge` command", - formatter_class=( - argparse.ArgumentDefaultsHelpFormatter - if default_help - else argparse.HelpFormatter - ), - ) - required = parser.add_argument_group("required arguments") - required.add_argument( - "-i", "--input", help="Path to normalized file(s)", required=True, nargs="+" - ) - - required.add_argument( - "--output", - help="Path to Parquet file containing ranked features.", - ) - - parser.add_argument( - "--features", - help="Features to include. If not specified, all features are used.", - nargs="*", - ) - parser.add_argument( - "--rank-method", - help="Method to rank features", - choices=["welch_t", "student_t", "mannwhitney"], - default="welch_t", - ) - parser.add_argument( - "--label-filter", - help="Expression to filter labels (e.g. barcode_Q_mean_0/barcode_Q_mean > 0.5)", - ) - parser.add_argument( - "--iqr-multiplier", - help="Include values between Q25 - multiplier * IQR and Q75 - multiplier * IQR", - type=float, - ) - - parser.add_argument( - "--perturbation", - help="Field name to group perturbations", - default="gene_symbol", - ) - parser.add_argument( - "--reference", - help="Reference value in `perturbation` to compare against.", - required=True, - ) - - parser.add_argument( - "--by", - help="Stratify by groups when ranking.", - nargs="*", - ) - - parser.add_argument( - "--min-labels", - help="Require at least `min-labels` to include perturbation", - default=10, - type=int, - ) - - parser.add_argument( - "--metadata", - help="Path to CVS or Parquet file containing metadata to join with merged data.", - ) - parser.add_argument( - "--join", - help="Field(s) to join on", - nargs="*", - ) - - dask_client_arg(parser) - dask_cluster_arg(parser) - force_arg(parser) - no_version_arg(parser) - _sort_groups(parser) - parser.set_defaults(func=_run_rank_features) diff --git a/scallops/features/agg.py b/scallops/features/agg.py index f8d37930..c85ef954 100644 --- a/scallops/features/agg.py +++ b/scallops/features/agg.py @@ -114,16 +114,15 @@ def weighted_agg(x): index=groups, ) obs = result.coords["obs"].to_dataframe() - obs = obs.join(group_counts, rsuffix="_1").reset_index(drop=True) + obs = obs.join(group_counts, rsuffix="_1") + # both index and column set if group_by_multi: - new_obs = pd.DataFrame(obs["obs"].tolist(), columns=by) - for c in obs.columns: - if c.startswith("count") and c not in new_obs.columns: - new_obs[c] = obs[c] - obs = new_obs + obs[by] = obs["obs"].apply(pd.Series) + obs.index = obs["obs"].map(lambda x: "-".join(map(str, x))) + obs = obs.drop("obs", axis=1) else: obs = obs.rename({"obs": by}, axis=1) - obs = obs.set_index(pd.RangeIndex(len(obs)).astype(str)) + return anndata.AnnData( X=X, obs=obs, diff --git a/scallops/features/decomposition.py b/scallops/features/decomposition.py index a15dcba6..d61b1dc7 100644 --- a/scallops/features/decomposition.py +++ b/scallops/features/decomposition.py @@ -1,89 +1,18 @@ import logging -from collections.abc import Sequence -from functools import partial import anndata -import dask import dask.array as da from array_api_compat import get_namespace from sklearn.utils import gen_batches -from scallops.features.util import _anndata_to_xr +from scallops.utils import tqdm_func logger = logging.getLogger("scallops") -def _centerscale( - data: anndata.AnnData, - min_std: float | None = 0, - standardize: bool = True, - standardize_by: str | Sequence[str] | None = None, - max_value: float | None = None, -) -> anndata.AnnData: - is_dask = isinstance(data.X, da.Array) - xp = get_namespace(data.X) - if standardize and standardize_by is not None: - xdata = _anndata_to_xr(data, standardize_by) - - def _standardize(x, min_std, max_value): - std = x.std(dim="obs") - if min_std is not None and min_std > 0: - std = std.where(std.data > min_std) - x = (x - x.mean(dim="obs")) / std - if max_value is not None: - x = x.clip(-max_value, max_value) - return x - - xdata = xdata.groupby(standardize_by).map( - partial(_standardize, min_std=min_std, max_value=max_value) - ) - X = xdata.data - no_nans_per_feature = xp.isnan(X).sum(axis=0) == 0 - - if is_dask: - no_nans_per_feature = no_nans_per_feature.compute() - X = X[:, no_nans_per_feature] - logger.info(f"# of features {X.shape[1]:,} / {data.X.shape[1]:,}") - return anndata.AnnData( - X=X, - obs=data.obs.loc[xdata.coords["obs"].values], - var=data.var[no_nans_per_feature], - ) - else: - X = data.X - var = data.var - means = None - stds = None - if standardize or min_std is not None: - means = X.mean(axis=0, keepdims=True) - stds = X.std(axis=0, keepdims=True) - if min_std is not None: - if is_dask: - means, stds = dask.compute(means, stds) - features_keep = stds > min_std - features_keep = features_keep.squeeze() - - X = X[:, features_keep] - - stds = stds[:, features_keep] - means = means[:, features_keep] - var = data.var[features_keep] - logger.info(f"# of features {X.shape[1]:,} / {data.X.shape[1]:,}") - if standardize: - X = (X - means) / stds - if max_value is not None: - X = xp.clip(X, -max_value, max_value) - - return anndata.AnnData(X=X, obs=data.obs.copy(), var=var) - - def pca( data: anndata.AnnData, n_components: int | float | None = None, - min_std: float | None = 0, - standardize: bool = True, - standardize_by: str | Sequence[str] | None = None, - max_value: float | None = None, batch_size: int | None = None, gpu: bool | None = None, whiten: bool = False, @@ -92,26 +21,13 @@ def pca( """Embed data using PCA. :param data: AnnData object. - :param standardize: Whether to standardize the data. - :param standardize_by: Standardize the data specified groups :param n_components: Number of PCA components. - :param min_std: Remove features with standard deviation <= `min_std` after - standardization. - :param max_value: Clip to this value after standardizing :param batch_size: Batch size for incremental PCA. :param gpu: Whether to use GPU. :param whiten: Whether to use whitening. :param progress: Whether to show progress bar for incremental PCA. :return: PCA Embedding """ - if standardize: - data = _centerscale( - data=data, - min_std=min_std, - standardize=standardize, - standardize_by=standardize_by, - max_value=max_value, - ) X = data.X is_dask = isinstance(data.X, da.Array) if gpu is None: @@ -132,15 +48,8 @@ def pca( d = IncrementalPCA(n_components=n_components, whiten=whiten, copy=not is_dask) batches = list(gen_batches(X.shape[0], batch_size, min_batch_size=n_components)) - - if progress: - try: - from tqdm import tqdm - except ImportError: - from scallops.utils import _tqdm_shim as tqdm - else: - from scallops.utils import _tqdm_shim as tqdm - for batch in tqdm(batches): + tqdm, progress_args = tqdm_func(progress) + for batch in tqdm(batches, **progress_args): X_batch = X[batch] if is_dask: X_batch = X_batch.compute() diff --git a/scallops/features/normalize.py b/scallops/features/normalize.py index a15d9952..151a863e 100644 --- a/scallops/features/normalize.py +++ b/scallops/features/normalize.py @@ -5,13 +5,18 @@ import anndata import dask.array as da import numpy as np +import pandas as pd import scipy import xarray as xr -from dask.delayed import delayed +from anndata._core.index import _normalize_index +from array_api_compat import get_namespace +from flox import rechunk_for_blockwise +from flox.lib import _issorted from sklearn.decomposition import PCA from sklearn.neighbors import NearestNeighbors -from scallops.features.util import _anndata_to_xr +from scallops.features.util import _slice_anndata +from scallops.utils import tqdm_func logger = logging.getLogger("scallops") @@ -25,70 +30,401 @@ def _convert_scale(mad_scale): return mad_scale -def _normalize_features_array( - values: np.ndarray | da.Array, - reference_values: np.ndarray | da.Array, - indices: np.ndarray | None, - mad_scale: float | str, - robust: bool, - scaling: bool, - centering: bool, - max_value: float | None, -): - """Normalize 2d labels by features array. - - :param values: Array of values to normalize - :param reference_values: Array of reference values - :param indices: Array of nearest neighbor indices for local-zscore - :param mad_scale: The numerical value of mad_scale will be divided out of the final - result of the median absolute deviation. The default is 1.0. The string - "normal" is also accepted,and results in `mad_scale` being the inverse of the - standard normal quantile function at 0.75, which is approximately 0.67449 - :param robust: Use robust statistics - :return: Array of normalized values +def normalize_features( + data: anndata.AnnData, + reference_query: str | None = None, + by: Sequence[str] | str | None = None, + normalize: Literal["zscore", "local-zscore"] = "zscore", + n_neighbors: int = 100, + neighbors_metric: str = "minkowski", + robust: bool = False, + mad_scale: float | str = "normal", + max_value: float | None = None, + centering: bool = True, + scaling: bool = True, + batch_size: int | None = 25000, + centroid_column_names: tuple[str, str] = ( + "Nuclei_AreaShape_Center_Y", + "Nuclei_AreaShape_Center_X", + ), +) -> anndata.AnnData: + """Normalize features + :param data: Annotated data matrix. + :param reference_query: Query to extract reference observations + (e.g. "gene_symbol=='NTC'") + :param by: Column(s) in `data.obs` to stratify by. + :param normalize: Normalization method to use where `local` uses nearest neighbors by location. + :param n_neighbors: Number of neighbors for local and nearest neighbor zscore. + :param neighbors_metric: Nearest neighbor metric to use when normalize is + `local-zscore`. + :param robust: Use robust statistics. + :param mad_scale: Numerical scale factor to divide median absolute deviation. The + string “normal” is also accepted, and results in scale being the inverse of the + standard normal quantile function at 0.75 + :param centering: Whether to center the data before scaling. + :param max_value: Truncate to this value after scaling + :param scaling: Whether to scale the data by dividing by the standard deviation. + :param batch_size: Batch size to use for local z-score scaling to conserve memory. + :param centroid_column_names: Columns for y and x centroids to use for local zscore. + :return: Normalized data """ - if isinstance(values, da.Array): - chunks = list(values.chunksize) - if chunks[0] != values.shape[0]: - chunks[0] = -1 - chunks = tuple(chunks) - if chunks != values.chunksize: - values = values.rechunk(chunks) - if reference_values is not None: - ref_chunks = list(reference_values.chunksize) - if ref_chunks[0] != reference_values.shape[0]: - ref_chunks[0] = -1 - if ref_chunks[1] != chunks[1]: - ref_chunks[1] = chunks[1] - ref_chunks = tuple(ref_chunks) - if ref_chunks != reference_values.chunksize: - reference_values = reference_values.rechunk(ref_chunks) - arrays = ( - (values, reference_values) if reference_values is not None else (values,) + + mad_scale = _convert_scale(mad_scale) + centroid_column_names = list(centroid_column_names) + is_dask = isinstance(data.X, da.Array) + has_sorted_groups = False + if by is not None: + by_multi = not isinstance(by, str) and isinstance(by, Sequence) + if by_multi: + by = list(by) + if len(by) == 1: + by = by[0] + by_multi = False + + groupby_values = ( + data.obs[by].apply(tuple, axis=1) if by_multi else data.obs[by].values + ) + series = pd.Series(groupby_values, dtype="category") + has_sorted_groups = ( + normalize == "local-zscore" + and is_dask + and len(data.X.chunks[0]) > 1 + and _issorted(series.cat.codes.values) + ) + + group_indices = series.groupby( + series, observed=True, sort=False, dropna=False + ).indices + else: + group_indices = {None: None} + if normalize == "zscore": + coords = {} + if by is not None: + coords["obs"] = groupby_values + xdata = xr.DataArray(data.X, dims=["obs", "var"], coords=coords) + x_ref_data = xdata + if reference_query is not None: + refererence_data = _slice_anndata( + data, data.obs.query(reference_query).index + ) + coords = dict() + if by is not None: + coords["obs"] = ( + refererence_data.obs[by].apply(tuple, axis=1) + if by_multi + else refererence_data.obs[by].values + ) + x_ref_data = xr.DataArray( + refererence_data.X, + dims=["obs", "var"], + coords=coords, + ) + kwargs = dict() + if by is not None: + grouped_ref = x_ref_data.groupby("obs") + grouped_values = ( + xdata.groupby("obs") if reference_query is not None else grouped_ref + ) + else: + kwargs["dim"] = "obs" + grouped_ref = x_ref_data + grouped_values = xdata + + means = None + stds = None + + if robust: + xp = get_namespace(data.X) + if centering: + means = grouped_ref.median(**kwargs) + if scaling: + diff = xp.abs(grouped_ref - means) + stds = diff.median(axis=0) / mad_scale + else: + if centering: + means = grouped_ref.mean(**kwargs) + if scaling: + stds = grouped_ref.std(**kwargs) + if centering: + grouped_values = grouped_values - means + + if scaling: + if by is not None and isinstance(groupby_values, xr.DataArray): + grouped_values = grouped_values.groupby("obs") + grouped_values = grouped_values / stds + if max_value is not None: + grouped_values = grouped_values.clip(-max_value, max_value) + return anndata.AnnData( + X=grouped_values.data, + obs=data.obs.copy(), + var=data.var.copy(), + uns=data.uns.copy(), + obsm=data.obsm.copy(), + varm=data.varm.copy(), ) + indices = [] if has_sorted_groups else None + results = [] if not has_sorted_groups else None + obs_list = [] if not has_sorted_groups else None + for key in group_indices.keys(): + group_indices_ = None + if by is not None: + group_indices_ = group_indices[key] + if not has_sorted_groups: + array_subset = group_indices_ + if np.all(np.diff(group_indices_) == 1): + array_subset = slice(group_indices_[0], group_indices_[-1] + 1) + x = data.X[array_subset] + df = data.obs.iloc[group_indices_] + else: + if not has_sorted_groups: + x = data.X + df = data.obs + + global_reference_indices = None + local_reference_indices = None + if reference_query is not None: + local_reference_indices = _normalize_index( + df.query(reference_query).index, df.index + ) + if has_sorted_groups: + global_reference_indices = ( + group_indices_[local_reference_indices] + if group_indices_ is not None + else local_reference_indices + ) + + if normalize == "local-zscore": + query_coordinates = df[centroid_column_names].values + reference_coordinates = ( + df.iloc[local_reference_indices][centroid_column_names].values + if local_reference_indices is not None + else query_coordinates + ) + nn_indices = _nearest_neighbors_indices( + query=query_coordinates, + reference=reference_coordinates, + n_neighbors=n_neighbors, + metric=neighbors_metric, + ) + if has_sorted_groups: + if global_reference_indices is not None: + nn_indices = global_reference_indices[nn_indices] + indices.append(nn_indices) + else: + if local_reference_indices is not None: + nn_indices = local_reference_indices[nn_indices] + if is_dask: + nn_indices = da.from_array(nn_indices) + + # memory = (x.shape[0] * x.shape[1] * n_neighbors) / batch_size + (x.shape[0] * x.shape[1]) + # memory *= 8 + result = _local_z_batched( + x=x, + nn_indices=nn_indices, # indices into x + robust=robust, + mad_scale=mad_scale, + centering=centering, + scaling=scaling, + max_value=max_value, + batch_size=batch_size, + ) - return da.map_blocks( - _normalize_features_np, - *arrays, - indices=delayed(indices), + # else: + # if has_sorted_groups: + # if global_reference_indices is not None: + # indices.append(global_reference_indices) + # else: + # if is_dask and local_reference_indices is not None: + # local_reference_indices = da.from_array(local_reference_indices) + # + # result = _normalize_features_array( + # values=x, + # reference_indices=local_reference_indices, + # robust=robust, + # mad_scale=mad_scale, + # centering=centering, + # scaling=scaling, + # max_value=max_value, + # local_zscore=False, + # ) + + if not has_sorted_groups: + results.append(result) + obs_list.append(df) + + if has_sorted_groups: + rechunked_data = rechunk_for_blockwise(data.X, 0, series.cat.codes.values)[1] + indices = np.concatenate(indices, axis=0) + chunks = [(rechunked_data.chunks[0])] + for s in indices.shape[1:]: + chunks.append((s,)) + + indices = da.from_array(indices, chunks=tuple(chunks)) + assert indices.shape[0] == rechunked_data.shape[0] + kwargs = dict( robust=robust, mad_scale=mad_scale, + centering=centering, scaling=scaling, + max_value=max_value, + ) + + result = da.map_blocks( + _local_z_batched, rechunked_data, indices, **kwargs, dtype=np.float64 + ) + return anndata.AnnData( + X=result, + obs=data.obs.copy(), + var=data.var.copy(), + uns=data.uns.copy(), + obsm=data.obsm.copy(), + varm=data.varm.copy(), + ) + return anndata.AnnData( + X=get_namespace(data.X).vstack(results), + obs=pd.concat(obs_list), + var=data.var.copy(), + uns=data.uns.copy(), + obsm=data.obsm.copy(), + varm=data.varm.copy(), + ) + + +def _local_z_batched( + x: np.ndarray | da.Array, + nn_indices: np.ndarray | da.Array, + scaling: bool = True, + centering: bool = True, + max_value: float | None = None, + mad_scale: float | str = "normal", + robust: bool = False, + batch_size: int | None = None, + progress: bool | str = False, + block_info=None, +): + if isinstance(x, da.Array): + batch_size = None + + if batch_size is None: + batch_size = x.shape[0] + result_arrays = [] + if block_info is not None: + array_location = block_info[0]["array-location"][0] + nn_indices = nn_indices - array_location[0] + tqdm, progress_args = tqdm_func(progress) + for batch in tqdm(range(0, x.shape[0], batch_size), **progress_args): + sl = slice(batch, batch + batch_size) + nn_indices_ = nn_indices[sl] + x_ = x[sl] + if isinstance(x, da.Array): + n_labels = nn_indices_.shape[0] + n_neighbors = nn_indices_.shape[1] + nn_indices_ = nn_indices_.flatten() + if not isinstance(nn_indices_, da.Array): + nn_indices_ = da.from_array(nn_indices_) + # (labels,neighbors,features) + reference_data_ = x[nn_indices_].reshape((n_labels, n_neighbors, -1)) + else: + reference_data_ = x[nn_indices_] + result = _normalize_features_array( + values=x_, + # reference_indices=reference_data_, + reference_values=reference_data_, + robust=robust, + mad_scale=mad_scale, centering=centering, + scaling=scaling, max_value=max_value, - meta=values._meta, + local_zscore=nn_indices is not None, + ) + result_arrays.append(result) + return ( + get_namespace(x).vstack(result_arrays) + if len(result_arrays) > 1 + else result_arrays[0] + ) + + +def _normalize_features_array( + values: np.ndarray | da.Array, + reference_indices: np.ndarray | da.Array | None = None, + reference_values: np.ndarray | da.Array | None = None, + scaling: bool = True, + centering: bool = True, + max_value: float | None = None, + local_zscore: bool = False, + mad_scale: float | str = "normal", + robust: bool = False, +): + mad_scale = _convert_scale(mad_scale) if robust else None + xp = get_namespace(values) + if reference_values is None: + reference_values = ( + values if reference_indices is None else values[reference_indices] ) + means = None + stds = None + if not local_zscore: + if robust: + if centering: + means = xp.nanmedian(reference_values, axis=0) + if scaling: + stds = ( + xp.nanmedian(xp.abs(reference_values - means), axis=0) / mad_scale + ) + else: + if centering: + means = xp.nanmean(reference_values, axis=0) + if scaling: + stds = xp.nanstd(reference_values, axis=0) + if centering: + means = xp.expand_dims(means, 0) + if scaling: + stds = xp.expand_dims(stds, 0) + else: + # reference_values dims are (labels,neighbors,features) + if robust: + means = xp.nanmedian(reference_values, axis=1) + + if scaling: + stds = ( + xp.nanmedian( + xp.abs(reference_values - xp.expand_dims(means, axis=1)), + axis=1, + ) + / mad_scale + ) + + else: + if centering: + means = xp.nanmean(reference_values, axis=1) + if scaling: + stds = xp.nanstd(reference_values, axis=1) + + if centering: + values = values - means + if scaling: + stds[stds == 0] = 1.0 + values = values / stds + if max_value is not None: + values = xp.clip(values, -max_value, max_value) + return values + - return _normalize_features_np( - values=values, - ref_values=reference_values, - indices=indices, - robust=robust, - mad_scale=mad_scale, - scaling=scaling, - max_value=max_value, - centering=centering, +def _nearest_neighbors_indices( + reference: np.ndarray, + query: np.ndarray, + n_neighbors: int = 100, + metric: str = "minkowski", +) -> np.ndarray: + if n_neighbors > len(reference): + raise ValueError(f"n_neighbors: {n_neighbors}, n points: {len(reference)}") + # shape is reference.shape[0], n_neighbors + return ( + NearestNeighbors(n_neighbors=n_neighbors, metric=metric) + .fit(reference) + .kneighbors(query, return_distance=False) ) @@ -112,19 +448,21 @@ def typical_variation_normalization( """ # Adapted from EFAAR_benchmarking _ X = data.X - ref_indices = data.obs.index.get_indexer_for(data.obs.query(reference_query).index) + reference_indices = data.obs.index.get_indexer_for( + data.obs.query(reference_query).index + ) X = _normalize_features_array( X, - X[ref_indices], - indices=None, + reference_values=X[reference_indices], robust=False, mad_scale="normal", centering=True, scaling=True, max_value=None, + local_zscore=False, ) d = PCA() - X = d.fit(X[ref_indices]).transform(X) + X = d.fit(X[reference_indices]).transform(X) components_ = d.components_ mean_ = d.mean_ variance_ratio = d.explained_variance_ratio_ @@ -134,11 +472,13 @@ def typical_variation_normalization( group_to_indices = data.obs.groupby(by, observed=True, sort=False).indices for group in group_to_indices.keys(): group_indices = group_to_indices[group] - group_control_indices = group_indices[np.isin(group_indices, ref_indices)] + group_control_indices = group_indices[ + np.isin(group_indices, reference_indices) + ] X[group_indices] = _normalize_features_array( X[group_indices], - X[group_control_indices], - indices=None, + reference_values=X[group_control_indices], + local_zscore=False, robust=False, mad_scale="normal", centering=True, @@ -146,13 +486,15 @@ def typical_variation_normalization( max_value=None, ) - target_cov = np.cov(X[ref_indices], rowvar=False, ddof=1) + 0.5 * np.eye( + target_cov = np.cov(X[reference_indices], rowvar=False, ddof=1) + 0.5 * np.eye( X.shape[1] ) for group in group_to_indices.keys(): group_indices = group_to_indices[group] - group_control_indices = group_indices[np.isin(group_indices, ref_indices)] + group_control_indices = group_indices[ + np.isin(group_indices, reference_indices) + ] source_cov = np.cov( X[group_control_indices], rowvar=False, ddof=1 @@ -167,8 +509,8 @@ def typical_variation_normalization( else: X = _normalize_features_array( X, - X[ref_indices], - indices=None, + reference_values=X[reference_indices], + local_zscore=False, robust=False, mad_scale="normal", centering=True, @@ -188,253 +530,3 @@ def typical_variation_normalization( } }, ) - - -def normalize_features( - data: anndata.AnnData, - reference_query: str | None = None, - by: Sequence[str] | str | None = None, - normalize: Literal["zscore", "local-zscore", "nn-zscore"] = "zscore", - n_neighbors: int | None = 100, - neighbors_metric: str = "minkowski", - robust: bool = False, - mad_scale: float | str = "normal", - max_value: float | None = None, - centering: bool = True, - scaling: bool = True, - batch_size: int | None = None, - centroid_column_names: tuple[str, str] = ( - "Nuclei_AreaShape_Center_Y", - "Nuclei_AreaShape_Center_X", - ), -) -> anndata.AnnData: - """Normalize features - - :param data: Annotated data matrix. - :param reference_query: Query to extract reference observations - (e.g. "gene_symbol=='NTC'") - :param by: Column(s) in `data.obs` to stratify by. - :param normalize: Normalization method to use where `local` uses nearest - neighbors by location and `nn` uses nearest neighbors by `neighbors_metric`. - :param n_neighbors: Number of neighbors for local and nearest neighbor zscore. - :param neighbors_metric: Nearest neighbor metric to use when normalize is - `nn-zscore`. - :param robust: Use robust statistics. - :param mad_scale: Numerical scale factor to divide median absolute deviation. The - string “normal” is also accepted, and results in scale being the inverse of the - standard normal quantile function at 0.75 - :param centering: Whether to center the data before scaling. - :param max_value: Truncate to this value after scaling - :param scaling: Whether to scale the data by dividing by the standard deviation. - :param batch_size: Batch size to use for local scaling to conserve memory. - :param centroid_column_names: Columns for y and x centroids to use for local zscore. - :return: Normalized data - """ - - mad_scale = _convert_scale(mad_scale) - xdata = _anndata_to_xr(data) - if by is not None: - group_result = xdata.groupby(by).map( - lambda x: _normalize_group( - x, - reference_query=reference_query, - normalize=normalize, - n_neighbors=n_neighbors, - neighbors_metric=neighbors_metric, - robust=robust, - max_value=max_value, - mad_scale=mad_scale, - centering=centering, - scaling=scaling, - batch_size=batch_size, - centroid_column_names=centroid_column_names, - ) - ) - - return anndata.AnnData( - X=group_result.data, - obs=data.obs.loc[group_result.coords["obs"].values], - var=data.var.copy(), - ) - - result = _normalize_group( - xdata, - reference_query=reference_query, - normalize=normalize, - n_neighbors=n_neighbors, - neighbors_metric=neighbors_metric, - robust=robust, - max_value=max_value, - mad_scale=mad_scale, - centering=centering, - scaling=scaling, - batch_size=batch_size, - centroid_column_names=centroid_column_names, - ) - return anndata.AnnData(X=result.data, obs=data.obs.copy(), var=data.var.copy()) - - -def _normalize_group( - data: xr.DataArray, - reference_query: str | None, - normalize: Literal["zscore", "local-zscore", "nn-zscore"], - n_neighbors: int | None, - neighbors_metric: str, - robust: bool, - mad_scale: float | str, - centering: bool, - max_value: float | None, - scaling: bool, - batch_size: int | None, - centroid_column_names: tuple[str, str] = ( - "Nuclei_AreaShape_Center_Y", - "Nuclei_AreaShape_Center_X", - ), -) -> xr.DataArray: - indices = None - reference_data = ( - data.query(dict(obs=reference_query)) if reference_query is not None else None - ) - - if reference_data is not None: - if reference_data.shape[0] == 0: - raise ValueError("No reference data found.") - if normalize == "nn-zscore": - # nearest neighbors in PCA space - nn_query = data.data - nn_ref = nn_query - if reference_data is not None: - nn_ref = reference_data.data - indices = _nearest_neighbors_indices( - nn_ref, nn_query, n_neighbors=n_neighbors, metric=neighbors_metric - ) - elif normalize == "local-zscore": - nn_query = np.stack( - ( - data.coords[centroid_column_names[0]].values, - data.coords[centroid_column_names[1]].values, - ), - axis=1, - ) - nn_ref = nn_query - if reference_data is not None: - nn_ref = np.stack( - ( - reference_data.coords[centroid_column_names[0]].values, - reference_data.coords[centroid_column_names[1]].values, - ), - axis=1, - ) - - indices = _nearest_neighbors_indices( - nn_ref, nn_query, n_neighbors=n_neighbors, metric=neighbors_metric - ) - if batch_size is not None and indices is not None and indices.shape[0] > batch_size: - value_list = [] - if reference_data is None: - reference_data = data - - for i in range(0, indices.shape[0], batch_size): - sl = slice(i, i + batch_size) - values = _normalize_features_array( - data.data[sl], - reference_data.data, - indices=indices[sl], - robust=robust, - mad_scale=mad_scale, - centering=centering, - scaling=scaling, - max_value=max_value, - ) - value_list.append(values) - values = np.concatenate(value_list) - else: - values = _normalize_features_array( - data.data, - reference_data.data if reference_data is not None else None, - indices=indices, - robust=robust, - mad_scale=mad_scale, - centering=centering, - scaling=scaling, - max_value=max_value, - ) - return data.copy(data=values, deep=False) - - -def _nearest_neighbors_indices( - reference: np.ndarray, - query: np.ndarray, - n_neighbors: int = 100, - metric: str = "minkowski", -) -> np.ndarray: - if n_neighbors > len(reference): - raise ValueError(f"n_neighbors: {n_neighbors}, n points: {len(reference)}") - return ( - NearestNeighbors(n_neighbors=n_neighbors, metric=metric) - .fit(reference) - .kneighbors(query, return_distance=False) - ) - - -def _normalize_features_np( - values: np.ndarray, - ref_values: np.ndarray | None = None, - indices: np.ndarray | None = None, - mad_scale: float | str = "normal", - centering: bool = True, - scaling: bool = True, - robust: bool = True, - max_value: float | None = None, -) -> np.ndarray: - mad_scale = _convert_scale(mad_scale) - - if ref_values is None: - ref_values = values - means = None - stds = None - if indices is None: - if robust: - if centering: - means = np.nanmedian(ref_values, axis=0) - if scaling: - stds = np.nanmedian(np.abs(ref_values - means), axis=0) / mad_scale - else: - if centering: - means = np.nanmean(ref_values, axis=0) - if scaling: - stds = np.nanstd(ref_values, axis=0) - if centering: - means = np.expand_dims(means, 0) - if scaling: - stds = np.expand_dims(stds, 0) - else: - ref_values = ref_values[indices] - # ref_values dims are (labels,neighbors,features) - if robust: - means = np.nanmedian(ref_values, axis=1) - - if scaling: - stds = ( - np.nanmedian( - np.abs(ref_values - np.expand_dims(means, axis=1)), - axis=1, - ) - / mad_scale - ) - - else: - if centering: - means = np.nanmean(ref_values, axis=1) - if scaling: - stds = np.nanstd(ref_values, axis=1) - - if centering: - values = values - means - if scaling: - stds[stds == 0] = 1.0 - values = values / stds - if max_value is not None: - values[values > max_value] = max_value - values[values < -max_value] = -max_value - return values diff --git a/scallops/features/preprocessing.py b/scallops/features/preprocessing.py index 3850d683..7a7b4483 100644 --- a/scallops/features/preprocessing.py +++ b/scallops/features/preprocessing.py @@ -59,17 +59,16 @@ def _transform_feature_group(x): def feature_variance( - data: anndata.AnnData, by: str | Sequence -) -> np.ndarray | da.Array: + data: anndata.AnnData, by: str | Sequence, scale: bool = True +) -> xr.DataArray: """Compute feature variance stratified by column(s) in `data.obs` :param data: AnnData object :param by: Column(s) in `data.obs` to stratify by when computing variance. - :return: Median feature variance + :param scale: Set to True to apply min-max scaling. + :return: Feature variance """ - xp = get_namespace(data.X) - if not isinstance(by, str) and isinstance(by, Sequence): # xarray outputs all combinations, even ones that don't exist # https://github.com/pydata/xarray/issues/11264 @@ -82,18 +81,29 @@ def feature_variance( by = "obs" else: xdata = _anndata_to_xr(data, by) + grouped = xdata.groupby(by) + if scale: - variance = xdata.groupby(by).var(skipna=False) # dims (by, 'var') - variance = xp.median(variance.data, axis=0) + def single_group(x): + min_value = x.min(skipna=False, dim="obs") + max_value = x.max(skipna=False, dim="obs") + x = (x - min_value) / (max_value - min_value) + return x.var(skipna=False, dim="obs") + + variance = grouped.map(single_group) + else: + variance = grouped.var(skipna=False) # dims (by, 'var') return variance def filter_data( data: anndata.AnnData, max_fraction_not_finite: float | None = 0.25, - min_variance: float | None = 0.1, + min_variance: float | None = None, max_variance: float | None = None, + n_features: int | None = None, by: str | Sequence | None = None, + scale: bool = True, ) -> anndata.AnnData: """Filter cells using `max_fraction_not_finite` then filter features using variance @@ -102,18 +112,21 @@ def filter_data( missing or infinite values :param min_variance: Keep features with variance >= `min_variance` :param max_variance: Keep features with variance <= `max_variance` + :param n_features: Keep `n_features` after applying any min/max filter :param by: Column(s) in `data.obs` to stratify by when computing variance. If provided, the median variance is used for filtering. + :param scale: Set to True to apply min-max scaling per group. :return: Filtered AnnData object """ xp = get_namespace(data.X) keep_cells = None - keep_features = None + variance = None + variance_indices = None if max_fraction_not_finite is not None: invalid_counts_per_cell = (~xp.isfinite(data.X)).sum(axis=1) max_counts = int(data.shape[1] * max_fraction_not_finite) keep_cells = invalid_counts_per_cell <= max_counts - if min_variance is not None or max_variance is not None: + if min_variance is not None or max_variance is not None or n_features is not None: if min_variance is None: min_variance = -np.inf if max_variance is None: @@ -124,17 +137,27 @@ def filter_data( data = _slice_anndata(data, keep_cells) keep_cells = None if by is not None: - variance = feature_variance(data, by) - + variance = feature_variance(data, by, scale) + variance = xp.median(variance.data, axis=0) else: - variance = xp.var(data.X, axis=0) + X = data.X + if scale: + min_value = xp.min(X, axis=0) + max_value = xp.max(X, axis=0) + X = (X - min_value) / (max_value - min_value) + variance = xp.var(X, axis=0) - keep_features = ( - (variance >= min_variance) - & (variance <= max_variance) - & (xp.isfinite(variance)) + if isinstance(data.X, da.Array): + keep_cells, variance = dask.compute(keep_cells, variance) + if variance is not None: + variance_indices = np.argsort(variance) # nans are at end + keep_variance_indices = ( + (variance[variance_indices] >= min_variance) + & (variance[variance_indices] <= max_variance) + & (xp.isfinite(variance[variance_indices])) ) + variance_indices = variance_indices[keep_variance_indices] + if n_features is not None: + variance_indices = variance_indices[:n_features] - if isinstance(data.X, da.Array): - keep_features, keep_cells = dask.compute(keep_features, keep_cells) - return _slice_anndata(data, keep_cells, keep_features) + return _slice_anndata(data, keep_cells, variance_indices) diff --git a/scallops/features/util.py b/scallops/features/util.py index d10c64b1..9d489674 100644 --- a/scallops/features/util.py +++ b/scallops/features/util.py @@ -13,7 +13,6 @@ from pandas.core.computation.parsing import BACKTICK_QUOTED_STRING, tokenize_string from scallops.features.constants import _metadata_columns_whitelist_str -from scallops.io import read_anndata_zarr logger = logging.getLogger("scallops") @@ -186,33 +185,6 @@ def _join_metadata( data.obs = data.obs.join(join_df, on=on) -def _read_data( - paths: Sequence[str] | str, features: Sequence[str] | None = None -) -> anndata.AnnData: - if isinstance(paths, str): - paths = [paths] - assert len(paths) == len(set(paths)), "Duplicate path" - data_arrays = [] - for path in paths: - if path.lower().endswith(".parquet") or path.lower().endswith(".pq"): - df = pd.read_parquet(path) - d = pandas_to_anndata(df, features) - else: - d = read_anndata_zarr(path, dask=True) - if features is not None and len(features) > 0: - d = d[:, features] - data_arrays.append(d) - if len(data_arrays) == 0: - raise RuntimeError("No data found.") - - data = ( - data_arrays[0] - if len(data_arrays) == 1 - else anndata.concat(data_arrays, index_unique="-") - ) - return data - - def _get_names_from_pd_query(source) -> set[str]: tokens = tokenize_string(source) result = set() diff --git a/scallops/io.py b/scallops/io.py index c9dbe121..c413a8ef 100644 --- a/scallops/io.py +++ b/scallops/io.py @@ -34,6 +34,7 @@ import dask.array as da import dask.dataframe as dd import fsspec +import h5py import numpy as np import ome_types import pandas as pd @@ -1514,18 +1515,66 @@ def _subset_include(x): return _subset_include -def read_anndata_zarr(store: StoreLike, dask: bool = False) -> anndata.AnnData: - """Read from a hierarchical Zarr array store. +def is_anndata(store: StoreLike) -> bool: + """Determines whether store is an AnnData Zarr or h5py file. + + :param store: Store to read from. + """ + try: + is_store_arg_h5_store = isinstance(store, h5py.Dataset | h5py.File | h5py.Group) + is_store_arg_h5_path = ( + isinstance(store, os.PathLike | str) and Path(store).suffix == ".h5ad" + ) + is_h5 = is_store_arg_h5_path or is_store_arg_h5_store + + if not is_h5: + import zarr + + if not isinstance(store, zarr.Group): + try: + f = zarr.open_consolidated(store, mode="r") + except ValueError: + f = zarr.open_group(store, mode="r") + else: + f = store + elif is_store_arg_h5_store: + f = store + else: + f = h5py.File(store, mode="r") + return f.get("layers") is not None + except: # noqa: E722 + return False + + +def read_anndata(store: StoreLike, dask: bool = False) -> anndata.AnnData: + """Read from a hierarchical Zarr or HDF5 array store. :param store: Store to read from. :param dask: Whether to use dask. :return: AnnData object. """ + is_store_arg_h5_store = isinstance(store, h5py.Dataset | h5py.File | h5py.Group) + is_store_arg_h5_path = ( + isinstance(store, os.PathLike | str) and Path(store).suffix == ".h5ad" + ) + is_h5 = is_store_arg_h5_path or is_store_arg_h5_store if not dask: - return anndata.read_zarr(store) + return anndata.read_h5ad(store) if is_h5 else anndata.read_zarr(store) + if not is_h5: + import zarr - f = zarr.open(store, mode="r") + if not isinstance(store, zarr.Group): + try: + f = zarr.open_consolidated(store, mode="r") + except ValueError: + f = zarr.open_group(store, mode="r") + else: + f = store + elif is_store_arg_h5_store: + f = store + else: + f = h5py.File(store, mode="r") def callback(func, elem_name: str, elem, iospec): if iospec.encoding_type in ( diff --git a/scallops/tests/conftest.py b/scallops/tests/conftest.py index abec19f9..e030ce1b 100644 --- a/scallops/tests/conftest.py +++ b/scallops/tests/conftest.py @@ -1,5 +1,7 @@ from pathlib import Path +import numpy as np +import pandas as pd import pytest from scallops.io import read_experiment, read_image @@ -18,6 +20,22 @@ ).exists(), "Test files not found. Please ensure you have Git LFS installed" +@pytest.fixture(scope="module", autouse=True) +def test_feature_table(): + return pd.DataFrame( + data=dict( + label=np.arange(6), + Cells_Intensity_feature_1=[1, 2, 4, 8, 16, 32], + Cells_Intensity_feature_2=[10, 20, 40, 80, 160, 320], + gene_symbol=["a", "NTC", "a", "NTC", "a", "NTC"], + well=["a", "a", "a", "b", "b", "b"], + plate=["a", "a", "a", "b", "b", "b"], + Nuclei_AreaShape_Center_Y=[1, 7, 12, 16, 19, 21], + Nuclei_AreaShape_Center_X=[1, 7, 12, 16, 19, 21], + ), + ) + + @pytest.fixture(scope="module", autouse=True) def experiment_c(): return read_experiment( diff --git a/scallops/tests/test_decomposition.py b/scallops/tests/test_decomposition.py index 4d2a3d92..6e423af0 100644 --- a/scallops/tests/test_decomposition.py +++ b/scallops/tests/test_decomposition.py @@ -19,10 +19,6 @@ def test_decomposition(): result = pca( data=adata, n_components=2, - min_std=0, - standardize=True, - standardize_by=["plate", "well"], - max_value=10, progress=False, batch_size=2, ) @@ -37,7 +33,6 @@ def test_decomposition_compare_numpy(): result = pca( data=adata, n_components=2, - standardize=False, progress=False, batch_size=2, ) diff --git a/scallops/tests/test_features_preprocessing.py b/scallops/tests/test_features_preprocessing.py index b6ed0a13..55c82358 100644 --- a/scallops/tests/test_features_preprocessing.py +++ b/scallops/tests/test_features_preprocessing.py @@ -5,7 +5,53 @@ import pytest from sklearn.preprocessing import PowerTransformer -from scallops.features.preprocessing import filter_data, transform_features_yj +from scallops.features.preprocessing import ( + feature_variance, + filter_data, + transform_features_yj, +) + + +@pytest.mark.parametrize("use_dask", [True, False]) +@pytest.mark.features +def test_feature_variance(use_dask): + rng = np.random.default_rng(0) + X = rng.random((20, 2)) + adata = anndata.AnnData( + da.from_array(X) if use_dask else X, + obs=pd.DataFrame( + data=dict( + pert=["pert1", "pert2"] * 10, + well=["well1", "well2"] * 10, + ) + ), + var=pd.DataFrame(index=["gene1", "gene2"]), + ) + if use_dask: + adata2 = adata.copy() + adata2.X = adata2.X.compute() + df = adata2.to_df().join(adata2.obs) + else: + df = adata.to_df().join(adata.obs) + + def single_group(x): + x = x.copy() + for gene in ["gene1", "gene2"]: + value = (x[gene].values - np.min(x[gene])) / ( + np.max(x[gene]) - np.min(x[gene]) + ) + value = np.var(value) + x[gene] = value + + return x.drop_duplicates(["gene1", "gene2"]) + + result_df = ( + df.groupby("well").apply(single_group, include_groups=False).reset_index() + ) + result = feature_variance(adata, by="well", scale=True) + if use_dask: + result = result.compute() + np.testing.assert_almost_equal(result.data, result_df[["gene1", "gene2"]].values) @pytest.mark.parametrize("by", [None, "well"]) @@ -29,18 +75,32 @@ def test_filter_data(use_dask, by): adata.X[0, 0] = np.nan # np.var(adata.X, axis=0) array([nan, 5.], dtype=float32) test_nan_filter = filter_data( - adata, max_fraction_not_finite=0, min_variance=None, max_variance=None + adata, + max_fraction_not_finite=0, + min_variance=None, + max_variance=None, + scale=False, ) assert test_nan_filter.shape == (3, 2) # np.var(adata.X, axis=0) # array([nan, 5.] # np.var(adata[adata.obs['well'] == 'well1'].X, axis=0) # array([nan, 4.]) # np.var(adata[adata.obs['well'] == 'well2'].X, axis=0) # array([2209., 4.] d1 = filter_data( - adata, max_fraction_not_finite=None, min_variance=0, max_variance=None, by=by + adata, + max_fraction_not_finite=None, + min_variance=0, + max_variance=None, + by=by, + scale=False, ) # np.var(adata[1:].X, axis=0) array([2006.2222, 2.6666667] d2 = filter_data( - adata, max_fraction_not_finite=0, min_variance=5, max_variance=None, by=by + adata, + max_fraction_not_finite=0, + min_variance=5, + max_variance=None, + by=by, + scale=False, ) assert d1.shape == (4, 1) diff --git a/scallops/tests/test_io.py b/scallops/tests/test_io.py index b3c3a860..e42c840d 100644 --- a/scallops/tests/test_io.py +++ b/scallops/tests/test_io.py @@ -22,9 +22,10 @@ _set_up_experiment, _to_parquet, get_image_spacing, + is_anndata, is_parquet_file, is_scallops_zarr, - read_anndata_zarr, + read_anndata, read_experiment, read_image, save_ome_tiff, @@ -33,7 +34,6 @@ from scallops.zarr_io import ( _write_zarr_image, _write_zarr_labels, - is_anndata_zarr, open_ome_zarr, read_ome_zarr_array, ) @@ -63,7 +63,7 @@ def test_is_anndata_zarr(tmp_path): ) path1 = tmp_path / "test1.zarr" d.write_zarr(path1, convert_strings_to_categoricals=False) - assert is_anndata_zarr(path1) + assert is_anndata(path1) @delayed def create_array(fail): @@ -83,7 +83,7 @@ def create_array(fail): d.write_zarr(path2, convert_strings_to_categoricals=False) except ValueError: pass - assert not is_anndata_zarr(path2) + assert not is_anndata(path2) @pytest.mark.io @@ -713,7 +713,7 @@ def test_anndata_zarr(tmp_path): obs=pd.DataFrame({"b": [1, 2, 3, 4]}), ) d.write_zarr(path, convert_strings_to_categoricals=False) - d2 = read_anndata_zarr(path, dask=True) + d2 = read_anndata(path, dask=True) np.testing.assert_equal(d2.X.compute(), d.X) pd.testing.assert_frame_equal(d.obs, d2.obs) pd.testing.assert_frame_equal(d.var, d2.var) diff --git a/scallops/tests/test_norm_rank_features.py b/scallops/tests/test_norm_rank_features.py index 88c07cf7..8b95e765 100644 --- a/scallops/tests/test_norm_rank_features.py +++ b/scallops/tests/test_norm_rank_features.py @@ -29,21 +29,9 @@ def client(): @pytest.fixture -def data(): - df = pd.DataFrame( - data=dict( - label=np.arange(6), - Cells_Intensity_feature_1=[1, 2, 4, 8, 16, 32], - Cells_Intensity_feature_2=[10, 20, 40, 80, 160, 320], - gene_symbol=["a", "NTC", "a", "NTC", "a", "NTC"], - well=["a", "a", "a", "b", "b", "b"], - plate=["a", "a", "a", "b", "b", "b"], - Nuclei_AreaShape_Center_Y=[1, 7, 12, 16, 19, 21], - Nuclei_AreaShape_Center_X=[1, 7, 12, 16, 19, 21], - ), - ) +def data(test_feature_table): return pandas_to_anndata( - df, ["Cells_Intensity_feature_1", "Cells_Intensity_feature_2"] + test_feature_table, ["Cells_Intensity_feature_1", "Cells_Intensity_feature_2"] ) @@ -97,8 +85,8 @@ def _diff_values(ds, normed_data, normalize, robust, reference, scaling, n_neigh values = values / std np.testing.assert_array_equal( values, - normed_data.X, - err_msg="Not equal", + normed_data[ds.obs.index].X, + err_msg="Expected values not equal", ) @@ -109,12 +97,15 @@ def _compare_anndata(data1: anndata.AnnData, data2: anndata.AnnData): pd.testing.assert_frame_equal(data1.var, data2.var) -@pytest.mark.parametrize("normalize", ["zscore", "local-zscore", "nn-zscore"]) +@pytest.mark.parametrize("normalize", ["local-zscore", "zscore"]) @pytest.mark.parametrize("reference", ["gene_symbol=='NTC'", None]) -@pytest.mark.parametrize("robust", [True, False]) +@pytest.mark.parametrize("robust", [False, True]) @pytest.mark.parametrize("by", [["plate", "well"], None]) +@pytest.mark.parametrize("sort", ["well", None]) @pytest.mark.features -def test_norm_features(client, data, normalize, by, robust, reference, tmp_path): +def test_norm_features(client, data, normalize, by, robust, reference, sort, tmp_path): + if sort is not None: + data = _slice_anndata(data, data.obs.sort_values(sort).index) n_neighbors = 2 if by is None else 1 scaling = n_neighbors > 1 normed_data = normalize_features( @@ -126,7 +117,29 @@ def test_norm_features(client, data, normalize, by, robust, reference, tmp_path) n_neighbors=n_neighbors, scaling=scaling, ) - if normalize in ("nn-zscore", "local-zscore"): + if by is not None: + indices = data.obs.groupby(by).indices + for name in indices: + query = [] + for i in range(len(by)): + query.append(f"{by[i]}=='{name[i]}'") + + _diff_values( + _slice_anndata(data, indices[name]), + _slice_anndata( + normed_data, normed_data.obs.query("&".join(query)).index + ), + normalize, + robust, + reference, + scaling, + n_neighbors, + ) + else: + _diff_values( + data, normed_data, normalize, robust, reference, scaling, n_neighbors + ) + if normalize == "local-zscore": normed_data2 = normalize_features( data, reference_query=reference, @@ -153,7 +166,7 @@ def test_norm_features(client, data, normalize, by, robust, reference, tmp_path) scaling=scaling, ) normed_data_dask.X = normed_data_dask.X.compute() - if normalize in ("nn-zscore", "local-zscore"): + if normalize == "local-zscore": normed_data_dask2 = normalize_features( dask_data, reference_query=reference, @@ -175,29 +188,6 @@ def test_norm_features(client, data, normalize, by, robust, reference, tmp_path) ) _compare_anndata(normed_data, normed_data_dask) - if by is not None: - indices = data.obs.groupby(by).indices - for name in indices: - query = [] - for i in range(len(by)): - query.append(f"{by[i]}=='{name[i]}'") - - _diff_values( - _slice_anndata(data, indices[name]), - _slice_anndata( - normed_data, normed_data.obs.query("&".join(query)).index - ), - normalize, - robust, - reference, - scaling, - n_neighbors, - ) - else: - _diff_values( - data, normed_data, normalize, robust, reference, scaling, n_neighbors - ) - @pytest.mark.parametrize("by", [None, ["well"]]) @pytest.mark.features @@ -338,7 +328,9 @@ def test_agg_features(by, weighted, agg_func, use_dask): agg_d.X = agg_d.X.compute() assert agg_d.shape == (2, 2) agg_df = agg_d.to_df().join(agg_d.obs).sort_values("pert").drop("count", axis=1) - pd.testing.assert_frame_equal(result_df[agg_df.columns], agg_df) + pd.testing.assert_frame_equal( + result_df[agg_df.columns].reset_index(drop=True), agg_df.reset_index(drop=True) + ) @pytest.mark.features diff --git a/scallops/tests/test_pert_map.py b/scallops/tests/test_pert_map.py new file mode 100644 index 00000000..e0e95ef2 --- /dev/null +++ b/scallops/tests/test_pert_map.py @@ -0,0 +1,55 @@ +from subprocess import check_call + +import pytest + +from scallops.features.util import pandas_to_anndata + + +@pytest.mark.parametrize("input_format", ["zarr", "parquet"]) +@pytest.mark.features +def test_map_filter(tmp_path, test_feature_table, input_format): + dataset_path = str(tmp_path / f"dataset_test.{input_format}") + output_path = tmp_path / "labels.zarr" + + if input_format == "parquet": + test_feature_table.to_parquet(dataset_path) + else: + d = pandas_to_anndata( + test_feature_table, + ["Cells_Intensity_feature_1", "Cells_Intensity_feature_2"], + ) + d.write_zarr(dataset_path, convert_strings_to_categoricals=False) + cmd = [ + "scallops", + "pert-map", + "filter", + "--input", + dataset_path, + "--output", + str(output_path), + ] + check_call(cmd) + assert output_path.exists() + + +@pytest.mark.parametrize("outut_format", ["zarr", "parquet"]) +@pytest.mark.features +def test_map_norm(tmp_path, test_feature_table, outut_format): + dataset_path = str(tmp_path / "dataset_test.zarr") + output_path = tmp_path / f"dataset_test.{outut_format}" + + d = pandas_to_anndata( + test_feature_table, ["Cells_Intensity_feature_1", "Cells_Intensity_feature_2"] + ) + d.write_zarr(dataset_path, convert_strings_to_categoricals=False) + cmd = [ + "scallops", + "pert-map", + "normalize", + "--input", + dataset_path, + "--output", + str(output_path), + ] + check_call(cmd) + assert output_path.exists() diff --git a/scallops/utils.py b/scallops/utils.py index 78e6227b..98ebc504 100644 --- a/scallops/utils.py +++ b/scallops/utils.py @@ -50,6 +50,20 @@ def _tqdm_shim(iterator, *args, **kwargs): return iterator +def tqdm_func(progress: bool | str = True): + progress_args = dict() + tqdm_ = _tqdm_shim + if progress != False: # noqa: E712 + try: + from tqdm import tqdm as tqdm_ + + if isinstance(progress, str): + progress_args["desc"] = progress + except ImportError: + pass + return tqdm_, progress_args + + def _fix_json(d): """Attempts to serialize and deserialize a dictionary to ensure it can be safely converted to JSON. diff --git a/scallops/zarr_io.py b/scallops/zarr_io.py index df30e011..156b548b 100644 --- a/scallops/zarr_io.py +++ b/scallops/zarr_io.py @@ -9,17 +9,29 @@ """ import logging -from collections.abc import Callable, Hashable +from collections.abc import Callable, Hashable, Mapping +from importlib.metadata import version from pathlib import Path +from types import MappingProxyType from typing import Any, Literal import dask import dask.array as da import fsspec +import h5py import numpy as np import ome_types import xarray as xr import zarr +from anndata._core.views import DaskArrayView +from anndata._io.specs import _REGISTRY +from anndata._io.specs.methods import ( + suppress_autoshard_warning, + zarr_v3_compressor_compat, + zarr_v3_sharding, +) +from anndata._io.specs.registry import IOSpec, Writer +from anndata.compat import DaskArray from dask.array import from_zarr from dask.delayed import Delayed from dask.graph_manipulation import bind @@ -29,6 +41,7 @@ from ome_zarr.io import parse_url from ome_zarr.types import JSONDict from ome_zarr.writer import write_image, write_multiscale +from packaging.version import Version from xarray.core.coordinates import DataArrayCoordinates from zarr.storage import StoreLike @@ -75,17 +88,6 @@ def _get_store_path(group: zarr.Group): return "" -def is_anndata_zarr(store: StoreLike) -> bool: - """Determines whether store is an AnnData Zarr . - - :param store: Zarr store - """ - try: - return isinstance(zarr.open(store, mode="r", path="layers"), zarr.Group) - except: # noqa: E722 - return False - - def is_ome_zarr_array(node: zarr.Group) -> bool: """Check if a Zarr node is an OME-Zarr array. @@ -899,3 +901,33 @@ def data(self) -> xr.DataArray: if self._data is None: self._data = read_ome_zarr_array(self._group, self._dask) return self._data + + +@_REGISTRY.register_write(zarr.Group, DaskArrayView, IOSpec("array", "0.2.0")) +@_REGISTRY.register_write(zarr.Group, DaskArray, IOSpec("array", "0.2.0")) +@suppress_autoshard_warning +def write_basic_dask_dask_dense( + f: zarr.Group | h5py.Group, + k: str, + elem: DaskArray, + *, + _writer: Writer, + dataset_kwargs: Mapping[str, Any] = MappingProxyType({}), +): + # https://github.com/scverse/anndata/pull/2584 + import dask.array as da + + dataset_kwargs = dict(dataset_kwargs) + if isinstance(f, h5py.Group): + g = f.require_dataset(k, shape=elem.shape, dtype=elem.dtype, **dataset_kwargs) + else: + dataset_kwargs = zarr_v3_compressor_compat(dataset_kwargs) + with zarr_v3_sharding( + dataset_kwargs, format=f.metadata.zarr_format + ) as dataset_kwargs: + g = f.require_array(k, shape=elem.shape, dtype=elem.dtype, **dataset_kwargs) + # use threaded scheduler with dask<=2025.3.0 avoid "Could not serialize object of type HighLevelGraph" error + if isinstance(f, h5py.Group) or Version(version("dask")) <= Version("2025.3.0"): + da.store(elem, g, scheduler="threads") + else: + da.store(elem, g)