mugo 0.1.1

Differentiable Combinatorial Optimization for Genomics

MUGO: Multi-Head Genomic Optimization

MUGO is a differentiable combinatorial optimization framework designed for discovering causal variants in the non-coding genome. By leveraging Gumbel-Softmax relaxation and Straight-Through Estimators (STE), MUGO enables end-to-end gradient-based optimization on discrete DNA sequences.

Key Features

• 🧬 Model-Agnostic: Compatible with Borzoi, Enformer, HyenaDNA, and other PyTorch-based genomic models.
• 🎯 Multi-Modal Objectives: Optimize for Gene Expression, Chromatin Accessibility (ATAC), or TF Binding.
• 📉 Variance Reduction: Built-in Multi-Head Consensus strategy to filter stochastic noise.
• 🚀 Production Ready: Easy-to-use Python API for high-performance computing.

Installation

pip install mugo

Quick Start

import torch
from mugo import MultiHeadSelector

# Initialize the optimizer
selector = MultiHeadSelector(num_snps=1000, snp_positions=positions, k=20)

# Optimization loop
for step in range(200):
    input_seq, mask, _ = selector(ref_seq, alt_seq, tau=1.0)
    # ... compute loss and backward ...

Documentation

Comprehensive documentation, tutorials, and API references are available at: 👉 https://mugo-framework.netlify.app

Citations

If you use MUGO in your research, please cite: @software{mugo2026, author = {SciML Team}, title = {MUGO: Differentiable Combinatorial Optimization for Genomics}, year = {2026}, url = {https://github.com/anonymous/mugo} }