fix normal dist for ellipse

Author: avantikalal

src/polygraph/models.py

@@ -41,7 +41,7 @@ def load_enformer():
     )
 
 
-def enformer_embed(sequences, model):
+def enformer_embed(sequences, model, device="cpu"):
     """
     Embed a batch of sequences using pretrained or fine-tuned enformer
 
@@ -52,6 +52,9 @@ def enformer_embed(sequences, model):
     Returns:
         np.array of shape (n_seqs x 3072)
     """
+    if isinstance(device, int):
+        device = torch.device(device)
+    sequences = str_to_one_hot(sequences).to(device)
     return model(sequences, return_only_embeddings=True).mean(1).cpu().detach().numpy()
 
 
@@ -128,7 +131,7 @@ def _get_embeddings(seqs, model, drop_last_layers=1, device="cpu", swapaxes=Fals
         np.array of shape (n_seqs x n_features)
     """
     if isinstance(model, Enformer):
-        return enformer_embed(seqs, model)
+        return enformer_embed(seqs, model, device=device)
     elif isinstance(model, EsmForMaskedLM):
         return nucleotide_transformer_embed(seqs, model)
     elif isinstance(model, nn.Sequential):

src/polygraph/motifs.py

@@ -50,7 +50,7 @@ def scan(seqs, meme_file, group_col="Group", pthresh=1e-3, rc=True):
         for m in match:
             out["MotifID"].append(motif.name.decode())
             out["SeqID"].append(m.source.accession.decode())
-            if m.strand=='+':
+            if m.strand == "+":
                 out["start"].append(m.start)
                 out["end"].append(m.stop)
             else:
@@ -122,6 +122,10 @@ def nmf(counts, seqs, reference_group, group_col="Group", n_components=10):
     H.index = factors
     H.columns = counts.columns
 
+    # Normalize W and H matrices
+    H = H.div(H.sum(axis=1), axis=0)
+    W = W * H.sum(1)
+
     # Add group IDs to W
     W[group_col] = seqs[group_col].tolist()
 

src/polygraph/visualize.py

@@ -146,7 +146,9 @@ def pca_plot(
         group_col (str): Column containing group IDs.
         components (list): PCA components to plot
         size (float): Size of points
-        show_ellipse (bool): Outline each group with an ellipse.
+        show_ellipse (bool): Fit each group with a multivariate normal
+            distribution and display an ellipse representing the 95%
+            confidence level.
         reference_group (str): Group to use as reference. This group will
             be plotted first.
     """
@@ -175,7 +177,7 @@ def pca_plot(
         + p9.theme_classic()
     )
     if show_ellipse:
-        g = g + p9.stat_ellipse()
+        g = g + p9.stat_ellipse(type="norm")
     return g