Improve xchr function(AlphaGenes/AlphaPeel#225)

src/tinypeel/Peeling/PeelingIO.py

@@ -214,6 +214,7 @@ def writeGenotypes(pedigree, genoProbFunc, isXChr):
                 writeCalledPhase(
                     pedigree,
                     genoProbFunc,
+                    isXChr,
                     args.out_file + ".hap_" + str(threshold) + ".txt",
                     threshold,
                 )
@@ -232,7 +233,7 @@ def writePhasedGenoProbs(pedigree, genoProbFunc, outputFile):
     """
     with open(outputFile, "w+") as f:
         for idx, ind in pedigree.writeOrder():
-            matrix = genoProbFunc(ind.idn)
+            matrix = genoProbFunc(ind.idn, ind.sex)
             f.write("\n")
             for i in range(matrix.shape[0]):
                 f.write(
@@ -253,7 +254,7 @@ def writeGenoProbs(pedigree, genoProbFunc, outputFile):
     """
     with open(outputFile, "w+") as f:
         for idx, ind in pedigree.writeOrder():
-            matrix = genoProbFunc(ind.idn)
+            matrix = genoProbFunc(ind.idn, ind.sex)
             f.write("\n")
             for i in range(matrix.shape[0]):
                 if i == 1:  # Add up probabilities for aA and Aa
@@ -310,10 +311,10 @@ def writeDosages(pedigree, genoProbFunc, isXChr, outputFile):
     with open(outputFile, "w+") as f:
         for idx, ind in pedigree.writeOrder():
             if isXChr and ind.sex == 0:
-                tmp = np.array([0, 1, 0, 0])
+                tmp = np.array([0, 0, 0, 1])
             else:
                 tmp = np.array([0, 1, 1, 2])
-            matrix = np.dot(tmp, genoProbFunc(ind.idn))
+            matrix = np.dot(tmp, genoProbFunc(ind.idn, ind.sex))
             f.write(ind.idx + " " + " ".join(map("{:.4f}".format, matrix)) + "\n")
 
 
@@ -334,7 +335,7 @@ def writeCalledGenotypes(pedigree, genoProbFunc, isXChr, outputFile, thresh):
     """
     with open(outputFile, "w+") as f:
         for idx, ind in pedigree.writeOrder():
-            matrix = genoProbFunc(ind.idn)
+            matrix = genoProbFunc(ind.idn, ind.sex)
             if isXChr and ind.sex == 0:
                 matrixCollapsedHets = np.array(
                     [matrix[0, :] + matrix[2, :], matrix[1, :] + matrix[3, :]],
@@ -351,7 +352,7 @@ def writeCalledGenotypes(pedigree, genoProbFunc, isXChr, outputFile, thresh):
             f.write(ind.idx + " " + " ".join(map(str, calledGenotypes)) + "\n")
 
 
-def writeCalledPhase(pedigree, genoProbFunc, outputFile, thresh):
+def writeCalledPhase(pedigree, genoProbFunc, isXChr, outputFile, thresh):
     """Writes out the called haplotypes to file.
 
     :param pedigree: pedigree information container
@@ -366,15 +367,20 @@ def writeCalledPhase(pedigree, genoProbFunc, outputFile, thresh):
     """
     with open(outputFile, "w+") as f:
         for idx, ind in pedigree.writeOrder():
-            matrix = genoProbFunc(ind.idn)
+            matrix = genoProbFunc(ind.idn, ind.sex)
 
             # Paternal
-            paternal_probs = np.array(
-                [matrix[0, :] + matrix[1, :], matrix[2, :] + matrix[3, :]],
-                dtype=np.float32,
-            )
-            paternal_haplotype = np.argmax(paternal_probs, axis=0)
-            setMissing(paternal_haplotype, paternal_probs, thresh)
+            if isXChr and ind.sex == 0:
+                paternal_haplotype = np.array(
+                    9 * np.ones(matrix.shape[1]), dtype=np.int8
+                )
+            else:
+                paternal_probs = np.array(
+                    [matrix[0, :] + matrix[1, :], matrix[2, :] + matrix[3, :]],
+                    dtype=np.float32,
+                )
+                paternal_haplotype = np.argmax(paternal_probs, axis=0)
+                setMissing(paternal_haplotype, paternal_probs, thresh)
             f.write(ind.idx + " " + " ".join(map(str, paternal_haplotype)) + "\n")
 
             # Maternal
@@ -403,10 +409,17 @@ def writeBinaryCalledGenotypes(pedigree, genoProbFunc, isXChr, outputFile, thres
     :return: None. Writes to the specified output file.
     """
     for idx, ind in pedigree.writeOrder():
-        matrix = genoProbFunc(ind.idn)
-        matrixCollapsedHets = np.array(
-            [matrix[0, :], matrix[1, :] + matrix[2, :], matrix[3, :]], dtype=np.float32
-        )
+        matrix = genoProbFunc(ind.idn, ind.sex)
+        if isXChr and ind.sex == 0:
+            matrixCollapsedHets = np.array(
+                [matrix[0, :], matrix[3, :]],
+                dtype=np.float32,
+            )
+        else:
+            matrixCollapsedHets = np.array(
+                [matrix[0, :], matrix[1, :] + matrix[2, :], matrix[3, :]],
+                dtype=np.float32,
+            )
         calledGenotypes = np.argmax(matrixCollapsedHets, axis=0)
         setMissing(calledGenotypes, matrixCollapsedHets, thresh)
         ind.genotypes = calledGenotypes.astype(np.int8)

src/tinypeel/Peeling/PeelingInfo.py

@@ -92,6 +92,27 @@ def createPeelingInfo(pedigree, args, createSeg=True, phaseFounder=False):
             XChrMaleFlag,
         )
 
