Update write_vcf

szhan · Feb 25, 2024 · e8daea2 · e8daea2
1 parent 07030d7
commit e8daea2
Showing 1 changed file with 16 additions and 8 deletions.
diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
@@ -724,8 +724,6 @@ def write_vcf(ref_ts, impdata, out_file, chr_name="1"):
     :param str chr_name: Chromosome name.
     :return: None
     """
-    x = impdata.num_sites
-    n = impdata.num_individuals
     _HEADER = [
         "##fileformat=VCFv4.2",
         "##INFO=<ID=AF,Number=A,Type=Float,"
@@ -762,24 +760,34 @@ def write_vcf(ref_ts, impdata, out_file, chr_name="1"):
             f.write(line + "\n")
         for col_name in _COLUMN_NAMES:
             f.write(col_name + "\t")
-        for i in range(x):
+        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)
             ar2 = compute_estimated_allelic_r_squared(gt_probs)
+            line_str = ""
+            line_str += chr_name + "\t"
+            line_str += impdata.site_pos[i] + "\t"
+            line_str += i + "\t"
+            REF = "A"  # TODO
+            ALT = "T"  # TODO
+            line_str += REF + "\t"
+            line_str += ALT + "\t"
             # QUAL
-            # qual_str = "."
+            qual_str = "."
+            line_str += qual_str + "\t"
             # FILTER
-            # filter_str = "PASS"
+            filter_str = "PASS"
+            line_str += filter_str + "\t"
             # INFO
             dr2 = 0
             af = 0
             info_str = f"AR2{ar2};DR2{dr2};AF{af}"
+            line_str += info_str + "\t"
             # FORMAT
             format_str = ""
-            for j in range(n):
+            for j in range(impdata.num_individuals):
                 format_str += impdata.alleles[j] + "|" + impdata.alleles[j + 1] + ":"
                 format_str += gt_probs[j] + ":"
                 format_str += dosages[j]
                 format_str += "\t"
-            print(info_str + "\t")
-            print(format_str + "\n")
+            line_str += format_str + "\n"