Optimizations to sex-ploidy adjustment and reducing groups during aggs

[mike-w-wilson]

Two changes that make compute_stats_per_ref_site (per-reference-site AN and coverage over a densified VDS) cheaper to run without changing its output. Both are needed by the gnomAD v5 AoU coverage/AN pipeline.

Breaking change: annotate_and_index_source_mt_for_sex_ploidy now takes the MatrixTable and karyotype expression and returns the annotated MT plus its column and row flag structs, instead of taking a locus expression and indexing the flags back through a cols()/rows() self-join. The only downstream caller is gnomad_qc v4 fix_freq_an.py (line 217). The old behavior is available unchanged under the new public name index_sex_ploidy_flags, so that call can be fixed by renaming it. adjust_sex_ploidy and adjusted_sex_ploidy_expr keep their signatures.

1. Sex-ploidy adjustment without a second densify. adjusted_sex_ploidy_expr looks up each sample's karyotype and each locus's PAR status by reaching back into the source MatrixTable's columns and rows, which Hail evaluates as a second full densify on a densified VDS. annotate_and_index_source_mt_for_sex_ploidy now annotates those flags directly onto the MT it is given, adjust_sex_ploidy is a thin wrapper over it, and compute_stats_per_ref_site uses that path. The flags themselves are defined once in get_sex_ploidy_col_flags_expr / get_sex_ploidy_row_flags_expr, and the genotype rules are shared with adjusted_sex_ploidy_expr via _sex_ploidy_case_expr. Temporary flag fields use collision-safe names.

2. Aggregate once per "cell" instead of once per group. Allele-number aggregation previously ran once for every freq_meta group, and the minimal-groups reduction could not help with downsampling groups. The new reduce_to_cells option groups samples by their exact pattern of group memberships. Cells don't overlap and every group is exactly a set of cells, so every group's value, downsamplings included, is recovered by summing its cells. It uses the same reduction globals as before. force_leaf_groups still works, groups with no samples are kept as real leaves, and a freq_meta entry without a "group" key or a force_leaf_groups target not in freq_meta raises up front.

Tests cover both changes, including that cell reconstruction matches both the full aggregation and find_minimal_strata_groups. The /review skill was run with Opus and Sonnet 5; both of its findings were accepted and addressed (the breaking change above is documented, and index_sex_ploidy_flags was made public).

[Copilot]

🟡 Changes recommended

Missing strata values can break cell reduction, and sex-ploidy temporary fields can overwrite caller data.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Optimizes per-reference-site statistics by avoiding repeated VDS densification and aggregating disjoint membership cells.

Changes:

• Adds in-place sex-ploidy adjustment using precomputed row/column flags.
• Adds cell-based stratification reduction, including downsampling groups.
• Adds correctness and integration tests for both optimizations.

File summaries

File	Description
gnomad/sample_qc/sex.py	Implements reusable in-place ploidy adjustment.
gnomad/utils/annotations.py	Adds cell-based group reduction and reconstruction metadata.
gnomad/utils/sparse_mt.py	Uses the optimized ploidy path after densification.
tests/sample_qc/test_sex.py	Tests ploidy adjustment behavior.
tests/utils/test_annotations.py	Tests cell discovery and expression scaling.
tests/utils/test_sparse_mt.py	Tests integration, downsampling, empty groups, and ploidy.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 2
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[ch-kr]

thank you for adding this! a few efficiency notes

[ch-kr]

I didn't realize this comment was copied and pasted. I'm not sure I understand what it means; maybe we should just remove it?

[ch-kr]

I know this was copied and pasted from below, but maybe this flag should be in_autosome_or_par, which would remove the need for .when(~row_flags.in_non_par, gt_expr) below

[ch-kr]

I changed my mind since this would also mean updating

gnomad_methods/gnomad/utils/annotations.py

Line 697 in af5ccc1

def annotate_and_index_source_mt_for_sex_ploidy(

for a pretty trivial change, but I left these as comments in case this comes up again

[ch-kr]

maybe it would be better to update annotate_and_index_source_mt_for_sex_ploidy directly rather than creating a second option to adjust sex ploidy? it is a little against our usual conventions for this repo, but that function could take the source MT as input (rather than deriving the source using Hail internal code) and return column/row expressions so that these extra joins are never computed, regardless of dataset size. this would remove redundant looking code in adjust_sex_ploidy but would be a breaking change

[mike-w-wilson]

Good call, I took your advice and shortened things up a bit with helpers too. Like you said its a breaking change for the annotate_and_index with the change in params but currently (internally) only fix_freq_an in gnomad_qc/v4 uses this. We can mark this as a breaking change anyways and its an easy caller update for anyone who actually uses it. Worth the optimization to me.

[ch-kr]

group will always be defined for gnomad, but gemini flagged this:

The Issue:
If any dictionary in freq_meta is missing the "group" key, m.get("group") will return None.
Later, when you attempt to unpack the dictionary into hl.struct(**{...}), Python will encounter {None: hl.agg.counter(...)}. In Python, unpacking a dictionary as keyword arguments strictly requires the keys to be strings. This will instantly throw a TypeError: keywords must be strings.
The Fix:
Ensure a default string fallback or filter out missing groups.

[ch-kr]

gemini:

If force_leaf_groups contains duplicate target metadata dictionaries, forced.extend(matches) will append the same indices multiple times. This leads to duplicate leaf_indices and redundant checks.
Fix: Use a set or deduplicate forced before proceeding: forced = list(set(forced_raw)).

[ch-kr]

maybe one of these genotypes should be het, just to make this test more thorough, since s2 and s3 are both XY

[mike-w-wilson]

Yes good point

[ch-kr]

minor, but this check could be combined with the for loop enumerating freq_meta below. if you find i in force_leaf_groups, add it to forced, and then check the difference between force_leaf_groups and forced at the end of the loop. this would do a little extra work, since you'd build labels and idx_by_label but hit an error if any entries in force_leaf_groups aren't in freq_meta, but it would clean this up slightly

[ch-kr]

from gemini:

In your _pattern_expr function, you are using a Python list comprehension ([... for i in ...]) instead of Hail's native .map(). You must convert the Python list of indices into a Hail array literal and map over it natively. This shifts the loop execution to the JVM:

# FIXED (Executes natively in Hail Engine)
def _pattern_expr(label: str) -> hl.expr.StringExpression:
    indices_array = hl.literal(idx_by_label[label])
    return hl.delimit(
        indices_array.map(
            lambda i: hl.if_else(hl.coalesce(ht.group_membership[i], False), "1", "0")
        ),
        "",
    )```

[ch-kr]

none of the tests below seem to test running find_strata_cells with force_leaf_groups set; this should probably be covered.

I also don't know if I missed this, but should there be a test ensuring that the output of find_strata_cells and find_minimal_strata_groups are equivalent?

[mike-w-wilson]

back to you @ch-kr ! I took your suggestion around updating the annotate_and _index function instead to remove bloat and rewrote the PR description to include the breaking change.

[mike-w-wilson]

Ive added broadinstitute/gnomad_qc#787 which fixes the QC check. I originally was not going to change it but figured its simple change is easier than getting around the action.

[ch-kr]

two comments about test, and one separate efficiency issue that claude flagged

[ch-kr]

from claude, though it noted this doesn't need to be fixed as part of this PR since it existed before 844

Issue. In agg_by_strata, an entry_agg_group_membership target that resolves to a parent (not a leaf) builds its sample index list inline in the per-row expression (gnomad/utils/annotations.py:3458), by filtering all samples against every child position. Leaf targets read the precomputed indices_by_group global instead (annotations.py:3388). Hail does not hoist global-only sub-expressions out of a row map, so the parent list is rebuilt at every reference site at a cost of n_samples × n_children. Quick check (global-only filter over a 100k array, inline in annotate vs precomputed with annotate_globals):

rows	inline in row expr	precomputed global
2,000	0.6 s	0.2 s
40,000	2.5 s	0.5 s

This path predates the PR, but reduce_to_cells makes it heavier: under cells {"group": "adj"} is never a leaf, so it decomposes into every adj cell. Today the only parent pin is gnomAD coverage_stats -> {"group": "adj"} with ~800 samples and ~20 cells, so the current cost is negligible. On the AoU run (~400k samples, hundreds of downsampling cells) any future pin to adj would add 400k × n_cells work per site and give back a large share of the cell savings. Forcing {"group": "adj"} as a leaf avoids the scan but adds an extra all-samples AN aggregation per site, so it is the wrong fix for AoU.

Suggestion. Compute the per-target parent index lists once as globals and have the row expression read them. Output is unchanged; also speeds up the existing minimal-groups parent path.

ht = ht.annotate_globals(**global_expr)
if entry_agg_group_membership:
    ht = ht.annotate_globals(
        **{
            f"_{ann}_targets": hl.struct(
                **dict(zip(("s_indices", "adj"), _per_target_indices_and_adj(targets)))
            )
            for ann, targets in entry_agg_group_membership.items()
        }
    )

def _agg_for(ann, f):
    if ann in entry_agg_group_membership:
        t = ht[f"_{ann}_targets"]
        return _agg_by_group(t.s_indices, t.adj, agg_func=f[1], ann_expr=ht[ann])
    ...

(_per_target_indices_and_adj needs to be defined before the annotate_globals, and the _*_targets globals dropped with cols at the end.)

[mike-w-wilson]

Added since this is already about optimizations.

[ch-kr]

I asked claude if there were any missing tests that were important:

A pinned parent with many cells. The production gnomAD path pins coverage_stats to {"group": "adj"}, which under cells is a parent of every adj cell. The only cells-mode pin tested is max_called -> {"group": "raw"} in test_sparse_mt.py:966, and the raw label collapses to one cell, so the multi-child parent path in agg_by_strata is never exercised under cells. Add a non-summable agg pinned to {"group": "adj"} or {"group": "adj", "gen_anc": "afr"} and compare to the full run. The hl.agg.hist parent test at test_sparse_mt.py:703 covers this under minimal groups only, and the IR-shape bug it guards is exactly the kind that could differ with hundreds of children.

[mike-w-wilson]

Added however "many" is only 8 here. We can increase it but I'm not sure its needed.

[ch-kr]

since this function was updated, should we add a test?

[mike-w-wilson]

Added!