Ruffing

malariagen · Jun 20, 2024 · 1ebea03 · 1ebea03
1 parent 1a30510
commit 1ebea03
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 12 deletions.
diff --git a/malariagen_data/anoph/fst.py b/malariagen_data/anoph/fst.py
@@ -480,7 +480,7 @@ def pairwise_average_fst(
     def plot_pairwise_average_fst(
         self,
         fst_df: fst_params.df_pairwise_fst,
-        annotate_se: fst_params.annotate_se = None,
+        annotation: fst_params.annotation = None,
         zmin: Optional[plotly_params.zmin] = 0.0,
         zmax: Optional[plotly_params.zmax] = None,
         text_auto: plotly_params.text_auto = ".3f",
@@ -501,11 +501,11 @@ def plot_pairwise_average_fst(
             col = fst_df.iloc[index_key]["cohort2"]
             fst = fst_df.iloc[index_key]["fst"]
             fig_df[index][col] = fst
-            if annotate_se:
-                if annotate_se is True or annotate_se == "standard error":
+            if annotation:
+                if annotation == "standard error":
                     se = fst_df.iloc[index_key]["se"]
                     fig_df[col][index] = se
-                elif annotate_se == "Z score":
+                elif annotation == "Z score":
                     zs = fst_df.iloc[index_key]["fst"] / fst_df.iloc[index_key]["se"]
                     fig_df[col][index] = zs
             else:

diff --git a/malariagen_data/anoph/fst_params.py b/malariagen_data/anoph/fst_params.py
@@ -1,6 +1,6 @@
 """Parameter definitions for Fst functions."""
 
-from typing import Optional, Union, Literal
+from typing import Optional, Literal
 
 import pandas as pd
 from typing_extensions import Annotated, TypeAlias
@@ -23,11 +23,11 @@
     """,
 ]
 
-annotate_se: TypeAlias = Annotated[
-    Optional[Union[Literal["standard error", "Z score"], bool]],
+annotation: TypeAlias = Annotated[
+    Optional[Literal["standard error", "Z score"]],
     """
-    How to annotate the upper-right corner of the plot. Default behaviour is using Fst, other options
-    are using the standard error (if annotate_se is True or 'standard error') or the Z score of the two
-    cohorts being the same (if annotate_se is 'Z score').
+    How to annotate the upper-right corner of the plot. Default behaviour (None) is using Fst, other options
+    are using the standard error (if annotation is 'standard error') or the Z score of the two
+    cohorts being the same (if annotation is 'Z score').
     """,
 ]
diff --git a/tests/anoph/test_fst.py b/tests/anoph/test_fst.py
@@ -206,9 +206,11 @@ def check_pairwise_average_fst(api: AnophelesFstAnalysis, fst_params):
     if len(fst_df) > 0:
         fig = api.plot_pairwise_average_fst(fst_df, show=False)
         assert isinstance(fig, go.Figure)
-        fig = api.plot_pairwise_average_fst(fst_df, annotate_se=True, show=False)
+        fig = api.plot_pairwise_average_fst(
+            fst_df, annotation="standard error", show=False
+        )
         assert isinstance(fig, go.Figure)
-        fig = api.plot_pairwise_average_fst(fst_df, annotate_se="Z score", show=False)
+        fig = api.plot_pairwise_average_fst(fst_df, annotation="Z score", show=False)
         assert isinstance(fig, go.Figure)