text: # Copyright 2017 Google LLC.
  #
  # Redistribution and use in source and binary forms, with or without
  # modification, are permitted provided that the following conditions
  # are met:
  #
  # 1. Redistributions of source code must retain the above copyright notice,
  #    this list of conditions and the following disclaimer.
  #
  # 2. Redistributions in binary form must reproduce the above copyright
  #    notice, this list of conditions and the following disclaimer in the
  #    documentation and/or other materials provided with the distribution.
  #
  # 3. Neither the name of the copyright holder nor the names of its
  #    contributors may be used to endorse or promote products derived from this
  #    software without specific prior written permission.
  #
  # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
  # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
  # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
  # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
  # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
  # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  # POSSIBILITY OF SUCH DAMAGE.
  """Tests for deepvariant .variant_labeler."""
  
  
  
  from absl.testing import absltest
  from absl.testing import parameterized
  import six
  from third_party.nucleus.io import vcf
  from third_party.nucleus.testing import test_utils
  from third_party.nucleus.util import ranges
  from deepvariant import testdata
  from deepvariant.labeler import variant_labeler
  
  
  def setUpModule():
    testdata.init()
  
  
  class PlaceholderVariantLabeler(variant_labeler.VariantLabeler):
    """A placeholder VariantLabeler.
  
    This class provides a label_variants implementation and so allows the base
    class to be instantiated and its methods tested.
    """
  
    def __init__(self, *pos, **kwargs):
      super(PlaceholderVariantLabeler, self).__init__(*pos, **kwargs)
  
    def label_variants(self, variants, region):
      raise NotImplementedError
  
  
  class VariantLabelerTest(parameterized.TestCase):
    snp = test_utils.make_variant(start=10, alleles=['A', 'C'], gt=[0, 1])
  
    def test_get_truth_variants(self):
      v1 = test_utils.make_variant(chrom='1', start=10)
      v2 = test_utils.make_variant(chrom='1', start=20)
      v3_filtered = test_utils.make_variant(chrom='1', start=30, filters=['FAIL'])
      v4_del = test_utils.make_variant(chrom='1', start=40, alleles=['AAAA', 'A'])
      v5_non_confident = test_utils.make_variant(chrom='1', start=150)
  
      variants = [v1, v2, v3_filtered, v4_del, v5_non_confident]
      reader = vcf.InMemoryVcfReader(variants=variants)
      confident_regions = ranges.RangeSet([ranges.make_range('1', 1, 100)])
      labeler = PlaceholderVariantLabeler(
          truth_vcf_reader=reader, confident_regions=confident_regions)
  
      # Check that we get v1 and v2 specifically when only they are covered by the
      # query.
      self.assertEqual(
          list(labeler._get_truth_variants(ranges.parse_literal('1:1-15'))), [v1])
      self.assertEqual(
          list(labeler._get_truth_variants(ranges.parse_literal('1:15-25'))),
          [v2])
  
      # We don't include filtered variants.
      self.assertEqual(
          list(labeler._get_truth_variants(ranges.parse_literal('1:25-35'))), [])
  
      # Check that we get all overlapping variants of our query.
      for del_query in ['1:35-45', '1:42-43', '1:38-42', '1:42-50']:
        self.assertEqual(
            list(labeler._get_truth_variants(ranges.parse_literal(del_query))),
            [v4_del])
  
      # Checks that a simple query gets all our non-filtered variants.
      self.assertEqual(
          list(labeler._get_truth_variants(ranges.parse_literal('1:1-100'))),
          [v1, v2, v4_del])
      # Even through our query covers v5, it's not confident, so we don't get it.
      self.assertEqual(
          list(labeler._get_truth_variants(ranges.parse_literal('1:1-1000'))),
          [v1, v2, v4_del])
  
    @parameterized.parameters(
        # Make sure we get the right alt counts for all diploid genotypes.
        (['A', 'C'], ['C'], ['A', 'C'], [0, 0], (0, 0), 0),
        (['A', 'C'], ['C'], ['A', 'C'], [0, 1], (0, 1), 1),
        (['A', 'C'], ['C'], ['A', 'C'], [1, 0], (0, 1), 1),
        (['A', 'C'], ['C'], ['A', 'C'], [1, 1], (1, 1), 2),
  
        # Make sure get back a zero alt count for a reference variant.
        (['A'], [], ['A'], [0, 0], (0, 0), 0),
  
        # Basic multi-allelic tests, without having to deal with simplifying
        # alleles as all of the alleles are SNPs. Our candidates have an extra
        # allele, but the true GT is A/C.
        (['A', 'C', 'G'], ['C'], ['A', 'C'], [0, 1], (0, 1), 1),
        (['A', 'C', 'G'], ['C'], ['A', 'C'], [1, 1], (1, 1), 2),
  
        # When considering A/G our answer should be 0 as we have no copies
        # of the G allele.
        (['A', 'C', 'G'], ['G'], ['A', 'C'], [0, 1], (0, 1), 0),
        (['A', 'C', 'G'], ['G'], ['A', 'C'], [1, 1], (1, 1), 0),
  
        # We are considering the het-alt configuration here of A vs. C+G. We've
        # got one copy of the C allele so our true genotype is het. If truth is
        # hom-var for the C, though, we again label the composite as hom_var as
        # we have two copies of the C/G alt.
        (['A', 'C', 'G'], ['C', 'G'], ['A', 'C'], [0, 1], (0, 1), 1),
        (['A', 'C', 'G'], ['C', 'G'], ['A', 'C'], [1, 1], (1, 1), 2),
  
        # Here we have an extra allele in truth, while candidate is bi-allelic.
        # This example 'G' is unused in truth, so we are simply the normal
        # bi-allelic result.
        (['A', 'C'], ['C'], ['A', 'C', 'G'], [0, 0], (0, 0), 0),
        (['A', 'C'], ['C'], ['A', 'C', 'G'], [0, 1], (0, 1), 1),
        (['A', 'C'], ['C'], ['A', 'C', 'G'], [1, 1], (1, 1), 2),
  
        # We check here that we get the bi-allelic result even when the extra
        # allele is in position 1 not 2.
        (['A', 'G'], ['G'], ['A', 'C', 'G'], [0, 0], (0, 0), 0),
        (['A', 'G'], ['G'], ['A', 'C', 'G'], [0, 2], (0, 1), 1),
        (['A', 'G'], ['G'], ['A', 'C', 'G'], [2, 2], (1, 1), 2),
  
        # Now for a real het-alt. We've got three alleles in both, and the true
        # genotype is 1/2.
        (['A', 'C', 'G'], ['C'], ['A', 'C', 'G'], [1, 2], (1, 2), 1),
        (['A', 'C', 'G'], ['G'], ['A', 'C', 'G'], [1, 2], (1, 2), 1),
        (['A', 'C', 'G'], ['C', 'G'], ['A', 'C', 'G'], [1, 2], (1, 2), 2),
  
        # Test all possible values in candidate against het-alt:
        (['A', 'C', 'G', 'T'], ['C'], ['A', 'C', 'G'], [1, 2], (1, 2), 1),
        (['A', 'C', 'G', 'T'], ['G'], ['A', 'C', 'G'], [1, 2], (1, 2), 1),
        (['A', 'C', 'G', 'T'], ['T'], ['A', 'C', 'G'], [1, 2], (1, 2), 0),
        (['A', 'C', 'G', 'T'], ['C', 'G'], ['A', 'C', 'G'], [1, 2], (1, 2), 2),
        (['A', 'C', 'G', 'T'], ['C', 'T'], ['A', 'C', 'G'], [1, 2], (1, 2), 1),
        (['A', 'C', 'G', 'T'], ['G', 'T'], ['A', 'C', 'G'], [1, 2], (1, 2), 1),
  
        # Simple start for indel alleles => exact matching works here.
        (['A', 'AC'], ['AC'], ['A', 'AC'], [0, 0], (0, 0), 0),
        (['A', 'AC'], ['AC'], ['A', 'AC'], [0, 1], (0, 1), 1),
        (['A', 'AC'], ['AC'], ['A', 'AC'], [1, 1], (1, 1), 2),
  
        # We've got a multi-allelic truth, but again exact matching is enough.
        (['A', 'AC'], ['AC'], ['A', 'AC', 'ACC'], [0, 0], (0, 0), 0),
        (['A', 'AC'], ['AC'], ['A', 'AC', 'ACC'], [0, 1], (0, 1), 1),
        (['A', 'AC'], ['AC'], ['A', 'AC', 'ACC'], [1, 1], (1, 1), 2),
        (['A', 'AC'], ['AC'], ['A', 'AC', 'ACC'], [0, 2], (0, 0), 0),
        (['A', 'AC'], ['AC'], ['A', 'AC', 'ACC'], [1, 2], (0, 1), 1),
        (['A', 'AC'], ['AC'], ['A', 'AC', 'ACC'], [2, 2], (0, 0), 0),
  
        # This case has an extra allele (A) in truth but the true genotype
        # corresponds to our candidate alleles exactly.
        (['A', 'AC'], ['AC'], ['AC', 'A', 'ACC'], [0, 2], (0, 1), 1),
        (['A', 'AC'], ['AC'], ['AC', 'A', 'ACC'], [2, 2], (1, 1), 2),
        # If the true genotype involved just the deletion (A) allele, we don't
        # have that allele in our candidate so we always get 0 copies.
        (['A', 'AC'], ['AC'], ['AC', 'A', 'ACC'], [0, 1], (0, 0), 0),
        (['A', 'AC'], ['AC'], ['AC', 'A', 'ACC'], [1, 1], (0, 0), 0),
        # If the truth is het-alt, we can't match the deletion A allele but we do
        # in fact have the A => AC allele as this matches the AC => ACC allele in
        # truth set.
        (['A', 'AC'], ['AC'], ['AC', 'A', 'ACC'], [1, 2], (0, 1), 1),
  
        # We have a multi-allelic candidate but a simple bi-allelic truth. Make
        # sure we match correctly. This is an key case, as we should expect that
        # our candidates frequently have extra alleles changing the represention
        # relative to our truth candidates.
        (['ACT', 'A', 'AACT'], ['A'], ['A', 'AA'], [0, 1], (0, 2), 0),
        (['ACT', 'A', 'AACT'], ['A'], ['A', 'AA'], [1, 1], (2, 2), 0),
        (['ACT', 'A', 'AACT'], ['AACT'], ['A', 'AA'], [0, 1], (0, 2), 1),
        (['ACT', 'A', 'AACT'], ['AACT'], ['A', 'AA'], [1, 1], (2, 2), 2),
        (['ACT', 'A', 'AACT'], ['A', 'AACT'], ['A', 'AA'], [0, 1], (0, 2), 1),
        (['ACT', 'A', 'AACT'], ['A', 'AACT'], ['A', 'AA'], [1, 1], (2, 2), 2),
  
        # The whole complexity: multi-allelic candidate and truth, all with
        # different allele representations.
        # True genotype here is A/AGTGT where ref is AGT [common
        # dinucleotide expansion]. Both candidate and truth have this but each
        # as a different ref so none of the alleles exactly match.
        #
        # Truth     : AGT   => A [1] + AGTGT [2]
        # Candidate : AGTGT => AGT [2] + AGTGTGT [3]
        (['AGTGT', 'A', 'AGT', 'AGTGTGT'
         ], ['A'], ['AGT', 'A', 'AGTGT', 'AGTGTGT'], [1, 2], (2, 3), 0),
        (['AGTGT', 'A', 'AGT', 'AGTGTGT'
         ], ['AGT'], ['AGT', 'A', 'AGTGT', 'AGTGTGT'], [1, 2], (2, 3), 1),
        (['AGTGT', 'A', 'AGT', 'AGTGTGT'
         ], ['AGTGTGT'], ['AGT', 'A', 'AGTGT', 'AGTGTGT'], [1, 2], (2, 3), 1),
        (['AGTGT', 'A', 'AGT', 'AGTGTGT'
         ], ['A', 'AGT'], ['AGT', 'A', 'AGTGT', 'AGTGTGT'], [1, 2], (2, 3), 1),
        (['AGTGT', 'A', 'AGT', 'AGTGTGT'], ['A', 'AGTGTGT'],
         ['AGT', 'A', 'AGTGT', 'AGTGTGT'], [1, 2], (2, 3), 1),
        (['AGTGT', 'A', 'AGT', 'AGTGTGT'], ['AGT', 'AGTGTGT'],
         ['AGT', 'A', 'AGTGT', 'AGTGTGT'], [1, 2], (2, 3), 2),
  
        # Misc. checks with block substititions.
        (['AT', 'A', 'GC'], ['A'], ['ATT', 'AT', 'A'], [0, 1], (0, 1), 1),
        (['AT', 'A', 'GT'], ['A'], ['A', 'G'], [0, 1], (0, 2), 0),
        (['AT', 'A', 'GT'], ['GT'], ['A', 'G'], [0, 1], (0, 2), 1),
    )
    def test_genotype_from_matched_truth(self, variant_alleles, alt_alleles,
                                         truth_alleles, truth_gt,
                                         expected_genotype, expected_label):
      variant = test_utils.make_variant(start=10, alleles=variant_alleles)
      truth_variant = test_utils.make_variant(
          start=10, alleles=truth_alleles, gt=truth_gt)
      self.assertEqual(
          expected_genotype,
          variant_labeler._genotype_from_matched_truth(variant, truth_variant))
      labeled = variant_labeler.VariantLabel(
          is_confident=True, variant=variant, genotype=expected_genotype)
      indices = [variant_alleles.index(alt) - 1 for alt in alt_alleles]
      self.assertEqual(labeled.label_for_alt_alleles(indices), expected_label)
  
    def test_genotype_from_matched_truth_none_truth_variant_raises(self):
      with six.assertRaisesRegex(self, ValueError,
                                 'truth_variant cannot be None'):
        variant_labeler._genotype_from_matched_truth(self.snp, None)
  
    def test_genotype_from_matched_truth_no_call_truth_variant_raises(self):
      with six.assertRaisesRegex(self, ValueError,
                                 'Expected exactly one VariantCal'):
        variant_labeler._genotype_from_matched_truth(
            self.snp, test_utils.make_variant(
                start=10,
                alleles=['A', 'C'],
            ))
  
    def test_genotype_from_matched_truth_no_gt_truth_variant_raises(self):
      with six.assertRaisesRegex(self, ValueError,
                                 'truth_variant needs genotypes'):
        variant_labeler._genotype_from_matched_truth(
            self.snp,
            test_utils.make_variant(
                start=10,
                alleles=['A', 'C'],
                gt=[-1, -1],
            ))
  
    def test_genotype_from_matched_truth_none_variant_raises(self):
      with six.assertRaisesRegex(self, ValueError, 'variant cannot be None'):
        variant_labeler._genotype_from_matched_truth(None, self.snp)
  
  
  if __name__ == '__main__':
    absltest.main()
  