Beispiel #1
0
    def test_dataset(self):
        test_resources = 'src/test/resources'

        vds = hc.import_vcf(test_resources + '/sample.vcf')
        vds2 = hc.import_vcf(test_resources + '/sample2.vcf')

        for (dataset, dataset2) in [(vds, vds2)]:
            gt = 'g.GT'

            dataset = dataset.cache()
            dataset2 = dataset2.persist()

            dataset.write('/tmp/sample.vds', overwrite=True)

            dataset.count()

            self.assertEqual(dataset.head(3).count_variants(), 3)

            dataset.query_variants(['variants.count()'])
            dataset.query_samples(['samples.count()'])

            (dataset.annotate_samples_expr(
                'sa.nCalled = gs.filter(g => isDefined({0})).count()'.format(
                    gt)).samples_table().select(['s', 'nCalled = sa.nCalled'
                                                 ]).export('/tmp/sa.tsv'))

            dataset.annotate_global_expr('global.foo = 5')
            dataset.annotate_global_expr(['global.foo = 5', 'global.bar = 6'])

            dataset = dataset.annotate_samples_table(
                hc.import_table(test_resources +
                                '/sampleAnnotations.tsv').key_by('Sample'),
                expr=
                'sa.isCase = table.Status == "CASE", sa.qPhen = table.qPhen')

            (dataset.annotate_variants_expr(
                'va.nCalled = gs.filter(g => isDefined({0})).count()'.format(
                    gt)).count())

            loci_tb = (
                hc.import_table(test_resources + '/sample2_loci.tsv').annotate(
                    'locus = Locus(chr, pos.toInt32())').key_by('locus'))
            (dataset.annotate_variants_table(loci_tb,
                                             root='va.locus_annot').count())

            variants_tb = (hc.import_table(
                test_resources + '/variantAnnotations.tsv'
            ).annotate(
                'variant = Variant(Chromosome, Position.toInt32(), Ref, Alt)').
                           key_by('variant'))
            (dataset.annotate_variants_table(variants_tb,
                                             root='va.table').count())

            (dataset.annotate_variants_vds(
                dataset, expr='va.good = va.info.AF == vds.info.AF').count())

            downsampled = dataset.sample_variants(0.10)
            downsampled.variants_table().select(
                ['chr = v.contig',
                 'pos = v.start']).export('/tmp/sample2_loci.tsv')
            downsampled.variants_table().select('v').export(
                '/tmp/sample2_variants.tsv')

            with open(test_resources + '/sample2.sample_list') as f:
                samples = [s.strip() for s in f]
            (dataset.filter_samples_list(samples).count()[0] == 56)

            locus_tb = (
                hc.import_table(test_resources + '/sample2_loci.tsv').annotate(
                    'locus = Locus(chr, pos.toInt32())').key_by('locus'))

            (dataset.annotate_variants_table(locus_tb,
                                             root='va.locus_annot').count())

            tb = (hc.import_table(
                test_resources + '/variantAnnotations.tsv'
            ).annotate(
                'variant = Variant(Chromosome, Position.toInt32(), Ref, Alt)').
                  key_by('variant'))
            (dataset.annotate_variants_table(tb, root='va.table').count())

            (dataset.annotate_variants_vds(
                dataset, expr='va.good = va.info.AF == vds.info.AF').count())

            dataset.export_vcf('/tmp/sample2.vcf.bgz')

            self.assertEqual(dataset.drop_samples().count()[0], 0)
            self.assertEqual(dataset.drop_variants().count()[1], 0)

            dataset_dedup = (hc.import_vcf([
                test_resources + '/sample2.vcf',
                test_resources + '/sample2.vcf'
            ]).deduplicate())
            self.assertEqual(dataset_dedup.count()[1], 735)

            (dataset.filter_samples_expr('pcoin(0.5)').samples_table().select(
                's').export('/tmp/sample2.sample_list'))

            (dataset.filter_variants_expr('pcoin(0.5)').variants_table().
             select('v').export('/tmp/sample2.variant_list'))

            (dataset.filter_variants_table(
                KeyTable.import_interval_list(
                    test_resources + '/annotinterall.interval_list')).count())

            dataset.filter_intervals(Interval.parse('1:100-end')).count()
            dataset.filter_intervals(map(Interval.parse,
                                         ['1:100-end', '3-22'])).count()

            (dataset.filter_variants_table(
                KeyTable.import_interval_list(
                    test_resources + '/annotinterall.interval_list')).count())

            self.assertEqual(
                dataset2.filter_variants_table(
                    hc.import_table(test_resources + '/sample2_variants.tsv',
                                    key='f0',
                                    impute=True,
                                    no_header=True)).count()[1], 21)

            m2 = {
                r.f0: r.f1
                for r in hc.import_table(test_resources +
                                         '/sample2_rename.tsv',
                                         no_header=True).collect()
            }
            self.assertEqual(
                dataset2.join(dataset2.rename_samples(m2)).count()[0], 200)

            dataset._typecheck()

            dataset.variants_table().export('/tmp/variants.tsv')
            self.assertTrue(
                (dataset.variants_table().annotate('va = json(va)')).same(
                    hc.import_table('/tmp/variants.tsv',
                                    impute=True).key_by('v')))

            dataset.samples_table().export('/tmp/samples.tsv')
            self.assertTrue((
                dataset.samples_table().annotate('s = s, sa = json(sa)')).same(
                    hc.import_table('/tmp/samples.tsv',
                                    impute=True).key_by('s')))

            gt_string = 'gt = g.GT, gq = g.GQ'
            gt_string2 = 'gt: g.GT, gq: g.GQ'

            cols = ['v = v', 'info = va.info']
            for s in dataset.sample_ids:
                cols.append('`{s}`.gt = va.G["{s}"].gt'.format(s=s))
                cols.append('`{s}`.gq = va.G["{s}"].gq'.format(s=s))

            dataset_table = (dataset.annotate_variants_expr(
                'va.G = index(gs.map(g => { s: s, %s }).collect(), s)' %
                gt_string2).variants_table().select(cols))

            dataset_table_typs = {
                fd.name: fd.typ
                for fd in dataset_table.schema.fields
            }
            dataset_table.export('/tmp/sample_kt.tsv')

            self.assertTrue((dataset.make_table(
                'v = v, info = va.info', gt_string, ['v'])).same(
                    hc.import_table('/tmp/sample_kt.tsv',
                                    types=dataset_table_typs).key_by('v')))

            dataset.annotate_variants_expr(
                "va.nHet = gs.filter(g => {0}.isHet()).count()".format(gt))

            dataset.make_table('v = v, info = va.info', 'gt = {0}'.format(gt),
                               ['v'])

            dataset.num_partitions()
            dataset.file_version()
            dataset.sample_ids[:5]
            dataset.variant_schema
            dataset.sample_schema

            self.assertEqual(dataset2.num_samples, 100)
            self.assertEqual(dataset2.count_variants(), 735)

            dataset.annotate_variants_table(dataset.variants_table(),
                                            root="va")

            kt = (dataset.variants_table().annotate("v2 = v").key_by(
                ["v", "v2"]))

            dataset.annotate_variants_table(kt,
                                            root='va.foo',
                                            vds_key=["v", "v"])

            self.assertEqual(kt.query('v.fraction(x => x == v2)'), 1.0)

            dataset.genotypes_table()

            ## This is very slow!!!
            variants_py = (dataset.annotate_variants_expr(
                'va.hets = gs.filter(g => {0}.isHet()).collect()'.format(
                    gt)).variants_table().filter('pcoin(0.1)').collect())

            expr = 'g.GT.isHet() && g.GQ > 20'

            (dataset.filter_genotypes(expr).genotypes_table().select([
                'v', 's', 'nNonRefAlleles = {0}.nNonRefAlleles()'.format(gt)
            ]).export('/tmp/sample2_genotypes.tsv'))

            self.assertTrue((dataset.repartition(16,
                                                 shuffle=False).same(dataset)))

            self.assertTrue(dataset.naive_coalesce(2).same(dataset))

            print(dataset.storage_level())

            dataset = dataset.unpersist()
            dataset2 = dataset2.unpersist()

            new_sample_order = dataset.sample_ids[:]
            random.shuffle(new_sample_order)
            self.assertEqual(
                vds.reorder_samples(new_sample_order).sample_ids,
                new_sample_order)

        sample = hc.import_vcf(test_resources + '/sample.vcf').cache()

        sample.summarize().report()
        sample.drop_samples().summarize().report()

        sample_split = sample.split_multi_hts()

        sample2 = hc.import_vcf(test_resources + '/sample2.vcf')
        sample2 = sample2.persist()

        sample2_split = sample2.split_multi_hts()

        sample.annotate_alleles_expr_hts(
            'va.gs = gs.map(g => g.GT).callStats(g => v)').count()
        sample.annotate_alleles_expr_hts(
            ['va.gs = gs.map(g => g.GT).callStats(g => v)',
             'va.foo = 5']).count()

        glob, concordance1, concordance2 = (
            sample2_split.concordance(sample2_split))
        print(glob[1][4])
        print(glob[4][0])
        print(glob[:][3])
        concordance1.write('/tmp/foo.kt', overwrite=True)
        concordance2.write('/tmp/foo.kt', overwrite=True)

        sample2_split.export_gen('/tmp/sample2.gen', 5)
        sample2_split.export_plink('/tmp/sample2')

        sample2.filter_variants_expr('v.isBiallelic').count()

        sample2.split_multi_hts().grm().export_gcta_grm_bin('/tmp/sample2.grm')

        sample2.hardcalls().count()

        sample2_split.ibd(min=0.2, max=0.6)

        sample2.split_multi_hts().impute_sex().variant_schema

        self.assertEqual(sample2.genotype_schema, Type.hts_schema())

        m2 = {
            r.f0: r.f1
            for r in hc.import_table(test_resources + '/sample2_rename.tsv',
                                     no_header=True,
                                     impute=True).collect()
        }
        self.assertEqual(
            sample2.join(sample2.rename_samples(m2)).count()[0], 200)

        cov = hc.import_table(test_resources + '/regressionLinear.cov',
                              types={
                                  'Cov1': TFloat64(),
                                  'Cov2': TFloat64()
                              }).key_by('Sample')

        phen1 = hc.import_table(test_resources + '/regressionLinear.pheno',
                                missing='0',
                                types={
                                    'Pheno': TFloat64()
                                }).key_by('Sample')
        phen2 = hc.import_table(test_resources +
                                '/regressionLogisticBoolean.pheno',
                                missing='0',
                                types={
                                    'isCase': TBoolean()
                                }).key_by('Sample')

        regression = (hc.import_vcf(
            test_resources +
            '/regressionLinear.vcf').split_multi_hts().annotate_samples_table(
                cov, root='sa.cov').annotate_samples_table(
                    phen1, root='sa.pheno.Pheno').annotate_samples_table(
                        phen2, root='sa.pheno.isCase'))

        (regression.linreg(['sa.pheno.Pheno'],
                           'g.GT.nNonRefAlleles()',
                           covariates=['sa.cov.Cov1',
                                       'sa.cov.Cov2 + 1 - 1']).count())

        (regression.logreg('wald',
                           'sa.pheno.isCase',
                           'g.GT.nNonRefAlleles()',
                           covariates=['sa.cov.Cov1',
                                       'sa.cov.Cov2 + 1 - 1']).count())

        vds_assoc = (regression.annotate_samples_expr(
            'sa.culprit = gs.filter(g => v == Variant("1", 1, "C", "T")).map(g => g.GT.gt).collect()[0]'
        ).annotate_samples_expr(
            'sa.pheno.PhenoLMM = (1 + 0.1 * sa.cov.Cov1 * sa.cov.Cov2) * sa.culprit'
        ))

        covariatesSkat = hc.import_table(test_resources + "/skat.cov",
                                         impute=True).key_by("Sample")

        phenotypesSkat = (hc.import_table(test_resources + "/skat.pheno",
                                          types={
                                              "Pheno": TFloat64()
                                          },
                                          missing="0").key_by("Sample"))

        intervalsSkat = KeyTable.import_interval_list(test_resources +
                                                      "/skat.interval_list")

        weightsSkat = (hc.import_table(test_resources + "/skat.weights",
                                       types={
                                           "locus": TLocus(),
                                           "weight": TFloat64()
                                       }).key_by("locus"))

        skatVds = (vds2.split_multi_hts().annotate_variants_table(
            intervalsSkat, root="va.gene").annotate_variants_table(
                weightsSkat, root="va.weight").annotate_samples_table(
                    phenotypesSkat, root="sa.pheno").annotate_samples_table(
                        covariatesSkat, root="sa.cov").annotate_samples_expr(
                            "sa.pheno = if (sa.pheno == 1.0) false else " +
                            "if (sa.pheno == 2.0) true else NA: Boolean"))

        (skatVds.skat(key_expr='va.gene',
                      weight_expr='va.weight',
                      y='sa.pheno',
                      x='g.GT.nNonRefAlleles()',
                      covariates=['sa.cov.Cov1', 'sa.cov.Cov2'],
                      logistic=False).count())

        (skatVds.skat(key_expr='va.gene',
                      weight_expr='va.weight',
                      y='sa.pheno',
                      x='plDosage(g.PL)',
                      covariates=['sa.cov.Cov1', 'sa.cov.Cov2'],
                      logistic=True).count())

        vds_kinship = vds_assoc.filter_variants_expr('v.start < 4')

        km = vds_kinship.rrm(False, False)

        ld_matrix_path = '/tmp/ldmatrix'
        ldMatrix = vds_kinship.ld_matrix()
        if os.path.isdir(ld_matrix_path):
            shutil.rmtree(ld_matrix_path)
        ldMatrix.write(ld_matrix_path)
        LDMatrix.read(ld_matrix_path).to_local_matrix()

        vds_assoc = vds_assoc.lmmreg(km, 'sa.pheno.PhenoLMM',
                                     'g.GT.nNonRefAlleles()',
                                     ['sa.cov.Cov1', 'sa.cov.Cov2'])

        vds_assoc.variants_table().select(['Variant = v', 'va.lmmreg.*'
                                           ]).export('/tmp/lmmreg.tsv')

        men, fam, ind, var = sample_split.mendel_errors(
            Pedigree.read(test_resources + '/sample.fam'))
        men.select(['fid', 's', 'code'])
        fam.select(['father', 'nChildren'])
        self.assertEqual(ind.key, ['s'])
        self.assertEqual(var.key, ['v'])
        sample_split.annotate_variants_table(var, root='va.mendel').count()

        sample_split.pca_of_normalized_genotypes()

        sample_split.tdt(Pedigree.read(test_resources + '/sample.fam'))

        sample2_split.variant_qc().variant_schema

        sample2.variants_table().export('/tmp/variants.tsv')
        self.assertTrue(
            (sample2.variants_table().annotate('va = json(va)')).same(
                hc.import_table('/tmp/variants.tsv', impute=True).key_by('v')))

        sample2.samples_table().export('/tmp/samples.tsv')
        self.assertTrue(
            (sample2.samples_table().annotate('s = s, sa = json(sa)')).same(
                hc.import_table('/tmp/samples.tsv', impute=True).key_by('s')))

        cols = ['v = v', 'info = va.info']
        for s in sample2.sample_ids:
            cols.append('{s}.gt = va.G["{s}"].gt'.format(s=s))
            cols.append('{s}.gq = va.G["{s}"].gq'.format(s=s))

        sample2_table = (sample2.annotate_variants_expr(
            'va.G = index(gs.map(g => { s: s, gt: g.GT, gq: g.GQ }).collect(), s)'
        ).variants_table().select(cols))

        sample2_table.export('/tmp/sample_kt.tsv')
        sample2_typs = {fd.name: fd.typ for fd in sample2_table.schema.fields}

        self.assertTrue((sample2.make_table(
            'v = v, info = va.info', 'gt = g.GT, gq = g.GQ', ['v'])).same(
                hc.import_table('/tmp/sample_kt.tsv',
                                types=sample2_typs).key_by('v')))

        sample_split.annotate_variants_expr(
            "va.nHet = gs.filter(g => g.GT.isHet()).count()")

        sample2.make_table('v = v, info = va.info', 'gt = g.GT', ['v'])

        sample.num_partitions()
        sample.file_version()
        sample.sample_ids[:5]

        sample2.filter_alleles_hts('pcoin(0.5)')

        sample_split.ld_prune(8).variants_table().select('v').export(
            "/tmp/testLDPrune.tsv")
        kt = (sample2.variants_table().annotate("v2 = v").key_by(["v", "v2"]))
        sample2.annotate_variants_table(kt, root="va.foo", vds_key=["v", "v"])

        self.assertEqual(kt.query('v.fraction(x => x == v2)'), 1.0)

        variants_py = (sample.annotate_variants_expr(
            'va.hets = gs.filter(g => g.GT.isHet).collect()').variants_table().
                       take(5))

        VariantDataset.from_table(sample.variants_table())
