refactor: classify isotype before EU-numbered Fc allotype detection

src/abcat/cdomain.py

@@ -10,12 +10,23 @@
 
 logger = logging.getLogger(__name__)
 
-# Data directory path
 DATA_DIR = Path(__file__).parent.parent.parent / "data"
 
 _C_GENE_DB: dict[str, Any] | None = None
 _ALLOTYPE_DB: dict[str, Any] | None = None
 
+# EU numbering origins for the constant-domain reference sequences used by the
+# allotype database. These are explicit so allotype detection never depends on
+# the legacy subclass-relative offsets stored in historical marker records.
+_EU_REFERENCE_STARTS: dict[tuple[str, str], int] = {
+    ("heavy", "IgG1"): 118,
+    ("heavy", "IgG2"): 118,
+    ("heavy", "IgG3"): 118,
+    ("heavy", "IgG4"): 118,
+    ("heavy", "IgE"): 99,
+    ("light", "IGKC"): 109,
+}
+
 
 def _load_databases() -> tuple[dict[str, Any], dict[str, Any]]:
     global _C_GENE_DB, _ALLOTYPE_DB
@@ -38,63 +49,146 @@ def _load_databases() -> tuple[dict[str, Any], dict[str, Any]]:
     return _C_GENE_DB, _ALLOTYPE_DB
 
 
+def _make_aligner() -> PairwiseAligner:
+    aligner = PairwiseAligner()
+    aligner.mode = "global"
+    aligner.match_score = 2
+    aligner.mismatch_score = -1
+    aligner.open_gap_score = -5
+    aligner.extend_gap_score = -1
+    return aligner
+
+
+def _isotype_candidates(
+    chain_category: str,
+    c_gene_db: dict[str, Any],
+) -> dict[str, dict[str, tuple[str, dict[str, Any]]]]:
+    """Group constant-gene references by isotype before subclass matching."""
+    grouped: dict[str, dict[str, tuple[str, dict[str, Any]]]] = {}
+    for gene_name, gene_info in c_gene_db.get(chain_category, {}).items():
+        isotype = gene_info.get("isotype", "Unknown")
+        grouped.setdefault(isotype, {})[gene_name] = (gene_name, gene_info)
+    return grouped
+
+
+def _classify_isotype(
+    c_seq: str,
+    chain_category: str,
+    c_gene_db: dict[str, Any],
+    aligner: PairwiseAligner,
+) -> tuple[str, float]:
+    """Determine isotype first, using the N-terminal constant-region signature."""
+    grouped = _isotype_candidates(chain_category, c_gene_db)
+    signature_len = min(70, len(c_seq))
+    signature = c_seq[:signature_len]
+
+    best_isotype = "Unknown"
+    best_score = float("-inf")
+    for isotype, genes in grouped.items():
+        iso_score = float("-inf")
+        for _, gene_info in genes.values():
+            ref_signature = gene_info["sequence"][:signature_len]
+            alignments = aligner.align(signature, ref_signature)
+            if alignments:
+                iso_score = max(iso_score, alignments[0].score)
+        if iso_score > best_score:
+            best_score = iso_score
+            best_isotype = isotype
+
+    return best_isotype, best_score
+
+
+def _select_subclass(
+    c_seq: str,
+    chain_category: str,
+    isotype: str,
+    c_gene_db: dict[str, Any],
+    aligner: PairwiseAligner,
+) -> tuple[str | None, dict[str, Any] | None, float, Any | None]:
+    """Match the constant sequence only against genes belonging to the chosen isotype."""
+    best_gene: str | None = None
+    best_info: dict[str, Any] | None = None
+    best_score = float("-inf")
+    best_alignment: Any | None = None
+
+    for gene_name, gene_info in c_gene_db.get(chain_category, {}).items():
+        if gene_info.get("isotype") != isotype:
+            continue
+        alignments = aligner.align(c_seq, gene_info["sequence"])
+        if alignments and alignments[0].score > best_score:
+            best_score = alignments[0].score
+            best_gene = gene_name
+            best_info = gene_info
+            best_alignment = alignments[0]
+
+    return best_gene, best_info, best_score, best_alignment
+
+
+def _get_eu_reference_start(chain_category: str, matched_subclass: str) -> int | None:
+    """Return the explicit EU-numbering origin for the selected reference."""
+    return _EU_REFERENCE_STARTS.get((chain_category, matched_subclass))
+
+
+def _get_query_index_for_ref_offset(
+    alignment: Any, ref_idx: int, default_len: int
+) -> int | None:
+    """Map a 0-indexed reference position to query position using aligned blocks."""
+    if hasattr(alignment, "aligned") and len(alignment.aligned) >= 2:
+        query_blocks, ref_blocks = alignment.aligned[0], alignment.aligned[1]
+        for (q_start, q_end), (r_start, r_end) in zip(query_blocks, ref_blocks):
+            if r_start <= ref_idx < r_end:
+                return q_start + (ref_idx - r_start)
+        return None
+    if 0 <= ref_idx < default_len:
+        return ref_idx
+    return None
+
+
+def _get_query_index_for_eu_position(
+    alignment: Any,
+    eu_position: int,
+    reference_eu_start: int | None,
+) -> int | None:
+    """Map a canonical EU position through reference and query pairwise alignment."""
+    if reference_eu_start is None:
+        return None
+    ref_idx = eu_position - reference_eu_start
+    if ref_idx < 0:
+        return None
+    return _get_query_index_for_ref_offset(alignment, ref_idx, 0)
+
+
 def analyze_cdomain(
     sequence: str,
     v_domain_len: int = 0,
     chain_hint: ChainType | None = None,
 ) -> CAnalysisResult:
-    """
-    Analyzes the Constant Domain of an antibody sequence.
-    Determines Isotype, Subclass, Isoallotype, and Allotype markers.
-    """
+    """Analyze the constant domain in order: isotype -> subclass -> EU allotype."""
     cleaned = clean_sequence(sequence)
-
-    # Extract constant domain sequence if full sequence is provided
-    if v_domain_len > 0 and len(cleaned) > v_domain_len:
-        c_seq = cleaned[v_domain_len:]
-    else:
-        c_seq = cleaned
+    c_seq = cleaned[v_domain_len:] if v_domain_len > 0 and len(cleaned) > v_domain_len else cleaned
 
     c_gene_db, allotype_db = _load_databases()
-    aligner = PairwiseAligner()
-    aligner.mode = "global"
-    aligner.match_score = 2
-    aligner.mismatch_score = -1
-    aligner.open_gap_score = -5
-    aligner.extend_gap_score = -1
-
-    best_match_key = None
-    best_c_info = None
-    best_score = -1e9
-    best_alignment = None
-    chain_category = "heavy"
+    aligner = _make_aligner()
 
-    # Filter targets based on chain hint if provided
-    search_groups = []
     if chain_hint == ChainType.HEAVY:
-        search_groups = [("heavy", c_gene_db.get("heavy", {}))]
+        search_groups = ["heavy"]
     elif chain_hint in (ChainType.KAPPA, ChainType.LAMBDA):
-        search_groups = [("light", c_gene_db.get("light", {}))]
+        search_groups = ["light"]
     else:
-        search_groups = [
-            ("heavy", c_gene_db.get("heavy", {})),
-            ("light", c_gene_db.get("light", {})),
-        ]
-
-    for cat_name, genes in search_groups:
-        for gene_name, gene_info in genes.items():
-            ref_seq = gene_info["sequence"]
-            alignments = aligner.align(c_seq, ref_seq)
-            if alignments:
-                top = alignments[0]
-                if top.score > best_score:
-                    best_score = top.score
-                    best_match_key = gene_name
-                    best_c_info = gene_info
-                    best_alignment = top
-                    chain_category = cat_name
-
-    if not best_match_key or not best_c_info or not best_alignment:
+        search_groups = ["heavy", "light"]
+
+    # Stage 1: determine isotype without letting subclass/allotype differences influence the class call.
+    best_chain: str | None = None
+    best_isotype = "Unknown"
+    best_iso_score = float("-inf")
+    for category in search_groups:
+        isotype, score = _classify_isotype(c_seq, category, c_gene_db, aligner)
+        if score > best_iso_score:
+            best_iso_score = score
+            best_isotype = isotype
+            best_chain = category
+
+    if best_chain is None or best_isotype == "Unknown":
         return CAnalysisResult(
             c_region_sequence=c_seq,
             isotype="Unknown",
@@ -105,27 +199,41 @@ def analyze_cdomain(
             isoallotypes=[],
         )
 
+    # Stage 2: subclass is matched only inside the selected isotype.
+    best_match_key, best_c_info, best_score, best_alignment = _select_subclass(
+        c_seq, best_chain, best_isotype, c_gene_db, aligner
+    )
+    if not best_match_key or not best_c_info or best_alignment is None:
+        return CAnalysisResult(
+            c_region_sequence=c_seq,
+            isotype=best_isotype,
+            subclass="Unknown",
+            matched_c_gene="None",
+            alignment_identity=0.0,
+            allotypes=[],
+            isoallotypes=[],
+        )
+
     ref_len = len(best_c_info["sequence"])
     identity = round(min(1.0, max(0.0, best_score / (2.0 * ref_len))), 4)
-
-    isotype = best_c_info["isotype"]
     subclass = best_c_info["subclass"]
 
-    # Call Allotypes and Isoallotypes
+    # Stage 3: determine allotypes from canonical EU-numbered constant/Fc positions.
     allotype_calls, isoallotype_calls = _call_allotypes(
         query_c_seq=c_seq,
         matched_subclass=subclass,
-        chain_category=chain_category,
+        chain_category=best_chain,
         allotype_db=allotype_db,
         alignment=best_alignment,
+        reference_c_seq=best_c_info["sequence"],
     )
 
     present_allotypes = [a.allotype for a in allotype_calls if a.status == "present"]
     summary_str = format_allotype_summary(present_allotypes)
 
     return CAnalysisResult(
         c_region_sequence=c_seq,
-        isotype=isotype,
+        isotype=best_isotype,
         subclass=subclass,
         matched_c_gene=best_match_key,
         alignment_identity=identity,
@@ -147,13 +255,9 @@ def format_allotype_summary(present_allotypes: list[str]) -> str:
         match = re.match(r"^([A-Za-z]+[0-9]*m|\w+[\*\_]?)(.*)$", allo)
         if match:
             prefix, suffix = match.groups()
-            if prefix not in groups:
-                groups[prefix] = []
-            groups[prefix].append(suffix)
+            groups.setdefault(prefix, []).append(suffix)
         else:
-            if "other" not in groups:
-                groups["other"] = []
-            groups["other"].append(allo)
+            groups.setdefault("other", []).append(allo)
 
     formatted_parts = []
     for prefix, suffixes in groups.items():
@@ -162,76 +266,63 @@ def format_allotype_summary(present_allotypes: list[str]) -> str:
         elif len(suffixes) == 1:
             formatted_parts.append(f"{prefix}{suffixes[0]}")
         else:
-            joined_suffixes = ",".join(suffixes)
-            formatted_parts.append(f"{prefix}{joined_suffixes}")
+            formatted_parts.append(f"{prefix}{','.join(suffixes)}")
 
     return ", ".join(formatted_parts)
 
 
-def _get_query_index_for_ref_offset(
-    alignment: Any, ref_idx: int, default_len: int
-) -> int | None:
-    """
-    Maps 0-indexed reference C-gene position (ref_idx) to 0-indexed position in query_c_seq using PairwiseAligner alignment.
-    alignment.aligned[0] is query blocks, alignment.aligned[1] is target (ref) blocks.
-    """
-    if hasattr(alignment, "aligned") and len(alignment.aligned) >= 2:
-        query_blocks, ref_blocks = alignment.aligned[0], alignment.aligned[1]
-        for (q_start, q_end), (r_start, r_end) in zip(query_blocks, ref_blocks):
-            if r_start <= ref_idx < r_end:
-                return q_start + (ref_idx - r_start)
-        return None
-    if 0 <= ref_idx < default_len:
-        return ref_idx
-    return None
-
-
 def _call_allotypes(
     query_c_seq: str,
     matched_subclass: str,
     chain_category: str,
     allotype_db: dict[str, Any],
     alignment: Any,
+    reference_c_seq: str,
 ) -> tuple[list[AllotypeMarkerCall], list[AllotypeMarkerCall]]:
-    """Scans key polymorphic position markers to assign Allotypes and Isoallotypes."""
+    """Call allotype/isoallotype markers using canonical EU positions."""
     allotypes: list[AllotypeMarkerCall] = []
     isoallotypes: list[AllotypeMarkerCall] = []
 
     subclass_markers = allotype_db.get(chain_category, {}).get(matched_subclass, [])
     if not subclass_markers:
         return allotypes, isoallotypes
 
-    # Align query C-region to reference
+    reference_eu_start = _get_eu_reference_start(chain_category, matched_subclass)
+    if reference_eu_start is None:
+        logger.warning("No EU reference origin for %s/%s", chain_category, matched_subclass)
+        return allotypes, isoallotypes
+
     for rule in subclass_markers:
         allotype_name = rule["allotype"]
         opposing = rule.get("opposing_allotype")
         domain = rule.get("domain", "C")
         category = rule.get("type", "allotype")
         markers = rule.get("markers", [])
 
-        matches = {}
+        matches: dict[str, str] = {}
         all_matched = True
+        any_observed = False
+
+        for marker in markers:
+            eu_pos = int(marker["eu_position"])
+            expected_aa = marker["amino_acid"]
+            query_idx = _get_query_index_for_eu_position(alignment, eu_pos, reference_eu_start)
+            if query_idx is None or not (0 <= query_idx < len(query_c_seq)):
+                all_matched = False
+                continue
 
-        for m in markers:
-            ref_offset = (
-                m["subclass_offset"] - 1
-            )  # 0-indexed offset in reference C-gene
-            expected_aa = m["amino_acid"]
-            eu_pos = m["eu_position"]
-
-            query_idx = _get_query_index_for_ref_offset(
-                alignment, ref_offset, len(query_c_seq)
-            )
-
-            if query_idx is not None and 0 <= query_idx < len(query_c_seq):
-                actual_aa = query_c_seq[query_idx]
-                matches[f"EU_{eu_pos}"] = actual_aa
-                if actual_aa != expected_aa:
-                    all_matched = False
-            else:
+            actual_aa = query_c_seq[query_idx]
+            matches[f"EU_{eu_pos}"] = actual_aa
+            any_observed = True
+            if actual_aa != expected_aa:
                 all_matched = False
 
-        status = "present" if all_matched else "absent"
+        if all_matched and any_observed:
+            status = "present"
+        elif any_observed:
+            status = "absent"
+        else:
+            status = "inconclusive"
 
         call = AllotypeMarkerCall(
             allotype=allotype_name,