[feat] Add gene expression and cell types over time

.gitignore

@@ -22,6 +22,8 @@ outputs/
 logs/
 data
 data/
+grn_data
+grn_data/
 
 # UV stuff
 .python-version

configs/ot_cfm/perturbation_meta.yaml

@@ -7,3 +7,6 @@ method:
 
 metrics:
   - name: MetaPerturbation
+  # - name: StackedBarPlot
+  # - name: StackedBarPlot
+  #   from_tp_zero: True

configs/ot_cfm/perturbation_set.yaml

@@ -13,9 +13,9 @@ metrics:
     perturbation_set_config:
       filter_cell_type: 'GC'
       filter_tp_idx: 5
+      gene_col_name: gene_symbols
       perturbations:
-        - gene_col_name: gene_symbols
-          knockin_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7']
+        - 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:
@@ -25,9 +25,9 @@ metrics:
     perturbation_set_config:
       filter_cell_type: 'GC'
       filter_tp_idx: 5
+      gene_col_name: gene_symbols
       perturbations:
-        - gene_col_name: gene_symbols
-          knockin_genes: ['STRA8', 'ZGLP1', 'ZIC1']
+        - knockin_genes: ['STRA8', 'ZGLP1', 'ZIC1']
           knockout_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7']
           timepoint_idx: 0
     trajectory_infer_model:
@@ -40,3 +40,23 @@ metrics:
     trajectory_infer_model:
       name: CellTypist
       renormalize: True
+
+  - name: PerturbationGeneExpression
+    perturbation_set_config:
+      filter_cell_type: 'PGC'
+      filter_tp_idx: 1
+      gene_col_name: gene_symbols
+      perturbations:
+        - 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: GlobalPerturbationGeneExpression
+    affected_genes: ['HDAC8', 'JUND']
+    perturbation_set_config:
+      gene_col_name: gene_symbols
+      knockin_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7']
+      knockout_genes: ['STRA8', 'ZGLP1', 'ZIC1']

configs/scNODE/perturbation_set.yaml

@@ -8,10 +8,11 @@ method:
 metrics:
   - name: PerturbationCellTypeProportion
     perturbation_set_config:
-      - gene_col_name: gene_symbols
-        knockin_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7']
-        knockout_genes: ['STRA8', 'ZGLP1', 'ZIC1']
-        timepoint_idx: 4
+      gene_col_name: gene_symbols
+      perturbations:
+        - knockin_genes: ['SOX4', 'EGR4', 'KLF6', 'KLF7']
+          knockout_genes: ['STRA8', 'ZGLP1', 'ZIC1']
+          timepoint_idx: 4
     trajectory_infer_model:
       name: CellTypist
       renormalize: True

methods/ot_cfm/run.py

@@ -20,6 +20,10 @@
 
 from scTimeBench.method_utils.method_runner import BaseMethod, main
 from scTimeBench.shared.constants import ObservationColumns
+from scTimeBench.shared.utils import (
+    undo_log_normalization,
+    log_normalize_to_counts,
+)
 
 try:
     from torchcfm.conditional_flow_matching import (
@@ -376,6 +380,15 @@ def generate_zero_to_end_pred_gex(self, first_tp_cells, all_tps) -> sc.AnnData:
             if self.embedding_space == "PCA":
                 predicted_x = self._pca_model.inverse_transform(predicted_x)
 
+            # now let's clip this and re-log normalize
+            predicted_x = np.clip(predicted_x, a_min=0, a_max=20)
+            # put it in an ann data and then grab it out
+            predicted_ann_data = sc.AnnData(predicted_x)
+            predicted_ann_data = log_normalize_to_counts(
+                undo_log_normalization(predicted_ann_data)
+            )
+            predicted_x = predicted_ann_data.X
+
             tp_ann_data = first_tp_cells.copy()
             tp_ann_data.X = np.asarray(predicted_x, dtype=np.float32)
             tp_ann_data.obs[time_col] = tp

src/scTimeBench/method_utils/method_runner.py

@@ -140,6 +140,13 @@ def generate(self, test_ann_data):
                 result = self._generate_perturbation(test_ann_data)
                 result.write_h5ad(output_file)
 
+            elif (
+                required_output
+                == RequiredOutputFiles.PERTURBED_TEST_ANN_DATA_T_TO_T_PLUS_ONE
+            ):
+                result = self._generate_perturbation_t_to_t1(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:
@@ -304,6 +311,55 @@ def _generate_perturbation(self, test_ann_data) -> sc.AnnData:
         self._check_from_first_tp(first_tp_cells, all_tps, all_gex)
         return all_gex
 
+    # ** NOTE: DO NOT OVERWRITE THIS FUNCTION **
+    def _generate_perturbation_t_to_t1(self, test_ann_data) -> sc.AnnData:
+        """
+        Generate predicted gene expression from timepoint t to timepoint t1 for a perturbation.
+        Returns: AnnData object with predicted gene expression from t to t1.
+        """
+        # here we generate the perturbation set
+        from scTimeBench.shared.perturbation_set import GlobalPerturbationSet
+        from scTimeBench.shared.constants import GLOBAL_PERTURBATION_SET_CONFIG_FILENAME
+
+        # read the yaml config
+        perturbation_config_path = os.path.join(
+            self.output_path, GLOBAL_PERTURBATION_SET_CONFIG_FILENAME
+        )
+        with open(perturbation_config_path, "r") as f:
+            perturbation_set_config = yaml.safe_load(f)
+        perturbation_set = GlobalPerturbationSet(perturbation_set_config)
+
+        # first let's get the timepoints
+        all_tps = sorted(test_ann_data.obs[ObservationColumns.TIMEPOINT.value].unique())
+
+        # now let's move everything from t to t + 1
+        test_ann_data = perturbation_set.apply_perturbation(test_ann_data)
+        final_ann_data = None
+        for t_idx in range(len(all_tps) - 1):
+            t = all_tps[t_idx]
+            t1 = all_tps[t_idx + 1]
+            print(f"Predicting cells from timepoint {t} to {t1}...")
+
+            ann_data_t = test_ann_data[
+                test_ann_data.obs[ObservationColumns.TIMEPOINT.value] == t
+            ].copy()
+            ann_data_t1 = self.generate_gex_from_t_to_t1(ann_data_t, t, t1)
+            if final_ann_data is None:
+                final_ann_data = ann_data_t1.copy()
+            else:
+                final_ann_data = sc.concat([final_ann_data, ann_data_t1], axis=0)
+
+        # we should have original - tn # of cells
+        num_tn = len(
+            test_ann_data[
+                test_ann_data.obs[ObservationColumns.TIMEPOINT.value] == all_tps[-1]
+            ]
+        )
+        assert (
+            final_ann_data.shape[0] == test_ann_data.shape[0] - num_tn
+        ), f"Expected {test_ann_data.shape[0] - num_tn} cells in the final perturbed data, but got {final_ann_data.shape[0]}"
+        return final_ann_data
+
 
 def main(method_class: BaseMethod):
     print(f"Starting train and testing for method...")