gbanyan/usher-exploring
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(03-01): add annotation DuckDB loader, CLI command, and tests

Browse Source 
- Create load_to_duckdb with provenance tracking and tier distribution stats
- Add query_poorly_annotated helper to find under-studied genes
- Register `evidence annotation` CLI command with checkpoint-restart pattern
- Add comprehensive unit tests (9 tests) covering GO extraction, NULL handling, tier classification, score normalization, weighting
- Add integration tests (6 tests) for pipeline, idempotency, checkpoint-restart, provenance, queries
- All 15 tests pass with proper NULL preservation and schema validation
This commit is contained in:
gbanyan
2026-02-11 19:03:10 +08:00
parent 0e389c7e41
commit d70239c4ce
5 changed files with 1625 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

227
tests/test_annotation.py Normal file
View File

@@ -0,0 +1,227 @@
"""Unit tests for annotation evidence layer."""
import polars as pl
import pytest
from unittest.mock import Mock, patch
from usher_pipeline.evidence.annotation import (
classify_annotation_tier,
normalize_annotation_score,
)
def test_go_count_extraction():
"""Test correct GO term counting by category."""
# Create synthetic data with different GO counts per category
df = pl.DataFrame({
"gene_id": ["ENSG001", "ENSG002", "ENSG003"],
"gene_symbol": ["GENE1", "GENE2", "GENE3"],
"go_term_count": [50, 15, 3],
"go_biological_process_count": [30, 10, 2],
"go_molecular_function_count": [15, 3, 1],
"go_cellular_component_count": [5, 2, 0],
"uniprot_annotation_score": [5, 4, 2],
"has_pathway_membership": [True, True, False],
})
# Verify counts sum correctly (BP + MF + CC should equal total)
for row in df.iter_rows(named=True):
expected_total = (
row["go_biological_process_count"]
+ row["go_molecular_function_count"]
+ row["go_cellular_component_count"]
)
assert row["go_term_count"] == expected_total
def test_null_go_handling():
"""Test that genes with no GO data get NULL counts."""
# Create data with NULL GO counts
df = pl.DataFrame({
"gene_id": ["ENSG001", "ENSG002"],
"gene_symbol": ["GENE1", "GENE2"],
"go_term_count": [20, None],
"go_biological_process_count": [15, None],
"go_molecular_function_count": [3, None],
"go_cellular_component_count": [2, None],
"uniprot_annotation_score": [4, 3],
"has_pathway_membership": [True, False],
})
# Verify NULL is preserved (not converted to 0)
assert df["go_term_count"][1] is None
assert df["go_biological_process_count"][1] is None
def test_tier_classification_well_annotated():
"""Test well_annotated tier: high GO + high UniProt."""
df = pl.DataFrame({
"gene_id": ["ENSG001", "ENSG002", "ENSG003"],
"gene_symbol": ["GENE1", "GENE2", "GENE3"],
"go_term_count": [25, 20, 22],
"go_biological_process_count": [15, 12, 13],
"go_molecular_function_count": [7, 6, 7],
"go_cellular_component_count": [3, 2, 2],
"uniprot_annotation_score": [5, 4, 4],
"has_pathway_membership": [True, True, False],
})
result = classify_annotation_tier(df)
# All should be well_annotated (GO >= 20 AND UniProt >= 4)
assert all(result["annotation_tier"] == "well_annotated")
def test_tier_classification_poorly_annotated():
"""Test poorly_annotated tier: low/NULL GO + low UniProt."""
df = pl.DataFrame({
"gene_id": ["ENSG001", "ENSG002", "ENSG003"],
"gene_symbol": ["GENE1", "GENE2", "GENE3"],
"go_term_count": [2, None, 0],
"go_biological_process_count": [1, None, 0],
"go_molecular_function_count": [1, None, 0],
"go_cellular_component_count": [0, None, 0],
"uniprot_annotation_score": [2, None, 1],
"has_pathway_membership": [False, None, False],
})
result = classify_annotation_tier(df)
# All should be poorly_annotated
assert all(result["annotation_tier"] == "poorly_annotated")
def test_tier_classification_partial():
"""Test partially_annotated tier: medium annotations."""
df = pl.DataFrame({
"gene_id": ["ENSG001", "ENSG002", "ENSG003"],
"gene_symbol": ["GENE1", "GENE2", "GENE3"],
"go_term_count": [10, 3, 15],
"go_biological_process_count": [7, 2, 10],
"go_molecular_function_count": [2, 1, 4],
"go_cellular_component_count": [1, 0, 1],
"uniprot_annotation_score": [3, 3, 2],
"has_pathway_membership": [True, False, True],
})
result = classify_annotation_tier(df)
# All should be partially_annotated (GO >= 5 OR UniProt >= 3)
assert all(result["annotation_tier"] == "partially_annotated")
def test_normalization_bounds():
"""Test that normalized scores are always in [0, 1] range."""
df = pl.DataFrame({
"gene_id": ["ENSG001", "ENSG002", "ENSG003", "ENSG004"],
"gene_symbol": ["GENE1", "GENE2", "GENE3", "GENE4"],
"go_term_count": [100, 50, 10, 1],
"go_biological_process_count": [60, 30, 7, 1],
"go_molecular_function_count": [30, 15, 2, 0],
"go_cellular_component_count": [10, 5, 1, 0],
"uniprot_annotation_score": [5, 4, 3, 1],
"has_pathway_membership": [True, True, False, False],
})
result = normalize_annotation_score(df)
# All non-NULL scores should be in [0, 1]
scores = result.filter(pl.col("annotation_score_normalized").is_not_null())["annotation_score_normalized"]
assert all(scores >= 0.0)
assert all(scores <= 1.0)
def test_normalization_null_preservation():
"""Test that all-NULL inputs produce NULL score."""
df = pl.DataFrame({
"gene_id": ["ENSG001"],
"gene_symbol": ["GENE1"],
"go_term_count": pl.Series([None], dtype=pl.Int64),
"go_biological_process_count": pl.Series([None], dtype=pl.Int64),
"go_molecular_function_count": pl.Series([None], dtype=pl.Int64),
"go_cellular_component_count": pl.Series([None], dtype=pl.Int64),
"uniprot_annotation_score": pl.Series([None], dtype=pl.Int64),
"has_pathway_membership": pl.Series([None], dtype=pl.Boolean),
})
result = normalize_annotation_score(df)
# Should get NULL score (not 0.0)
assert result["annotation_score_normalized"][0] is None
def test_normalization_with_pathway():
"""Test that pathway membership contributes to score."""
# Two genes with identical GO/UniProt, different pathway membership
df = pl.DataFrame({
"gene_id": ["ENSG001", "ENSG002"],
"gene_symbol": ["GENE1", "GENE2"],
"go_term_count": [10, 10],
"go_biological_process_count": [7, 7],
"go_molecular_function_count": [2, 2],
"go_cellular_component_count": [1, 1],
"uniprot_annotation_score": [3, 3],
"has_pathway_membership": [True, False],
})
result = normalize_annotation_score(df)
# Gene with pathway should have higher score
assert result["annotation_score_normalized"][0] > result["annotation_score_normalized"][1]
def test_composite_weighting():
"""Test that composite score follows 0.5/0.3/0.2 weight distribution."""
# Create gene with only GO data (should contribute 50% weight)
df_go_only = pl.DataFrame({
"gene_id": ["ENSG001"],
"gene_symbol": ["GENE1"],
"go_term_count": [100], # Max GO to get full GO component
"go_biological_process_count": [60],
"go_molecular_function_count": [30],
"go_cellular_component_count": [10],
"uniprot_annotation_score": pl.Series([None], dtype=pl.Int64),
"has_pathway_membership": pl.Series([None], dtype=pl.Boolean),
})
# Create gene with only UniProt data (should contribute 30% weight)
df_uniprot_only = pl.DataFrame({
"gene_id": ["ENSG002"],
"gene_symbol": ["GENE2"],
"go_term_count": pl.Series([None], dtype=pl.Int64),
"go_biological_process_count": pl.Series([None], dtype=pl.Int64),
"go_molecular_function_count": pl.Series([None], dtype=pl.Int64),
"go_cellular_component_count": pl.Series([None], dtype=pl.Int64),
"uniprot_annotation_score": [5], # Max UniProt score
"has_pathway_membership": pl.Series([None], dtype=pl.Boolean),
})
# Create gene with only pathway data (should contribute 20% weight)
df_pathway_only = pl.DataFrame({
"gene_id": ["ENSG003"],
"gene_symbol": ["GENE3"],
"go_term_count": pl.Series([None], dtype=pl.Int64),
"go_biological_process_count": pl.Series([None], dtype=pl.Int64),
"go_molecular_function_count": pl.Series([None], dtype=pl.Int64),
"go_cellular_component_count": pl.Series([None], dtype=pl.Int64),
"uniprot_annotation_score": pl.Series([None], dtype=pl.Int64),
"has_pathway_membership": [True],
})
# Normalize each separately (need same GO max, so combine first)
df_combined = pl.concat([df_go_only, df_uniprot_only, df_pathway_only])
result = normalize_annotation_score(df_combined)
# Check approximate weights (allowing for small rounding)
go_score = result["annotation_score_normalized"][0]
uniprot_score = result["annotation_score_normalized"][1]
pathway_score = result["annotation_score_normalized"][2]
# GO component should be ~0.5 (full weight)
assert abs(go_score - 0.5) < 0.01
# UniProt component should be 0.3 (full score * weight)
assert abs(uniprot_score - 0.3) < 0.01
# Pathway component should be 0.2 (full weight)
assert abs(pathway_score - 0.2) < 0.01