feat(fine-mapping): planner for fine-map study combinations.

[project-defiant]

✨ Context

FineMappingPlanGeneratorStep reads the study index and emits a FineMappingPlanner dataset that routes each study to a fine-mapping method based on a set of per-study eligibility constraints. This PR implements the constraint model, the MultiSuSiE constraint set, and the step itself, wired into the Hydra config surface.

🛠 What does this PR implement

• MethodConstraint model (method/fine_mapping/constraint.py): a Protocol-based per-study constraint abstraction (IsAllowedStudyType, HasSumstats, HasMappedTrait, PassSumstatQC, HasAllowedAnalysisFlags, HasAllowedMajorAncestry), each annotating a StudyIndex with a boolean verdict via ConstraintResult.
• MultiSuSiEConstraintSet (method/fine_mapping/constraint_set.py): composes the constraints above, resolves per-study eligibility, selects one representative study per (trait, major ancestry) group by effective sample size, and groups eligible studies sharing a trait across ancestries into a deterministic runId for multi-ancestry runs.
  • resolve() is decomposed into named, independently-readable stages (constraint resolution, n_eff/major-ancestry derivation, representative-study selection, run-id assignment) instead of one long chained expression.
  • The two independent branches (constraint resolution vs. representative-selection inputs) are materialized concurrently via a small thread pool, since Spark's own scheduler doesn't parallelize them on its own at this data scale.
  • Fixed a severe performance bug in the runId assignment: an unbounded window aggregation was evaluated over the (large) null-runId partition unnecessarily; scoping it to non-null rows only cut a real-data run from ~12-15 minutes to ~23 seconds.
  • Added a study-count sanity check (input vs. output distinct studyId) that raises if resolution ever drops or duplicates studies.
• FineMappingPlanner dataset (dataset/fine_mapping.py): the output schema (studyId, runId, constraints, route), plus FineMappingRoute and a unionByName-based __add__ for combining multiple constraint sets' plans (documented as unique per (studyId, route), not studyId alone).
• FineMappingPlanGeneratorStep (finemapping_planner.py): reads a study index, resolves the registered constraint set(s), and writes the combined plan as parquet partitioned by route, sorted by (route, runId, studyId).
• Hydra wiring: FineMappingPlanGeneratorConfig registered in config.py so the step is runnable via the standard CLI (step=fine_mapping_plan_generator).
• Fixed pre-existing bugs found along the way:
  • common/spark.py's all_struct_fields_in_array built invalid SQL via f-string-interpolated Column objects; its own doctest was failing at HEAD.
  • study_index.py: minor validation fixes surfaced while building validate_ccs()-dependent tests.
• Docs: added docs/python_api/datasets/fine_mapping.md, docs/python_api/methods/fine_mapping_constraints.md, docs/python_api/steps/fine_mapping_plan_generator.md.
• Tests: doctests for each constraint's happy path, pytest coverage for edge cases (null/empty traits, ancestry ordering, eligibility gating) and MultiSuSiEConstraintSet.resolve()'s grouping/run-id logic, a fast unit test of the step's orchestration, and a mocked integration test exercising the real Spark write path end-to-end.

🙈 Missing

• Only the MultiSuSiE route is implemented; PICS and SuSiE-inf routes are separate future work.
• The multi-ancestry LD population scheme (gnomAD-shaped vs. Pan-UKBB-native) is an open question — raised separately for team discussion, not resolved in this PR.

🚦 Before submitting

• Do these changes cover one single feature (one change at a time)?
• Did you read the contributor guideline?
• Did you make sure to update the documentation with your changes?
• Did you make sure there is no commented out code in this PR?
• Did you follow conventional commits standards in PR title and commit messages?
• Did you make sure the branch is up-to-date with the dev branch?
• Did you write any new necessary tests?
• Did you make sure the changes pass local tests (make test)?
• Did you make sure the changes pass pre-commit rules (e.g uv run pre-commit run --all-files)?

Note: the "branch up-to-date with dev" checkbox is left unchecked — worth verifying/rebasing before merge, since I didn't check that as part of this session's work.

[addramir]

I don't understand this function. Isn't it always present? And why not to use DiseaseIDs?
Is it happening before study index validation?

[project-defiant]

This checks if the traitFromSourceMappedIds are in the studyIndex. In the staging bucket there are a few cases where the study does not have any efos. Examples below.

+--------------+--------------------------+
|   studyId    | traitFromSourceMappedIds |
+--------------+--------------------------+
| GCST90667999 | NULL                     |
| GCST90727244 | NULL                     |
| GCST90668000 | NULL                     |
| GCST90667998 | NULL                     |
+--------------+--------------------------+

I raised it to GWAS Catalog a new months ago, but the staging still has these issues on the slack with them

| STUDY ACCESSION |              DISEASE/TRAIT              | MAPPED_TRAIT |
|-----------------|-----------------------------------------|--------------|
| GCST90727244    | Peripheral venous and lymphatic disease | NULL         |
| GCST90667998    | Kidney cortex volume                    | NULL         |
| GCST90667999    | Kidney medulla volume                   | NULL         |
| GCST90668000    | Kidney sinus volume                     | NULL         |

AS to your other question, we can not use DIseaseIds, as these are populated only after the study validation step, which is part of the Pipeline run every 3 months, but this process shall be decoupled from that, as I do not want to run the fine-mapping during Pipeline (this would increase release process way too much). The clear scenario is to use traitFromSourceMappedIds at this stage and later determine which credible sets are derived from truely correct efo term studies during validation.

[addramir]

Is it the same we used in SuSiE?