Skip to content

michael-denyer/jamma

Repository files navigation

CI PyPI Python 3.11+ NumPy Hypothesis License: GPL-3.0 Buy Me a Coffee

JAMMA

JAMMA (Highly-Accelerated Multi-method Mixed-Model Association) -- a modern Python and C reimplementation of GEMMA for large-scale GWAS.

  • Drop-in GEMMA replacement: Same CLI flags, same file formats, same results. Change one word in your pipeline.
  • Numerical equivalence: Validated against GEMMA -- 100% significance agreement, 100% effect direction agreement
  • Fast: Up to 30x faster than GEMMA 0.98.5 (LOCO mode); 12-17x on single-pass LMM
  • Memory-safe: Pre-flight memory checks prevent OOM crashes before allocation
  • Cross-platform: Runs on Linux, macOS, and Windows with NumPy and vendor BLAS
  • Optimized for Intel: Best performance on Intel CPUs with MKL BLAS. Runs well on Apple Silicon (Accelerate BLAS). Other architectures (AMD, ARM Linux) work correctly but with less BLAS optimization
  • Pure Python + C extensions (OpenMP SIMD): NumPy stack with vendor BLAS dispatch (MKL-ILP64, Accelerate-ILP64) via jlinalg C layer for eigendecomposition and OpenMP-parallel Wald tests
  • Large-scale ready: Optional numpy-mkl ILP64 wheels (numpy 2.4.4) for >46k sample eigendecomposition

Installation

macOS (13.3+)

pip install jamma

That's it. macOS Accelerate BLAS handles large matrices natively (Accelerate-ILP64).

Windows (10+), Windows Server (2016+) and Linux (Intel/AMD)

Install numpy-mkl first -- standard numpy uses 32-bit BLAS integers which overflow at ~46k samples. Pre-built ILP64 wheels are available for Python 3.11-3.14:

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

From Git (latest development version):

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

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 BLAS ILP64 Notes
Linux x86_64 MKL (optimal) numpy-mkl Best performance
ARM Linux OpenBLAS -- Works correctly
ARM Mac (M1+) Accelerate native Excellent performance
Intel Mac (macOS 13.3+) Accelerate native Full support
Windows x86_64 (10+) MKL (optimal) numpy-mkl Best performance
Windows Server x86_64 (2016+) MKL (optimal) numpy-mkl Best performance

See the User Guide for BLAS backend details.

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.

GEMMA CLI Parity

JAMMA supports GEMMA's core GWAS flags (-gk, -lmm, -bfile, -k, -c, -o, -n, -loco, -snps, -hwe) with identical names and semantics. Existing GEMMA commands work by changing gemma to jamma:

GEMMA JAMMA
gemma -gk 1 -bfile study -o out jamma -gk 1 -bfile study -o out
gemma -lmm 1 -bfile study -k kinship.cXX.txt -o results jamma -lmm 1 -bfile study -k kinship.cXX.txt -o results
gemma -lmm 4 -bfile study -k k.txt -c covars.txt -o results jamma -lmm 4 -bfile study -k k.txt -c covars.txt -o results
  • Reads and writes GEMMA .assoc.txt and .cXX.txt formats
  • Accepts PLINK binary .bed/.bim/.fam files (same as GEMMA)
  • Output columns match GEMMA (mode-dependent -- see User Guide)
  • Also supports binary .npy format for kinship (faster I/O); use --legacy-text for GEMMA text format

Python API

The gwas() function handles the full pipeline -- data loading, kinship computation, eigendecomposition, and LMM association -- in a single call. You don't need to compute a kinship matrix separately unless you want to reuse it across runs.

from jamma import gwas

# Simplest usage: computes kinship internally, no separate kinship step needed
result = gwas("data/my_study")
print(f"Tested {result.n_snps_tested} SNPs in {result.timing['total_s']:.1f}s")

# Or supply a pre-computed kinship matrix to skip recomputation
result = gwas("data/my_study", kinship_file="data/kinship.cXX.npy")

# Compute kinship from scratch and save it for reuse
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)

# LOCO with eigen caching: writes a per-chromosome eigen cache to output_dir
result = gwas("data/my_study", loco=True, write_eigen=True, output_dir="output")
# Reusing a LOCO eigen cache on a later run is CLI-only — the cache is a set of
# per-chromosome files keyed by directory, so point --eigen-dir at the same dir:
#   jamma -lmm 1 -bfile data/my_study -loco --eigen-dir output -o result

# 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)

See the User Guide for the low-level component API (kinship, eigendecomposition, LMM runners).

Memory Safety

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

  • 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. See the User Guide for the programmatic memory estimation API.

Performance

Benchmark on mouse_hs1940 (1,940 samples x 12,226 SNPs), Apple M2, GEMMA 0.98.5. Best-of runs, end-to-end wall clock:

