boltzmann-generators 0.2.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- boltzmann_generators-0.2.1/.coverage +0 -0
- boltzmann_generators-0.2.1/.github/workflows/cd.yml +107 -0
- boltzmann_generators-0.2.1/.github/workflows/ci.yml +35 -0
- boltzmann_generators-0.2.1/.gitignore +27 -0
- boltzmann_generators-0.2.1/.pre-commit-config.yaml +19 -0
- boltzmann_generators-0.2.1/.python-version +1 -0
- boltzmann_generators-0.2.1/CHANGELOG.md +50 -0
- boltzmann_generators-0.2.1/CITATION.cff +28 -0
- boltzmann_generators-0.2.1/CONTRIBUTING.md +38 -0
- boltzmann_generators-0.2.1/LICENSE +21 -0
- boltzmann_generators-0.2.1/MIGRATION.md +76 -0
- boltzmann_generators-0.2.1/PKG-INFO +146 -0
- boltzmann_generators-0.2.1/README.md +104 -0
- boltzmann_generators-0.2.1/RELEASE.md +55 -0
- boltzmann_generators-0.2.1/examples/01_quickstart_realnvp.ipynb +52 -0
- boltzmann_generators-0.2.1/examples/02_reweighting_and_ess.ipynb +56 -0
- boltzmann_generators-0.2.1/examples/03_cnf_flow_matching.ipynb +54 -0
- boltzmann_generators-0.2.1/examples/04_checkpoint_and_resume.ipynb +62 -0
- boltzmann_generators-0.2.1/pyproject.toml +109 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/__init__.py +18 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/analysis.py +36 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/base/__init__.py +6 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/base/density.py +36 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/base/energy.py +20 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/energies/__init__.py +5 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/energies/dipeptide.py +81 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/energies/double_well.py +56 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/energies/muller.py +65 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/extensions/__init__.py +1 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/extensions/equivariant.py +7 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/extensions/molecular.py +7 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/extensions/transferable.py +7 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/flows/__init__.py +20 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/flows/base.py +68 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/flows/cnf.py +266 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/flows/coupling.py +89 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/flows/periodic.py +27 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/flows/realnvp.py +79 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/io.py +34 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/losses.py +70 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/mcmc.py +68 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/py.typed +0 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/sampling.py +54 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/services/__init__.py +7 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/services/analysis.py +65 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/services/checkpoint.py +43 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/services/sampling.py +74 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/train.py +40 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/training/__init__.py +19 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/training/loss_strategies.py +113 -0
- boltzmann_generators-0.2.1/src/boltzmann_generators/training/trainer.py +94 -0
- boltzmann_generators-0.2.1/tests/conftest.py +22 -0
- boltzmann_generators-0.2.1/tests/test_analysis.py +41 -0
- boltzmann_generators-0.2.1/tests/test_base.py +40 -0
- boltzmann_generators-0.2.1/tests/test_cnf.py +24 -0
- boltzmann_generators-0.2.1/tests/test_energies.py +33 -0
- boltzmann_generators-0.2.1/tests/test_integration.py +38 -0
- boltzmann_generators-0.2.1/tests/test_io.py +42 -0
- boltzmann_generators-0.2.1/tests/test_losses.py +38 -0
- boltzmann_generators-0.2.1/tests/test_mcmc.py +15 -0
- boltzmann_generators-0.2.1/tests/test_periodic.py +21 -0
- boltzmann_generators-0.2.1/tests/test_realnvp.py +21 -0
- boltzmann_generators-0.2.1/tests/test_sampling.py +27 -0
- boltzmann_generators-0.2.1/tests/test_services.py +101 -0
- boltzmann_generators-0.2.1/tests/test_train.py +20 -0
- boltzmann_generators-0.2.1/tests/test_trainer.py +47 -0
- boltzmann_generators-0.2.1/uv.lock +3475 -0
|
Binary file
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
name: CD
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main]
|
|
6
|
+
|
|
7
|
+
permissions:
|
|
8
|
+
contents: write
|
|
9
|
+
id-token: write
|
|
10
|
+
|
|
11
|
+
jobs:
|
|
12
|
+
detect-version-bump:
|
|
13
|
+
runs-on: ubuntu-latest
|
|
14
|
+
outputs:
|
|
15
|
+
should_release: ${{ steps.detect.outputs.should_release }}
|
|
16
|
+
version: ${{ steps.detect.outputs.version }}
|
|
17
|
+
steps:
|
|
18
|
+
- uses: actions/checkout@v4
|
|
19
|
+
with:
|
|
20
|
+
fetch-depth: 2
|
|
21
|
+
- name: Detect version bump
|
|
22
|
+
id: detect
|
|
23
|
+
run: |
|
|
24
|
+
CURRENT=$(grep '^version = ' pyproject.toml | head -1 | sed 's/version = "\(.*\)"/\1/')
|
|
25
|
+
PREVIOUS=$(git show HEAD~1:pyproject.toml 2>/dev/null | grep '^version = ' | head -1 | sed 's/version = "\(.*\)"/\1/' || true)
|
|
26
|
+
|
|
27
|
+
if [ -z "$CURRENT" ]; then
|
|
28
|
+
echo "Could not read version from pyproject.toml"
|
|
29
|
+
exit 1
|
|
30
|
+
fi
|
|
31
|
+
|
|
32
|
+
echo "Current version: $CURRENT"
|
|
33
|
+
echo "Previous version: ${PREVIOUS:-<none>}"
|
|
34
|
+
|
|
35
|
+
if [ "$CURRENT" != "$PREVIOUS" ]; then
|
|
36
|
+
echo "should_release=true" >> "$GITHUB_OUTPUT"
|
|
37
|
+
echo "version=$CURRENT" >> "$GITHUB_OUTPUT"
|
|
38
|
+
else
|
|
39
|
+
echo "should_release=false" >> "$GITHUB_OUTPUT"
|
|
40
|
+
echo "No version bump detected; skipping release."
|
|
41
|
+
fi
|
|
42
|
+
|
|
43
|
+
validate-changelog:
|
|
44
|
+
needs: detect-version-bump
|
|
45
|
+
if: needs.detect-version-bump.outputs.should_release == 'true'
|
|
46
|
+
runs-on: ubuntu-latest
|
|
47
|
+
steps:
|
|
48
|
+
- uses: actions/checkout@v4
|
|
49
|
+
- name: Verify changelog contains release version
|
|
50
|
+
env:
|
|
51
|
+
VERSION: ${{ needs.detect-version-bump.outputs.version }}
|
|
52
|
+
run: |
|
|
53
|
+
if ! grep -q "\[${VERSION}\]" CHANGELOG.md; then
|
|
54
|
+
echo "CHANGELOG.md is missing section for version ${VERSION}"
|
|
55
|
+
exit 1
|
|
56
|
+
fi
|
|
57
|
+
|
|
58
|
+
release:
|
|
59
|
+
needs: [detect-version-bump, validate-changelog]
|
|
60
|
+
if: needs.detect-version-bump.outputs.should_release == 'true'
|
|
61
|
+
runs-on: ubuntu-latest
|
|
62
|
+
environment: pypi
|
|
63
|
+
steps:
|
|
64
|
+
- uses: actions/checkout@v4
|
|
65
|
+
with:
|
|
66
|
+
fetch-depth: 0
|
|
67
|
+
- uses: actions/setup-python@v5
|
|
68
|
+
with:
|
|
69
|
+
python-version: "3.11"
|
|
70
|
+
- name: Install uv
|
|
71
|
+
uses: astral-sh/setup-uv@v3
|
|
72
|
+
- name: Install build tools
|
|
73
|
+
run: uv sync --extra dev
|
|
74
|
+
- name: Build package
|
|
75
|
+
run: uv run python -m build
|
|
76
|
+
- name: Check distribution
|
|
77
|
+
run: uv run twine check dist/*
|
|
78
|
+
- name: Create git tag
|
|
79
|
+
env:
|
|
80
|
+
VERSION: ${{ needs.detect-version-bump.outputs.version }}
|
|
81
|
+
run: |
|
|
82
|
+
TAG="v${VERSION}"
|
|
83
|
+
if git rev-parse "$TAG" >/dev/null 2>&1; then
|
|
84
|
+
echo "Tag ${TAG} already exists; skipping tag creation."
|
|
85
|
+
exit 0
|
|
86
|
+
fi
|
|
87
|
+
git config user.name "github-actions[bot]"
|
|
88
|
+
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
|
|
89
|
+
git tag -a "$TAG" -m "Release ${TAG}"
|
|
90
|
+
git push origin "$TAG"
|
|
91
|
+
- name: Publish to PyPI
|
|
92
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
93
|
+
with:
|
|
94
|
+
password: ${{ secrets.PYPI_TOKEN }}
|
|
95
|
+
- name: Create GitHub Release
|
|
96
|
+
env:
|
|
97
|
+
VERSION: ${{ needs.detect-version-bump.outputs.version }}
|
|
98
|
+
GH_TOKEN: ${{ github.token }}
|
|
99
|
+
run: |
|
|
100
|
+
TAG="v${VERSION}"
|
|
101
|
+
if gh release view "$TAG" >/dev/null 2>&1; then
|
|
102
|
+
echo "GitHub release ${TAG} already exists; skipping."
|
|
103
|
+
exit 0
|
|
104
|
+
fi
|
|
105
|
+
gh release create "$TAG" \
|
|
106
|
+
--title "v${VERSION}" \
|
|
107
|
+
--notes "See [CHANGELOG.md](https://github.com/${{ github.repository }}/blob/main/CHANGELOG.md#${VERSION//./}) for details."
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
name: CI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
pull_request:
|
|
5
|
+
branches: [main]
|
|
6
|
+
|
|
7
|
+
jobs:
|
|
8
|
+
test:
|
|
9
|
+
runs-on: ubuntu-latest
|
|
10
|
+
strategy:
|
|
11
|
+
fail-fast: false
|
|
12
|
+
matrix:
|
|
13
|
+
python-version: ["3.11", "3.12"]
|
|
14
|
+
|
|
15
|
+
steps:
|
|
16
|
+
- uses: actions/checkout@v4
|
|
17
|
+
- uses: actions/setup-python@v5
|
|
18
|
+
with:
|
|
19
|
+
python-version: ${{ matrix.python-version }}
|
|
20
|
+
- name: Install uv
|
|
21
|
+
uses: astral-sh/setup-uv@v3
|
|
22
|
+
- name: Install dependencies
|
|
23
|
+
run: uv sync --extra dev
|
|
24
|
+
- name: Lint
|
|
25
|
+
run: uv run ruff check .
|
|
26
|
+
- name: Format check
|
|
27
|
+
run: uv run ruff format --check .
|
|
28
|
+
- name: Type check
|
|
29
|
+
run: uv run mypy src
|
|
30
|
+
- name: Test with coverage
|
|
31
|
+
run: uv run pytest
|
|
32
|
+
- name: Build package
|
|
33
|
+
run: uv run python -m build
|
|
34
|
+
- name: Check distribution
|
|
35
|
+
run: uv run twine check dist/*
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
# Python
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[oc]
|
|
4
|
+
build/
|
|
5
|
+
dist/
|
|
6
|
+
wheels/
|
|
7
|
+
*.egg-info
|
|
8
|
+
|
|
9
|
+
# Virtual envs
|
|
10
|
+
.venv
|
|
11
|
+
|
|
12
|
+
# Jupyter
|
|
13
|
+
.ipynb_checkpoints/
|
|
14
|
+
*.ipynb_checkpoints
|
|
15
|
+
|
|
16
|
+
# Data + model artifacts
|
|
17
|
+
data/
|
|
18
|
+
*.pt
|
|
19
|
+
*.pth
|
|
20
|
+
*.npz
|
|
21
|
+
*.h5
|
|
22
|
+
|
|
23
|
+
# Reference repos (cloned, not authored here)
|
|
24
|
+
references/paper_boltzmann_generators/
|
|
25
|
+
|
|
26
|
+
# OS
|
|
27
|
+
.DS_Store
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
repos:
|
|
2
|
+
- repo: https://github.com/astral-sh/ruff-pre-commit
|
|
3
|
+
rev: v0.15.17
|
|
4
|
+
hooks:
|
|
5
|
+
- id: ruff
|
|
6
|
+
- id: ruff-format
|
|
7
|
+
- repo: https://github.com/pre-commit/pre-commit-hooks
|
|
8
|
+
rev: v5.0.0
|
|
9
|
+
hooks:
|
|
10
|
+
- id: end-of-file-fixer
|
|
11
|
+
- id: trailing-whitespace
|
|
12
|
+
- id: check-yaml
|
|
13
|
+
- repo: https://github.com/pre-commit/mirrors-mypy
|
|
14
|
+
rev: v1.15.0
|
|
15
|
+
hooks:
|
|
16
|
+
- id: mypy
|
|
17
|
+
additional_dependencies: [types-setuptools]
|
|
18
|
+
args: [--config-file=pyproject.toml]
|
|
19
|
+
files: ^src/
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
3.11
|
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
# Changelog
|
|
2
|
+
|
|
3
|
+
All notable changes to this project are documented in this file.
|
|
4
|
+
|
|
5
|
+
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
|
|
6
|
+
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
|
7
|
+
|
|
8
|
+
## [Unreleased]
|
|
9
|
+
|
|
10
|
+
## [0.2.1] - 2026-06-21
|
|
11
|
+
|
|
12
|
+
### Changed
|
|
13
|
+
|
|
14
|
+
- CI now runs on pull requests to `main` only.
|
|
15
|
+
- CD runs on merge to `main` and publishes to PyPI when the version is bumped.
|
|
16
|
+
- Release automation creates git tags and GitHub releases on version bump.
|
|
17
|
+
|
|
18
|
+
## [0.2.0] - 2026-06-20
|
|
19
|
+
|
|
20
|
+
### Added
|
|
21
|
+
|
|
22
|
+
- OOP core abstractions: `EnergyModel`, `BaseDensityModel`, and `DensityModel`.
|
|
23
|
+
- Class-based training API via `boltzmann_generators.training.Trainer`.
|
|
24
|
+
- Loss strategy classes: `ForwardKLLoss`, `ReverseKLLoss`, `MixedLossStrategy`.
|
|
25
|
+
- Service classes: `SamplingEngine`, `AnalysisSuite`, `CheckpointManager`.
|
|
26
|
+
- PEP 561 typing marker (`py.typed`) and mypy configuration.
|
|
27
|
+
- Coverage gate (`>=75%`) in pytest/CI.
|
|
28
|
+
- PyPI release workflow with changelog validation and trusted publishing.
|
|
29
|
+
- Example notebooks under `examples/notebooks/`.
|
|
30
|
+
- Release checklist in `RELEASE.md` and migration guide in `MIGRATION.md`.
|
|
31
|
+
|
|
32
|
+
### Changed
|
|
33
|
+
|
|
34
|
+
- `FlowModel` and `CNFFlowModel` now inherit from `BaseDensityModel`.
|
|
35
|
+
- Energy classes now inherit from `EnergyModel`.
|
|
36
|
+
- `train_bg()` is deprecated in favor of `Trainer.fit()`.
|
|
37
|
+
- Package version bumped to `0.2.0`.
|
|
38
|
+
|
|
39
|
+
### Fixed
|
|
40
|
+
|
|
41
|
+
- Removed stale `src/bg` wheel package reference that could break builds.
|
|
42
|
+
- README quickstart no longer calls `.to(device)` on non-module energy objects.
|
|
43
|
+
|
|
44
|
+
## [0.1.0] - 2026-06-16
|
|
45
|
+
|
|
46
|
+
### Added
|
|
47
|
+
|
|
48
|
+
- Initial open-source packaging of the Boltzmann Generators PyTorch implementation.
|
|
49
|
+
- Split package layout with `boltzmann_generators` as the primary import path.
|
|
50
|
+
- OSS project metadata, CI, contribution docs, and testing scaffold.
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
cff-version: 1.2.0
|
|
2
|
+
title: boltzmann-generators
|
|
3
|
+
message: "If you use this software, please cite this repository and the original papers."
|
|
4
|
+
type: software
|
|
5
|
+
authors:
|
|
6
|
+
- family-names: Martin
|
|
7
|
+
given-names: Manu
|
|
8
|
+
repository-code: "https://github.com/manumartinm/boltzmann-generators"
|
|
9
|
+
license: MIT
|
|
10
|
+
version: 0.2.1
|
|
11
|
+
date-released: 2026-06-21
|
|
12
|
+
preferred-citation:
|
|
13
|
+
type: article
|
|
14
|
+
authors:
|
|
15
|
+
- family-names: Noe
|
|
16
|
+
given-names: Frank
|
|
17
|
+
- family-names: Olsson
|
|
18
|
+
given-names: Simon
|
|
19
|
+
- family-names: Kohler
|
|
20
|
+
given-names: Jonas
|
|
21
|
+
- family-names: Wu
|
|
22
|
+
given-names: Hao
|
|
23
|
+
title: "Boltzmann Generators: Sampling Equilibrium States of Many-Body Systems with Deep Learning"
|
|
24
|
+
journal: Science
|
|
25
|
+
year: 2019
|
|
26
|
+
volume: 365
|
|
27
|
+
issue: 6457
|
|
28
|
+
doi: 10.1126/science.aaw1147
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
# Contributing
|
|
2
|
+
|
|
3
|
+
Thanks for your interest in contributing to `boltzmann-generators`.
|
|
4
|
+
|
|
5
|
+
## Development setup
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
uv sync --extra dev --extra notebook
|
|
9
|
+
pre-commit install
|
|
10
|
+
```
|
|
11
|
+
|
|
12
|
+
## Running checks
|
|
13
|
+
|
|
14
|
+
```bash
|
|
15
|
+
uv run ruff check .
|
|
16
|
+
uv run ruff format --check .
|
|
17
|
+
uv run mypy src
|
|
18
|
+
uv run pytest
|
|
19
|
+
uv run python -m build
|
|
20
|
+
uv run twine check dist/*
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
## Pull request guidelines
|
|
24
|
+
|
|
25
|
+
- Open PRs against `main`; CI runs automatically on each PR.
|
|
26
|
+
- To release: bump version in `pyproject.toml` + `__init__.py`, update `CHANGELOG.md`, merge — CD publishes on merge.
|
|
27
|
+
- Keep changes focused and scoped to a clear problem.
|
|
28
|
+
- Add or update tests for behavior changes.
|
|
29
|
+
- Update docs when public APIs change.
|
|
30
|
+
- Add a `CHANGELOG.md` entry for user-visible changes.
|
|
31
|
+
- Keep notebooks reproducible and deterministic where possible.
|
|
32
|
+
|
|
33
|
+
## Code style
|
|
34
|
+
|
|
35
|
+
- Python 3.11+.
|
|
36
|
+
- Ruff handles linting and import sorting.
|
|
37
|
+
- Prefer explicit, typed APIs for public modules.
|
|
38
|
+
- Use OOP service classes (`Trainer`, `SamplingEngine`, etc.) for new features.
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Manu Martin
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
# Migration Guide (0.1.x → 0.2.x)
|
|
2
|
+
|
|
3
|
+
Version 0.2.0 introduces a class-based OOP API while keeping backward-compatible function wrappers.
|
|
4
|
+
|
|
5
|
+
## Training
|
|
6
|
+
|
|
7
|
+
**Before (0.1.x):**
|
|
8
|
+
|
|
9
|
+
```python
|
|
10
|
+
from boltzmann_generators.train import TrainConfig, train_bg
|
|
11
|
+
|
|
12
|
+
history = train_bg(model, energy, x_data=x_data, config=config, device=device)
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
**After (0.2.x, recommended):**
|
|
16
|
+
|
|
17
|
+
```python
|
|
18
|
+
from boltzmann_generators.training import TrainConfig, Trainer
|
|
19
|
+
|
|
20
|
+
trainer = Trainer(model, energy, config, device=device)
|
|
21
|
+
history = trainer.fit(x_data)
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
`train_bg()` still works but emits a `DeprecationWarning`.
|
|
25
|
+
|
|
26
|
+
## Sampling and analysis
|
|
27
|
+
|
|
28
|
+
**Before:**
|
|
29
|
+
|
|
30
|
+
```python
|
|
31
|
+
from boltzmann_generators.sampling import sample_with_weights, effective_sample_size
|
|
32
|
+
from boltzmann_generators.analysis import basin_populations
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
**After (recommended service API):**
|
|
36
|
+
|
|
37
|
+
```python
|
|
38
|
+
from boltzmann_generators.services import SamplingEngine, AnalysisSuite
|
|
39
|
+
|
|
40
|
+
engine = SamplingEngine(model, energy)
|
|
41
|
+
x, log_w, log_q = engine.sample_with_weights(1024, device="cpu")
|
|
42
|
+
ess = engine.effective_sample_size(log_w)
|
|
43
|
+
|
|
44
|
+
suite = AnalysisSuite(engine)
|
|
45
|
+
pops = suite.basin_populations(x, {"left": left_fn, "right": right_fn}, log_w=log_w)
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
Function imports in `boltzmann_generators.sampling` and `boltzmann_generators.analysis` remain available.
|
|
49
|
+
|
|
50
|
+
## Checkpoints
|
|
51
|
+
|
|
52
|
+
**Before:**
|
|
53
|
+
|
|
54
|
+
```python
|
|
55
|
+
from boltzmann_generators.io import save_checkpoint, load_checkpoint
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
**After (recommended):**
|
|
59
|
+
|
|
60
|
+
```python
|
|
61
|
+
from boltzmann_generators.services import CheckpointManager
|
|
62
|
+
|
|
63
|
+
mgr = CheckpointManager()
|
|
64
|
+
mgr.save(model, "model.pt", config={"lr": 1e-3})
|
|
65
|
+
model, meta = mgr.load("model.pt", model)
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
## Type interfaces
|
|
69
|
+
|
|
70
|
+
- Energies should subclass or conform to `boltzmann_generators.base.EnergyModel`.
|
|
71
|
+
- Density models should subclass `boltzmann_generators.base.BaseDensityModel`.
|
|
72
|
+
- Loss functions accept any object matching the `DensityModel` protocol.
|
|
73
|
+
|
|
74
|
+
## Example notebooks
|
|
75
|
+
|
|
76
|
+
See runnable examples in [`examples/notebooks/`](examples/notebooks/).
|
|
@@ -0,0 +1,146 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: boltzmann-generators
|
|
3
|
+
Version: 0.2.1
|
|
4
|
+
Summary: Boltzmann Generators in PyTorch: flows, energies, losses, and sampling utilities.
|
|
5
|
+
Project-URL: Homepage, https://github.com/manumartinm/boltzmann-generators
|
|
6
|
+
Project-URL: Repository, https://github.com/manumartinm/boltzmann-generators
|
|
7
|
+
Project-URL: Issues, https://github.com/manumartinm/boltzmann-generators/issues
|
|
8
|
+
Project-URL: Changelog, https://github.com/manumartinm/boltzmann-generators/blob/main/CHANGELOG.md
|
|
9
|
+
Author: Manu Martin
|
|
10
|
+
License-Expression: MIT
|
|
11
|
+
License-File: LICENSE
|
|
12
|
+
Keywords: boltzmann-generators,molecular-dynamics,normalizing-flows,pytorch
|
|
13
|
+
Classifier: Development Status :: 3 - Alpha
|
|
14
|
+
Classifier: Intended Audience :: Science/Research
|
|
15
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
16
|
+
Classifier: Programming Language :: Python :: 3
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Classifier: Topic :: Scientific/Engineering :: Physics
|
|
21
|
+
Classifier: Typing :: Typed
|
|
22
|
+
Requires-Python: >=3.11
|
|
23
|
+
Requires-Dist: numpy>=2.0
|
|
24
|
+
Requires-Dist: torch>=2.0
|
|
25
|
+
Requires-Dist: tqdm>=4.60
|
|
26
|
+
Provides-Extra: dev
|
|
27
|
+
Requires-Dist: build>=1.2; extra == 'dev'
|
|
28
|
+
Requires-Dist: mypy>=1.10; extra == 'dev'
|
|
29
|
+
Requires-Dist: pre-commit>=3.7; extra == 'dev'
|
|
30
|
+
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
|
|
31
|
+
Requires-Dist: pytest>=8.0; extra == 'dev'
|
|
32
|
+
Requires-Dist: ruff>=0.5; extra == 'dev'
|
|
33
|
+
Requires-Dist: twine>=5.0; extra == 'dev'
|
|
34
|
+
Provides-Extra: notebook
|
|
35
|
+
Requires-Dist: jupyterlab>=4.0; extra == 'notebook'
|
|
36
|
+
Requires-Dist: matplotlib>=3.8; extra == 'notebook'
|
|
37
|
+
Requires-Dist: nbconvert>=7.0; extra == 'notebook'
|
|
38
|
+
Requires-Dist: pandas>=2.0; extra == 'notebook'
|
|
39
|
+
Provides-Extra: ode
|
|
40
|
+
Requires-Dist: torchdiffeq>=0.2; extra == 'ode'
|
|
41
|
+
Description-Content-Type: text/markdown
|
|
42
|
+
|
|
43
|
+
# boltzmann-generators
|
|
44
|
+
|
|
45
|
+
[](https://github.com/manumartinm/boltzmann-generators/actions/workflows/ci.yml)
|
|
46
|
+
[](https://pypi.org/project/boltzmann-generators/)
|
|
47
|
+
[](https://www.python.org/downloads/)
|
|
48
|
+
[](LICENSE)
|
|
49
|
+
|
|
50
|
+
PyTorch tools for Boltzmann Generators: invertible flows, benchmark energy functions, BG losses, and reweighting diagnostics (ESS and free-energy differences). The project is designed for reproducible experiments on toy molecular systems and as a base for transferable BG research.
|
|
51
|
+
|
|
52
|
+
## References
|
|
53
|
+
|
|
54
|
+
1. Noé, Olsson, Köhler, Wu, *Boltzmann Generators: Sampling Equilibrium States of Many-Body Systems with Deep Learning* (Science, 2019): <https://arxiv.org/abs/1812.01729>
|
|
55
|
+
2. Klein, Noé, *Transferable Boltzmann Generators* (NeurIPS, 2024): <https://openreview.net/forum?id=AYq6GxxrrY>
|
|
56
|
+
3. Original reference implementation: <https://github.com/noegroup/paper_boltzmann_generators>
|
|
57
|
+
|
|
58
|
+
## Installation
|
|
59
|
+
|
|
60
|
+
```bash
|
|
61
|
+
pip install boltzmann-generators
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
For notebooks and plots:
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
pip install "boltzmann-generators[notebook]"
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
For development:
|
|
71
|
+
|
|
72
|
+
```bash
|
|
73
|
+
pip install "boltzmann-generators[dev]"
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
For higher-accuracy CNF ODE integration:
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
pip install "boltzmann-generators[ode]"
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
## Quickstart
|
|
83
|
+
|
|
84
|
+
```python
|
|
85
|
+
import torch
|
|
86
|
+
from boltzmann_generators.energies import DoubleWell2D
|
|
87
|
+
from boltzmann_generators.flows import FlowModel, GaussianPrior, RealNVP
|
|
88
|
+
from boltzmann_generators.training import TrainConfig, Trainer
|
|
89
|
+
|
|
90
|
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
91
|
+
|
|
92
|
+
energy = DoubleWell2D()
|
|
93
|
+
prior = GaussianPrior(dim=2).to(device)
|
|
94
|
+
flow = RealNVP(dim=2, num_layers=8, hidden_dim=64, mask="halves").to(device)
|
|
95
|
+
model = FlowModel(prior, flow).to(device)
|
|
96
|
+
|
|
97
|
+
config = TrainConfig(n_epochs=200, batch_size=256, lr=1e-3, w_ml=0.0, w_kl=1.0)
|
|
98
|
+
trainer = Trainer(model, energy, config, device=device)
|
|
99
|
+
history = trainer.fit(x_data=None)
|
|
100
|
+
print("Final loss:", history[-1]["loss"])
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
## Project layout
|
|
104
|
+
|
|
105
|
+
```text
|
|
106
|
+
src/boltzmann_generators/
|
|
107
|
+
base/ # EnergyModel and BaseDensityModel ABCs
|
|
108
|
+
training/ # Trainer and loss strategy classes
|
|
109
|
+
services/ # SamplingEngine, AnalysisSuite, CheckpointManager
|
|
110
|
+
energies/ # double-well, Muller-Brown, Ramachandran-like energy
|
|
111
|
+
flows/ # base flow API, RealNVP, CNF
|
|
112
|
+
losses.py # KL_x, KL_z, mixed loss (function wrappers)
|
|
113
|
+
sampling.py # weighted samples, ESS, free-energy difference
|
|
114
|
+
mcmc.py # Metropolis baseline sampler
|
|
115
|
+
train.py # deprecated train_bg wrapper
|
|
116
|
+
analysis.py # basin and population utilities
|
|
117
|
+
|
|
118
|
+
examples/notebooks/ # curated runnable examples
|
|
119
|
+
notebooks/ # end-to-end reproductions (02-07)
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
## Reproduced results
|
|
123
|
+
|
|
124
|
+
| Notebook | System | Metric | Value |
|
|
125
|
+
|---|---|---|---|
|
|
126
|
+
| 02 | 4-mode Gaussian mixture | Final NLL | 2.23 nats (target entropy: 2.13) |
|
|
127
|
+
| 03 | Double-well 2D (4 kT barrier) | ESS / DeltaF | 98% / +0.004 kT (target 0) |
|
|
128
|
+
| 04 | Muller-Brown (3 basins) | ESS / populations | 96% / target [80.37, 7.74, 11.89]%, BG [80.23, 7.71, 12.05]% |
|
|
129
|
+
| 05 | Double-well paper comparison | ESS | 92% |
|
|
130
|
+
| 06 | Gaussian mixture (CFM) | Visual mode coverage | All 4 modes recovered |
|
|
131
|
+
| 07 | Synthetic Ramachandran | Visual + basin populations | Reweighted BG close to ground truth |
|
|
132
|
+
|
|
133
|
+
## Documentation
|
|
134
|
+
|
|
135
|
+
- Example notebooks: [`examples/notebooks/`](examples/notebooks/)
|
|
136
|
+
- Migration guide (0.1 → 0.2): [`MIGRATION.md`](MIGRATION.md)
|
|
137
|
+
- Changelog: [`CHANGELOG.md`](CHANGELOG.md)
|
|
138
|
+
- Release process: [`RELEASE.md`](RELEASE.md)
|
|
139
|
+
- Contribution guide: [`CONTRIBUTING.md`](CONTRIBUTING.md)
|
|
140
|
+
- Research notebooks: `notebooks/`
|
|
141
|
+
|
|
142
|
+
If you need the original Noé group reference repository locally, clone it into:
|
|
143
|
+
|
|
144
|
+
```bash
|
|
145
|
+
git clone --depth 1 https://github.com/noegroup/paper_boltzmann_generators.git references/paper_boltzmann_generators
|
|
146
|
+
```
|
|
@@ -0,0 +1,104 @@
|
|
|
1
|
+
# boltzmann-generators
|
|
2
|
+
|
|
3
|
+
[](https://github.com/manumartinm/boltzmann-generators/actions/workflows/ci.yml)
|
|
4
|
+
[](https://pypi.org/project/boltzmann-generators/)
|
|
5
|
+
[](https://www.python.org/downloads/)
|
|
6
|
+
[](LICENSE)
|
|
7
|
+
|
|
8
|
+
PyTorch tools for Boltzmann Generators: invertible flows, benchmark energy functions, BG losses, and reweighting diagnostics (ESS and free-energy differences). The project is designed for reproducible experiments on toy molecular systems and as a base for transferable BG research.
|
|
9
|
+
|
|
10
|
+
## References
|
|
11
|
+
|
|
12
|
+
1. Noé, Olsson, Köhler, Wu, *Boltzmann Generators: Sampling Equilibrium States of Many-Body Systems with Deep Learning* (Science, 2019): <https://arxiv.org/abs/1812.01729>
|
|
13
|
+
2. Klein, Noé, *Transferable Boltzmann Generators* (NeurIPS, 2024): <https://openreview.net/forum?id=AYq6GxxrrY>
|
|
14
|
+
3. Original reference implementation: <https://github.com/noegroup/paper_boltzmann_generators>
|
|
15
|
+
|
|
16
|
+
## Installation
|
|
17
|
+
|
|
18
|
+
```bash
|
|
19
|
+
pip install boltzmann-generators
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
For notebooks and plots:
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
pip install "boltzmann-generators[notebook]"
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
For development:
|
|
29
|
+
|
|
30
|
+
```bash
|
|
31
|
+
pip install "boltzmann-generators[dev]"
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
For higher-accuracy CNF ODE integration:
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
pip install "boltzmann-generators[ode]"
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
## Quickstart
|
|
41
|
+
|
|
42
|
+
```python
|
|
43
|
+
import torch
|
|
44
|
+
from boltzmann_generators.energies import DoubleWell2D
|
|
45
|
+
from boltzmann_generators.flows import FlowModel, GaussianPrior, RealNVP
|
|
46
|
+
from boltzmann_generators.training import TrainConfig, Trainer
|
|
47
|
+
|
|
48
|
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
49
|
+
|
|
50
|
+
energy = DoubleWell2D()
|
|
51
|
+
prior = GaussianPrior(dim=2).to(device)
|
|
52
|
+
flow = RealNVP(dim=2, num_layers=8, hidden_dim=64, mask="halves").to(device)
|
|
53
|
+
model = FlowModel(prior, flow).to(device)
|
|
54
|
+
|
|
55
|
+
config = TrainConfig(n_epochs=200, batch_size=256, lr=1e-3, w_ml=0.0, w_kl=1.0)
|
|
56
|
+
trainer = Trainer(model, energy, config, device=device)
|
|
57
|
+
history = trainer.fit(x_data=None)
|
|
58
|
+
print("Final loss:", history[-1]["loss"])
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
## Project layout
|
|
62
|
+
|
|
63
|
+
```text
|
|
64
|
+
src/boltzmann_generators/
|
|
65
|
+
base/ # EnergyModel and BaseDensityModel ABCs
|
|
66
|
+
training/ # Trainer and loss strategy classes
|
|
67
|
+
services/ # SamplingEngine, AnalysisSuite, CheckpointManager
|
|
68
|
+
energies/ # double-well, Muller-Brown, Ramachandran-like energy
|
|
69
|
+
flows/ # base flow API, RealNVP, CNF
|
|
70
|
+
losses.py # KL_x, KL_z, mixed loss (function wrappers)
|
|
71
|
+
sampling.py # weighted samples, ESS, free-energy difference
|
|
72
|
+
mcmc.py # Metropolis baseline sampler
|
|
73
|
+
train.py # deprecated train_bg wrapper
|
|
74
|
+
analysis.py # basin and population utilities
|
|
75
|
+
|
|
76
|
+
examples/notebooks/ # curated runnable examples
|
|
77
|
+
notebooks/ # end-to-end reproductions (02-07)
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
## Reproduced results
|
|
81
|
+
|
|
82
|
+
| Notebook | System | Metric | Value |
|
|
83
|
+
|---|---|---|---|
|
|
84
|
+
| 02 | 4-mode Gaussian mixture | Final NLL | 2.23 nats (target entropy: 2.13) |
|
|
85
|
+
| 03 | Double-well 2D (4 kT barrier) | ESS / DeltaF | 98% / +0.004 kT (target 0) |
|
|
86
|
+
| 04 | Muller-Brown (3 basins) | ESS / populations | 96% / target [80.37, 7.74, 11.89]%, BG [80.23, 7.71, 12.05]% |
|
|
87
|
+
| 05 | Double-well paper comparison | ESS | 92% |
|
|
88
|
+
| 06 | Gaussian mixture (CFM) | Visual mode coverage | All 4 modes recovered |
|
|
89
|
+
| 07 | Synthetic Ramachandran | Visual + basin populations | Reweighted BG close to ground truth |
|
|
90
|
+
|
|
91
|
+
## Documentation
|
|
92
|
+
|
|
93
|
+
- Example notebooks: [`examples/notebooks/`](examples/notebooks/)
|
|
94
|
+
- Migration guide (0.1 → 0.2): [`MIGRATION.md`](MIGRATION.md)
|
|
95
|
+
- Changelog: [`CHANGELOG.md`](CHANGELOG.md)
|
|
96
|
+
- Release process: [`RELEASE.md`](RELEASE.md)
|
|
97
|
+
- Contribution guide: [`CONTRIBUTING.md`](CONTRIBUTING.md)
|
|
98
|
+
- Research notebooks: `notebooks/`
|
|
99
|
+
|
|
100
|
+
If you need the original Noé group reference repository locally, clone it into:
|
|
101
|
+
|
|
102
|
+
```bash
|
|
103
|
+
git clone --depth 1 https://github.com/noegroup/paper_boltzmann_generators.git references/paper_boltzmann_generators
|
|
104
|
+
```
|