added annotate option

okkevaneck · May 23, 2023 · 7c1a30b · 7c1a30b
1 parent dab1da0
commit 7c1a30b
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 27 deletions.
diff --git a/prospr/__init__.py b/prospr/__init__.py
@@ -8,7 +8,7 @@
 from .datasets import load_vanEck250, load_vanEck1000, load_vanEck_hratio
 from .visualize import plot_protein
 
-__version__ = "0.2a42"
+__version__ = "0.2a43"
 
 __all__ = [
     "__version__",

diff --git a/prospr/visualize.py b/prospr/visualize.py
@@ -102,7 +102,7 @@ def _plot_aminos_2d_paper(
         label="P",
     )
 
-    # Plot first point with a subscript 1.
+    # Plot first point with a green edge color.
     if annotate_first:
         # Plot first point with a different color.
         if df.iloc[0]["Type"] == "H":
@@ -245,31 +245,32 @@ def _plot_aminos_3d_paper(
         depthshade=False,
     )
 
-    # Plot first point with a different color.
-    if df.iloc[0]["Type"] == "H":
-        ax.scatter(
-            df.iloc[0]["x"],
-            df.iloc[0]["y"],
-            df.iloc[0]["z"],
-            marker="o",
-            fc="royalblue",
-            ec="#00ce00",
-            lw=1.3,
-            s=markersize,
-            zorder=2,
-        )
-    else:
-        ax.scatter(
-            df.iloc[0]["x"],
-            df.iloc[0]["y"],
-            df.iloc[0]["z"],
-            marker="o",
-            fc="white",
-            ec="#00ce00",
-            lw=1.3,
-            s=markersize,
-            zorder=2,
-        )
+    # Plot first point with a green edge color.
+    if annotate_first:
+        if df.iloc[0]["Type"] == "H":
+            ax.scatter(
+                df.iloc[0]["x"],
+                df.iloc[0]["y"],
+                df.iloc[0]["z"],
+                marker="o",
+                fc="royalblue",
+                ec="#00ce00",
+                lw=1.3,
+                s=markersize,
+                zorder=2,
+            )
+        else:
+            ax.scatter(
+                df.iloc[0]["x"],
+                df.iloc[0]["y"],
+                df.iloc[0]["z"],
+                marker="o",
+                fc="white",
+                ec="#00ce00",
+                lw=1.3,
+                s=markersize,
+                zorder=2,
+            )
 
     # Plot dotted lines between the aminos that increase the stability.
     pairs = get_scoring_pairs(protein)