diff --git a/.github/workflows/install_and_test.yml b/.github/workflows/install_and_test.yml index 37c0c22..ce381b0 100644 --- a/.github/workflows/install_and_test.yml +++ b/.github/workflows/install_and_test.yml @@ -6,6 +6,17 @@ on: - synchronize - opened - reopened + # Only run CI when a change can affect its outcome: library code, tests, + # dependency metadata, the workflow/config it consumes, or lint config. + # Docs/notebooks/plan-file-only changes skip CI. + paths: + - 'grad_dft/**' + - 'tests/**' + - 'requirements.txt' + - 'setup.py' + - '.github/**' + - '.pylintrc' + - '.pre-commit-config.yaml' jobs: build: @@ -14,7 +25,7 @@ jobs: PYSCF_CONFIG_FILE: ".github/workflows/pyscf_conf.py" strategy: matrix: - python-version: ["3.9", "3.10", "3.11"] + python-version: ["3.12"] os: [ubuntu-latest] steps: - uses: actions/checkout@v3 diff --git a/.gitignore b/.gitignore index 2f05116..1496eff 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,8 @@ *__pycache__/ data/training/dissociation/ runs/ -*/__pycache__/ \ No newline at end of file +*/__pycache__/ +.vscode/ +*.egg-info/ +# Agent workflow plan files (single source of truth: mirrored issue comments) +tasks/ diff --git a/.opencode/skills/code-review/SKILL.md b/.opencode/skills/code-review/SKILL.md new file mode 100644 index 0000000..1a72652 --- /dev/null +++ b/.opencode/skills/code-review/SKILL.md @@ -0,0 +1,46 @@ +--- +name: code-review +description: Use after all subtasks are committed and before integrating a task. Perform a whole-repo review of the task's changes against the integration branch, iterate to zero findings. +--- + +# Code Review + +Run after all subtasks are committed, before the quality gate and integration +into the integration branch. Review the entire repo's diff for the task and +iterate until there are zero findings. + +## Procedure + +1. Determine the task's changes against the integration branch + (`INTEGRATION=$(git config --get graddft.integrationBranch || echo agent_settings)`): + `git diff --stat "$INTEGRATION"...HEAD` and `git diff "$INTEGRATION"...HEAD`. +2. Review against the checklist below. +3. For every finding, fix it in a follow-up commit on the same feature branch, + using the commit format from the `subtask-loop` skill. +4. Repeat until zero findings. + +## Checklist + +- [ ] **Correctness** — every clause of the task issue (#N) is traceable to + committed code. +- [ ] **Tests** — new code is covered by `tests/unit/` or + `tests/integration/`; no test was weakened or skipped; the suite passes + (see `run-tests`). +- [ ] **Docs** — the documentation surface (`README.md`, `examples/`, + docstrings) was updated per the mapping in the `planning` skill, and every + path, file name, and API name it mentions exists. +- [ ] **Duplication** — no copy-pasted or near-identical logic; shared helpers + belong in `grad_dft/utils/`. +- [ ] **Style** — `black` and `pylint` (per `.pre-commit-config.yaml`) pass + clean; types come from `grad_dft/utils/types.py` where applicable. +- [ ] **Dead code** — no unused imports, functions, or debug artifacts. +- [ ] **Scope** — changes touch only what the task requires; no unrelated edits. + +Zero findings means all boxes pass. Do not merge with known findings. + +## Aggregate final review (stacked mode) + +Immediately before the final pull request to `main`, review the whole +`git diff main...$INTEGRATION` (all tasks accumulated on the integration +branch) with the same checklist, and re-run the quality gate. This catches +interactions between tasks that per-task reviews cannot see. diff --git a/.opencode/skills/planning/SKILL.md b/.opencode/skills/planning/SKILL.md new file mode 100644 index 0000000..27bdd00 --- /dev/null +++ b/.opencode/skills/planning/SKILL.md @@ -0,0 +1,216 @@ +--- +name: planning +description: Use when planning tasks: creating global plans across multiple tasks, decomposing a task into atomic subtasks, or authoring new task descriptions. Covers multi-task planning, detailed plan format, atomic subtask requirements, documentation surface mapping, and task authoring guidelines. +--- + +# Planning + +## Multi-task Planning + +When the user asks to work on a **set of related tasks**, do NOT jump directly +into implementation. First set up the batch, then plan it: + +1. Create a **hub** issue (label `hub`) for the batch and one issue per task + (label `task`), and attach each task issue to the hub as a GitHub + **sub-issue** via the REST API (`POST repos/$REPO/issues//sub_issues` + with the task's database `id`); a link in the hub body is **not** a + sub-issue. The exact commands live in the `workflow-management` skill. +2. Create the high-level global plan in `tasks/global_plan.md`, naming the hub + and listing each task with its issue number, then mirror it to the hub as a + comment whose first line is ``. + +The global plan must: + +- Name the **hub** issue it belongs to. +- List all tasks to be done with their IDs, issue numbers (`#N`), brief + descriptions, and a **Status** (`[done #N]` once the task is integrated). +- Record the integration branch the batch targets (`agent_settings` by + default; see the resolution command in the `workflow-management` skill). +- Identify dependencies between tasks (which must be done before which). +- Identify potential conflicts or overlapping changes (e.g., two tasks + modifying the same file). +- Identify shared infrastructure that multiple tasks need, and place it in a + single shared module instead of duplicating it across tasks. The usual home + for it is `grad_dft/utils/` — type aliases in `grad_dft/utils/types.py`, + array helpers in `grad_dft/utils/utils.py`, memory helpers in + `grad_dft/utils/chunk.py`. +- Identify existing reusable abstractions (see "Reuse analysis" below) that new + tasks should use rather than reinvent. +- Keep plotting, reporting, and experiment-driver code out of `grad_dft/`; it + belongs in `examples/` (`advanced_scripts/` for scripts, + `article_experiments/` for paper reproductions) so rendering and library + logic stay in separate modules from the start. +- Propose an execution order that minimizes rework and avoids conflicts. +- Align tasks with the project architecture: `README.md` describes the + intended usage flow (`Molecule`/`Solid` -> energy densities and coefficients + -> `NeuralFunctional` -> training), and the tests under `tests/unit/` and + `tests/integration/` show how each layer is exercised. + +After the global plan is created and mirrored to the hub, proceed with the +normal working loop: one task at a time, a feature branch per task created from +the integration branch, a detailed plan in `tasks/detailed_plan.md` for each. +After each task integrates, mark it `[done #]` in `tasks/global_plan.md` and +update the hub's `` comment. The local file is the single +source of truth; the hub comment is a pure mirror (see the `workflow-management` +skill). + +## Detailed Plan (Atomic Subtasks) + +`tasks/detailed_plan.md` MUST decompose the task into atomic subtasks. Each +subtask: + +- Is small enough to complete in a single focused work session (typically + 30–90 minutes). +- Produces a compilable, testable increment — no partial implementations left + uncommitted. +- Has a unique identifier (e.g., "S1", "S2") used in commit messages and plan + tracking. + +If a task is ambiguous or underspecified **during planning** (conflicting +requirements, unclear scope, missing constraints), do not guess. Ask the user +for clarification, then transfer the guidance to the task per the +`user-guidance-transfer` skill before proceeding with decomposition. + +### Reuse analysis + +Before writing the detailed plan, search the repository for functions, types, +and documentation sections that can be reused or generalized rather than +reimplemented: `grep`/`glob` across `grad_dft/`, `tests/`, and `examples/`, +starting from `grad_dft/utils/` and the neighbouring module of the code being +changed. Record what is reused in each subtask's **Code** section. Reusing +`grad_dft/utils/types.py` aliases (`Array`, `Key`, `PyTree`, `Scalar`, ...) in +new signatures counts as reuse. + +### Subtask Format + +Each subtask in `tasks/detailed_plan.md` MUST include these four sections. **A +subtask with missing Code, Tests, or Docs sections is incomplete and must not +be executed.** + +``` +### SN: + +**Code:** <files to create or modify, types and functions to add> +**Tests:** <test files to create or modify, specific test approaches: + pytest unit test in tests/unit/, integration test in + tests/integration/... , etc.> +**Docs:** <README.md sections, examples/ notebooks or scripts, and docstrings + to create or update — see "Documentation surface" below.> + +**Spec:** +- <detailed implementation specification> +``` + +Example: + +``` +### S1: Add a prediction test for the DM21 functional + +**Code:** New `tests/integration/molecules/test_predict_DM21.py` (reuse the + molecule construction and energy-comparison helpers from + `tests/integration/molecules/test_predict_B88.py`) +**Tests:** The new file itself is the test: assert the neural functional + predicts the reference DM21 energies within tolerance, and that the + loss is finite over a few SCF iterations +**Docs:** No user-facing change needed (test-only); keep the README untouched + +**Spec:** +- Reuse the `B3LYP_WITH_VWN5`-style PySCF configuration path ... +``` + +### Documentation surface + +This repository has no Sphinx site. The documentation surface is: + +- **`README.md`** — `Functionality`, `Install`, `Use example`, `Bibtex`. + Update it when a task adds or changes public API, installation, or the + intended usage flow. +- **`examples/`** — `basic_notebooks/` (first use of the library), + `intermediate_notebooks/` (specific training/functional techniques), + `advanced_scripts/` (scripts such as `examples/advanced_scripts/train_scf_loop.py`, + plus `examples/advanced_scripts/readme.md`), `article_experiments/` (paper + reproductions). Update or add a notebook/script when a task changes how the + library is used. +- **Docstrings** — numpydoc-style docstrings on the public API in + `grad_dft/`. New public functions, classes, and parameters get them; keep + examples in docstrings consistent with the README snippets. + +Use this mapping to fill each subtask's **Docs** section. A subtask whose change +is invisible in all three places writes "No user-facing change needed" and +explains why — it is never left blank. + +### Granularity + +If a subtask cannot be committed as a self-contained increment, it is too large +— split it further. The commit message format lives in the `subtask-loop` +skill; it is the single source of truth for commit format and procedure. + +### Bounded uncommitted work + +Uncommitted work on a feature branch must never exceed one atomic subtask. If a +session is interrupted, the loss is bounded to that single subtask. + +## Post-Implementation Design Notes + +When a task hits algorithmic limitations — partially completed, with skipped +tests or remaining work — append a `## Design Notes` section to +`tasks/detailed_plan.md`. This serves as persistent design knowledge for future +task refinement. The section must be structured as follows: + +``` +## Design Notes (discovered during implementation) + +### <Topic Title> + +Design rationale, coordinate spaces, invariants — as confirmed by the user. + +### <Failure Topic Title> + +- Root causes with concrete examples (e.g., "for input `mol` with 2 basis + functions, the generalized eigenvalue problem returns 0 occupied states + because...") +- What was attempted and why it didn't fully work +- Remaining work: concrete, actionable items +- Skipped tests: list them and the reason +``` + +Requirements: + +- Every algorithmic limitation MUST be traceable to a concrete input, a + concrete location in the data structure, and a concrete execution path in + the code. +- Never write vague descriptions like "some gradients are missing" — specify + which input, which state/range, and which path should produce them. +- Remaining work items must be actionable (e.g., "propagate the electron count + through the SCF loop by adding it to the `Molecule` dataclass") — not vague + goals. +- If the user provided design guidance (e.g., ansatz choice, loss + decomposition), record it verbatim in the `### <Topic>` section as the + authoritative reference. + +## Task Authoring Guidelines + +When creating a new task issue (`gh issue create`, label `task`), follow +these rules (they apply to the issue body): + +- **Specify output format upfront**. If the task produces artifacts (plots, + tables, checkpoint files, trained models), include the exact layout, units, + and formatting rules in the task description. +- **Keep tasks single-responsibility**. A task should do one thing. If it + requires more than 5 sub-items or spans multiple unrelated concerns, split it + into multiple tasks. +- **Specify equivalence requirements**. For any new algorithm variant, + explicitly state "must produce results identical to X" (e.g., identical to + the non-jitted implementation, or to the PySCF reference energy) so + equivalence tests are built from the start. +- **Specify type genericity**. If a module must handle arbitrary inputs, state + it explicitly (e.g., "generic over the `Molecule` and `Solid` classes", not + molecules only). +- **Specify reuse expectations**. If the task builds on existing infrastructure + (e.g., "reuse the array aliases from `grad_dft/utils/types.py`"), name the + dependencies. This prevents reinvention. + +## Task Completeness Verification + +See the `subtask-loop` skill — it is the single source of truth for verifying +task completion before a task is considered done. diff --git a/.opencode/skills/quality-gates/SKILL.md b/.opencode/skills/quality-gates/SKILL.md new file mode 100644 index 0000000..0f4b327 --- /dev/null +++ b/.opencode/skills/quality-gates/SKILL.md @@ -0,0 +1,44 @@ +--- +name: quality-gates +description: Use before integrating a task. Defines the hard gate that must pass (tests + style) and how to interpret its result. References the CI workflow and the pre-commit config for the exact commands. +--- + +# Quality Gates + +The hard gate a task must pass before integration. It has exactly two terminal +states: **PASS** or **BLOCKED**. There is no "pass with exceptions". The same +gate is run again over the whole `main...$INTEGRATION` diff immediately before +the final pull request to `main`. + +## What the gate is + +The gate is the combination of the test suite and the style/lint checks: + +- **Tests** — the per-file `pytest -v` invocations in + `.github/workflows/install_and_test.yml` (see the `run-tests` skill). 0 + failures, 0 skipped. +- **Style/lint** — the hooks in `.pre-commit-config.yaml`: `black` (rev pinned + there) and `pylint -rn -sn --rcfile=.pylintrc`. + +There is **no docs build gate** — this repository has no Sphinx site; the +documentation surface is `README.md` plus `examples/`, checked by review, not +by a build. There is also no type-check step and no coverage threshold. + +## Procedure + +1. Install the package and extras (`pip install -e .`, + `pip install -e ".[examples]"`), then run the tests. 0 failures, 0 skipped. +2. Run `pre-commit run --all-files` (or `black` and `pylint -rn -sn + --rcfile=.pylintrc` directly). No errors. +3. Interpret the result: + - All clean → **PASS**. Proceed to merge into the integration branch. + - Any failure → **BLOCKED**. + +## On BLOCKED + +- STOP. Do not merge. Do not mark the task done. +- Do **not** assess whether a failure is pre-existing or unrelated to your + changes — fix it regardless. +- Do **not** weaken, skip, or comment out failing tests to make the suite green + (see the Blocked Work Protocol in the `subtask-loop` skill). +- Fix every failure and re-run until **PASS**. diff --git a/.opencode/skills/run-tests/SKILL.md b/.opencode/skills/run-tests/SKILL.md new file mode 100644 index 0000000..08d2d47 --- /dev/null +++ b/.opencode/skills/run-tests/SKILL.md @@ -0,0 +1,41 @@ +--- +name: run-tests +description: Use when running the GradDFT test suite (pytest). Thin pointer to .github/workflows/install_and_test.yml, which holds the exact commands and working directories, plus the machine-specific pitfalls. +--- + +# Run tests + +The test pipeline and the exact commands live in the CI workflow steps in +`.github/workflows/install_and_test.yml`. Read that workflow before running +anything — do not reconstruct the command list from memory. + +## Setup + +The suite needs the package and its example extras installed: + +```bash +pip install -e . +pip install -e ".[examples]" +``` + +## Commands + +CI runs `pytest -v <file>` from the repository root, once per test file (two +unit files, then the molecules and solids integration files). Run the same +invocations locally so local and CI results are comparable. Example: + +```bash +PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py \ + pytest -v tests/unit/test_eigenproblem.py +``` + +## Notes + +- `PYSCF_CONFIG_FILE` must point at `.github/workflows/pyscf_conf.py` for the + functional-prediction tests; CI sets it as a job-level `env` var. Without it + PySCF builds functionals with its own defaults and results differ from CI. +- `tests/unit/` is fast; `tests/integration/` compiles SCF loops and is slow. + Prefer the unit file that covers your change while iterating, then run the + full CI list before calling a subtask done. +- No coverage gate is configured; the suite is expected to report 0 failures + and 0 skipped. diff --git a/.opencode/skills/subtask-loop/SKILL.md b/.opencode/skills/subtask-loop/SKILL.md new file mode 100644 index 0000000..b2ee4da --- /dev/null +++ b/.opencode/skills/subtask-loop/SKILL.md @@ -0,0 +1,248 @@ +--- +name: subtask-loop +description: Use when executing an atomic subtask from the detailed plan: implement -> test -> document -> pre-commit checks -> quality checks -> commit -> mark done. Covers the full execution cycle, commit format, blocked work protocol, and completion tracking. +--- + +# Subtask Execution Loop + +Each atomic subtask from `tasks/detailed_plan.md` is executed as a +self-contained cycle. + +## Cycle Steps + +Execute these steps in order. **Do not skip steps. Do not proceed past a step +until it is verified complete.** + +## Documentation-Only Subtasks + +When a subtask modifies only `.md`/`.ipynb` files (no `.py` files), the +following cycle steps are adapted: + +| Step | Action | +|------|--------| +| 1. Implement | Write documentation | +| 2. Write Tests | **Skip** — no code to test | +| 3. Update Docs | The implementation IS the documentation; verify README links and referenced example paths exist | +| 4. Pre-Commit Check | **Skip** — the `black`/`pylint` hooks in `.pre-commit-config.yaml` only match Python files | +| 5. Code Quality Checks | **Skip** — no code to check | +| 6. Commit | **Follow exactly** — one commit per subtask, single SN identifier | +| 7. Mark Done | **Follow exactly** | + +The absence of code changes **never** justifies batching multiple subtasks +into a single commit. + +### 1. Implement + +Write the code specified in the subtask's **Code** section. Follow the existing +structure and conventions of the module being changed (Apache-2.0 header, +import grouping, `grad_dft/utils/types.py` aliases). + +### 2. Write Tests + +Write the tests specified in the subtask's **Tests** section. See the +`run-tests` skill for how to run them. Unit tests live in `tests/unit/`, +integration tests in `tests/integration/molecules/` or +`tests/integration/solids/`. + +### 3. Update Documentation + +Map the source changes onto the documentation surface (`README.md`, +`examples/`, docstrings) exactly as defined in the **Documentation surface** +section of the `planning` skill — that mapping is the single source of truth. + +**Hard gate — this step is not complete until:** + +- [ ] The subtask's **Docs** section items are done, or an explicit + "No user-facing change needed" justification is in the plan +- [ ] Every README/example path, file name, and function name mentioned in the + updated docs actually exists (`glob`/`grep` to confirm) — a stale + example path is a documentation defect + +### 4. Pre-Commit Check + +Run `pre-commit run --all-files` (or `black` and +`pylint -rn -sn --rcfile=.pylintrc` directly, per +`.pre-commit-config.yaml`) and fix everything it reports. Do not commit until it +passes. + +### 5. Code Quality Checks + +- **Duplication check**: scan for accidental code duplication (same logic under + different names, copy-pasted blocks). Consolidate into + `grad_dft/utils/` if found. +- **Genericity check**: verify new code reuses existing abstractions and the + shared types from `grad_dft/utils/types.py` instead of bespoke + re-implementations. +- **Equivalence test check**: if the subtask is a variant of an existing + algorithm, ensure an equivalence test exists comparing it to the reference + implementation (e.g., a jitted path against the eager path, a new functional + against the PySCF reference energy). +- **Separation check**: verify library modules (`grad_dft/`) do not contain + plotting, file I/O, or ad-hoc string generation that belongs in + `examples/advanced_scripts/`. + +### 6. Commit + +**Commit format** (the single source of truth for commit format in this +repository): + +- Subject line: `(<issue>-S<n>): <imperative summary>` — e.g. + `(42-S2): Reuse `chunk` for the two-electron integrals`. This exact prefix is + what task-completion detection greps for, so never reword it. +- Body: explain **why** the change was made and what was rejected, not a + restatement of the diff. Reference the subtask and any design decision + recorded in `## Design Notes`. +- The **last** subtask of a task adds `Closes #<issue>` as a standalone line + (the task's own issue). The last task of a batch also adds `Closes #<hub>`. + +Before committing, verify `git status` shows only the files this subtask +touches, and that the work is not left uncommitted across subtasks. + +### 7. Mark Completed + +Mark the subtask as completed in `tasks/detailed_plan.md` (`[done]` suffix + +commit hash on the subtask heading, and tick its cycle steps in the same +plan), then mirror the file onto the plan comment of the task issue: find the +comment by its `<!-- detailed-plan -->` first line and PATCH its body to the +current file content: + +```bash +REPO=$(gh repo view --json nameWithOwner -q .nameWithOwner) +CID=$(gh api "repos/$REPO/issues/<N>/comments" \ + --jq '[.[] | select(.body | startswith("<!-- detailed-plan -->"))] | last.id') +{ echo '<!-- detailed-plan -->'; cat tasks/detailed_plan.md; } > /tmp/plan_comment.md +gh api -X PATCH "repos/$REPO/issues/comments/$CID" \ + --input <(jq -Rs '{body: .}' /tmp/plan_comment.md) +``` + +The local file is the single source of truth; the comment is a pure mirror. +The per-subtask cycle checklist required by "Per-Subtask Execution Tracking" +below lives in this same plan file — there is no separate tracker. + +If at any point in steps 1–7 you hit an unresolvable problem that prevents 100% +completion, **STOP the cycle immediately** and follow the Blocked Work Protocol +below. Do NOT attempt to "complete" the subtask with partial results, reverted +work, or known limitations. Do NOT proceed to the next subtask. + +## Subtask Outcome + +A subtask has exactly two valid outcomes: + +- **Resolved**: implemented, tested, docs updated, committed. Record the commit + hash in `tasks/detailed_plan.md`. +- **Blocked**: an algorithmic or design problem prevents 100% completion. Do NOT + commit partial work. Do NOT proceed to the next subtask. Follow the Blocked + Work Protocol. + +There is no third state. "Reverted and left as a known limitation" is not a +valid outcome — it means the subtask is blocked. Report it. + +Never silently skip a subtask. If a subtask was attempted, reverted, and its +planned changes were not committed, the subtask is incomplete. Do not mark it +done. Do not proceed. Report it as blocked. + +## Per-Subtask Execution Tracking + +Track the cycle steps as a checklist inside the subtask's own section of +`tasks/detailed_plan.md`, e.g. for S1: + +``` +- [x] S1: Implement +- [x] S1: Write tests +- [x] S1: Update documentation +- [x] S1: Pre-Commit Check (black + pylint) +- [x] S1: Quality checks +- [x] S1: Commit (<hash>) +``` + +**No subtask may be committed with any step still unticked.** + +## Multi-Subtask Discipline + +When a task has multiple subtasks (S1, S2, S3, ...), execute them **strictly +sequentially**: + +1. Complete all cycle steps for S1 (Implement → Tests → Docs → Pre-Commit Check + → Quality → Commit → Mark Done). +2. Only after S1 is committed, start S2. +3. Never tick multiple subtasks' steps before committing each individually. + +A plan listing "S1: Implement [x], S2: Implement [x], S1: Write tests [x]" +indicates skipped commits — each subtask must be fully committed before the next +begins. + +## Blocked Work Protocol + +If you encounter an algorithmic problem that you cannot resolve to 100% +correctness, **STOP**. Do not commit. Do not merge. Do not comment out or weaken +failing tests to make the suite green. Instead: + +1. Stay on the feature branch. +2. Report the problem concretely to the user: + - Which tests fail and why. + - What algorithmic gap exists (e.g., "gradients are missing for the + range-separated Hartree-Fock term in `grad_dft/external/_hf_density.py`"). + - What you've tried and what remains unresolved. +3. Ask the user for guidance: additional subtasks, algorithmic hints, descoping, + or splitting the task. +4. **Transfer user guidance to the task** per the `user-guidance-transfer` + skill — post it verbatim as a comment on the task issue. +5. Append a `## Design Notes` section to `tasks/detailed_plan.md`. See the + `planning` skill for the full template. Minimum required content: + + - **Correct Design**: algorithmic design as confirmed by the user — + ansatz, loss decomposition, invariants. Quote the user's design guidance + verbatim where available. + - **Blocked Subtasks**: which subtasks are blocked and why. + - **Root Causes**: why each failure occurs, with concrete examples. Every + limitation MUST be traceable to a concrete input, a concrete location in + the data structure, and a concrete execution path in the code. Never write + vague descriptions. + - **Approaches Tried**: what was attempted and why it didn't fully work. + - **Remaining Work**: concrete, actionable items (e.g., "propagate the + density matrix through the SCF loop by adding a field to the `Molecule` + dataclass") — not vague goals. + - **Skipped Tests**: list any tests skipped and the reason. + + Commit this summary so the plan serves as a persistent design record for + future task refinement. +6. Post the block report (the `## Design Notes` content) as a comment on the + task issue, so the block is visible in the task record. + +## Task Completion Verification + +This section is the **single source of truth** for what "done" means. Before +a task is considered done — its last subtask's commit carries `Closes #<N>` +(the task's own issue) and is merged to the integration branch +(`git config --get graddft.integrationBranch || echo agent_settings`) — +verify: + +- [ ] Every clause in the task description is traceable to implemented and + committed code. +- [ ] No subtask was reverted without resolution. +- [ ] No requirement was silently skipped or deferred. +- [ ] All tests pass (0 failures, 0 skipped). +- [ ] All quality gates pass (see the `quality-gates` skill). +- [ ] If any of the above fails, the task is NOT done — it is blocked. Follow + the Blocked Work Protocol. Do NOT close it with known unresolved + limitations. + +Partial completion is not completion. "All passing tests are for the parts I +did" does not mean the task is done if other parts were reverted. + +## Marking Complete + +A task is complete when its last subtask's commit carries `Closes #<N>` (the +task's own issue number) as a standalone line and that commit is on the +integration branch. The closing keyword takes effect when the commit reaches +`main`: because work is stacked on the integration branch, the issue stays open +until the final pull request to `main` (opened on the user's explicit request) +lands. Never edit the issue body — it is user-authored and immutable. + +Commit hashes recorded in `tasks/detailed_plan.md` may change after an +integration-branch rebase; the `(<issue>-S<n>)` commit subject is the stable +key. + +Complete means COMPLETE: every requirement met, every test passing, every +edge case handled. Never close a task with known failures or unresolved +limitations. diff --git a/.opencode/skills/user-guidance-transfer/SKILL.md b/.opencode/skills/user-guidance-transfer/SKILL.md new file mode 100644 index 0000000..6c25945 --- /dev/null +++ b/.opencode/skills/user-guidance-transfer/SKILL.md @@ -0,0 +1,31 @@ +--- +name: user-guidance-transfer +description: Use when the user provides design guidance, descoping, algorithmic hints, or clarification during a task. Record it verbatim as a comment on the task issue so it is never lost or paraphrased. +--- + +# User Guidance Transfer + +When the user gives guidance mid-task (design decisions, algorithmic hints, +descoping, clarification of ambiguity), transfer it into the task record +verbatim. + +## Procedure + +1. Post the guidance as a comment on the task issue: + `gh issue comment <N> --body '[USER GUIDANCE]: "<verbatim quote>"'`. + Never edit the issue body — it is user-authored and immutable. +2. Quote the user's guidance **verbatim**. Do not paraphrase, summarize, or + interpret. +3. If the guidance is algorithmic and belongs in the persistent design record, + also record it verbatim in the `### <Topic>` section of the `## Design + Notes` block in `tasks/detailed_plan.md` (see the `planning` skill). + +Example comment: + +``` +[USER GUIDANCE]: "Keep the electron density in atomic units and reuse +`grad_dft/utils/types.py` aliases; do not introduce a second dtype switch in +`grad_dft/functional.py`." +``` + +Only transfer guidance the user actually gave; do not invent or pad it. diff --git a/.opencode/skills/workflow-management/SKILL.md b/.opencode/skills/workflow-management/SKILL.md new file mode 100644 index 0000000..c1c52d0 --- /dev/null +++ b/.opencode/skills/workflow-management/SKILL.md @@ -0,0 +1,144 @@ +--- +name: workflow-management +description: Load first at the start of every session, before doing anything else. Drives the overall working loop for a task or set of tasks on the integration branch (a personal, stacked branch): global plan, one task at a time, feature branch, detailed plan, subtask execution, code review, quality gate, integration, final pull request. +--- + +# Workflow Management + +The top-level loop that orchestrates all other workflow skills. For a +single task, load this skill first; it points to the skills that handle each +step. + +## Core Rules + +- Log all tasks as GitHub issues with the `task` label (`gh issue create + --label task`). Use the description exactly as the user provided it — minimal + changes, only splitting and numbering. The task ID is the issue number. +- A batch of **related** tasks is tracked by a single **hub** issue (label + `hub`). Hubs carry the global plan and are never chosen as tasks; task issues + keep the `task` label and are attached to the hub as GitHub **sub-issues** + through the REST API (`POST repos/$REPO/issues/<HUB>/sub_issues`). Listing the + task numbers inside the hub body is **not** a substitute for the sub-issue + relationship. +- The global plan for the active batch lives in `tasks/global_plan.md` (single + source of truth) and is mirrored to the hub issue as a comment whose first + line is `<!-- global-plan -->`. Task progress is marked in the file and the + mirror is updated after each task integrates, exactly as the detailed plan is + mirrored to a task issue (see the `subtask-loop` skill). +- Resolve the **integration branch** once per session: + + ```bash + INTEGRATION=$(git config --get graddft.integrationBranch || echo agent_settings) + ``` + + This snippet is the single source of truth for the integration branch; + other skills reuse it verbatim. This repository works in **stacked mode**: + the integration branch is a long-lived personal branch (`agent_settings` by + default) sitting on top of `main`. Never commit directly to it. +- Do tasks strictly **one at a time**. Each task gets its own feature branch + (created from the integration branch), its own detailed plan, and its own + integration step. Never combine multiple tasks in a single feature branch. +- A task is integrated into the integration branch by rebase and fast-forward + merge, preserving the per-subtask commits. **No pull request is opened**; the + single pull request into `main` is opened only on the user's explicit request + (see "Finalize to main"). +- Each decision must be documented before implementation. Documentation must + be detailed enough to reproduce the project from scratch and understand why + each decision was made. +- Commit messages must be detailed enough to understand the reasons for + changes (see the commit format in the `subtask-loop` skill). +- Documentation-only tasks (no `.py` files changed) skip code-specific gates + (tests, `black`, `pylint`) because the pre-commit hooks only match Python + files. All other workflow rules still apply: one task per branch, one commit + per subtask, code review. + +## Working Loop + +0. If the user requests multiple related tasks at once, set up the batch before + any implementation: + a. Create the **hub** issue (`gh issue create --label hub`); it tracks the + batch and is not itself a task. + b. Create one issue per task (`gh issue create --label task`), then attach + each task issue to the hub as a GitHub **sub-issue** (mandatory, not a + body link): + + ```bash + REPO=$(gh repo view --json nameWithOwner -q .nameWithOwner) + CHILD_ID=$(gh api "repos/$REPO/issues/<N>" --jq .id) + gh api -X POST "repos/$REPO/issues/<HUB>/sub_issues" -F sub_issue_id="$CHILD_ID" + ``` + + Verify the parent: `gh api "repos/$REPO/issues/<N>/parent" --jq .number` + must print the hub number. + c. Write the global plan in `tasks/global_plan.md` (see the `planning` + skill), naming the hub and listing each task with its issue number. + d. Post the global plan as a comment on the hub whose first line is + `<!-- global-plan -->` (a pure mirror of the file). A related task added + later is attached to the same hub and added to the plan. +1. Ensure user-defined tasks, the global plan, and project architecture are + aligned. Start from `README.md` and the module layout under `grad_dft/`. +2. Choose exactly ONE open `task`-labeled issue that is not yet done: a task + is done when its subtask commits are on the integration branch + (`git log "$INTEGRATION" --format=%s | grep -cE '\(<N>-S[0-9]+\):'` > 0, + where N is the issue number). List candidates with + `gh issue list --label task --state open`. Hubs carry `hub`, not `task`, so + this list never returns them. +3. Rebase the integration branch onto `origin/main` (task-boundary refresh), + then create the feature branch for this single task from it: + `git checkout -b <task-slug> "$INTEGRATION"`. +4. Generate a detailed plan in `tasks/detailed_plan.md`, decomposing the task + into atomic subtasks (see `planning`), then post it as a comment on the + task issue; the first line of the comment is the marker + `<!-- detailed-plan -->`. +5. Load the `subtask-loop` skill, then execute each subtask using its cycle. +5a. Verify all subtasks are complete and unblocked. Check + `tasks/detailed_plan.md`: + - If any subtask is marked `[blocked]` or `[deferred]`, STOP immediately. + The task is NOT complete. Do NOT proceed to code review or the merge. + Report blocking subtasks to the user and await guidance. + - If a subtask was attempted, not committed, and its work reverted, the + subtask is NOT complete. Do not silently skip it. +5b. When continuing a partially-done task, analyze current state before any + code changes: + - Verify `git branch --show-current` is the correct feature branch. + - Review committed subtasks: `git log --oneline` on the feature branch. + - Read `tasks/detailed_plan.md` and `tasks/global_plan.md`. + - Cross-reference committed file changes with planned subtasks: + `git diff --stat HEAD.."$INTEGRATION"`. + - Report status to the user: "S1-S3 committed, S4 pending, ..." +6. After all subtasks are done, perform code review on the entire repo (see + the `code-review` skill). Iteratively detect and fix problems until zero + findings. +7. Load the `quality-gates` skill and run the gate. It must show PASS. If + BLOCKED, do not assess whether failures are pre-existing or unrelated to + your changes; fix every failure and re-run until PASS. +8. Integrate the task: rebase the feature branch onto the updated integration + branch and fast-forward merge it, then delete the feature branch. Do **not** + push for a PR and do **not** open a PR. Advance the integration branch only + after the user confirms the merge. +9. For a batch, update its progress: mark the integrated task `[done #<N>]` in + `tasks/global_plan.md`, then mirror the file to the hub's global-plan + comment using the same procedure as the detailed plan (step 7 of the + `subtask-loop` skill), substituting the hub number for the task number and + the `<!-- global-plan -->` marker for `<!-- detailed-plan -->`. +10. Verify the last subtask's commit carries `Closes #<N>` (the task's own + issue) as a standalone line. For the last task of a batch, the same commit + also carries `Closes #<hub>` once every sub-issue is resolved. The issues + close when that commit reaches `main` (see the Task Completeness + Verification in the `subtask-loop` skill, the single source of truth for + what "done" means). +11. Return to step 2 for the next task, accumulating completed tasks on + `$INTEGRATION`. Open the request to `main` only via the finalize step, on + the user's explicit request. + +## Finalize to main (explicit request only) + +When — and only when — the user explicitly asks to open the request to `main`: + +1. Rebase `$INTEGRATION` onto `origin/main` and re-run the aggregated code + review and quality gate over the whole `main...$INTEGRATION` diff (see + `code-review` and `quality-gates`). +2. Push `$INTEGRATION` and open one pull request with base `main` and head + `$INTEGRATION`. +3. The user reviews and merges the PR with rebase and merge. After the merge, + sync `main`; never delete the long-lived integration branch. diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..aa0d8c0 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,51 @@ +# AGENTS.md + +This file is a short TOC of the agent workflow. The project itself is +described in `README.md`, which is the entrypoint for navigating the +repository. + +## Start here + +1. Load the `workflow-management` skill + (`.opencode/skills/workflow-management/SKILL.md`) first, before doing + anything else. +2. Read `README.md` for project navigation, then `examples/` for worked usage. + +## Main Principles + +* This file is a short entry point for fast cold errors-free start. +* Only one source of truth. No duplicates. Each thing (in both code and documentation) described exactly once. Use generalization (especially for code), cross-references, links, other similar techniques to avoid duplicates and reuse staff. +* Source-of-truth hierarchy: code, scripts, CI configs > docs. If a fact can + be extracted from code or scripts (e.g., CI configs), it is not duplicated + in docs. Docs hold only what cannot be unambiguously reconstructed from + code: design decisions, non-trivial constraints. Skills stay thin pointers + to docs, code, or CI — they never re-describe them. +* Tools, not instructions. If you can do something with existing tool --- do it. No thinking, no long instructions, no manual analysis. You want to analyze code coverage? Just run coverage tool and analyze report. No workaround for regular tasks. If there is a tool for regular task it must be installed and configured appropriately. +* Never stay silent about problems. Any problem discovered or detected must be addressed. If it is in current scope or the fix is easy (e.g. a misprint), fix it — even if pre-existing. If it is significantly out of scope, report it to the user and propose creating an issue. +* Always learn, never forget — encode patterns before session ends. + +## Repository Workflow Configuration + +The integration branch is the long-lived personal branch `agent_settings` +(stacked on `main`); resolve it with the command in the `workflow-management` +skill, which is the single source of truth. Feature branches are cut from it, +one per task. + +## Skills + +### Domain + +| Skill | When to use | +|---|---| +| `.opencode/skills/run-tests` | Running the test suite (`pytest`, per `.github/workflows/install_and_test.yml`) | +| `.opencode/skills/quality-gates` | The pre-merge gate that must pass | + +### Workflow / process + +| Skill | When to use | +|---|---| +| `.opencode/skills/workflow-management` | Driving the overall task loop | +| `.opencode/skills/planning` | Global plans, atomic subtasks, task authoring | +| `.opencode/skills/subtask-loop` | Executing one atomic subtask | +| `.opencode/skills/user-guidance-transfer` | Recording user guidance verbatim | +| `.opencode/skills/code-review` | Whole-repo review before merge | diff --git a/examples/advanced_scripts/train_complex_model.py b/examples/advanced_scripts/train_complex_model.py index 763c8da..6c5681d 100644 --- a/examples/advanced_scripts/train_complex_model.py +++ b/examples/advanced_scripts/train_complex_model.py @@ -28,7 +28,7 @@ import warnings from grad_dft import ( - train_kernel, + train_kernel, energy_predictor, DispersionFunctional, NeuralFunctional, @@ -39,15 +39,16 @@ dm21_combine_densities, dm21_hfgrads_cinputs, dm21_hfgrads_densities, - loader + loader, ) from torch.utils.tensorboard import SummaryWriter import jax -from jax.config import config +from jax import config + config.update("jax_enable_x64", True) -config.update('jax_debug_nans', True) +config.update("jax_debug_nans", True) warnings.warn("This script takes a long time to run.") # In this example we explain how to replicate the experiments that train @@ -70,7 +71,9 @@ squash_offset = 1e-4 layer_widths = [width_layers] * n_layers nlc_layer_widths = [width_layers // 4] * (n_layers // 2) -out_features = 20 # 2 for each spin x 2 for exchange/correlation x 4 for MGGA + 4 for HF +out_features = ( + 20 # 2 for each spin x 2 for exchange/correlation x 4 for MGGA + 4 for HF +) sigmoid_scale_factor = 2.0 activation = gelu loadcheckpoint = False @@ -88,7 +91,9 @@ def nn_coefficients(instance, rhoinputs, *_, **__): instance.sow("intermediates", "tanh", x) # 6 Residual blocks with 256-features dense layer and layer norm - for features, i in zip(layer_widths, range(len(layer_widths))): # layer_widths = [256]*6 + for features, i in zip( + layer_widths, range(len(layer_widths)) + ): # layer_widths = [256]*6 res = x x = instance.dense(features=features)(x) instance.sow("intermediates", "residual_dense_" + str(i), x) @@ -172,7 +177,9 @@ def combine_densities(densities, ehf): ), combine_densities=combine_densities, coefficient_inputs=dm21_coefficient_inputs, - nograd_coefficient_inputs=lambda molecule, *_, **__: molecule.HF_energy_density(omegas), + nograd_coefficient_inputs=lambda molecule, *_, **__: molecule.HF_energy_density( + omegas + ), coefficient_input_grads=lambda self, params, molecule, nograd_cinputs, grad_cinputs, densities, *_, **__: dm21_hfgrads_cinputs( self, params, molecule, nograd_cinputs, grad_cinputs, densities, omegas ), @@ -212,7 +219,10 @@ def combine_densities(densities, ehf): ckpt_dir = os.path.join(dirpath, "ckpts/", "checkpoint_" + str(checkpoint_step) + "/") if loadcheckpoint: train_state = functional.load_checkpoint( - tx=tx, ckpt_dir=ckpt_dir, step=checkpoint_step, orbax_checkpointer=orbax_checkpointer + tx=tx, + ckpt_dir=ckpt_dir, + step=checkpoint_step, + orbax_checkpointer=orbax_checkpointer, ) params = train_state.params tx = train_state.tx @@ -264,7 +274,9 @@ def train_epoch(state, training_files, training_data_dirpath): load = loader(fname=fpath, randomize=True, training=True, config_omegas=omegas) for _, system in tqdm(load, "Molecules/reactions per file"): - params, opt_state, cost_val, metrics = kernel(params, opt_state, system, system.energy) + params, opt_state, cost_val, metrics = kernel( + params, opt_state, system, system.energy + ) del system # Logging the resulting metrics @@ -289,7 +301,9 @@ def train_epoch(state, training_files, training_data_dirpath): # Run an optimization step over a training batch state = params, opt_state, cost_val - state, metrics, epoch_metrics = train_epoch(state, training_files, training_data_dirpath) + state, metrics, epoch_metrics = train_epoch( + state, training_files, training_data_dirpath + ) params, opt_state, cost_val = state # Save metrics and checkpoint @@ -299,6 +313,8 @@ def train_epoch(state, training_files, training_data_dirpath): for metric in epoch_metrics.keys(): writer.add_scalar(f"/{metric}/train", epoch_metrics[metric], epoch) writer.flush() - functional.save_checkpoints(params, tx, step=epoch, orbax_checkpointer=orbax_checkpointer) + functional.save_checkpoints( + params, tx, step=epoch, orbax_checkpointer=orbax_checkpointer + ) print(f"-------------\n") print(f"\n") diff --git a/examples/advanced_scripts/train_scf_loop.py b/examples/advanced_scripts/train_scf_loop.py index 77762de..9c31f73 100644 --- a/examples/advanced_scripts/train_scf_loop.py +++ b/examples/advanced_scripts/train_scf_loop.py @@ -24,20 +24,22 @@ from grad_dft.interface.pyscf import molecule_from_pyscf from grad_dft.interface.pyscf import loader -from grad_dft.functional import NeuralFunctional, canonicalize_inputs, dm21_coefficient_inputs, densities +from grad_dft.functional import ( + NeuralFunctional, + canonicalize_inputs, + dm21_coefficient_inputs, + densities, +) from jax.nn import gelu from orbax.checkpoint import PyTreeCheckpointer from torch.utils.tensorboard import SummaryWriter -from grad_dft import ( - train_kernel, - energy_predictor, - diff_scf_loop -) +from grad_dft import train_kernel, energy_predictor, diff_scf_loop + +from jax import config -from jax.config import config config.update("jax_enable_x64", True) -config.update('jax_debug_nans', True) +config.update("jax_debug_nans", True) orbax_checkpointer = PyTreeCheckpointer() @@ -83,7 +85,9 @@ def coefficients(instance, rhoinputs, *_, **__): instance.sow("intermediates", "tanh", x) # 6 Residual blocks with 256-features dense layer and layer norm - for features, i in zip(layer_widths, range(len(layer_widths))): # layer_widths = [256]*6 + for features, i in zip( + layer_widths, range(len(layer_widths)) + ): # layer_widths = [256]*6 res = x x = instance.dense(features=features)(x) instance.sow("intermediates", "residual_dense_" + str(i), x) @@ -169,7 +173,9 @@ def train_epoch(state, training_files, training_data_dirpath): load = loader(fname=fpath, randomize=True, training=True, config_omegas=[]) for _, system in tqdm.tqdm(load): - params, opt_state, cost_val, metrics = kernel(params, opt_state, system, system.energy) + params, opt_state, cost_val, metrics = kernel( + params, opt_state, system, system.energy + ) del system for k in metrics.keys(): @@ -193,7 +199,9 @@ def train_epoch(state, training_files, training_data_dirpath): # Run an optimization step over a training batch state = params, opt_state, cost_val - state, metrics, epoch_metrics = train_epoch(state, training_files, training_data_dirpath) + state, metrics, epoch_metrics = train_epoch( + state, training_files, training_data_dirpath + ) params, opt_state, cost_val = state # Save metrics and checkpoint @@ -203,7 +211,9 @@ def train_epoch(state, training_files, training_data_dirpath): for metric in epoch_metrics.keys(): writer.add_scalar(f"/{metric}/train", epoch_metrics[metric], epoch) writer.flush() - functional.save_checkpoints(params, tx, step=epoch, orbax_checkpointer=orbax_checkpointer) + functional.save_checkpoints( + params, tx, step=epoch, orbax_checkpointer=orbax_checkpointer + ) print(f"-------------\n") print(f"\n") diff --git a/examples/basic_notebooks/example_lda_functional_02.ipynb b/examples/basic_notebooks/example_lda_functional_02.ipynb index a9e2a2c..e0d34d2 100644 --- a/examples/basic_notebooks/example_lda_functional_02.ipynb +++ b/examples/basic_notebooks/example_lda_functional_02.ipynb @@ -111,7 +111,7 @@ " # Molecule can compute the density matrix.\n", " rho = molecule.density()\n", " # To avoid numerical issues in JAX we limit too small numbers.\n", - " rho = jnp.clip(rho, a_min=clip_cte)\n", + " rho = jnp.clip(rho, min=clip_cte)\n", " # Now we can implement the LDA energy density equation in the paper.\n", " lda_e = -3 / 2 * (3 / (4 * jnp.pi)) ** (1 / 3) * (rho ** (4 / 3)).sum(axis=1, keepdims=True)\n", " # For simplicity we do not include the exchange polarization correction\n", diff --git a/examples/basic_notebooks/example_neural_functional_03.ipynb b/examples/basic_notebooks/example_neural_functional_03.ipynb index 8e0e1fb..9b795ae 100644 --- a/examples/basic_notebooks/example_neural_functional_03.ipynb +++ b/examples/basic_notebooks/example_neural_functional_03.ipynb @@ -88,7 +88,7 @@ "def energy_densities(molecule: gd.Molecule, clip_cte: float = 1e-30, *_, **__):\n", " r\"\"\"Auxiliary function to generate the features of LSDA.\"\"\"\n", " # Molecule can compute the density matrix.\n", - " rho = jnp.clip(molecule.density(), a_min=clip_cte)\n", + " rho = jnp.clip(molecule.density(), min=clip_cte)\n", " # Now we can implement the LDA energy density equation in the paper.\n", " lda_e = -3/2 * (3/(4*jnp.pi)) ** (1/3) * (rho**(4/3)).sum(axis = 1, keepdims = True)\n", " # For simplicity we do not include the exchange polarization correction\n", diff --git a/examples/basic_notebooks/example_scf_neural_functional_04.ipynb b/examples/basic_notebooks/example_scf_neural_functional_04.ipynb index 6e50e61..5d18880 100644 --- a/examples/basic_notebooks/example_scf_neural_functional_04.ipynb +++ b/examples/basic_notebooks/example_scf_neural_functional_04.ipynb @@ -76,7 +76,7 @@ "def energy_densities(molecule: gd.Molecule, clip_cte: float = 1e-30, *_, **__):\n", " r\"\"\"Auxiliary function to generate the features of LSDA.\"\"\"\n", " # Molecule can compute the density matrix.\n", - " rho = jnp.clip(molecule.density(), a_min=clip_cte)\n", + " rho = jnp.clip(molecule.density(), min=clip_cte)\n", " # Now we can implement the LDA energy density equation in the paper.\n", " lda_e = -3/2 * (3/(4*jnp.pi)) ** (1/3) * (rho**(4/3)).sum(axis = 1, keepdims = True)\n", " # For simplicity we do not include the exchange polarization correction\n", diff --git a/examples/intermediate_notebooks/DM21_02.ipynb b/examples/intermediate_notebooks/DM21_02.ipynb index 149e917..71c7105 100644 --- a/examples/intermediate_notebooks/DM21_02.ipynb +++ b/examples/intermediate_notebooks/DM21_02.ipynb @@ -63,7 +63,7 @@ "from tqdm import tqdm\n", "\n", "\n", - "from jax.config import config\n", + "from jax import config\n", "config.update(\"jax_enable_x64\", True)" ] }, diff --git a/examples/intermediate_notebooks/neural_metaGGA_01.ipynb b/examples/intermediate_notebooks/neural_metaGGA_01.ipynb index d8d3561..e0b7635 100644 --- a/examples/intermediate_notebooks/neural_metaGGA_01.ipynb +++ b/examples/intermediate_notebooks/neural_metaGGA_01.ipynb @@ -60,7 +60,7 @@ "from optax import adam, apply_updates\n", "from jax.nn import gelu\n", "\n", - "from jax.config import config\n", + "from jax import config\n", "config.update(\"jax_enable_x64\", True)\n", "\n", "import grad_dft as gd\n", diff --git a/examples/intermediate_notebooks/periodic_systems_bz_sampling_05.ipynb b/examples/intermediate_notebooks/periodic_systems_bz_sampling_05.ipynb index 060c590..b9ddfb6 100644 --- a/examples/intermediate_notebooks/periodic_systems_bz_sampling_05.ipynb +++ b/examples/intermediate_notebooks/periodic_systems_bz_sampling_05.ipynb @@ -260,7 +260,7 @@ } ], "source": [ - "from jax.config import config\n", + "from jax import config\n", "config.update(\"jax_enable_x64\", True)\n", "import grad_dft as gd\n", "\n", diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py index 7c62fa2..5a8c62e 100644 --- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py +++ b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py @@ -16,7 +16,7 @@ import enum import os -from typing import Generator, Optional, Sequence, Tuple, Union +from typing import Generator, List, Optional, Sequence, Tuple, Union import attr import numpy as np @@ -24,6 +24,24 @@ from pyscf import gto from pyscf.dft import numint import tensorflow.compat.v1 as tf + +try: + import pkg_resources # pylint: disable=unused-import +except ModuleNotFoundError: # setuptools >= 81 removed pkg_resources + # tensorflow_hub.__init__ still imports pkg_resources.parse_version to + # check the installed tensorflow version. Provide the one symbol it needs + # when pkg_resources is no longer part of the environment. + import sys as _sys + from types import SimpleNamespace as _SimpleNamespace + + def _parse_version(version): + """Returns a comparable key for a dotted numeric version string.""" + return tuple(int(part) for part in version.split(".") if part.isdigit()) + + _sys.modules.setdefault( + "pkg_resources", _SimpleNamespace(parse_version=_parse_version) + ) + import tensorflow_hub as hub import grad_dft.external.density_functional_approximation_dm21.density_functional_approximation_dm21.compute_hfx_density as compute_hfx_density @@ -161,7 +179,8 @@ class NeuralNumInt(numint.NumInt): In order to supply the local Hartree-Fock features required for the DM21 functionals, we lightly wrap the NumInt class. The actual evaluation of the - exchange-correlation functional is performed in NeuralNumInt.eval_xc. + exchange-correlation functional is performed in NeuralNumInt.eval_xc, which + PySCF's nr_rks/nr_uks grid loops reach through the eval_xc_eff override. Usage: mf = dft.RKS(...) # dft.ROKS and dft.UKS are also supported. @@ -216,7 +235,11 @@ def _build_graph(self, batch_dim: Optional[int] = None): library. """ - self._functional = hub.Module(spec=self._model_path) + # TF-Hub 0.16 removed the TF1 Module API (hub.Module). The checkpoints + # are SavedModels, so load them with hub.load; the "default" signature's + # keyword names keep the "tensor_dict$" prefix used below (only the + # TensorSpec name= fields are sanitised to "_"). + self._functional = hub.load(self._model_path).signatures["default"] grid_coords = tf.placeholder(tf.float32, shape=[batch_dim, 3], name="grid_coords") grid_weights = tf.placeholder(tf.float32, shape=[batch_dim], name="grid_weights") @@ -260,7 +283,7 @@ def _build_graph(self, batch_dim: Optional[int] = None): } tensor_dict = {f"tensor_dict${k}": v for k, v in features.items()} - predictions = self._functional(tensor_dict, as_dict=True) + predictions = self._functional(**tensor_dict) local_xc = predictions["grid_contribution"] weighted_local_xc = local_xc * grid_weights unweighted_xc = tf.reduce_sum(local_xc, axis=0) @@ -311,18 +334,10 @@ def _build_graph(self, batch_dim: Optional[int] = None): grid_weights=grid_weights, ) - outputs = { - "vxc": self._vxc, - "vrho": tf.stack(self._vrho), - "vsigma": tf.stack(self._vsigma), - "vtau": tf.stack(self._vtau), - "vhf": tf.stack(self._vhf), - } - # Create the signature for TF-Hub, including both the energy and functional - # derivatives. - # This is a no-op if _build_graph is called outside of - # hub.create_module_spec. - hub.add_signature(inputs=attr.asdict(self._placeholders), outputs=outputs) + # Note: TF-Hub 0.16 removed hub.add_signature/create_module_spec, so no + # signature is registered here. The derivatives (self._vrho, ...) stay + # graph attributes consumed by eval_xc via self._session.run; only the + # export path below needed them in a signature, and it now raises. def export_functional_and_derivatives( self, @@ -346,15 +361,24 @@ def export_functional_and_derivatives( batch_dim: the batch dimension of the grid to use in the model. Default: None (determine at runtime). This should only be set if the exported model is to be ahead-of-time compiled into a standalone library. + + Raises: + NotImplementedError: always. TF-Hub 0.16 removed the TF1 Module API + (hub.create_module_spec, hub.Module.export) this method was built + on; exporting would require a TF2 tf.saved_model.save + reimplementation. Only the vendored export_saved_model.py and + neural_numint_test.py use this method, and neither is part of the + test suite. """ - with tf.Graph().as_default(): - spec = hub.create_module_spec( - self._build_graph, tags_and_args=[(set(), {"batch_dim": batch_dim})] - ) - functional_and_derivatives = hub.Module(spec=spec) - with tf.Session() as session: - session.run(tf.global_variables_initializer()) - functional_and_derivatives.export(export_path, session) + del export_path, batch_dim # unused: always raises + raise NotImplementedError( + "export_functional_and_derivatives requires the TF1 Hub Module API " + "(hub.create_module_spec / hub.Module.export) that tensorflow-hub " + "0.16 removed, and would need a TF2 tf.saved_model.save " + "reimplementation instead. Only the vendored export_saved_model.py " + "and neural_numint_test.py use this method; neither is part of the " + "test suite." + ) # DM21* functionals include the hybrid term directly, so set the # range-separated and hybrid parameters expected by PySCF to 0 so PySCF @@ -802,3 +826,129 @@ def eval_xc( exc = exc.astype(np.float64) vxc = tuple(v.astype(np.float64) for v in (vxc_0, vxc_1, vxc_2, vxc_3)) return exc, vxc, fxc, kxc + + def eval_xc_eff( # pylint: disable=too-many-arguments,too-many-positional-arguments + self, + xc_code: str, + rho: np.ndarray, + deriv: int = 1, + omega: Optional[float] = None, + xctype: Optional[str] = None, + verbose=None, + spin: Optional[int] = None, + ) -> Union[List[np.ndarray], Tuple[np.ndarray, np.ndarray, None, None]]: + """Evaluates the XC energy and derivatives against density parameters. + + PySCF's nr_rks/nr_uks grid loops call this method rather than eval_xc. + It differs from eval_xc only in the derivative convention: eval_xc + returns derivatives with respect to sigma = |nabla rho|^2, while + eval_xc_eff returns them with respect to the density parameters + [rho, nabla rho, tau]. This override applies that chain rule to + eval_xc's outputs, so the DM21 evaluation itself (including the local + Hartree-Fock side effects on self._vmat_hf) stays in eval_xc. + + See pyscf.dft.numint.NumInt.eval_xc_eff for more details on the + interface. The layouts below were verified against that method + (max |diff| = 0.0 on mgga_x_tpss for both spin channels). + + Args: + xc_code: unused (see eval_xc). + rho: density and density derivatives at each grid point, shape + (5, N) [rho, nabla_x, nabla_y, nabla_z, tau], or (2, 5, N) for + spin-polarized calculations; a laplacian-carrying (6, N) / + (2, 6, N) rho is accepted as well. eval_rho runs with + MGGA_DENSITY_LAPL off, so rho arrives without the laplacian row + the graph's (6, N) placeholders expect; a zero row is inserted + before evaluation (the graph unstacks and discards it). + deriv: derivative order. 1 (the default) is supported; 0 evaluates + the energy and the local Hartree-Fock side effects but returns no + potential; greater than 1 raises. + omega: RSH parameter. None (the default) is passed through to + eval_xc, which raises for an explicit value: DM21 has no + range-separation support. + xctype: unused. The DM21 functional is always a meta-GGA. + verbose: unused. + spin: 0 for a spin-unpolarized (restricted) calculation, 1 for + spin-polarized. Inferred from rho's leading dimension when None, + the same rule as pyscf.dft.numint.NumInt.eval_xc_eff. + + Returns: + exc, vxc, fxc, kxc, where: + exc is the XC energy density at each grid point, shape (N). + vxc has shape (5, N) for spin=0, with rows + [dE/d rho, dE/d nabla rho (3 components), dE/d tau], or shape + (2, 5, N) for spin=1 with the same rows for the alpha and beta + channels (rows [rho, nabla rho, tau] of each). + fxc is None. (Second derivatives are not implemented.) + kxc is None. (Third derivatives are not implemented.) + For deriv=0 the return is [exc, None, None, None]. + + Raises: + NotImplementedError: if deriv > 1 (DM21 provides no second or higher + derivatives; the base class would fail inside transform_xc). + """ + if omega is None: + omega = self.omega + if deriv > 1: + raise NotImplementedError( + "DM21 functionals provide no second or higher derivatives, so " + "eval_xc_eff cannot be called with deriv > 1." + ) + del xctype # unused: the DM21 functional is always a meta-GGA + + rho = np.asarray(rho, order="C", dtype=np.float64) + if spin is None: + spin = 1 if (rho.ndim >= 2 and rho.shape[0] == 2) else 0 + if rho.shape[-2] == 5: + rho = np.insert(rho, 4, 0.0, axis=-2) + + exc, (vrho, vsigma, _vlapl, vtau), _, _ = self.eval_xc( + xc_code, rho, spin=spin, deriv=deriv, omega=omega, verbose=verbose + ) + if deriv < 1: + # Evaluated for exc and the _vmat_hf side effect; no potential is + # requested at derivative order 0. + return [exc, None, None, None] + + # Chain rule from eval_xc's sigma derivatives (dE/d sigma, with sigma + # the PySCF/libxc (sigma_aa, sigma_ab, sigma_bb) components) to the + # density-parameter derivatives PySCF contracts with the AOs. To + # re-derive it, follow PySCF's MGGA branches: nr_rks does + # wv = weight * vxc, wv[0] *= .5, wv[4] *= .5, + # _scale_ao_sparse(ao[:4], wv[:4], ...) and _tau_dot_sparse(ao, ao, + # wv[4], ...); nr_uks does the same per spin row (wv[:, 0] and + # wv[:, 4]) — i.e. row 0 is dE/d rho, rows 1:4 are dE/d nabla rho and + # row 4 is dE/d tau. dE/d nabla rho = 2 (dE/d sigma) nabla rho for a + # single sigma, and the pairwise expansion below for sigma_ab. + if spin == 0: + # sigma = nabla rho . nabla rho, so the gradient rows pick up a + # factor of 2. + vxc = np.concatenate( + [ + vrho[None, :], + 2.0 * vsigma[None, :] * rho[1:4], + vtau[None, :], + ], + axis=0, + ) + else: + nabla_a, nabla_b = rho[0, 1:4], rho[1, 1:4] + vsigma_aa, vsigma_ab, vsigma_bb = vsigma[:, 0], vsigma[:, 1], vsigma[:, 2] + vxc_a = np.concatenate( + [ + vrho[:, 0][None, :], + 2.0 * vsigma_aa[None, :] * nabla_a + vsigma_ab[None, :] * nabla_b, + vtau[:, 0][None, :], + ], + axis=0, + ) + vxc_b = np.concatenate( + [ + vrho[:, 1][None, :], + 2.0 * vsigma_bb[None, :] * nabla_b + vsigma_ab[None, :] * nabla_a, + vtau[:, 1][None, :], + ], + axis=0, + ) + vxc = np.stack([vxc_a, vxc_b], axis=0) + return exc, vxc, None, None diff --git a/grad_dft/functional.py b/grad_dft/functional.py index 6e5ffef..53d2afb 100644 --- a/grad_dft/functional.py +++ b/grad_dft/functional.py @@ -591,10 +591,10 @@ def dm21_densities( grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1) # LDA preprocessing data - log_rho = jnp.log2(jnp.clip(rho, a_min=clip_cte)) + log_rho = jnp.log2(jnp.clip(rho, min=clip_cte)) # GGA preprocessing data - log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, a_min=clip_cte)) / 2 + log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, min=clip_cte)) / 2 log_x_sigma = log_grad_rho_norm - 4 / 3.0 * log_rho log_u_sigma = jnp.where( jnp.greater(log_rho, jnp.log2(clip_cte)), @@ -603,7 +603,7 @@ def dm21_densities( ) # MGGA preprocessing data - log_tau = jnp.log2(jnp.clip(tau, a_min=clip_cte)) + log_tau = jnp.log2(jnp.clip(tau, min=clip_cte)) log_1t_sigma = -( 5 / 3.0 * log_rho - log_tau + 2 / 3.0 * jnp.log2(6 * jnp.pi**2) + jnp.log2(3 / 5.0) ) @@ -851,45 +851,52 @@ def generate_DM21_weights( path = os.path.dirname(os.path.dirname(__file__)) folder = os.path.join(path, folder) - variables = tf.saved_model.load(folder).variables - - def tf_tensor_to_jax(tf_tensor: tf.Tensor) -> Array: - return jnp.asarray(tf_tensor.numpy()) - - def vars_to_params(variables: List[tf.Variable]) -> PyTree: + # The vendored checkpoint is a TensorFlow 1 SavedModel, which + # tf.saved_model.load can no longer import on TF >= 2.13 (the V1 + # restore machinery it relies on was removed). Read the variables + # directly from the checkpoint files instead; the ":0" op suffix is + # appended so the name patterns below match what variables used to be + # called by tf.saved_model.load(...).variables. + checkpoint_path = os.path.join(folder, "variables", "variables") + checkpoint = tf.train.load_checkpoint(checkpoint_path) + variables = [ + (name + ":0", checkpoint.get_tensor(name)) + for name, _ in tf.train.list_variables(checkpoint_path) + ] + + def tf_tensor_to_jax(value: jnp.ndarray) -> Array: + return jnp.asarray(value) + + def vars_to_params(variables: List[tuple]) -> PyTree: import re params = {} - for var in variables: - if "ResidualBlock_" in var.name: - number = int(re.findall("ResidualBlock_[0-9]", var.name)[0][-1]) + 1 - elif "ResidualBlock/" in var.name: + for name, value in variables: + if "ResidualBlock_" in name: + number = int(re.findall("ResidualBlock_[0-9]", name)[0][-1]) + 1 + elif "ResidualBlock/" in name: number = 1 - elif "Squash" in var.name: + elif "Squash" in name: number = 0 - elif "Output" in var.name: + elif "Output" in name: number = 7 else: - raise ValueError("Unknown variable name.") + raise ValueError(f"Unknown variable name {name}.") - if "/linear/" in var.name: + if "/linear/" in name: if "Dense_" + str(number) not in params.keys(): params["Dense_" + str(number)] = {} - if "/w:" in var.name: - params["Dense_" + str(number)]["kernel"] = tf_tensor_to_jax(var.value()) - elif "/b:" in var.name: - params["Dense_" + str(number)]["bias"] = tf_tensor_to_jax(var.value()) - elif "/layer_norm/" in var.name: + if "/w:" in name: + params["Dense_" + str(number)]["kernel"] = tf_tensor_to_jax(value) + elif "/b:" in name: + params["Dense_" + str(number)]["bias"] = tf_tensor_to_jax(value) + elif "/layer_norm/" in name: if "LayerNorm_" + str(number - 1) not in params.keys(): params["LayerNorm_" + str(number - 1)] = {} - if "gamma:" in var.name: - params["LayerNorm_" + str(number - 1)]["scale"] = tf_tensor_to_jax( - var.value() - ) - elif "beta:" in var.name: - params["LayerNorm_" + str(number - 1)]["bias"] = tf_tensor_to_jax( - var.value() - ) + if "gamma:" in name: + params["LayerNorm_" + str(number - 1)]["scale"] = tf_tensor_to_jax(value) + elif "beta:" in name: + params["LayerNorm_" + str(number - 1)]["bias"] = tf_tensor_to_jax(value) return params example_features = normal(rng, shape=(2, n_input_features)) @@ -1005,7 +1012,7 @@ def correlation_polarization_correction( The ready to be integrated electronic energy density. """ - log_rho = jnp.log2(jnp.clip(rho.sum(axis=1), a_min=clip_cte)) + log_rho = jnp.log2(jnp.clip(rho.sum(axis=1), min=clip_cte)) # assert not jnp.isnan(log_rho).any() and not jnp.isinf(log_rho).any() log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0 @@ -1033,7 +1040,7 @@ def fzeta(z): # assert not jnp.isnan(alphac).any() and not jnp.isinf(alphac).any() fz = fzeta(zeta) #jnp.round(fzeta(zeta), int(math.log10(clip_cte))) - z4 = zeta**4 #jnp.round(2 ** (4 * jnp.log2(jnp.clip(zeta, a_min=clip_cte))), int(math.log10(clip_cte))) + z4 = zeta**4 #jnp.round(2 ** (4 * jnp.log2(jnp.clip(zeta, min=clip_cte))), int(math.log10(clip_cte))) e_tilde = ( e_tilde_PF[:, 0] @@ -1105,10 +1112,10 @@ def densities( grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1) # LDA preprocessing data - log_rho = jnp.log2(jnp.clip(rho, a_min=clip_cte)) + log_rho = jnp.log2(jnp.clip(rho, min=clip_cte)) # GGA preprocessing data - log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, a_min=clip_cte)) / 2 + log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, min=clip_cte)) / 2 log_x_sigma = log_grad_rho_norm - 4 / 3.0 * log_rho log_u_sigma = jnp.where( jnp.greater(log_rho, jnp.log2(clip_cte)), @@ -1117,7 +1124,7 @@ def densities( ) # MGGA preprocessing data - log_tau = jnp.log2(jnp.clip(tau, a_min=clip_cte)) + log_tau = jnp.log2(jnp.clip(tau, min=clip_cte)) log_1t_sigma = log_tau - 5 / 3.0 * log_rho log_w_sigma = jnp.where( jnp.greater(log_rho, jnp.log2(clip_cte)), @@ -1136,8 +1143,8 @@ def densities( ######### Correlation features ############### grad_rho_norm_sq_ss = jnp.sum((grad_rho.sum(axis=1)) ** 2, axis=-1) - log_grad_rho_norm_ss = jnp.log2(jnp.clip(grad_rho_norm_sq_ss, a_min=clip_cte)) / 2 - log_rho_ss = jnp.log2(jnp.clip(rho.sum(axis=1), a_min=clip_cte)) + log_grad_rho_norm_ss = jnp.log2(jnp.clip(grad_rho_norm_sq_ss, min=clip_cte)) / 2 + log_rho_ss = jnp.log2(jnp.clip(rho.sum(axis=1), min=clip_cte)) log_x_ss = log_grad_rho_norm_ss - 4 / 3.0 * log_rho_ss log_u_ss = jnp.where( @@ -1154,7 +1161,7 @@ def densities( log_u_c = jnp.stack((log_u_ss, log_u_ab), axis=1) - log_tau_ss = jnp.log2(jnp.clip(tau.sum(axis=1), a_min=clip_cte)) + log_tau_ss = jnp.log2(jnp.clip(tau.sum(axis=1), min=clip_cte)) log_1t_ss = log_tau_ss - 5 / 3.0 * log_rho_ss log_w_ss = jnp.where( jnp.greater(log_rho.sum(axis=1), jnp.log2(clip_cte)), @@ -1177,7 +1184,7 @@ def densities( beta3 = jnp.array([[1.6382, 3.3662]]) beta4 = jnp.array([[0.49294, 0.62517]]) - log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), a_min=clip_cte)) + log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), min=clip_cte)) log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0 brs_1_2 = 2 ** (log_rs / 2 + jnp.log2(beta1)) ars = 2 ** (log_rs + jnp.log2(alpha1)) diff --git a/grad_dft/popular_functionals.py b/grad_dft/popular_functionals.py index 69d4c03..c11fe59 100644 --- a/grad_dft/popular_functionals.py +++ b/grad_dft/popular_functionals.py @@ -38,7 +38,7 @@ def lsda_x_e(rho: Float[Array, "grid spin"], clip_cte) -> Float[Array, "grid"]: ------- Float[Array, "grid"] """ - rho = jnp.clip(rho, a_min=clip_cte) + rho = jnp.clip(rho, min=clip_cte) lda_es = ( -3.0 / 4.0 @@ -69,14 +69,14 @@ def b88_x_e(rho: Float[Array, "grid spin"], grad_rho: Float[Array, "grid spin di beta = 0.0042 - rho = jnp.clip(rho, a_min=clip_cte) + rho = jnp.clip(rho, min=clip_cte) # LDA preprocessing data: Note that we duplicate the density to sum and divide in the last eq. - log_rho = jnp.log2(jnp.clip(rho, a_min=clip_cte)) + log_rho = jnp.log2(jnp.clip(rho, min=clip_cte)) grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1) - log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, a_min=clip_cte)) / 2 + log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, min=clip_cte)) / 2 # GGA preprocessing data log_x_sigma = log_grad_rho_norm - 4 / 3.0 * log_rho @@ -124,7 +124,7 @@ def pw92_c_e(rho: Float[Array, "grid spin"], clip_cte: float = 1e-30) -> Float[A beta3 = jnp.array([[1.6382, 3.3662]]) beta4 = jnp.array([[0.49294, 0.62517]]) - log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), a_min=clip_cte)) + log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), min=clip_cte)) log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0 brs_1_2 = 2 ** (log_rs / 2 + jnp.log2(beta1)) ars = 2 ** (log_rs + jnp.log2(alpha1)) @@ -161,7 +161,7 @@ def vwn_c_e(rho: Float[Array, "grid spin"], clip_cte: float = 1e-30) -> Float[Ar x0 = jnp.array([[-0.10498, -0.325]]) rho = jnp.where(rho > clip_cte, rho, 0.0) - log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), a_min=clip_cte)) + log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), min=clip_cte)) # assert not jnp.isnan(log_rho).any() and not jnp.isinf(log_rho).any() log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0 log_x = log_rs / 2 @@ -232,7 +232,7 @@ def lyp_c_e(rho: Float[Array, "grid spin"], grad_rho: Float[Array, "grid spin 3" d = 0.349 CF = (3 / 10) * (3 * jnp.pi**2) ** (2 / 3) - rho = jnp.clip(rho, a_min=clip_cte) + rho = jnp.clip(rho, min=clip_cte) grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1) diff --git a/grad_dft/train.py b/grad_dft/train.py index e81680f..6fbcc2a 100644 --- a/grad_dft/train.py +++ b/grad_dft/train.py @@ -420,7 +420,7 @@ def dm21_grad_regularization(molecule: Molecule, F: Float[Array, "spin ao ao"]) prefactors = numerator / safe_denominator - dE = jnp.clip(0.5 * jnp.sum(prefactors * factors), a_min=-10, a_max=10) + dE = jnp.clip(0.5 * jnp.sum(prefactors * factors), min=-10, max=10) return dE**2 @@ -579,10 +579,10 @@ def sq_electron_err_int( Scalar: the value epsilon described above ---------- """ - pred_density = jnp.clip(pred_density, a_min=clip_cte) - truth_density = jnp.clip(truth_density, a_min=clip_cte) - diff_up = jnp.clip(jnp.clip(pred_density[:, 0] - truth_density[:, 0], a_min=clip_cte) ** 2, a_min=clip_cte) - diff_dn = jnp.clip(jnp.clip(pred_density[:, 1] - truth_density[:, 1], a_min=clip_cte) ** 2, a_min=clip_cte) + pred_density = jnp.clip(pred_density, min=clip_cte) + truth_density = jnp.clip(truth_density, min=clip_cte) + diff_up = jnp.clip(jnp.clip(pred_density[:, 0] - truth_density[:, 0], min=clip_cte) ** 2, min=clip_cte) + diff_dn = jnp.clip(jnp.clip(pred_density[:, 1] - truth_density[:, 1], min=clip_cte) ** 2, min=clip_cte) err_int = jnp.sum(diff_up * atoms.grid.weights) + jnp.sum(diff_dn * atoms.grid.weights) return err_int diff --git a/requirements.txt b/requirements.txt index a8f10c0..212febf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,10 @@ -jax>=0.4.14 -jaxlib>=0.4.14 +jax>=0.4.27 +jaxlib>=0.4.27 pyscf>=2.3.0 attrs>=23.1.0 flax>=0.7.2 tensorflow>=2.13.0 -tensorflow-hub>=0.14.0 +tensorflow-hub>=0.16.0 typeguard==2.13.3 typing_extensions>=4.8.0 jaxtyping diff --git a/setup.py b/setup.py index 199197f..42ed03b 100644 --- a/setup.py +++ b/setup.py @@ -42,11 +42,12 @@ "openpyxl" ], }, + python_requires=">=3.12", classifiers=[ "Intended Audience :: Developers", - "Intended Audience :: Science/Research" - "License :: OSI Approved :: Apache 2.0", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: Apache 2.0", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3 :: Only", ], ) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/molecules/__init__.py b/tests/integration/molecules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/molecules/test_functional_implementations.py b/tests/integration/molecules/test_functional_implementations.py index b12f1ec..7dd5244 100644 --- a/tests/integration/molecules/test_functional_implementations.py +++ b/tests/integration/molecules/test_functional_implementations.py @@ -182,4 +182,31 @@ def test_dm21(mol): dm21diff = (ground_truth_energy - predicted_e) * Hartree2kcalmol assert not jnp.isnan(fock).any() - assert jnp.allclose(dm21diff, 0, atol=1) \ No newline at end of file + assert jnp.allclose(dm21diff, 0, atol=1) + + +def test_dm21_rks(): + """Restricted SCF must agree with unrestricted on the closed-shell molecule. + + DM21's restricted path maps inputs to rho_a = rho_b = rho/2, and + get_hf_density splits a restricted density matrix into dm/2 per spin, so on + a closed-shell molecule the restricted and unrestricted self-consistent + solutions must agree; the unrestricted path (spin=1) is already validated + end-to-end by test_dm21. A wrong spin=0 derivative layout in eval_xc_eff + therefore shows up as an energy mismatch (or a diverged/NaN SCF) here, not + as a silent pass. + """ + mol_rks = mols[0] # closed-shell HF; the open-shell Li atom (mols[1]) is + # spin-polarized, so not a valid restricted Kohn-Sham molecule. + + mf_uks = dft.UKS(mol_rks) + mf_uks._numint = NeuralNumInt(Functional.DM21) # pylint: disable=protected-access + e_uks = mf_uks.kernel() + + mf_rks = dft.RKS(mol_rks) + mf_rks._numint = NeuralNumInt(Functional.DM21) # pylint: disable=protected-access + e_rks = mf_rks.kernel() + + assert not jnp.isnan(e_uks) + assert not jnp.isnan(e_rks) + assert jnp.allclose((e_rks - e_uks) * Hartree2kcalmol, 0, atol=1) \ No newline at end of file diff --git a/tests/integration/molecules/test_non_xc_energy.py b/tests/integration/molecules/test_non_xc_energy.py index 3c75ebc..50636ba 100644 --- a/tests/integration/molecules/test_non_xc_energy.py +++ b/tests/integration/molecules/test_non_xc_energy.py @@ -146,7 +146,8 @@ def molecule_and_energies(geom: str) -> tuple[float, float]: """ mf = dft.RKS(geom) mf.xc = "0.00*LDA" # quick way of having no XC energy in PySCF - E_pyscf = mf.kernel(max_cycle=SCF_ITERS) + mf.max_cycle = SCF_ITERS + E_pyscf = mf.kernel() molecule = molecule_from_pyscf(mf, scf_iteration=SCF_ITERS) E_gdft = molecule.nonXC() return E_pyscf, E_gdft diff --git a/tests/integration/molecules/test_predict_B3LYP.py b/tests/integration/molecules/test_predict_B3LYP.py index fdab7d3..5611653 100644 --- a/tests/integration/molecules/test_predict_B3LYP.py +++ b/tests/integration/molecules/test_predict_B3LYP.py @@ -19,7 +19,8 @@ import pytest # This only works on startup! -from jax.config import config +from jax import config + config.update("jax_enable_x64", True) from grad_dft import molecule_from_pyscf, scf_loop @@ -56,6 +57,7 @@ # This test is now NOT included in the CI because of implementation differences between B3LYP in Grad DFT # versus PySCF. See above. + @pytest.mark.parametrize("mol_and_name", [(MOL_WATER, "water"), (MOL_LI, "Li")]) def test_predict(mol_and_name: tuple[gto.Mole, str]) -> None: r"""Compare the total energy predicted by Grad-DFT for the B3LYP functional versus PySCF. @@ -79,4 +81,6 @@ def test_predict(mol_and_name: tuple[gto.Mole, str]) -> None: molecule_out = iterator(PARAMS, molecule) e_XND = molecule_out.energy kcalmoldiff = (e_XND - e_DM) * Hartree2kcalmol - assert jnp.allclose(kcalmoldiff, 0, atol=10), f"Energy difference with PySCF for B3LYP on {name} exceeds the threshold." \ No newline at end of file + assert jnp.allclose( + kcalmoldiff, 0, atol=10 + ), f"Energy difference with PySCF for B3LYP on {name} exceeds the threshold." diff --git a/tests/integration/molecules/test_predict_B88.py b/tests/integration/molecules/test_predict_B88.py index a5b5b44..fe227c1 100644 --- a/tests/integration/molecules/test_predict_B88.py +++ b/tests/integration/molecules/test_predict_B88.py @@ -19,14 +19,11 @@ import pytest # This only works on startup! -from jax.config import config +from jax import config + config.update("jax_enable_x64", True) -from grad_dft import ( - molecule_from_pyscf, - scf_loop, - diff_scf_loop -) +from grad_dft import molecule_from_pyscf, scf_loop, diff_scf_loop from grad_dft.utils.types import Hartree2kcalmol from grad_dft.popular_functionals import B88 @@ -50,6 +47,7 @@ SCF_ITERS = 10 + @pytest.mark.parametrize("mol_and_name", [(MOL_WATER, "water"), (MOL_LI, "Li")]) def test_predict(mol_and_name: tuple[gto.Mole, str]) -> None: r"""Compare the total energy predicted by Grad-DFT for the B88 functional versus PySCF. @@ -63,11 +61,11 @@ def test_predict(mol_and_name: tuple[gto.Mole, str]) -> None: mf = dft.RKS(mol) else: mf = dft.UKS(mol) - + mf.xc = "B88" mf.max_cycle = 10 e_DM = mf.kernel() - + # Start from Non-SCF density molecule = molecule_from_pyscf(mf, energy=e_DM, omegas=[], scf_iteration=0) @@ -75,7 +73,9 @@ def test_predict(mol_and_name: tuple[gto.Mole, str]) -> None: molecule_out = iterator(PARAMS, molecule) e_XND = molecule_out.energy kcalmoldiff = (e_XND - e_DM) * Hartree2kcalmol - assert jnp.allclose(kcalmoldiff, 0, atol=1e-6), f"Energy difference with PySCF for B88 on {name} exceeds the threshold." + assert jnp.allclose( + kcalmoldiff, 0, atol=1e-6 + ), f"Energy difference with PySCF for B88 on {name} exceeds the threshold." @pytest.mark.parametrize("mol_and_name", [(MOL_WATER, "water"), (MOL_LI, "Li")]) @@ -87,23 +87,27 @@ def test_jit(mol_and_name: tuple[gto.Mole, str]) -> None: mol_and_name (tuple[gto.Mole, str]): PySCF molecule object and the name of the molecule. """ mol, name = mol_and_name - if mol.spin == 0: mf = dft.RKS(mol) - else: mf = dft.UKS(mol) - + if mol.spin == 0: + mf = dft.RKS(mol) + else: + mf = dft.UKS(mol) + mf.xc = "B88" mf.max_cycle = 0 mf.kernel() - + # Start from Non-SCF density molecule = molecule_from_pyscf(mf, omegas=[], scf_iteration=0) iterator = scf_loop(FUNCTIONAL, verbose=2, cycles=10) molecule_out = iterator(PARAMS, molecule) e_XND = molecule_out.energy - + iterator = diff_scf_loop(FUNCTIONAL, cycles=10) molecule_out = iterator(PARAMS, molecule) e_XND_jit = molecule_out.energy kcalmoldiff = (e_XND - e_XND_jit) * Hartree2kcalmol - assert jnp.allclose(kcalmoldiff, 0, atol=1e-6), f"Energy difference with between jitted and non-jitted SCF for B88 on {name} exceeds the threshold." \ No newline at end of file + assert jnp.allclose( + kcalmoldiff, 0, atol=1e-6 + ), f"Energy difference with between jitted and non-jitted SCF for B88 on {name} exceeds the threshold." diff --git a/tests/integration/molecules/test_predict_DM21.py b/tests/integration/molecules/test_predict_DM21.py index c260559..c237bf8 100644 --- a/tests/integration/molecules/test_predict_DM21.py +++ b/tests/integration/molecules/test_predict_DM21.py @@ -18,7 +18,8 @@ import pytest # This only works on startup! -from jax.config import config +from jax import config + config.update("jax_enable_x64", True) dirpath = os.path.dirname(os.path.dirname(__file__)) @@ -32,11 +33,7 @@ learning_rate = 1e-3 -from grad_dft import ( - molecule_from_pyscf, - scf_loop, - DM21 -) +from grad_dft import molecule_from_pyscf, scf_loop, DM21 from grad_dft.utils.types import Hartree2kcalmol from grad_dft.external import NeuralNumInt @@ -69,7 +66,9 @@ def test_predict(mol): ## Load the molecule, RKS warnings.warn("Remember to set the grid level to 3 in the config file!") - molecule = molecule_from_pyscf(mf, energy=energy, omegas=[0.0, 0.4], scf_iteration=0) + molecule = molecule_from_pyscf( + mf, energy=energy, omegas=[0.0, 0.4], scf_iteration=0 + ) # tx = adam(learning_rate = learning_rate) # iterator = mol_orb_optimizer(functional, tx, omegas = [0., 0.4], verbose = 2, functional_type = 'DM21') @@ -206,9 +205,7 @@ def test_uks(): molecule_name = "CoC" mol = gto.Mole() mol.atom = [["Co", [0, 0, 0]], ["C", [1.56, 0, 0]]] # def2-tzvp -mol.basis = ( - "def2-tzvp" # basis_set_exchange.api.get_basis(name='cc-pvdz', fmt='nwchem', elements='Co') -) +mol.basis = "def2-tzvp" # basis_set_exchange.api.get_basis(name='cc-pvdz', fmt='nwchem', elements='Co') mol.spin = 1 mol.unit = "angstrom" mol.build() diff --git a/tests/integration/molecules/test_training.py b/tests/integration/molecules/test_training.py index d9e2a20..425c20e 100644 --- a/tests/integration/molecules/test_training.py +++ b/tests/integration/molecules/test_training.py @@ -73,7 +73,8 @@ for mol in PYSCF_MOLS: mf = pyscf.dft.UKS(mol) mf.xc = "LDA" - E_pyscf = mf.kernel(max_cycle=SCF_ITERS) + mf.max_cycle = SCF_ITERS + E_pyscf = mf.kernel() TRUTH_ENERGIES.append(E_pyscf) molecule = molecule_from_pyscf(mf) MOLECULES.append(molecule) @@ -95,14 +96,14 @@ # Define a simple neural functional and its initial parameters def coefficient_inputs(molecule: Molecule, clip_cte: float = 1e-30, *_, **__): - rho = jnp.clip(molecule.density(), a_min = clip_cte) + rho = jnp.clip(molecule.density(), min=clip_cte) return jnp.concatenate((rho, ), axis = 1) def energy_densities(molecule: Molecule, clip_cte: float = 1e-30, *_, **__): r"""Auxiliary function to generate the features of LSDA.""" rho = molecule.density() # To avoid numerical issues in JAX we limit too small numbers. - rho = jnp.clip(rho, a_min = clip_cte) + rho = jnp.clip(rho, min=clip_cte) # Now we can implement the LSDA exchange energy density lda_e = -3/2 * (3/(4*jnp.pi)) ** (1/3) * (rho**(4/3)).sum(axis = 1, keepdims = True) return lda_e diff --git a/tests/integration/solids/__init__.py b/tests/integration/solids/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/solids/test_non_xc_energy.py b/tests/integration/solids/test_non_xc_energy.py index 4dff83e..f2a9645 100644 --- a/tests/integration/solids/test_non_xc_energy.py +++ b/tests/integration/solids/test_non_xc_energy.py @@ -88,7 +88,8 @@ def solid_and_energies(geom) -> tuple[float, float]: """ kmf = dft.KRKS(geom, kpts=geom.make_kpts(KPTS)) kmf.xc = "0.00*LDA" # quick way of having no XC energy in PySCF - E_pyscf = kmf.kernel(max_cycle=SCF_ITERS) + kmf.max_cycle = SCF_ITERS + E_pyscf = kmf.kernel() sol = solid_from_pyscf(kmf) E_gdft = sol.nonXC() return E_pyscf, E_gdft diff --git a/tests/integration/solids/test_training.py b/tests/integration/solids/test_training.py index ace766a..4825e80 100644 --- a/tests/integration/solids/test_training.py +++ b/tests/integration/solids/test_training.py @@ -90,7 +90,8 @@ for sol in PYSCF_SOLS: kmf = scf.KRKS(sol, kpts=sol.make_kpts(KPTS)) kmf.xc = "LDA" - E_pyscf = kmf.kernel(max_cycle=SCF_ITERS) + kmf.max_cycle = SCF_ITERS + E_pyscf = kmf.kernel() solid = solid_from_pyscf(kmf) SOLIDS.append(solid) @@ -102,7 +103,8 @@ # DFT calculations used for their grids to calculate MP2 densities kmf_dft_dummy = scf.KRKS(sol, kpts=sol.make_kpts(KPTS)) - kmf_dft_dummy.kernel(max_cycle=1) + kmf_dft_dummy.max_cycle = 1 + kmf_dft_dummy.kernel() grad_dft_sol_dummy = solid_from_pyscf(kmf_dft_dummy) dft_kccsd_rdm1 = grad_dft_sol_dummy.replace(rdm1=jnp.asarray(mp2_rdm1)) # Works because we use the same AOs for DFT and MP2 @@ -114,14 +116,14 @@ # Define a simple neural functional and its initial parameters def coefficient_inputs(solid: Solid, clip_cte: float = 1e-30, *_, **__): - rho = jnp.clip(solid.density(), a_min = clip_cte) + rho = jnp.clip(solid.density(), min=clip_cte) return jnp.concatenate((rho, ), axis = 1) def energy_densities(solid: Solid, clip_cte: float = 1e-30, *_, **__): r"""Auxiliary function to generate the features of LSDA.""" rho = solid.density() # To avoid numerical issues in JAX we limit too small numbers. - rho = jnp.clip(rho, a_min = clip_cte) + rho = jnp.clip(rho, min=clip_cte) # Now we can implement the LSDA exchange energy density lda_e = -3/2 * (3/(4*jnp.pi)) ** (1/3) * (rho**(4/3)).sum(axis = 1, keepdims = True) return lda_e diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..e69de29