Operation GEMMA (OpenBLAS) GEMMA (Accelerate) JAMMA NumPy JAMMA NumPy+C JAMMA NumPy+C (stream) C speedup vs GEMMA (OB) vs GEMMA (Accel)
Kinship (-gk 1) 2.1s 1.7s 262ms 262ms -- 1.0x 8.0x 6.5x
LMM Wald (-lmm 1) 11.0s 7.6s 4.1s 879ms 1.1s 4.7x 12.5x 8.7x
LMM All (-lmm 4) 20.5s 13.9s 6.0s 1.3s 1.4s 4.7x 16.0x 10.9x
LMM Wald+4cov (-lmm 1 -c) 40.8s 18.8s 9.1s 2.4s 2.6s 3.8x 17.0x 7.8x
LOCO Wald (-loco) 3m30s 2m26s -- 7.1s -- -- 29.6x 20.6x

See Performance for benchmark methodology and large-scale (125k) results.

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
  • LOCO eigen caching -- --eigen-dir saves/loads per-chromosome eigen files across runs
  • 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.4)
  • 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
  • In-place mean imputation for missing genotypes (per-chunk, zero-copy)
  • Early sample filtering -- kinship accumulated at filtered size when phenotype missingness is present
  • jlinalg C layer: vendor BLAS dispatch for eigendecomposition (DSYEVD default, DSYEVR O(n) workspace fallback under memory pressure), DSYRK, DGEMM
  • Optional C extension: OpenMP-parallel Wald tests (auto-fallback to pure Python)

Planned

  • Multivariate LMM (mvLMM)

Architecture

JAMMA uses NumPy for data loading and kinship. Eigendecomposition uses jlinalg.eigh which dispatches to vendor DSYEVD (default) or DSYEVR (O(n) workspace, under memory pressure) via the jlinalg C layer. LMM association uses a NumPy backend with an optional C extension for OpenMP-parallel Wald/Score/LRT tests. Mode is auto-selected based on available memory: batch runner when genotypes fit in RAM, streaming runner (two-pass disk I/O) for large datasets.

flowchart TD
    subgraph ENTRY["ENTRY"]
        CLI["CLI / gwas()"]
        PIPE["PipelineRunner"]
        CLI --> PIPE
    end

    subgraph IO["DATA LOADING"]
        LOAD["Load PLINK +<br/>Phenotypes"]
    end

    subgraph CORE["CORE COMPUTATION"]
        KIN["Kinship<br/>(DGEMM, chunked)"]
        EIG["Eigendecomposition<br/>(jlinalg.eigh → DSYEVD/DSYEVR)"]
        KIN --> EIG
    end

    subgraph ASSOC["ASSOCIATION TESTING"]
        MEM{"Memory<br/>budget?"}
        NP["Batch Runner<br/>(genotypes in RAM)"]
        NPS["Streaming Runner<br/>(two-pass disk I/O)"]
        CEXT{"C extension?"}
        C["C Extension<br/>OpenMP + SIMD"]
        PY["Pure Python<br/>fallback"]
        MEM -->|fits| NP
        MEM -->|large| NPS
        NP --> CEXT
        NPS --> CEXT
        CEXT -->|yes| C
        CEXT -->|no| PY
    end

    RES["AssocResult<br/>(.assoc.txt)"]

    PIPE --> LOAD --> CORE
    EIG --> ASSOC
    C --> RES
    PY --> RES

    style ENTRY fill:#1a1a2e,stroke:#53a8b6,color:#eee,stroke-width:2px
    style IO fill:#1a1a2e,stroke:#53a8b6,color:#eee,stroke-width:2px
    style CORE fill:#0f3460,stroke:#f5b461,color:#eee,stroke-width:2px
    style ASSOC fill:#0f3460,stroke:#e94560,color:#eee,stroke-width:2px

    style CLI fill:#53a8b6,stroke:#3d8a96,color:#fff
    style PIPE fill:#53a8b6,stroke:#3d8a96,color:#fff
    style LOAD fill:#53a8b6,stroke:#3d8a96,color:#fff

    style KIN fill:#f5b461,stroke:#d4943f,color:#1a1a2e
    style EIG fill:#f5b461,stroke:#d4943f,color:#1a1a2e

    style MEM fill:#e94560,stroke:#c73550,color:#fff
    style NP fill:#7b68ae,stroke:#5a4d8a,color:#fff
    style NPS fill:#7b68ae,stroke:#5a4d8a,color:#fff
    style CEXT fill:#e94560,stroke:#c73550,color:#fff
    style C fill:#2ecc71,stroke:#27ae60,color:#1a1a2e
    style PY fill:#95a5a6,stroke:#7f8c8d,color:#1a1a2e

    style RES fill:#2ecc71,stroke:#27ae60,color:#1a1a2e
Loading

Core algorithms (likelihood.py, prepare_common.py) are shared between batch and streaming runners. See jlinalg Architecture for the C vendor BLAS dispatch layer.

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

Documentation

Requirements

  • Python 3.11+
  • NumPy 2.0+

License

GPL-3.0 (same as GEMMA).

About

JAMMA (Highly-Accelerated Multi-method Mixed-model Association) -- a fast and modern Python and C reimplementation of GEMMA for large-scale GWAS.

Resources

License

Contributing

Stars

Watchers

Forks

Packages

 
 
 

Contributors