Type glycans and non-canonical residues, and drop the deprecated AtomWorks call

TokenType carried only five members while the RFD4 repository has seven, and nothing ever assigned GLYCAN or NCAA. A glycan or non-canonical residue was therefore typed as protein, unknown or atomized, which hands the decoder the wrong per-token-type residue restriction - on a release whose point is that it handles them. Ports the two assignments from datapipes/bonds.py along with the 13 glycan and 100 NCAA residue names.

Replaces get_annotation_2d_categories() with get_annotation_categories(n_body=2) at the two call sites that still used it, matching main and silencing the DeprecationWarning AtomWorks 3.0.3 emits on every run.

check_vendored.py now compares vendored content against the source by default, rather than only comparing the recorded commit against the current one. The pin check cannot catch a header that was updated without the content being reconciled, which is how the above slipped through.

Output-neutral on the nine example inputs: all 720 tensors bit-identical. Glycan and NCAA typing is newly reachable but is not exercised by those examples.

src/rfd4_proteina_prep/_cif/datapipes_bonds.py

@@ -27,6 +27,33 @@ from .datapipes_schema_proteina import TokenType
 
 logger = logging.getLogger(__name__)
 
+# Residue names the unified vocabulary tokenizes as glycans and non-canonical amino acids,
+# from the RFD4 repository's RESIDUES_BY_TOKEN_TYPE. A residue missing from these is typed as
+# protein/RNA/DNA or atomized instead, which gives the decoder the wrong residue restriction.
+GLYCAN_NAMES = frozenset(
+    (
+        "A2G", "BGC", "BMA", "FUC", "GAL", "GLA", "GLC", "MAN",
+        "NAG", "NDG", "NGA", "SIA", "XYP",
+    )
+)  # fmt: skip
+NCAA_NAMES = frozenset(
+    (
+        "004", "0AF", "0TD", "2MR", "3FG", "4BF", "4D4", "4HH",
+        "5CT", "5OH", "ABA", "AGM", "AIB", "ALC", "ALN", "ALO",
+        "ALY", "API", "ASA", "CAF", "CAS", "CCS", "CGU", "CIR",
+        "CME", "CSD", "CSO", "CSS", "CSU", "CSX", "CYG", "D2T",
+        "DAB", "DAH", "DDE", "DDZ", "DPP", "EEP", "FTR", "GLZ",
+        "HIC", "HIP", "HRG", "HY3", "HYP", "IIL", "ILX", "IYR",
+        "KCR", "KCX", "KPI", "LLP", "LYR", "M3L", "MEN", "MEQ",
+        "MGN", "MHO", "MHS", "MIS", "MK8", "ML3", "MLY", "MLZ",
+        "MP8", "NEP", "NIY", "NLE", "NMM", "OAS", "OCS", "OMT",
+        "OMY", "ORN", "P1L", "PBF", "PHD", "PHI", "PPN", "PTR",
+        "SCH", "SCY", "SEB", "SEC", "SEP", "SLL", "SMC", "SME",
+        "SNC", "TBG", "TOX", "TPO", "TPQ", "TRQ", "TRX", "TYI",
+        "TYQ", "TYS", "V5N", "YCM",
+    )
+)  # fmt: skip
+
 _UNKNOWN_RES_BY_TOKEN_TYPE = {TokenType.PROTEIN: UNKNOWN_AA, TokenType.RNA: UNKNOWN_RNA, TokenType.DNA: UNKNOWN_DNA}
 
 _AROMATIC_BOND_TYPES = (
@@ -54,6 +81,8 @@ def _token_type_at(atom_array: AtomArray, atom_idxs: np.ndarray) -> np.ndarray:
     types[atom_array.get_annotation("is_standard_unknown_or_masked_dna")[atom_idxs]] = TokenType.DNA
     types[atom_array.get_annotation("is_standard_unknown_or_masked_rna")[atom_idxs]] = TokenType.RNA
     types[atom_array.get_annotation("is_standard_unknown_or_masked_aa")[atom_idxs]] = TokenType.PROTEIN
+    types[np.isin(atom_array.res_name[atom_idxs], list(NCAA_NAMES))] = TokenType.NCAA
+    types[np.isin(atom_array.res_name[atom_idxs], list(GLYCAN_NAMES))] = TokenType.GLYCAN
     types[atom_array.atomize[atom_idxs]] = TokenType.ATOMIZED
     return types
 

src/rfd4_proteina_prep/_cif/datapipes_pipeline.py

@@ -510,7 +510,7 @@ class ConvertToAtom37Format(Transform):
         na_ss_pair = np.zeros((n_tokens, n_tokens), dtype=np.int16)
         na_ss_pair_mask = np.zeros((n_tokens, n_tokens), dtype=bool)
 
-        if C_BP.full_name in atom_array.get_annotation_2d_categories():
+        if C_BP.full_name in atom_array.get_annotation_categories(n_body=2):
             pair_annotation = C_BP.annotation(atom_array, default="raise")
             atom_to_token = np.full(atom_array.array_length(), -1, dtype=np.int32)
             atom_to_token[token_starts] = np.arange(n_tokens, dtype=np.int32)

src/rfd4_proteina_prep/_cif/datapipes_schema_proteina.py

@@ -42,6 +42,8 @@ class TokenType:
     DNA = 2
     ATOMIZED = 3
     UNKNOWN = 4
+    GLYCAN = 5
+    NCAA = 6
 
 
 SPECS: dict[str, tuple] = {

src/rfd4_proteina_prep/_cif/transforms_na_geom.py

@@ -58,7 +58,7 @@ class FeaturizeNucleicAcidSecondaryStructure(Transform):
         token_starts = get_token_starts(atom_array)
         n_tokens = len(token_starts)
         if (
-            C_BP.full_name not in atom_array.get_annotation_2d_categories()
+            C_BP.full_name not in atom_array.get_annotation_categories(n_body=2)
             or C_NA_SS_STATE.full_name not in atom_array.get_annotation_categories()
         ):
             pair_annotation = AnnotationList2D(

tools/check_vendored.py

@@ -1,13 +1,23 @@
 #!/usr/bin/env python3
 """Report drift between the vendored `_cif` modules and the RFD4 repository they came from.
 
-Every vendored file records the commit it was adapted from. This re-reads those sources at the
-recorded commit, and again at the checkout's current commit, and reports which ones moved. It
-does not rewrite anything: the vendored copies deliberately omit training-only code, so the
-diffs need a human. The point is that drift shows up as a diff rather than as a silent
-behavioural difference months later.
+Two checks, because they catch different things:
+
+* `--mode pin` re-reads each source at the commit the vendored file records and at the
+  checkout's current commit, and reports which sources moved upstream. This answers "has main
+  changed since someone last looked", and it is only as truthful as the recorded commit.
+* `--mode content` (the default) compares the vendored file against the source itself, and
+  reports lines that exist in both but differ. This answers "does the copy actually match",
+  and unlike the pin check it cannot be defeated by a header that was updated without the
+  content being reconciled. That mistake has already been made here once.
+
+Neither rewrites anything. The vendored copies deliberately omit training-only code, so removals
+are expected and are not reported; only lines that correspond and disagree are. Some of those are
+legitimate inference-only adaptations, so the output is a triage list, not a defect list.
 
     python tools/check_vendored.py --rfd4 /path/to/RFD4-Proteina-dev
+    python tools/check_vendored.py --rfd4 ../RFD4-Proteina-dev --mode pin
+    python tools/check_vendored.py --rfd4 ../RFD4-Proteina-dev --only transforms_na_geom.py --show-diff
 """
 
 from __future__ import annotations
@@ -43,18 +53,76 @@ SOURCES = {
 }
 
 
-def show(repo: Path, commit: str, path: str) -> list[str] | None:
+def show_file(repo: Path, commit: str, path: str) -> list[str] | None:
     result = subprocess.run(
         ["git", "-C", str(repo), "show", f"{commit}:{path}"], capture_output=True, text=True
     )
     return result.stdout.splitlines() if result.returncode == 0 else None
 
 
+def content_drift(vendored: Path, repo: Path, commit: str, only: str | None, show: bool) -> int:
+    """Compare each vendored copy against its source, reporting lines that correspond but differ.
+
+    Removals are skipped: the vendored copies drop training-only code on purpose, so a missing
+    block is expected. What matters is a line that exists on both sides and says something
+    different — a renamed call, a changed constant, a deprecated API the copy still uses.
+    """
+    noise = re.compile(r"^\s*(#|from\s|import\s|\"\"\"|\'\'\')")
+    total, reported = 0, 0
+    print(f"comparing vendored copies against {commit[:12]}\n")
+    for name, source in sorted(SOURCES.items()):
+        if only and name != only:
+            continue
+        local = vendored / name
+        if not local.is_file():
+            continue
+        theirs = show_file(repo, commit, source)
+        if theirs is None:
+            print(f"  {name:<40} source unreadable at {commit[:12]}  <- {source}")
+            continue
+        ours = local.read_text().splitlines()
+        differing = []
+        matcher = difflib.SequenceMatcher(None, ours, theirs, autojunk=False)
+        for tag, i1, i2, j1, j2 in matcher.get_opcodes():
+            if tag != "replace":
+                continue
+            for mine, theirs_line in zip(ours[i1:i2], theirs[j1:j2]):
+                left, right = mine.strip(), theirs_line.strip()
+                if not left or not right or left == right:
+                    continue
+                if noise.match(mine) and noise.match(theirs_line):
+                    continue
+                differing.append((left, right))
+        if differing:
+            reported += 1
+            total += len(differing)
+            print(f"  {name:<40} {len(differing):>4} differing line(s)  <- {source}")
+            if show:
+                for left, right in differing:
+                    print(f"      copy: {left}")
+                    print(f"      main: {right}")
+    if not total:
+        print("\nno differing lines: every vendored copy matches its source.")
+        return 0
+    print(f"\n{total} differing line(s) across {reported} file(s).")
+    print("Some are deliberate inference-only adaptations. Re-run with --show-diff to triage,")
+    print("and only update a file's recorded commit once its content has actually been reconciled.")
+    return 1
+
+
 def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
     parser.add_argument("--rfd4", type=Path, required=True, help="Checkout of baker-laboratory/RFD4-Proteina-dev.")
     parser.add_argument("--against", default="HEAD", help="Commit to compare against (default: HEAD).")
-    parser.add_argument("--show-diff", action="store_true", help="Print the diff, not just the line count.")
+    parser.add_argument("--show-diff", action="store_true", help="Print the differing lines, not just a count.")
+    parser.add_argument(
+        "--mode",
+        choices=("content", "pin"),
+        default="content",
+        help="content: compare the vendored copy against the source. pin: compare the source at the recorded "
+        "commit against the current one. Default: content.",
+    )
+    parser.add_argument("--only", help="Restrict to one vendored file name.")
     args = parser.parse_args()
 
     vendored = Path(__file__).resolve().parent.parent / "src" / "rfd4_proteina_prep" / "_cif"
@@ -65,6 +133,9 @@ def main() -> int:
         print(f"error: {args.rfd4} is not a readable git checkout", file=sys.stderr)
         return 2
 
+    if args.mode == "content":
+        return content_drift(vendored, args.rfd4, current, args.only, args.show_diff)
+
     pins, moved, unpinned = set(), [], []
     for name, source in sorted(SOURCES.items()):
         local = vendored / name
@@ -75,7 +146,7 @@ def main() -> int:
             unpinned.append(name)
             continue
         pins.add(match.group(1))
-        was, now = show(args.rfd4, match.group(1), source), show(args.rfd4, current, source)
+        was, now = show_file(args.rfd4, match.group(1), source), show_file(args.rfd4, current, source)
         if was is None or now is None:
             moved.append((name, source, -1))
             continue