dftax 0.1.0__tar.gz

dftax-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,89 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v8.2.0
+        with:
+          python-version: "3.13"
+          cache-dependency-glob: "pyproject.toml"
+      - name: Sync (test deps)
+        run: uv sync --extra test
+      - name: Ruff
+        run: uv run ruff check dftax scripts examples tests
+
+  test:
+    runs-on: ubuntu-latest
+    env:
+      OMP_NUM_THREADS: "2"
+      XLA_PYTHON_CLIENT_PREALLOCATE: "false"
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v8.2.0
+        with:
+          python-version: "3.13"
+          cache-dependency-glob: "pyproject.toml"
+      - name: Sync (test deps)
+        run: uv sync --extra test
+      - name: Unit tests (fast subset)
+        run: uv run --extra test pytest tests/unit -q -m "not slow"
+
+  # Confirm the package installs and runs on the lower bound of requires-python.
+  # Core deps only (no PySCF), so this is a cheap import + tiny-calc smoke.
+  compat:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v8.2.0
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache-dependency-glob: "pyproject.toml"
+      - name: Sync (core only)
+        run: uv sync --no-dev
+      - name: Import + tiny-calc smoke
+        run: >
+          uv run --no-dev python -c "import jax; jax.config.update('jax_enable_x64', True);
+          import dftax; from dftax import run_ks; from dftax.system import Molecule;
+          from dftax.energy.xc import LDA;
+          r = run_ks(Molecule.from_xyz('H 0 0 0; H 0 0 0.74', 'sto-3g'), LDA());
+          assert r.converged; print('compat OK', dftax.__version__, r.converged)"
+
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v8.2.0
+        with:
+          cache-dependency-glob: "pyproject.toml"
+      - name: Build sdist + wheel
+        run: uv build
+      - name: Check metadata
+        run: uvx twine check dist/*
+      - uses: actions/upload-artifact@v7
+        with:
+          name: dist
+          path: dist/
+
+  docs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v8.2.0
+        with:
+          python-version: "3.13"
+          cache-dependency-glob: "pyproject.toml"
+      - name: Sync (docs deps)
+        run: uv sync --extra docs
+      - name: Build docs
+        run: uv run mkdocs build

dftax-0.1.0/.github/workflows/docs.yml

@@ -0,0 +1,34 @@
+name: Deploy docs
+
+# Builds the mkdocs site and publishes it to the gh-pages branch on every push to
+# main. Enable GitHub Pages once (Settings -> Pages -> Source: gh-pages branch) so
+# https://andresguzco.github.io/dftax/ serves it. PRs only build (in ci.yml), they
+# do not deploy.
+on:
+  push:
+    branches: [main]
+    paths:
+      - "docs/**"
+      - "mkdocs.yml"
+      - "dftax/**"
+      - ".github/workflows/docs.yml"
+  workflow_dispatch:
+
+permissions:
+  contents: write
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+        with:
+          fetch-depth: 0
+      - uses: astral-sh/setup-uv@v8.2.0
+        with:
+          python-version: "3.13"
+          cache-dependency-glob: "pyproject.toml"
+      - name: Sync (docs deps)
+        run: uv sync --extra docs
+      - name: Deploy to gh-pages
+        run: uv run mkdocs gh-deploy --force

dftax-0.1.0/.github/workflows/nightly.yml

@@ -0,0 +1,28 @@
+name: Nightly (full suite)
+
+# The full test suite (including the modules auto-marked `slow`: SCF sweeps,
+# finite-difference properties, CPHF, density fitting and streaming, the d-shell
+# ERI compile) is too heavy for every PR, so it runs once a day and on demand.
+# The conftest per-test JAX cache release keeps memory bounded on the runner.
+on:
+  schedule:
+    - cron: "0 6 * * *"   # 06:00 UTC daily
+  workflow_dispatch:
+
+jobs:
+  full:
+    runs-on: ubuntu-latest
+    timeout-minutes: 120
+    env:
+      OMP_NUM_THREADS: "2"
+      XLA_PYTHON_CLIENT_PREALLOCATE: "false"
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v8.2.0
+        with:
+          python-version: "3.13"
+          cache-dependency-glob: "pyproject.toml"
+      - name: Sync (test deps)
+        run: uv sync --extra test
+      - name: Full test suite (including slow)
+        run: uv run --extra test pytest tests/unit -q

dftax-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,39 @@
+name: Publish to PyPI
+
+# Fires only on a version tag (e.g. `git tag v0.1.0 && git push --tags`).
+# Uses PyPI Trusted Publishing (OIDC), so no API token is stored in the repo.
+on:
+  push:
+    tags: ["v*"]
+
+jobs:
+  # Gate the release on the suite: a tag must never publish a failing build.
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v8.2.0
+        with:
+          cache-dependency-glob: "pyproject.toml"
+      - name: Sync (test extra)
+        run: uv sync --extra test
+      - name: Lint
+        run: uv run ruff check .
+      - name: Test
+        run: uv run pytest tests/unit -m "not slow" -q
+
+  publish:
+    needs: [test]
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write          # required for OIDC trusted publishing
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v8.2.0
+        with:
+          cache-dependency-glob: "pyproject.toml"
+      - name: Build
+        run: uv build
+      - name: Publish
+        uses: pypa/gh-action-pypi-publish@release/v1

dftax-0.1.0/.gitignore

@@ -0,0 +1,224 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*.egg-info
+
+# results
+results/
+
+# Checkpoints and model weights
+*.eqx
+*.pkl
+*.pt
+*.pth
+*.h5
+*.hdf5
+*.npz
+*.npy
+*.ckpt
+
+# docs
+docs/_build/
+/.dir-locals.el
+*bazel-*
+*tmp*
+*viztracer*
+*profiler*
+/compile_commands.json
+*.whl
+*.DS_Store
+*.ipynb_checkpoints
+__pycache__/
+# results/
+
+
+# somewhat against the advice but doing this anyway to support gpu as an optional dependency
+uv.lock
+
+# IDE
+.vscode
+
+# jupyter-book derived files
+docs/_autosummary
+docs/_build
+.benchmarks/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+#
+wandb/
+*.out
+output/
+
+.claude/
+# FL working-tree leftovers (not part of the engine branch)
+data/
+data
+paper/
+third_party/
+
+# Vendored engine data (required at runtime): Lebedev angular-grid tables and
+# Cartesian->spherical transforms. Re-include the dir first, then the file, to
+# override data/ and *.npz above.
+!dftax/grid/data/
+!dftax/grid/data/lebedev.npz
+!dftax/basis/data/
+!dftax/basis/data/cart2sph.npz

dftax-0.1.0/.python-version

@@ -0,0 +1 @@
+3.13

dftax-0.1.0/CHANGELOG.md

@@ -0,0 +1,45 @@
+# Changelog
+
+All notable changes to dftax are documented here. The format follows
+[Keep a Changelog](https://keepachangelog.com/), and the project aims to adhere
+to [Semantic Versioning](https://semver.org/).
+
+## [0.1.0] - 2026-06-30
+
+First public release: a differentiable Kohn-Sham DFT engine in pure JAX/Equinox,
+validated against PySCF and on GPU. The whole calculation is differentiable, so
+forces, response properties, and converged-density gradients all come from one
+autodiff engine.
+
+### Added
+- **RKS + UKS** total energies as differentiable `E(P)` functionals; the KS Fock
+  is `sym(∂E/∂P)` by autodiff (no hand-coded XC potential).
+- Solvers: on-device DIIS **SCF** (`rks_scf`/`uks_scf`) with optional
+  **level-shifting**, and Adam **direct minimization** (`rks_minimize`/`uks_minimize`).
+- **Analytic nuclear forces** via autodiff (`rks_forces`/`uks_forces`), FD-checked.
+- **Response properties**: dipole, polarizability, Hessian, vibrational
+  frequencies, IR and Raman intensities, and alchemical derivatives (∂E/∂Z).
+- **Implicit differentiation** of the SCF fixed point (CPHF) for gradients of
+  converged quantities, used for the analytic polarizability.
+- **Batched** energies and forces over many geometries via `vmap`
+  (`run_rks_batched`/`run_uks_batched`/`run_ks_batched`).
+- **Integrals** (Obara-Saika / McMurchie-Davidson): overlap, kinetic, nuclear
+  attraction, nuclear repulsion, and 2-/3-/4-center ERIs; up to **g (l=4)**;
+  spherical-harmonic (`cart2sph`) support; shell-pair-batched builders.
+- **XC functionals**: LDA (Slater + VWN5), PBE, PBE0, B3LYP; all closed- **and**
+  open-shell.
+- **Coulomb/exchange backends**: exact 4-center ERI; RI density fitting (RI-J /
+  RI-K); Cauchy-Schwarz screening.
+- **Memory-light paths**: streamed exact J/K, streamed + Schwarz-screened RI-J,
+  orbital-chunk streamed RI-K (with an exact `custom_vjp`), and streamed XC grid,
+  removing the O(N⁴)/nao²×naux materialization on every RKS/UKS path.
+- **Native pipeline**: a `Molecule` loader (Basis Set Exchange) and a native
+  Becke/Lebedev integration grid, so the compute path needs no PySCF, libcint, or
+  libxc at runtime.
+- **Validation/benchmark harnesses**: `scripts/gpu/validate_gpu.py` (+
+  `GPU_VALIDATION.md`) and `scripts/bench/benchmark.py` (+ `BENCHMARKS.md`).
+
+### Notes
+- Accuracy vs PySCF/libxc: LDA/B3LYP ~machine precision; PBE/PBE0 ~1e-5 Ha (the
+  hand-rolled GGA enhancement factors).
+- GPU is opt-in: `pip install dftax[cuda12]` (Linux). PySCF is a test-only oracle.

dftax-0.1.0/CITATION.cff

@@ -0,0 +1,24 @@
+cff-version: 1.2.0
+message: "If you use dftax in academic work, please cite it as below."
+title: "dftax: a differentiable Kohn-Sham DFT engine in JAX"
+abstract: >-
+  A differentiable Kohn-Sham density functional theory engine in pure JAX/Equinox.
+  Restricted and unrestricted DFT with autodiff Fock, analytic forces, response
+  properties, implicit-differentiation SCF response, density fitting, and a native
+  integration grid.
+type: software
+authors:
+  - family-names: "Guzmán-Cordero"
+    given-names: "Andrés"
+    email: "andres.guzman-cordero@mila.quebec"
+version: "0.1.0"
+license: Apache-2.0
+repository-code: "https://github.com/andresguzco/dftax"
+url: "https://github.com/andresguzco/dftax"
+keywords:
+  - density functional theory
+  - quantum chemistry
+  - electronic structure
+  - JAX
+  - differentiable
+  - automatic differentiation

dftax-0.1.0/CONTRIBUTING.md

@@ -0,0 +1,39 @@
+# Contributing to dftax
+
+Thanks for your interest! dftax is a small, dependency-light, differentiable
+Kohn-Sham DFT engine in pure JAX. A few conventions keep it that way.
+
+## Development setup
+
+Managed with [uv](https://docs.astral.sh/uv/) (Python ≥ 3.13):
+
+```bash
+uv sync --extra test          # engine + test deps (pytest, scipy, pyscf)
+uv sync --extra cuda12        # add the CUDA-12 jaxlib for GPU work (Linux)
+```
+
+## Running tests
+
+```bash
+uv run --extra test pytest tests/unit -q                 # full unit suite
+uv run --extra test pytest tests/unit -q -m "not slow"   # skip heavy compiles
+uv run --extra test pytest tests/unit/test_scf_rks.py -q # one file
+```
+
+Tests require float64 (`conftest.py` enables `jax_enable_x64`). **PySCF is a
+reference oracle only**. It must never enter the compute path; use it to build a
+reference energy/grid and compare. GPU correctness is validated interactively
+(`scripts/gpu/validate_gpu.py`), not in CPU CI.
+
+## Style
+
+- `ruff check .` and `ruff format .` (line length 100).
+- Match the surrounding code: typed `jaxtyping` shapes, small pure functions,
+  `eqx.Module` for stateful objects, chunked `vmap`/`lax.scan` for memory.
+- New functionals / integrals: validate against PySCF/libxc to the documented
+  tolerance and add a unit test.
+
+## Pull requests
+
+Keep PRs focused, include a test that fails without the change, and run the
+(non-slow) suite + ruff before submitting.