Feat: add genome/cai.py — per-gene Codon Adaptation Index

- Extract shared codon utilities to genome/core.py:
  AA_TABLE, iter_codons(), make_aa_to_codons(), calculate_cai()
- Refactor genome/codon.py to use shared functions from genome/core.py
- New genome/cai.py: computes per-gene CAI against a reference
  of highly expressed genes (unlike codon --cai which gives one global value)
- 6 new tests, all pass

Zero new dependencies — pure biopython + stdlib.

Co-Authored-By: Claude <noreply@anthropic.com>

Author: imjiaoyuan

docs/en/module-genome.md

@@ -1,6 +1,6 @@
 # jsrc genome
 
-Genome-level analysis tools. `jsrc genome` covers genome statistics, feature detection, comparative analysis, evolutionary analysis, and annotation utilities.
+Genome-level analysis tools. `jsrc genome` covers genome statistics, feature detection, codon analysis, comparative analysis, evolutionary analysis, and annotation utilities.
 
 ## cpg
 
@@ -112,6 +112,19 @@ jsrc genome window -fa genome.fa
 jsrc genome window -fa genome.fa --window 50000 --step 10000 --json
 ```
 
+## cai
+
+Codon Adaptation Index measures how well a gene's codon usage matches a reference set of highly expressed genes. CAI ranges from 0 to 1; higher values indicate stronger similarity to the reference (and potentially higher expression).
+
+Unlike `codon --cai` which computes a single global CAI for all input sequences combined, this command computes per-gene CAI values. Each gene in the query FASTA gets an individual score, making it suitable for genome-wide CAI profiling.
+
+The reference should be a FASTA of highly expressed genes (e.g., ribosomal proteins, elongation factors). Both query and reference should be CDS sequences — introns and UTRs are not handled.
+
+```bash
+jsrc genome cai -fa all_genes.fa --reference highly_expressed.fa
+jsrc genome cai -fa all_genes.fa --reference highly_expressed.fa --json
+```
+
 ## codon
 
 Codon usage frequency and RSCU (Relative Synonymous Codon Usage) analysis. Input CDS sequences in FASTA format to count codon occurrences and calculate RSCU.

docs/zh/module-genome.md

@@ -1,6 +1,6 @@
 # jsrc genome
 
-基因组级别的分析功能。`jsrc genome` 涵盖了基因组统计、特征检测、比较分析、进化分析和注释辅助等常用功能。
+基因组级别的分析功能。`jsrc genome` 涵盖了基因组统计、特征检测、密码子分析、比较分析、进化分析和注释辅助等常用功能。
 
 ## cpg
 
@@ -112,6 +112,19 @@ jsrc genome window -fa genome.fa
 jsrc genome window -fa genome.fa --window 50000 --step 10000 --json
 ```
 
+## cai
+
+密码子适应指数（CAI），衡量一个基因的密码子使用模式与参考高表达基因集的匹配程度。CAI 范围 0 到 1，值越高表示与参考集越相似（通常意味着更高的表达潜力）。
+
+与 `codon --cai`（对所有输入序列计算一个全局 CAI 值）不同，这个命令对每个基因单独计算 CAI，适合全基因组范围的 CAI 分析。
+
+参考序列应使用高表达基因的 CDS（如核糖体蛋白、延伸因子等）。输入的查询和参考都应该是 CDS 序列，不处理内含子和 UTR。
+
+```bash
+jsrc genome cai -fa all_genes.fa --reference highly_expressed.fa
+jsrc genome cai -fa all_genes.fa --reference highly_expressed.fa --json
+```
+
 ## codon
 
 密码子使用频率和 RSCU（相对同义密码子使用度）分析。输入 CDS 序列的 FASTA，统计每个密码子的出现次数并计算 RSCU。

src/jsrc/genome/__init__.py

@@ -15,6 +15,7 @@
     "gc-skew": ("jsrc.genome.gc_skew", "Cumulative GC skew for replication origin"),
     "window": ("jsrc.genome.window", "Sliding-window GC and AT skew"),
     "codon": ("jsrc.genome.codon", "Codon usage and RSCU analysis"),
+    "cai": ("jsrc.genome.cai", "Codon Adaptation Index"),
     "distance": ("jsrc.genome.distance", "Calculate pairwise genetic distances"),
     "kaks": ("jsrc.genome.kaks", "Calculate Ka/Ks ratio for two aligned CDS sequences"),
     "density": ("jsrc.genome.density", "Calculate gene/feature density along genome"),

src/jsrc/genome/cai.py