+        if peelingInfo.isXChr:
+            if ind.sex == 0:
+                # male the segregation probabilities are 0.5 for pp and pm
+                peelingInfo.pointSeg[ind.idn, 0, :] = 0.5
+                peelingInfo.pointSeg[ind.idn, 1, :] = 0.5
+                peelingInfo.pointSeg[ind.idn, 2, :] = 0
+                peelingInfo.pointSeg[ind.idn, 3, :] = 0
+                peelingInfo.segregation[ind.idn, 0, :] = 0.5
+                peelingInfo.segregation[ind.idn, 1, :] = 0.5
+                peelingInfo.segregation[ind.idn, 2, :] = 0
+                peelingInfo.segregation[ind.idn, 3, :] = 0
+            else:
+                # female the segregation probabilities are 0.5 for mp and mm
+                peelingInfo.pointSeg[ind.idn, 0, :] = 0
+                peelingInfo.pointSeg[ind.idn, 1, :] = 0
+                peelingInfo.pointSeg[ind.idn, 2, :] = 0.5
+                peelingInfo.pointSeg[ind.idn, 3, :] = 0.5
+                peelingInfo.segregation[ind.idn, 0, :] = 0
+                peelingInfo.segregation[ind.idn, 1, :] = 0
+                peelingInfo.segregation[ind.idn, 2, :] = 0.5
+                peelingInfo.segregation[ind.idn, 3, :] = 0.5
         if ind.phenotype is not None:
             # If penetrance is yet updated by genotype inputs, use uniform distribution of 0.25 for all genotypes established in initialisation.
             # TODO: Update for if multiple phenotypes in input or multiple loci in genotypes.
@@ -402,11 +423,13 @@ def construct(self, createSeg=True):
         self.maf = np.full((self.nLoci), 0.5, dtype=np.float32)
         self.transmissionRate = np.full((self.nLoci - 1), 0, dtype=np.float32)
 
-    def getGenoProbs(self, idn):
+    def getGenoProbs(self, idn, sex=None):
         """Estimates the genotype probabilities for a given individual.
 
         :param idn: Internal number for an individual in the pedigree.
         :type idn: int
+        :param sex: 0 is male; 1 is female
+        :type sex: int
         :return: genoProbs: the genotype probabilities for the individual
         :rtype: 2D numpy array of float32 with shape 4 x nLoci
         """
@@ -415,6 +438,12 @@ def getGenoProbs(self, idn):
             * self.posterior[idn, :, :]
             * self.penetrance[idn, :, :]
         )
+        if self.isXChr and sex == 0:  # male
+            genoProbs[0, :] = genoProbs[0, :] + genoProbs[2, :]
+            genoProbs[3, :] = genoProbs[1, :] + genoProbs[3, :]
+            genoProbs[1, :] = 0
+            genoProbs[2, :] = 0
+
         genoProbs = genoProbs / np.sum(genoProbs, 0)
         return genoProbs
 

src/tinypeel/tinypeel.py

@@ -361,9 +361,6 @@ def generateSingleLocusSegregation(peelingInfo, pedigree, args):
                 + (1 - distance[i]) * seg[:, :, segLoc1]
             )
 
-    else:
-        peelingInfo.segregation[:, :, :] = 0.25
-
 
 def get_probability_options():
     """Collects potential user inputs for genotype error rate and sequencing error rate,