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.
Files changed (42) hide show
  1. off-0.1.0/LICENSE +21 -0
  2. off-0.1.0/PKG-INFO +154 -0
  3. off-0.1.0/README.md +123 -0
  4. off-0.1.0/off/__init__.py +23 -0
  5. off-0.1.0/off/atom_energies.py +151 -0
  6. off-0.1.0/off/config/_config.py +108 -0
  7. off-0.1.0/off/dft_distrax/__init__.py +27 -0
  8. off-0.1.0/off/dft_distrax/dft_distrax.py +216 -0
  9. off-0.1.0/off/flow/__init__.py +29 -0
  10. off-0.1.0/off/flow/equiv_flows.py +99 -0
  11. off-0.1.0/off/functionals/__init__.py +35 -0
  12. off-0.1.0/off/functionals/core_correction.py +84 -0
  13. off-0.1.0/off/functionals/exchange_correlation.py +174 -0
  14. off-0.1.0/off/functionals/external.py +49 -0
  15. off-0.1.0/off/functionals/functional.py +129 -0
  16. off-0.1.0/off/functionals/hartree.py +62 -0
  17. off-0.1.0/off/functionals/kinetic.py +87 -0
  18. off-0.1.0/off/main.py +172 -0
  19. off-0.1.0/off/ode_solver/__init__.py +32 -0
  20. off-0.1.0/off/ode_solver/eqx_ode.py +76 -0
  21. off-0.1.0/off/plot_binding_csv.py +63 -0
  22. off-0.1.0/off/plot_pes_ema.py +259 -0
  23. off-0.1.0/off/plot_pes_mpl.py +280 -0
  24. off-0.1.0/off/promolecular/__init__.py +27 -0
  25. off-0.1.0/off/promolecular/promolecular_dist.py +465 -0
  26. off-0.1.0/off/quadrature.py +261 -0
  27. off-0.1.0/off/quadrature_scan.py +188 -0
  28. off-0.1.0/off/scan_pes.py +133 -0
  29. off-0.1.0/off/test_fwd_rev.py +290 -0
  30. off-0.1.0/off/train/__init__.py +44 -0
  31. off-0.1.0/off/train/loop.py +228 -0
  32. off-0.1.0/off/train/loss.py +149 -0
  33. off-0.1.0/off/train/utils.py +38 -0
  34. off-0.1.0/off/utils.py +618 -0
  35. off-0.1.0/off.egg-info/PKG-INFO +154 -0
  36. off-0.1.0/off.egg-info/SOURCES.txt +40 -0
  37. off-0.1.0/off.egg-info/dependency_links.txt +1 -0
  38. off-0.1.0/off.egg-info/entry_points.txt +3 -0
  39. off-0.1.0/off.egg-info/requires.txt +17 -0
  40. off-0.1.0/off.egg-info/top_level.txt +1 -0
  41. off-0.1.0/pyproject.toml +50 -0
  42. 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
+ )