Beispiel #2
0
    def test_dataset(self):
        test_resources = 'src/test/resources'

        vds = hc.import_vcf(test_resources + '/sample.vcf')
        vds2 = hc.import_vcf(test_resources + '/sample2.vcf')
        gds = hc.import_vcf(test_resources + '/sample.vcf', generic=True)
        gds2 = hc.import_vcf(test_resources + '/sample2.vcf', generic=True)

        for (dataset, dataset2) in [(vds, vds2), (gds, gds2)]:

            if dataset._is_generic_genotype:
                gt = 'g.GT'
            else:
                gt = 'g'

            dataset.cache()
            dataset2.persist()

            dataset.write('/tmp/sample.vds', overwrite=True)

            dataset.count()

            dataset.query_variants(['variants.count()'])
            dataset.query_samples(['samples.count()'])

            (dataset.annotate_samples_expr(
                'sa.nCalled = gs.filter(g => {0}.isCalled()).count()'.format(
                    gt)).export_samples('/tmp/sa.tsv',
                                        's = s, nCalled = sa.nCalled'))

            dataset.annotate_global_expr('global.foo = 5')
            dataset.annotate_global_expr(['global.foo = 5', 'global.bar = 6'])

            dataset = dataset.annotate_samples_table(
                hc.import_table(test_resources +
                                '/sampleAnnotations.tsv').key_by('Sample'),
                expr=
                'sa.isCase = table.Status == "CASE", sa.qPhen = table.qPhen')

            (dataset.annotate_variants_expr(
                'va.nCalled = gs.filter(g => {0}.isCalled()).count()'.format(
                    gt)).count())

            loci_tb = (
                hc.import_table(test_resources + '/sample2_loci.tsv').annotate(
                    'locus = Locus(chr, pos.toInt())').key_by('locus'))
            (dataset.annotate_variants_table(loci_tb,
                                             root='va.locus_annot').count())

            variants_tb = (hc.import_table(
                test_resources + '/variantAnnotations.tsv').annotate(
                    'variant = Variant(Chromosome, Position.toInt(), Ref, Alt)'
                ).key_by('variant'))
            (dataset.annotate_variants_table(variants_tb,
                                             root='va.table').count())

            (dataset.annotate_variants_vds(
                dataset, expr='va.good = va.info.AF == vds.info.AF').count())

            downsampled = dataset.sample_variants(0.10)
            downsampled.export_variants('/tmp/sample2_loci.tsv',
                                        'chr = v.contig, pos = v.start')
            downsampled.export_variants('/tmp/sample2_variants.tsv', 'v')

            with open(test_resources + '/sample2.sample_list') as f:
                samples = [s.strip() for s in f]
            (dataset.filter_samples_list(samples).count()[0] == 56)

            locus_tb = (
                hc.import_table(test_resources + '/sample2_loci.tsv').annotate(
                    'locus = Locus(chr, pos.toInt())').key_by('locus'))

            (dataset.annotate_variants_table(locus_tb,
                                             root='va.locus_annot').count())

            tb = (hc.import_table(
                test_resources + '/variantAnnotations.tsv').annotate(
                    'variant = Variant(Chromosome, Position.toInt(), Ref, Alt)'
                ).key_by('variant'))
            (dataset.annotate_variants_table(tb, root='va.table').count())

            (dataset.annotate_variants_vds(
                dataset, expr='va.good = va.info.AF == vds.info.AF').count())

            dataset.export_vcf('/tmp/sample2.vcf.bgz')

            self.assertEqual(dataset.drop_samples().count()[0], 0)
            self.assertEqual(dataset.drop_variants().count()[1], 0)

            dataset_dedup = (hc.import_vcf([
                test_resources + '/sample2.vcf',
                test_resources + '/sample2.vcf'
            ]).deduplicate())
            self.assertEqual(dataset_dedup.count()[1], 735)

            (dataset.filter_samples_expr('pcoin(0.5)').export_samples(
                '/tmp/sample2.sample_list', 's'))

            (dataset.filter_variants_expr('pcoin(0.5)').export_variants(
                '/tmp/sample2.variant_list', 'v'))

            (dataset.filter_variants_table(
                KeyTable.import_interval_list(
                    test_resources + '/annotinterall.interval_list')).count())

            dataset.filter_intervals(Interval.parse('1:100-end')).count()
            dataset.filter_intervals(map(Interval.parse,
                                         ['1:100-end', '3-22'])).count()

            (dataset.filter_variants_table(
                KeyTable.import_interval_list(
                    test_resources + '/annotinterall.interval_list')).count())

            self.assertEqual(
                dataset2.filter_variants_table(
                    hc.import_table(test_resources + '/sample2_variants.tsv',
                                    key='f0',
                                    impute=True,
                                    no_header=True)).count()[1], 21)

            m2 = {
                r.f0: r.f1
                for r in hc.import_table(test_resources +
                                         '/sample2_rename.tsv',
                                         no_header=True).collect()
            }
            self.assertEqual(
                dataset2.join(dataset2.rename_samples(m2)).count()[0], 200)

            dataset._typecheck()

            dataset.export_variants('/tmp/variants.tsv', 'v = v, va = va')
            self.assertTrue(
                (dataset.variants_table().annotate('va = json(va)')).same(
                    hc.import_table('/tmp/variants.tsv',
                                    impute=True).key_by('v')))

            dataset.export_samples('/tmp/samples.tsv', 's = s, sa = sa')
            self.assertTrue((
                dataset.samples_table().annotate('s = s, sa = json(sa)')).same(
                    hc.import_table('/tmp/samples.tsv',
                                    impute=True).key_by('s')))

            if dataset._is_generic_genotype:
                gt_string = 'gt = g.GT, gq = g.GQ'
                gt_string2 = 'gt: g.GT, gq: g.GQ'
            else:
                gt_string = 'gt = g.gt, gq = g.gq'
                gt_string2 = 'gt: g.gt, gq: g.gq'

            cols = ['v = v, info = va.info']
            for s in dataset.sample_ids:
                cols.append(
                    '{s}.gt = va.G["{s}"].gt, {s}.gq = va.G["{s}"].gq'.format(
                        s=s))

            (dataset.annotate_variants_expr(
                'va.G = index(gs.map(g => { s: s, %s }).collect(), s)' %
                gt_string2).export_variants('/tmp/sample_kt.tsv',
                                            ','.join(cols)))

            ((dataset.make_table(
                'v = v, info = va.info', gt_string, ['v'])).same(
                    hc.import_table('/tmp/sample_kt.tsv').key_by('v')))

            dataset.annotate_variants_expr(
                "va.nHet = gs.filter(g => {0}.isHet()).count()".format(gt))

            dataset.aggregate_by_key(
                "Variant = v",
                "nHet = g.map(g => {0}.isHet().toInt()).sum().toLong()".format(
                    gt))
            dataset.aggregate_by_key(["Variant = v"], [
                "nHet = g.map(g => {0}.isHet().toInt()).sum().toLong()".format(
                    gt)
            ])

            dataset.make_table('v = v, info = va.info', 'gt = {0}'.format(gt),
                               ['v'])

            dataset.num_partitions()
            dataset.file_version()
            dataset.sample_ids[:5]
            dataset.variant_schema
            dataset.sample_schema

            self.assertEqual(dataset2.num_samples, 100)
            self.assertEqual(dataset2.count_variants(), 735)

            dataset.annotate_variants_table(dataset.variants_table(),
                                            root="va")

            kt = (dataset.variants_table().annotate("v2 = v").key_by(
                ["v", "v2"]))

            dataset.annotate_variants_table(kt,
                                            root='va.foo',
                                            vds_key=["v", "v"])

            self.assertEqual(kt.query('v.fraction(x => x == v2)'), 1.0)

            dataset.genotypes_table()

            ## This is very slow!!!
            variants_py = (dataset.annotate_variants_expr(
                'va.hets = gs.filter(g => {0}.isHet()).collect()'.format(
                    gt)).variants_table().filter('pcoin(0.1)').collect())

            if dataset._is_generic_genotype:
                expr = 'g.GT.isHet() && g.GQ > 20'
            else:
                expr = 'g.isHet() && g.gq > 20'

            (dataset.filter_genotypes(expr).export_genotypes(
                '/tmp/sample2_genotypes.tsv',
                'v, s, {0}.nNonRefAlleles()'.format(gt)))

            self.assertTrue((dataset.repartition(16,
                                                 shuffle=False).same(dataset)))

            print(dataset.storage_level())
            dataset.unpersist()
            dataset2.unpersist()

        sample = hc.import_vcf(test_resources + '/sample.vcf')
        sample.cache()

        sample.summarize().report()

        sample_split = sample.split_multi()

        sample2 = hc.import_vcf(test_resources + '/sample2.vcf')
        sample2.persist()

        sample2_split = sample2.split_multi()

        sample.annotate_alleles_expr('va.gs = gs.callStats(g => v)').count()
        sample.annotate_alleles_expr(
            ['va.gs = gs.callStats(g => v)', 'va.foo = 5']).count()

        glob, concordance1, concordance2 = (
            sample2_split.concordance(sample2_split))
        print(glob[1][4])
        print(glob[4][0])
        print(glob[:][3])
        concordance1.write('/tmp/foo.kt', overwrite=True)
        concordance2.write('/tmp/foo.kt', overwrite=True)

        sample2_split.export_gen('/tmp/sample2.gen', 5)
        sample2_split.export_plink('/tmp/sample2')

        sample2.filter_multi().count()

        sample2.split_multi().grm().export_gcta_grm_bin('/tmp/sample2.grm')

        sample2.hardcalls().count()

        sample2_split.ibd(min=0.2, max=0.6)

        sample2.split_multi().impute_sex().variant_schema

        self.assertTrue(isinstance(sample2.genotype_schema, TGenotype))

        m2 = {
            r.f0: r.f1
            for r in hc.import_table(test_resources + '/sample2_rename.tsv',
                                     no_header=True,
                                     impute=True).collect()
        }
        self.assertEqual(
            sample2.join(sample2.rename_samples(m2)).count()[0], 200)

        cov = hc.import_table(test_resources + '/regressionLinear.cov',
                              types={
                                  'Cov1': TDouble(),
                                  'Cov2': TDouble()
                              }).key_by('Sample')

        phen1 = hc.import_table(test_resources + '/regressionLinear.pheno',
                                missing='0',
                                types={
                                    'Pheno': TDouble()
                                }).key_by('Sample')
        phen2 = hc.import_table(test_resources +
                                '/regressionLogisticBoolean.pheno',
                                missing='0',
                                types={
                                    'isCase': TBoolean()
                                }).key_by('Sample')

        regression = (hc.import_vcf(
            test_resources +
            '/regressionLinear.vcf').split_multi().annotate_samples_table(
                cov, root='sa.cov').annotate_samples_table(
                    phen1, root='sa.pheno.Pheno').annotate_samples_table(
                        phen2, root='sa.pheno.isCase'))

        (regression.linreg('sa.pheno.Pheno',
                           covariates=['sa.cov.Cov1',
                                       'sa.cov.Cov2 + 1 - 1']).count())

        (regression.logreg('wald',
                           'sa.pheno.isCase',
                           covariates=['sa.cov.Cov1',
                                       'sa.cov.Cov2 + 1 - 1']).count())

        vds_assoc = (regression.annotate_samples_expr(
            'sa.culprit = gs.filter(g => v == Variant("1", 1, "C", "T")).map(g => g.gt).collect()[0]'
        ).annotate_samples_expr(
            'sa.pheno.PhenoLMM = (1 + 0.1 * sa.cov.Cov1 * sa.cov.Cov2) * sa.culprit'
        ))

        vds_kinship = vds_assoc.filter_variants_expr('v.start < 4')

        km = vds_kinship.rrm(False, False)
        ldMatrix = vds_kinship.ld_matrix()
        vds_assoc = vds_assoc.lmmreg(km, 'sa.pheno.PhenoLMM',
                                     ['sa.cov.Cov1', 'sa.cov.Cov2'])

        vds_assoc.export_variants('/tmp/lmmreg.tsv',
                                  'Variant = v, va.lmmreg.*')

        men, fam, ind, var = sample_split.mendel_errors(
            Pedigree.read(test_resources + '/sample.fam'))
        men.select(['fid', 's', 'code'])
        fam.select(['father', 'nChildren'])
        self.assertEqual(ind.key, ['s'])
        self.assertEqual(var.key, ['v'])
        sample_split.annotate_variants_table(var, root='va.mendel').count()

        sample_split.pca('sa.scores')

        sample_split.tdt(Pedigree.read(test_resources + '/sample.fam'))

        sample2_split.variant_qc().variant_schema

        sample2.export_variants('/tmp/variants.tsv', 'v = v, va = va')
        self.assertTrue(
            (sample2.variants_table().annotate('va = json(va)')).same(
                hc.import_table('/tmp/variants.tsv', impute=True).key_by('v')))

        sample2.export_samples('/tmp/samples.tsv', 's = s, sa = sa')
        self.assertTrue(
            (sample2.samples_table().annotate('s = s, sa = json(sa)')).same(
                hc.import_table('/tmp/samples.tsv', impute=True).key_by('s')))

        cols = ['v = v, info = va.info']
        for s in sample2.sample_ids:
            cols.append(
                '{s}.gt = va.G["{s}"].gt, {s}.gq = va.G["{s}"].gq'.format(s=s))

        (sample2.annotate_variants_expr(
            'va.G = index(gs.map(g => { s: s, gt: g.gt, gq: g.gq }).collect(), s)'
        ).export_variants('/tmp/sample_kt.tsv', ','.join(cols)))

        ((sample2.make_table(
            'v = v, info = va.info', 'gt = g.gt, gq = g.gq',
            ['v'])).same(hc.import_table('/tmp/sample_kt.tsv').key_by('v')))

        sample_split.annotate_variants_expr(
            "va.nHet = gs.filter(g => g.isHet()).count()")

        sample_split.aggregate_by_key(
            "Variant = v",
            "nHet = g.map(g => g.isHet().toInt()).sum().toLong()")
        sample_split.aggregate_by_key(
            ["Variant = v"],
            ["nHet = g.map(g => g.isHet().toInt()).sum().toLong()"])

        sample2.make_table('v = v, info = va.info', 'gt = g.gt', ['v'])

        sample.num_partitions()
        sample.file_version()
        sample.sample_ids[:5]

        self.assertFalse(sample2.was_split())

        self.assertTrue(sample_split.was_split())

        sample2.filter_alleles('pcoin(0.5)')

        gds.annotate_genotypes_expr('g = g.GT.toGenotype()').split_multi()

        sample_split.ld_prune().export_variants("/tmp/testLDPrune.tsv", "v")
        kt = (sample2.variants_table().annotate("v2 = v").key_by(["v", "v2"]))
        sample2.annotate_variants_table(kt, root="va.foo", vds_key=["v", "v"])

        self.assertEqual(kt.query('v.fraction(x => x == v2)'), 1.0)

        variants_py = (sample.annotate_variants_expr(
            'va.hets = gs.filter(g => g.isHet).collect()').variants_table().
                       collect())

        VariantDataset.from_table(sample.variants_table())