testing densification methods

src/cpg_seqr_loader/hail_scripts/sparse_mt.py

@@ -412,8 +412,8 @@ def default_compute_info(
         for VQSR.
     :param ac_filter_groups: Optional dictionary of sample filter expressions to compute
         additional groupings of ACs. Default is None.
-    :return: Table with info fields
-    :rtype: Table
+    :return: Tuple of (Table with info fields, set of allele-specific field names in the info struct)
+    :rtype: tuple[Table, set[str]]
     """
     if not site_annotations and not as_annotations and not quasi_as_annotations:
         raise ValueError(
@@ -452,8 +452,10 @@ def default_compute_info(
         # Add allele specific pab_max
         info_expr = info_expr.annotate(AS_pab_max=pab_max_expr(mt.LGT, mt.LAD, mt.LA, hl.len(mt.alleles)))
 
+    site_fields = set()
     if site_annotations:
         site_expr = get_site_info_expr(mt)
+        site_fields = {f for f, _ in site_expr.dtype.items()}
         info_expr = site_expr if info_expr is None else info_expr.annotate(**site_expr)
 
     # Add 'AC' and 'AC_raw' for each allele count filter group requested.
@@ -483,6 +485,8 @@ def default_compute_info(
         **{f'AC{"_" + f if f else f}': grp.map(lambda i: hl.int32(i.get(True, 0))) for f, grp in grp_ac_expr.items()},
     )
 
+    as_info_fields = {f for f, _ in info_expr.dtype.items()} - site_fields
+
     ann_expr = {'info': info_expr}
     if quasi_info_expr is not None:
         ann_expr['quasi_info'] = quasi_info_expr
@@ -511,4 +515,4 @@ def default_compute_info(
     if n_partitions is not None:
         info_ht = info_ht.naive_coalesce(n_partitions)
 
-    return info_ht
+    return info_ht, as_info_fields

src/cpg_seqr_loader/scripts/densify_VDS_to_MT.py

@@ -26,6 +26,26 @@
 from cpg_seqr_loader.utils import read_bed_file_as_intervals
 
 
+def adjust_info_after_split(mt: hl.MatrixTable, as_fields: set[str]) -> hl.MatrixTable:
+    """Post-split: convert allele-specific INFO arrays to per-allele scalars using a_index."""
+
+    info_updates = {}
+    for field in as_fields:
+        if field not in mt.info:
+            continue
+        if field == 'AS_SB_TABLE':
+            info_updates[field] = [mt.info.AS_SB_TABLE[0], mt.info.AS_SB_TABLE[mt.a_index]]
+        else:
+            info_updates[field] = mt.info[field][mt.a_index - 1]
+
+    mt = mt.annotate_rows(info=mt.info.annotate(**info_updates))
+
+    if 'AS_lowqual' in mt.row:
+        mt = mt.annotate_rows(AS_lowqual=mt.AS_lowqual[mt.a_index - 1])
+
+    return mt
+
+
 def densify(vds_path: str, checkpoint_path: str) -> hl.MatrixTable:
     """
     Segregating the densification logic out here - this method reads in a VDS, densifies, repartitions, and checkpoints
@@ -108,7 +128,7 @@ def main(
             mt = densify(vds_in, dense_checkpoint_path)
 
         # content shared with large_cohort.site_only_vcf.py
-        info_ht = default_compute_info(mt, site_annotations=True, n_partitions=mt.n_partitions())
+        info_ht, as_info_fields = default_compute_info(mt, site_annotations=True, n_partitions=mt.n_partitions())
         info_ht = info_ht.annotate(info=info_ht.info.annotate(DP=mt.rows()[info_ht.key].site_dp))
 
         info_ht = adjust_vcf_incompatible_types(
@@ -130,6 +150,9 @@ def main(
         loguru.logger.info('Splitting multiallelics, in a sparse way')
         mt = hl.experimental.sparse_split_multi(mt)
 
+        loguru.logger.info('Adjusting allele-specific INFO fields for split alleles')
+        mt = adjust_info_after_split(mt, as_info_fields)
+
         loguru.logger.info(f'Writing fresh data into {dense_mt_out}')
         mt.write(dense_mt_out, overwrite=True)