NNMM.jl - Neural Network Mixed Models

NNMM.jl is a Julia package for Mixed Effect Neural Networks that extend traditional linear mixed models to multilayer neural networks for genomic prediction and genome-wide association studies.

Key Features

• Generalizes Traditional Methods: Traditional BayesC and other Bayesian alphabet methods are special cases of NNMM
• Neural Network Architecture: Extend linear mixed models to multilayer neural networks
• Intermediate Omics Integration: Incorporate known intermediate omics features (e.g., gene expression) in the middle layer
• Flexible Missing Data Handling: Allow any patterns of missing data in intermediate layers
• Bayesian Framework: Full Bayesian inference using MCMC and Hamiltonian Monte Carlo
• Multi-threaded Parallelism: Efficient parallel computing for large-scale analyses

Network Architecture

NNMM models relationships between:

• Input Layer: Genotypes (SNP markers)
• Middle Layer: Intermediate traits (e.g., gene expression, metabolomics)
• Output Layer: Phenotypes

The package supports:

• Fully-connected neural networks
• Partial-connected neural networks
• User-defined nonlinear activation functions

Installation

using Pkg
Pkg.add("NNMM")

Or for the latest development version:

Pkg.add(url="https://github.com/reworkhow/NNMM.jl")

Requirements: Julia 1.9 or later.

Quick Start

This example uses built-in simulated data to demonstrate the NNMM workflow:

using NNMM
using NNMM.Datasets
using DataFrames, CSV
using Statistics
using Random

# Set seed for reproducibility
Random.seed!(42)

# Load built-in simulated dataset
geno_path = Datasets.dataset("genotypes_1000snps.txt", dataset_name="simulated_omics_data")
pheno_path = Datasets.dataset("phenotypes_sim.txt", dataset_name="simulated_omics_data")

# Read phenotype file and create separate omics/trait files
pheno_df = CSV.read(pheno_path, DataFrame)

# Create omics file (10 omics features)
omics_df = pheno_df[:, vcat(:ID, [Symbol("omic$i") for i in 1:10])]
omics_path = "omics_temp.csv"
CSV.write(omics_path, omics_df; missingstring="NA")

# Create trait file
trait_df = pheno_df[:, [:ID, :trait1]]
trait_path = "trait_temp.csv"
CSV.write(trait_path, trait_df; missingstring="NA")

# Define network layers (3-layer architecture)
layers = [
    # Layer 1: Genotypes (input layer)
    Layer(layer_name="geno", data_path=[geno_path]),
    # Layer 2: Omics (hidden/intermediate layer)  
    Layer(layer_name="omics", data_path=omics_path, missing_value="NA"),
    # Layer 3: Phenotypes (output layer)
    Layer(layer_name="phenotypes", data_path=trait_path, missing_value="NA")
]

# Define equations connecting layers
equations = [
    # Equation 1: Genotypes → Omics
    Equation(
        from_layer_name = "geno",
        to_layer_name = "omics",
        equation = "omics = intercept + geno",
        omics_name = ["omic$i" for i in 1:10],
        method = "BayesC",
        estimatePi = true
    ),
    # Equation 2: Omics → Phenotypes  
    Equation(
        from_layer_name = "omics",
        to_layer_name = "phenotypes",
        equation = "phenotypes = intercept + omics",
        phenotype_name = ["trait1"],
        method = "BayesC",
        activation_function = "linear"  # or "sigmoid", "tanh", "relu"
    )
]

# Run NNMM analysis
results = runNNMM(layers, equations;
    chain_length = 5000,        # Total MCMC iterations
    burnin = 1000,              # Burn-in iterations to discard
    output_folder = "nnmm_results"
)

# Access results
ebv = results["EBV_NonLinear"]  # Estimated breeding values
println("Mean EBV: ", mean(Float64.(ebv.EBV)))

# Cleanup temporary files
rm(omics_path, force=true)
rm(trait_path, force=true)

API Summary Table

Supporting and Citing

If you use NNMM.jl for your research, please cite:

Zhao, T., Zeng, J., & Cheng, H. (2022). Extend mixed models to multilayer neural networks for genomic prediction including intermediate omics data. GENETICS. https://doi.org/10.1093/genetics/iyac034

Zhao, T., Fernando, R., & Cheng, H. (2021). Interpretable artificial neural networks incorporating Bayesian alphabet models for genome-wide prediction and association studies. G3 Genes|Genomes|Genetics. https://doi.org/10.1093/g3journal/jkab228

Please star the repository here to help demonstrate community involvement.

Getting Help

• Open an issue on GitHub
• Contact: <qtlcheng@ucdavis.edu>

Tutorials

NNMM Tutorials

• Part 1: Introduction to NNMM
• • 1. Overview
  • 2. Extend Linear Mixed Model to Multilayer Neural Networks
  • 3. Flexibility
  • 4. Multi-threaded Parallelism
  • 5. Quick Start
  • 6. Tutorial Overview
• Part 2: Mixed Effect Neural Network (NNMM)
• • Model Architecture
  • Example: Fully-Connected Neural Network with Unobserved Latent Traits
  • Output Files
  • Different Number of Latent Nodes
  • Different Activation Functions
  • Tips for Best Results
• Part 3: NNMM with Intermediate Omics Features
• • Model Architecture
  • Example: Fully-Connected Network with Observed Omics
  • Including Residual Polygenic Effects
  • Handling Missing Omics Data
  • Output Files
  • Tips
• Part 4: Partial-Connected Neural Networks
• • Model Architecture
  • Example: Partial-Connected Network with Observed Omics
  • Example: Partial-Connected Network with Latent Traits
  • Key Points
  • Comparison: Fully-Connected vs Partial-Connected
  • Output Files
• Part 5: User-Defined Nonlinear Functions
• • Overview
  • Current Status
  • Workaround: Use Built-in Activation Functions
  • Example: Using tanh as Alternative to Custom Function
  • Comparison: Built-in Activation Functions
  • Future: Planned User-Defined Function Support
  • Tips
• Part 6: NNMM as Traditional Genomic Prediction
• • Overview
  • The Key Insight
  • Example: Single-Trait BayesC via NNMM
  • Current Limitations
  • Comparison: Direct BayesC vs NNMM as BayesC
  • Why Use NNMM for Traditional Analysis?
  • Other Bayesian Methods
  • Summary

Manual

• Get Started
• • Installation
  • Multi-threaded Parallelism
  • Access Documentation
  • Quick Example
  • Run Your Analysis
  • Key Concepts
  • Next Steps
• Step-by-Step Tutorial
• • Overview
  • Step 1: Load Packages and Data
  • Step 2: Prepare Data Files
  • Step 3: Define the Network Architecture
  • Step 4: Define the Equations
  • Step 5: Run the Analysis
  • Step 6: Examine Results
  • Step 7: Post-Analysis with GWAS
  • Complete Script
  • Next Steps
  • Troubleshooting
• Missing Data Handling
• • Overview
  • NNMM vs “Classic” Bayesian Neural Networks (BNN)
  • Why "Deterministic Latent" Doesn't Work in NNMM
  • Summary Table
  • Detailed Breakdown
  • Detailed: L2 Sampling for Test Individuals
  • Prediction for Test Individuals
  • Visual Summary
  • Implementation Details
  • Example: Train/Test Split
  • See Also
• Skip Connection (Genotypes → Phenotypes) and Multi-Class Priors
• • Model overview
  • Step-by-step example
  • Which output should you compare to JWAS?
  • Omics-class scaling for BayesC priors (JWAS-like)
  • How this differs from the older “all-missing node” pattern
  • Real-data reference script (repository)
• Public API Reference
• • Index
  • Core Types
  • Main Functions
  • Data Reading Functions
  • Post-Analysis Functions
  • Built-in Datasets
  • Pedigree Functions
  • Parameter Reference Tables

Legacy API

# OLD API - DEPRECATED
genotypes = get_genotypes("genotypes.csv", separator=',', method="BayesC")
model = build_model("y = intercept + genotypes", num_hidden_nodes=3, nonlinear_function="tanh")
out = runMCMC(model, phenotypes, chain_length=5000)