jamma 2.11.0

High-performance Multi-method Mixed-Model Association for large-scale GWAS

JAMMA (High-performance Multi-method Mixed-Model Association) — a modern Python and C reimplementation of GEMMA for large-scale 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 0.98.5 at scale
• Memory-safe: Pre-flight memory checks prevent OOM crashes before allocation
• Cross-platform: Runs on Linux, macOS, and Windows — NumPy backend works everywhere, JAX adds batch acceleration on Linux and ARM Mac
• Pure Python + optional C extensions: NumPy + optional JAX stack; C extensions for DSYEVR eigendecomposition and OpenMP-parallel Wald tests, JAX for batch MLE optimization
• Large-scale ready: Optional numpy-mkl ILP64 wheels (numpy 2.4.2) for >46k sample eigendecomposition

Installation

macOS (Intel or ARM)

pip install jamma          # NumPy backend
pip install 'jamma[jax]'   # + JAX acceleration (ARM Mac only)

That's it. macOS Accelerate BLAS handles large matrices natively.

Linux / Windows / Intel Mac

For small datasets (<46k samples), the standard install works:

pip install jamma          # NumPy backend
pip install 'jamma[jax]'   # + JAX acceleration

For large-scale GWAS (>46k samples) on x86_64 (Linux or Intel Mac), install numpy-mkl first — standard numpy uses 32-bit BLAS integers which overflow at ~46k samples. MKL is x86_64-only; ARM Mac and Windows users are limited to <46k samples. Pre-built ILP64 wheels are available for Python 3.11–3.14:

NumPy backend only:

pip install numpy \
  --extra-index-url https://michael-denyer.github.io/numpy-mkl \
  --force-reinstall --upgrade
pip install jamma --no-deps
pip install psutil loguru threadpoolctl click progressbar2 bed-reader

With JAX acceleration:

pip install numpy \
  --extra-index-url https://michael-denyer.github.io/numpy-mkl \
  --force-reinstall --upgrade
pip install 'jamma[jax]' --no-deps
pip install psutil loguru threadpoolctl click progressbar2 bed-reader \
  jax jaxlib jaxtyping

From Git (latest development version):

pip install numpy \
  --extra-index-url https://michael-denyer.github.io/numpy-mkl \
  --force-reinstall --upgrade
pip install git+https://github.com/michael-denyer/jamma.git --no-deps
pip install psutil loguru threadpoolctl click progressbar2 bed-reader

Why --no-deps? JAMMA depends on numpy>=2.0.0, so a normal pip install jamma will pull in standard numpy and overwrite the ILP64 build. --no-deps prevents this; you install the runtime dependencies manually instead.

See the User Guide for ILP64 verification steps.

Platform Support

Platform	pip install jamma	pip install jamma[jax]	Notes
Linux x86_64	JAX (auto-included)	—	Full support; ILP64 for >46k samples
ARM Mac (M1+)	JAX (auto-included)	—	Full support
Intel Mac	NumPy only	Not available	JAX dropped Intel Mac; ILP64 for >46k samples
Windows	NumPy only	Not available	JAX dropped Windows support

JAX is auto-included on Linux and ARM Mac via platform markers. Force a specific backend with --backend numpy or --backend jax.

Quick Start

# Compute kinship matrix (centered relatedness)
jamma -gk 1 -bfile data/my_study -o output
# Output: output/output.cXX.npy (binary, fast)
# Add --legacy-text for GEMMA-compatible text format

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

# Multiple phenotypes (eigendecomp computed once, reused)
jamma -lmm 1 -bfile data/my_study -k output/output.cXX.npy -n "1 2 3" -o results

Output files:

• output.cXX.npy — Kinship matrix (binary NumPy format; .cXX.txt with --legacy-text)
• 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

The reader auto-detects format, so existing .cXX.txt files still work as -k input.

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 (Python API)
result = gwas("data/my_study", write_eigen=True, phenotype_column=1)
result = gwas("data/my_study", eigenvalue_file="output/result.eigenD.npy",
              eigenvector_file="output/result.eigenU.npy", phenotype_column=2)
# Or use the CLI for automatic multi-phenotype: jamma -lmm 1 ... -n "1 2 3"

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

Low-level API (JAX backend)

import numpy as np

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 and phenotypes
data = load_plink_binary("data/my_study")
phenotypes = np.loadtxt("data/my_study.pheno")  # loaded separately from .fam or phenotype file

