Add AxoMEME v1 weights, config, and model card

README.md

@@ -0,0 +1,138 @@
+---
+license: mit
+library_name: axomeme
+tags:
+- biology
+- bioinformatics
+- phylogenetics
+- molecular-evolution
+- positive-selection
+- dn-ds
+---
+
+# AxoMEME
+
+**Ultra-fast neural inference of episodic positive selection in molecular sequences.**
+
+AxoMEME detects site-level episodic diversifying positive selection (dN/dS > 1) across
+arbitrary codon alignments and phylogenetic trees. It replaces numerical Maximum Likelihood
+Estimation with an axial geometric transformer using Continuous 4D Tree Rotary Position
+Embeddings (Tree-RoPE), reaching a 100x-1,000x speedup over HyPhy MEME and PAML CodeML while
+matching or exceeding their empirical statistical power at a strictly controlled false
+positive rate.
+
+- **Repository:** https://github.com/veg/axomeme
+- **Architecture notes:** [ARCHITECTURE.md](https://github.com/veg/axomeme/blob/main/ARCHITECTURE.md)
+- **Developed by:** Sergei L. Kosakovsky Pond and collaborators
+- **License:** MIT
+
+## Model details
+
+| | |
+| :--- | :--- |
+| Parameters | 1,909,404 |
+| Embedding dimension | 384 |
+| Layers | 6 |
+| Attention heads | 12 |
+| Ordinal thresholds (CORAL head) | 16 |
+| Precision | float32 |
+| Checkpoint format | safetensors |
+
+The model consumes a codon alignment plus a phylogeny and emits, per codon site, a likelihood
+ratio test statistic for episodic diversifying selection. Asymptotic p-values follow the
+standard MEME null mixture, `0.5 * delta(0) + 0.5 * chi^2(1)`.
+
+## Usage
+
+Install the package, which pulls these weights automatically on first use:
+
+```bash
+pip install git+https://github.com/veg/axomeme.git
+```
+
+### Command line
+
+```bash
+axomeme predict \
+  --alignment my_genes.fasta \
+  --tree my_genes.nwk \
+  --output results.json \
+  --csv results.csv
+```
+
+If `--tree` is omitted, AxoMEME extracts an embedded tree from a NEXUS or FASTA alignment.
+
+### Python
+
+```python
+from axomeme import PhyloAxialTransformer, load_alignment_and_tree
+
+model = PhyloAxialTransformer.from_pretrained("datamonkey/axomeme").eval()
+codons, aas, dists, mds, invariable, taxa, num_sites = load_alignment_and_tree(
+    "my_genes.fasta", "my_genes.nwk"
+)
+```
+
+See the repository README for a complete inference loop, including site batching and p-value
+computation.
+
+## Inputs and outputs
+
+**Inputs.** An in-frame codon alignment (FASTA or NEXUS) and a phylogenetic tree (Newick or
+NEXUS) with branch lengths. Trees lacking branch lengths are estimated via HyPhy under HKY85
+when HyPhy is on the PATH; otherwise strictly positive lower bounds of 1e-4 are enforced,
+which degrades accuracy.
+
+**Outputs.** Per codon site: an LRT statistic, an asymptotic p-value, and an invariable-site
+flag. Sites that are invariable across the alignment are assigned an LRT of 0.
+
+## Evaluation
+
+Taking HyPhy MEME (p <= 0.10, asymptotic LRT >= 4.605) as ground truth across 84 empirical
+datasets from 9 independent literature studies (43,302 codons, up to 476 taxa):
+
+| Literature study & system | Datasets | Codons | ROC-AUC | PR-AUC | PPV | FPR | Spearman rho | HyPhy MLE | AxoMEME CPU | Speedup |
+| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| Abdul et al. (2018), SMC5/6 complex | 9 | 7,073 | 0.990 | 0.752 | 100.0% | 0.00% | 0.874 | 561.0 s | 3.84 s | 171.3x |
+| Nisson et al. (2025), CCDC137 (HIV Vpr) | 1 | 290 | 0.940 | 0.650 | 75.0% | 0.35% | 0.833 | 80.0 s | 0.21 s | 380.9x |
+| Le Corf et al. (2026), GBP5 GTPase | 2 | 1,223 | 0.958 | 0.584 | 68.2% | 0.49% | 0.782 | 1,174.0 s | 1.68 s | 693.3x |
+| D'Oliviera et al. (2025), TRMT1 cleavage | 2 | 1,613 | 0.945 | 0.512 | 62.5% | 0.31% | 0.671 | 723.0 s | 1.30 s | 543.9x |
+| Lytras et al. (2023), bat OAS1 factor | 1 | 351 | 0.904 | 0.628 | 87.5% | 0.35% | 0.675 | 352.0 s | 0.22 s | 1,566.3x |
+| Wisotsky et al. (2020), benchmark suite | 12 | 4,290 | 0.965 | 0.618 | 68.8% | 0.58% | 0.615 | 8,603.0 s | 55.49 s | 608.8x |
+| Hilbert & Elde (2023), Siglec / C-lectins | 57 | 28,462 | 0.892 | 0.224 | 44.8% | 0.17% | 0.376 | 14,203.0 s | 16.12 s | 859.1x |
+| **Global aggregate** | **84** | **43,302** | **0.914** | **0.286** | **50.6%** | **0.191%** | **0.489** | **25,696.0 s** | **78.87 s** | **325.8x** |
+
+Validation metrics recorded in the training checkpoint: Spearman rho 0.5974, MSE
+3.6411, at epoch 16.
+
+## Limitations and intended use
+
+AxoMEME is a research tool for comparative molecular evolution, intended as a fast screen for
+sites under episodic diversifying selection.
+
+- It is trained to approximate HyPhy MEME, so it inherits that method's assumptions and is
+  evaluated against it rather than against ground-truth selection. It is an approximation of
+  MEME, not an independent test of it.
+- Precision falls on large, diverse gene families. On the Hilbert & Elde Siglec/C-lectin set
+  the positive predictive value is 44.8%, so roughly half of flagged sites are not MEME
+  positives. Treat hits as candidates for confirmatory analysis, not conclusions.
+- Aggregate PR-AUC of 0.286 reflects strong class imbalance; positive sites are rare.
+- Accuracy depends on branch lengths being present and meaningful. Topology-only trees are
+  handled but degrade results.
+- Alignments must be in frame. Frameshifted or misaligned input produces silently wrong
+  output.
+- The model has not been evaluated on non-coding sequence, recombinant alignments, or
+  alignments with pervasive sequencing error.
+
+## Citation
+
+A publication is in preparation. Until it appears, cite the repository:
+
+```bibtex
+@software{axomeme,
+  title = {AxoMEME: Ultra-Fast Neural Inference of Episodic Positive Selection in Molecular Sequences},
+  author = {Kosakovsky Pond, Sergei L. and collaborators},
+  url = {https://github.com/veg/axomeme},
+  year = {2026}
+}
+```

config.json

@@ -0,0 +1,11 @@
+{
+  "dropout": 0.1,
+  "embed_dim": 384,
+  "max_k": 32,
+  "max_species": 256,
+  "num_heads": 12,
+  "num_layers": 6,
+  "num_thresholds": 16,
+  "num_tokens": 66,
+  "window_size": 1
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db59db0d13aeca06fddbec3c56280a69915bd6c4fd77a19220c4e82365a5e5bc
+size 7655016