Adding main script

HakimBenkirane · HakimBenkirane · commit 500039ee5d9b · 2024-06-26T11:50:49.000+02:00
diff --git a/config/benchmark_union_cpu.yml b/config/benchmark_union_cpu.yml
@@ -0,0 +1,17 @@
+dataset:
+  data_dir: '../../data/spatial'
+  metadata-file: 'metadata_2024_06_21.csv'
+  tissues: ['ovarian']
+  mode: 'union'
+
+params:
+  model_names: ['SEDR', 'STAGATE', 'SpaceFlow', 'GraphST']
+  hidden_dim: 64
+  batch_key: 'ID'
+  device: 'cpu'
+  fast_dev_run: True
+  n_clusters: 7
+
+  
+
+
diff --git a/main.py b/main.py
@@ -0,0 +1,58 @@
+import argparse
+import yaml
+
+from novae_benchmark import AnnDataset
+from novae_benchmark import get_model
+
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-c', '--config', help='config file to use', type=str)
+
+args = parser.parse_args()
+
+def load_yaml(file_path):
+    with open(file_path, 'r') as file:
+        data = yaml.safe_load(file)
+    return data
+
+config = load_yaml(args.config)
+
+print(config)
+
+if __name__ == "__main__":
+    print("------------ Loading Dataset ----------------\n")
+
+    print('Tissues considered : ')
+    for tissue in config['dataset']['tissues']:
+        print('----- ', tissue, '\n')
+
+        dataset = AnnDataset(data_dir='../../data/spatial', metadata_filename='metadata_2024_06_21.csv')
+        adataset = dataset.load_data(tissue_types=[tissue], mode=config['dataset']['mode'])
+
+        print("------------ Dataset Loaded ! ----------------\n")
+
+        results = {model_name: [] for model_name in config['params']['model_names']}
+
+        for model_name in config['params']['model_names']:
+            print("------------ Loading {} Model ----------------\n".format(model_name))
+            model = get_model(model_name=model_name, hidden_dim=config['params']['hidden_dim'])
+            print("------------ Model Loaded ! ----------------\n")
+            if config['dataset']['mode'] == 'union':
+                for adata in adataset:
+                    model(adata=adata, n_clusters=config['params']['n_clusters'], batch_key=config['params']['batch_key'], 
+                        device=config['params']['device'], fast_dev_run=config['params']['fast_dev_run'])
+            
+                    results[model_name].append(model.model_performances)
+            else:
+                    model(adata=adataset, n_clusters=config['params']['n_clusters'], batch_key=config['params']['batch_key'], 
+                        device=config['params']['device'], fast_dev_run=config['params']['fast_dev_run'])
+                    results[model_name].append(model.model_performances)
+
+            print(results)
+            
+        
+
+
+
+    
diff --git a/novae_benchmark/dataset.py b/novae_benchmark/dataset.py
@@ -1,5 +1,4 @@
 import os
-import numpy as np
 import pandas as pd
 import scanpy as sc
 import anndata
@@ -10,7 +9,7 @@ def __init__(self, data_dir, metadata_filename):
         self.metadata_file = os.path.join(data_dir, metadata_filename)
         self.metadata = pd.read_csv(self.metadata_file)
         
-    def load_data(self, tissue_types, use_common_genes=True, multi_slide=False):
+    def load_data(self, tissue_types, mode='union'):
         anndata_list = []
         anndata_list_original = []
         
@@ -22,7 +21,6 @@ def load_data(self, tissue_types, use_common_genes=True, multi_slide=False):
                 adata = sc.read_h5ad(file_path)
                 anndata_list_original.append(adata)
 
-                
                 # Convert gene names to lowercase to handle case insensitivity
                 adata.var.index = adata.var.index.str.lower()
                 
@@ -32,7 +30,7 @@ def load_data(self, tissue_types, use_common_genes=True, multi_slide=False):
                 anndata_list.append(adata)
         
         if anndata_list:
-            if use_common_genes:
+            if mode=='inter':
                 # Find common genes across all datasets
                 common_genes = set(anndata_list[0].var.index)
                 for adata in anndata_list[1:]:
