From 79a19054e830a1aeac5837907c2d6bccc184da85 Mon Sep 17 00:00:00 2001 From: Anya497 Date: Mon, 5 Oct 2026 22:10:56 +0300 Subject: [PATCH 01/17] Agent v1. --- .opencode/skills/code-review/SKILL.md | 42 ++++ .opencode/skills/planning/SKILL.md | 185 +++++++++++++++ .opencode/skills/quality-gates/SKILL.md | 50 ++++ .opencode/skills/run-tests/SKILL.md | 21 ++ .opencode/skills/subtask-loop/SKILL.md | 223 ++++++++++++++++++ .../skills/user-guidance-transfer/SKILL.md | 30 +++ .opencode/skills/workflow-management/SKILL.md | 144 +++++++++++ AGENTS.md | 49 ++++ 8 files changed, 744 insertions(+) create mode 100644 .opencode/skills/code-review/SKILL.md create mode 100644 .opencode/skills/planning/SKILL.md create mode 100644 .opencode/skills/quality-gates/SKILL.md create mode 100644 .opencode/skills/run-tests/SKILL.md create mode 100644 .opencode/skills/subtask-loop/SKILL.md create mode 100644 .opencode/skills/user-guidance-transfer/SKILL.md create mode 100644 .opencode/skills/workflow-management/SKILL.md create mode 100644 AGENTS.md diff --git a/.opencode/skills/code-review/SKILL.md b/.opencode/skills/code-review/SKILL.md new file mode 100644 index 0000000..616a833 --- /dev/null +++ b/.opencode/skills/code-review/SKILL.md @@ -0,0 +1,42 @@ +--- +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 main)`): + `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 + (one commit per subtask/fix as per `git-workflow`). +4. Repeat until zero findings. + +## Checklist + +- [ ] **Correctness** — every clause of the task issue (#N) is traceable to + committed code. +- [ ] **Tests** — new code is covered; no test was weakened or skipped; the + suite passes (see `run-tests`). +- [ ] **Docs** — docs updated per the `documentation` skill; docstrings follow + its conventions. +- [ ] **Duplication** — no copy-pasted or near-identical logic (see `reusing`). +- [ ] **Style** — `code-style` pass clean; types consistent with the codebase. +- [ ] **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..8d338d1 --- /dev/null +++ b/.opencode/skills/planning/SKILL.md @@ -0,0 +1,185 @@ +--- +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, 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 (see the "Contribution + guidelines" section of `docs/developer.rst`; `main` by default). +- 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 (types, helpers, utilities) that multiple + tasks need. Create shared modules to avoid duplication across tasks. +- Identify existing reusable abstractions (types, helper functions, generators) + that new tasks should use rather than reinvent. +- **Reuse analysis**: load the `reusing` skill and run the reuse checklist + against existing code and documentation before creating new material. +- For each task that involves rendering/visualization, note that rendering and + algorithm logic must be in separate modules from the start. +- Propose an execution order that minimizes rework and avoids conflicts. +- Align tasks with the project architecture. + +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 #N]` 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. + +Before writing the detailed plan, load the `reusing` skill and run the reuse +checklist: search existing documentation and code for functions, types, and doc +sections that can be reused or generalized. Record what is reused in each +subtask's **Code** section. + +### 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: + golden, property-based, pytest, etc.> +**Docs:** <doc files to create or update. Use the `documentation` skill + to determine which docs are affected.> + +**Spec:** +- <detailed implementation specification> +``` + +Example: + +``` +### S1: Add ONNX export for the path-selector model + +**Code:** New `export_onnx` function in `AIAgent/ml/inference.py` (reuse the + existing model loading from `AIAgent/ml/models/`) +**Tests:** New `AIAgent/tests/test_onnx_export.py` verifying the exported + graph produces the same outputs as the torch model +**Docs:** Add to the matching autosummary list under `docs/reference/`; + numpydoc docstring with `Examples` + +**Spec:** +- Serialize the trained path-selector model to ONNX with dynamic batch ... +``` + +### Documentation Mapping + +For the mapping of source changes to required doc actions, load the +`documentation` skill — it is the single source of truth. Use it when writing +the **Docs** section of each subtask. + +### 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 "Contribution +guidelines" section of `docs/developer.rst`; the operational commit procedure +lives in the `git-workflow` skill. + +### 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 `aa`, the edge + (5,0)→(7,2) has no entries 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 entries are missing" — specify + which input, which state/range, and which path should produce them. +- Remaining work items must be actionable (e.g., "Track origin state through + the BFS queue by adding a field to the queue item") — not vague goals. +- If the user provided design guidance (e.g., decomposition schema, coordinate + system), 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 report artifacts + (plots, tables, exported 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" so equivalence tests + are built from the start. +- **Specify type genericity**. If a module must handle arbitrary types, state + it explicitly (e.g., "generic over the graph and feature types"). +- **Specify reuse expectations**. If the task builds on existing infrastructure + (e.g., "reuse the dataset abstractions from `AIAgent/ml/dataset.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..723ab91 --- /dev/null +++ b/.opencode/skills/quality-gates/SKILL.md @@ -0,0 +1,50 @@ +--- +name: quality-gates +description: Use before integrating a task. Defines the hard gate that must pass (tests + style + docs build) and how to interpret its result. References the CI workflows 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". In stacked +mode the same gate is also run 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, the style/lint checks, and the +docs build: + +- **Tests** — see `.github/workflows/python_tests.yaml` for the exact command + and working directories (see the `run-tests` skill). 0 failures, 0 skipped. +- **Style/lint** — see `.github/workflows/python_linting.yaml` for the exact + commands (`ruff check`, `ruff format --check`; see the `code-style` skill). +- **Docs build** — see `.github/workflows/docs.yaml`. Sphinx builds under the + no-warnings policy (`-W --keep-going`): any warning fails the build, so the + exit code is sufficient. + +The CI workflows are the source of truth for the commands they run; this skill +only defines the gate semantics. There is no type-check step and no coverage +threshold in this project. + +## Procedure + +1. Run the tests (`poetry run pytest tests -sv`) from every component the + change touches (`AIAgent/`, `tools/compstrat/`, `tools/runstrat/`). 0 + failures, 0 skipped. +2. Run `ruff check` and `ruff format --check` (see + `.github/workflows/python_linting.yaml`). No errors. +3. Build the docs (see `.github/workflows/docs.yaml`). It must exit 0 + (no-warnings policy). +4. Interpret the result: + - All clean → **PASS**. Proceed to merge (see `git-workflow`). + - 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..0b7c7cb --- /dev/null +++ b/.opencode/skills/run-tests/SKILL.md @@ -0,0 +1,21 @@ +--- +name: run-tests +description: Use when running the GradDFT test suite (pytest). Thin pointer to .github/workflows/install_and_test.yaml, which holds the exact command and working directories, plus the machine-specific pitfalls. +--- + +# Run tests + +What the test pipeline is and where the exact commands live (the CI workflow +steps) are in `.github/workflows/install_and_test.yaml`. Read that workflow before +running anything. + +## Commands + +CI runs pytest from root directory. Example: + + pytest -v tests/unit/test_eigenproblem.py + +## Notes + +- 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..85e5efe --- /dev/null +++ b/.opencode/skills/subtask-loop/SKILL.md @@ -0,0 +1,223 @@ +--- +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, code quality checks, commit rules, 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` 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 navigation links are updated | +| 4. Pre-Commit Check | **Skip** — no source files to format or lint | +| 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. + +### 2. Write Tests + +Write the tests specified in the subtask's **Tests** section. See the +`run-tests` skill for how to run them. + +### 3. Update Documentation + +**Load the `documentation` skill before updating any docs.** Update all +task-related documentation per that skill. + +**Hard gate — this step is not complete until:** + +- [ ] At least one documentation file was created or updated for this subtask +- [ ] Documentation completeness is verified per the `documentation` skill's + procedure + +### 4. Pre-Commit Check + +Run formatting and linting per the `code-style` skill. 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 if found. +- **Genericity check**: verify new code reuses existing abstractions and uses + appropriate typing (e.g., `typing` generics) instead of bespoke + re-implementations. +- **Equivalence test check**: if the subtask is a variant of an existing + algorithm, ensure a property-based equivalence test exists comparing it to + the reference implementation. +- **Separation check**: verify library modules do not contain ad-hoc string + generation or file I/O that belongs in a readwrite module. + +### 6. Commit + +**Load the `git-workflow` skill before committing.** The commit message +format lives in the "Contribution guidelines" section of `docs/developer.rst`; +the operational procedure (validation, merge) lives in the skill. + +### 7. Mark Completed + +Mark the subtask as completed in `tasks/detailed_plan.md` (`[done]` suffix + +commit hash on the subtask heading), 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. + +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 + +For each subtask, use the `todowrite` tool to track cycle steps as separate +items. **No subtask may be committed with any step still `pending`.** + +Example for subtask S1: + +``` +- "S1: Implement" -> in_progress -> completed +- "S1: Write tests" -> in_progress -> completed +- "S1: Update documentation" -> in_progress -> completed +- "S1: Pre-Commit Check (format + lint)" -> in_progress -> completed +- "S1: Quality checks" -> in_progress -> completed +- "S1: Commit" -> in_progress -> completed +``` + +## 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 mark multiple subtasks `completed` in `todowrite` before committing + each individually. + +A `todowrite` listing "S1: Implement [completed], S2: Implement [completed], S1: +Write tests [completed]" 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., "LR goto entries missing for nested + nonterminal calls"). + - 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 — + coordinate spaces, invariants, decomposition schema. 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., "Track origin state + through the BFS queue by adding a field to the queue item") — 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 its integration branch (`main`, or the +personal branch set via `graddft.integrationBranch`) — 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`: in direct mode the task's own pull request carries it; in stacked mode +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..26c281a --- /dev/null +++ b/.opencode/skills/user-guidance-transfer/SKILL.md @@ -0,0 +1,30 @@ +--- +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]: "Use CSR representation; the matrix indices must be the same +node order as the DFA from task 2." +``` + +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..314e3f0 --- /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 an integration branch (main or a personal 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: + `INTEGRATION=$(git config --get graddft.integrationBranch || echo main)`. + It is `main` in direct mode and a personal, long-lived branch in stacked + mode. Never commit directly to the integration branch. +- 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. +- In direct mode a task is integrated into `main` through its own pull + request. In stacked mode a task is rebased and fast-forward merged into the + integration branch and **no pull request is opened**; a single pull request + into `main` is opened only on the user's explicit request. +- 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. +- Documentation-only tasks (no `.py` files changed) skip code-specific gates + (tests, lint, format) but still run the docs build — for documentation + changes it is the relevant gate — and follow all other workflow rules: + 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. +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, see + `git-workflow`), then create a feature branch from the integration branch for + this single task (branching model: the "Contribution guidelines" section of + `docs/developer.rst`; procedure: `git-workflow`). +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 PR. + 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 (see `git-workflow`): + - **Direct mode** (`$INTEGRATION` is `main`): push the feature branch and + open a pull request into `main`, merged with rebase and merge + (fast-forward) so the per-subtask commits are preserved and `main` stays + linear. The user reviews and merges the PR manually; proceed only after + they confirm the merge. + - **Stacked mode** (`$INTEGRATION` is a personal branch): rebase and + fast-forward merge the feature branch into `$INTEGRATION`, then delete the + feature branch. Do **not** push for a PR and do **not** open a PR. +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 — the issue closes when the commit reaches + `main`. For the last task of a batch, the same commit also carries + `Closes #<hub>` once every sub-issue is resolved, closing the hub when the + 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. In stacked mode, keep accumulating + completed tasks on `$INTEGRATION`; open the request to `main` only via the + finalize step below, on the user's explicit request. + +## Finalize to main (stacked mode, explicit request only) + +When — and only when — the user explicitly asks to open the request to `main`: + +1. Rebase `$INTEGRATION` onto `origin/main` (see `git-workflow`) 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` (see `git-workflow`). +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..0b88fb0 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,49 @@ +# 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. + +## 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 + +## Skills + +### Domain + +| Skill | When to use | +|---|---| +| `.opencode/skills/project-navigation` | Starting navigation before exploring code | +| `.opencode/skills/run-tests` | Running the test suite | +| `.opencode/skills/code-style` | Formatting / linting before commit | + +### 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/git-workflow` | Branching, commits, merging | +| `.opencode/skills/reusing` | Finding existing code/docs to reuse, not duplicate | +| `.opencode/skills/user-guidance-transfer` | Recording user guidance verbatim | +| `.opencode/skills/documentation` | Mapping code changes to doc updates | +| `.opencode/skills/quality-gates` | The pre-merge gate that must pass | +| `.opencode/skills/code-review` | Whole-repo review before merge | \ No newline at end of file From 19bf2e02456ae79bfc60a30bdf06498a8bb6c4d8 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Mon, 5 Oct 2026 22:36:43 +0300 Subject: [PATCH 02/17] (#4-S1): Adapt agent skills and AGENTS.md to GradDFT The workflow-management skill, AGENTS.md, and the six remaining skills were copied from another project. They pointed at facts and files that do not exist here, so following them literally was impossible: every one of them referenced `docs/developer.rst`, `docs/reference/` autosummary lists, `.github/workflows/ python_tests.yaml`, `python_linting.yaml`, `docs.yaml`, `poetry`, component directories (`AIAgent/`, `tools/compstrat/`), the ONNX/path-selector worked example, and the `todowrite` tool. All of that was replaced with what this repository actually has. Facts corrected to match the repository: - CI workflow is `.github/workflows/install_and_test.yml` (not `.yaml`) and runs `pytest -v <file>` per file from the root, with `PYSCF_CONFIG_FILE` set to `.github/workflows/pyscf_conf.py`. That env var is now called out because the functional-prediction tests silently differ from CI without it. - Style gate is `.pre-commit-config.yaml` (`black` + `pylint -rn -sn --rcfile=.pylintrc`). There is no linting CI workflow. - Test layout is `tests/unit/`, `tests/integration/molecules/`, `tests/integration/solids/`. - Installation is `pip install -e .` / `pip install -e ".[examples]"`. - The docs build gate is gone: this repository has no Sphinx site, so the quality gate is tests + style only. Documentation surface replaces the Sphinx docs: `README.md`, `examples/` (`basic_notebooks`, `intermediate_notebooks`, `advanced_scripts`, `article_experiments`), and numpydoc-style docstrings. The mapping is now defined once, in the `planning` skill, and referenced from `subtask-loop` and `code-review` so it is not restated. Dangling skill references resolved by pointing at the owning artifact instead of loading a skill that does not exist here: - `code-style` -> `.pre-commit-config.yaml` + `.pylintrc` - `documentation` -> the documentation surface section in `planning` - `reusing` -> an explicit grep/glob reuse search over `grad_dft/`, `tests/`, `examples/`, seeded with `grad_dft/utils/` - `git-workflow` -> the commit format is now the single source of truth in `subtask-loop` step 6 (`(<issue>-S<n>): <summary>`); branching and integration stay in `workflow-management` - `todowrite` -> per-subtask cycle checklists inside `tasks/detailed_plan.md`, since that file is already the single source of truth and the plan is mirrored to the issue comment AGENTS.md listed 12 skills but only 7 exist; its table now lists exactly the 7 present skills, and a short "Repository Workflow Configuration" section records the integration branch so a fresh clone does not depend on local git config. Mode simplified to stacked: the integration branch is the personal branch `agent_settings`, so the direct-mode/PR-per-task branches were removed and every reference now resolves `INTEGRATION` as `git config --get graddft.integrationBranch || echo agent_settings`. That snippet is stated once in `workflow-management` and reused verbatim in `code-review` and `subtask-loop`. The local repo config was set to match. Worked examples in the skills are now GradDFT-specific (`utils/types.py` aliases, `utils/chunk.py`, a DM21 prediction test mirroring `test_predict_B88.py`) instead of ONNX export and graph-algorithm blockers. Not committed: untracked `.vscode/` and `grad_dft.egg-info/` build/editor artifacts, which are out of scope and not covered by `.gitignore`. Closes #4 --- .opencode/skills/code-review/SKILL.md | 22 +-- .opencode/skills/planning/SKILL.md | 133 ++++++++++------- .opencode/skills/quality-gates/SKILL.md | 44 +++--- .opencode/skills/run-tests/SKILL.md | 32 +++- .opencode/skills/subtask-loop/SKILL.md | 137 +++++++++++------- .../skills/user-guidance-transfer/SKILL.md | 5 +- .opencode/skills/workflow-management/SKILL.md | 80 +++++----- AGENTS.md | 22 +-- 8 files changed, 276 insertions(+), 199 deletions(-) diff --git a/.opencode/skills/code-review/SKILL.md b/.opencode/skills/code-review/SKILL.md index 616a833..1a72652 100644 --- a/.opencode/skills/code-review/SKILL.md +++ b/.opencode/skills/code-review/SKILL.md @@ -12,23 +12,27 @@ iterate until there are zero findings. ## Procedure 1. Determine the task's changes against the integration branch - (`INTEGRATION=$(git config --get graddft.integrationBranch || echo main)`): + (`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 - (one commit per subtask/fix as per `git-workflow`). +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; no test was weakened or skipped; the - suite passes (see `run-tests`). -- [ ] **Docs** — docs updated per the `documentation` skill; docstrings follow - its conventions. -- [ ] **Duplication** — no copy-pasted or near-identical logic (see `reusing`). -- [ ] **Style** — `code-style` pass clean; types consistent with the codebase. +- [ ] **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. diff --git a/.opencode/skills/planning/SKILL.md b/.opencode/skills/planning/SKILL.md index 8d338d1..27bdd00 100644 --- a/.opencode/skills/planning/SKILL.md +++ b/.opencode/skills/planning/SKILL.md @@ -1,6 +1,6 @@ --- 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, and task authoring guidelines. +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 @@ -24,26 +24,32 @@ 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 (see the "Contribution - guidelines" section of `docs/developer.rst`; `main` by default). +- 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 (types, helpers, utilities) that multiple - tasks need. Create shared modules to avoid duplication across tasks. -- Identify existing reusable abstractions (types, helper functions, generators) - that new tasks should use rather than reinvent. -- **Reuse analysis**: load the `reusing` skill and run the reuse checklist - against existing code and documentation before creating new material. -- For each task that involves rendering/visualization, note that rendering and - algorithm logic must be in separate modules from the start. +- 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. +- 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 #N]` in `tasks/global_plan.md` and +After each task integrates, mark it `[done #<N>]` in `tasks/global_plan.md` and update the hub's `<!-- global-plan -->` comment. The local file is the single source of truth; the hub comment is a pure mirror (see the `workflow-management` skill). @@ -65,10 +71,15 @@ 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. -Before writing the detailed plan, load the `reusing` skill and run the reuse -checklist: search existing documentation and code for functions, types, and doc -sections that can be reused or generalized. Record what is reused in each -subtask's **Code** section. +### 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 @@ -81,9 +92,10 @@ be executed.** **Code:** <files to create or modify, types and functions to add> **Tests:** <test files to create or modify, specific test approaches: - golden, property-based, pytest, etc.> -**Docs:** <doc files to create or update. Use the `documentation` skill - to determine which docs are affected.> + 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> @@ -92,31 +104,46 @@ be executed.** Example: ``` -### S1: Add ONNX export for the path-selector model +### S1: Add a prediction test for the DM21 functional -**Code:** New `export_onnx` function in `AIAgent/ml/inference.py` (reuse the - existing model loading from `AIAgent/ml/models/`) -**Tests:** New `AIAgent/tests/test_onnx_export.py` verifying the exported - graph produces the same outputs as the torch model -**Docs:** Add to the matching autosummary list under `docs/reference/`; - numpydoc docstring with `Examples` +**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:** -- Serialize the trained path-selector model to ONNX with dynamic batch ... +- Reuse the `B3LYP_WITH_VWN5`-style PySCF configuration path ... ``` -### Documentation Mapping +### 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. -For the mapping of source changes to required doc actions, load the -`documentation` skill — it is the single source of truth. Use it when writing -the **Docs** section of each subtask. +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 "Contribution -guidelines" section of `docs/developer.rst`; the operational commit procedure -lives in the `git-workflow` skill. +— 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 @@ -139,8 +166,9 @@ Design rationale, coordinate spaces, invariants — as confirmed by the user. ### <Failure Topic Title> -- Root causes with concrete examples (e.g., "for input `aa`, the edge - (5,0)→(7,2) has no entries because...") +- 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 @@ -151,12 +179,13 @@ 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 entries are missing" — specify +- 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., "Track origin state through - the BFS queue by adding a field to the queue item") — not vague goals. -- If the user provided design guidance (e.g., decomposition schema, coordinate - system), record it verbatim in the `### <Topic>` section as the +- 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 @@ -164,20 +193,22 @@ Requirements: 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 report artifacts - (plots, tables, exported models), include the exact layout, units, and - formatting rules in the task description. +- **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" so equivalence tests - are built from the start. -- **Specify type genericity**. If a module must handle arbitrary types, state - it explicitly (e.g., "generic over the graph and feature types"). + 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 dataset abstractions from `AIAgent/ml/dataset.py`"), name - the dependencies. This prevents reinvention. + (e.g., "reuse the array aliases from `grad_dft/utils/types.py`"), name the + dependencies. This prevents reinvention. ## Task Completeness Verification diff --git a/.opencode/skills/quality-gates/SKILL.md b/.opencode/skills/quality-gates/SKILL.md index 723ab91..0f4b327 100644 --- a/.opencode/skills/quality-gates/SKILL.md +++ b/.opencode/skills/quality-gates/SKILL.md @@ -1,43 +1,37 @@ --- name: quality-gates -description: Use before integrating a task. Defines the hard gate that must pass (tests + style + docs build) and how to interpret its result. References the CI workflows for the exact commands. +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". In stacked -mode the same gate is also run over the whole `main...integration` diff -immediately before the final pull request to `main`. +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, the style/lint checks, and the -docs build: +The gate is the combination of the test suite and the style/lint checks: -- **Tests** — see `.github/workflows/python_tests.yaml` for the exact command - and working directories (see the `run-tests` skill). 0 failures, 0 skipped. -- **Style/lint** — see `.github/workflows/python_linting.yaml` for the exact - commands (`ruff check`, `ruff format --check`; see the `code-style` skill). -- **Docs build** — see `.github/workflows/docs.yaml`. Sphinx builds under the - no-warnings policy (`-W --keep-going`): any warning fails the build, so the - exit code is sufficient. +- **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`. -The CI workflows are the source of truth for the commands they run; this skill -only defines the gate semantics. There is no type-check step and no coverage -threshold in this project. +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. Run the tests (`poetry run pytest tests -sv`) from every component the - change touches (`AIAgent/`, `tools/compstrat/`, `tools/runstrat/`). 0 - failures, 0 skipped. -2. Run `ruff check` and `ruff format --check` (see - `.github/workflows/python_linting.yaml`). No errors. -3. Build the docs (see `.github/workflows/docs.yaml`). It must exit 0 - (no-warnings policy). -4. Interpret the result: - - All clean → **PASS**. Proceed to merge (see `git-workflow`). +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 diff --git a/.opencode/skills/run-tests/SKILL.md b/.opencode/skills/run-tests/SKILL.md index 0b7c7cb..08d2d47 100644 --- a/.opencode/skills/run-tests/SKILL.md +++ b/.opencode/skills/run-tests/SKILL.md @@ -1,21 +1,41 @@ --- name: run-tests -description: Use when running the GradDFT test suite (pytest). Thin pointer to .github/workflows/install_and_test.yaml, which holds the exact command and working directories, plus the machine-specific pitfalls. +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 -What the test pipeline is and where the exact commands live (the CI workflow -steps) are in `.github/workflows/install_and_test.yaml`. Read that workflow before -running anything. +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 from root directory. Example: +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: - pytest -v tests/unit/test_eigenproblem.py +```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 index 85e5efe..b2ee4da 100644 --- a/.opencode/skills/subtask-loop/SKILL.md +++ b/.opencode/skills/subtask-loop/SKILL.md @@ -1,6 +1,6 @@ --- 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, code quality checks, commit rules, blocked work protocol, and completion tracking. +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 @@ -15,15 +15,15 @@ until it is verified complete.** ## Documentation-Only Subtasks -When a subtask modifies only `.md` files (no `.py` files), the following cycle -steps are adapted: +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 navigation links are updated | -| 4. Pre-Commit Check | **Skip** — no source files to format or lint | +| 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** | @@ -33,54 +33,78 @@ into a single commit. ### 1. Implement -Write the code specified in the subtask's **Code** section. +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. +`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 -**Load the `documentation` skill before updating any docs.** Update all -task-related documentation per that skill. +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:** -- [ ] At least one documentation file was created or updated for this subtask -- [ ] Documentation completeness is verified per the `documentation` skill's - procedure +- [ ] 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 formatting and linting per the `code-style` skill. Do not commit until it +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 if found. -- **Genericity check**: verify new code reuses existing abstractions and uses - appropriate typing (e.g., `typing` generics) instead of bespoke + 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 a property-based equivalence test exists comparing it to - the reference implementation. -- **Separation check**: verify library modules do not contain ad-hoc string - generation or file I/O that belongs in a readwrite module. + 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 -**Load the `git-workflow` skill before committing.** The commit message -format lives in the "Contribution guidelines" section of `docs/developer.rst`; -the operational procedure (validation, merge) lives in the skill. +**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), 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: +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) @@ -92,6 +116,8 @@ gh api -X PATCH "repos/$REPO/issues/comments/$CID" \ ``` 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 @@ -117,20 +143,20 @@ done. Do not proceed. Report it as blocked. ## Per-Subtask Execution Tracking -For each subtask, use the `todowrite` tool to track cycle steps as separate -items. **No subtask may be committed with any step still `pending`.** - -Example for subtask S1: +Track the cycle steps as a checklist inside the subtask's own section of +`tasks/detailed_plan.md`, e.g. for S1: ``` -- "S1: Implement" -> in_progress -> completed -- "S1: Write tests" -> in_progress -> completed -- "S1: Update documentation" -> in_progress -> completed -- "S1: Pre-Commit Check (format + lint)" -> in_progress -> completed -- "S1: Quality checks" -> in_progress -> completed -- "S1: Commit" -> in_progress -> completed +- [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 @@ -139,12 +165,11 @@ 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 mark multiple subtasks `completed` in `todowrite` before committing - each individually. +3. Never tick multiple subtasks' steps before committing each individually. -A `todowrite` listing "S1: Implement [completed], S2: Implement [completed], S1: -Write tests [completed]" indicates skipped commits — each subtask must be fully -committed before the next begins. +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 @@ -155,8 +180,8 @@ 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., "LR goto entries missing for nested - nonterminal calls"). + - 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. @@ -166,21 +191,21 @@ failing tests to make the suite green. Instead: `planning` skill for the full template. Minimum required content: - **Correct Design**: algorithmic design as confirmed by the user — - coordinate spaces, invariants, decomposition schema. Quote the user's - design guidance verbatim where available. + 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., "Track origin state - through the BFS queue by adding a field to the queue item") — not vague - goals. + - **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. + 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. @@ -188,8 +213,9 @@ failing tests to make the suite green. Instead: 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 its integration branch (`main`, or the -personal branch set via `graddft.integrationBranch`) — verify: +(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. @@ -209,10 +235,9 @@ did" does not mean the task is done if other parts were reverted. 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`: in direct mode the task's own pull request carries it; in stacked mode -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. +`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 diff --git a/.opencode/skills/user-guidance-transfer/SKILL.md b/.opencode/skills/user-guidance-transfer/SKILL.md index 26c281a..6c25945 100644 --- a/.opencode/skills/user-guidance-transfer/SKILL.md +++ b/.opencode/skills/user-guidance-transfer/SKILL.md @@ -23,8 +23,9 @@ verbatim. Example comment: ``` -[USER GUIDANCE]: "Use CSR representation; the matrix indices must be the same -node order as the DFA from task 2." +[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 index 314e3f0..c1c52d0 100644 --- a/.opencode/skills/workflow-management/SKILL.md +++ b/.opencode/skills/workflow-management/SKILL.md @@ -1,6 +1,6 @@ --- 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 an integration branch (main or a personal branch): global plan, one task at a time, feature branch, detailed plan, subtask execution, code review, quality gate, integration, final pull request. +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 @@ -26,25 +26,31 @@ step. 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: - `INTEGRATION=$(git config --get graddft.integrationBranch || echo main)`. - It is `main` in direct mode and a personal, long-lived branch in stacked - mode. Never commit directly to the integration branch. + + ```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. -- In direct mode a task is integrated into `main` through its own pull - request. In stacked mode a task is rebased and fast-forward merged into the - integration branch and **no pull request is opened**; a single pull request - into `main` is opened only on the user's explicit request. +- 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. + changes (see the commit format in the `subtask-loop` skill). - Documentation-only tasks (no `.py` files changed) skip code-specific gates - (tests, lint, format) but still run the docs build — for documentation - changes it is the relevant gate — and follow all other workflow rules: - one task per branch, one commit per subtask, code review. + (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 @@ -70,17 +76,16 @@ step. `<!-- 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. + 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, see - `git-workflow`), then create a feature branch from the integration branch for - this single task (branching model: the "Contribution guidelines" section of - `docs/developer.rst`; procedure: `git-workflow`). +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 @@ -89,7 +94,7 @@ step. 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 PR. + 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. @@ -107,38 +112,33 @@ step. 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 (see `git-workflow`): - - **Direct mode** (`$INTEGRATION` is `main`): push the feature branch and - open a pull request into `main`, merged with rebase and merge - (fast-forward) so the per-subtask commits are preserved and `main` stays - linear. The user reviews and merges the PR manually; proceed only after - they confirm the merge. - - **Stacked mode** (`$INTEGRATION` is a personal branch): rebase and - fast-forward merge the feature branch into `$INTEGRATION`, then delete the - feature branch. Do **not** push for a PR and do **not** open a PR. +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 — the issue closes when the commit reaches - `main`. For the last task of a batch, the same commit also carries - `Closes #<hub>` once every sub-issue is resolved, closing the hub when the - 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. In stacked mode, keep accumulating - completed tasks on `$INTEGRATION`; open the request to `main` only via the - finalize step below, on the user's explicit request. + 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 (stacked mode, explicit request only) +## 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` (see `git-workflow`) and re-run the - aggregated code review and quality gate over the whole `main...integration` - diff (see `code-review` and `quality-gates`). +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` (see `git-workflow`). + `$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 index 0b88fb0..aa0d8c0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -9,7 +9,7 @@ repository. 1. Load the `workflow-management` skill (`.opencode/skills/workflow-management/SKILL.md`) first, before doing anything else. -2. Read `README.md` for project navigation. +2. Read `README.md` for project navigation, then `examples/` for worked usage. ## Main Principles @@ -22,7 +22,14 @@ repository. 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 +* 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 @@ -30,9 +37,8 @@ repository. | Skill | When to use | |---|---| -| `.opencode/skills/project-navigation` | Starting navigation before exploring code | -| `.opencode/skills/run-tests` | Running the test suite | -| `.opencode/skills/code-style` | Formatting / linting before commit | +| `.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 @@ -41,9 +47,5 @@ repository. | `.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/git-workflow` | Branching, commits, merging | -| `.opencode/skills/reusing` | Finding existing code/docs to reuse, not duplicate | | `.opencode/skills/user-guidance-transfer` | Recording user guidance verbatim | -| `.opencode/skills/documentation` | Mapping code changes to doc updates | -| `.opencode/skills/quality-gates` | The pre-merge gate that must pass | -| `.opencode/skills/code-review` | Whole-repo review before merge | \ No newline at end of file +| `.opencode/skills/code-review` | Whole-repo review before merge | From 7a3e4b02c25375969489cc2321cb3fa7793a4f23 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Mon, 5 Oct 2026 22:47:57 +0300 Subject: [PATCH 03/17] (#4-S2): Ignore editor and build artifacts `.vscode/` and `grad_dft.egg-info/` showed up as untracked noise in every `git status`, because `.gitignore` covered neither. The existing entries (`*__pycache__/`, `data/training/dissociation/`, `runs/`, `*/__pycache__/`) are all narrow, so these follow the same style: `.vscode/` for the editor directory and the generalized `*.egg-info/` rather than the single `grad_dft.egg-info/`, since `pip install -e .` regenerates that directory under whatever the package is named and would otherwise reappear as untracked noise. Also drop a trailing blank line from `.gitignore`; the file ended without a newline, which makes every future diff on it noisier than it needs to be. Closes #4 --- .gitignore | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 2f05116..cfb38a2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,6 @@ *__pycache__/ data/training/dissociation/ runs/ -*/__pycache__/ \ No newline at end of file +*/__pycache__/ +.vscode/ +*.egg-info/ From b2e598b523c1ce2bd504f3bfbb4ea224986bf235 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Mon, 5 Oct 2026 23:10:34 +0300 Subject: [PATCH 04/17] (#5-S1): Use the supported jax config import in tests and examples `from jax.config import config` no longer resolves: JAX removed the `jax.config` submodule, so every module that used it raised `ModuleNotFoundError: No module named 'jax.config'` at import. That took out `test_predict_B88.py`, which CI *does* run, so the suite could not pass at all on a fresh install from `requirements.txt` (verified on jax/jaxlib 0.10.2). `jax.config` still exists as an attribute of the `jax` package, which is why `grad_dft/utils/types.py` (`jax.config.x64_enabled`) keeps working and needs no change. That makes the supported spelling `from jax import config` -- already used by 12 of the 17 files in this repo that call `config.update("jax_enable_x64", ...)`, including `tests/integration/molecules/test_Harris.py` and `tests/unit/test_eigenproblem.py`. This change adopts that existing convention rather than inventing one, so the call sites are untouched. Fixed in the three `test_predict_*.py` modules and, deliberately, in `examples/advanced_scripts/{train_scf_loop,train_complex_model}.py` as well: those are the user-facing documentation surface and carry the identical defect, so following the README-adjacent example scripts would fail the same way. `test_predict_B3LYP.py` stays out of CI (Grad-DFT's B3LYP differs from PySCF's, as that file documents) and `test_predict_DM21.py` stays out because it is broken independently of JAX; both are fixed only so they are not left importing a module that no longer exists. Measured, full-tree `pytest --collect-only`: 7 collection errors / 60 tests collected before, 4 / 69 after. The 4 that remain are pre-existing and untouched by this change -- 24 removed `jnp.clip(a_min=...)` call sites and PySCF's `get_init_guess(max_cycle=...)`. They are recorded as S2's scope with their root causes, not silently left behind. `black` was applied to the five touched files (it also cleared 13 pre-existing findings there: trailing whitespace, a missing final newline, collapsible imports). Pylint on them went 124 -> 111 messages with no genuinely new message. The repo-wide state -- 51 of 59 tracked Python files failing `black --check`, with no lint job in CI -- is recorded in the plan's Design Notes and deliberately not fixed here, since reformatting ~50 files would bury this compatibility fix. --- .../advanced_scripts/train_complex_model.py | 38 +++-- examples/advanced_scripts/train_scf_loop.py | 34 ++-- tasks/detailed_plan.md | 161 ++++++++++++++++++ .../molecules/test_predict_B3LYP.py | 8 +- .../integration/molecules/test_predict_B88.py | 34 ++-- .../molecules/test_predict_DM21.py | 17 +- 6 files changed, 242 insertions(+), 50 deletions(-) create mode 100644 tasks/detailed_plan.md 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/tasks/detailed_plan.md b/tasks/detailed_plan.md new file mode 100644 index 0000000..a2b1939 --- /dev/null +++ b/tasks/detailed_plan.md @@ -0,0 +1,161 @@ +# Detailed Plan — #5 Support newer JAX + +Integration branch: `agent_settings` (stacked on `main`). +Feature branch: `support_newer_jax`. + +Problem: `from jax.config import config` was removed from JAX, so the molecule +test suite fails at collection on a current JAX (verified on jax/jaxlib 0.10.2). +`requirements.txt` also declares `jax>=0.4.14`, a floor the code cannot satisfy, +because `grad_dft/utils/chunk.py` imports `jax.extend.linear_util`. + +B3LYP stays excluded from CI by design (see the issue); DM21 stays excluded +because it is broken independently of JAX (see the issue). + +## Reuse analysis + +- 12 of the 17 files that call `config.update("jax_enable_x64", True)` already + use the supported import `from jax import config` (e.g. + `tests/integration/molecules/test_Harris.py`, `tests/unit/test_eigenproblem.py`, + `examples/advanced_scripts/test_constraints.py`). The 5 broken files are fixed + to that existing convention; no new pattern is introduced. +- `grad_dft/utils/types.py` (`jax.config.x64_enabled`) and + `grad_dft/utils/chunk.py` (`jax.api_util.argnums_partial`, + `jax.extend.linear_util`) were probed on jax 0.10.2 and resolve, so the library + needs no import surgery. + +### S1: Replace the removed `jax.config` module import + +**Code:** `from jax.config import config` -> `from jax import config` in + `tests/integration/molecules/test_predict_B88.py`, + `tests/integration/molecules/test_predict_B3LYP.py`, + `tests/integration/molecules/test_predict_DM21.py`, + `examples/advanced_scripts/train_scf_loop.py`, + `examples/advanced_scripts/train_complex_model.py` + (the two example scripts are the user-facing documentation surface and + carry the identical defect; `B3LYP` and `DM21` are fixed for the same + reason even though they stay out of CI) +**Tests:** `pytest --collect-only` over the whole `tests/` tree, before vs after. + Criterion: the three `test_predict_*.py` modules stop erroring. Measured: + 7 collection errors / 60 tests collected before, 4 errors / 69 tests after. + The 4 remaining errors are pre-existing, in files this subtask does not + touch, and belong to S2 (see Design Notes). No test assertion changes. +**Docs:** No user-facing change needed — no API, install, or usage-flow change; + the example scripts keep their current behaviour, they just stop raising + `ModuleNotFoundError` at import. + +**Spec:** +- One-line change per file; `config.update(...)` calls stay exactly as they are. +- `jax_enable_x64` must still be set before any array is created (the comment + "This only works on startup!" stays). +- Success criterion: `pytest --collect-only tests/` reports no collection errors + for the three `test_predict_*.py` modules. + +- [x] S1: Implement +- [x] S1: Write tests (collect-only before/after: 7 errors/60 tests -> 4/69) +- [x] S1: Update documentation (no user-facing change needed, justified above) +- [x] S1: Pre-Commit Check (black clean on all 5 files; pylint 124 -> 111, + no new messages) +- [x] S1: Quality checks (reused the existing `from jax import config` + convention; no tolerance or assertion touched; library modules untouched) +- [x] S1: Commit + +### S2: Revalidate the suite on modern JAX and fix library breakage + +**Code:** whatever `grad_dft/` changes the suite proves necessary. None is + anticipated from the import probe; the diff is driven by test results + only. +**Tests:** the full CI list from `.github/workflows/install_and_test.yml`, run with + `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py`. 0 failures, 0 + skipped. `B3LYP` and `DM21` are not in that list and stay out. +**Docs:** No user-facing change needed unless a fix changes public behaviour. + +**Spec:** +- No tolerance may be widened, and no assertion weakened or skipped, to make a + test pass. A failure that cannot be fixed within the JAX update is a block, not + a licence to adjust the test. +- The B88 jitted-vs-non-jitted equivalence test and the energy-vs-PySCF + comparisons are the numerical guard: they must pass unchanged. + +### S3: Make the declared JAX bound honest + +**Code:** `requirements.txt` — raise the `jax` and `jaxlib` floors from `>=0.4.14` + to `>=0.4.16`, the oldest release that provides the + `jax.extend.linear_util` module `grad_dft/utils/chunk.py` imports. +**Tests:** none — a dependency bound. Verified against the published artifacts: + `jax/extend/linear_util.py` is absent in the 0.4.14 and 0.4.15 sdists and + present from 0.4.16 on. +**Docs:** No user-facing change needed — `README.md` states no JAX version. + +**Spec:** +- `jax` and `jaxlib` move together; they are released in lockstep and pip does + not resolve them independently in practice. +- Keep the change to those two lines; do not touch unrelated pins. +- Do not raise the floor further: CI still runs Python 3.9, so the floor must stay + within what JAX supports there. + +## Design Notes (discovered during implementation) + +### The style hooks have never been enforced on this repository + +`.pre-commit-config.yaml` pins `black` 22.3.0 and runs `pylint -rn -sn +--rcfile=.pylintrc`, but **51 of the 59 tracked Python files fail `black +--check`**, and there is no lint job in `.github/workflows/`. The hook was +configured and never run over the codebase. Two consequences for this task: + +- `black` was applied to the five files this task touches (it also cleaned up 13 + pre-existing findings there: trailing whitespace, a missing final newline, + collapsible imports). Pylint on those files went from 124 to 111 messages, and + a normalized before/after diff shows **no genuinely new message** — the four + apparent additions are the pre-existing `wrong-import-order` violation on the + same import line, re-rendered with a different module string + (`jax.config.config` before, `jax.config` after). +- Reformatting the whole repository is deliberately **not** done here. It would + bury a JAX compatibility fix under ~50 files of churn. It belongs in its own + task, together with adding a lint job to CI so the hook cannot rot again. + +`black` 22.3.0 also cannot be run over multiple files on this interpreter: it +calls `asyncio.get_event_loop()`, which raises under the installed `uvloop` +policy on Python 3.12. Single-file invocations work. + +### Root causes behind the 4 remaining collection errors + +Both are API removals, unrelated to `jax.config`, and both are S2's work: + +1. **`jnp.clip(..., a_min=..., a_max=...)` was removed from JAX.** 24 call sites: + 19 in `grad_dft/` (`functional.py`, `popular_functionals.py`, `train.py`) and + 5 in `tests/`. The replacement keywords are `min=` / `max=`. This is what + breaks `tests/integration/molecules/test_training.py` at import, from the + `coefficient_inputs` helper at line 98. +2. **`PySCF`'s `get_init_guess()` no longer accepts `max_cycle`**, breaking + `tests/integration/solids/test_training.py` at import. + +The two remaining `solids/` collection errors are order-dependent: those files +collect cleanly on their own (18 tests) and only fail in a full-tree run, so the +cause is import-time state shared across the `solids` test modules rather than +the files themselves. S2 must determine it; do not assume it is fixed by 1 and 2. + +### `jax.config` as a module is gone; the package attribute is not + +`from jax.config import config` fails on jax 0.10.2 because the submodule was +removed, while `jax.config` remains reachable as an attribute of the `jax` +package. Hence `from jax import config` works and `jax.config.x64_enabled` in +`grad_dft/utils/types.py` still resolves. This is why the fix is an import-line +change only, with no call-site changes. + +### The old floor was never satisfiable + +`requirements.txt` declared `jax>=0.4.14` while `grad_dft/utils/chunk.py` imports +`jax.extend.linear_util`. Inspecting the published sdists shows +`jax/extend/linear_util.py` is **absent in 0.4.14 and 0.4.15** and **present from +0.4.16** onwards (in 0.4.16 it is the real module, re-exporting +`jax._src.linear_util`; `jax/extend/__init__.py` alone already exists in 0.4.14, +which is why the breakage is easy to miss). An installation that resolved to the +oldest allowed JAX would fail at import of `grad_dft`, so the floor is +documentation of a real constraint, not a policy choice. The fix is `>=0.4.16` +for both `jax` and `jaxlib`, which must move together. + +This floor is deliberately kept low rather than raised to a modern release: CI +still runs Python 3.9, and JAX releases that support 3.9 stop well before the +current ones. "Newer JAX" therefore means the code must work across the whole +range the matrix resolves to, which is exactly what the `from jax import config` +change achieves (valid on both ends). \ No newline at end of file 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() From fc92fe8791d0872d03f60083ed70aa064b3d6837 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 00:50:44 +0300 Subject: [PATCH 05/17] (#5-S2): Restore the suite on modern JAX and PySCF S1 fixed the import error that stopped the suite collecting at all. This subtask fixes the API removals it uncovered, so the suite can actually be run end to end: `pytest --collect-only tests/` goes from 4 collection errors / 69 tests collected to 0 errors / 127 collected. Three independent causes, none of them JAX's fault: 1. `jnp.clip(..., a_min=..., a_max=...)` was removed. On jax 0.10.2 the keyword is simply gone -- `jnp.clip(x, a_min=0.5)` raises `TypeError: clip() got an unexpected keyword argument 'a_min'`. 28 call sites now pass `min=` / `max=`: 24 in `grad_dft/{functional,popular_functionals,train}.py`, 4 in the two `test_training.py` modules. Pylint, measured against this commit's parent: 20 errors down to 9, all 11 removed being exactly the `E1123 unexpected-keyword-arg 'a_min'` ones, and no genuinely new message (total 197 -> 186; the six apparent additions are pre-existing `W0621` warnings renumbered by the one line added at each `max_cycle` site). 2. PySCF's `get_init_guess()` no longer takes `max_cycle`, so passing it through `mf.kernel(max_cycle=...)` leaks it into `get_init_guess(**kwargs)` and raises `TypeError`. It is an attribute on the SCF class, so the 5 call sites set it and then call `kernel()`. 3. `test_training.py`, `test_non_xc_energy.py` and `test_functional_implementations.py` each exist under both `molecules/` and `solids/`. Pytest derives a module name from the basename, so the second import collides with the first and fails with "import file mismatch". This is why the `solids/` errors looked order-dependent: each file collected alone and only broke once both copies were collected. The `tests/**/__init__.py` files make each directory a package so the derived names differ. CI never saw this, because it runs one file per process and never collects both copies together. The same two renames are applied to the six example notebooks that carry them, so the documented workflows stop raising on their first cell. Gate: the 10 per-file invocations from `.github/workflows/install_and_test.yml`, run with `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py`, are 9/10 green: 2, 3, 35, 12+2failed, 6, 4, 26, 14, 10, 8 tests. The 2 failures are `test_dm21[mol0]`/`[mol1]`, and they are NOT caused by this task -- they fail identically on the parent commit, because two upstream breaks are stacked: tensorflow-hub 0.16 deleted the TF1 `hub.Module` API that the vendored DeepMind code calls, and PySCF 2.13's `eval_xc_eff` now dispatches through `eval_xc1`, which `NeuralNumInt` does not override, so its `eval_xc` is never called at all. Neither is a JAX change. The gate is reported BLOCKED on this rather than reported green; both root causes, a working prototype for the first, and the derivative-layout analysis for the second are written up in #6. `black` is deliberately not applied here: reformatting the 7 touched modules costs ~737 changed lines against this 44-line fix. The repo-wide state -- 46 of 64 tracked Python files failing `black --check`, no lint job in CI -- is filed as #7. --- .../example_lda_functional_02.ipynb | 2 +- .../example_neural_functional_03.ipynb | 2 +- .../example_scf_neural_functional_04.ipynb | 2 +- examples/intermediate_notebooks/DM21_02.ipynb | 2 +- .../neural_metaGGA_01.ipynb | 2 +- .../periodic_systems_bz_sampling_05.ipynb | 2 +- grad_dft/functional.py | 24 +-- grad_dft/popular_functionals.py | 14 +- grad_dft/train.py | 10 +- tasks/detailed_plan.md | 197 ++++++++++++++---- tests/__init__.py | 0 tests/integration/__init__.py | 0 tests/integration/molecules/__init__.py | 0 .../molecules/test_non_xc_energy.py | 3 +- tests/integration/molecules/test_training.py | 7 +- tests/integration/solids/__init__.py | 0 .../integration/solids/test_non_xc_energy.py | 3 +- tests/integration/solids/test_training.py | 10 +- tests/unit/__init__.py | 0 19 files changed, 204 insertions(+), 76 deletions(-) create mode 100644 tests/__init__.py create mode 100644 tests/integration/__init__.py create mode 100644 tests/integration/molecules/__init__.py create mode 100644 tests/integration/solids/__init__.py create mode 100644 tests/unit/__init__.py 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/functional.py b/grad_dft/functional.py index 6e5ffef..cc27296 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) ) @@ -1005,7 +1005,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 +1033,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 +1105,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 +1117,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 +1136,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 +1154,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 +1177,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/tasks/detailed_plan.md b/tasks/detailed_plan.md index a2b1939..3fb172e 100644 --- a/tasks/detailed_plan.md +++ b/tasks/detailed_plan.md @@ -61,13 +61,28 @@ because it is broken independently of JAX (see the issue). ### S2: Revalidate the suite on modern JAX and fix library breakage -**Code:** whatever `grad_dft/` changes the suite proves necessary. None is - anticipated from the import probe; the diff is driven by test results - only. +**Code:** driven by test results, in four groups. + 1. `jnp.clip(..., a_min=/a_max=)` → `min=/max=` at 28 sites: 24 in + `grad_dft/{functional,popular_functionals,train}.py`, 4 in + `tests/integration/{molecules,solids}/test_training.py`. + 2. `mf.kernel(max_cycle=N)` → `mf.max_cycle = N; mf.kernel()` at 5 + sites in the same two test modules. `max_cycle` is an attribute on + the PySCF class; passing it as a keyword leaks into + `get_init_guess(**kwargs)` and raises `TypeError`. + 3. `tests/**/__init__.py` (5 new files). `test_training.py`, + `test_non_xc_energy.py` and `test_functional_implementations.py` + each exist twice, once under `molecules/` and once under `solids/`. + Without packages, pytest imports them under one basename and errors + with "import file mismatch" in any run that sees both. CI hid this by + invoking one file per process; the tree could not be validated as a + whole. + 4. Notebooks: 3 × the S1 import fix, 3 × the `a_min` rename. **Tests:** the full CI list from `.github/workflows/install_and_test.yml`, run with - `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py`. 0 failures, 0 - skipped. `B3LYP` and `DM21` are not in that list and stay out. -**Docs:** No user-facing change needed unless a fix changes public behaviour. + `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py`, one file per + invocation exactly as CI does. 0 failures, 0 skipped. `B3LYP` and `DM21` + are not in that list and stay out. +**Docs:** No user-facing change needed — the fixed notebooks keep their behaviour, + they stop raising at their first cell. **Spec:** - No tolerance may be widened, and no assertion weakened or skipped, to make a @@ -76,42 +91,69 @@ because it is broken independently of JAX (see the issue). - The B88 jitted-vs-non-jitted equivalence test and the energy-vs-PySCF comparisons are the numerical guard: they must pass unchanged. +- [x] S2: Implement (items 1-4 above) +- [x] S2: Write tests (`pytest --collect-only tests/`: 4 errors / 69 tests -> + **0 errors / 127 tests**) +- [x] S2: Update documentation (no user-facing change needed, justified above) +- [x] S2: Pre-Commit Check (pylint 197 -> 186 messages on the 7 touched modules, + 11 `E1123 unexpected-keyword-arg 'a_min'` errors eliminated and no genuinely + new message; `black` deliberately **not** applied here — see Design Notes) +- [x] S2: Quality checks (existing conventions reused; no tolerance or assertion + touched) +- [ ] S2: Commit +- [x] S2: Quality gate — **9 of 10 CI files green**; `test_dm21` blocked, see + "DM21 is blocked by two independent upstream breaks" below + ### S3: Make the declared JAX bound honest **Code:** `requirements.txt` — raise the `jax` and `jaxlib` floors from `>=0.4.14` - to `>=0.4.16`, the oldest release that provides the - `jax.extend.linear_util` module `grad_dft/utils/chunk.py` imports. -**Tests:** none — a dependency bound. Verified against the published artifacts: - `jax/extend/linear_util.py` is absent in the 0.4.14 and 0.4.15 sdists and - present from 0.4.16 on. + to `>=0.4.27`. +**Tests:** none — a dependency bound. Verified against the published artifacts, see + Design Notes ("The true floor is 0.4.27, not 0.4.16"). **Docs:** No user-facing change needed — `README.md` states no JAX version. **Spec:** - `jax` and `jaxlib` move together; they are released in lockstep and pip does not resolve them independently in practice. - Keep the change to those two lines; do not touch unrelated pins. -- Do not raise the floor further: CI still runs Python 3.9, so the floor must stay - within what JAX supports there. +- Do not raise the floor further: CI still runs Python 3.9, and the last JAX + release supporting 3.9 is 0.4.30, so a floor above 0.4.30 would make the 3.9 job + uninstallable. + +- [ ] S3: Implement +- [ ] S3: Write tests (n/a — dependency bound) +- [ ] S3: Update documentation (n/a) +- [ ] S3: Pre-Commit Check +- [ ] S3: Quality checks +- [ ] S3: Commit ## Design Notes (discovered during implementation) ### The style hooks have never been enforced on this repository `.pre-commit-config.yaml` pins `black` 22.3.0 and runs `pylint -rn -sn ---rcfile=.pylintrc`, but **51 of the 59 tracked Python files fail `black ---check`**, and there is no lint job in `.github/workflows/`. The hook was -configured and never run over the codebase. Two consequences for this task: +--rcfile=.pylintrc`, but **46 of the 64 tracked Python files fail `black +--check`** (40 of 53 first-party, plus 6 of 11 under `grad_dft/external/`), and +there is no lint job in `.github/workflows/` — the single `build` job installs and +runs `pytest` and never invokes `pre-commit`. The hook was configured and never run +over the codebase. Two consequences for this task: -- `black` was applied to the five files this task touches (it also cleaned up 13 +- `black` was applied to the five files S1 touches (it also cleaned up 13 pre-existing findings there: trailing whitespace, a missing final newline, collapsible imports). Pylint on those files went from 124 to 111 messages, and a normalized before/after diff shows **no genuinely new message** — the four apparent additions are the pre-existing `wrong-import-order` violation on the same import line, re-rendered with a different module string (`jax.config.config` before, `jax.config` after). +- `black` was **not** applied to S2's modules. Reformatting those 7 files costs + ~737 changed lines against a 44-line compatibility fix — 17:1 noise that would + make the change unreviewable and mix two unrelated concerns into one commit. + Pylint was run there and shows the fix removing 11 real errors. - Reformatting the whole repository is deliberately **not** done here. It would - bury a JAX compatibility fix under ~50 files of churn. It belongs in its own - task, together with adding a lint job to CI so the hook cannot rot again. + bury a JAX compatibility fix under ~50 files of churn, and `grad_dft/external/` is + vendored verbatim from `deepmind-research`, so reformatting it also creates a + local divergence from upstream. Both are filed as follow-up **#7**, together with + adding a lint job to CI so the hook cannot rot again. `black` 22.3.0 also cannot be run over multiple files on this interpreter: it calls `asyncio.get_event_loop()`, which raises under the installed `uvloop` @@ -119,20 +161,27 @@ policy on Python 3.12. Single-file invocations work. ### Root causes behind the 4 remaining collection errors -Both are API removals, unrelated to `jax.config`, and both are S2's work: +All three are API removals, unrelated to `jax.config`, and all are S2's work: -1. **`jnp.clip(..., a_min=..., a_max=...)` was removed from JAX.** 24 call sites: - 19 in `grad_dft/` (`functional.py`, `popular_functionals.py`, `train.py`) and - 5 in `tests/`. The replacement keywords are `min=` / `max=`. This is what +1. **`jnp.clip(..., a_min=..., a_max=...)` was removed from JAX.** 28 call sites: + 24 in `grad_dft/` (`functional.py`, `popular_functionals.py`, `train.py`) and + 4 in `tests/`. The replacement keywords are `min=` / `max=`. This is what breaks `tests/integration/molecules/test_training.py` at import, from the - `coefficient_inputs` helper at line 98. + `coefficient_inputs` helper. 2. **`PySCF`'s `get_init_guess()` no longer accepts `max_cycle`**, breaking - `tests/integration/solids/test_training.py` at import. + `tests/integration/solids/test_training.py` at import. `max_cycle` is an + attribute on the SCF class, so the fix sets it rather than passing it. +3. **Duplicate test-module basenames.** `test_training.py`, + `test_non_xc_energy.py` and `test_functional_implementations.py` each exist + under both `molecules/` and `solids/`. Pytest derives a module name from the + basename, so the second import collides with the first and raises "import file + mismatch". This is why the `solids/` errors were order-dependent: each file + collected alone, and only failed once both copies were collected. The + `tests/**/__init__.py` files make each directory a package, so the derived + module names become distinct. -The two remaining `solids/` collection errors are order-dependent: those files -collect cleanly on their own (18 tests) and only fail in a full-tree run, so the -cause is import-time state shared across the `solids` test modules rather than -the files themselves. S2 must determine it; do not assume it is fixed by 1 and 2. +Cause 3 was invisible to CI, which runs one file per `pytest` process and so +never has both copies in a single collection. ### `jax.config` as a module is gone; the package attribute is not @@ -151,11 +200,85 @@ change only, with no call-site changes. `jax._src.linear_util`; `jax/extend/__init__.py` alone already exists in 0.4.14, which is why the breakage is easy to miss). An installation that resolved to the oldest allowed JAX would fail at import of `grad_dft`, so the floor is -documentation of a real constraint, not a policy choice. The fix is `>=0.4.16` -for both `jax` and `jaxlib`, which must move together. - -This floor is deliberately kept low rather than raised to a modern release: CI -still runs Python 3.9, and JAX releases that support 3.9 stop well before the -current ones. "Newer JAX" therefore means the code must work across the whole -range the matrix resolves to, which is exactly what the `from jax import config` -change achieves (valid on both ends). \ No newline at end of file +documentation of a real constraint, not a policy choice. `jax` and `jaxlib` are +released in lockstep and must move together. + +### But 0.4.16 is not enough — the real floor is 0.4.27 + +`jax.extend.linear_util` sets the floor at 0.4.16, but S2's `min=` / `max=` +rewrite sets a higher one. Reading `jnp.clip`'s published signature per release: + +| release | `a_min` / `a_max` | `min` / `max` | +|---|---|---| +| ≤ 0.4.26 | yes | **no** | +| 0.4.27 – 0.4.30 | yes (deprecated) | yes | +| ≥ 0.4.31 | **removed** | yes | + +Since S2 removes the last `a_min` / `a_max` call site, the declared floor must be +a release where `min` / `max` exist: **0.4.27**. Below that, the requirements +would advertise a JAX the code cannot run on. Above 0.4.30, CI's Python 3.9 job +becomes uninstallable. 0.4.27 is therefore both the lowest correct bound and the +lowest that keeps the existing CI matrix installable. + +### DM21 is blocked by two independent upstream breaks + +`tests/integration/molecules/test_functional_implementations.py::test_dm21` fails +on two counts, neither related to JAX. Verified by stashing this task's changes and +re-running on pristine `agent_settings`: it fails identically there, so this task +did not cause it and cannot be finished by finishing S2. + +1. **TF-Hub 0.16 removed the TF1 `Module` API.** `hub.Module` raised + `AttributeError` at `neural_numint.py:219`. `tensorflow_hub` 0.15.0 still ships + it, but cannot be used as a fix: its `__init__.py` imports + `tensorflow_hub.estimator`, which does `from tensorflow.compat.v1 import + estimator`, and that submodule no longer exists in TF 2.21. Pinning TF-Hub + below 0.16 therefore drags TF back too — the opposite of this task's goal, and + it fails against the `tensorflow>=2.13.0` already declared. +2. **PySCF 2.13's `eval_xc_eff` no longer routes through `eval_xc`.** This is the + real blocker. `numint.nr_uks` calls `ni.eval_xc_eff(...)`, which now calls + `self.eval_xc1(...)` — a method `NeuralNumInt` does not override, so it falls + through to `libxc`, which then tries to resolve the DM21 functional as a real + libxc name and dies on `ValueError: cannot reshape array of size 240720 into + shape (2,1,24072)`. `NeuralNumInt.eval_xc` is now dead code: **it is never + called**. + +A `hub.load`-based port was prototyped to prove point 1 is fixable in isolation +(`hub.load(path).signatures["default"]` with `$`→`_` keyword mangling, and the +signature's `(batch, 1)` shapes match the code's `features` dict exactly). It +does clear break 1, and then break 2 is reached. It was reverted rather than +committed: a partial port of vendored third-party numerics that still leaves the +test red, and that changes the failure message from the clear +`AttributeError: hub.Module` to a misleading reshape error, is worse than leaving +the file pristine. + +Fixing break 2 means reimplementing PySCF's derivative-tensor contract for a +neural meta-GGA inside vendored DeepMind code — reshaping `eval_xc`'s +`(vrho, vsigma, vlapl, vtau)` outputs into the `(spin, ncomp, ngrids)` layout +`nr_uks` indexes. That is research-grade numerics work, unrelated to JAX, and it +belongs in its own task rather than inside a JAX compatibility update. + +**Decision:** the user chose to keep it out of this task and file it separately. +Follow-up issue: **#6**. The gate is therefore reported as BLOCKED on file 4, not +quietly reported green. + +## Gate status + +Per-file `pytest -v` invocations from `.github/workflows/install_and_test.yml`, +with `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py`: + +| # | file | result | +|---|---|---| +| 1 | `tests/unit/test_eigenproblem.py` | 2 passed | +| 2 | `tests/unit/test_loss.py` | 3 passed | +| 3 | `tests/integration/molecules/test_non_xc_energy.py` | 35 passed | +| 4 | `tests/integration/molecules/test_functional_implementations.py` | 12 passed, **2 failed** (`test_dm21`) | +| 5 | `tests/integration/molecules/test_Harris.py` | 6 passed | +| 6 | `tests/integration/molecules/test_predict_B88.py` | 4 passed | +| 7 | `tests/integration/molecules/test_training.py` | 26 passed | +| 8 | `tests/integration/solids/test_training.py` | 14 passed | +| 9 | `tests/integration/solids/test_non_xc_energy.py` | 10 passed | +| 10 | `tests/integration/solids/test_functional_implementations.py` | 8 passed | + +**BLOCKED** on file 4, per the rule "do not assess whether a failure is +pre-existing — fix it regardless". The fix is out of scope for this task and needs +a decision; see the DM21 section above. \ No newline at end of file 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_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_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 From 6668628bc187bce88c3e5248a6d9ac325c098389 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 00:50:49 +0300 Subject: [PATCH 06/17] (#5-S3): Raise the declared JAX floor to 0.4.27 `requirements.txt` declared `jax>=0.4.14`, a floor the code has never been able to satisfy. Two independent reasons, and the higher of the two is what this commit sets: 1. `grad_dft/utils/chunk.py` imports `jax.extend.linear_util`. Inspecting the published sdists, `jax/extend/linear_util.py` is absent in 0.4.14 and 0.4.15 and present from 0.4.16. (`jax/extend/__init__.py` alone already exists in 0.4.14, which is why this is easy to miss.) An install that resolved to the oldest allowed JAX would fail at `import grad_dft`. 2. #5-S2 removed the last `jnp.clip(a_min=...)` call site in favour of `min=`/`max=`. Reading `jnp.clip`'s published signature per release: `min`/`max` do not exist at or below 0.4.26, exist alongside the deprecated `a_min`/`a_max` from 0.4.27 to 0.4.30, and are the only spelling from 0.4.31 on. So the code cannot run below 0.4.27 either. 0.4.27 is also the lowest value that keeps CI's existing matrix installable. The matrix still runs Python 3.9, and 0.4.30 is the last JAX release supporting 3.9, so any floor above 0.4.30 would make that job unresolvable. 0.4.27-0.4.30 all accept both spellings, so the 3.9 job is safe even when pip picks the newest JAX it can install there. `jax` and `jaxlib` move together because they are released in lockstep. No other pin is touched. `README.md` states no JAX version, so no documentation changes. --- requirements.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/requirements.txt b/requirements.txt index a8f10c0..0523d68 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ -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 From 43791942637ddfbdb64162bda3d5cd0224a87a50 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 00:50:56 +0300 Subject: [PATCH 07/17] (#5-S3): Record the S2/S3 gate results and decisions in the detailed plan Tick the remaining cycle steps and fold in what was measured while running the gate: the pylint before/after (20 errors -> 9), the corrected 0.4.27 floor with the per-release jnp.clip signature table, the full 10-file gate table, and the DM21 root-cause analysis with its decision to defer to #6 and the black debt to #7. --- tasks/detailed_plan.md | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/tasks/detailed_plan.md b/tasks/detailed_plan.md index 3fb172e..8f9d581 100644 --- a/tasks/detailed_plan.md +++ b/tasks/detailed_plan.md @@ -95,12 +95,13 @@ because it is broken independently of JAX (see the issue). - [x] S2: Write tests (`pytest --collect-only tests/`: 4 errors / 69 tests -> **0 errors / 127 tests**) - [x] S2: Update documentation (no user-facing change needed, justified above) -- [x] S2: Pre-Commit Check (pylint 197 -> 186 messages on the 7 touched modules, - 11 `E1123 unexpected-keyword-arg 'a_min'` errors eliminated and no genuinely - new message; `black` deliberately **not** applied here — see Design Notes) +- [x] S2: Pre-Commit Check (pylint, measured against this subtask's parent: + 197 -> 186 messages and **20 errors -> 9**, all 11 removed being the + `E1123 unexpected-keyword-arg 'a_min'` ones, with no genuinely new message; + `black` deliberately **not** applied here — see Design Notes) - [x] S2: Quality checks (existing conventions reused; no tolerance or assertion touched) -- [ ] S2: Commit +- [x] S2: Commit - [x] S2: Quality gate — **9 of 10 CI files green**; `test_dm21` blocked, see "DM21 is blocked by two independent upstream breaks" below @@ -120,12 +121,12 @@ because it is broken independently of JAX (see the issue). release supporting 3.9 is 0.4.30, so a floor above 0.4.30 would make the 3.9 job uninstallable. -- [ ] S3: Implement -- [ ] S3: Write tests (n/a — dependency bound) -- [ ] S3: Update documentation (n/a) -- [ ] S3: Pre-Commit Check -- [ ] S3: Quality checks -- [ ] S3: Commit +- [x] S3: Implement (`jax>=0.4.27`, `jaxlib>=0.4.27`) +- [x] S3: Write tests (n/a — dependency bound; evidence in Design Notes) +- [x] S3: Update documentation (n/a — `README.md` states no JAX version) +- [x] S3: Pre-Commit Check (no Python file touched) +- [x] S3: Quality checks (no other pin changed; 0.4.27 keeps the CI 3.9 job installable) +- [x] S3: Commit ## Design Notes (discovered during implementation) From 3fe0b91b2a086ed5af7911d484c2813b248dfa83 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 18:28:09 +0300 Subject: [PATCH 08/17] (#9-S1): restrict CI matrix to Python 3.12 --- .github/workflows/install_and_test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/install_and_test.yml b/.github/workflows/install_and_test.yml index 37c0c22..815747d 100644 --- a/.github/workflows/install_and_test.yml +++ b/.github/workflows/install_and_test.yml @@ -14,7 +14,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 From 918db16b0f3a44d401c6d28d8e6db0a9386b3af4 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 18:28:09 +0300 Subject: [PATCH 09/17] (#9-S2): declare python_requires >=3.12 and fix classifiers --- setup.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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", ], ) From 48abc2e3bf68cebb1b15853a5a25ea6147a02521 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 18:28:09 +0300 Subject: [PATCH 10/17] (#9-S3): drop stale Python 3.9 claims from the #5 plan --- tasks/detailed_plan.md | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tasks/detailed_plan.md b/tasks/detailed_plan.md index 8f9d581..2b5ec02 100644 --- a/tasks/detailed_plan.md +++ b/tasks/detailed_plan.md @@ -117,9 +117,11 @@ because it is broken independently of JAX (see the issue). - `jax` and `jaxlib` move together; they are released in lockstep and pip does not resolve them independently in practice. - Keep the change to those two lines; do not touch unrelated pins. -- Do not raise the floor further: CI still runs Python 3.9, and the last JAX - release supporting 3.9 is 0.4.30, so a floor above 0.4.30 would make the 3.9 job - uninstallable. +- Do not raise the floor further: 0.4.27 is the lowest release whose signature + accepts the call sites this code uses (see Design Notes), and a higher floor + would narrow support for no corresponding fix. The *upper* bound that CI's + Python 3.9 job imposed — last JAX supporting 3.9 being 0.4.30 — was removed + along with that job when CI moved to Python 3.12 only (#9). - [x] S3: Implement (`jax>=0.4.27`, `jaxlib>=0.4.27`) - [x] S3: Write tests (n/a — dependency bound; evidence in Design Notes) @@ -217,9 +219,9 @@ rewrite sets a higher one. Reading `jnp.clip`'s published signature per release: Since S2 removes the last `a_min` / `a_max` call site, the declared floor must be a release where `min` / `max` exist: **0.4.27**. Below that, the requirements -would advertise a JAX the code cannot run on. Above 0.4.30, CI's Python 3.9 job -becomes uninstallable. 0.4.27 is therefore both the lowest correct bound and the -lowest that keeps the existing CI matrix installable. +would advertise a JAX the code cannot run on. That release declares +`Requires-Python: >=3.9`, so it also admits the Python 3.12 floor introduced in +#9. 0.4.27 is therefore the lowest correct bound. ### DM21 is blocked by two independent upstream breaks From b25afd2466bd7049090592cea3a289fa7bf43cb4 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 19:10:33 +0300 Subject: [PATCH 11/17] (#6-S1): port neural_numint to the TF-Hub 0.16 hub.load API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The vendored DeepMind code built its graph with the TF1 Hub Module API (`hub.Module(spec=...)`, `hub.add_signature`, `hub.create_module_spec`), all of which tensorflow-hub 0.16 removed, so `test_dm21` died at import-time graph construction with `AttributeError: hub.Module`. Three changes, no behaviour change intended: 1. `_build_graph` now loads the checkpoint with `hub.load(self._model_path).signatures["default"]` and calls it as `self._functional(**tensor_dict)`. The keyword names keep the `tensor_dict$` prefix (only the TensorSpec `name=` fields are sanitised to `_`), so the existing `tensor_dict = {f"tensor_dict${k}": ...}` construction and the `grid_contribution` output key are untouched, as is the whole `tf.gradients` derivative machinery downstream. 2. The `outputs` dict and the trailing `hub.add_signature(...)` call are gone with the API. Their only consumer was the spec-based export, which now raises, so the derivative tensors stay plain graph attributes consumed by `eval_xc` via `self._session.run`. 3. `export_functional_and_derivatives` gets an explicit outcome instead of a TF1 reimplementation: a documented `NotImplementedError` (TF-Hub >= 0.16 removed the Module/export API it was built on; a TF2 `tf.saved_model.save` port would be a new feature). Its only callers are the vendored `export_saved_model.py` / `neural_numint_test.py`, neither in the suite. `requirements.txt` raises the floor to `tensorflow-hub>=0.16.0`: 0.14/0.15 import `tensorflow_hub.estimator`, which pulls `tensorflow.compat.v1.estimator` — gone in TF 2.21 — so older hubs cannot even be imported in this environment (verified from the published wheels). Gate for this subtask, `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py pytest -v tests/integration/molecules/test_functional_implementations.py`: 12 passed, 2 failed — exactly the pre-#6 baseline, with `test_dm21` now failing at the second documented break (`ValueError: cannot reshape array of size 240720 into shape (2,1,24072)`, the PySCF 2.13 `eval_xc_eff` routing that S2 fixes) instead of `AttributeError: hub.Module`. Pylint vs the subtask parent: no new messages (the 4 baseline `E1101 tensorflow_hub no-member` errors are gone with the removed API). Black is deliberately not applied: the baseline already fails `black --check`, and a file-wide reformat of this vendored file is out of scope (task #5 precedent). --- .../neural_numint.py | 49 +- requirements.txt | 2 +- tasks/detailed_plan.md | 519 +++++++++--------- 3 files changed, 273 insertions(+), 297 deletions(-) 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..732e0b3 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 @@ -216,7 +216,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 +264,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 +315,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 +342,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 diff --git a/requirements.txt b/requirements.txt index 0523d68..212febf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,7 @@ 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/tasks/detailed_plan.md b/tasks/detailed_plan.md index 2b5ec02..5406068 100644 --- a/tasks/detailed_plan.md +++ b/tasks/detailed_plan.md @@ -1,287 +1,258 @@ -# Detailed Plan — #5 Support newer JAX +# Detailed Plan — #6 Restore the DM21 test: port the vendored numint to TF-Hub 0.16 and PySCF 2.13 Integration branch: `agent_settings` (stacked on `main`). -Feature branch: `support_newer_jax`. - -Problem: `from jax.config import config` was removed from JAX, so the molecule -test suite fails at collection on a current JAX (verified on jax/jaxlib 0.10.2). -`requirements.txt` also declares `jax>=0.4.14`, a floor the code cannot satisfy, -because `grad_dft/utils/chunk.py` imports `jax.extend.linear_util`. - -B3LYP stays excluded from CI by design (see the issue); DM21 stays excluded -because it is broken independently of JAX (see the issue). +Feature branch: `restore_dm21_test` (created from the rebased integration branch). + +Problem: `tests/integration/molecules/test_functional_implementations.py::test_dm21` +fails on the current dependency set (`tensorflow 2.21.0`, `tensorflow-hub 0.16.1`, +`pyscf 2.13.1`) with two independent upstream breaks stacked in the vendored +DeepMind code `grad_dft/external/density_functional_approximation_dm21/.../neural_numint.py`: + +1. **TF-Hub 0.16 removed the TF1 `Module` API** — `hub.Module` raised + `AttributeError` at `neural_numint.py:219` (reproduced on the feature branch + before any change). +2. **PySCF 2.13's `eval_xc_eff` no longer routes through `eval_xc`** — + `nr_rks`/`nr_uks` now call `ni.eval_xc_eff(...)`, which is implemented as + `eval_xc1` + `xc_deriv.transform_xc`. `NeuralNumInt` overrides `eval_xc` but + not `eval_xc1`, so the call falls through to libxc, which resolves the + SCF's `xc_code` (the default `LDA,VWN`) against a meta-GGA-shaped density and + dies with `ValueError: cannot reshape array of size 240720 into shape + (2,1,24072)`. `NeuralNumInt.eval_xc` is dead code today. + +The issue's acceptance criteria: `test_dm21` passes with no tolerance change; +the `spin=0` path is either covered by a new test or explicitly rejected with a +clear `NotImplementedError`; the `hub.load` port and the PySCF shim are separate +commits; `export_functional_and_derivatives()` gets an explicit outcome. + +## Verification performed during planning (prototypes in /tmp/opencode, no repo changes) + +- **`hub.load` port is viable.** `hub.load(checkpoint).signatures["default"]` + loads the vendored TF1 SavedModel under `tf.disable_v2_behavior()` graph mode. + The signature's keyword keys keep the `$` separator (`tensor_dict$rho_a`, + only the TensorSpec `name=` fields are sanitised to `_`), so + `self._functional(**{f"tensor_dict${k}": v ...})` works as-is — the issue's + `$`→`_` mangling claim is not needed. The output key is `grid_contribution` + (unchanged). `tf.gradients` through the resulting `PartitionedCall` works and + returns finite gradients with the expected shapes (`vrho` 2×(N,1), `vsigma` + 3×(N,1), `vtau` 2×(N,1), `vhf` 2×(N,2)) — i.e. the whole existing + derivative machinery survives the port untouched. +- **The `eval_xc_eff` chain rule is settled, not guessed.** The layout PySCF + expects was validated against PySCF's own `NumInt.eval_xc_eff` on the real + libxc functional `mgga_x_tpss` (spin=1 and spin=0): building the effective + density-parameter derivatives from `eval_xc`'s `(vrho, vsigma, vtau)` with + + | spin | row layout of `vxc` | + |---|---| + | 0 | `(5, N)`: `[vrho, 2·vsigma·∇ρ, vtau]` | + | 1 | `(2, 5, N)`: `[:,0]=vrho`, `[0,1:4]=2·vsigma_aa·∇ρ_a + vsigma_ab·∇ρ_b`, `[1,1:4]=2·vsigma_bb·∇ρ_b + vsigma_ab·∇ρ_a`, `[:,4]=vtau` | + + reproduces `eval_xc_eff`'s output **exactly (max |diff| = 0.0)** on both spins. + (The alternative of feeding `eval_xc`'s outputs into + `pyscf.dft.xc_deriv.transform_xc` directly was prototyped too and produced + garbage — the `out` row order follows the xcfun convention and `transform_xc` + reads past the initialised rows for our input — so the explicit chain rule is + the design of record.) +- **Row counts.** `eval_rho` for MGGA runs with `with_lapl=False` + (`MGGA_DENSITY_LAPL = False`), so rho arrives with 5 rows (or `(2,5,N)`), + while the graph's placeholders expect 6 rows; a zero laplacian row is + inserted at index 4 (the graph discards index 4 anyway). +- **`tensorflow-hub<0.16` is unusable on the declared TF.** Verified from the + wheels: 0.14.0 and 0.15.0 both do + `from tensorflow_hub.estimator import ...` at `__init__.py:90`, which imports + `tensorflow.compat.v1.estimator`, absent from TF 2.21. The honest floor for + this repo is `tensorflow-hub>=0.16.0` (root `setup.py` reads + `requirements.txt`, so that one line is the whole dependency change). +- **The restricted path's local-HF features are spin-consistent:** + `compute_hfx_density.get_hf_density` splits a restricted `dm` into `dm/2` per + spin internally, so the RKS and UKS paths evaluate the same functional at + `ρ_a = ρ_b`; an RKS-vs-UKS energy comparison is therefore a sound end-to-end + test of the `spin=0` layout. ## Reuse analysis -- 12 of the 17 files that call `config.update("jax_enable_x64", True)` already - use the supported import `from jax import config` (e.g. - `tests/integration/molecules/test_Harris.py`, `tests/unit/test_eigenproblem.py`, - `examples/advanced_scripts/test_constraints.py`). The 5 broken files are fixed - to that existing convention; no new pattern is introduced. -- `grad_dft/utils/types.py` (`jax.config.x64_enabled`) and - `grad_dft/utils/chunk.py` (`jax.api_util.argnums_partial`, - `jax.extend.linear_util`) were probed on jax 0.10.2 and resolve, so the library - needs no import surgery. - -### S1: Replace the removed `jax.config` module import - -**Code:** `from jax.config import config` -> `from jax import config` in - `tests/integration/molecules/test_predict_B88.py`, - `tests/integration/molecules/test_predict_B3LYP.py`, - `tests/integration/molecules/test_predict_DM21.py`, - `examples/advanced_scripts/train_scf_loop.py`, - `examples/advanced_scripts/train_complex_model.py` - (the two example scripts are the user-facing documentation surface and - carry the identical defect; `B3LYP` and `DM21` are fixed for the same - reason even though they stay out of CI) -**Tests:** `pytest --collect-only` over the whole `tests/` tree, before vs after. - Criterion: the three `test_predict_*.py` modules stop erroring. Measured: - 7 collection errors / 60 tests collected before, 4 errors / 69 tests after. - The 4 remaining errors are pre-existing, in files this subtask does not - touch, and belong to S2 (see Design Notes). No test assertion changes. -**Docs:** No user-facing change needed — no API, install, or usage-flow change; - the example scripts keep their current behaviour, they just stop raising - `ModuleNotFoundError` at import. +- **`NeuralNumInt.eval_xc` (existing)** — the shim does not re-implement any + DM21 evaluation: `eval_xc_eff` routes to it, reusing its derivative + computation, its float64 casts, and its `_vmat_hf` (local-HF potential) + side effects exactly as before the PySCF upgrade. +- **PySCF's own contract** — the layout table above is PySCF's + `NumInt.eval_xc_eff` contract, not a new invention; it was confirmed against + `pyscf.dft.numint.NumInt.eval_xc_eff` empirically (see Verification). +- **Test file conventions** — the new test reuses the module-level `mols` list, + `dft.RKS`/`dft.UKS` + `NeuralNumInt` setup of the existing `test_dm21`, the + `Hartree2kcalmol` conversion already imported at + `tests/.../test_functional_implementations.py:24`, and the file's `atol=1` + kcal/mol tolerance style. No new helpers. +- No `grad_dft/utils/` additions are needed; nothing else in the repo calls + `eval_xc`/`eval_xc_eff` (verified by grep: only `neural_numint.py` itself and + `grad_dft/interface/pyscf.py:1050`, which only monkey-patches `mf._numint`). + +### S1: Port `neural_numint.py` to the TF-Hub 0.16 API (`hub.load`) and raise the declared hub floor [done] + +**Code:** `grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py` +(`_build_graph`, `export_functional_and_derivatives`), `requirements.txt` +(`tensorflow-hub>=0.14.0` → `tensorflow-hub>=0.16.0`) +**Tests:** run the new test only via construction smoke test +(`python -c "from grad_dft.external import NeuralNumInt, Functional; +NeuralNumInt(Functional.DM21)"` must build the graph), then +`pytest -v tests/integration/molecules/test_functional_implementations.py::test_dm21` +— expected to fail at the **documented second break** +(`ValueError: cannot reshape array of size 240720 into shape (2,1,24072)`), not +at `AttributeError: hub.Module`; that error is the S2 input. The suite must +otherwise be unchanged (all other tests in the file still pass). +**Docs:** No README/example change needed (TF-Hub internals are invisible to +users). The `export_functional_and_derivatives` docstring is updated in-code to +state the new `NotImplementedError` behaviour and why; the class docstring's +usage example is untouched. **Spec:** -- One-line change per file; `config.update(...)` calls stay exactly as they are. -- `jax_enable_x64` must still be set before any array is created (the comment - "This only works on startup!" stays). -- Success criterion: `pytest --collect-only tests/` reports no collection errors - for the three `test_predict_*.py` modules. - -- [x] S1: Implement -- [x] S1: Write tests (collect-only before/after: 7 errors/60 tests -> 4/69) -- [x] S1: Update documentation (no user-facing change needed, justified above) -- [x] S1: Pre-Commit Check (black clean on all 5 files; pylint 124 -> 111, - no new messages) -- [x] S1: Quality checks (reused the existing `from jax import config` - convention; no tolerance or assertion touched; library modules untouched) +- In `_build_graph`: + - Replace `self._functional = hub.Module(spec=self._model_path)` with + `self._functional = hub.load(self._model_path).signatures["default"]`. + - Keep `tensor_dict = {f"tensor_dict${k}": v for k, v in features.items()}` + and change `predictions = self._functional(tensor_dict, as_dict=True)` to + `predictions = self._functional(**tensor_dict)`; `local_xc = + predictions["grid_contribution"]` stays (verified output key). + - Delete the `outputs` dict and the trailing + `hub.add_signature(inputs=..., outputs=...)` call (API removed in 0.16; the + only consumer was the spec-based export, which now raises — see below), and + drop the "no-op if called outside hub.create_module_spec" comment. + - Everything else (placeholders, feature construction, `tf.gradients`, + session) is untouched — verified to work through the port. +- Replace the body of `export_functional_and_derivatives` with a + `raise NotImplementedError` whose message says: TF-Hub >= 0.16 removed the + TF1 `Module`/`create_module_spec` API this export was built on; exporting + would require a TF2 `tf.saved_model.save` reimplementation; the method is + only used by the vendored `export_saved_model.py` / `neural_numint_test.py`, + which are not part of the test suite. Keep the original docstring and append + the Raises section. +- `requirements.txt`: raise the floor to `tensorflow-hub>=0.16.0` (0.14/0.15 + import `tensorflow_hub.estimator` → `tensorflow.compat.v1.estimator`, which no + longer exists in TF 2.21; verified from the published wheels). +- Style: `black` is **not** applied to this vendored file — the baseline already + fails `black --check`, and a file-wide reformat is out of scope (task #5 + precedent: edits stay diff-only to limit churn/upstream divergence). `pylint + -rn -sn --rcfile=.pylintrc` must show no new messages vs the subtask parent. +- Commit: `(#6-S1): port neural_numint to the TF-Hub 0.16 hub.load API`. + +- [x] S1: Implement (`hub.load` port, `outputs`/`hub.add_signature` removal, + `NotImplementedError` export, hub floor `>=0.16.0`) +- [x] S1: Write tests (no new test in this subtask; smoke test builds the graph + and `WrappedFunction`, and with `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py` + the file gives 12 passed / 2 failed — `test_dm21` alone, now at the + documented second break, the reshape `ValueError`, instead of + `AttributeError: hub.Module`) +- [x] S1: Update documentation (no user-facing change needed — TF-Hub internals + are invisible to users; the `Raises` section added to + `export_functional_and_derivatives`' docstring is the explicit outcome) +- [x] S1: Pre-Commit Check (`black` deliberately not applied: the baseline + already fails `black --check`, so the file stays as-is per the plan's + style note; `pylint -rn -sn --rcfile=.pylintrc` vs subtask parent: no new + messages, and the 4 baseline `E1101 tensorflow_hub no-member` errors are + gone) +- [x] S1: Quality checks (no duplication: single evaluation path retained; + nothing added to library modules outside the vendored shim; no tolerance + or assertion touched) - [x] S1: Commit -### S2: Revalidate the suite on modern JAX and fix library breakage - -**Code:** driven by test results, in four groups. - 1. `jnp.clip(..., a_min=/a_max=)` → `min=/max=` at 28 sites: 24 in - `grad_dft/{functional,popular_functionals,train}.py`, 4 in - `tests/integration/{molecules,solids}/test_training.py`. - 2. `mf.kernel(max_cycle=N)` → `mf.max_cycle = N; mf.kernel()` at 5 - sites in the same two test modules. `max_cycle` is an attribute on - the PySCF class; passing it as a keyword leaks into - `get_init_guess(**kwargs)` and raises `TypeError`. - 3. `tests/**/__init__.py` (5 new files). `test_training.py`, - `test_non_xc_energy.py` and `test_functional_implementations.py` - each exist twice, once under `molecules/` and once under `solids/`. - Without packages, pytest imports them under one basename and errors - with "import file mismatch" in any run that sees both. CI hid this by - invoking one file per process; the tree could not be validated as a - whole. - 4. Notebooks: 3 × the S1 import fix, 3 × the `a_min` rename. -**Tests:** the full CI list from `.github/workflows/install_and_test.yml`, run with - `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py`, one file per - invocation exactly as CI does. 0 failures, 0 skipped. `B3LYP` and `DM21` - are not in that list and stay out. -**Docs:** No user-facing change needed — the fixed notebooks keep their behaviour, - they stop raising at their first cell. +### S2: Add the `eval_xc_eff` override that routes PySCF's derivative contract to `eval_xc` + +**Code:** new method `NeuralNumInt.eval_xc_eff` in +`grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py` +(placed directly after `eval_xc`); one class-docstring sentence updated +("The actual evaluation ... is performed in NeuralNumInt.eval_xc" → mention +that PySCF reaches it through the `eval_xc_eff` override) +**Tests:** `pytest -v tests/integration/molecules/test_functional_implementations.py` +— `test_dm21` (both parametrized molecules) must pass with the **unchanged** +`atol=1` tolerance, and the other 9 tests in the file must stay green; also run +`pytest -v tests/integration/molecules/test_predict_B88.py +tests/integration/molecules/test_non_xc_energy.py tests/integration/molecules/test_Harris.py` +to confirm no collateral damage. +**Docs:** numpydoc docstring on `eval_xc_eff` documenting the argument +semantics (delegates to `eval_xc`), the returned layout table for `spin=0`/ +`spin=1`, and the Raises conditions; no README/example change (public usage of +`NeuralNumInt` is unchanged). **Spec:** -- No tolerance may be widened, and no assertion weakened or skipped, to make a - test pass. A failure that cannot be fixed within the JAX update is a block, not - a licence to adjust the test. -- The B88 jitted-vs-non-jitted equivalence test and the energy-vs-PySCF - comparisons are the numerical guard: they must pass unchanged. - -- [x] S2: Implement (items 1-4 above) -- [x] S2: Write tests (`pytest --collect-only tests/`: 4 errors / 69 tests -> - **0 errors / 127 tests**) -- [x] S2: Update documentation (no user-facing change needed, justified above) -- [x] S2: Pre-Commit Check (pylint, measured against this subtask's parent: - 197 -> 186 messages and **20 errors -> 9**, all 11 removed being the - `E1123 unexpected-keyword-arg 'a_min'` ones, with no genuinely new message; - `black` deliberately **not** applied here — see Design Notes) -- [x] S2: Quality checks (existing conventions reused; no tolerance or assertion - touched) -- [x] S2: Commit -- [x] S2: Quality gate — **9 of 10 CI files green**; `test_dm21` blocked, see - "DM21 is blocked by two independent upstream breaks" below - -### S3: Make the declared JAX bound honest - -**Code:** `requirements.txt` — raise the `jax` and `jaxlib` floors from `>=0.4.14` - to `>=0.4.27`. -**Tests:** none — a dependency bound. Verified against the published artifacts, see - Design Notes ("The true floor is 0.4.27, not 0.4.16"). -**Docs:** No user-facing change needed — `README.md` states no JAX version. +- Signature mirrors PySCF's base method exactly: + `def eval_xc_eff(self, xc_code, rho, deriv=1, omega=None, xctype=None, verbose=None, spin=None)`. +- Behaviour: + 1. Mirror the base class preamble: `omega = self.omega if omega is None else + omega` (`NumInt.omega` is `None`, so a plain call passes `omega=None` + through to `eval_xc`, which accepts `None` and raises a clear + `NotImplementedError` for an explicit RSH value — unchanged semantics); + normalise `rho` with `np.asarray(rho, order="C")`; infer `spin` from + `rho.shape[0] == 2` when it is not given (same rule as the base class). + 2. `deriv == 0` → still evaluate (the `eval_xc` call carries the `_vmat_hf` + side effect) but return `[exc, None, None, None]`; `deriv > 1` → raise + `NotImplementedError` ("DM21 provides no second derivatives"; the base + class would die inside `transform_xc`). + 3. Re-insert the laplacian row: if `rho.shape[-2] == 5`, insert a zero row at + index 4 (graph placeholders are `(6, N)`; index 4 is discarded by the + graph). If 6 rows arrive, pass through. + 4. `exc, (vrho, vsigma, vlapl, vtau), _, _ = self.eval_xc(xc_code, rho6, + spin=spin, deriv=deriv, omega=omega, verbose=verbose)` — full reuse of the + existing evaluation, including `_vmat_hf` accumulation. + 5. Build the effective derivatives with the validated chain rule + (outputs of `eval_xc` are `(N, ncomp)` — transpose to `(ncomp, N)` for the + spin-polarised case): + - `spin == 0`: `vxc = [vrho, 2·vsigma[:,None]·rho[1:4], zeros, vtau]` + stacked to shape `(5, N)` (row 2 = laplacian slot, zeros as PySCF does). + - `spin == 1`: shape `(2, 5, N)` with + `vxc[:,0] = vrho.T`, `vxc[0,1:4] = 2·vsigma[:,0]·∇ρ_a + vsigma[:,1]·∇ρ_b`, + `vxc[1,1:4] = 2·vsigma[:,2]·∇ρ_b + vsigma[:,1]·∇ρ_a`, + `vxc[:,4] = vtau.T`, where `∇ρ_a = rho[0,1:4]`, `∇ρ_b = rho[1,1:4]`. + 6. Return `(exc, vxc, None, None)` (float64 — `eval_xc` already casts). +- The docstring states the layout contract and notes the chain rule was + verified against `pyscf.dft.numint.NumInt.eval_xc_eff`. +- Comment where the chain rule lives, pointing at the PySCF call sites + (`nr_rks` MGGA branch: `wv[0]/wv[4]` halving, `_scale_ao_sparse(ao[:4], + wv[:4])`, `_tau_dot_sparse(..., wv[4])`; `nr_uks` the same per spin row) so a + future reader can re-derive it. +- Commit: `(#6-S2): add an eval_xc_eff shim routing PySCF 2.13 back to eval_xc`. + +### S3: Cover the restricted (`spin=0`) path with an RKS-vs-UKS consistency test + +**Code:** new test `test_dm21_rks` in +`tests/integration/molecules/test_functional_implementations.py` (after +`test_dm21`), reusing the module-level `mols[0]` closed-shell molecule and the +existing imports +**Tests:** the new test itself: build `dft.RKS(mol)` and `dft.UKS(mol)` on the +closed-shell HF molecule, give both `mf._numint = NeuralNumInt(Functional.DM21)` +(identical default grids, mirroring `test_dm21`'s setup exactly), run both +`mf.kernel()` calls, and assert `(e_rks - e_uks) * Hartree2kcalmol` is +`allclose(..., 0, atol=1)` and that neither energy is NaN. The whole file must +stay green. +**Docs:** No user-facing change needed (test-only; nothing in README/examples +documents individual test names). **Spec:** -- `jax` and `jaxlib` move together; they are released in lockstep and pip does - not resolve them independently in practice. -- Keep the change to those two lines; do not touch unrelated pins. -- Do not raise the floor further: 0.4.27 is the lowest release whose signature - accepts the call sites this code uses (see Design Notes), and a higher floor - would narrow support for no corresponding fix. The *upper* bound that CI's - Python 3.9 job imposed — last JAX supporting 3.9 being 0.4.30 — was removed - along with that job when CI moved to Python 3.12 only (#9). - -- [x] S3: Implement (`jax>=0.4.27`, `jaxlib>=0.4.27`) -- [x] S3: Write tests (n/a — dependency bound; evidence in Design Notes) -- [x] S3: Update documentation (n/a — `README.md` states no JAX version) -- [x] S3: Pre-Commit Check (no Python file touched) -- [x] S3: Quality checks (no other pin changed; 0.4.27 keeps the CI 3.9 job installable) -- [x] S3: Commit - -## Design Notes (discovered during implementation) - -### The style hooks have never been enforced on this repository - -`.pre-commit-config.yaml` pins `black` 22.3.0 and runs `pylint -rn -sn ---rcfile=.pylintrc`, but **46 of the 64 tracked Python files fail `black ---check`** (40 of 53 first-party, plus 6 of 11 under `grad_dft/external/`), and -there is no lint job in `.github/workflows/` — the single `build` job installs and -runs `pytest` and never invokes `pre-commit`. The hook was configured and never run -over the codebase. Two consequences for this task: - -- `black` was applied to the five files S1 touches (it also cleaned up 13 - pre-existing findings there: trailing whitespace, a missing final newline, - collapsible imports). Pylint on those files went from 124 to 111 messages, and - a normalized before/after diff shows **no genuinely new message** — the four - apparent additions are the pre-existing `wrong-import-order` violation on the - same import line, re-rendered with a different module string - (`jax.config.config` before, `jax.config` after). -- `black` was **not** applied to S2's modules. Reformatting those 7 files costs - ~737 changed lines against a 44-line compatibility fix — 17:1 noise that would - make the change unreviewable and mix two unrelated concerns into one commit. - Pylint was run there and shows the fix removing 11 real errors. -- Reformatting the whole repository is deliberately **not** done here. It would - bury a JAX compatibility fix under ~50 files of churn, and `grad_dft/external/` is - vendored verbatim from `deepmind-research`, so reformatting it also creates a - local divergence from upstream. Both are filed as follow-up **#7**, together with - adding a lint job to CI so the hook cannot rot again. - -`black` 22.3.0 also cannot be run over multiple files on this interpreter: it -calls `asyncio.get_event_loop()`, which raises under the installed `uvloop` -policy on Python 3.12. Single-file invocations work. - -### Root causes behind the 4 remaining collection errors - -All three are API removals, unrelated to `jax.config`, and all are S2's work: - -1. **`jnp.clip(..., a_min=..., a_max=...)` was removed from JAX.** 28 call sites: - 24 in `grad_dft/` (`functional.py`, `popular_functionals.py`, `train.py`) and - 4 in `tests/`. The replacement keywords are `min=` / `max=`. This is what - breaks `tests/integration/molecules/test_training.py` at import, from the - `coefficient_inputs` helper. -2. **`PySCF`'s `get_init_guess()` no longer accepts `max_cycle`**, breaking - `tests/integration/solids/test_training.py` at import. `max_cycle` is an - attribute on the SCF class, so the fix sets it rather than passing it. -3. **Duplicate test-module basenames.** `test_training.py`, - `test_non_xc_energy.py` and `test_functional_implementations.py` each exist - under both `molecules/` and `solids/`. Pytest derives a module name from the - basename, so the second import collides with the first and raises "import file - mismatch". This is why the `solids/` errors were order-dependent: each file - collected alone, and only failed once both copies were collected. The - `tests/**/__init__.py` files make each directory a package, so the derived - module names become distinct. - -Cause 3 was invisible to CI, which runs one file per `pytest` process and so -never has both copies in a single collection. - -### `jax.config` as a module is gone; the package attribute is not - -`from jax.config import config` fails on jax 0.10.2 because the submodule was -removed, while `jax.config` remains reachable as an attribute of the `jax` -package. Hence `from jax import config` works and `jax.config.x64_enabled` in -`grad_dft/utils/types.py` still resolves. This is why the fix is an import-line -change only, with no call-site changes. - -### The old floor was never satisfiable - -`requirements.txt` declared `jax>=0.4.14` while `grad_dft/utils/chunk.py` imports -`jax.extend.linear_util`. Inspecting the published sdists shows -`jax/extend/linear_util.py` is **absent in 0.4.14 and 0.4.15** and **present from -0.4.16** onwards (in 0.4.16 it is the real module, re-exporting -`jax._src.linear_util`; `jax/extend/__init__.py` alone already exists in 0.4.14, -which is why the breakage is easy to miss). An installation that resolved to the -oldest allowed JAX would fail at import of `grad_dft`, so the floor is -documentation of a real constraint, not a policy choice. `jax` and `jaxlib` are -released in lockstep and must move together. - -### But 0.4.16 is not enough — the real floor is 0.4.27 - -`jax.extend.linear_util` sets the floor at 0.4.16, but S2's `min=` / `max=` -rewrite sets a higher one. Reading `jnp.clip`'s published signature per release: - -| release | `a_min` / `a_max` | `min` / `max` | -|---|---|---| -| ≤ 0.4.26 | yes | **no** | -| 0.4.27 – 0.4.30 | yes (deprecated) | yes | -| ≥ 0.4.31 | **removed** | yes | - -Since S2 removes the last `a_min` / `a_max` call site, the declared floor must be -a release where `min` / `max` exist: **0.4.27**. Below that, the requirements -would advertise a JAX the code cannot run on. That release declares -`Requires-Python: >=3.9`, so it also admits the Python 3.12 floor introduced in -#9. 0.4.27 is therefore the lowest correct bound. - -### DM21 is blocked by two independent upstream breaks - -`tests/integration/molecules/test_functional_implementations.py::test_dm21` fails -on two counts, neither related to JAX. Verified by stashing this task's changes and -re-running on pristine `agent_settings`: it fails identically there, so this task -did not cause it and cannot be finished by finishing S2. - -1. **TF-Hub 0.16 removed the TF1 `Module` API.** `hub.Module` raised - `AttributeError` at `neural_numint.py:219`. `tensorflow_hub` 0.15.0 still ships - it, but cannot be used as a fix: its `__init__.py` imports - `tensorflow_hub.estimator`, which does `from tensorflow.compat.v1 import - estimator`, and that submodule no longer exists in TF 2.21. Pinning TF-Hub - below 0.16 therefore drags TF back too — the opposite of this task's goal, and - it fails against the `tensorflow>=2.13.0` already declared. -2. **PySCF 2.13's `eval_xc_eff` no longer routes through `eval_xc`.** This is the - real blocker. `numint.nr_uks` calls `ni.eval_xc_eff(...)`, which now calls - `self.eval_xc1(...)` — a method `NeuralNumInt` does not override, so it falls - through to `libxc`, which then tries to resolve the DM21 functional as a real - libxc name and dies on `ValueError: cannot reshape array of size 240720 into - shape (2,1,24072)`. `NeuralNumInt.eval_xc` is now dead code: **it is never - called**. - -A `hub.load`-based port was prototyped to prove point 1 is fixable in isolation -(`hub.load(path).signatures["default"]` with `$`→`_` keyword mangling, and the -signature's `(batch, 1)` shapes match the code's `features` dict exactly). It -does clear break 1, and then break 2 is reached. It was reverted rather than -committed: a partial port of vendored third-party numerics that still leaves the -test red, and that changes the failure message from the clear -`AttributeError: hub.Module` to a misleading reshape error, is worse than leaving -the file pristine. - -Fixing break 2 means reimplementing PySCF's derivative-tensor contract for a -neural meta-GGA inside vendored DeepMind code — reshaping `eval_xc`'s -`(vrho, vsigma, vlapl, vtau)` outputs into the `(spin, ncomp, ngrids)` layout -`nr_uks` indexes. That is research-grade numerics work, unrelated to JAX, and it -belongs in its own task rather than inside a JAX compatibility update. - -**Decision:** the user chose to keep it out of this task and file it separately. -Follow-up issue: **#6**. The gate is therefore reported as BLOCKED on file 4, not -quietly reported green. - -## Gate status - -Per-file `pytest -v` invocations from `.github/workflows/install_and_test.yml`, -with `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py`: - -| # | file | result | -|---|---|---| -| 1 | `tests/unit/test_eigenproblem.py` | 2 passed | -| 2 | `tests/unit/test_loss.py` | 3 passed | -| 3 | `tests/integration/molecules/test_non_xc_energy.py` | 35 passed | -| 4 | `tests/integration/molecules/test_functional_implementations.py` | 12 passed, **2 failed** (`test_dm21`) | -| 5 | `tests/integration/molecules/test_Harris.py` | 6 passed | -| 6 | `tests/integration/molecules/test_predict_B88.py` | 4 passed | -| 7 | `tests/integration/molecules/test_training.py` | 26 passed | -| 8 | `tests/integration/solids/test_training.py` | 14 passed | -| 9 | `tests/integration/solids/test_non_xc_energy.py` | 10 passed | -| 10 | `tests/integration/solids/test_functional_implementations.py` | 8 passed | - -**BLOCKED** on file 4, per the rule "do not assess whether a failure is -pre-existing — fix it regardless". The fix is out of scope for this task and needs -a decision; see the DM21 section above. \ No newline at end of file +- Rationale (recorded in the test as a comment): DM21 maps restricted inputs to + `ρ_a = ρ_b = ρ/2` and `get_hf_density` splits a restricted `dm` into `dm/2` + per spin, so on a closed-shell molecule the restricted and unrestricted + self-consistent solutions must agree; UKS is already validated end-to-end by + `test_dm21`. A wrong `spin=0` derivative layout therefore shows up as an + energy mismatch (or a diverged/NaN SCF), not as a silent pass. +- Only the closed-shell molecule (`mols[0]`, HF) is used — RKS on the open-shell + Li atom (`mols[1]`) is not a valid restricted calculation. +- Fallback (decided in advance, per the issue): if this test demonstrates that + the `spin=0` path cannot be made to agree (i.e. the failure is not the layout + shim's to fix), replace the RKS path with an explicit + `raise NotImplementedError("The restricted (spin=0) path ...")` in + `eval_xc_eff` and change the test to assert that raise. The test-coverage + route is tried first because the restricted route is the documented usage in + the class docstring (`mf = dft.RKS(...)`). +- Commit (last subtask): `(#6-S3): cover the DM21 restricted path with an + RKS-vs-UKS consistency test` plus a standalone `Closes #6` line. + +## Acceptance criteria mapping + +| Criterion (issue #6) | Where | +|---|---| +| `test_dm21` passes, no tolerance change | S2 Tests (`atol=1` untouched) | +| `spin=0` covered by a new test or explicitly rejected | S3 (test first, `NotImplementedError` fallback pre-decided) | +| `hub.load` port and PySCF shim in separate commits | S1 / S2 commits | +| `export_functional_and_derivatives()` explicit outcome | S1 (`NotImplementedError` + docstring) | + +After S3: whole-repo code review (`code-review` skill), then the quality gate +(`quality-gates` skill: CI-file pytest invocations + `black` + `pylint`), then +integration into `agent_settings` after user confirmation. From 57edba1e14f89f580914bec1576cf5e8e2175685 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 19:41:52 +0300 Subject: [PATCH 12/17] (#6-S2): add an eval_xc_eff shim routing PySCF 2.13 back to eval_xc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PySCF 2.x's nr_rks/nr_uks grid loops call `ni.eval_xc_eff(...)` instead of `ni.eval_xc(...)`. NeuralNumInt only overrode eval_xc, so the call fell through to the base class, which evaluated with libxc and died on the meta-GGA-shaped density (`ValueError: cannot reshape array of size 240720 into shape (2,1,24072)`) — and NeuralNumInt.eval_xc was dead code. The new override routes the whole evaluation back to eval_xc (the DM21 graph run plus the local Hartree-Fock `_vmat_hf` side effects) and applies the chain rule from eval_xc's sigma-derivative convention to the density-parameter convention eval_xc_eff promises: - spin=0 -> (5, N): [dE/d rho, dE/d nabla rho, dE/d tau] - spin=1 -> (2, 5, N): the same rows per spin channel, with the cross term dE/d sigma_ab distributing between the two gradient rows. The layout is not guessed: it reproduces `pyscf.dft.numint.NumInt.eval_xc_eff` exactly (max |diff| = 0.0 on mgga_x_tpss for both spins). The alternative of feeding eval_xc's outputs through `xc_deriv.transform_xc` was prototyped first and produced garbage (its out row order is the xcfun convention), so the explicit chain rule is the design of record; the code comments point at PySCF's MGGA contract sites (nr_rks: wv[0]/wv[4] halving, _scale_ao_sparse on ao[:4]/wv[:4], _tau_dot_sparse on wv[4]; nr_uks the same per spin row). Rho arrives from eval_rho without the laplacian row (MGGA_DENSITY_LAPL is off) while the graph's placeholders are (6, N); a zero laplacian row is inserted before evaluation (the graph discards it). The preamble mirrors the base method (omega fallback, asarray, spin inference from rho.shape[0]), deriv == 0 still evaluates for the _vmat_hf side effect and returns [exc, None, None, None], and deriv > 1 raises NotImplementedError (DM21 has no second derivatives; the base class would fail inside transform_xc). The signature mirrors the base exactly, which is why the new method carries an inline pylint suppression for the two too-many-arguments checks. With this commit `test_dm21`'s PySCF leg — `mf.kernel()` on both parametrized molecules — completes for the first time on these dependencies; the test as a whole still cannot pass until S3, because its second leg (`generate_DM21_weights`) hits an independent break (TF1 SavedModel loading on TF 2.21). Gate for this subtask, with PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py: test_functional_implementations.py gives 12 passed / 2 failed (both failures now at test_dm21's second leg), and the collateral files test_predict_B88.py, test_non_xc_energy.py, test_Harris.py give 45 passed. Pylint vs the S1 parent: message multiset identical (the eval_xc_eff too-many-arguments pair is suppressed inline). Black is not applied to the vendored file, same precedent as S1 (baseline already fails `black --check`). --- .../neural_numint.py | 131 +++++++++++++++++- tasks/detailed_plan.md | 82 +++++++++-- 2 files changed, 201 insertions(+), 12 deletions(-) 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 732e0b3..a1dc42c 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 @@ -161,7 +161,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. @@ -807,3 +808,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/tasks/detailed_plan.md b/tasks/detailed_plan.md index 5406068..a652f1f 100644 --- a/tasks/detailed_plan.md +++ b/tasks/detailed_plan.md @@ -154,7 +154,7 @@ usage example is untouched. or assertion touched) - [x] S1: Commit -### S2: Add the `eval_xc_eff` override that routes PySCF's derivative contract to `eval_xc` +### S2: Add the `eval_xc_eff` override that routes PySCF's derivative contract to `eval_xc` [done] **Code:** new method `NeuralNumInt.eval_xc_eff` in `grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py` @@ -162,8 +162,11 @@ usage example is untouched. ("The actual evaluation ... is performed in NeuralNumInt.eval_xc" → mention that PySCF reaches it through the `eval_xc_eff` override) **Tests:** `pytest -v tests/integration/molecules/test_functional_implementations.py` -— `test_dm21` (both parametrized molecules) must pass with the **unchanged** -`atol=1` tolerance, and the other 9 tests in the file must stay green; also run +— with the shim, `test_dm21`'s PySCF leg (`mf.kernel()` on both parametrized +molecules) completes and the other 9 tests stay green; the test as a whole +**cannot pass until S3** because its second leg, +`generate_DM21_weights`, hits an independent break (TF1 `tf.saved_model.load` +on TF 2.21 — see S3). The `atol=1` tolerance is untouched. Also run `pytest -v tests/integration/molecules/test_predict_B88.py tests/integration/molecules/test_non_xc_energy.py tests/integration/molecules/test_Harris.py` to confirm no collateral damage. @@ -195,8 +198,12 @@ semantics (delegates to `eval_xc`), the returned layout table for `spin=0`/ 5. Build the effective derivatives with the validated chain rule (outputs of `eval_xc` are `(N, ncomp)` — transpose to `(ncomp, N)` for the spin-polarised case): - - `spin == 0`: `vxc = [vrho, 2·vsigma[:,None]·rho[1:4], zeros, vtau]` - stacked to shape `(5, N)` (row 2 = laplacian slot, zeros as PySCF does). + - `spin == 0`: `vxc = np.concatenate([vrho[None, :], + 2·vsigma[None, :]·rho[1:4], vtau[None, :]])` → shape `(5, N)`, rows + `[dE/dρ, dE/d∇ρ (3), dE/dτ]` — **no laplacian row** (this line + previously claimed a `zeros` row at index 2; re-validated against + `NumInt.eval_xc_eff` on `mgga_x_tpss`, max |diff| = 0.0, so the layout + table in Verification is the correct one). - `spin == 1`: shape `(2, 5, N)` with `vxc[:,0] = vrho.T`, `vxc[0,1:4] = 2·vsigma[:,0]·∇ρ_a + vsigma[:,1]·∇ρ_b`, `vxc[1,1:4] = 2·vsigma[:,2]·∇ρ_b + vsigma[:,1]·∇ρ_a`, @@ -210,7 +217,62 @@ semantics (delegates to `eval_xc`), the returned layout table for `spin=0`/ future reader can re-derive it. - Commit: `(#6-S2): add an eval_xc_eff shim routing PySCF 2.13 back to eval_xc`. -### S3: Cover the restricted (`spin=0`) path with an RKS-vs-UKS consistency test +- [x] S2: Implement (`eval_xc_eff` override + class-docstring sentence) +- [x] S2: Write tests (chain rule vs `NumInt.eval_xc_eff` on `mgga_x_tpss`: + max |diff| = 0.0 for both spins; `test_dm21`'s PySCF leg — `mf.kernel()` + on both molecules — completes; collateral files green: 45 passed across + `test_predict_B88.py`, `test_non_xc_energy.py`, `test_Harris.py`; the + full `test_dm21` is deferred to S3, see Tests) +- [x] S2: Update documentation (numpydoc docstring on `eval_xc_eff` with the + layout contract and raises; class docstring updated; no user-facing change) +- [x] S2: Pre-Commit Check (`black` not applied to the vendored file, same + precedent as S1; `pylint -rn -sn --rcfile=.pylintrc` vs S1's parent: no + new messages) +- [x] S2: Quality checks (no duplication: single evaluation path reused via + `eval_xc`; no tolerance or assertion touched) +- [x] S2: Commit + +### S3: Read the vendored TF1 DM21 checkpoint directly in `generate_DM21_weights` + +**Discovered while running S2's tests** — not part of the original issue +analysis: with the `eval_xc_eff` shim in place, `mf.kernel()` completes, and +`test_dm21` moves to its second leg, `DM21().generate_DM21_weights()` +(`grad_dft/functional.py`), which loads the same vendored checkpoint with +`tf.saved_model.load(folder)`. On TF 2.21 the V1 SavedModel restore machinery +removed in TF 2.x fails inside `load_v1_in_v2.restore_variables` with +`TypeError: Binding inputs to tf.function failed due to 'too many positional +arguments'` (the pruned restore function has an empty signature but is called +with the variables path). The dependency floor already forces TF >= 2.13, so +`tf.saved_model.load` on a V1 SavedModel is not a supported path. + +**Code:** `grad_dft/functional.py`, `generate_DM21_weights` +**Tests:** `test_dm21` (both parametrized molecules) must now pass end-to-end +unchanged (`atol=1`), which also validates the entire PySCF+shim leg against +the independently implemented JAX port of DM21; the whole +`test_functional_implementations.py` file must stay green. +**Docs:** No user-facing change needed — `generate_DM21_weights`'s interface +(`folder` default, return value) is unchanged. + +**Spec:** +- Replace `variables = tf.saved_model.load(folder).variables` with a direct + checkpoint read: `tf.train.load_checkpoint(os.path.join(folder, + "variables", "variables"))` plus `tf.train.list_variables(...)`, appending + the `":0"` op suffix to each name so the existing name patterns in + `vars_to_params` (`"/w:"`, `"/b:"`, `"gamma:"`, `"beta:"`) still match. +- `vars_to_params` iterates `(name, value)` pairs; `tf_tensor_to_jax` becomes + a plain `jnp.asarray(value)` (the checkpoint reader returns numpy arrays). +- Comment in code: the vendored checkpoint is a TF1 SavedModel that + `tf.saved_model.load` can no longer import on TF >= 2.13. +- Commit: `(#6-S3): read the vendored TF1 DM21 checkpoint directly`. + +- [ ] S3: Implement +- [ ] S3: Write tests (`test_dm21` end-to-end; file stays green) +- [ ] S3: Update documentation (no user-facing change needed, justified above) +- [ ] S3: Pre-Commit Check (black + pylint, no new messages vs S2's parent) +- [ ] S3: Quality checks (no duplication: reader is local to the method) +- [ ] S3: Commit + +### S4: Cover the restricted (`spin=0`) path with an RKS-vs-UKS consistency test **Code:** new test `test_dm21_rks` in `tests/integration/molecules/test_functional_implementations.py` (after @@ -241,18 +303,18 @@ documents individual test names). `eval_xc_eff` and change the test to assert that raise. The test-coverage route is tried first because the restricted route is the documented usage in the class docstring (`mf = dft.RKS(...)`). -- Commit (last subtask): `(#6-S3): cover the DM21 restricted path with an +- Commit (last subtask): `(#6-S4): cover the DM21 restricted path with an RKS-vs-UKS consistency test` plus a standalone `Closes #6` line. ## Acceptance criteria mapping | Criterion (issue #6) | Where | |---|---| -| `test_dm21` passes, no tolerance change | S2 Tests (`atol=1` untouched) | -| `spin=0` covered by a new test or explicitly rejected | S3 (test first, `NotImplementedError` fallback pre-decided) | +| `test_dm21` passes, no tolerance change | S2 (shim) + S3 (checkpoint reader) — `atol=1` untouched | +| `spin=0` covered by a new test or explicitly rejected | S4 (test first, `NotImplementedError` fallback pre-decided) | | `hub.load` port and PySCF shim in separate commits | S1 / S2 commits | | `export_functional_and_derivatives()` explicit outcome | S1 (`NotImplementedError` + docstring) | -After S3: whole-repo code review (`code-review` skill), then the quality gate +After S4: whole-repo code review (`code-review` skill), then the quality gate (`quality-gates` skill: CI-file pytest invocations + `black` + `pylint`), then integration into `agent_settings` after user confirmation. From 256c6ebde3e91a9c737c7dd218a231beef62c4a9 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 19:49:17 +0300 Subject: [PATCH 13/17] (#6-S3): read the vendored TF1 DM21 checkpoint directly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With the eval_xc_eff shim (S2) in place, test_dm21 moved to its second leg: `DM21.generate_DM21_weights()` loads the same vendored checkpoint with `tf.saved_model.load(folder)`. That is a TensorFlow 1 SavedModel, and on TF >= 2.13 (the declared floor) its V1 restore machinery in load_v1_in_v2.restore_variables fails inside the lifted restore function (`TypeError: Binding inputs to tf.function failed due to 'too many positional arguments'` — the pruned function has an empty signature but is called with the variables path), so the test could never complete. Instead of importing the SavedModel, read the variables straight out of its checkpoint files with `tf.train.load_checkpoint(folder/variables/variables)` + `tf.train.list_variables`, appending the ":0" op suffix so the existing name patterns in `vars_to_params` ("/w:", "/b:", "gamma:", "beta:") still match. `vars_to_params` now iterates `(name, value)` pairs and `tf_tensor_to_jax` is a plain `jnp.asarray` (the reader returns numpy arrays directly, no .numpy() call needed). The method's interface (folder default, return value) is unchanged, and the whole inspection stays local to the method. This completes test_dm21 for the first time on these dependencies: both parametrized molecules pass end-to-end with the unchanged `atol=1`, validating the S1+S2 PySCF leg against grad_dft's independently implemented JAX port of DM21. Gate for this subtask, with PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py: tests/integration/molecules/test_functional_implementations.py -> 14 passed (0 failures). Pylint vs the S2 parent: symbol multiset identical (no new messages; the line-shifted redefined-outer-name/too-many-statements are the pre-existing ones). Black introduces no new debt: the baseline already fails `black --check` (task #7's repo-wide reformat). --- grad_dft/functional.py | 61 +++++++++++++++++++++++------------------- tasks/detailed_plan.md | 18 ++++++++----- 2 files changed, 45 insertions(+), 34 deletions(-) diff --git a/grad_dft/functional.py b/grad_dft/functional.py index cc27296..53d2afb 100644 --- a/grad_dft/functional.py +++ b/grad_dft/functional.py @@ -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)) diff --git a/tasks/detailed_plan.md b/tasks/detailed_plan.md index a652f1f..28e628c 100644 --- a/tasks/detailed_plan.md +++ b/tasks/detailed_plan.md @@ -232,7 +232,7 @@ semantics (delegates to `eval_xc`), the returned layout table for `spin=0`/ `eval_xc`; no tolerance or assertion touched) - [x] S2: Commit -### S3: Read the vendored TF1 DM21 checkpoint directly in `generate_DM21_weights` +### S3: Read the vendored TF1 DM21 checkpoint directly in `generate_DM21_weights` [done] **Discovered while running S2's tests** — not part of the original issue analysis: with the `eval_xc_eff` shim in place, `mf.kernel()` completes, and @@ -265,12 +265,16 @@ the independently implemented JAX port of DM21; the whole `tf.saved_model.load` can no longer import on TF >= 2.13. - Commit: `(#6-S3): read the vendored TF1 DM21 checkpoint directly`. -- [ ] S3: Implement -- [ ] S3: Write tests (`test_dm21` end-to-end; file stays green) -- [ ] S3: Update documentation (no user-facing change needed, justified above) -- [ ] S3: Pre-Commit Check (black + pylint, no new messages vs S2's parent) -- [ ] S3: Quality checks (no duplication: reader is local to the method) -- [ ] S3: Commit +- [x] S3: Implement (direct checkpoint read; `vars_to_params` on `(name, + value)` pairs; interface unchanged) +- [x] S3: Write tests (`test_dm21` end-to-end: 2 passed; whole file 14 passed, + `atol=1` untouched) +- [x] S3: Update documentation (no user-facing change needed, justified above) +- [x] S3: Pre-Commit Check (pylint vs S2's parent: symbol multiset identical; + black: no new debt — baseline `black --check` already fails, task #7) +- [x] S3: Quality checks (no duplication: reader stays local to the method; + no tolerance or assertion touched) +- [x] S3: Commit ### S4: Cover the restricted (`spin=0`) path with an RKS-vs-UKS consistency test From 840e2d8b6616d5a1a726edc00b1bfb0d425368a3 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 20:01:47 +0300 Subject: [PATCH 14/17] (#6-S4): cover the DM21 restricted path with an RKS-vs-UKS consistency test Closes #6 --- tasks/detailed_plan.md | 16 +++++++++- .../test_functional_implementations.py | 29 ++++++++++++++++++- 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/tasks/detailed_plan.md b/tasks/detailed_plan.md index 28e628c..9fb68d5 100644 --- a/tasks/detailed_plan.md +++ b/tasks/detailed_plan.md @@ -276,7 +276,7 @@ the independently implemented JAX port of DM21; the whole no tolerance or assertion touched) - [x] S3: Commit -### S4: Cover the restricted (`spin=0`) path with an RKS-vs-UKS consistency test +### S4: Cover the restricted (`spin=0`) path with an RKS-vs-UKS consistency test [done] **Code:** new test `test_dm21_rks` in `tests/integration/molecules/test_functional_implementations.py` (after @@ -310,6 +310,20 @@ documents individual test names). - Commit (last subtask): `(#6-S4): cover the DM21 restricted path with an RKS-vs-UKS consistency test` plus a standalone `Closes #6` line. +- [x] S4: Implement (new `test_dm21_rks` after `test_dm21`, reusing `mols[0]` + and the existing imports; local `mol_rks` avoids shadowing the module-level + `mol`) +- [x] S4: Write tests (`test_dm21_rks` passes standalone — 1 passed in 144 s; + whole file 15 passed in 303 s — the RKS and UKS energies agree within + `atol=1`, neither NaN) +- [x] S4: Update documentation (no user-facing change needed — test-only; the + rationale is recorded in the test's docstring) +- [x] S4: Pre-Commit Check (pylint vs S3's parent: symbol multiset identical; + black: no new debt — baseline `black --check` already fails, task #7) +- [x] S4: Quality checks (no duplication: reuses `mols[0]`, `NeuralNumInt`, + `Functional.DM21`; no tolerance or assertion touched) +- [x] S4: Commit (Closes #6) + ## Acceptance criteria mapping | Criterion (issue #6) | Where | 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 From f9691722781dd47fa0b3cb338bbcbd65285c8664 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Tue, 6 Oct 2026 22:40:37 +0300 Subject: [PATCH 15/17] (#10-S1): shim pkg_resources.parse_version for tensorflow-hub imports setuptools >= 81 removed pkg_resources, and a fresh pip install resolves setuptools (an unpinned requirement of pyscf) to 84.0.0, so CI crashed at collection with ModuleNotFoundError: No module named 'pkg_resources'. tensorflow-hub 0.16.1 (the newest release) still imports pkg_resources.parse_version in _ensure_tf_install() at import time, reached via the vendored neural_numint. Provide a stdlib-only stand-in for that one symbol on the import path; no dependency pins or new deps. Closes #10 --- .../neural_numint.py | 18 + tasks/detailed_plan.md | 429 ++++-------------- 2 files changed, 118 insertions(+), 329 deletions(-) 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 a1dc42c..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 @@ -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 diff --git a/tasks/detailed_plan.md b/tasks/detailed_plan.md index 9fb68d5..8712646 100644 --- a/tasks/detailed_plan.md +++ b/tasks/detailed_plan.md @@ -1,338 +1,109 @@ -# Detailed Plan — #6 Restore the DM21 test: port the vendored numint to TF-Hub 0.16 and PySCF 2.13 +# Detailed Plan — #10 CI collection fails: tensorflow-hub imports pkg_resources, removed in setuptools >= 81 Integration branch: `agent_settings` (stacked on `main`). -Feature branch: `restore_dm21_test` (created from the rebased integration branch). - -Problem: `tests/integration/molecules/test_functional_implementations.py::test_dm21` -fails on the current dependency set (`tensorflow 2.21.0`, `tensorflow-hub 0.16.1`, -`pyscf 2.13.1`) with two independent upstream breaks stacked in the vendored -DeepMind code `grad_dft/external/density_functional_approximation_dm21/.../neural_numint.py`: - -1. **TF-Hub 0.16 removed the TF1 `Module` API** — `hub.Module` raised - `AttributeError` at `neural_numint.py:219` (reproduced on the feature branch - before any change). -2. **PySCF 2.13's `eval_xc_eff` no longer routes through `eval_xc`** — - `nr_rks`/`nr_uks` now call `ni.eval_xc_eff(...)`, which is implemented as - `eval_xc1` + `xc_deriv.transform_xc`. `NeuralNumInt` overrides `eval_xc` but - not `eval_xc1`, so the call falls through to libxc, which resolves the - SCF's `xc_code` (the default `LDA,VWN`) against a meta-GGA-shaped density and - dies with `ValueError: cannot reshape array of size 240720 into shape - (2,1,24072)`. `NeuralNumInt.eval_xc` is dead code today. - -The issue's acceptance criteria: `test_dm21` passes with no tolerance change; -the `spin=0` path is either covered by a new test or explicitly rejected with a -clear `NotImplementedError`; the `hub.load` port and the PySCF shim are separate -commits; `export_functional_and_derivatives()` gets an explicit outcome. - -## Verification performed during planning (prototypes in /tmp/opencode, no repo changes) - -- **`hub.load` port is viable.** `hub.load(checkpoint).signatures["default"]` - loads the vendored TF1 SavedModel under `tf.disable_v2_behavior()` graph mode. - The signature's keyword keys keep the `$` separator (`tensor_dict$rho_a`, - only the TensorSpec `name=` fields are sanitised to `_`), so - `self._functional(**{f"tensor_dict${k}": v ...})` works as-is — the issue's - `$`→`_` mangling claim is not needed. The output key is `grid_contribution` - (unchanged). `tf.gradients` through the resulting `PartitionedCall` works and - returns finite gradients with the expected shapes (`vrho` 2×(N,1), `vsigma` - 3×(N,1), `vtau` 2×(N,1), `vhf` 2×(N,2)) — i.e. the whole existing - derivative machinery survives the port untouched. -- **The `eval_xc_eff` chain rule is settled, not guessed.** The layout PySCF - expects was validated against PySCF's own `NumInt.eval_xc_eff` on the real - libxc functional `mgga_x_tpss` (spin=1 and spin=0): building the effective - density-parameter derivatives from `eval_xc`'s `(vrho, vsigma, vtau)` with - - | spin | row layout of `vxc` | - |---|---| - | 0 | `(5, N)`: `[vrho, 2·vsigma·∇ρ, vtau]` | - | 1 | `(2, 5, N)`: `[:,0]=vrho`, `[0,1:4]=2·vsigma_aa·∇ρ_a + vsigma_ab·∇ρ_b`, `[1,1:4]=2·vsigma_bb·∇ρ_b + vsigma_ab·∇ρ_a`, `[:,4]=vtau` | - - reproduces `eval_xc_eff`'s output **exactly (max |diff| = 0.0)** on both spins. - (The alternative of feeding `eval_xc`'s outputs into - `pyscf.dft.xc_deriv.transform_xc` directly was prototyped too and produced - garbage — the `out` row order follows the xcfun convention and `transform_xc` - reads past the initialised rows for our input — so the explicit chain rule is - the design of record.) -- **Row counts.** `eval_rho` for MGGA runs with `with_lapl=False` - (`MGGA_DENSITY_LAPL = False`), so rho arrives with 5 rows (or `(2,5,N)`), - while the graph's placeholders expect 6 rows; a zero laplacian row is - inserted at index 4 (the graph discards index 4 anyway). -- **`tensorflow-hub<0.16` is unusable on the declared TF.** Verified from the - wheels: 0.14.0 and 0.15.0 both do - `from tensorflow_hub.estimator import ...` at `__init__.py:90`, which imports - `tensorflow.compat.v1.estimator`, absent from TF 2.21. The honest floor for - this repo is `tensorflow-hub>=0.16.0` (root `setup.py` reads - `requirements.txt`, so that one line is the whole dependency change). -- **The restricted path's local-HF features are spin-consistent:** - `compute_hfx_density.get_hf_density` splits a restricted `dm` into `dm/2` per - spin internally, so the RKS and UKS paths evaluate the same functional at - `ρ_a = ρ_b`; an RKS-vs-UKS energy comparison is therefore a sound end-to-end - test of the `spin=0` layout. +Feature branch: `fix_ci_pkg_resources` (created from the rebased integration branch). + +Problem: CI (PR #8 head `agent_settings` @ 840e2d8, Python 3.12 / ubuntu-latest) fails at +collection with `ModuleNotFoundError: No module named 'pkg_resources'`. A fresh +`pip install -e .` resolves `setuptools` (an unpinned requirement of +`pyscf>=2.3.0`) to 84.0.0, which removed `pkg_resources`; `tensorflow-hub +0.16.1` (the newest release, and the floor `>=0.16.0` we set in #6) still does +`from pkg_resources import parse_version` inside `_ensure_tf_install()` at +import time, so any `import tensorflow_hub` — reached via the vendored +`neural_numint.py` — crashes during test collection. + +The local environment (setuptools 81.0.0, `pkg_resources` still shipped) could +not reproduce this, which is why the #6 gate passed while CI failed. Verified +against the real CI env: installing `setuptools==84.0.0` removes +`pkg_resources` locally and reproduces the exact failure. ## Reuse analysis -- **`NeuralNumInt.eval_xc` (existing)** — the shim does not re-implement any - DM21 evaluation: `eval_xc_eff` routes to it, reusing its derivative - computation, its float64 casts, and its `_vmat_hf` (local-HF potential) - side effects exactly as before the PySCF upgrade. -- **PySCF's own contract** — the layout table above is PySCF's - `NumInt.eval_xc_eff` contract, not a new invention; it was confirmed against - `pyscf.dft.numint.NumInt.eval_xc_eff` empirically (see Verification). -- **Test file conventions** — the new test reuses the module-level `mols` list, - `dft.RKS`/`dft.UKS` + `NeuralNumInt` setup of the existing `test_dm21`, the - `Hartree2kcalmol` conversion already imported at - `tests/.../test_functional_implementations.py:24`, and the file's `atol=1` - kcal/mol tolerance style. No new helpers. -- No `grad_dft/utils/` additions are needed; nothing else in the repo calls - `eval_xc`/`eval_xc_eff` (verified by grep: only `neural_numint.py` itself and - `grad_dft/interface/pyscf.py:1050`, which only monkey-patches `mf._numint`). - -### S1: Port `neural_numint.py` to the TF-Hub 0.16 API (`hub.load`) and raise the declared hub floor [done] - -**Code:** `grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py` -(`_build_graph`, `export_functional_and_derivatives`), `requirements.txt` -(`tensorflow-hub>=0.14.0` → `tensorflow-hub>=0.16.0`) -**Tests:** run the new test only via construction smoke test -(`python -c "from grad_dft.external import NeuralNumInt, Functional; -NeuralNumInt(Functional.DM21)"` must build the graph), then -`pytest -v tests/integration/molecules/test_functional_implementations.py::test_dm21` -— expected to fail at the **documented second break** -(`ValueError: cannot reshape array of size 240720 into shape (2,1,24072)`), not -at `AttributeError: hub.Module`; that error is the S2 input. The suite must -otherwise be unchanged (all other tests in the file still pass). -**Docs:** No README/example change needed (TF-Hub internals are invisible to -users). The `export_functional_and_derivatives` docstring is updated in-code to -state the new `NotImplementedError` behaviour and why; the class docstring's -usage example is untouched. - -**Spec:** -- In `_build_graph`: - - Replace `self._functional = hub.Module(spec=self._model_path)` with - `self._functional = hub.load(self._model_path).signatures["default"]`. - - Keep `tensor_dict = {f"tensor_dict${k}": v for k, v in features.items()}` - and change `predictions = self._functional(tensor_dict, as_dict=True)` to - `predictions = self._functional(**tensor_dict)`; `local_xc = - predictions["grid_contribution"]` stays (verified output key). - - Delete the `outputs` dict and the trailing - `hub.add_signature(inputs=..., outputs=...)` call (API removed in 0.16; the - only consumer was the spec-based export, which now raises — see below), and - drop the "no-op if called outside hub.create_module_spec" comment. - - Everything else (placeholders, feature construction, `tf.gradients`, - session) is untouched — verified to work through the port. -- Replace the body of `export_functional_and_derivatives` with a - `raise NotImplementedError` whose message says: TF-Hub >= 0.16 removed the - TF1 `Module`/`create_module_spec` API this export was built on; exporting - would require a TF2 `tf.saved_model.save` reimplementation; the method is - only used by the vendored `export_saved_model.py` / `neural_numint_test.py`, - which are not part of the test suite. Keep the original docstring and append - the Raises section. -- `requirements.txt`: raise the floor to `tensorflow-hub>=0.16.0` (0.14/0.15 - import `tensorflow_hub.estimator` → `tensorflow.compat.v1.estimator`, which no - longer exists in TF 2.21; verified from the published wheels). -- Style: `black` is **not** applied to this vendored file — the baseline already - fails `black --check`, and a file-wide reformat is out of scope (task #5 - precedent: edits stay diff-only to limit churn/upstream divergence). `pylint - -rn -sn --rcfile=.pylintrc` must show no new messages vs the subtask parent. -- Commit: `(#6-S1): port neural_numint to the TF-Hub 0.16 hub.load API`. - -- [x] S1: Implement (`hub.load` port, `outputs`/`hub.add_signature` removal, - `NotImplementedError` export, hub floor `>=0.16.0`) -- [x] S1: Write tests (no new test in this subtask; smoke test builds the graph - and `WrappedFunction`, and with `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py` - the file gives 12 passed / 2 failed — `test_dm21` alone, now at the - documented second break, the reshape `ValueError`, instead of - `AttributeError: hub.Module`) -- [x] S1: Update documentation (no user-facing change needed — TF-Hub internals - are invisible to users; the `Raises` section added to - `export_functional_and_derivatives`' docstring is the explicit outcome) -- [x] S1: Pre-Commit Check (`black` deliberately not applied: the baseline - already fails `black --check`, so the file stays as-is per the plan's - style note; `pylint -rn -sn --rcfile=.pylintrc` vs subtask parent: no new - messages, and the 4 baseline `E1101 tensorflow_hub no-member` errors are - gone) -- [x] S1: Quality checks (no duplication: single evaluation path retained; - nothing added to library modules outside the vendored shim; no tolerance - or assertion touched) -- [x] S1: Commit - -### S2: Add the `eval_xc_eff` override that routes PySCF's derivative contract to `eval_xc` [done] - -**Code:** new method `NeuralNumInt.eval_xc_eff` in -`grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py` -(placed directly after `eval_xc`); one class-docstring sentence updated -("The actual evaluation ... is performed in NeuralNumInt.eval_xc" → mention -that PySCF reaches it through the `eval_xc_eff` override) -**Tests:** `pytest -v tests/integration/molecules/test_functional_implementations.py` -— with the shim, `test_dm21`'s PySCF leg (`mf.kernel()` on both parametrized -molecules) completes and the other 9 tests stay green; the test as a whole -**cannot pass until S3** because its second leg, -`generate_DM21_weights`, hits an independent break (TF1 `tf.saved_model.load` -on TF 2.21 — see S3). The `atol=1` tolerance is untouched. Also run -`pytest -v tests/integration/molecules/test_predict_B88.py -tests/integration/molecules/test_non_xc_energy.py tests/integration/molecules/test_Harris.py` -to confirm no collateral damage. -**Docs:** numpydoc docstring on `eval_xc_eff` documenting the argument -semantics (delegates to `eval_xc`), the returned layout table for `spin=0`/ -`spin=1`, and the Raises conditions; no README/example change (public usage of -`NeuralNumInt` is unchanged). - -**Spec:** -- Signature mirrors PySCF's base method exactly: - `def eval_xc_eff(self, xc_code, rho, deriv=1, omega=None, xctype=None, verbose=None, spin=None)`. -- Behaviour: - 1. Mirror the base class preamble: `omega = self.omega if omega is None else - omega` (`NumInt.omega` is `None`, so a plain call passes `omega=None` - through to `eval_xc`, which accepts `None` and raises a clear - `NotImplementedError` for an explicit RSH value — unchanged semantics); - normalise `rho` with `np.asarray(rho, order="C")`; infer `spin` from - `rho.shape[0] == 2` when it is not given (same rule as the base class). - 2. `deriv == 0` → still evaluate (the `eval_xc` call carries the `_vmat_hf` - side effect) but return `[exc, None, None, None]`; `deriv > 1` → raise - `NotImplementedError` ("DM21 provides no second derivatives"; the base - class would die inside `transform_xc`). - 3. Re-insert the laplacian row: if `rho.shape[-2] == 5`, insert a zero row at - index 4 (graph placeholders are `(6, N)`; index 4 is discarded by the - graph). If 6 rows arrive, pass through. - 4. `exc, (vrho, vsigma, vlapl, vtau), _, _ = self.eval_xc(xc_code, rho6, - spin=spin, deriv=deriv, omega=omega, verbose=verbose)` — full reuse of the - existing evaluation, including `_vmat_hf` accumulation. - 5. Build the effective derivatives with the validated chain rule - (outputs of `eval_xc` are `(N, ncomp)` — transpose to `(ncomp, N)` for the - spin-polarised case): - - `spin == 0`: `vxc = np.concatenate([vrho[None, :], - 2·vsigma[None, :]·rho[1:4], vtau[None, :]])` → shape `(5, N)`, rows - `[dE/dρ, dE/d∇ρ (3), dE/dτ]` — **no laplacian row** (this line - previously claimed a `zeros` row at index 2; re-validated against - `NumInt.eval_xc_eff` on `mgga_x_tpss`, max |diff| = 0.0, so the layout - table in Verification is the correct one). - - `spin == 1`: shape `(2, 5, N)` with - `vxc[:,0] = vrho.T`, `vxc[0,1:4] = 2·vsigma[:,0]·∇ρ_a + vsigma[:,1]·∇ρ_b`, - `vxc[1,1:4] = 2·vsigma[:,2]·∇ρ_b + vsigma[:,1]·∇ρ_a`, - `vxc[:,4] = vtau.T`, where `∇ρ_a = rho[0,1:4]`, `∇ρ_b = rho[1,1:4]`. - 6. Return `(exc, vxc, None, None)` (float64 — `eval_xc` already casts). -- The docstring states the layout contract and notes the chain rule was - verified against `pyscf.dft.numint.NumInt.eval_xc_eff`. -- Comment where the chain rule lives, pointing at the PySCF call sites - (`nr_rks` MGGA branch: `wv[0]/wv[4]` halving, `_scale_ao_sparse(ao[:4], - wv[:4])`, `_tau_dot_sparse(..., wv[4])`; `nr_uks` the same per spin row) so a - future reader can re-derive it. -- Commit: `(#6-S2): add an eval_xc_eff shim routing PySCF 2.13 back to eval_xc`. - -- [x] S2: Implement (`eval_xc_eff` override + class-docstring sentence) -- [x] S2: Write tests (chain rule vs `NumInt.eval_xc_eff` on `mgga_x_tpss`: - max |diff| = 0.0 for both spins; `test_dm21`'s PySCF leg — `mf.kernel()` - on both molecules — completes; collateral files green: 45 passed across - `test_predict_B88.py`, `test_non_xc_energy.py`, `test_Harris.py`; the - full `test_dm21` is deferred to S3, see Tests) -- [x] S2: Update documentation (numpydoc docstring on `eval_xc_eff` with the - layout contract and raises; class docstring updated; no user-facing change) -- [x] S2: Pre-Commit Check (`black` not applied to the vendored file, same - precedent as S1; `pylint -rn -sn --rcfile=.pylintrc` vs S1's parent: no - new messages) -- [x] S2: Quality checks (no duplication: single evaluation path reused via - `eval_xc`; no tolerance or assertion touched) -- [x] S2: Commit - -### S3: Read the vendored TF1 DM21 checkpoint directly in `generate_DM21_weights` [done] - -**Discovered while running S2's tests** — not part of the original issue -analysis: with the `eval_xc_eff` shim in place, `mf.kernel()` completes, and -`test_dm21` moves to its second leg, `DM21().generate_DM21_weights()` -(`grad_dft/functional.py`), which loads the same vendored checkpoint with -`tf.saved_model.load(folder)`. On TF 2.21 the V1 SavedModel restore machinery -removed in TF 2.x fails inside `load_v1_in_v2.restore_variables` with -`TypeError: Binding inputs to tf.function failed due to 'too many positional -arguments'` (the pruned restore function has an empty signature but is called -with the variables path). The dependency floor already forces TF >= 2.13, so -`tf.saved_model.load` on a V1 SavedModel is not a supported path. - -**Code:** `grad_dft/functional.py`, `generate_DM21_weights` -**Tests:** `test_dm21` (both parametrized molecules) must now pass end-to-end -unchanged (`atol=1`), which also validates the entire PySCF+shim leg against -the independently implemented JAX port of DM21; the whole -`test_functional_implementations.py` file must stay green. -**Docs:** No user-facing change needed — `generate_DM21_weights`'s interface -(`folder` default, return value) is unchanged. - -**Spec:** -- Replace `variables = tf.saved_model.load(folder).variables` with a direct - checkpoint read: `tf.train.load_checkpoint(os.path.join(folder, - "variables", "variables"))` plus `tf.train.list_variables(...)`, appending - the `":0"` op suffix to each name so the existing name patterns in - `vars_to_params` (`"/w:"`, `"/b:"`, `"gamma:"`, `"beta:"`) still match. -- `vars_to_params` iterates `(name, value)` pairs; `tf_tensor_to_jax` becomes - a plain `jnp.asarray(value)` (the checkpoint reader returns numpy arrays). -- Comment in code: the vendored checkpoint is a TF1 SavedModel that - `tf.saved_model.load` can no longer import on TF >= 2.13. -- Commit: `(#6-S3): read the vendored TF1 DM21 checkpoint directly`. - -- [x] S3: Implement (direct checkpoint read; `vars_to_params` on `(name, - value)` pairs; interface unchanged) -- [x] S3: Write tests (`test_dm21` end-to-end: 2 passed; whole file 14 passed, - `atol=1` untouched) -- [x] S3: Update documentation (no user-facing change needed, justified above) -- [x] S3: Pre-Commit Check (pylint vs S2's parent: symbol multiset identical; - black: no new debt — baseline `black --check` already fails, task #7) -- [x] S3: Quality checks (no duplication: reader stays local to the method; - no tolerance or assertion touched) -- [x] S3: Commit - -### S4: Cover the restricted (`spin=0`) path with an RKS-vs-UKS consistency test [done] - -**Code:** new test `test_dm21_rks` in -`tests/integration/molecules/test_functional_implementations.py` (after -`test_dm21`), reusing the module-level `mols[0]` closed-shell molecule and the -existing imports -**Tests:** the new test itself: build `dft.RKS(mol)` and `dft.UKS(mol)` on the -closed-shell HF molecule, give both `mf._numint = NeuralNumInt(Functional.DM21)` -(identical default grids, mirroring `test_dm21`'s setup exactly), run both -`mf.kernel()` calls, and assert `(e_rks - e_uks) * Hartree2kcalmol` is -`allclose(..., 0, atol=1)` and that neither energy is NaN. The whole file must -stay green. -**Docs:** No user-facing change needed (test-only; nothing in README/examples -documents individual test names). - -**Spec:** -- Rationale (recorded in the test as a comment): DM21 maps restricted inputs to - `ρ_a = ρ_b = ρ/2` and `get_hf_density` splits a restricted `dm` into `dm/2` - per spin, so on a closed-shell molecule the restricted and unrestricted - self-consistent solutions must agree; UKS is already validated end-to-end by - `test_dm21`. A wrong `spin=0` derivative layout therefore shows up as an - energy mismatch (or a diverged/NaN SCF), not as a silent pass. -- Only the closed-shell molecule (`mols[0]`, HF) is used — RKS on the open-shell - Li atom (`mols[1]`) is not a valid restricted calculation. -- Fallback (decided in advance, per the issue): if this test demonstrates that - the `spin=0` path cannot be made to agree (i.e. the failure is not the layout - shim's to fix), replace the RKS path with an explicit - `raise NotImplementedError("The restricted (spin=0) path ...")` in - `eval_xc_eff` and change the test to assert that raise. The test-coverage - route is tried first because the restricted route is the documented usage in - the class docstring (`mf = dft.RKS(...)`). -- Commit (last subtask): `(#6-S4): cover the DM21 restricted path with an - RKS-vs-UKS consistency test` plus a standalone `Closes #6` line. - -- [x] S4: Implement (new `test_dm21_rks` after `test_dm21`, reusing `mols[0]` - and the existing imports; local `mol_rks` avoids shadowing the module-level - `mol`) -- [x] S4: Write tests (`test_dm21_rks` passes standalone — 1 passed in 144 s; - whole file 15 passed in 303 s — the RKS and UKS energies agree within - `atol=1`, neither NaN) -- [x] S4: Update documentation (no user-facing change needed — test-only; the - rationale is recorded in the test's docstring) -- [x] S4: Pre-Commit Check (pylint vs S3's parent: symbol multiset identical; - black: no new debt — baseline `black --check` already fails, task #7) -- [x] S4: Quality checks (no duplication: reuses `mols[0]`, `NeuralNumInt`, - `Functional.DM21`; no tolerance or assertion touched) -- [x] S4: Commit (Closes #6) +- `pkg_resources` is referenced in `tensorflow_hub` exactly once at import time + (`tensorflow_hub/__init__.py:61`, `parse_version` used only to compare + `tf.__version__ >= 1.15.0`). No other package in our import chain uses it. +- The vendored `neural_numint.py` is the only repo module that imports + `tensorflow_hub`; it is already patched by #6 (hub.load port), so the shim + fits the established "surgical edits to the vendored file" precedent. +- No newer `tensorflow-hub` release exists (0.16.1 is latest; pip resolved it + in CI). Pinning `setuptools<81` in `requirements.txt` is rejected: it fixes + only this symbol via a build-tool downgrade, adds a dependency-resolution + constraint that rots, and does not future-proof the import. + +## Subtasks + +### S1: Shim `pkg_resources.parse_version` on the tensorflow-hub import path + +**Code:** in +`grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py`, +immediately before `import tensorflow_hub as hub`: + +```python +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 +``` + +Notes: +- `parse_version` semantics: `pkg_resources.parse_version("2.21.0")` returns a + comparable version key; a dotted-numeric `int` tuple compares identically for + the only use (`tf.__version__` vs `"1.15.0"`, both clean release tags). +- No external deps (e.g. `packaging`) added; zero-dependency and + setuptools-version-proof. +- `_sys.modules.setdefault` keeps a pre-existing `pkg_resources` untouched and + only fills the gap. + +**Tests:** reproduce the CI failure and the fix: +- `/tmp/opencode/repro_pkg_resources.py` (a `sys.meta_path` finder that raises + `ModuleNotFoundError` for `pkg_resources`, then imports `grad_dft.external`): + without the shim → the CI error; with it → `IMPORT OK`. +- Real-env check: `pip install setuptools==84.0.0` (removes `pkg_resources`, + matching CI) → `pytest -v tests/unit/test_eigenproblem.py` passes. +- Full CI suite green under setuptools 84. +- `pylint -rn -sn --rcfile=.pylintrc` on the vendored file: symbol multiset + identical to subtask parent (no new messages — verified: the shim adds none). +- `black`: vendored file not reformatted (same precedent as #6; baseline debt + unchanged). + +**Docs:** no user-facing change — nothing in README/examples describes the +vendored import internals. The shim comment records the reason. + +**Quality checks:** no duplication (the shim lives on the single import path); +no tolerance or assertion touched; scope limited to the import shim. + +**Commit:** `(#10-S1): shim pkg_resources.parse_version for tensorflow-hub imports` +plus a standalone `Closes #10`. + +- [x] S1: Implement (stdlib-only shim before `import tensorflow_hub`; comment + documents the setuptools >= 81 removal) +- [x] S1: Write tests (meta-path repro fails pre-fix / passes post-fix; real + `setuptools==84.0.0` env: `test_eigenproblem.py` 2 passed; full CI suite + 123 passed) +- [x] S1: Pre-Commit Check (pylint vs parent: symbol multiset identical; black: + vendored file untouched, no new debt) +- [x] S1: Quality checks (no duplication — single import path; zero new deps; + no tolerance/assertion touched) +- [x] S1: Commit (Closes #10) ## Acceptance criteria mapping -| Criterion (issue #6) | Where | +| Criterion | Where | |---|---| -| `test_dm21` passes, no tolerance change | S2 (shim) + S3 (checkpoint reader) — `atol=1` untouched | -| `spin=0` covered by a new test or explicitly rejected | S4 (test first, `NotImplementedError` fallback pre-decided) | -| `hub.load` port and PySCF shim in separate commits | S1 / S2 commits | -| `export_functional_and_derivatives()` explicit outcome | S1 (`NotImplementedError` + docstring) | - -After S4: whole-repo code review (`code-review` skill), then the quality gate -(`quality-gates` skill: CI-file pytest invocations + `black` + `pylint`), then -integration into `agent_settings` after user confirmation. +| CI collection error gone (`test_eigenproblem.py` collects on a fresh 3.12 install) | S1: shim before `import tensorflow_hub` | +| No new lint/debt | S1: pylint multiset identical; black untouched (vendored) | +| Robust to future setuptools | S1: shim depends only on stdlib; `setuptools<81` pin rejected | \ No newline at end of file From 818500b734ab64ff96a0c44bd62bd0b77fd1ade8 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Wed, 7 Oct 2026 09:22:43 +0300 Subject: [PATCH 16/17] Untrack the tasks/ plan directory and gitignore it The agent workflow's plan files (tasks/detailed_plan.md) are personal workflow state whose single source of truth is the plan comment mirrored onto the task issue (<!-- detailed-plan -->). Tracking them in the public repository adds per-subtask churn to every PR diff (e.g. PR #8) and to maintainers' history. tasks/global_plan.md already worked this way (untracked); detailed_plan.md now follows suit. The file stays on disk and keeps being mirrored to issues exactly as before. --- .gitignore | 2 + tasks/detailed_plan.md | 109 ----------------------------------------- 2 files changed, 2 insertions(+), 109 deletions(-) delete mode 100644 tasks/detailed_plan.md diff --git a/.gitignore b/.gitignore index cfb38a2..1496eff 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,5 @@ runs/ */__pycache__/ .vscode/ *.egg-info/ +# Agent workflow plan files (single source of truth: mirrored issue comments) +tasks/ diff --git a/tasks/detailed_plan.md b/tasks/detailed_plan.md deleted file mode 100644 index 8712646..0000000 --- a/tasks/detailed_plan.md +++ /dev/null @@ -1,109 +0,0 @@ -# Detailed Plan — #10 CI collection fails: tensorflow-hub imports pkg_resources, removed in setuptools >= 81 - -Integration branch: `agent_settings` (stacked on `main`). -Feature branch: `fix_ci_pkg_resources` (created from the rebased integration branch). - -Problem: CI (PR #8 head `agent_settings` @ 840e2d8, Python 3.12 / ubuntu-latest) fails at -collection with `ModuleNotFoundError: No module named 'pkg_resources'`. A fresh -`pip install -e .` resolves `setuptools` (an unpinned requirement of -`pyscf>=2.3.0`) to 84.0.0, which removed `pkg_resources`; `tensorflow-hub -0.16.1` (the newest release, and the floor `>=0.16.0` we set in #6) still does -`from pkg_resources import parse_version` inside `_ensure_tf_install()` at -import time, so any `import tensorflow_hub` — reached via the vendored -`neural_numint.py` — crashes during test collection. - -The local environment (setuptools 81.0.0, `pkg_resources` still shipped) could -not reproduce this, which is why the #6 gate passed while CI failed. Verified -against the real CI env: installing `setuptools==84.0.0` removes -`pkg_resources` locally and reproduces the exact failure. - -## Reuse analysis - -- `pkg_resources` is referenced in `tensorflow_hub` exactly once at import time - (`tensorflow_hub/__init__.py:61`, `parse_version` used only to compare - `tf.__version__ >= 1.15.0`). No other package in our import chain uses it. -- The vendored `neural_numint.py` is the only repo module that imports - `tensorflow_hub`; it is already patched by #6 (hub.load port), so the shim - fits the established "surgical edits to the vendored file" precedent. -- No newer `tensorflow-hub` release exists (0.16.1 is latest; pip resolved it - in CI). Pinning `setuptools<81` in `requirements.txt` is rejected: it fixes - only this symbol via a build-tool downgrade, adds a dependency-resolution - constraint that rots, and does not future-proof the import. - -## Subtasks - -### S1: Shim `pkg_resources.parse_version` on the tensorflow-hub import path - -**Code:** in -`grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py`, -immediately before `import tensorflow_hub as hub`: - -```python -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 -``` - -Notes: -- `parse_version` semantics: `pkg_resources.parse_version("2.21.0")` returns a - comparable version key; a dotted-numeric `int` tuple compares identically for - the only use (`tf.__version__` vs `"1.15.0"`, both clean release tags). -- No external deps (e.g. `packaging`) added; zero-dependency and - setuptools-version-proof. -- `_sys.modules.setdefault` keeps a pre-existing `pkg_resources` untouched and - only fills the gap. - -**Tests:** reproduce the CI failure and the fix: -- `/tmp/opencode/repro_pkg_resources.py` (a `sys.meta_path` finder that raises - `ModuleNotFoundError` for `pkg_resources`, then imports `grad_dft.external`): - without the shim → the CI error; with it → `IMPORT OK`. -- Real-env check: `pip install setuptools==84.0.0` (removes `pkg_resources`, - matching CI) → `pytest -v tests/unit/test_eigenproblem.py` passes. -- Full CI suite green under setuptools 84. -- `pylint -rn -sn --rcfile=.pylintrc` on the vendored file: symbol multiset - identical to subtask parent (no new messages — verified: the shim adds none). -- `black`: vendored file not reformatted (same precedent as #6; baseline debt - unchanged). - -**Docs:** no user-facing change — nothing in README/examples describes the -vendored import internals. The shim comment records the reason. - -**Quality checks:** no duplication (the shim lives on the single import path); -no tolerance or assertion touched; scope limited to the import shim. - -**Commit:** `(#10-S1): shim pkg_resources.parse_version for tensorflow-hub imports` -plus a standalone `Closes #10`. - -- [x] S1: Implement (stdlib-only shim before `import tensorflow_hub`; comment - documents the setuptools >= 81 removal) -- [x] S1: Write tests (meta-path repro fails pre-fix / passes post-fix; real - `setuptools==84.0.0` env: `test_eigenproblem.py` 2 passed; full CI suite - 123 passed) -- [x] S1: Pre-Commit Check (pylint vs parent: symbol multiset identical; black: - vendored file untouched, no new debt) -- [x] S1: Quality checks (no duplication — single import path; zero new deps; - no tolerance/assertion touched) -- [x] S1: Commit (Closes #10) - -## Acceptance criteria mapping - -| Criterion | Where | -|---|---| -| CI collection error gone (`test_eigenproblem.py` collects on a fresh 3.12 install) | S1: shim before `import tensorflow_hub` | -| No new lint/debt | S1: pylint multiset identical; black untouched (vendored) | -| Robust to future setuptools | S1: shim depends only on stdlib; `setuptools<81` pin rejected | \ No newline at end of file From a4f1f89534451937d32c5ad6e4717b44c2d2bf62 Mon Sep 17 00:00:00 2001 From: Anya497 <chi.vinny0702@gmail.com> Date: Wed, 7 Oct 2026 09:28:02 +0300 Subject: [PATCH 17/17] (#11-S1): run CI only when the change can affect its outcome The pull_request trigger fired on every push (synchronize/opened/reopened), so non-code changes (.gitignore, the removed tasks/ plan churn) wasted a full ~15-minute suite run. Add a paths whitelist matching exactly the workflow's inputs: grad_dft/**, tests/**, requirements.txt, setup.py, .github/** (the workflow itself plus the pyscf_conf.py it consumes), and the lint configs that #7 will put into CI. Whitelist semantics: docs/notebook/ plan-file-only PRs skip CI entirely. .github/** is in the filter, so this very change re-triggers CI on push, self-validating the filter. Closes #11 --- .github/workflows/install_and_test.yml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.github/workflows/install_and_test.yml b/.github/workflows/install_and_test.yml index 815747d..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: