diff --git a/CLAUDE.md b/CLAUDE.md index 6075a631..6f37d533 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -17,8 +17,6 @@ The package includes a C extension (`_tsinfer`) built from `lib/` sources via se uv run pytest tests/ -v # Run all tests uv run pytest tests/test_matching.py # Run a single test file uv run pytest tests/test_matching.py::TestFoo::test_bar -v # Run a single test -uv run pytest --skip-slow # Skip slow tests - uv run ruff check --fix # Lint Python code (auto-fix) uv run ruff format # Format Python code ``` @@ -81,7 +79,7 @@ Source in `lib/`. Three main classes exposed to Python: - `AncestorBuilder` — builds inferred ancestors from genotype data - `AncestorMatcher` — Li & Stephens HMM matching algorithm -When changes are made to the C library, ensure that the ``_tskit`` module is rebuilt +When changes are made to the C library, ensure that the ``_tsinfer`` module is rebuilt before running Python tests. Vendored dependencies in `lib/subprojects/`: tskit C library and kastore. @@ -98,9 +96,22 @@ Sample VCZ → `infer_ancestors` → Ancestor VCZ → `match` → raw `tskit.Tre - Do not be overly defensive - defend only against circumstances that can occur within the current codebase. +- Do not make production code more complex for the sake of minimising + changes to the test suite. Simplicity and clarity of the production code + is imperative. +- Do not combine multiple complex operations in a single statement. Prefer + to keep a single operation per statement, and use intermediate variables + as a form of documentation. For example: + ```python + # Bad — multiple operations in one expression + result = sorted(k for k, v in mapping.items() if v in set(x.name for x in sources)) + + # Good — intermediate variable makes intent clear + source_names = {x.name for x in sources} + result = sorted(k for k, v in mapping.items() if v in source_names) + ``` - Prefer dataclasses over tuples when returning multiple values. - Use explicit `None` comparisons: `if x is not None` not `if x`. -- Zarr v3 is now used (dependency: `zarr>=3`). - Import all modules at the top of the file, not inside functions or methods. - Prefer importing a module and using module.function instead of using ``from module import function``. This applies to intra-package @@ -112,6 +123,7 @@ Sample VCZ → `infer_ancestors` → Ancestor VCZ → `match` → raw `tskit.Tre - When a parameter has a computed default derived from another parameter, compute it once at the point of use (the leaf function), not at every layer in the call chain. Pass `None` through intermediate layers. +- Zarr v3 is used (dependency: `zarr>=3`). Do not use Zarr v2 APIs. - Use PEP 604 union syntax: `int | None`, not `Optional[int]`. - One `logger = logging.getLogger(__name__)` per module at top level. @@ -125,3 +137,5 @@ Sample VCZ → `infer_ancestors` → Ancestor VCZ → `match` → raw `tskit.Tre - Test helpers are in `tests/helpers.py` (e.g., `make_sample_vcz`, `make_ancestor_vcz`) - `tests/algorithm.py` contains Python reference implementations used to verify C code - `msprime` is used to simulate test data +- Run the test suite with coverage before committing to ensure that new code is + fully covered by tests. diff --git a/docs/inference.md b/docs/inference.md index 81fbb715..207fc2c8 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -225,9 +225,10 @@ paths, "path compression". Matching ancestors is dependent on the time allocated to each ancestor; an ancestor can only copy from any older ancestor. For each ancestor, -we find the most likely path through older ancestors: that is the path that -maximises the product of the probabilities of recombination and mismatch -over all sites. +we find a copying path through older ancestors using a Li & Stephens +hidden Markov model. This determines, at each site, which older ancestor +the current ancestor copies from, allowing for both recombination +(switching between ancestors) and mismatch (mutations). :::{todo} Schematic of the ancestors copying process. @@ -279,4 +280,4 @@ The final phase of a `tsinfer` inference consists of a number steps: section 2. Describe the structure of the output tree sequences; how the nodes are mapped, what the time values mean, etc. -::: \ No newline at end of file +::: diff --git a/tests/test_config.py b/tests/test_config.py index 64f6973a..8769f80b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -159,7 +159,7 @@ def test_basic(self): }, output="out.trees", ) - assert m.path_compression is True + assert m.path_compression is False assert m.reference_ts is None assert m.sources["cohort"].node_flags == 1 assert m.sources["cohort"].create_individuals is True diff --git a/tests/test_grouping.py b/tests/test_grouping.py index 274fd0cb..b722740c 100644 --- a/tests/test_grouping.py +++ b/tests/test_grouping.py @@ -29,7 +29,7 @@ import numpy as np import pytest -from tsinfer import grouping +from tsinfer import config, grouping @dataclasses.dataclass @@ -558,7 +558,7 @@ class TestAssignGroups: """Test that assign_groups takes ungrouped MatchJobs and assigns groups.""" def _make_job(self, haplotype_index, time, start=0, end=100): - return grouping.MatchJob( + return config.MatchJob( haplotype_index=haplotype_index, source="test", sample_id=f"s{haplotype_index}", @@ -617,7 +617,7 @@ def test_sorted_by_group_then_haplotype_index(self): ) def test_preserves_job_fields(self): - job = grouping.MatchJob( + job = config.MatchJob( haplotype_index=0, source="my_source", sample_id="my_sample", diff --git a/tests/test_matcher_fixtures.py b/tests/test_matcher_fixtures.py index 7aa405a9..a7d159bd 100644 --- a/tests/test_matcher_fixtures.py +++ b/tests/test_matcher_fixtures.py @@ -29,8 +29,7 @@ import numpy as np import pytest -from tsinfer import matching, vcz -from tsinfer.grouping import MatchJob +from tsinfer import config, matching, vcz # --------------------------------------------------------------------------- # Helpers @@ -52,8 +51,20 @@ def read_haplotype(self, job): def _match(ts, positions, query, allele_mapper): """Run the current Matcher and return (path, mutations) as tuples.""" - matcher = matching.Matcher(ts, positions) - results = list(matcher.match([None], _ArrayReader([query]))) + job = config.MatchJob( + haplotype_index=0, + source="test", + sample_id="s0", + ploidy_index=0, + time=0, + start_position=0, + end_position=int(ts.sequence_length), + group=0, + ) + matcher = matching.Matcher( + ts, positions, source_parameters={"test": config.MatchSourceConfig()} + ) + results = list(matcher.match([job], _ArrayReader([query]))) _, r = results[0] path = [(s.left, s.right, s.parent) for s in r.path] mutations = [(m.position, m.derived_state) for m in r.mutations] @@ -62,7 +73,7 @@ def _match(ts, positions, query, allele_mapper): def _add_node(ts, time, path_segments, mutations, allele_mapper): """Add a single ancestor node to *ts* via extend_ts.""" - job = MatchJob( + job = config.MatchJob( haplotype_index=0, source="test", sample_id="s0", diff --git a/tests/test_matching.py b/tests/test_matching.py index 041993a0..31c911bb 100644 --- a/tests/test_matching.py +++ b/tests/test_matching.py @@ -25,7 +25,7 @@ import numpy as np import tskit -from tsinfer import grouping, matching, vcz +from tsinfer import config, matching, vcz # --------------------------------------------------------------------------- # Helpers @@ -92,8 +92,8 @@ def read_haplotype(self, job): def _jobs(n): - """Return *n* opaque dummy jobs (the reader ignores them).""" - return [None] * n + """Return *n* dummy MatchJobs (the reader ignores their content).""" + return [_make_job(haplotype_index=i) for i in range(n)] def _make_job( @@ -110,7 +110,7 @@ def _make_job( population_id=None, ): """Create a MatchJob with sensible defaults for testing.""" - return grouping.MatchJob( + return config.MatchJob( haplotype_index=haplotype_index, source=source, sample_id=sample_id, @@ -125,6 +125,9 @@ def _make_job( ) +_SOURCE_PARAMS = {"test": config.MatchSourceConfig()} + + # --------------------------------------------------------------------------- # TestMakeRootTs # --------------------------------------------------------------------------- @@ -252,7 +255,7 @@ def test_match_identical_haplotype(self): """Matching an identical haplotype returns a MatchResult with path edges.""" hap = np.array([0, 1, 0], dtype=np.int8) ts = self._make_ts_with_one_ancestor(hap) - matcher = matching.Matcher(ts, self.positions) + matcher = matching.Matcher(ts, self.positions, source_parameters=_SOURCE_PARAMS) results = list(matcher.match(_jobs(1), _ArrayReader([hap]))) assert len(results) == 1 _, r = results[0] @@ -265,7 +268,7 @@ def test_match_with_mutation(self): ts = self._make_ts_with_one_ancestor(ancestor_hap) # Query haplotype has derived state at site 1 query_hap = np.array([0, 1, 0], dtype=np.int8) - matcher = matching.Matcher(ts, self.positions) + matcher = matching.Matcher(ts, self.positions, source_parameters=_SOURCE_PARAMS) results = list(matcher.match(_jobs(1), _ArrayReader([query_hap]))) _, r = results[0] # There should be a mutation at position 20 (site index 1) @@ -277,7 +280,7 @@ def test_match_all_missing(self): hap = np.array([0, 1, 0], dtype=np.int8) ts = self._make_ts_with_one_ancestor(hap) missing_hap = np.array([-1, -1, -1], dtype=np.int8) - matcher = matching.Matcher(ts, self.positions) + matcher = matching.Matcher(ts, self.positions, source_parameters=_SOURCE_PARAMS) results = list(matcher.match(_jobs(1), _ArrayReader([missing_hap]))) assert len(results) == 1 @@ -286,7 +289,7 @@ def test_match_partial_missing(self): hap = np.array([0, 1, 0], dtype=np.int8) ts = self._make_ts_with_one_ancestor(hap) query_hap = np.array([-1, 1, 0], dtype=np.int8) - matcher = matching.Matcher(ts, self.positions) + matcher = matching.Matcher(ts, self.positions, source_parameters=_SOURCE_PARAMS) results = list(matcher.match(_jobs(1), _ArrayReader([query_hap]))) assert len(results) == 1 _, r = results[0] @@ -299,7 +302,7 @@ def test_match_returns_one_result_per_haplotype(self): hap = np.array([0, 1, 0], dtype=np.int8) ts = self._make_ts_with_one_ancestor(hap) haplotypes = np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]], dtype=np.int8) - matcher = matching.Matcher(ts, self.positions) + matcher = matching.Matcher(ts, self.positions, source_parameters=_SOURCE_PARAMS) results = list(matcher.match(_jobs(3), _ArrayReader(haplotypes))) assert len(results) == 3 @@ -307,7 +310,7 @@ def test_match_result_types(self): """MatchResult should contain PathSegment and Mutation objects.""" hap = np.array([0, 1, 0], dtype=np.int8) ts = self._make_ts_with_one_ancestor(hap) - matcher = matching.Matcher(ts, self.positions) + matcher = matching.Matcher(ts, self.positions, source_parameters=_SOURCE_PARAMS) results = list(matcher.match(_jobs(1), _ArrayReader([hap]))) _, r = results[0] assert isinstance(r.path, list) @@ -322,7 +325,7 @@ def test_path_segments_have_valid_coordinates(self): """PathSegment left < right for every segment.""" hap = np.array([0, 1, 0], dtype=np.int8) ts = self._make_ts_with_one_ancestor(hap) - matcher = matching.Matcher(ts, self.positions) + matcher = matching.Matcher(ts, self.positions, source_parameters=_SOURCE_PARAMS) results = list(matcher.match(_jobs(1), _ArrayReader([hap]))) _, r = results[0] for seg in r.path: @@ -333,7 +336,7 @@ def test_mutations_have_valid_fields(self): hap = np.array([0, 0, 0], dtype=np.int8) ts = self._make_ts_with_one_ancestor(hap) query = np.array([1, 1, 0], dtype=np.int8) - matcher = matching.Matcher(ts, self.positions) + matcher = matching.Matcher(ts, self.positions, source_parameters=_SOURCE_PARAMS) results = list(matcher.match(_jobs(1), _ArrayReader([query]))) _, r = results[0] assert len(r.mutations) > 0 @@ -342,6 +345,19 @@ def test_mutations_have_valid_fields(self): assert isinstance(m.position, float) assert isinstance(m.derived_state, int) + def test_match_with_custom_recombination_and_mismatch(self): + """Custom per-source recombination/mismatch overrides are used.""" + hap = np.array([0, 1, 0], dtype=np.int8) + ts = self._make_ts_with_one_ancestor(hap) + source_params = { + "test": config.MatchSourceConfig(recombination=0.5, mismatch=0.1), + } + matcher = matching.Matcher(ts, self.positions, source_parameters=source_params) + results = list(matcher.match(_jobs(1), _ArrayReader([hap]))) + assert len(results) == 1 + _, r = results[0] + assert len(r.path) > 0 + # --------------------------------------------------------------------------- # TestExtendTs @@ -673,7 +689,7 @@ def _build_root_ts(self): def _match_and_pair(self, ts, haplotypes): """Match haplotypes against ts and return paired_results list.""" - matcher = matching.Matcher(ts, self.positions) + matcher = matching.Matcher(ts, self.positions, source_parameters=_SOURCE_PARAMS) results = list(matcher.match(_jobs(len(haplotypes)), _ArrayReader(haplotypes))) return results @@ -705,7 +721,9 @@ def test_matching_after_two_extends(self): # Now match a sample against ts2 sample_hap = np.array([0, 0, 1, 1, 0], dtype=np.int8) - matcher2 = matching.Matcher(ts2, self.positions) + matcher2 = matching.Matcher( + ts2, self.positions, source_parameters=_SOURCE_PARAMS + ) results2 = list(matcher2.match(_jobs(1), _ArrayReader([sample_hap]))) _, r = results2[0] assert len(r.path) > 0 diff --git a/tests/test_vcz.py b/tests/test_vcz.py index e9f5f65d..0ba8d748 100644 --- a/tests/test_vcz.py +++ b/tests/test_vcz.py @@ -30,7 +30,7 @@ import pytest import zarr -from tsinfer import config, grouping, vcz +from tsinfer import config, vcz # --------------------------------------------------------------------------- # Fixtures @@ -346,7 +346,7 @@ def test_ancestor_haplotype(self): schedule=schedule, ) # a0 has genotype [0, 1, 0] - job = grouping.MatchJob( + job = config.MatchJob( haplotype_index=1, source="ancestors", sample_id="a0", @@ -371,7 +371,7 @@ def test_ancestor_haplotype_second(self): schedule=schedule, ) # a1 has genotype [1, 0, 1] - job = grouping.MatchJob( + job = config.MatchJob( haplotype_index=2, source="ancestors", sample_id="a1", @@ -396,7 +396,7 @@ def test_sample_haplotype_encoding(self): schedule=schedule, ) # sample_0: gt [0, 1, 1] → encoded [0=anc, 1=derived, 1=derived] - job = grouping.MatchJob( + job = config.MatchJob( haplotype_index=3, source="test", sample_id="sample_0", @@ -421,7 +421,7 @@ def test_sample_haplotype_second(self): schedule=schedule, ) # sample_1: gt [1, 0, 0] → encoded [1=derived, 0=anc, 0=anc] - job = grouping.MatchJob( + job = config.MatchJob( haplotype_index=4, source="test", sample_id="sample_1", @@ -467,7 +467,7 @@ def test_missing_genotype_encoded_as_minus_one(self): anc_alleles, schedule=schedule, ) - job = grouping.MatchJob( + job = config.MatchJob( haplotype_index=2, source="test", sample_id="sample_0", @@ -519,7 +519,7 @@ def test_multiple_sources(self): schedule = [("src_a", "sample_0", 0), ("src_b", "sample_0", 0)] reader = vcz.HaplotypeReader(sources, positions, anc_alleles, schedule=schedule) - job_a = grouping.MatchJob( + job_a = config.MatchJob( haplotype_index=2, source="src_a", sample_id="sample_0", @@ -529,7 +529,7 @@ def test_multiple_sources(self): end_position=200, group=2, ) - job_b = grouping.MatchJob( + job_b = config.MatchJob( haplotype_index=3, source="src_b", sample_id="sample_0", @@ -1475,7 +1475,7 @@ def test_facade_returns_alleles_after_reads(self): reader = vcz.HaplotypeReader(sources, positions, anc_alleles, schedule=schedule) # Read from both sources to trigger allele discovery - job_a = grouping.MatchJob( + job_a = config.MatchJob( haplotype_index=0, source="src_a", sample_id="sample_0", @@ -1485,7 +1485,7 @@ def test_facade_returns_alleles_after_reads(self): end_position=200, group=0, ) - job_b = grouping.MatchJob( + job_b = config.MatchJob( haplotype_index=1, source="src_b", sample_id="sample_0", diff --git a/tsinfer/config.py b/tsinfer/config.py index d7b48898..1d2af67a 100644 --- a/tsinfer/config.py +++ b/tsinfer/config.py @@ -157,6 +157,26 @@ class MatchSourceConfig: node_flags: int = DEFAULT_NODE_FLAGS # tskit.NODE_IS_SAMPLE create_individuals: bool = True + recombination: float | None = None + mismatch: float | None = None + + +@dataclasses.dataclass +class MatchJob: + """One row of the compute-groups output — one per haplotype.""" + + haplotype_index: int + source: str + sample_id: str + ploidy_index: int + time: float + start_position: int + end_position: int + group: int + node_flags: int = 1 + individual_id: int | None = None + population_id: int | None = None + sample_chunk: int = 0 @dataclasses.dataclass @@ -165,7 +185,7 @@ class MatchConfig: sources: dict[str, MatchSourceConfig] output: str | pathlib.Path - path_compression: bool = True + path_compression: bool = False reference_ts: str | pathlib.Path | None = None workdir: str | pathlib.Path | None = None keep_intermediates: bool = False @@ -396,7 +416,12 @@ def from_toml(cls, path: str | pathlib.Path) -> Config: "keep_intermediates", } -_KNOWN_MATCH_SOURCE_KEYS = {"node_flags", "create_individuals"} +_KNOWN_MATCH_SOURCE_KEYS = { + "node_flags", + "create_individuals", + "recombination", + "mismatch", +} _KNOWN_INDIVIDUAL_METADATA_KEYS = {"fields", "population"} @@ -512,16 +537,20 @@ def _parse_match(raw: dict) -> MatchConfig: src_val, _KNOWN_MATCH_SOURCE_KEYS, ) + raw_recomb = src_val.get("recombination") + raw_mismatch = src_val.get("mismatch") sources[src_name] = MatchSourceConfig( node_flags=int(src_val.get("node_flags", DEFAULT_NODE_FLAGS)), create_individuals=bool(src_val.get("create_individuals", True)), + recombination=float(raw_recomb) if raw_recomb is not None else None, + mismatch=float(raw_mismatch) if raw_mismatch is not None else None, ) else: sources[src_name] = MatchSourceConfig() return MatchConfig( sources=sources, output=_resolve_path(entry["output"]), - path_compression=bool(entry.get("path_compression", True)), + path_compression=bool(entry.get("path_compression", False)), reference_ts=_resolve_path(entry.get("reference_ts")), workdir=_resolve_path(entry.get("workdir")), keep_intermediates=bool(entry.get("keep_intermediates", False)), diff --git a/tsinfer/grouping.py b/tsinfer/grouping.py index 7652bcc3..f33613b6 100644 --- a/tsinfer/grouping.py +++ b/tsinfer/grouping.py @@ -30,29 +30,13 @@ import numba import numpy as np +from . import config + logging.getLogger("numba").setLevel(logging.WARNING) logger = logging.getLogger(__name__) -@dataclasses.dataclass -class MatchJob: - """One row of the compute-groups output — one per haplotype.""" - - haplotype_index: int - source: str - sample_id: str - ploidy_index: int - time: float - start_position: int - end_position: int - group: int - node_flags: int = 1 - individual_id: int | None = None - population_id: int | None = None - sample_chunk: int = 0 - - -def assign_groups(jobs: list[MatchJob]) -> list[MatchJob]: +def assign_groups(jobs: list[config.MatchJob]) -> list[config.MatchJob]: """Assign group indices to MatchJobs, return sorted by (group, haplotype_index).""" if len(jobs) == 0: return jobs diff --git a/tsinfer/matching.py b/tsinfer/matching.py index 477a22f6..cb1c97f2 100644 --- a/tsinfer/matching.py +++ b/tsinfer/matching.py @@ -33,7 +33,7 @@ import _tsinfer -from . import arg_ops, grouping, vcz +from . import arg_ops, config, vcz logger = logging.getLogger(__name__) @@ -156,10 +156,12 @@ def __init__( ts: tskit.TreeSequence, positions: np.ndarray, # (num_sites,) int32 — inference site positions num_alleles: np.ndarray | None = None, + source_parameters: dict[str, config.MatchSourceConfig] | None = None, ): self._positions = np.asarray(positions, dtype=np.int32) self._num_sites = len(positions) self._sequence_length = ts.sequence_length + self._source_parameters = source_parameters or {} logger.info( "Creating Matcher: ts=%d nodes, %d edges, %.1f MiB, RSS=%.1f MiB", @@ -177,17 +179,27 @@ def __init__( time_.monotonic() - t0, _rss_mib(), ) - self._rho = np.full(self._num_sites, 1e-2) - self._mu = np.full(self._num_sites, 1e-20) def _match_one(self, job, reader) -> tuple: """Match a single haplotype: read, run HMM, convert to result.""" num_sites = self._num_sites pos = self._positions.astype(np.float64) + # Look up per-source recombination/mismatch overrides + source = job.source + src_cfg = self._source_parameters[source] + recombination = 1e-2 + if src_cfg.recombination is not None: + recombination = src_cfg.recombination + mismatch = 1e-20 + if src_cfg.mismatch is not None: + mismatch = src_cfg.mismatch + rho = np.full(num_sites, recombination) + mu = np.full(num_sites, mismatch) + # Each call gets its own matcher so threads don't share state t0 = time_.monotonic() - matcher = AncestorMatcher(self._matcher_indexes, self._rho, self._mu) + matcher = AncestorMatcher(self._matcher_indexes, rho, mu) t_init = time_.monotonic() - t0 t0 = time_.monotonic() @@ -274,7 +286,7 @@ def match( def extend_ts( ts: tskit.TreeSequence, *, - paired_results: list[tuple[grouping.MatchJob, MatchResult]], + paired_results: list[tuple[config.MatchJob, MatchResult]], allele_mapper: vcz.AlleleMapper, provenance_record: str | None = None, ) -> tskit.TreeSequence: diff --git a/tsinfer/pipeline.py b/tsinfer/pipeline.py index 8fcbe2ec..0cfaef46 100644 --- a/tsinfer/pipeline.py +++ b/tsinfer/pipeline.py @@ -64,7 +64,7 @@ def _build_jobs_with_metadata(cfg): # Build MatchJobs directly from VCZ stores, assigning individual IDs # in the same pass. - jobs: list[grouping.MatchJob] = [] + jobs: list[config.MatchJob] = [] ind_key_to_id: dict[tuple[str, str], int] = {} hap_idx = 0 @@ -159,7 +159,7 @@ def _build_jobs_with_metadata(cfg): for p in range(ploidy): jobs.append( - grouping.MatchJob( + config.MatchJob( haplotype_index=hap_idx, source=source.name, sample_id=sid, @@ -292,10 +292,10 @@ def compute_groups_json(cfg: config.Config) -> str: return json_str -def _load_match_jobs(path) -> list[grouping.MatchJob]: +def _load_match_jobs(path) -> list[config.MatchJob]: """Read a match-jobs JSON file and return a list of MatchJob objects.""" records = json.loads(pathlib.Path(path).read_text()) - return [grouping.MatchJob(**rec) for rec in records] + return [config.MatchJob(**rec) for rec in records] def _setup_workdir(workdir, cfg): @@ -356,6 +356,7 @@ def _process_group( match_fh, config_str: str, progress: bool, + source_parameters: dict[str, config.MatchSourceConfig] | None = None, ) -> tskit.TreeSequence: """Match one group of haplotypes and return the extended tree sequence.""" with provenance.TimingAndMemory() as tm: @@ -363,6 +364,7 @@ def _process_group( ts, positions, num_alleles=reader.get_num_alleles(), + source_parameters=source_parameters, ) job_list = [job for _, job in group_jobs] match_iter = m.match(job_list, reader, num_threads=num_threads) @@ -592,6 +594,7 @@ def match( match_fh=match_fh, config_str=config_str, progress=progress, + source_parameters=cfg.match.sources, ) completed_haps += num_in_group rss = psutil.Process().memory_info().rss / (1024 * 1024)