@@ -41,20 +39,39 @@ def load_data(self, tissue_types, use_common_genes=True, multi_slide=False):
                 # Filter each AnnData to include only the common genes
                 anndata_list = [adata[:, list(common_genes)] for adata in anndata_list]
             
-            combined_adata = anndata.concat(
-                anndata_list, 
-                axis=0,
-                join='inner', 
-                label='slide_id', 
-                keys=[adata.obs['slide_id'][0] for adata in anndata_list],
-                pairwise=True
-            )
-        
-        else:
-            combined_adata = None
-        
-        if multi_slide:
-            return combined_adata
+            if mode=='inter':
+                combined_adata = anndata.concat(
+                    anndata_list, 
+                    axis=0,
+                    join='inner', 
+                    label='slide_id', 
+                    keys=[adata.obs['slide_id'][0] for adata in anndata_list],
+                    pairwise=True
+                )
+                return combined_adata
+            else:
+                # Group by gene panels
+                gene_panels = {}
+                for adata in anndata_list:
+                    genes = tuple(sorted(adata.var.index))
+                    if genes not in gene_panels:
+                        gene_panels[genes] = []
+                    gene_panels[genes].append(adata)
+                
+                # Concatenate within each group
+                concatenated_adatas = []
+                for genes, adatas in gene_panels.items():
+                    
+                    concatenated_adata = anndata.concat(
+                        adatas,
+                        axis=0,
+                        join='inner', 
+                        label='slide_id', 
+                        keys=[adata.obs['slide_id'][0] for adata in adatas],
+                        pairwise=True
+                    )
+                    concatenated_adatas.append(concatenated_adata)
+                
+                return concatenated_adatas
         else:
