Merge pull request #538 from malariagen/536-z-score-fst

Option to annotate upper-right corner of Fst plots with Z scores.
malariagen · Jun 21, 2024 · 6ebd049 · 6ebd049
2 parents 54e7a45 + 1efc4bc
commit 6ebd049
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 7 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: bool = False,
+        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",
@@ -500,9 +500,13 @@ def plot_pairwise_average_fst(
             index = fst_df.iloc[index_key]["cohort1"]
             col = fst_df.iloc[index_key]["cohort2"]
             fst = fst_df.iloc[index_key]["fst"]
-            if annotate_se is True:
+            fig_df[index][col] = fst
+            if annotation == "standard error":
                 se = fst_df.iloc[index_key]["se"]
-                fig_df.loc[index, col] = se
+                fig_df[col][index] = se
+            elif annotation == "Z score":
+                zs = fst_df.iloc[index_key]["fst"] / fst_df.iloc[index_key]["se"]
+                fig_df[col][index] = zs
             else:
                 fig_df.loc[index, col] = fst
 

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
+from typing import Optional, Literal
 
 import pandas as pd
 from typing_extensions import Annotated, TypeAlias
@@ -22,3 +22,12 @@
     A dataframe of pairwise Fst and standard error values.
     """,
 ]
+
+annotation: TypeAlias = Annotated[
+    Optional[Literal["standard error", "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/notebooks/plot_pairwise_average_fst.ipynb b/notebooks/plot_pairwise_average_fst.ipynb
@@ -94,7 +94,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ag3.plot_pairwise_average_fst(pairwise_fst_df, annotate_se=True)"
+    "ag3.plot_pairwise_average_fst(pairwise_fst_df, annotation=\"standard error\")"
    ]
   },
   {
@@ -104,7 +104,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ag3.plot_pairwise_average_fst(pairwise_fst_df, annotate_se=False)"
+    "ag3.plot_pairwise_average_fst(pairwise_fst_df, annotation=\"Z score\")"
    ]
   },
   {

diff --git a/tests/anoph/test_fst.py b/tests/anoph/test_fst.py
@@ -204,7 +204,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, annotation="Z score", show=False)
         assert isinstance(fig, go.Figure)