add working option for scvi model

dgea/dgea_scvi.py

@@ -11,25 +11,26 @@
 def scanvi_dgea(adata:ad.AnnData, groupby:str, reference:str, alternative:str):
 
     directory_model = "model_scvi_test/"
-    model_path = os.path.join(directory_model, "model.pt")
-    print(reference, alternative)
+    scanvi_subdirectory = os.path.join(directory_model, "scanvi_model")
 
-    loaded_model = SCVI.load(directory_model, adata=None)
-
-    # Check the type of the loaded model
+    try:
+        loaded_model = SCVI.load(directory_model, adata=None)
+    except Exception as e:
+        loaded_model = None
+
     if isinstance(loaded_model, SCVI):
+        os.makedirs(scanvi_subdirectory, exist_ok=True)
         scanvi_model = SCANVI.from_scvi_model(loaded_model, unlabeled_category = 'Unknown', labels_key="cell_type")
+        scanvi_model.save(scanvi_subdirectory, overwrite=True)
+        directory_model = scanvi_subdirectory
         print('is scvi')
-    elif isinstance(loaded_model, SCANVI):
-        scanvi_model = loaded_model
-        print('is scanvi')
-    else:
-        raise ValueError("The model is of an unknown type.")
 
-    scanvi_model.prepare_query_anndata(adata = adata, reference_model=directory_model)
+    else:
+        print('is scanvi')
+
+    SCANVI.prepare_query_anndata(adata = adata, reference_model=directory_model)
 
-    scanvi_model = scanvi_model.load_query_data(adata, directory_model)
-    print(type(scanvi_model))
+    scanvi_model = SCANVI.load_query_data(adata, directory_model)
 
     groups = np.array(adata.obs[groupby].unique())
 
@@ -55,5 +56,6 @@ def get_normalized_counts(adata):
     df_counts = pd.DataFrame(dense_matrix, index=adata.obs_names, columns=adata.var_names)
     return df_counts
 
-adata = sc.read_h5ad('model_scvi_test/adata.h5ad')
-dge_test = scanvi_dgea(adata, "cell_type", "Endothelial", "Epithelial")
+#adata = sc.read_h5ad('model_scvi_test/adata.h5ad')
+#dge_test = scanvi_dgea(adata, "cell_type", "Endothelial", "Epithelial")
+#print(dge_test.head())