# Compute kinship and eigendecompose (treat kinship as consumed after this)
kinship = compute_centered_kinship(data.genotypes)
eigenvalues, eigenvectors = eigendecompose_kinship(kinship)

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

Low-level API (NumPy backend)

import numpy as np

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

data = load_plink_binary("data/my_study")
phenotypes = np.loadtxt("data/my_study.pheno")
kinship = compute_centered_kinship(data.genotypes)
eigenvalues, eigenvectors = eigendecompose_kinship(kinship)

snp_info = [
    {"chr": str(data.chromosome[i]), "rs": data.sid[i],
     "pos": int(data.bp_position[i]), "a1": data.allele_1[i], "a0": data.allele_2[i]}
    for i in range(data.n_snps)
]

# Returns list[AssocResult] — write to disk via IncrementalAssocWriter
results = run_lmm_association_numpy(
    genotypes=data.genotypes,
    phenotypes=phenotypes,
    kinship=None,  # Not needed when eigenvalues/eigenvectors provided
    snp_info=snp_info,
    eigenvalues=eigenvalues,
    eigenvectors=eigenvectors,
    lmm_mode=1,
)

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 (AC power), GEMMA 0.98.5. Best of multiple runs, end-to-end wall clock:

Operation	GEMMA 0.98.5	JAMMA NumPy+C	JAMMA JAX (batch)	JAMMA JAX (streaming)	vs GEMMA
Kinship (-gk 1)	2.1s	259ms	259ms	—	8.1x
LMM Wald (-lmm 1)	11.1s	1.0s	2.0s	2.7s	11.1x
LMM All (-lmm 4)	20.6s	5.1s	2.8s	4.3s	7.3x

NumPy+C uses a C extension with OpenMP for Wald-only (-lmm 1) — REML optimization is compute-bound and parallelizes well across SNPs. JAX (batch) pulls ahead on all-tests (-lmm 4) because the additional MLE optimization per SNP benefits from jax.vmap batching. JAX (streaming) reads genotypes from disk in chunks and is the production code path for large datasets that don't fit in memory.

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)
• Binary .npy I/O — default for kinship and eigen files; --legacy-text for GEMMA text format
• Multi-phenotype support — -n "1 2 3" with single eigendecomposition reuse
• Eigendecomposition reuse — manual via -d/-u/-eigen, automatic in multi-phenotype mode
• 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) with automatic 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
• Optional C extensions: DSYEVR eigendecomposition (O(n) workspace, enables >100k samples) and OpenMP-parallel Wald tests (auto-fallback to pure Python)

Planned

• Multivariate LMM (mvLMM)

Architecture

JAMMA uses NumPy for data loading and kinship. Eigendecomposition defaults to DSYEVD (via numpy) but falls back to DSYEVR (C extension, O(n) workspace) under memory pressure — critical for >100k samples. At LMM it splits into a JAX backend (JIT, vmap, sharding) or a NumPy backend with an optional C extension for OpenMP-parallel Wald tests.

flowchart TD
    CLI["CLI / gwas()"] --> PIPE["PipelineRunner"]
    PIPE --> LOAD["Load PLINK + Phenotypes<br>(NumPy)"]
    LOAD --> KIN["Kinship<br>(NumPy matmul)"]
    KIN --> EIGMEM{"DSYEVD fits<br>in memory?"}
    EIGMEM -->|yes| EIGD["Eigendecomposition<br>(LAPACK DSYEVD · O(n²) workspace)"]
    EIGMEM -->|no| EIGR["Eigendecomposition<br>(LAPACK DSYEVR · O(n) workspace)"]
    EIGD --> DET{"detect_backend()"}
    EIGR --> DET
    DET -->|"jax"| JAX["JAX Streaming Runner<br>JIT + vmap + sharding"]
    DET -->|"numpy"| NP["NumPy Batch Runner"]
    NP --> CEXT{"C LMM extension<br>available?"}
    CEXT -->|yes| C["C Extension<br>OpenMP + SIMD"]
    CEXT -->|no| PY["Pure Python<br>fallback"]
    JAX --> RES["AssocResult"]
    C --> RES
    PY --> RES

Both backends share the same core algorithms (likelihood.py, prepare_common.py) and produce identical results. Backend-specific files follow a naming convention: *_jax.py / *_numpy.py.

See Code Map for the full architecture diagram with source links.

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+
• NumPy 2.0+
• JAX 0.5.0+ (auto-included on Linux/ARM Mac; explicit extra on other platforms: pip install 'jamma[jax]')

License

GPL-3.0 (same as GEMMA)