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.
Files changed (67) hide show
  1. boltzmann_generators-0.2.1/.coverage +0 -0
  2. boltzmann_generators-0.2.1/.github/workflows/cd.yml +107 -0
  3. boltzmann_generators-0.2.1/.github/workflows/ci.yml +35 -0
  4. boltzmann_generators-0.2.1/.gitignore +27 -0
  5. boltzmann_generators-0.2.1/.pre-commit-config.yaml +19 -0
  6. boltzmann_generators-0.2.1/.python-version +1 -0
  7. boltzmann_generators-0.2.1/CHANGELOG.md +50 -0
  8. boltzmann_generators-0.2.1/CITATION.cff +28 -0
  9. boltzmann_generators-0.2.1/CONTRIBUTING.md +38 -0
  10. boltzmann_generators-0.2.1/LICENSE +21 -0
  11. boltzmann_generators-0.2.1/MIGRATION.md +76 -0
  12. boltzmann_generators-0.2.1/PKG-INFO +146 -0
  13. boltzmann_generators-0.2.1/README.md +104 -0
  14. boltzmann_generators-0.2.1/RELEASE.md +55 -0
  15. boltzmann_generators-0.2.1/examples/01_quickstart_realnvp.ipynb +52 -0
  16. boltzmann_generators-0.2.1/examples/02_reweighting_and_ess.ipynb +56 -0
  17. boltzmann_generators-0.2.1/examples/03_cnf_flow_matching.ipynb +54 -0
  18. boltzmann_generators-0.2.1/examples/04_checkpoint_and_resume.ipynb +62 -0
  19. boltzmann_generators-0.2.1/pyproject.toml +109 -0
  20. boltzmann_generators-0.2.1/src/boltzmann_generators/__init__.py +18 -0
  21. boltzmann_generators-0.2.1/src/boltzmann_generators/analysis.py +36 -0
  22. boltzmann_generators-0.2.1/src/boltzmann_generators/base/__init__.py +6 -0
  23. boltzmann_generators-0.2.1/src/boltzmann_generators/base/density.py +36 -0
  24. boltzmann_generators-0.2.1/src/boltzmann_generators/base/energy.py +20 -0
  25. boltzmann_generators-0.2.1/src/boltzmann_generators/energies/__init__.py +5 -0
  26. boltzmann_generators-0.2.1/src/boltzmann_generators/energies/dipeptide.py +81 -0
  27. boltzmann_generators-0.2.1/src/boltzmann_generators/energies/double_well.py +56 -0
  28. boltzmann_generators-0.2.1/src/boltzmann_generators/energies/muller.py +65 -0
  29. boltzmann_generators-0.2.1/src/boltzmann_generators/extensions/__init__.py +1 -0
  30. boltzmann_generators-0.2.1/src/boltzmann_generators/extensions/equivariant.py +7 -0
  31. boltzmann_generators-0.2.1/src/boltzmann_generators/extensions/molecular.py +7 -0
  32. boltzmann_generators-0.2.1/src/boltzmann_generators/extensions/transferable.py +7 -0
  33. boltzmann_generators-0.2.1/src/boltzmann_generators/flows/__init__.py +20 -0
  34. boltzmann_generators-0.2.1/src/boltzmann_generators/flows/base.py +68 -0
  35. boltzmann_generators-0.2.1/src/boltzmann_generators/flows/cnf.py +266 -0
  36. boltzmann_generators-0.2.1/src/boltzmann_generators/flows/coupling.py +89 -0
  37. boltzmann_generators-0.2.1/src/boltzmann_generators/flows/periodic.py +27 -0
  38. boltzmann_generators-0.2.1/src/boltzmann_generators/flows/realnvp.py +79 -0
  39. boltzmann_generators-0.2.1/src/boltzmann_generators/io.py +34 -0
  40. boltzmann_generators-0.2.1/src/boltzmann_generators/losses.py +70 -0
  41. boltzmann_generators-0.2.1/src/boltzmann_generators/mcmc.py +68 -0
  42. boltzmann_generators-0.2.1/src/boltzmann_generators/py.typed +0 -0
  43. boltzmann_generators-0.2.1/src/boltzmann_generators/sampling.py +54 -0
  44. boltzmann_generators-0.2.1/src/boltzmann_generators/services/__init__.py +7 -0
  45. boltzmann_generators-0.2.1/src/boltzmann_generators/services/analysis.py +65 -0
  46. boltzmann_generators-0.2.1/src/boltzmann_generators/services/checkpoint.py +43 -0
  47. boltzmann_generators-0.2.1/src/boltzmann_generators/services/sampling.py +74 -0
  48. boltzmann_generators-0.2.1/src/boltzmann_generators/train.py +40 -0
  49. boltzmann_generators-0.2.1/src/boltzmann_generators/training/__init__.py +19 -0
  50. boltzmann_generators-0.2.1/src/boltzmann_generators/training/loss_strategies.py +113 -0
  51. boltzmann_generators-0.2.1/src/boltzmann_generators/training/trainer.py +94 -0
  52. boltzmann_generators-0.2.1/tests/conftest.py +22 -0
  53. boltzmann_generators-0.2.1/tests/test_analysis.py +41 -0
  54. boltzmann_generators-0.2.1/tests/test_base.py +40 -0
  55. boltzmann_generators-0.2.1/tests/test_cnf.py +24 -0
  56. boltzmann_generators-0.2.1/tests/test_energies.py +33 -0
  57. boltzmann_generators-0.2.1/tests/test_integration.py +38 -0
  58. boltzmann_generators-0.2.1/tests/test_io.py +42 -0
  59. boltzmann_generators-0.2.1/tests/test_losses.py +38 -0
  60. boltzmann_generators-0.2.1/tests/test_mcmc.py +15 -0
  61. boltzmann_generators-0.2.1/tests/test_periodic.py +21 -0
  62. boltzmann_generators-0.2.1/tests/test_realnvp.py +21 -0
  63. boltzmann_generators-0.2.1/tests/test_sampling.py +27 -0
  64. boltzmann_generators-0.2.1/tests/test_services.py +101 -0
  65. boltzmann_generators-0.2.1/tests/test_train.py +20 -0
  66. boltzmann_generators-0.2.1/tests/test_trainer.py +47 -0
  67. 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
+ [![CI](https://github.com/manumartinm/boltzmann-generators/actions/workflows/ci.yml/badge.svg)](https://github.com/manumartinm/boltzmann-generators/actions/workflows/ci.yml)
46
+ [![PyPI version](https://img.shields.io/pypi/v/boltzmann-generators)](https://pypi.org/project/boltzmann-generators/)
47
+ [![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
48
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](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
+ [![CI](https://github.com/manumartinm/boltzmann-generators/actions/workflows/ci.yml/badge.svg)](https://github.com/manumartinm/boltzmann-generators/actions/workflows/ci.yml)
4
+ [![PyPI version](https://img.shields.io/pypi/v/boltzmann-generators)](https://pypi.org/project/boltzmann-generators/)
5
+ [![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
6
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](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
+ ```