take first element of numpy.unique()

since it returns an `ndarray` rather than a scalar
alexlancaster · Sep 9, 2023 · 190bf4a · 190bf4a
1 parent 0b0d177
commit 190bf4a
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/PyPop/Haplo.py b/src/PyPop/Haplo.py
@@ -752,25 +752,25 @@ def _compute_LD(haplos, freqs, compute_ALD=False, debug=False):
         F_2 = 0.0
         F_1_2 = 0.0
         for i in numpy.unique(alleles1):
-           af_1 = numpy.unique(a_freq1[alleles1==i])
+           af_1 = numpy.unique(a_freq1[alleles1==i])[0]  # take the first element of ndarray
            F_1 = F_1 + af_1**2
            F_2_1 = F_2_1 + ((hap_prob[alleles1==i]**2)/af_1).sum()
         for i in numpy.unique(alleles2):
-           af_2 = numpy.unique(a_freq2[alleles2==i])
+           af_2 = numpy.unique(a_freq2[alleles2==i])[0]
            F_2 = F_2 + af_2**2
            F_1_2 = F_1_2 + ((hap_prob[alleles2==i]**2)/af_2).sum()
         if F_2 == 1.0:
            F_2_1_prime = numpy.nan  
            ALD_2_1 = numpy.nan
         else:
            F_2_1_prime = (F_2_1 - F_2)/(1 - F_2)
-           ALD_2_1 = math.sqrt(F_2_1_prime[0])
+           ALD_2_1 = math.sqrt(F_2_1_prime)
         if F_1 == 1:
            F_1_2_prime = numpy.nan
            ALD_1_2 = numpy.nan
         else:
            F_1_2_prime = (F_1_2 - F_1)/(1 - F_1)
-           ALD_1_2 = math.sqrt(F_1_2_prime[0])
+           ALD_1_2 = math.sqrt(F_1_2_prime)
         if debug:
             print ("ALD_1_2:", ALD_1_2)
             print ("ALD_2_1:", ALD_2_1)