[fix] nucleotide diversity - account also for reference base

cbg-ethz · Mar 16, 2023 · 507246a · 507246a
1 parent 72452b3
commit 507246a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/workflow/scripts/compute_diversity_measures.py b/workflow/scripts/compute_diversity_measures.py
@@ -51,8 +51,11 @@ def population_nucleotide_diversity(df_mutations, length):
         if N == 0:
             continue
         freq = df_temp["frequency"].to_numpy()
+        ref_freq = 1 - freq.sum()
+
         position_pnd = freq**2
-        postion_pi = ( 1 - position_pnd.sum() )
+        postion_pi = ( 1 - (position_pnd.sum() + ref_freq**2))* N / (N-1 )
+
         pi += postion_pi
 
     return float(pi / length)