diff --git a/protected/java/src/org/broadinstitute/sting/gatk/downsampling/AlleleBiasedDownsamplingUtils.java b/protected/java/src/org/broadinstitute/sting/gatk/downsampling/AlleleBiasedDownsamplingUtils.java index 1a7b4da51..a61614481 100755 --- a/protected/java/src/org/broadinstitute/sting/gatk/downsampling/AlleleBiasedDownsamplingUtils.java +++ b/protected/java/src/org/broadinstitute/sting/gatk/downsampling/AlleleBiasedDownsamplingUtils.java @@ -102,17 +102,24 @@ public class AlleleBiasedDownsamplingUtils { } private static int scoreAlleleCounts(final int[] alleleCounts) { - final int maxIndex = MathUtils.maxElementIndex(alleleCounts); - final int maxCount = alleleCounts[maxIndex]; + if ( alleleCounts.length < 2 ) + return 0; - int nonMaxCount = 0; - for ( int i = 0; i < 4; i++ ) { - if ( i != maxIndex ) - nonMaxCount += alleleCounts[i]; - } + // sort the counts (in ascending order) + final int[] alleleCountsCopy = alleleCounts.clone(); + Arrays.sort(alleleCountsCopy); - // try to get the best score: in the het case the counts should be equal and in the hom case the non-max should be zero - return Math.min(Math.abs(maxCount - nonMaxCount), Math.abs(nonMaxCount)); + final int maxCount = alleleCountsCopy[alleleCounts.length - 1]; + final int nextBestCount = alleleCountsCopy[alleleCounts.length - 2]; + + int remainderCount = 0; + for ( int i = 0; i < alleleCounts.length - 2; i++ ) + remainderCount += alleleCountsCopy[i]; + + // try to get the best score: + // - in the het case the counts should be equal with nothing else + // - in the hom case the non-max should be zero + return Math.min(maxCount - nextBestCount + remainderCount, Math.abs(nextBestCount + remainderCount)); } /** diff --git a/protected/java/test/org/broadinstitute/sting/gatk/downsampling/AlleleBiasedDownsamplingUtilsUnitTest.java b/protected/java/test/org/broadinstitute/sting/gatk/downsampling/AlleleBiasedDownsamplingUtilsUnitTest.java index 7c2f5619a..be19d3ef4 100755 --- a/protected/java/test/org/broadinstitute/sting/gatk/downsampling/AlleleBiasedDownsamplingUtilsUnitTest.java +++ b/protected/java/test/org/broadinstitute/sting/gatk/downsampling/AlleleBiasedDownsamplingUtilsUnitTest.java @@ -95,7 +95,7 @@ public class AlleleBiasedDownsamplingUtilsUnitTest extends BaseTest { actualCounts[3] += addT; final int[] results = AlleleBiasedDownsamplingUtils.runSmartDownsampling(actualCounts, (int)(pileupSize * contaminationFraction)); - Assert.assertTrue(countsAreEqual(actualCounts, targetCounts)); + Assert.assertTrue(countsAreEqual(results, targetCounts)); } private static boolean countsAreEqual(final int[] counts1, final int[] counts2) {