-            return anndata_list_original
-
+            return None
diff --git a/novae_benchmark/model/GraphST/graphst_model.py b/novae_benchmark/model/GraphST/graphst_model.py
@@ -92,10 +92,10 @@ def __init__(self,
         
         fix_seed(self.random_seed)
 
-        if batch_key:
-            list_adatas = [self.adata[self.adata.obs[batch_key] == b].copy() for b in self.adata.obs[batch_key].unique()]
-        else:
+        if batch_key is None:
             list_adatas = [self.adata]
+        else:
+            list_adatas = [self.adata[self.adata.obs[batch_key] == b].copy() for b in self.adata.obs[batch_key].unique()]
         
         for adata in list_adatas:
             if 'highly_variable' not in adata.var.keys():
diff --git a/novae_benchmark/model/SpaceFlow/spaceflow_model.py b/novae_benchmark/model/SpaceFlow/spaceflow_model.py
@@ -11,6 +11,7 @@
 import torch.nn as nn
 import matplotlib.pyplot as plt
 from scipy.spatial import distance_matrix
+from scipy.sparse import block_diag
 from torch_geometric.nn import GCNConv, DeepGraphInfomax
 from sklearn.neighbors import kneighbors_graph
 
@@ -115,7 +116,7 @@ def prepare_figure(self, rsz=4., csz=4., wspace=.4, hspace=.5, left=0.125, right
         plt.subplots_adjust(wspace=wspace, hspace=hspace, left=left, right=right, bottom=bottom, top=top)
         return fig, axs
 
-    def preprocessing_data(self, n_top_genes=None, n_neighbors=10):
+    def preprocessing_data(self, n_top_genes=None, n_neighbors=10, batch_key=None):
         """
         Preprocessing the spatial transcriptomics data
         Generates:  `self.adata_filtered`: (n_cells, n_locations) `numpy.ndarray`
@@ -131,19 +132,28 @@ def preprocessing_data(self, n_top_genes=None, n_neighbors=10):
         :return: a geometry-aware spatial proximity graph of the spatial spots of cells
         :rtype: class:`scipy.sparse.csr_matrix`
         """
-        adata = self.adata
-        if not adata:
+        adatas = self.adata
+        sc.pp.normalize_total(adatas, target_sum=1e4)
+        sc.pp.log1p(adatas)
+        sc.pp.highly_variable_genes(adatas, n_top_genes=n_top_genes, flavor='cell_ranger', subset=True)
+        sc.pp.pca(adatas)
+        if batch_key is None:
+            list_adatas = [adatas]
+        else:
+            list_adatas = [adatas[self.adata.obs[batch_key] == b].copy() for b in adatas.obs[batch_key].unique()]
+        if not adatas:
             print("No annData object found, please run SpaceFlow.SpaceFlow(expr_data, spatial_locs) first!")
             return
-        sc.pp.normalize_total(adata, target_sum=1e4)
-        sc.pp.log1p(adata)
-        sc.pp.highly_variable_genes(adata, n_top_genes=n_top_genes, flavor='cell_ranger', subset=True)
-        sc.pp.pca(adata)
-        spatial_locs = adata.obsm['spatial']
-        spatial_graph = self.graph_alpha(spatial_locs, n_neighbors=n_neighbors)
-
-        self.adata_preprocessed = adata
-        self.spatial_graph = spatial_graph
+        spatial_graphs = []
+        adatas_preprocessed = []
+        for adata in list_adatas:
+            spatial_locs = adata.obsm['spatial']
+            spatial_graphs.append(self.graph_alpha(spatial_locs, n_neighbors=n_neighbors))
+            adatas_preprocessed.append(adata)
+
+
+        self.adata_preprocessed = sc.concat(adatas_preprocessed)
+        self.spatial_graph = block_diag(spatial_graphs, format='csr')
 
     def graph_alpha(self, spatial_locs, n_neighbors=10):
         """
diff --git a/novae_benchmark/model/build.py b/novae_benchmark/model/build.py
@@ -53,25 +53,24 @@ def __call__(
         batch_key: str | None = None,
         device: str = "cpu",
         fast_dev_run: bool = False,
-        multi_slide:bool = False,
     ) -> tuple[np.ndarray, pd.Series]:
         """
         Runs all steps, i.e preprocessing -> training -> inference -> clustering.
 
         Returns:
             A numpy array of shape (n_cells, hidden_dim) and a pandas Series with the cluster labels.
         """
-        print("--------------- {}: Preprocessing Started-------------------".format(self.model_name))
+        print("--------------- {}: Preprocessing Started-------------------\n".format(self.model_name))
         self.preprocess(adata)
-        print("--------------- {}: Preprocessing Finished-------------------".format(self.model_name))
-        print("--------------- {}: Training Started-------------------".format(self.model_name))
+        print("--------------- {}: Preprocessing Finished-------------------\n".format(self.model_name))
+        print("--------------- {}: Training Started-------------------\n".format(self.model_name))
         self.train(adata, batch_key=batch_key, device=device, fast_dev_run=fast_dev_run)
-        print("--------------- {}: Training Finished-------------------".format(self.model_name))
-        print("--------------- {}: Clustering Started-------------------".format(self.model_name))
+        print("--------------- {}: Training Finished-------------------\n".format(self.model_name))
+        print("--------------- {}: Clustering Started-------------------\n".format(self.model_name))
         self.cluster(adata, n_clusters)
-        print("--------------- {}: Clustering Finished-------------------".format(self.model_name))
+        print("--------------- {}: Clustering Finished-------------------\n".format(self.model_name))
         self.evaluate(adata, batch_key, n_clusters)
-        print("--------------- {}: Evaluation completed-------------------".format(self.model_name))
+        print("--------------- {}: Evaluation completed-------------------\n".format(self.model_name))
         print(self.model_performances)
 
 
@@ -135,7 +134,7 @@ def preprocess(self, adata: AnnData):
 
     def train(self, adata: AnnData, batch_key: str | None = None, device: str = "cpu", fast_dev_run: bool = False):
         spaceflow_net = SpaceFlow.Spaceflow(adata=adata)
-        spaceflow_net.preprocessing_data(n_top_genes=self.N_TOP_GENES)
+        spaceflow_net.preprocessing_data(n_top_genes=self.N_TOP_GENES, batch_key=batch_key)
         spaceflow_embedding = spaceflow_net.train(z_dim=self.hidden_dim, epochs=2 if fast_dev_run else 1000)
         adata.obsm[self.model_name] = spaceflow_embedding
 
diff --git a/scripts/train.py b/scripts/train.py
diff --git a/tests/dummy_dataset.py b/tests/dummy_dataset.py
