jamma 2.5.5

JAMMA: JAX-Accelerated Mixed Model Association

JAX-Accelerated Mixed Model Association — A modern Python reimplementation of GEMMA for genome-wide association studies (GWAS).

• GEMMA-compatible: Drop-in replacement with identical CLI flags and output formats
• Numerical equivalence: Validated against GEMMA — 100% significance agreement, 100% effect direction agreement
• Fast: Up to 11x faster than GEMMA on kinship and 6x faster on LMM association
• Memory-safe: Pre-flight memory checks prevent OOM crashes before allocation
• Pure Python: JAX + NumPy stack, no C++ compilation required
• Large-scale ready: Optional numpy-mkl ILP64 wheels (numpy 2.4.2) for >46k sample eigendecomposition

Installation

pip install jamma

Or with uv:

uv add jamma

Quick Start

# Compute kinship matrix (centered relatedness)
jamma -gk 1 -bfile data/my_study -o output

# Run LMM association (Wald test)
jamma -lmm 1 -bfile data/my_study -k output/output.cXX.txt -o results

Output files match GEMMA format exactly:

• output.cXX.txt — Kinship matrix
• results.assoc.txt — Association results (chr, rs, ps, n_miss, allele1, allele0, af, beta, se, logl_H1, l_remle, p_wald)
• results.log.txt — Run log

Python API

One-call GWAS (recommended)

from jamma import gwas

# Full pipeline: load data → kinship → eigendecomp → LMM → results
result = gwas("data/my_study", kinship_file="data/kinship.cXX.txt")
print(f"Tested {result.n_snps_tested} SNPs in {result.timing['total_s']:.1f}s")

# Compute kinship from scratch and save it
result = gwas("data/my_study", save_kinship=True, output_dir="output")

# With covariates and LRT test
result = gwas("data/my_study", kinship_file="k.txt", covariate_file="covars.txt", lmm_mode=2)

# LOCO analysis (leave-one-chromosome-out)
result = gwas("data/my_study", loco=True)

# Multi-phenotype with eigendecomp reuse
result = gwas("data/my_study", write_eigen=True, phenotype_column=1)
result = gwas("data/my_study", eigenvalue_file="output/result.eigenD.txt",
              eigenvector_file="output/result.eigenU.txt", phenotype_column=2)

# SNP filtering
result = gwas("data/my_study", kinship_file="k.txt", snps_file="snps.txt", hwe=0.001)

Low-level API

from jamma.io import load_plink_binary
from jamma.kinship import compute_centered_kinship
from jamma.lmm import run_lmm_association_streaming
from jamma.lmm.eigen import eigendecompose_kinship

# Load PLINK data
data = load_plink_binary("data/my_study")

# Compute kinship
kinship = compute_centered_kinship(data.genotypes)

# Eigendecompose for LMM
eigenvalues, eigenvectors = eigendecompose_kinship(kinship)

# Run association (streaming from disk)
results = run_lmm_association_streaming(
    bed_path="data/my_study",
    phenotypes=phenotypes,
    kinship=kinship,
    chunk_size=5000,
)

Memory Safety

Unlike GEMMA, JAMMA includes pre-flight memory checks that prevent out-of-memory crashes:

from jamma.core.memory import estimate_workflow_memory

# Check memory requirements BEFORE loading data
estimate = estimate_workflow_memory(n_samples=200_000, n_snps=95_000)
print(f"Peak memory: {estimate.total_gb:.1f}GB")
print(f"Available: {estimate.available_gb:.1f}GB")
print(f"Sufficient: {estimate.sufficient}")

Key features:

• Pre-flight checks before large allocations (eigendecomposition, genotype loading)
• RSS memory logging at workflow boundaries
• Incremental result writing (no memory accumulation)
• Safe chunk size defaults with hard caps

GEMMA will silently OOM and get killed by the OS. JAMMA fails fast with clear error messages.

Performance

Benchmark on mouse_hs1940 (1,940 samples × 12,226 SNPs), Apple M2:

Operation	GEMMA	JAMMA	Speedup
Kinship (-gk 1)	26.5s	2.4s	11.0x
LMM (-lmm 1)	27.6s	4.5s	6.1x
Total	54.1s	6.9s	7.8x

Supported Features

Current

• Kinship matrix computation — centered (-gk 1) and standardized (-gk 2)
• Univariate LMM Wald test (-lmm 1)
• Likelihood ratio test (-lmm 2)
• Score test (-lmm 3)
• All tests mode (-lmm 4)
• LOCO kinship — leave-one-chromosome-out analysis (-loco)
• Eigendecomposition reuse — multi-phenotype workflows (-d/-u/-eigen)
• Phenotype column selection (-n)
• SNP subset selection for association and kinship (-snps/-ksnps)
• HWE QC filtering (-hwe)
• Pre-computed kinship input (-k)
• Covariate support (-c)
• PLINK binary format (.bed/.bim/.fam) with input dimension validation
• Large-scale streaming I/O (>100k samples via numpy-mkl ILP64 — numpy 2.4.2)
• JAX acceleration (CPU/GPU) with automatic CPU device sharding
• XLA profiling traces (--profile-dir) for TensorBoard/Perfetto
• Lambda optimization bounds (-lmin/-lmax)
• Individual weights for kinship (-widv)
• Categorical covariates with one-hot encoding (-cat)
• Pre-flight memory checks (fail-fast before OOM)
• RSS memory logging at workflow boundaries
• Incremental result writing

Planned

• Multivariate LMM (mvLMM)

Documentation

• Why JAMMA? — Key differentiators from GEMMA
• User Guide — Installation, usage examples, CLI reference
• Code Map — Architecture diagrams and source navigation
• Equivalence Proof — Mathematical proofs and empirical validation against GEMMA
• GEMMA Divergences — Known differences from GEMMA
• Performance — Bottleneck analysis, scale validation, configuration guide
• Contributing — Development setup, testing, and PR guidelines
• Changelog — Version history

Requirements

• Python 3.11+
• JAX 0.8.0+
• NumPy 2.0+

License

GPL-3.0 (same as GEMMA)