Setting obs_names for corrected adata #4

Author: Mohsen Naghipourfar

scgen/models/util.py

@@ -302,6 +302,7 @@ def batch_removal(network, adata, batch_key="batch", cell_label_key="cell_type")
         corrected = anndata.AnnData(network.reconstruct(all_shared_ann.X, use_data=True))
         corrected.obs = all_shared_ann.obs.copy(deep=True)
         corrected.var_names = adata.var_names.tolist()
+        corrected.obs_names = adata.obs_names.tolist()
         return corrected
     else:
         all_not_shared_ann = anndata.AnnData.concatenate(*not_shared_ct, batch_key="concat_batch")
@@ -311,6 +312,7 @@ def batch_removal(network, adata, batch_key="batch", cell_label_key="cell_type")
         corrected = anndata.AnnData(network.reconstruct(all_corrected_data.X, use_data=True), )
         corrected.obs = pd.concat([all_shared_ann.obs, all_not_shared_ann.obs])
         corrected.var_names = adata.var_names.tolist()
+        corrected.obs_names = adata.obs_names.tolist()
         return corrected
 
 

tests/test_util.py

@@ -3,11 +3,13 @@
 
 
 def test_batch_removal():
-    train = sc.read("./tests/data/pancreas.h5ad", backup_url="https://goo.gl/V29FNk")
+    train = sc.read("./data/pancreas.h5ad", backup_url="https://goo.gl/V29FNk")
     train.obs["cell_type"] = train.obs["celltype"].tolist()
     network = scgen.VAEArith(x_dimension=train.shape[1], model_path="./models/batch")
-    network.train(train_data=train, n_epochs=0)
+    network.train(train_data=train, n_epochs=1, verbose=1)
     corrected_adata = scgen.batch_removal(network, train)
+    print(corrected_adata.obs)
     network.sess.close()
 
 
+test_batch_removal()