Add individual names to column line

szhan · Feb 25, 2024 · 9889011 · 9889011
1 parent ea33a77
commit 9889011
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
@@ -746,7 +746,6 @@ def write_vcf(ref_ts, impdata, out_file, chr_name="1"):
         + 'Description="Estimated Genotype Probability">',
     ]
     _COLUMN_NAMES = [
-        "#",
         "CHROM",
         "POS",
         "ID",
@@ -760,8 +759,10 @@ def write_vcf(ref_ts, impdata, out_file, chr_name="1"):
     with open(out_file, "w") as f:
         for line in _HEADER:
             f.write(line + "\n")
-        for col_name in _COLUMN_NAMES:
-            f.write(col_name + "\t")
+        col_str = "#"
+        col_str += "\t".join(_COLUMN_NAMES)
+        col_str += "\t".join(impdata.individual_names)
+        f.write(col_str + "\n")
         for i in range(impdata.num_sites):
             a1, ap1, a2, ap2 = impdata.get_alleles_at_site(i)
             gt_probs, dosages = compute_individual_scores(a1, ap1, a2, ap2)