This repository provides the official implementation of Coarse-Grained Boltzmann Generators (CG-BGs). Unlike traditional BGs that operate in all-atom space, CG-BGs act directly in a coarse-grained coordinate space. By leveraging Conditional Flow Matching and Enhanced Sampling Force Matching, this method enables exact equilibrium sampling and reweighting for CG molecular systems.

This project uses pixi for package management to ensure strict version control and reproducibility across different systems.

pixi install --frozen

Six standard tasks are pre-configured:

• mb_flow_unbiased
• mb_flow_biased
• aldp_HA_flow_unbiased
• aldp_HA_flow_biased
• aldp_CB_flow_unbiased
• aldp_CB_flow_biased

Run a task using the following command:

pixi run <task_name>

Configuration Overrides All configurations are managed by hydra-zen. Parameters can be modified via the Command Line Interface (CLI) or configuration files.

Example: Running specific stages on a designated GPU

# stage: 1-Training, 2-Sampling, 3-Reweighting, 4-Analysis
pixi run <task_name> device=<gpu_id> stage=234 hydra.run.dir=<output_dir>

If you use CG-BGs please cite:

@misc{chen2026coarsegrainedboltzmanngenerators,
      title={Coarse-Grained Boltzmann Generators}, 
      author={Weilong Chen and Bojun Zhao and Jan Eckwert and Julija Zavadlav},
      year={2026},
      eprint={2602.10637},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2602.10637}, 
}