Phase 3: taxon milestone computation#12
Draft
fchen13 wants to merge 1 commit into
Draft
Conversation
Add per-taxon assembly milestone tracking, computing when each taxon (at species and every higher rank) first reached four milestones — any assembly, EBP-affiliated assembly, meets-metric, and affiliated-and-meets-metric — from the version-tracked assembly TSVs that GoaT itself does not retain. - flows/lib/load_taxonomy.py: dev/test NCBI taxdump parser building the taxonomy contract (taxid -> scientific_name, rank, parent, canonical-rank lineage). Cycle-guarded parent-chain walks. Production uses upstream lineage. - flows/lib/compute_taxon_milestones.py: single chronological sweep producing species and higher-rank milestones; first-in-ranks recorded as a by-product; per-species latest_assembly / current / total counts. Emits compute.taxon_milestones.completed. Writes taxon_milestone_summary.tsv. - tests/test_taxon_milestones.py: 38 tests covering the parser, resolve_to_species, the sweep (each milestone earliest, tie-break, empty/None dates, in-ranks, latest-assembly), output schema, and validation invariants. - tests/README_phase_3_taxon_milestones.md: run/test guide and schema.
Contributor
Reviewer's GuideImplements Phase 3 of the assembly version tracking system: a taxonomy loader and a single-pass taxon milestone computation flow that reads version-tracked assembly TSVs, resolves taxids to species + canonical lineage, computes four milestone “firsts” across all ranks in one chronological sweep, and writes taxon_milestone_summary.tsv, with a focused test suite and documentation describing usage and invariants. Sequence diagram for the compute_taxon_milestones flowsequenceDiagram
participant Driver as compute_taxon_milestones
participant Taxdump as taxdump_dir
participant Taxonomy as build_taxonomy
participant Assemblies as assembly_TSVs
participant Milestones as compute_milestones
participant Output as build_output_rows
participant Storage as taxon_milestone_summary.tsv
participant Events as emit_event
Driver->>Taxonomy: build_taxonomy(taxdump_path)
Taxonomy-->>Driver: taxonomy
Driver->>Assemblies: load_assemblies(assembly_current.tsv, assembly_historical.tsv)
Assemblies-->>Driver: rows
alt rows is empty
Driver->>Events: emit_event(event="compute.taxon_milestones.completed")
Driver-->>Driver: return
else rows is non_empty
Driver->>Milestones: compute_milestones(rows, taxonomy)
Milestones-->>Driver: taxa, species_extra
Driver->>Output: build_output_rows(taxa, species_extra)
Output-->>Driver: output_rows
Driver->>Storage: csv.DictWriter(writerows(output_rows))
Driver->>Events: emit_event(event="compute.taxon_milestones.completed")
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- The console
printusage incompute_taxon_milestones(including per-row skip messages) may get very noisy at scale; consider routing these through the existing logging/Prefect facilities or gating verbose output behind a flag. - In
compute_milestones, theempty_datecounter is initialized and then immediately reassigned after filtering; you can drop the initial assignment and only compute it once fromlen(enriched) - len(dated)to avoid confusion.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The console `print` usage in `compute_taxon_milestones` (including per-row skip messages) may get very noisy at scale; consider routing these through the existing logging/Prefect facilities or gating verbose output behind a flag.
- In `compute_milestones`, the `empty_date` counter is initialized and then immediately reassigned after filtering; you can drop the initial assignment and only compute it once from `len(enriched) - len(dated)` to avoid confusion.
## Individual Comments
### Comment 1
<location path="tests/test_taxon_milestones.py" line_range="255" />
<code_context>
+ taxa, _ = compute_milestones(rows, taxonomy)
+ # 'None' string must not count as metric; first metric is GCA_002.1.
+ assert taxa[ASELLUS_AQUATICUS]["firsts"]["metric"][1] == "GCA_002.1"
+ assert "metric" not in {k for k in taxa[ASELLUS_AQUATICUS]["firsts"] if False}
+
+ def test_ebp_metric_requires_both(self, taxonomy):
</code_context>
<issue_to_address>
**issue (testing):** The assertion in `test_metric_predicate_and_none_string` is effectively a no-op and does not test anything.
Because the comprehension `{k for k in taxa[ASELLUS_AQUATICUS]["firsts"] if False}` is always empty, this assertion will always succeed and never fail based on the data. If you want to verify the `firsts` content, consider asserting directly on the keys, e.g. `assert set(taxa[ASELLUS_AQUATICUS]["firsts"].keys()) == {"metric"}` or whatever set of milestones you expect.
</issue_to_address>
### Comment 2
<location path="tests/test_taxon_milestones.py" line_range="219" />
<code_context>
+# The chronological sweep
+# ---------------------------------------------------------------------------
+
+class TestSweep:
+ def test_first_assembly_is_earliest_row(self, taxonomy):
+ rows = [
</code_context>
<issue_to_address>
**suggestion (testing):** Missing test coverage for rows with non-integer or malformed `taxId` values being safely skipped.
`compute_milestones` treats rows with non‑integer `taxId` values as `skipped`, but this behavior is currently untested. Please add a test (e.g., with `taxId="abc"` or empty) that confirms these rows are excluded from the output and do not raise errors.
Suggested implementation:
```python
class TestSweep:
def test_first_assembly_is_earliest_row(self, taxonomy):
rows = [
make_row("GCA_002.1", ASELLUS_AQUATICUS, "2020-01-01"),
make_row("GCA_001.1", ASELLUS_AQUATICUS, "2018-05-05"),
]
taxa, _ = compute_milestones(rows, taxonomy)
assert taxa[ASELLUS_AQUATICUS]["firsts"]["assembly"] == ("2018-05-05", "GCA_001.1")
def test_non_integer_taxid_rows_are_skipped(self, taxonomy):
rows = [
make_row("GCA_BAD.1", "abc", "2019-01-02"),
make_row("GCA_001.1", ASELLUS_AQUATICUS, "2018-05-05"),
]
taxa, skipped = compute_milestones(rows, taxonomy)
# Rows with non-integer taxId should be safely skipped and not appear in the taxa output
assert ASELLUS_AQUATICUS in taxa
assert "abc" not in taxa
assert any(row["accession"] == "GCA_BAD.1" for row in skipped)
```
Depending on the exact structure of elements in `skipped`, you may need to adjust the last assertion:
- If `skipped` is a list of dicts, `row["accession"]` is appropriate.
- If `skipped` contains row objects or tuples, update the accessor accordingly (e.g., `row.accession` or `row["assembly_accession"]`).
You may also want to add a variant using an empty string taxId (e.g., `taxId=""`) following the same pattern if `compute_milestones` distinguishes between different malformed values.
</issue_to_address>
### Comment 3
<location path="tests/test_taxon_milestones.py" line_range="312-321" />
<code_context>
+ assert species_extra[ASELLUS_AQUATICUS]["total_assemblies"] == 2
+ assert species_extra[ASELLUS_AQUATICUS]["current_assemblies"] == 1
+
+ def test_latest_is_current_version(self, taxonomy):
+ # Latest = the current (non-superseded) version, even if an out-of-order
+ # superseded row has a newer date it must not win over a current one.
+ rows = [
+ make_row("GCA_001.1", ASELLUS_AQUATICUS, "2018-01-01", version_status="superseded"),
+ make_row("GCA_001.2", ASELLUS_AQUATICUS, "2021-01-01", version_status="current"),
+ ]
+ _, species_extra = compute_milestones(rows, taxonomy)
+ latest = species_extra[ASELLUS_AQUATICUS]["latest"]
+ assert (latest["date"], latest["accession"]) == ("2021-01-01", "GCA_001.2")
+
+ def test_latest_prefers_current_over_newer_superseded(self, taxonomy):
</code_context>
<issue_to_address>
**suggestion (testing):** Consider adding a test for the interplay between undated current assemblies and dated superseded ones in latest-selection logic.
Current tests cover dated current vs dated superseded rows and superseded-only species, but not the case where a current row has an empty `releaseDate` and a superseded row has a non-empty date. Please add a test that creates this scenario (current with empty date, superseded with a later date) and asserts that the current assembly is still selected as `latest`, to guard the `(is_current, date, accession)` tie-breaking logic against regressions.
</issue_to_address>
### Comment 4
<location path="flows/lib/compute_taxon_milestones.py" line_range="169" />
<code_context>
+ return None
+
+
+def compute_milestones(rows: list[dict], taxonomy: dict) -> tuple[dict, dict]:
+ """Run the single chronological sweep over all assembly rows.
+
</code_context>
<issue_to_address>
**issue (complexity):** Consider refactoring latest-assembly tracking and the milestone sweep into small helpers to simplify `compute_milestones` and `build_output_rows` without changing behavior.
You can reduce complexity in `compute_milestones` and `build_output_rows` with a couple of small, local refactors that keep behavior identical:
### 1. Simplify latest-assembly tracking
The nested `latest` dict with a `key` tuple adds indirection. You can track the candidate tuple directly and derive `date`/`accession` only when needed:
```python
def _better_latest(current: tuple | None, candidate: tuple) -> bool:
"""Return True if candidate is a better 'latest' than current."""
if current is None:
return True
return candidate > current # (is_current, date, accession)
```
Use it in `compute_milestones`:
```python
# inside species_extra.setdefault(...)
"latest": None, # (is_current, date, accession)
# ...
candidate = (is_current, item["date"], item["accession"])
if _better_latest(extra["latest"], candidate):
extra["latest"] = candidate
```
And in `build_output_rows`:
```python
if is_species:
extra = species_extra.get(taxid, {})
latest = extra.get("latest")
if latest:
_, latest_date, latest_accession = latest
row["latest_assembly_date"] = latest_date
row["latest_assembly_accession"] = latest_accession
else:
row["latest_assembly_date"] = ""
row["latest_assembly_accession"] = ""
```
This removes the `latest["key"]` indirection and makes the comparison logic more obvious without changing semantics.
### 2. Extract the milestone sweep into a helper
The nested loops over `MILESTONES` and `LEVELS` make the core sweep dense. A small helper reduces nesting and centralizes the “apply milestone for a row” logic:
```python
def _apply_milestones_for_row(
item: dict,
taxonomy: dict,
taxa: dict,
species_extra: dict,
seen: dict[str, set[int]],
) -> None:
sp = item["species_taxid"]
lineage = item["lineage"]
def ensure_taxon(taxid: int):
if taxid not in taxa:
node = taxonomy.get(taxid, {})
taxa[taxid] = {
"rank": node.get("rank", ""),
"scientific_name": node.get("scientific_name", ""),
"firsts": {},
}
for milestone in MILESTONES:
key = milestone["key"]
if not milestone_predicate(item["row"], key):
continue
for level in LEVELS:
taxid = sp if level == "species" else lineage.get(level)
if not taxid or taxid in seen[key]:
continue
seen[key].add(taxid)
ensure_taxon(taxid)
taxa[taxid]["firsts"][key] = (item["date"], item["accession"])
if milestone["in_ranks"] and level != "species":
species_extra[sp]["in_ranks"][key].append(level)
```
Then `compute_milestones` becomes:
```python
taxa: dict[int, dict] = {}
seen = {m["key"]: set() for m in MILESTONES}
for item in dated:
_apply_milestones_for_row(item, taxonomy, taxa, species_extra, seen)
for sp in species_extra:
# ensure species rows exist
if sp not in taxa:
node = taxonomy.get(sp, {})
taxa[sp] = {
"rank": node.get("rank", ""),
"scientific_name": node.get("scientific_name", ""),
"firsts": {},
}
```
This keeps the chronological sweep and all milestone behavior intact, but flattens the nested control flow and makes the milestone logic easier to follow and test in isolation.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| taxa, _ = compute_milestones(rows, taxonomy) | ||
| # 'None' string must not count as metric; first metric is GCA_002.1. | ||
| assert taxa[ASELLUS_AQUATICUS]["firsts"]["metric"][1] == "GCA_002.1" | ||
| assert "metric" not in {k for k in taxa[ASELLUS_AQUATICUS]["firsts"] if False} |
Contributor
There was a problem hiding this comment.
issue (testing): The assertion in test_metric_predicate_and_none_string is effectively a no-op and does not test anything.
Because the comprehension {k for k in taxa[ASELLUS_AQUATICUS]["firsts"] if False} is always empty, this assertion will always succeed and never fail based on the data. If you want to verify the firsts content, consider asserting directly on the keys, e.g. assert set(taxa[ASELLUS_AQUATICUS]["firsts"].keys()) == {"metric"} or whatever set of milestones you expect.
| # The chronological sweep | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
| class TestSweep: |
Contributor
There was a problem hiding this comment.
suggestion (testing): Missing test coverage for rows with non-integer or malformed taxId values being safely skipped.
compute_milestones treats rows with non‑integer taxId values as skipped, but this behavior is currently untested. Please add a test (e.g., with taxId="abc" or empty) that confirms these rows are excluded from the output and do not raise errors.
Suggested implementation:
class TestSweep:
def test_first_assembly_is_earliest_row(self, taxonomy):
rows = [
make_row("GCA_002.1", ASELLUS_AQUATICUS, "2020-01-01"),
make_row("GCA_001.1", ASELLUS_AQUATICUS, "2018-05-05"),
]
taxa, _ = compute_milestones(rows, taxonomy)
assert taxa[ASELLUS_AQUATICUS]["firsts"]["assembly"] == ("2018-05-05", "GCA_001.1")
def test_non_integer_taxid_rows_are_skipped(self, taxonomy):
rows = [
make_row("GCA_BAD.1", "abc", "2019-01-02"),
make_row("GCA_001.1", ASELLUS_AQUATICUS, "2018-05-05"),
]
taxa, skipped = compute_milestones(rows, taxonomy)
# Rows with non-integer taxId should be safely skipped and not appear in the taxa output
assert ASELLUS_AQUATICUS in taxa
assert "abc" not in taxa
assert any(row["accession"] == "GCA_BAD.1" for row in skipped)Depending on the exact structure of elements in skipped, you may need to adjust the last assertion:
- If
skippedis a list of dicts,row["accession"]is appropriate. - If
skippedcontains row objects or tuples, update the accessor accordingly (e.g.,row.accessionorrow["assembly_accession"]).
You may also want to add a variant using an empty string taxId (e.g.,taxId="") following the same pattern ifcompute_milestonesdistinguishes between different malformed values.
Comment on lines
+312
to
+321
| def test_latest_is_current_version(self, taxonomy): | ||
| # Latest = the current (non-superseded) version, even if an out-of-order | ||
| # superseded row has a newer date it must not win over a current one. | ||
| rows = [ | ||
| make_row("GCA_001.1", ASELLUS_AQUATICUS, "2018-01-01", version_status="superseded"), | ||
| make_row("GCA_001.2", ASELLUS_AQUATICUS, "2021-01-01", version_status="current"), | ||
| ] | ||
| _, species_extra = compute_milestones(rows, taxonomy) | ||
| latest = species_extra[ASELLUS_AQUATICUS]["latest"] | ||
| assert (latest["date"], latest["accession"]) == ("2021-01-01", "GCA_001.2") |
Contributor
There was a problem hiding this comment.
suggestion (testing): Consider adding a test for the interplay between undated current assemblies and dated superseded ones in latest-selection logic.
Current tests cover dated current vs dated superseded rows and superseded-only species, but not the case where a current row has an empty releaseDate and a superseded row has a non-empty date. Please add a test that creates this scenario (current with empty date, superseded with a later date) and asserts that the current assembly is still selected as latest, to guard the (is_current, date, accession) tie-breaking logic against regressions.
| return None | ||
|
|
||
|
|
||
| def compute_milestones(rows: list[dict], taxonomy: dict) -> tuple[dict, dict]: |
Contributor
There was a problem hiding this comment.
issue (complexity): Consider refactoring latest-assembly tracking and the milestone sweep into small helpers to simplify compute_milestones and build_output_rows without changing behavior.
You can reduce complexity in compute_milestones and build_output_rows with a couple of small, local refactors that keep behavior identical:
1. Simplify latest-assembly tracking
The nested latest dict with a key tuple adds indirection. You can track the candidate tuple directly and derive date/accession only when needed:
def _better_latest(current: tuple | None, candidate: tuple) -> bool:
"""Return True if candidate is a better 'latest' than current."""
if current is None:
return True
return candidate > current # (is_current, date, accession)Use it in compute_milestones:
# inside species_extra.setdefault(...)
"latest": None, # (is_current, date, accession)
# ...
candidate = (is_current, item["date"], item["accession"])
if _better_latest(extra["latest"], candidate):
extra["latest"] = candidateAnd in build_output_rows:
if is_species:
extra = species_extra.get(taxid, {})
latest = extra.get("latest")
if latest:
_, latest_date, latest_accession = latest
row["latest_assembly_date"] = latest_date
row["latest_assembly_accession"] = latest_accession
else:
row["latest_assembly_date"] = ""
row["latest_assembly_accession"] = ""This removes the latest["key"] indirection and makes the comparison logic more obvious without changing semantics.
2. Extract the milestone sweep into a helper
The nested loops over MILESTONES and LEVELS make the core sweep dense. A small helper reduces nesting and centralizes the “apply milestone for a row” logic:
def _apply_milestones_for_row(
item: dict,
taxonomy: dict,
taxa: dict,
species_extra: dict,
seen: dict[str, set[int]],
) -> None:
sp = item["species_taxid"]
lineage = item["lineage"]
def ensure_taxon(taxid: int):
if taxid not in taxa:
node = taxonomy.get(taxid, {})
taxa[taxid] = {
"rank": node.get("rank", ""),
"scientific_name": node.get("scientific_name", ""),
"firsts": {},
}
for milestone in MILESTONES:
key = milestone["key"]
if not milestone_predicate(item["row"], key):
continue
for level in LEVELS:
taxid = sp if level == "species" else lineage.get(level)
if not taxid or taxid in seen[key]:
continue
seen[key].add(taxid)
ensure_taxon(taxid)
taxa[taxid]["firsts"][key] = (item["date"], item["accession"])
if milestone["in_ranks"] and level != "species":
species_extra[sp]["in_ranks"][key].append(level)Then compute_milestones becomes:
taxa: dict[int, dict] = {}
seen = {m["key"]: set() for m in MILESTONES}
for item in dated:
_apply_milestones_for_row(item, taxonomy, taxa, species_extra, seen)
for sp in species_extra:
# ensure species rows exist
if sp not in taxa:
node = taxonomy.get(sp, {})
taxa[sp] = {
"rank": node.get("rank", ""),
"scientific_name": node.get("scientific_name", ""),
"firsts": {},
}This keeps the chronological sweep and all milestone behavior intact, but flattens the nested control flow and makes the milestone logic easier to follow and test in isolation.
fchen13
marked this pull request as draft
July 16, 2026 15:17
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Phase 3 of the assembly version tracking system. Computes per-taxon assembly milestone firsts from the version-tracked TSVs produced by Phases 0–2, writing
taxon_milestone_summary.tsv.GoaT stores only the current assembly version and has no concept of "first" — a genome that first reached a milestone but was later superseded leaves no trace in it. Phase 3 supplies that missing time axis: when a taxon (at species and every higher rank) first reached each milestone, and which assembly/species got there first.
What's included
Two modules + tests + docs (4 files):
flows/lib/load_taxonomy.py— dev/test NCBI taxdump parser building the taxonomy contract (taxid -> scientific_name, rank, parent, canonical-rank lineage). Pure parse/build, no download or cache. Cycle-guarded parent-chain walks. In production the lineage is attached upstream; this module is dev/test only.flows/lib/compute_taxon_milestones.py— a single chronological sweep producing species and higher-rank milestones together; first-in-ranks recorded as a by-product of the higher-rank step; per-specieslatest_assembly/current_assemblies/total_assemblies. Runnable withSKIP_PREFECT=true; emitscompute.taxon_milestones.completed.tests/test_taxon_milestones.py— 38 tests (parser,resolve_to_species, the sweep, output schema, validation invariants).tests/README_phase_3_taxon_milestones.md— run/test guide and output schema.The four milestones
first_assemblyfirst_ebp_assemblyPRJNA533106inbioProjectAccessionfirst_metricebpStandardDatenon-emptyfirst_ebp_metricfirst_assembly_in_ranks/first_metric_in_ranksrecord the canonical ranks at which a species was the earliest in its clade to reach the milestone.Wiring
Phase 3 triggers off
update.assembly_versions.completed(parallel with Phase 2), same as documented in the README. The Prefect deployment wiring is intentionally not included here — left to the workflow owner.Testing
```
SKIP_PREFECT=true python3 -m pytest tests/test_taxon_milestones.py -v
```
38 passing. Validation invariants (nested milestone-date ordering;
in_ranks ⊆ CANONICAL_RANKS) asserted in the suite.Summary by Sourcery
Introduce Phase 3 taxon milestone computation that derives per-taxon assembly milestone firsts from version-tracked assembly TSVs and writes a consolidated milestone summary TSV.
New Features:
Enhancements:
Documentation:
Tests: