off 0.1.0__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.
- off-0.1.0/LICENSE +21 -0
- off-0.1.0/PKG-INFO +154 -0
- off-0.1.0/README.md +123 -0
- off-0.1.0/off/__init__.py +23 -0
- off-0.1.0/off/atom_energies.py +151 -0
- off-0.1.0/off/config/_config.py +108 -0
- off-0.1.0/off/dft_distrax/__init__.py +27 -0
- off-0.1.0/off/dft_distrax/dft_distrax.py +216 -0
- off-0.1.0/off/flow/__init__.py +29 -0
- off-0.1.0/off/flow/equiv_flows.py +99 -0
- off-0.1.0/off/functionals/__init__.py +35 -0
- off-0.1.0/off/functionals/core_correction.py +84 -0
- off-0.1.0/off/functionals/exchange_correlation.py +174 -0
- off-0.1.0/off/functionals/external.py +49 -0
- off-0.1.0/off/functionals/functional.py +129 -0
- off-0.1.0/off/functionals/hartree.py +62 -0
- off-0.1.0/off/functionals/kinetic.py +87 -0
- off-0.1.0/off/main.py +172 -0
- off-0.1.0/off/ode_solver/__init__.py +32 -0
- off-0.1.0/off/ode_solver/eqx_ode.py +76 -0
- off-0.1.0/off/plot_binding_csv.py +63 -0
- off-0.1.0/off/plot_pes_ema.py +259 -0
- off-0.1.0/off/plot_pes_mpl.py +280 -0
- off-0.1.0/off/promolecular/__init__.py +27 -0
- off-0.1.0/off/promolecular/promolecular_dist.py +465 -0
- off-0.1.0/off/quadrature.py +261 -0
- off-0.1.0/off/quadrature_scan.py +188 -0
- off-0.1.0/off/scan_pes.py +133 -0
- off-0.1.0/off/test_fwd_rev.py +290 -0
- off-0.1.0/off/train/__init__.py +44 -0
- off-0.1.0/off/train/loop.py +228 -0
- off-0.1.0/off/train/loss.py +149 -0
- off-0.1.0/off/train/utils.py +38 -0
- off-0.1.0/off/utils.py +618 -0
- off-0.1.0/off.egg-info/PKG-INFO +154 -0
- off-0.1.0/off.egg-info/SOURCES.txt +40 -0
- off-0.1.0/off.egg-info/dependency_links.txt +1 -0
- off-0.1.0/off.egg-info/entry_points.txt +3 -0
- off-0.1.0/off.egg-info/requires.txt +17 -0
- off-0.1.0/off.egg-info/top_level.txt +1 -0
- off-0.1.0/pyproject.toml +50 -0
- off-0.1.0/setup.cfg +4 -0
off-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 AlexandreDeCamargo
|
|
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.
|
off-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: off
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A library for learning the ground state energy using orbital free density functional theory with continuous normalizing flows.
|
|
5
|
+
Author-email: Alexandre de Camargo <decamara@mcmaster.ca>
|
|
6
|
+
Project-URL: Homepage, https://github.com/AlexandreDeCamargo/of_flows
|
|
7
|
+
Project-URL: Repository, https://github.com/AlexandreDeCamargo/of_flows
|
|
8
|
+
Classifier: Intended Audience :: Developers
|
|
9
|
+
Classifier: Intended Audience :: Science/Research
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
11
|
+
Requires-Python: >=3.9
|
|
12
|
+
Description-Content-Type: text/markdown
|
|
13
|
+
License-File: LICENSE
|
|
14
|
+
Requires-Dist: jax>=0.6.0
|
|
15
|
+
Requires-Dist: jaxlib>=0.6.0
|
|
16
|
+
Requires-Dist: equinox>=0.12.0
|
|
17
|
+
Requires-Dist: diffrax>=0.7.0
|
|
18
|
+
Requires-Dist: optax>=0.2.0
|
|
19
|
+
Requires-Dist: scipy>=1.15.0
|
|
20
|
+
Requires-Dist: chex>=0.1.89
|
|
21
|
+
Requires-Dist: jaxtyping>=0.3.0
|
|
22
|
+
Requires-Dist: distrax>=0.1.5
|
|
23
|
+
Requires-Dist: pyscf>=2.7.0
|
|
24
|
+
Requires-Dist: pandas>=2.0.3
|
|
25
|
+
Requires-Dist: matplotlib>=3.10.0
|
|
26
|
+
Provides-Extra: examples
|
|
27
|
+
Requires-Dist: matplotlib>=3.10.0; extra == "examples"
|
|
28
|
+
Requires-Dist: chex>=0.1.89; extra == "examples"
|
|
29
|
+
Requires-Dist: numpy>=2.3.0; extra == "examples"
|
|
30
|
+
Dynamic: license-file
|
|
31
|
+
|
|
32
|
+
# OFF: An Orbital-Free Density Functional Theory Python library using Normalizing Flows
|
|
33
|
+
|
|
34
|
+
OFF is a [JAX](https://github.com/google/jax)-based library for **orbital-free density
|
|
35
|
+
functional theory (OF-DFT)** in which the electron density is represented by a
|
|
36
|
+
**continuous normalizing flow (CNF)** and the ground-state energy is obtained by
|
|
37
|
+
*variationally minimizing* a density functional with Monte-Carlo gradient estimates.
|
|
38
|
+
The density is normalized by construction (it is a probability flow), and the physical
|
|
39
|
+
density is recovered as `ρ(x) = Ne · ρ_φ(x)`.
|
|
40
|
+
|
|
41
|
+
OFF is built entirely on the JAX ecosystem — automatic differentiation, JIT
|
|
42
|
+
compilation, vectorization, and GPU acceleration — with
|
|
43
|
+
[Diffrax](https://github.com/patrick-kidger/diffrax) for the flow ODEs,
|
|
44
|
+
[Equinox](https://github.com/patrick-kidger/equinox) for the models and functionals,
|
|
45
|
+
[Distrax](https://github.com/google-deepmind/distrax) for the base distribution, and
|
|
46
|
+
[Optax](https://github.com/google-deepmind/optax) for the optimization.
|
|
47
|
+
|
|
48
|
+
## Functionality
|
|
49
|
+
|
|
50
|
+
The current version of the library has the following capabilities:
|
|
51
|
+
|
|
52
|
+
* Provides an implementation of OF-DFT methods with continuous normalizing flows.
|
|
53
|
+
* A modular library of density functionals:
|
|
54
|
+
* **kinetic**: Thomas–Fermi, von Weizsäcker, and TF-λW;
|
|
55
|
+
* **exchange**: LDA and B88;
|
|
56
|
+
* **correlation**: VWN and PW92;
|
|
57
|
+
* **nuclear attraction**,
|
|
58
|
+
* **Hartree**, and
|
|
59
|
+
* **nuclear-cusp corrections** (Kato, Hutcheon).
|
|
60
|
+
* Evaluates density functionals using Monte Carlo estimators.
|
|
61
|
+
* In addition to Monte Carlo estimators, it also has a deterministic **grid (quadrature) readout** of the energy after training, using a
|
|
62
|
+
[PySCF](https://github.com/pyscf/pyscf) Becke grid converted to `jax.Array`.
|
|
63
|
+
|
|
64
|
+
## Install
|
|
65
|
+
|
|
66
|
+
A core dependency is [PySCF](https://pyscf.org), which needs `cmake` available on the
|
|
67
|
+
`PATH`. In a fresh environment, from the repository root:
|
|
68
|
+
|
|
69
|
+
```bash
|
|
70
|
+
pip install -e .
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
The `db_sir` prior additionally requires AtomDB; install it only if you use that prior.
|
|
74
|
+
|
|
75
|
+
## Use
|
|
76
|
+
|
|
77
|
+
Two stages: (1) **train** a normalizing flow for a given molecule and functional, then
|
|
78
|
+
(2) read out the energy on a **quadrature grid**. Both are exposed as command-line
|
|
79
|
+
tools (after `pip install -e .`) and as a Python API.
|
|
80
|
+
|
|
81
|
+
### 1. Train a flow
|
|
82
|
+
|
|
83
|
+
```bash
|
|
84
|
+
off-train --mol_name H2 --bond_length 1.4 \
|
|
85
|
+
--kin tf_w --lam 1/5 --x lda_b88_x --c none \
|
|
86
|
+
--hart coulomb --prior promolecular \
|
|
87
|
+
--solver dopri8 --epochs 500 --bs 512
|
|
88
|
+
```
|
|
89
|
+
(equivalently `python -m off.main ...`). This minimizes the OF-DFT energy with the
|
|
90
|
+
Monte-Carlo estimator and writes everything under a method-tagged directory:
|
|
91
|
+
|
|
92
|
+
```
|
|
93
|
+
Results/H2/tf_w_lam0.2_none_lda_b88_x_none_dopri8_promolecular_sched_mix/bl_1.40/
|
|
94
|
+
Checkpoints/checkpoint_*.eqx # the trained flow
|
|
95
|
+
training_metrics_ema.csv # EMA energy trace
|
|
96
|
+
job_params.json # everything needed to rebuild the run
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
Key options: `--kin {tf,w,tf_w}`, `--x {lda,b88_x,lda_b88_x}`, `--c {vwn_c,pw92_c,none}`,
|
|
100
|
+
`--cc {kato,hutcheon,none}`, `--hart {coulomb,coulomb_}` (all-pairs / element-wise),
|
|
101
|
+
`--prior {promolecular,db_sir}`, `--solver {dopri5,tsit5,dopri8}`.
|
|
102
|
+
|
|
103
|
+
### 2. Evaluate the energy on a grid
|
|
104
|
+
|
|
105
|
+
After training, point the quadrature tool at the run directory. It rebuilds the flow
|
|
106
|
+
from `job_params.json`, builds a PySCF grid, evaluates `ρ_φ` and its score there, and
|
|
107
|
+
integrates every energy term:
|
|
108
|
+
|
|
109
|
+
```bash
|
|
110
|
+
off-quadrature Results/H2/tf_w_lam0.2_none_lda_b88_x_none_dopri8_promolecular_sched_mix/bl_1.40 --grid_level 3
|
|
111
|
+
```
|
|
112
|
+
(equivalently `python -m off.quadrature ...`). It prints the per-term energies
|
|
113
|
+
(`T, V_N, V_H, E_X, E_C, E_CC, E_NN, E_total`) and the `∫ρ` check, and caches the
|
|
114
|
+
result in `energy_summary.json`. The same call from Python:
|
|
115
|
+
|
|
116
|
+
```python
|
|
117
|
+
from off import grid_energy_from_checkpoint
|
|
118
|
+
|
|
119
|
+
e = grid_energy_from_checkpoint(
|
|
120
|
+
"Results/H2/.../bl_1.40", grid_level=3)
|
|
121
|
+
print(e["E_total"], e["Ne_integral"])
|
|
122
|
+
```
|
|
123
|
+
|
|
124
|
+
### Build a grid
|
|
125
|
+
|
|
126
|
+
```python
|
|
127
|
+
from off import getGrid
|
|
128
|
+
|
|
129
|
+
h2_geom = "H 0 0 0; H 0 0 1.4"
|
|
130
|
+
w_grid, x_grid = getGrid(h2_geom, level=3, units="bohr")
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
## Package layout
|
|
134
|
+
|
|
135
|
+
```
|
|
136
|
+
off/
|
|
137
|
+
main.py # training entry point (off-train)
|
|
138
|
+
quadrature.py # grid-energy readout (off-quadrature)
|
|
139
|
+
flow/ # the continuous normalizing flow (CNF)
|
|
140
|
+
ode_solver/ # Diffrax forward/reverse ODE solves
|
|
141
|
+
promolecular/ # base distributions (promolecular, AtomDB)
|
|
142
|
+
functionals/ # kinetic, exchange-correlation, nuclear, Hartree, cusp
|
|
143
|
+
train/ # loss (Monte-Carlo energy) and the optimization loop
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
## Citation
|
|
147
|
+
|
|
148
|
+
```bibtex
|
|
149
|
+
@article{off,
|
|
150
|
+
title = {Orbital-Free DFT with Normalizing Flows},
|
|
151
|
+
author = {de Camargo, Alexandre and others},
|
|
152
|
+
year = {2026},
|
|
153
|
+
}
|
|
154
|
+
```
|
off-0.1.0/README.md
ADDED
|
@@ -0,0 +1,123 @@
|
|
|
1
|
+
# OFF: An Orbital-Free Density Functional Theory Python library using Normalizing Flows
|
|
2
|
+
|
|
3
|
+
OFF is a [JAX](https://github.com/google/jax)-based library for **orbital-free density
|
|
4
|
+
functional theory (OF-DFT)** in which the electron density is represented by a
|
|
5
|
+
**continuous normalizing flow (CNF)** and the ground-state energy is obtained by
|
|
6
|
+
*variationally minimizing* a density functional with Monte-Carlo gradient estimates.
|
|
7
|
+
The density is normalized by construction (it is a probability flow), and the physical
|
|
8
|
+
density is recovered as `ρ(x) = Ne · ρ_φ(x)`.
|
|
9
|
+
|
|
10
|
+
OFF is built entirely on the JAX ecosystem — automatic differentiation, JIT
|
|
11
|
+
compilation, vectorization, and GPU acceleration — with
|
|
12
|
+
[Diffrax](https://github.com/patrick-kidger/diffrax) for the flow ODEs,
|
|
13
|
+
[Equinox](https://github.com/patrick-kidger/equinox) for the models and functionals,
|
|
14
|
+
[Distrax](https://github.com/google-deepmind/distrax) for the base distribution, and
|
|
15
|
+
[Optax](https://github.com/google-deepmind/optax) for the optimization.
|
|
16
|
+
|
|
17
|
+
## Functionality
|
|
18
|
+
|
|
19
|
+
The current version of the library has the following capabilities:
|
|
20
|
+
|
|
21
|
+
* Provides an implementation of OF-DFT methods with continuous normalizing flows.
|
|
22
|
+
* A modular library of density functionals:
|
|
23
|
+
* **kinetic**: Thomas–Fermi, von Weizsäcker, and TF-λW;
|
|
24
|
+
* **exchange**: LDA and B88;
|
|
25
|
+
* **correlation**: VWN and PW92;
|
|
26
|
+
* **nuclear attraction**,
|
|
27
|
+
* **Hartree**, and
|
|
28
|
+
* **nuclear-cusp corrections** (Kato, Hutcheon).
|
|
29
|
+
* Evaluates density functionals using Monte Carlo estimators.
|
|
30
|
+
* In addition to Monte Carlo estimators, it also has a deterministic **grid (quadrature) readout** of the energy after training, using a
|
|
31
|
+
[PySCF](https://github.com/pyscf/pyscf) Becke grid converted to `jax.Array`.
|
|
32
|
+
|
|
33
|
+
## Install
|
|
34
|
+
|
|
35
|
+
A core dependency is [PySCF](https://pyscf.org), which needs `cmake` available on the
|
|
36
|
+
`PATH`. In a fresh environment, from the repository root:
|
|
37
|
+
|
|
38
|
+
```bash
|
|
39
|
+
pip install -e .
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
The `db_sir` prior additionally requires AtomDB; install it only if you use that prior.
|
|
43
|
+
|
|
44
|
+
## Use
|
|
45
|
+
|
|
46
|
+
Two stages: (1) **train** a normalizing flow for a given molecule and functional, then
|
|
47
|
+
(2) read out the energy on a **quadrature grid**. Both are exposed as command-line
|
|
48
|
+
tools (after `pip install -e .`) and as a Python API.
|
|
49
|
+
|
|
50
|
+
### 1. Train a flow
|
|
51
|
+
|
|
52
|
+
```bash
|
|
53
|
+
off-train --mol_name H2 --bond_length 1.4 \
|
|
54
|
+
--kin tf_w --lam 1/5 --x lda_b88_x --c none \
|
|
55
|
+
--hart coulomb --prior promolecular \
|
|
56
|
+
--solver dopri8 --epochs 500 --bs 512
|
|
57
|
+
```
|
|
58
|
+
(equivalently `python -m off.main ...`). This minimizes the OF-DFT energy with the
|
|
59
|
+
Monte-Carlo estimator and writes everything under a method-tagged directory:
|
|
60
|
+
|
|
61
|
+
```
|
|
62
|
+
Results/H2/tf_w_lam0.2_none_lda_b88_x_none_dopri8_promolecular_sched_mix/bl_1.40/
|
|
63
|
+
Checkpoints/checkpoint_*.eqx # the trained flow
|
|
64
|
+
training_metrics_ema.csv # EMA energy trace
|
|
65
|
+
job_params.json # everything needed to rebuild the run
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
Key options: `--kin {tf,w,tf_w}`, `--x {lda,b88_x,lda_b88_x}`, `--c {vwn_c,pw92_c,none}`,
|
|
69
|
+
`--cc {kato,hutcheon,none}`, `--hart {coulomb,coulomb_}` (all-pairs / element-wise),
|
|
70
|
+
`--prior {promolecular,db_sir}`, `--solver {dopri5,tsit5,dopri8}`.
|
|
71
|
+
|
|
72
|
+
### 2. Evaluate the energy on a grid
|
|
73
|
+
|
|
74
|
+
After training, point the quadrature tool at the run directory. It rebuilds the flow
|
|
75
|
+
from `job_params.json`, builds a PySCF grid, evaluates `ρ_φ` and its score there, and
|
|
76
|
+
integrates every energy term:
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
off-quadrature Results/H2/tf_w_lam0.2_none_lda_b88_x_none_dopri8_promolecular_sched_mix/bl_1.40 --grid_level 3
|
|
80
|
+
```
|
|
81
|
+
(equivalently `python -m off.quadrature ...`). It prints the per-term energies
|
|
82
|
+
(`T, V_N, V_H, E_X, E_C, E_CC, E_NN, E_total`) and the `∫ρ` check, and caches the
|
|
83
|
+
result in `energy_summary.json`. The same call from Python:
|
|
84
|
+
|
|
85
|
+
```python
|
|
86
|
+
from off import grid_energy_from_checkpoint
|
|
87
|
+
|
|
88
|
+
e = grid_energy_from_checkpoint(
|
|
89
|
+
"Results/H2/.../bl_1.40", grid_level=3)
|
|
90
|
+
print(e["E_total"], e["Ne_integral"])
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
### Build a grid
|
|
94
|
+
|
|
95
|
+
```python
|
|
96
|
+
from off import getGrid
|
|
97
|
+
|
|
98
|
+
h2_geom = "H 0 0 0; H 0 0 1.4"
|
|
99
|
+
w_grid, x_grid = getGrid(h2_geom, level=3, units="bohr")
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
## Package layout
|
|
103
|
+
|
|
104
|
+
```
|
|
105
|
+
off/
|
|
106
|
+
main.py # training entry point (off-train)
|
|
107
|
+
quadrature.py # grid-energy readout (off-quadrature)
|
|
108
|
+
flow/ # the continuous normalizing flow (CNF)
|
|
109
|
+
ode_solver/ # Diffrax forward/reverse ODE solves
|
|
110
|
+
promolecular/ # base distributions (promolecular, AtomDB)
|
|
111
|
+
functionals/ # kinetic, exchange-correlation, nuclear, Hartree, cusp
|
|
112
|
+
train/ # loss (Monte-Carlo energy) and the optimization loop
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
## Citation
|
|
116
|
+
|
|
117
|
+
```bibtex
|
|
118
|
+
@article{off,
|
|
119
|
+
title = {Orbital-Free DFT with Normalizing Flows},
|
|
120
|
+
author = {de Camargo, Alexandre and others},
|
|
121
|
+
year = {2026},
|
|
122
|
+
}
|
|
123
|
+
```
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
"""OFF — orbital-free DFT with continuous normalizing flows."""
|
|
2
|
+
|
|
3
|
+
__version__ = "0.1.0"
|
|
4
|
+
|
|
5
|
+
from .functionals.functional import (
|
|
6
|
+
Functional,
|
|
7
|
+
CompositeFunctional,
|
|
8
|
+
EnergyFunctional,
|
|
9
|
+
FunctionalInputs,
|
|
10
|
+
unit_coefficient,
|
|
11
|
+
)
|
|
12
|
+
from .functionals import (
|
|
13
|
+
tf, weizsacker, tf_weizsacker,
|
|
14
|
+
lda, b88, vwn, pw92, lda_b88,
|
|
15
|
+
NuclearPotential, CoulombPotential, CoulombPotential_,
|
|
16
|
+
KatoCondition, HutcheonCuspCondition,
|
|
17
|
+
)
|
|
18
|
+
from .quadrature import (
|
|
19
|
+
getGrid, get_grid, build_grid,
|
|
20
|
+
grid_energy, grid_energy_from_checkpoint,
|
|
21
|
+
)
|
|
22
|
+
from .train.loss import build_energy_functional, create_loss_function
|
|
23
|
+
from .train.loop import training
|
|
@@ -0,0 +1,151 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Grid-integrated total energy vs EMA energy, side by side, for a set of atoms.
|
|
3
|
+
|
|
4
|
+
For each atom it:
|
|
5
|
+
1. grid-integrates the LAST checkpoint via ``quadrature.grid_energy_from_checkpoint``
|
|
6
|
+
-> E_grid (= E_total; for a single atom E_NN = 0, so this is the electronic
|
|
7
|
+
total energy).
|
|
8
|
+
2. averages the last --window rows of training_metrics_ema.csv (E + CC) -> E_ema.
|
|
9
|
+
3. writes both numbers side by side to a CSV (+ prints a table).
|
|
10
|
+
|
|
11
|
+
Because E_NN = 0 for an isolated atom, E_grid and E_ema are the *same* physical
|
|
12
|
+
quantity (total atomic energy); the only difference is grid quadrature vs the
|
|
13
|
+
Monte-Carlo EMA estimate during training.
|
|
14
|
+
|
|
15
|
+
Directory layout assumed (same as main.py):
|
|
16
|
+
{results_root}/{atom}/{method}/bl_0.0000/
|
|
17
|
+
Checkpoints/checkpoint_*.eqx
|
|
18
|
+
training_metrics_ema.csv
|
|
19
|
+
job_params.json
|
|
20
|
+
|
|
21
|
+
Usage
|
|
22
|
+
-----
|
|
23
|
+
python atom_energies.py --method tf_w_lam0.2_none_lda_none_dopri8_db_sir_sched_MIX
|
|
24
|
+
python atom_energies.py --method <tag> --atoms B Be C F H He Li N Ne O \
|
|
25
|
+
--window 1000 --recompute --results_root /path/to/Results
|
|
26
|
+
"""
|
|
27
|
+
|
|
28
|
+
import sys, os
|
|
29
|
+
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
30
|
+
|
|
31
|
+
import argparse
|
|
32
|
+
from pathlib import Path
|
|
33
|
+
|
|
34
|
+
import pandas as pd
|
|
35
|
+
|
|
36
|
+
from quadrature import grid_energy_from_checkpoint
|
|
37
|
+
|
|
38
|
+
# atomic number == electron count for a neutral atom
|
|
39
|
+
Z_TABLE = {"H": 1, "He": 2, "Li": 3, "Be": 4, "B": 5,
|
|
40
|
+
"C": 6, "N": 7, "O": 8, "F": 9, "Ne": 10}
|
|
41
|
+
|
|
42
|
+
# ── CLI ───────────────────────────────────────────────────────────────────────
|
|
43
|
+
parser = argparse.ArgumentParser(
|
|
44
|
+
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
45
|
+
parser.add_argument("--method", type=str, required=True,
|
|
46
|
+
help="Method directory name, e.g. "
|
|
47
|
+
"tf_w_lam0.2_none_lda_none_dopri8_db_sir_sched_MIX")
|
|
48
|
+
parser.add_argument("--atoms", type=str, nargs="+",
|
|
49
|
+
default=["B", "Be", "C", "F", "H", "He", "Li", "N", "Ne", "O"],
|
|
50
|
+
help="Atoms to process (default: the full set).")
|
|
51
|
+
parser.add_argument("--results_root", type=str, default="Results",
|
|
52
|
+
help="Root dir holding {atom}/{method}/bl_0.0000 (default: Results)")
|
|
53
|
+
parser.add_argument("--window", type=int, default=1000,
|
|
54
|
+
help="Average the last N rows of training_metrics_ema.csv (default: 1000)")
|
|
55
|
+
parser.add_argument("--bs", type=int, default=256, help="Grid chunk size")
|
|
56
|
+
parser.add_argument("--grid_level", type=int, default=3, help="PySCF grid level")
|
|
57
|
+
parser.add_argument("--recompute", action="store_true",
|
|
58
|
+
help="Re-run grid integration even if energy_summary.json is cached")
|
|
59
|
+
parser.add_argument("--out", type=str, default="atom_energies.csv",
|
|
60
|
+
help="Output CSV path (default: atom_energies.csv)")
|
|
61
|
+
args = parser.parse_args()
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def read_last_ema(bl_dir: Path, window: int):
|
|
65
|
+
"""Mean of the last `window` rows of training_metrics_ema.csv.
|
|
66
|
+
E = E + CC (no nuclear repulsion; for an atom this is the total energy)."""
|
|
67
|
+
csv = bl_dir / "training_metrics_ema.csv"
|
|
68
|
+
if not csv.exists():
|
|
69
|
+
return None, None
|
|
70
|
+
try:
|
|
71
|
+
df = pd.read_csv(csv)
|
|
72
|
+
except pd.errors.EmptyDataError:
|
|
73
|
+
return None, None
|
|
74
|
+
if df.empty:
|
|
75
|
+
return None, None
|
|
76
|
+
tail = df.tail(window)
|
|
77
|
+
E = float(tail["E"].mean())
|
|
78
|
+
if "CC" in tail.columns:
|
|
79
|
+
E += float(tail["CC"].mean())
|
|
80
|
+
epoch = int(df.iloc[-1]["epoch"])
|
|
81
|
+
return E, epoch
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
# ── main loop ─────────────────────────────────────────────────────────────────
|
|
85
|
+
root = Path(args.results_root).resolve()
|
|
86
|
+
print(f"Results root : {root}")
|
|
87
|
+
print(f"Method : {args.method}")
|
|
88
|
+
print(f"EMA window : last {args.window} rows\n")
|
|
89
|
+
|
|
90
|
+
rows = []
|
|
91
|
+
for atom in args.atoms:
|
|
92
|
+
atom_dir = root / atom / args.method / "bl_0.0000"
|
|
93
|
+
print(f"[{atom}] {atom_dir}")
|
|
94
|
+
|
|
95
|
+
if not (atom_dir / "job_params.json").exists():
|
|
96
|
+
print(" SKIP — no job_params.json (directory missing?)\n")
|
|
97
|
+
continue
|
|
98
|
+
|
|
99
|
+
# grid integration of the last checkpoint (via the OFF quadrature module)
|
|
100
|
+
try:
|
|
101
|
+
data = grid_energy_from_checkpoint(
|
|
102
|
+
atom_dir, grid_level=args.grid_level, chunk=args.bs, recompute=args.recompute)
|
|
103
|
+
E_grid = data["E_total"]
|
|
104
|
+
grid_epoch = data["epoch"]
|
|
105
|
+
Ne_int = data["Ne_integral"]
|
|
106
|
+
except Exception as e:
|
|
107
|
+
print(f" grid integration FAILED: {e}")
|
|
108
|
+
E_grid = grid_epoch = Ne_int = None
|
|
109
|
+
|
|
110
|
+
# EMA mean of last `window` rows
|
|
111
|
+
E_ema, ema_epoch = read_last_ema(atom_dir, args.window)
|
|
112
|
+
if E_ema is None:
|
|
113
|
+
print(" no training_metrics_ema.csv")
|
|
114
|
+
|
|
115
|
+
diff = (E_grid - E_ema) if (E_grid is not None and E_ema is not None) else None
|
|
116
|
+
rows.append({
|
|
117
|
+
"atom": atom,
|
|
118
|
+
"Ne": Z_TABLE.get(atom),
|
|
119
|
+
"E_grid_Ha": E_grid,
|
|
120
|
+
"grid_epoch": grid_epoch,
|
|
121
|
+
"E_ema_Ha": E_ema,
|
|
122
|
+
"ema_epoch": ema_epoch,
|
|
123
|
+
"diff_Ha": diff,
|
|
124
|
+
"Ne_int": Ne_int,
|
|
125
|
+
})
|
|
126
|
+
if E_grid is not None and E_ema is not None:
|
|
127
|
+
print(f" E_grid={E_grid:+.6f} E_ema={E_ema:+.6f} Δ={diff:+.6f} Ha")
|
|
128
|
+
print()
|
|
129
|
+
|
|
130
|
+
if not rows:
|
|
131
|
+
raise RuntimeError("No atoms processed — check --results_root / --method.")
|
|
132
|
+
|
|
133
|
+
df = pd.DataFrame(rows)
|
|
134
|
+
|
|
135
|
+
# ── print + save ──────────────────────────────────────────────────────────────
|
|
136
|
+
print("=" * 80)
|
|
137
|
+
print(f"{'atom':>4} {'Ne':>3} {'E_grid [Ha]':>15} {'E_ema [Ha]':>15} "
|
|
138
|
+
f"{'Δ(grid-ema)':>13} {'∫ρ':>8}")
|
|
139
|
+
print("-" * 80)
|
|
140
|
+
for _, r in df.iterrows():
|
|
141
|
+
eg = f"{r['E_grid_Ha']:+15.6f}" if pd.notna(r['E_grid_Ha']) else f"{'—':>15}"
|
|
142
|
+
em = f"{r['E_ema_Ha']:+15.6f}" if pd.notna(r['E_ema_Ha']) else f"{'—':>15}"
|
|
143
|
+
dd = f"{r['diff_Ha']:+13.6f}" if pd.notna(r['diff_Ha']) else f"{'—':>13}"
|
|
144
|
+
ni = f"{r['Ne_int']:8.4f}" if pd.notna(r['Ne_int']) else f"{'—':>8}"
|
|
145
|
+
ne = f"{int(r['Ne'])}" if pd.notna(r['Ne']) else "?"
|
|
146
|
+
print(f"{r['atom']:>4} {ne:>3} {eg} {em} {dd} {ni}")
|
|
147
|
+
print("=" * 80)
|
|
148
|
+
|
|
149
|
+
out_path = Path(args.out).resolve()
|
|
150
|
+
df.to_csv(out_path, index=False, float_format="%.8f")
|
|
151
|
+
print(f"\nSaved → {out_path}")
|
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Global configuration module for storing runtime parameters and directories.
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
class Config:
|
|
6
|
+
"""Global configuration class to store all runtime parameters."""
|
|
7
|
+
|
|
8
|
+
# Model parameters
|
|
9
|
+
mol_name = None
|
|
10
|
+
epochs = None
|
|
11
|
+
bs = None
|
|
12
|
+
hl = None
|
|
13
|
+
lr = None
|
|
14
|
+
prior = None
|
|
15
|
+
|
|
16
|
+
# Functionals
|
|
17
|
+
kin = None
|
|
18
|
+
nuc = None
|
|
19
|
+
hart = None
|
|
20
|
+
x = None
|
|
21
|
+
c = None
|
|
22
|
+
cc = None
|
|
23
|
+
|
|
24
|
+
# Training settings
|
|
25
|
+
sched = None
|
|
26
|
+
solver = None
|
|
27
|
+
ckpt_freq = None
|
|
28
|
+
|
|
29
|
+
# Directories (set during runtime)
|
|
30
|
+
results_dir = None
|
|
31
|
+
ckpt_dir = None
|
|
32
|
+
|
|
33
|
+
@classmethod
|
|
34
|
+
def from_args(cls, args):
|
|
35
|
+
"""
|
|
36
|
+
Initialize configuration from argparse arguments.
|
|
37
|
+
|
|
38
|
+
Args:
|
|
39
|
+
args: argparse.Namespace object containing parsed arguments
|
|
40
|
+
"""
|
|
41
|
+
# Model parameters
|
|
42
|
+
cls.mol_name = args.mol_name
|
|
43
|
+
cls.epochs = args.epochs
|
|
44
|
+
cls.bs = args.bs
|
|
45
|
+
cls.hl = args.hl
|
|
46
|
+
cls.lr = args.lr
|
|
47
|
+
cls.prior = args.prior
|
|
48
|
+
|
|
49
|
+
# Functionals
|
|
50
|
+
cls.kin = args.kin
|
|
51
|
+
cls.nuc = args.nuc
|
|
52
|
+
cls.hart = args.hart
|
|
53
|
+
cls.x = args.x
|
|
54
|
+
cls.c = args.c
|
|
55
|
+
cls.cc = args.cc
|
|
56
|
+
|
|
57
|
+
# Training settings
|
|
58
|
+
cls.sched = args.sched
|
|
59
|
+
cls.solver = args.solver
|
|
60
|
+
cls.ckpt_freq = args.ckpt_freq
|
|
61
|
+
|
|
62
|
+
@classmethod
|
|
63
|
+
def set_directories(cls, results_dir, ckpt_dir):
|
|
64
|
+
"""
|
|
65
|
+
Set the results and checkpoint directories.
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
results_dir: Path to results directory
|
|
69
|
+
ckpt_dir: Path to checkpoint directory
|
|
70
|
+
"""
|
|
71
|
+
cls.results_dir = results_dir
|
|
72
|
+
cls.ckpt_dir = ckpt_dir
|
|
73
|
+
|
|
74
|
+
@classmethod
|
|
75
|
+
def get_model_params(cls):
|
|
76
|
+
"""Return dictionary of model parameters."""
|
|
77
|
+
return {
|
|
78
|
+
'mol_name': cls.mol_name,
|
|
79
|
+
'epochs': cls.epochs,
|
|
80
|
+
'bs': cls.bs,
|
|
81
|
+
'hl': cls.hl,
|
|
82
|
+
'lr': cls.lr,
|
|
83
|
+
'prior': cls.prior
|
|
84
|
+
}
|
|
85
|
+
|
|
86
|
+
@classmethod
|
|
87
|
+
def get_functionals(cls):
|
|
88
|
+
"""Return dictionary of functional parameters."""
|
|
89
|
+
return {
|
|
90
|
+
'kin': cls.kin,
|
|
91
|
+
'nuc': cls.nuc,
|
|
92
|
+
'hart': cls.hart,
|
|
93
|
+
'x': cls.x,
|
|
94
|
+
'c': cls.c,
|
|
95
|
+
'cc': cls.cc
|
|
96
|
+
}
|
|
97
|
+
|
|
98
|
+
@classmethod
|
|
99
|
+
def __repr__(cls):
|
|
100
|
+
"""String representation of configuration."""
|
|
101
|
+
return (
|
|
102
|
+
f"Config(\n"
|
|
103
|
+
f" Model: mol_name={cls.mol_name}, epochs={cls.epochs}, "
|
|
104
|
+
f"bs={cls.bs}, hl={cls.hl}, lr={cls.lr}, prior={cls.prior}\n"
|
|
105
|
+
f" Functionals: kin={cls.kin}, x={cls.x}, c={cls.c}, cc={cls.cc}\n"
|
|
106
|
+
f" Directories: results={cls.results_dir}, ckpt={cls.ckpt_dir}\n"
|
|
107
|
+
f")"
|
|
108
|
+
)
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
# MIT License
|
|
2
|
+
|
|
3
|
+
# Copyright (c) 2025 AlexandreDeCamargo
|
|
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.
|
|
22
|
+
|
|
23
|
+
__version__ = "0.1.0"
|
|
24
|
+
|
|
25
|
+
from .dft_distrax import (
|
|
26
|
+
DFTDistribution
|
|
27
|
+
)
|