update

README.md

@@ -13,14 +13,14 @@ tags:
 # CI-FM: Cellular Interaction Foundation Model
 
 ## Overview
-This is the PyTorch implementation of the CI-FM model -- an AI model that can simulate the cellular activities within a living tissue (AI virtual tissue). The signature functions of CI-FM are:
+This is the PyTorch implementation of the CI-FM model -- an AI model that can simulate the biological activities within a living tissue (AI virtual tissue). The signature functions of CI-FM are:
 - **Embedding** of celllular microenvironments via ```embeddings = model.embed(adata)``` (Figure below panel D top);
 - **Inference** of cellular gene expressions at a certain microenvironment via ```expressions = model.predict_cells_at_locations(adata, rand_loc)``` (Figure below panel D bottom).
 
 The detailed usage of the model can be found in the [tutorial](https://huggingface.co/ynyou/CIFM/blob/main/test.ipynb).
 Before running the tutorial, please set up an environment following the [environment instruction](https://huggingface.co/ynyou/CIFM#environment).
 
-More information about the model can be found in the [preprint]().
+<!-- More information about the model can be found in the [preprint](). -->
 
 ![](./figures/cifm.png)
 

models_cifm/cifm.py

@@ -33,7 +33,7 @@ class CIFM(
         self.radius_spatial_graph = args.radius_spatial_graph
     
     def channel_matching(self, channel2ensembl_ids_target, channel2ensembl_ids_source, zero_init_for_unmatched_genes=True):
-        # channel2ensembl_ids_target = [[i] for i in adata.var.index.tolist()]
+        device = next(self.parameters()).device
 
         linear_in = nn.Linear(len(channel2ensembl_ids_target), self.hidden_dim, bias=False)
         linear_out1 = nn.Linear(self.hidden_dim, len(channel2ensembl_ids_target), bias=False)
@@ -72,9 +72,9 @@ class CIFM(
 
             num_matching += 1
         
-        self.gene_encoder.layers[0] = linear_in
-        self.mask_cell_expression.layers[-1] = linear_out1
-        self.mask_cell_dropout.layers[-1] = linear_out2
+        self.gene_encoder.layers[0] = linear_in.to(device)
+        self.mask_cell_expression.layers[-1] = linear_out1.to(device)
+        self.mask_cell_dropout.layers[-1] = linear_out2.to(device)
 
         unmatched_channels = list(set(unmatched_channels))
         print('matching', num_matching, 'gene channels out of', len(channel2ensembl_ids_target), '; unmatched channels:', unmatched_channels)
@@ -113,10 +113,11 @@ class CIFM(
 
     def predict_cells_at_locations(self, adata, locations):
         device = next(self.parameters()).device
-        locations = torch.tensor(locations, dtype=torch.float32).to(device)
+    
+        locations = torch.tensor(locations, dtype=torch.float32)
 
         expressions = torch.tensor(adata.X.toarray(), dtype=torch.float32).to(device)
-        expressions = torch.cat([expressions, torch.zeros(locations.shape[0], expressions.shape[1])], dim=0)
+        expressions = torch.cat([expressions, torch.zeros(locations.shape[0], expressions.shape[1]).to(device)], dim=0)
         coordinates = torch.tensor(adata.obsm['spatial'], dtype=torch.float32)
         coordinates = torch.cat([coordinates, locations], dim=0)
         coordinates = torch.cat([coordinates, torch.zeros(coordinates.shape[0], 1)], dim=1).to(device)