off 0.1.0__py3-none-any.whl

off-0.1.0.dist-info/METADATA

@@ -0,0 +1,154 @@
+Metadata-Version: 2.4
+Name: off
+Version: 0.1.0
+Summary: A library for learning the ground state energy using orbital free density functional theory with continuous normalizing flows.
+Author-email: Alexandre de Camargo <decamara@mcmaster.ca>
+Project-URL: Homepage, https://github.com/AlexandreDeCamargo/of_flows
+Project-URL: Repository, https://github.com/AlexandreDeCamargo/of_flows
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.11
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.6.0
+Requires-Dist: jaxlib>=0.6.0
+Requires-Dist: equinox>=0.12.0
+Requires-Dist: diffrax>=0.7.0
+Requires-Dist: optax>=0.2.0
+Requires-Dist: scipy>=1.15.0
+Requires-Dist: chex>=0.1.89
+Requires-Dist: jaxtyping>=0.3.0
+Requires-Dist: distrax>=0.1.5
+Requires-Dist: pyscf>=2.7.0
+Requires-Dist: pandas>=2.0.3
+Requires-Dist: matplotlib>=3.10.0
+Provides-Extra: examples
+Requires-Dist: matplotlib>=3.10.0; extra == "examples"
+Requires-Dist: chex>=0.1.89; extra == "examples"
+Requires-Dist: numpy>=2.3.0; extra == "examples"
+Dynamic: license-file
+
+# OFF: An Orbital-Free Density Functional Theory Python library using Normalizing Flows
+
+OFF is a [JAX](https://github.com/google/jax)-based library for **orbital-free density
+functional theory (OF-DFT)** in which the electron density is represented by a
+**continuous normalizing flow (CNF)** and the ground-state energy is obtained by
+*variationally minimizing* a density functional with Monte-Carlo gradient estimates.
+The density is normalized by construction (it is a probability flow), and the physical
+density is recovered as `ρ(x) = Ne · ρ_φ(x)`.
+
+OFF is built entirely on the JAX ecosystem — automatic differentiation, JIT
+compilation, vectorization, and GPU acceleration — with
+[Diffrax](https://github.com/patrick-kidger/diffrax) for the flow ODEs,
+[Equinox](https://github.com/patrick-kidger/equinox) for the models and functionals,
+[Distrax](https://github.com/google-deepmind/distrax) for the base distribution, and
+[Optax](https://github.com/google-deepmind/optax) for the optimization.
+
+## Functionality
+
+The current version of the library has the following capabilities: 
+
+* Provides an implementation of OF-DFT methods with continuous normalizing flows. 
+* A modular library of density functionals:
+  * **kinetic**: Thomas–Fermi, von Weizsäcker, and TF-λW;
+  * **exchange**: LDA and B88;
+  * **correlation**: VWN and PW92;
+  * **nuclear attraction**, 
+  * **Hartree**, and 
+  * **nuclear-cusp corrections** (Kato, Hutcheon).
+* Evaluates density functionals using Monte Carlo estimators. 
+* In addition to Monte Carlo estimators, it also has a deterministic **grid (quadrature) readout** of the energy after training, using a
+  [PySCF](https://github.com/pyscf/pyscf) Becke grid converted to `jax.Array`.
+
+## Install
+
+A core dependency is [PySCF](https://pyscf.org), which needs `cmake` available on the
+`PATH`. In a fresh environment, from the repository root:
+
+```bash
+pip install -e .
+```
+
+The `db_sir` prior additionally requires AtomDB; install it only if you use that prior.
+
+## Use
+
+Two stages: (1) **train** a normalizing flow for a given molecule and functional, then
+(2) read out the energy on a **quadrature grid**. Both are exposed as command-line
+tools (after `pip install -e .`) and as a Python API.
+
+### 1. Train a flow
+
+```bash
+off-train --mol_name H2 --bond_length 1.4 \
+          --kin tf_w --lam 1/5 --x lda_b88_x --c none \
+          --hart coulomb --prior promolecular \
+          --solver dopri8 --epochs 500 --bs 512
+```
+(equivalently `python -m off.main ...`). This minimizes the OF-DFT energy with the
+Monte-Carlo estimator and writes everything under a method-tagged directory:
+
+```
+Results/H2/tf_w_lam0.2_none_lda_b88_x_none_dopri8_promolecular_sched_mix/bl_1.40/
+    Checkpoints/checkpoint_*.eqx     # the trained flow
+    training_metrics_ema.csv         # EMA energy trace
+    job_params.json                  # everything needed to rebuild the run
+```
+
+Key options: `--kin {tf,w,tf_w}`, `--x {lda,b88_x,lda_b88_x}`, `--c {vwn_c,pw92_c,none}`,
+`--cc {kato,hutcheon,none}`, `--hart {coulomb,coulomb_}` (all-pairs / element-wise),
+`--prior {promolecular,db_sir}`, `--solver {dopri5,tsit5,dopri8}`.
+
+### 2. Evaluate the energy on a grid
+
+After training, point the quadrature tool at the run directory. It rebuilds the flow
+from `job_params.json`, builds a PySCF grid, evaluates `ρ_φ` and its score there, and
+integrates every energy term:
+
+```bash
+off-quadrature Results/H2/tf_w_lam0.2_none_lda_b88_x_none_dopri8_promolecular_sched_mix/bl_1.40 --grid_level 3
+```
+(equivalently `python -m off.quadrature ...`). It prints the per-term energies
+(`T, V_N, V_H, E_X, E_C, E_CC, E_NN, E_total`) and the `∫ρ` check, and caches the
+result in `energy_summary.json`. The same call from Python:
+
+```python
+from off import grid_energy_from_checkpoint
+
+e = grid_energy_from_checkpoint(
+    "Results/H2/.../bl_1.40", grid_level=3)
+print(e["E_total"], e["Ne_integral"])
+```
+
+### Build a grid
+
+```python
+from off import getGrid
+
+h2_geom = "H 0 0 0; H 0 0 1.4"               
+w_grid, x_grid = getGrid(h2_geom, level=3, units="bohr")
+```
+
+## Package layout
+
+```
+off/
+  main.py            # training entry point   (off-train)
+  quadrature.py      # grid-energy readout    (off-quadrature)
+  flow/              # the continuous normalizing flow (CNF)
+  ode_solver/        # Diffrax forward/reverse ODE solves
+  promolecular/      # base distributions (promolecular, AtomDB)
+  functionals/       # kinetic, exchange-correlation, nuclear, Hartree, cusp
+  train/             # loss (Monte-Carlo energy) and the optimization loop
+```
+
+## Citation
+
+```bibtex
+@article{off,
+  title  = {Orbital-Free DFT with Normalizing Flows},
+  author = {de Camargo, Alexandre and others},
+  year   = {2026},
+}
+```

off-0.1.0.dist-info/RECORD

@@ -0,0 +1,37 @@
+off/__init__.py,sha256=ALVYNGD585q811-3aL94_cx8veoBFIE_g49tMp_NHu0,648
+off/atom_energies.py,sha256=67Zj77F6gYomMBzuAAFMPpw6_MEwjEZPKZ6XOanIkSk,6458
+off/main.py,sha256=BL2w7WmlvrAc9HJjxaumqqWhDl2EVzFyoI26vxJNGLY,6199
+off/plot_binding_csv.py,sha256=A8f68xUGdtFEiBYCclm0AYNJQSOfzYfQEuYtxtGpO2w,2369
+off/plot_pes_ema.py,sha256=GgdRHMLrHKenfP4hh1UqJH4EwFqEXj9C4XhYCWe2I5o,11391
+off/plot_pes_mpl.py,sha256=OMT1S8Du0SNZJk_5uXQBfjUT32bWPVXnLuun2eFTNAI,11305
+off/quadrature.py,sha256=1RenXdFdQw2yM4oyxJxONg1jMpJtKNiBJ5_7JBptgt0,11649
+off/quadrature_scan.py,sha256=PJ8PswqAnL4oECyDTZVbQJzG6ToEvu9IsFNjiunmass,7639
+off/scan_pes.py,sha256=bDHf1i_G-H2RRwn_54TOQlBuiy0UPPmN_PYHzqhXonQ,5554
+off/test_fwd_rev.py,sha256=XkxV2sgsmHwx0Q5hmeC-h42sWNA-UF56n73Xh4C6nZE,13056
+off/utils.py,sha256=NQgzFm1nZ5kS4ysN_FLaKUENzbw9iCAiEw6IYukKnU0,23844
+off/config/_config.py,sha256=PwL2hUIjM9ruZ12S7xH5gFwXyor6AW8vKEaj4Ix1eXo,2754
+off/dft_distrax/__init__.py,sha256=bZphcl4p3iJm7W7KLCCjXXSF_XEu_YfYUpEWDbSE-yI,1181
+off/dft_distrax/dft_distrax.py,sha256=Yx1cFYwpFiG43BH7okykx2mT3B1bYAK9I8w6FL6GHPk,5994
+off/flow/__init__.py,sha256=573x3ngNaxBEXsvorthaNrMcJBSzhYhaOVzxgKjxWx4,1196
+off/flow/equiv_flows.py,sha256=KExeU8NNP1DmQSCYH0fdWEui-YPGLVvRVpw_Jj19yJ4,2953
+off/functionals/__init__.py,sha256=v5gdzgt_IqtyP_qBe3-358p7RB6RI-P3PL3qgaNawnE,1522
+off/functionals/core_correction.py,sha256=UKUHTGq3C2mYy8xDeICnq1U-JT7by0vSkhyTy5UIhEg,2990
+off/functionals/exchange_correlation.py,sha256=44XBulYXmpvMwcMOd91BahcRAumRVMA2mCgZRO3LYeg,6034
+off/functionals/external.py,sha256=m7UgLKo-oRpRcuGL_OHfaB-glyGkod95tEFD1mucaOM,1389
+off/functionals/functional.py,sha256=KokOZwXBbHZPfafDadzR5EF52wQfKDfXA15Ax8F3OIU,5194
+off/functionals/hartree.py,sha256=dwJKZ2_FSc9qkJCgQAhkfd2ExJiDw_hSiIK9ezbv9Yg,2226
+off/functionals/kinetic.py,sha256=196RNHEo92hBO9IooIq8U44YL3nLtHQ-B8tKQ04mkig,2948
+off/ode_solver/__init__.py,sha256=mYZkRmXbxN5fC95LSWrQjsxVWqjO_4O-HFn4I7R6x38,1228
+off/ode_solver/eqx_ode.py,sha256=nlcg2u7_o-Y6RTMKc-RzKHisorhU4Rmza4Oyhpv-g9s,2792
+off/promolecular/__init__.py,sha256=f53ZFDEDgGF_sop-vQuPPn7rJeqNuC9zh5LM-8p_dB0,1191
+off/promolecular/promolecular_dist.py,sha256=03C5OEt5ZQZSIp8dtys_4k3h5Hi1rk9VSfU8XdQTzGY,17960
+off/train/__init__.py,sha256=WB-GNBSR4Q1Dx3GWBukTsKTsmumSL8Gmb2KiVrHOrJg,1387
+off/train/loop.py,sha256=nDqj62YvgXHeTnOyaRJ8qAtilVMc9gKYBCQdbIzZh1A,6981
+off/train/loss.py,sha256=X-cVUgQJaUdaWPzSIrKfx-NV4ZftQbI2UkyEsnsuMjg,4634
+off/train/utils.py,sha256=M2dyC_cF1y2ypy_AYxNbeIpi-MGV4vp8oHfhmpAcX5Y,1171
+off-0.1.0.dist-info/licenses/LICENSE,sha256=G-KAKMtNilNsYiaAdjC4OR9CAH5DCm9Cbkk1iCN6dpA,1075
+off-0.1.0.dist-info/METADATA,sha256=ySA7S66SiRbQ5ZHuj5R6_UNiNM-mEC6js7CYKR60Go4,5902
+off-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+off-0.1.0.dist-info/entry_points.txt,sha256=A77cJ29nE_z1jmTLWkYxzjhbVMyj1W-OIHWZW0KM9ho,81
+off-0.1.0.dist-info/top_level.txt,sha256=tHC4TPDNMYnhBhW__qfVy9n2ETS6HCz9UT3dRbd2pfY,4
+off-0.1.0.dist-info/RECORD,,

off-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

off-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+off-quadrature = off.quadrature:main
+off-train = off.main:main

off-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 AlexandreDeCamargo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

off-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+off