@@ -0,0 +1,78 @@
+import json
+import logging
+from argparse import Namespace
+from collections import Counter
+from typing import Any
+
+from Bio import SeqIO
+
+from jsrc.core import DataFormatError
+from jsrc.genome.core import (
+    AA_TABLE,
+    calculate_cai,
+    iter_codons,
+    make_aa_to_codons,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def cmd(args: Namespace) -> None:
+    aa_to_codons = make_aa_to_codons(AA_TABLE)
+
+    ref_counts: Counter[str] = Counter()
+    ref_records = list(SeqIO.parse(args.reference, "fasta"))
+    if not ref_records:
+        raise DataFormatError("No sequences found in reference FASTA")
+    for rec in ref_records:
+        for codon in iter_codons(str(rec.seq)):
+            if AA_TABLE.get(codon) != "*":
+                ref_counts[codon] += 1
+    if not ref_counts:
+        raise DataFormatError("No valid codons found in reference FASTA")
+    logger.info(
+        "Reference: %d sequences, %d codons",
+        len(ref_records),
+        sum(ref_counts.values()),
+    )
+
+    query_records = list(SeqIO.parse(args.fa, "fasta"))
+    if not query_records:
+        raise DataFormatError("No sequences found in query FASTA")
+
+    results = []
+    for rec in query_records:
+        gene_counts: Counter[str] = Counter()
+        for codon in iter_codons(str(rec.seq)):
+            if AA_TABLE.get(codon) != "*":
+                gene_counts[codon] += 1
+        cai = calculate_cai(gene_counts, ref_counts, aa_to_codons)
+        results.append(
+            {
+                "id": rec.id,
+                "codon_count": sum(gene_counts.values()),
+                "cai": round(cai, 6),
+            }
+        )
+
+    if args.json:
+        print(json.dumps(results, ensure_ascii=False, indent=2))
+        return
+    print("id\tcodon_count\tcai")
+    for r in results:
+        print(f"{r['id']}\t{r['codon_count']}\t{r['cai']:.6f}")
+    logger.info("Computed CAI for %d genes", len(results))
+
+
+def register(subparsers: Any) -> None:
+    p = subparsers.add_parser(
+        "cai", help="Codon Adaptation Index for each gene"
+    )
+    p.add_argument("-fa", required=True, help="Query CDS FASTA file")
+    p.add_argument(
+        "--reference",
+        required=True,
+        help="Reference CDS FASTA file (highly expressed genes)",
+    )
+    p.add_argument("--json", action="store_true", help="Print JSON output")
+    p.set_defaults(func=cmd)

src/jsrc/genome/codon.py

@@ -1,112 +1,19 @@
 import json
 import logging
-import math
 from argparse import Namespace
 from collections import Counter, defaultdict
-from collections.abc import Iterator
 from typing import Any
 
 from Bio import SeqIO
 
-logger = logging.getLogger(__name__)
+from jsrc.genome.core import (
+    AA_TABLE,
+    calculate_cai,
+    iter_codons,
+    make_aa_to_codons,
+)
 
-AA_TABLE = {
-    "TTT": "F",
-    "TTC": "F",
-    "TTA": "L",
-    "TTG": "L",
-    "CTT": "L",
-    "CTC": "L",
-    "CTA": "L",
-    "CTG": "L",
-    "ATT": "I",
-    "ATC": "I",
-    "ATA": "I",
-    "ATG": "M",
-    "GTT": "V",
-    "GTC": "V",
-    "GTA": "V",
-    "GTG": "V",
-    "TCT": "S",
-    "TCC": "S",
-    "TCA": "S",
-    "TCG": "S",
-    "CCT": "P",
-    "CCC": "P",
-    "CCA": "P",
-    "CCG": "P",
-    "ACT": "T",
-    "ACC": "T",
-    "ACA": "T",
-    "ACG": "T",
-    "GCT": "A",
-    "GCC": "A",
-    "GCA": "A",
-    "GCG": "A",
-    "TAT": "Y",
-    "TAC": "Y",
-    "TAA": "*",
-    "TAG": "*",
-    "CAT": "H",
-    "CAC": "H",
-    "CAA": "Q",
-    "CAG": "Q",
-    "AAT": "N",
-    "AAC": "N",
-    "AAA": "K",
-    "AAG": "K",
-    "GAT": "D",
-    "GAC": "D",
-    "GAA": "E",
-    "GAG": "E",
-    "TGT": "C",
-    "TGC": "C",
-    "TGA": "*",
-    "TGG": "W",
-    "CGT": "R",
-    "CGC": "R",
-    "CGA": "R",
-    "CGG": "R",
-    "AGT": "S",
-    "AGC": "S",
-    "AGA": "R",
-    "AGG": "R",
-    "GGT": "G",
-    "GGC": "G",
-    "GGA": "G",
-    "GGG": "G",
-}
-
-
-def _iter_codons(seq: str) -> Iterator[str]:
-    seq = seq.upper().replace("U", "T")
-    for i in range(0, len(seq) - 2, 3):
-        c = seq[i : i + 3]
-        if len(c) == 3 and set(c) <= {"A", "C", "G", "T"}:
-            yield c
-
-
-def _calculate_cai(
-    counts: Counter[str], ref_counts: Counter[str], aa_to_codons: dict[str, list[str]]
-) -> float:
-    w_values = {}
-    for _aa, codons in aa_to_codons.items():
-        max_count = max((ref_counts[c] for c in codons), default=0)
-        if max_count == 0:
-            for c in codons:
-                w_values[c] = 1.0
-        else:
-            for c in codons:
-                w_values[c] = ref_counts[c] / max_count
-
-    log_sum = 0.0
-    total = 0
-    for codon, count in counts.items():
-        if codon in w_values and count > 0:
-            log_sum += count * math.log(w_values[codon]) if w_values[codon] > 0 else 0
-            total += count
-
-    return math.exp(log_sum / total) if total > 0 else 0.0
+logger = logging.getLogger(__name__)
 
 
 def _calculate_enc(counts: Counter[str], aa_to_codons: dict[str, list[str]]) -> float:
@@ -144,14 +51,11 @@ def homozygosity(codons: list[str]) -> float:
 
 def cmd(args: Namespace) -> None:
     counts: Counter[str] = Counter()
-    aa_to_codons: dict[str, list[str]] = defaultdict(list)
-    for codon, aa in AA_TABLE.items():
-        if aa != "*":
-            aa_to_codons[aa].append(codon)
+    aa_to_codons = make_aa_to_codons(AA_TABLE)
 
     total_codons = 0
     for rec in SeqIO.parse(args.fa, "fasta"):
-        for codon in _iter_codons(str(rec.seq)):
+        for codon in iter_codons(str(rec.seq)):
             if AA_TABLE.get(codon) == "*":
                 continue
             counts[codon] += 1
@@ -172,10 +76,10 @@ def cmd(args: Namespace) -> None:
     if args.cai:
         ref_counts: Counter[str] = Counter()
         for rec in SeqIO.parse(args.cai, "fasta"):
-            for codon in _iter_codons(str(rec.seq)):
+            for codon in iter_codons(str(rec.seq)):
                 if AA_TABLE.get(codon) != "*":
                     ref_counts[codon] += 1
-        cai_value = _calculate_cai(counts, ref_counts, aa_to_codons)
+        cai_value = calculate_cai(counts, ref_counts, aa_to_codons)
 
     enc_value = None
     if args.enc:

src/jsrc/genome/core.py

@@ -1,7 +1,75 @@
+import math
+from collections import Counter, defaultdict
+from collections.abc import Iterator
+
+
 def normalize_sequence(seq: str) -> str:
     return seq.upper().replace("U", "T")
 
 
+AA_TABLE: dict[str, str] = {
+    "TTT": "F", "TTC": "F", "TTA": "L", "TTG": "L",
+    "CTT": "L", "CTC": "L", "CTA": "L", "CTG": "L",
+    "ATT": "I", "ATC": "I", "ATA": "I", "ATG": "M",
+    "GTT": "V", "GTC": "V", "GTA": "V", "GTG": "V",
+    "TCT": "S", "TCC": "S", "TCA": "S", "TCG": "S",
+    "CCT": "P", "CCC": "P", "CCA": "P", "CCG": "P",
+    "ACT": "T", "ACC": "T", "ACA": "T", "ACG": "T",
+    "GCT": "A", "GCC": "A", "GCA": "A", "GCG": "A",
+    "TAT": "Y", "TAC": "Y", "TAA": "*", "TAG": "*",
+    "CAT": "H", "CAC": "H", "CAA": "Q", "CAG": "Q",
+    "AAT": "N", "AAC": "N", "AAA": "K", "AAG": "K",
+    "GAT": "D", "GAC": "D", "GAA": "E", "GAG": "E",
+    "TGT": "C", "TGC": "C", "TGA": "*", "TGG": "W",
+    "CGT": "R", "CGC": "R", "CGA": "R", "CGG": "R",
+    "AGT": "S", "AGC": "S", "AGA": "R", "AGG": "R",
+    "GGT": "G", "GGC": "G", "GGA": "G", "GGG": "G",
+}
+
+
+def iter_codons(seq: str) -> Iterator[str]:
+    seq = seq.upper().replace("U", "T")
+    for i in range(0, len(seq) - 2, 3):
+        c = seq[i : i + 3]
+        if len(c) == 3 and set(c) <= {"A", "C", "G", "T"}:
+            yield c
+
+
+def make_aa_to_codons(codon_table: dict[str, str] | None = None) -> dict[str, list[str]]:
+    if codon_table is None:
+        codon_table = AA_TABLE
+    result: dict[str, list[str]] = defaultdict(list)
+    for codon, aa in codon_table.items():
+        if aa != "*":
+            result[aa].append(codon)
+    return dict(result)
+
+
+def calculate_cai(
+    counts: Counter[str],
+    ref_counts: Counter[str],
+    aa_to_codons: dict[str, list[str]],
+) -> float:
+    w_values: dict[str, float] = {}
+    for _aa, codons in aa_to_codons.items():
+        max_count = max((ref_counts[c] for c in codons), default=0)
+        if max_count == 0:
+            for c in codons:
+                w_values[c] = 1.0
+        else:
+            for c in codons:
+                w_values[c] = ref_counts[c] / max_count
+
+    log_sum = 0.0
+    total = 0
+    for codon, count in counts.items():
+        if codon in w_values and count > 0:
+            log_sum += count * math.log(w_values[codon]) if w_values[codon] > 0 else 0
+            total += count
+
+    return math.exp(log_sum / total) if total > 0 else 0.0
+
+
 def gc_content(seq: str) -> float:
     seq = normalize_sequence(seq)
     g = seq.count("G")