diff --git a/configs/ot_cfm/perturbation_meta.yaml b/configs/ot_cfm/perturbation_meta.yaml new file mode 100644 index 0000000..7e30ed4 --- /dev/null +++ b/configs/ot_cfm/perturbation_meta.yaml @@ -0,0 +1,9 @@ +datasets: + - tag: defaultGarciaAlonsoPseudotimeEvenCells + +method: + name: ot_cfm + train_and_test_script: ./methods/ot_cfm/train_and_test.sh + +metrics: + - name: MetaPerturbation diff --git a/configs/ot_cfm/perturbation_set.yaml b/configs/ot_cfm/perturbation_set.yaml index a72d38c..7927f75 100644 --- a/configs/ot_cfm/perturbation_set.yaml +++ b/configs/ot_cfm/perturbation_set.yaml @@ -1,5 +1,6 @@ datasets: - - tag: GarciaAlonsoGC + # - tag: GarciaAlonsoGC + - tag: defaultGarciaAlonsoPseudotimeEvenCells method: name: ot_cfm @@ -7,24 +8,35 @@ method: metrics: - name: PerturbationCellTypeProportion + # this specifies the cell type we want to perturb + # and we'll be filtering out the cells perturbation_set_config: - - gene_col_name: gene_symbols - knockin_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7'] - knockout_genes: ['STRA8', 'ZGLP1', 'ZIC1'] - timepoint_idx: 0 + filter_cell_type: 'GC' + filter_tp_idx: 5 + perturbations: + - gene_col_name: gene_symbols + knockin_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7'] + knockout_genes: ['STRA8', 'ZGLP1', 'ZIC1'] + timepoint_idx: 0 # this should be building off of filter_tp_idx trajectory_infer_model: name: CellTypist renormalize: True - name: PerturbationCellTypeProportion perturbation_set_config: - - gene_col_name: gene_symbols - knockin_genes: ['STRA8', 'ZGLP1', 'ZIC1'] - knockout_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7'] - timepoint_idx: 0 + filter_cell_type: 'GC' + filter_tp_idx: 5 + perturbations: + - gene_col_name: gene_symbols + knockin_genes: ['STRA8', 'ZGLP1', 'ZIC1'] + knockout_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7'] + timepoint_idx: 0 trajectory_infer_model: name: CellTypist renormalize: True - name: PerturbationCellTypeProportion + perturbation_set_config: + filter_cell_type: 'GC' + filter_tp_idx: 5 trajectory_infer_model: name: CellTypist renormalize: True diff --git a/pyproject.toml b/pyproject.toml index c15e315..9508e7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ where = ["src"] # Tells Python: "The actual code lives inside the src folder" scTimeBench = [ "shared/dataset/*.yaml", "shared/dataset/cell_lineages/**/*.txt", + "shared/dataset/perturbation_lineages/**/*.yaml", ] [project.optional-dependencies] diff --git a/src/scTimeBench/config.py b/src/scTimeBench/config.py index 201ccb6..752dab2 100644 --- a/src/scTimeBench/config.py +++ b/src/scTimeBench/config.py @@ -8,6 +8,8 @@ import argparse import logging import os +import random +import numpy as np import yaml from enum import Enum @@ -110,6 +112,13 @@ def __init__(self): help="Directory to save generated plots (default: plots)", ) + parser.add_argument( + "--random_seed", + type=int, + default=42, + help="Random seed for reproducibility (default: 42)", + ) + parser.add_argument( "--clear_tables", action="store_true", @@ -368,5 +377,9 @@ def __init__(self): for metric in self.metrics_skiplist ] + # finally set the random seed + random.seed(self.random_seed) + np.random.seed(self.random_seed) + logging.info("Configuration successfully loaded") logging.debug("Configuration details: %s", self.__dict__) diff --git a/src/scTimeBench/main.py b/src/scTimeBench/main.py index 27a5d9d..d20b3a2 100644 --- a/src/scTimeBench/main.py +++ b/src/scTimeBench/main.py @@ -16,7 +16,11 @@ scTimeBench.metrics # to avoid unused import warning scTimeBench.shared.dataset # to avoid unused import warning -from scTimeBench.metrics.base import METRIC_REGISTRY, BaseMetric +from scTimeBench.metrics.base import ( + METRIC_REGISTRY, + BaseMetric, + create_submetric_instance, +) from scTimeBench.metrics.method_manager import MethodManager from scTimeBench.shared.dataset.base import DATASET_REGISTRY @@ -53,13 +57,7 @@ def run_metrics(config: Config): db_manager = database.DatabaseManager(config) for metric in config.metrics: - metric_name = metric["name"] - if metric_name not in METRIC_REGISTRY: - raise ValueError(f"Metric {metric_name} not found in registry.") - metric_class = METRIC_REGISTRY[metric_name] - metric_instance = metric_class( - config=config, db_manager=db_manager, metric_config=metric - ) + metric_instance = create_submetric_instance(config, db_manager, metric) metric_instance.eval() db_manager.close() diff --git a/src/scTimeBench/method_utils/method_runner.py b/src/scTimeBench/method_utils/method_runner.py index eaeb259..b3aee5b 100644 --- a/src/scTimeBench/method_utils/method_runner.py +++ b/src/scTimeBench/method_utils/method_runner.py @@ -11,9 +11,15 @@ import numpy as np import pandas as pd import scanpy as sc +import scipy as sp from scTimeBench.shared.constants import RequiredOutputFiles from scTimeBench.shared.constants import ObservationColumns +from scTimeBench.shared.utils import ( + is_raw, + undo_log_normalization, + log_normalize_to_counts, +) def get_parser(): @@ -131,12 +137,14 @@ def generate(self, test_ann_data): result.write_h5ad(output_file) elif required_output == RequiredOutputFiles.PERTURBED_TEST_ANN_DATA: - # here we generate the perturbation set - first_tp_cells, all_tps = self._prep_from_first_tp(test_ann_data) - result = self._generate_perturbation(first_tp_cells, all_tps) - self._check_from_first_tp(first_tp_cells, all_tps, result) + result = self._generate_perturbation(test_ann_data) result.write_h5ad(output_file) + elif required_output == RequiredOutputFiles.META_FLAG: + # this is just a placeholder file to indicate that the meta metric has been run + with open(output_file, "w") as f: + f.write("This file indicates that the meta metric has been run.") + else: raise ValueError(f"Unknown required output: {required_output}") @@ -238,12 +246,12 @@ def generate_gex_from_t_to_t1(self, test_ann_data, t, t1) -> sc.AnnData: raise NotImplementedError("Subclasses should implement this method.") # ** NOTE: DO NOT OVERWRITE THIS FUNCTION ** - def _generate_perturbation(self, test_ann_data, all_tps) -> sc.AnnData: + def _generate_perturbation(self, test_ann_data) -> sc.AnnData: """ Generate predicted gene expression across all timepoints for a perturbation. Returns: AnnData object with predicted gene expression across all timepoints. """ - # now we need to read from the perturbation config to create the perturbation set + # here we generate the perturbation set from scTimeBench.shared.perturbation_set import PerturbationSet from scTimeBench.shared.constants import PERTURBATION_SET_CONFIG_FILENAME @@ -255,8 +263,9 @@ def _generate_perturbation(self, test_ann_data, all_tps) -> sc.AnnData: perturbation_set_config = yaml.safe_load(f) perturbation_set = PerturbationSet(perturbation_set_config) - # this is the gex at timepoint tp - gex_tp = test_ann_data.copy() + # then do the initial filtering + first_tp_cells, all_tps = perturbation_set.apply_intial_filter(test_ann_data) + gex_tp = first_tp_cells.copy() all_gex = None for t_idx in range(len(all_tps)): @@ -268,6 +277,14 @@ def _generate_perturbation(self, test_ann_data, all_tps) -> sc.AnnData: # let's print out the gene variable columns names: gex_tp = perturbation_set.apply_perturbation(gex_tp, t_idx) + # data handling to ensure that the output gex is properly normalized + if not is_raw(gex_tp): + if sp.sparse.issparse(gex_tp.X): + gex_tp.X.data = np.clip(gex_tp.X.data, a_min=0, a_max=20) + else: + gex_tp.X = np.clip(gex_tp.X, a_min=0, a_max=20) + gex_tp = log_normalize_to_counts(undo_log_normalization(gex_tp)) + # then we need to save this out to the final output file if all_gex is None: all_gex = gex_tp.copy() @@ -284,6 +301,7 @@ def _generate_perturbation(self, test_ann_data, all_tps) -> sc.AnnData: gex_tp = self.generate_gex_from_t_to_t1(gex_tp, t, t1) print("Finished generating perturbation across all timepoints.") + self._check_from_first_tp(first_tp_cells, all_tps, all_gex) return all_gex diff --git a/src/scTimeBench/metrics/base.py b/src/scTimeBench/metrics/base.py index f604ffb..7624171 100644 --- a/src/scTimeBench/metrics/base.py +++ b/src/scTimeBench/metrics/base.py @@ -290,7 +290,7 @@ def _eval(self): # finally, we evaluate on the test data (ground truth) # and the predicted data from the method - self._submetric_eval( + return self._submetric_eval( **self._prep_kwargs_for_submetric_eval(output_path, dataset, method) ) @@ -464,7 +464,7 @@ def eval(self): ) submetric_instance.eval() else: - self._eval() + return self._eval() # ** PREPROCESSING DATASET SECTION ** def _resolve_tag(self, to_match, dataset_tag): @@ -735,3 +735,18 @@ def _init_datasets(self): for dataset in self.datasets: logging.debug("-" * 100) logging.debug(dataset) + + +def create_submetric_instance( + config: Config, db_manager: DatabaseManager, metric_config: dict +) -> BaseMetric: + """ + Factory that creates an instance of a submetric. + """ + metric_name = metric_config["name"] + if metric_name not in METRIC_REGISTRY: + raise ValueError(f"Metric {metric_name} not found in registry.") + + return METRIC_REGISTRY[metric_name]( + config=config, db_manager=db_manager, metric_config=metric_config + ) diff --git a/src/scTimeBench/metrics/meta/__init__.py b/src/scTimeBench/metrics/meta/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/scTimeBench/metrics/meta/base.py b/src/scTimeBench/metrics/meta/base.py new file mode 100644 index 0000000..beaf1df --- /dev/null +++ b/src/scTimeBench/metrics/meta/base.py @@ -0,0 +1,355 @@ +""" +Meta-Based Metrics. + +This is a metric that operates by running multiple submetrics. + +It takes in as its parameters a set of submetrics to run, and returns this +to the metric to figure out what to do with that information. + +This is useful for things such as calculating the error bar of a metric, +by running it multiple times with different seeds, or also for calculating +the perturbation's baselines as well. +""" +from scTimeBench.metrics.base import BaseMetric, create_submetric_instance +from scTimeBench.shared.constants import RequiredOutputFiles, ObservationColumns +from scTimeBench.shared.utils import load_test_dataset +from scTimeBench.shared.dataset.registry import GarciaAlonsoDataset + +import logging +import json +import os +import random + + +class MetaMetric(BaseMetric): + def _setup_supported_datasets(self): + # ** NOTE: must define the following two attributes, though each subclass ** + # ** Must also define required_feature_specs and output_path_name individually, as they likely require ** + # ** different output files. ** + self.supported_datasets = [ + GarciaAlonsoDataset.__name__, + ] + self.default_dataset_group = "ontology_based" + + # get the path to the shared default datasets config + self.default_datasets_path = os.path.join( + os.path.dirname(__file__), "..", "shared", "default_datasets.yaml" + ) + + self.optional_datasets_path = os.path.join( + os.path.dirname(__file__), "..", "shared", "optional_datasets.yaml" + ) + + def _defaults(self): + """The default parameters for meta-based metrics.""" + return {} + + def _setup_method_output_requirements(self): + """Skip this, as it's a higher level class.""" + # because we're running submetrics, we just ignore this + self.required_outputs = [RequiredOutputFiles.META_FLAG] + + def _prep_kwargs_for_submetric_eval(self, output_path, dataset, method): + return {"output_path": output_path, "dataset": dataset, "method": method} + + def _run_submetric(self, submetric_config): + """ + Run a submetric and return the result from eval. + """ + # here, we'll change config to force a rerun because + # the submetric will likely require different outputs and require + # redoing the metric entirely + self.config.force_rerun = True + submetric_instance = create_submetric_instance( + self.config, self.db_manager, submetric_config + ) + return submetric_instance.eval() + + +class MetaPerturbation(MetaMetric): + def _defaults(self): + return { + "random_trials": 1, + } + + def _handle_transition( + self, transition, cell_type_to_timepoint, gene_col_name, gene_list + ): + """ + Given transition of the format: + - start: 'GC' + targets: + - end: 'oogonia' + genes: ['STRA8', 'ZGLP1', 'ZIC1'] + - end: 'pre_spermatogonia' + genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7'] + + We run and get the difference in the submetric for each of these targets + and return the submetric configuration which should look like: + name: PerturbationCellTypeProportion + perturbation_set_config: + - gene_col_name: gene_symbols + knockin_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7'] + knockout_genes: ['STRA8', 'ZGLP1', 'ZIC1'] + timepoint_idx: 0 + trajectory_infer_model: + name: CellTypist + renormalize: True + except we infer the timepoint_idx to be the timepoint of the start cell type, + with the largest number of cells in the test dataset. + + We also return the cell types that are used. + """ + # should have start and targets + start = transition["start"] + targets = transition["targets"] + + submetrics = [] + timepoint_idx = cell_type_to_timepoint.get(start) + assert ( + timepoint_idx is not None + ), f"Could not infer timepoint idx for cell type {start}." + + cell_types = [] + all_genes = set() + + def generate_submetric(alias, knockin_genes=[], knockout_genes=[]): + perturbations = [] + if len(knockin_genes) > 0 or len(knockout_genes) > 0: + perturbations.append( + { + "gene_col_name": gene_col_name + if gene_col_name is not None + else None, + "knockin_genes": knockin_genes, + "knockout_genes": knockout_genes, + "timepoint_idx": 0, + } + ) + return { + "name": "PerturbationCellTypeProportion", + "alias": alias, + "perturbation_set_config": { + "filter_cell_type": start, + "filter_tp_idx": timepoint_idx, + "perturbations": perturbations, + }, + "trajectory_infer_model": { + "name": "CellTypist", + "renormalize": True, + }, + } + + max_knockin = 0 + max_knockout = 0 + + for target in targets: + logging.debug(f"Running submetric for transition {start} -> {target}") + end = target["end"] + cell_types.append(end) + knockin_genes = target["genes"] + all_genes.update(knockin_genes) + max_knockin = max(max_knockin, len(knockin_genes)) + logging.debug(f"Genes for transition {start} -> {end}: {knockin_genes}") + + # knockout genes should be the ones not included here, but are in other targets + # for example, if we have A to B: [g1, g2], A to C: [g3, g4], + # then A to B knockout genes should be [g3, g4], and A to C knockout genes should be [g1, g2] + knockout_genes = [] + for other_target in targets: + if other_target["end"] != end: + knockout_genes.extend(other_target["genes"]) + max_knockout = max(max_knockout, len(knockout_genes)) + logging.debug( + f"Knockout genes for transition {start} -> {end}: {knockout_genes}" + ) + + submetrics.append(generate_submetric(end, knockin_genes, knockout_genes)) + + # add a baseline one + submetrics.append(generate_submetric("baseline")) + + # finally, also add a random one to check, where we perturb 5 random genes on + # and 5 random genes off, that are not in the original perturbation set + # we need to sort it so that the random genes are the same across different runs + # redo the random seed to ensure consistency across multiple runs + random.seed(self.config.random_seed) + + for i in range(self.params["random_trials"]): + random_genes = sorted(list(set(gene_list) - all_genes)) + random_knockin_genes = random.sample( + random_genes, min(max_knockin, len(random_genes)) + ) + remaining_genes = sorted( + list(set(random_genes) - set(random_knockin_genes)) + ) + random_knockout_genes = random.sample( + remaining_genes, min(max_knockout, len(remaining_genes)) + ) + logging.debug(f"Random knockin genes: {random_knockin_genes}") + logging.debug(f"Random knockout genes: {random_knockout_genes}") + submetrics.append( + generate_submetric( + f"random_{i}", random_knockin_genes, random_knockout_genes + ) + ) + + return submetrics, cell_types + + def _cell_type_to_timepoint_idx(self, test_ann_data): + """ + Returns a mapping of the cell type to the inferred timepoint idx, which is the timepoint with the largest number of cells for that cell type. + """ + cell_types = test_ann_data.obs[ObservationColumns.CELL_TYPE.value].unique() + timepoint_mapping = {} + tps = sorted(test_ann_data.obs[ObservationColumns.TIMEPOINT.value].unique()) + for cell_type in cell_types: + cells = test_ann_data[ + test_ann_data.obs[ObservationColumns.CELL_TYPE.value] == cell_type + ] + timepoint_counts = cells.obs[ + ObservationColumns.TIMEPOINT.value + ].value_counts() + # get the timepoint with the largest number of cells + inferred_timepoint = timepoint_counts.idxmax() + logging.debug( + f"Inferred timepoint for cell type {cell_type}: {inferred_timepoint}" + ) + timepoint_idx = tps.index(inferred_timepoint) + logging.debug( + f"Inferred timepoint idx for cell type {cell_type}: {timepoint_idx}" + ) + timepoint_mapping[cell_type] = timepoint_idx + return timepoint_mapping + + def _submetric_eval(self, output_path, dataset, method): + logging.debug(f"Dataset cell lineage genes: {dataset.cell_lineage_genes}") + + test_ann_data = load_test_dataset(output_path) + cell_type_to_timepoint = self._cell_type_to_timepoint_idx(test_ann_data) + + eval = {} + + gene_col_name = dataset.cell_lineage_genes.get("gene_col_name", None) + + if gene_col_name is not None: + gene_list = test_ann_data.var[gene_col_name].tolist() + else: + gene_list = test_ann_data.var_names.tolist() + + for transition in dataset.cell_lineage_genes["cell_lineage_genes"]: + logging.debug(f"Handling transition: {transition}") + submetrics, cell_types = self._handle_transition( + transition, + cell_type_to_timepoint, + gene_col_name, + gene_list, + ) + logging.debug( + f"Submetrics for transition {transition['start']} -> {[t['end'] for t in transition['targets']]}: {submetrics}" + ) + + results = {} + # now let's run the submetric + for submetric in submetrics: + result = self._run_submetric(submetric) + logging.debug(f"Result of submetric {submetric['alias']}: {result}") + results[submetric["alias"]] = result + + # now we do for each cell type, we're going to calculate + # 1. difference to baseline + # 2. average difference to others + # which is the average of the difference for the same cell type + # for the baseline, we'll report the raw score + + # here we calculate the differences of percentage for a specific cell type + # as described above + for cell_type in cell_types: + logging.debug( + f"Results for cell type {cell_type}: {results[cell_type]}" + ) + baseline_result = results["baseline"][cell_type] + + increase_cell_type_result = results[cell_type][cell_type] + baseline_delta = increase_cell_type_result - baseline_result + + # actually let's not create the average but instead have each difference be calculated + other_cells_delta = {} + + for other_cell_type in cell_types: + if other_cell_type == cell_type: + continue + other_cells_delta[other_cell_type] = ( + results[other_cell_type][cell_type] - baseline_result + ) + + # now we get the random baseline delta, which is the average of the random trials + random_deltas = [] + for i in range(self.params["random_trials"]): + random_delta = results[f"random_{i}"][cell_type] - baseline_result + random_deltas.append(random_delta) + avg_random_delta = sum(random_deltas) / len(random_deltas) + min_random_delta = min(random_deltas) + max_random_delta = max(random_deltas) + + eval[f'{transition["start"]}->{cell_type}'] = { + "baseline": baseline_result, + "baseline_delta": baseline_delta, + "other_cells_delta": other_cells_delta, + "avg_random_delta": avg_random_delta, + "min_random_delta": min_random_delta, + "max_random_delta": max_random_delta, + } + + logging.debug(f"Results of evaluation: {eval}") + + # now we can try to aggregate this information to get an overall score + # we can do: + # 1. a) Accuracy of predicting the correct direction for increase + # b) Accuracy of predicting correct direction for decrease + # 2. Average increase in the perturbed cell type proportion compared to baseline + # 3. Average percentage increase compared to the random baseline + # 4. Accuracy of better increase compared to random baseline + aggregate_scores = { + "pos_direction_accuracy": 0, + "neg_direction_accuracy": 0, + "avg_increase": 0, + "avg_increase_random_baseline": 0, + "beat_random_accuracy": 0, + } + total_other_cells = 0 + for _, result in eval.items(): + if result["baseline_delta"] > 0: + aggregate_scores["pos_direction_accuracy"] += 1 + + # now we iterate through all the other cells + for other_cell_type, delta in result["other_cells_delta"].items(): + if delta < 0: + aggregate_scores["neg_direction_accuracy"] += 1 + total_other_cells += 1 + + aggregate_scores["avg_increase"] += result["baseline_delta"] + aggregate_scores["avg_increase_random_baseline"] += ( + result["baseline_delta"] - result["avg_random_delta"] + ) + if result["baseline_delta"] > result["avg_random_delta"]: + aggregate_scores["beat_random_accuracy"] += 1 + + num_cell_types = len(eval) + aggregate_scores = { + k: v + / (total_other_cells if k == "neg_direction_accuracy" else num_cell_types) + for k, v in aggregate_scores.items() + } + + logging.debug(f"Aggregate scores: {aggregate_scores}") + + final_result = { + "aggregate": aggregate_scores, + "eval": eval, + } + + final_json = json.dumps(final_result, sort_keys=True) + self.db_manager.insert_eval( + method, self.__class__.__name__, self._get_param_encoding(), final_json + ) diff --git a/src/scTimeBench/metrics/perturbation/base.py b/src/scTimeBench/metrics/perturbation/base.py index bd3e545..52c97fe 100644 --- a/src/scTimeBench/metrics/perturbation/base.py +++ b/src/scTimeBench/metrics/perturbation/base.py @@ -1,5 +1,5 @@ """ -Ontology-Based Metrics. +Perturbation-Based Metrics. """ from scTimeBench.metrics.base import BaseMetric from scTimeBench.shared.constants import ( @@ -47,7 +47,7 @@ def _setup_method_output_requirements(self): """Skip this, as it's a higher level class.""" # build the required outputs based on the perturbation set defined self.perturbation_set = PerturbationSet( - self.metric_config.get("perturbation_set_config", []) + self.metric_config.get("perturbation_set_config", {}) ) self.required_outputs = [RequiredOutputFiles.PERTURBED_TEST_ANN_DATA] @@ -73,4 +73,4 @@ def _submetric_eval_setup(self, eval_output_path): eval_output_path, PERTURBATION_SET_CONFIG_FILENAME ) with open(perturbation_config_path, "w") as f: - yaml.safe_dump(self.metric_config.get("perturbation_set_config", []), f) + yaml.safe_dump(self.metric_config.get("perturbation_set_config", {}), f) diff --git a/src/scTimeBench/metrics/perturbation/cell_type_proportion.py b/src/scTimeBench/metrics/perturbation/cell_type_proportion.py index e5830c2..4f92d2d 100644 --- a/src/scTimeBench/metrics/perturbation/cell_type_proportion.py +++ b/src/scTimeBench/metrics/perturbation/cell_type_proportion.py @@ -84,3 +84,5 @@ def _submetric_eval(self, trajectory, method): self.db_manager.insert_eval( method, self.__class__.__name__, self._get_param_encoding(), eval ) + + return cell_type_dist diff --git a/src/scTimeBench/shared/constants.py b/src/scTimeBench/shared/constants.py index 57ac695..39095a7 100644 --- a/src/scTimeBench/shared/constants.py +++ b/src/scTimeBench/shared/constants.py @@ -22,6 +22,8 @@ class RequiredOutputFiles(Enum): FROM_ZERO_TO_END_PRED_GEX = "from_zero_to_end_predicted_gene_expression.h5ad" # perturbation ann data PERTURBED_TEST_ANN_DATA = "perturbed_test_ann_data.h5ad" + # flag to make sure that the meta metric can run + META_FLAG = "meta_flag.txt" DATASET_DIR = "datasets" diff --git a/src/scTimeBench/shared/dataset/base.py b/src/scTimeBench/shared/dataset/base.py index ae8e388..4dd3622 100644 --- a/src/scTimeBench/shared/dataset/base.py +++ b/src/scTimeBench/shared/dataset/base.py @@ -3,9 +3,11 @@ data, so this base class will define the necessary interface for dataset preprocessing. """ from scTimeBench.shared.constants import ObservationColumns, DATASET_DIR +from scTimeBench.shared.helpers import _resolve_shared_resource_path import json import hashlib import os +import yaml # ** DATASET PREPROCESSOR SECTION ** DATASET_PREPROCESSOR_REGISTRY = {} @@ -86,6 +88,18 @@ def __init__( self.TRAIN_PROCESSED_DATA_FILE = "train_processed_data.h5ad" self.TEST_PROCESSED_DATA_FILE = "test_processed_data.h5ad" + # load into this the cell lineages genes if it exists + # which specify the genes involved in the lineage transitions + if "cell_lineage_genes_file" in self.dataset_dict: + with open( + _resolve_shared_resource_path( + self.dataset_dict["cell_lineage_genes_file"] + ) + ) as f: + self.cell_lineage_genes = yaml.safe_load(f) + else: + self.cell_lineage_genes = None + def __init_subclass__(cls): register_dataset(cls) @@ -134,12 +148,17 @@ def encode_dataset_dict(self): # and the preprocessors are encoded elsewhere # then, whether the user decides to include the train dataset caching or not # should not affect the hash of the dataset + # same goes for the cell lineage genes file + + # TODO: consider changing the blocklist to an includes list + # TODO: I really think it should be the other way around... blocklist = [ "data_path", "requires_caching", "data_preprocessing_steps", "tag", "equiv_train_dataset_tag", + "cell_lineage_genes_file", ] return json.dumps( { diff --git a/src/scTimeBench/shared/dataset/default_datasets.yaml b/src/scTimeBench/shared/dataset/default_datasets.yaml index 11e9828..58459cc 100644 --- a/src/scTimeBench/shared/dataset/default_datasets.yaml +++ b/src/scTimeBench/shared/dataset/default_datasets.yaml @@ -51,6 +51,7 @@ datasets: num_tps: 15 - name: LogNormPreprocessor - name: CopyTrainTest + cell_lineage_genes_file: ./perturbation_lineages/garcia_alonso.yaml - name: SuoDataset tag: defaultSuoPseudotimeEvenCells diff --git a/src/scTimeBench/shared/dataset/perturbation_lineages/garcia_alonso.yaml b/src/scTimeBench/shared/dataset/perturbation_lineages/garcia_alonso.yaml new file mode 100644 index 0000000..5b24a1e --- /dev/null +++ b/src/scTimeBench/shared/dataset/perturbation_lineages/garcia_alonso.yaml @@ -0,0 +1,25 @@ +gene_col_name: "gene_symbols" +cell_lineage_genes: + - start: 'PGC' + targets: + - end: 'oogonia' + genes: ['STRA8', 'ZGLP1', 'ZIC1'] + - end: 'pre_spermatogonia' + genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7'] + + - start: 'GC' + targets: + - end: 'oogonia' + genes: ['STRA8', 'ZGLP1', 'ZIC1'] + - end: 'pre_spermatogonia' + genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7'] + + - start: 'oogonia' + targets: + - end: 'pre_oocyte' + genes: ['DMRTC2', 'ZNF711', 'DMRTB1'] + + - start: 'pre_oocyte' + targets: + - end: 'oocyte' + genes: ['TP63', 'ZHX3'] diff --git a/src/scTimeBench/shared/perturbation_set.py b/src/scTimeBench/shared/perturbation_set.py index 4c0a81c..4809a7a 100644 --- a/src/scTimeBench/shared/perturbation_set.py +++ b/src/scTimeBench/shared/perturbation_set.py @@ -1,15 +1,16 @@ import os -import yaml import json import hashlib import logging -from typing import List, Dict +from typing import Dict from scTimeBench.shared.utils import ( is_raw, undo_log_normalization, log_normalize_to_counts, ) +from scTimeBench.shared.constants import ObservationColumns +import scanpy as sc class PerturbationSet: @@ -28,19 +29,29 @@ class PerturbationSet: Which we should convert to a dictionary of the form: { - timepoint_idx: { - "knockin_genes": [...], - "knockout_genes": [...], - "gene_col_name": ... - }, - ... + "filter_cell_type": ..., + "filter_tp": ..., + "perturbations": [ + "timepoint_idx": { + "knockin_genes": [...], + "knockout_genes": [...], + "gene_col_name": ... + }, + ... + ] } so we can easily access the perturbations for each timepoint. """ - def __init__(self, perturbations: List[Dict]): + def __init__(self, perturbation_config: Dict): + self.perturbation_config = perturbation_config + logging.debug(f"Perturbation config: {self.perturbation_config}") + self.perturbations = {} - for perturbation in perturbations: + self.filter_cell_type = perturbation_config.get("filter_cell_type", None) + self.filter_tp_idx = perturbation_config.get("filter_tp_idx", 0) + + for perturbation in self.perturbation_config.get("perturbations", []): timepoint_idx = perturbation["timepoint_idx"] self.perturbations[timepoint_idx] = { "knockin_genes": perturbation.get("knockin_genes", []), @@ -48,12 +59,57 @@ def __init__(self, perturbations: List[Dict]): "gene_col_name": perturbation.get("gene_col_name", None), } + def apply_intial_filter(self, ann_data) -> sc.AnnData: + """ + Apply the initial filter to the data if specified, and then return + the filtered data, and the list of all timepoints in the data (after filtering). + """ + all_tps = sorted(ann_data.obs[ObservationColumns.TIMEPOINT.value].unique()) + + assert self.filter_tp_idx < len( + all_tps + ), f"Filter timepoint index {self.filter_tp_idx} is out of range for the number of timepoints {len(all_tps)} in the data." + + # filter for all timepoints after the specified timepoint (inclusive) + filter_tp = all_tps[self.filter_tp_idx] + tps = [tp for tp in all_tps if tp >= filter_tp] + print( + f"Filtering for timepoints {tps} (filtering for timepoint index {self.filter_tp_idx} which corresponds to timepoint {filter_tp})." + ) + + ann_data = ann_data[ + ann_data.obs[ObservationColumns.TIMEPOINT.value] == filter_tp + ] + + if self.filter_cell_type is not None: + # for the filtered timepoint, only select the cell type specified + assert ( + self.filter_cell_type + in ann_data.obs[ObservationColumns.CELL_TYPE.value].unique() + ), f"Filter cell type {self.filter_cell_type} not found in the data." + ann_data = ann_data[ + ann_data.obs[ObservationColumns.CELL_TYPE.value] + == self.filter_cell_type + ] + + print( + f"After filtering for timepoint {filter_tp} and cell type {self.filter_cell_type}, we have {ann_data.n_obs} cells." + ) + return ann_data, tps + def apply_perturbation(self, ann_data, timepoint_idx): if timepoint_idx not in self.perturbations: print(f"No perturbation specified for timepoint {timepoint_idx}.") return ann_data perturb = self.perturbations[timepoint_idx] + knockout_genes = perturb["knockout_genes"] + knockin_genes = perturb["knockin_genes"] + if len(knockout_genes) == 0 and len(knockin_genes) == 0: + print( + f"Knockout and knockin genes are both empty for timepoint {timepoint_idx}, skipping perturbation." + ) + return ann_data if perturb["gene_col_name"] is None: gene_names = list(ann_data.var_names) @@ -69,8 +125,6 @@ def apply_perturbation(self, ann_data, timepoint_idx): ann_data = undo_log_normalization(ann_data) # track average knockout change per gene for debugging - knockout_genes = perturb["knockout_genes"] - knockin_genes = perturb["knockin_genes"] for gene in knockout_genes: if gene not in gene_to_index: continue @@ -106,16 +160,8 @@ def encode(self): Encode the perturbation set into a unique hash string that will be saved under /perturbations/.yaml. """ - unique_string = json.dumps(self.perturbations, sort_keys=True) + unique_string = json.dumps(self.perturbation_config, sort_keys=True) return hashlib.sha256(unique_string.encode()).hexdigest() - def save_file(self, output_dir): - """ - Save the configuration yaml to the output path - """ - output_path = os.path.join(output_dir, "perturbation.yaml") - with open(output_path, "w") as f: - yaml.dump(self.perturbations, f) - def perturbation_path(self): return os.path.join("perturbations", self.encode())