Implement exact analytic DeePHF nuclear forces and force-aware training

## Summary

Implement strict analytic nuclear forces for perturbative DeePHF and make those forces usable in training, validation, and inference.

The repository already supports:

- perturbative DeePHF energy training and energy-only testing;
- analytic forces for self-consistent DeePKS;
- a force loss of the form `-dE/ddesc * ddesc/dR` using the stored `grad_vx` field.

However, the current `grad_vx` is the explicit derivative of the projected-density descriptor at fixed AO density matrix. That is sufficient for the variational/self-consistent DeePKS force expression, but it is not the complete derivative of a perturbative DeePHF energy evaluated on geometry-dependent HF/KS orbitals. Exact DeePHF forces also require the response of the converged reference density/orbitals to nuclear displacement.

This issue proposes a complete molecular PySCF implementation based on CP-HF/CP-KS for data generation and a Z-vector formulation for efficient inference.

## Current gap

The current non-SCF test path only evaluates energy:

- `deepks/model/test.py` reads `lb_e` and `eig`, calls `model(data)`, and reports energy errors.

The existing force-training path is:

```python
gev = dE_corr / d(desc)
f_corr = -einsum(grad_vx, gev)
```

in `deepks/model/train.py`. The required `grad_vx` is generated in `deepks/scf/grad.py` from derivatives of the molecular/projector overlap integrals while the AO density matrix passed to `t_make_grad_eig_x` is held fixed.

The self-consistent force path in `deepks/scf/grad.py` correctly adds the DeePKS correction to a variational SCF gradient. It must not be reused unchanged for perturbative DeePHF because the DeePHF correction does not participate in the reference HF/KS stationarity equations.

The DeePHF initialization examples also currently generate and train energy-only data.

## Mathematical target

Let the converged reference density be `P0(R)`, satisfying the HF/KS stationarity equations, and let the DeePHF descriptor be

```text
q(R) = q(P0(R), R).
```

The perturbative DeePHF energy is

```text
E_DeePHF(R) = E_ref(P0(R), R) + E_theta(q(P0(R), R)).
```

The exact force is

```text
F_DeePHF = F_ref - g_q . J_relaxed

g_q = dE_theta / dq
J_relaxed = dq/dR
= (partial q/partial R)_P
+ (partial q/partial P) : dP0/dR.
```

The current `grad_vx` represents only the first, explicit term. The new implementation must include the orbital/density-response term obtained from CP-HF or CP-KS.

Two mathematically equivalent backends should be provided:

1. **Direct response:** solve the coupled-perturbed equations for every nuclear perturbation, construct `dP0/dR`, and form the complete descriptor Jacobian `J_relaxed`. This is the preferred data-generation and reference implementation because `J_relaxed` is independent of the neural-network parameters and can be stored once for differentiable force training.
2. **Z-vector response:** form the orbital gradient induced by the DeePHF correction potential, solve one adjoint response equation, and contract it with all nuclear perturbation right-hand sides. This is the preferred inference implementation because it avoids `3 * natom` separate response solves for a scalar energy.

Both backends must agree with each other and with central finite differences of the complete perturbative DeePHF energy.

## Scope

### Required

- Molecular PySCF backend.
- Restricted and unrestricted references.
- HF and KS references supported by the current molecular runner:
- RHF;
- UHF;
- RKS/CPKS when `xc != "HF"`;
- UKS/CPKS when `xc != "HF"`.
- Exact DeePHF correction and total forces.
- Exact force training using a relaxed descriptor Jacobian.
- Python API, CLI/configuration integration, data dumping, statistics, tests, and documentation.
- Geometry-scanner support so the method can be used by PySCF geometry optimizers.

### Out of scope for this issue

- Periodic/ABACUS response theory.
- Analytic DeePHF Hessians.
- Nonadiabatic derivatives.
- Differentiation through the reference SCF procedure with PyTorch.

Periodic support and analytic Hessians can be follow-up issues after the molecular implementation and data semantics are stable.

## Proposed architecture

### 1. Add a non-self-consistent DeePHF method object

Add a dedicated package such as:

```text
deepks/deephf/
__init__.py
method.py # DeePHF energy object and scanner
response.py # CP-HF/CP-KS response adapters
grad.py # direct-response and Z-vector gradients
```

The method should use composition around a converged native PySCF mean-field object instead of subclassing `DSCF`. This prevents the correction potential from accidentally entering the Fock iterations.

Proposed Python API:

```python
from pyscf import gto, scf
from deepks.deephf import DeePHF

mol = gto.M(...)
mf = scf.RHF(mol).run()
phf = DeePHF(mf, model="model.pth", proj_basis=...)

e_tot = phf.kernel()
grad = phf.nuc_grad_method().kernel() # dE/dR
force = -grad
```

Required properties/methods:

- `e_ref`, `e_corr`, and `e_tot`;
- `make_rdm1()` delegated to the reference calculation;
- descriptor and correction-potential evaluation;
- `nuc_grad_method(response_method="zvector")`;
- `as_scanner()` that refreshes reference SCF, projector integrals, descriptors, and response intermediates for every geometry;
- explicit errors when the reference SCF or response equations do not converge.

### 2. Extract reusable descriptor evaluation

The descriptor code currently lives inside the self-consistent SCF mixins. Extract or wrap the following operations so DeePKS and DeePHF use one implementation:

- projector molecule and overlap construction;
- projected density matrices;
- descriptor eigenvalues;
- `dE_corr/dP` (the correction potential);
- `partial q/partial P`;
- the explicit projector/AO-overlap contribution `(partial q/partial R)_P`;
- cache invalidation after geometry changes.

Suggested module: `deepks/scf/descriptor.py` or `deepks/model/descriptor.py`.

Existing public behavior of `DSCF`, `UDSCF`, `make_eig`, and `make_grad_eig_x` should remain compatible.

### 3. Direct CP-HF/CP-KS descriptor response

Implement a response adapter around PySCF's analytic response infrastructure rather than duplicating integral derivatives.

For each supported reference:

1. Build the nuclear first-order Fock/core and overlap perturbations using the corresponding PySCF Hessian implementation.
2. Solve the first-order MO equations in coordinate blocks using the PySCF CPHF/UCPHF solver and `mf.gen_response(...)`.
3. Construct the AO first-order density matrices with the correct restricted/unrestricted occupation factors and overlap/orthonormality contribution.
4. Form the response contribution

```text
J_response = (partial q/partial P) : dP0/dR.
```

5. Add the existing explicit projector derivative:

```text
J_relaxed = J_explicit + J_response.
```

6. Return the same shape currently consumed by the force loss:

```text
(nframe, natom_raw, 3, natom_descriptor, ndescriptor).
```

The implementation should reuse PySCF response batching and expose:

- `atmlst` for atom subsets;
- `max_memory`/coordinate block size;
- `conv_tol_cpscf`;
- `max_cycle_cpscf`;
- optional response level shift;
- response residual and convergence diagnostics.

RKS/UKS must use the matching CPKS response kernel, including the XC kernel supplied by PySCF. The correction descriptor itself has no grid term, but the reference density response must be the response of the selected KS reference.

### 4. Z-vector force backend

Implement the efficient scalar-energy gradient through an adjoint response solve:

1. Evaluate `E_corr` and the AO correction potential `V_corr = dE_corr/dP` using PyTorch autograd.
2. Transform `V_corr` to the occupied-virtual orbital-gradient right-hand side using PySCF's RHF/UHF/RKS/UKS conventions.
3. Solve the transpose coupled-perturbed equation for the Z-vector. Do not assume symmetry without verifying the metric and spin convention used by the selected PySCF solver.
4. Contract the Z-vector with the same nuclear perturbation right-hand sides used by the direct-response implementation.
5. Add:
- the native reference analytic gradient;
- the explicit DeePHF projector/AO-overlap derivative;
- the orbital-response contribution.

The direct CP-HF/CP-KS force should remain available as `response_method="direct"` for debugging and as a fallback. `response_method="zvector"` should become the default only after the direct/Z-vector equivalence tests pass for every supported reference type.

The response operator and nuclear right-hand-side construction must be shared between the two backends so sign, occupation, and overlap conventions cannot drift independently.

### 5. Define unambiguous data fields

Do not silently change the meaning of existing datasets.

Keep the current field as a legacy/unrelaxed quantity:

```text
grad_vx # existing behavior; explicit/fixed-density derivative
```

Add:

```text
grad_vx_explicit # explicit alias with unambiguous name
grad_vx_response # optional diagnostic response contribution
grad_vx_deephf # complete relaxed DeePHF descriptor Jacobian
f_corr # predicted DeePHF correction force
f_tot # reference force + DeePHF correction force
```

Store data provenance in a metadata file, for example:

```yaml
descriptor_gradient:
semantics: deephf_relaxed
response: cphf
reference: RHF
version: 1
```

Training in strict DeePHF mode must fail with a clear message when only legacy `grad_vx` is present. There must be no silent fallback from `grad_vx_deephf` to the unrelaxed derivative.

### 6. Integrate exact force training

The existing training expression can be retained once it consumes the complete relaxed Jacobian:

```python
eig.requires_grad_(True)
e_corr = model(eig)
g_desc = torch.autograd.grad(
e_corr, eig,
grad_outputs=torch.ones_like(e_corr),
create_graph=True,
)[0]
f_corr = -torch.einsum("...bxap,...ap->...bx", grad_vx_deephf, g_desc)
```

This remains differentiable with respect to all model parameters without differentiating through PySCF because `grad_vx_deephf` depends on the reference method, geometry, and basis, but not on the DeePHF model parameters.

Add an explicit training mode, for example:

```yaml
data_args:
force_mode: deephf_exact
gvx_name: grad_vx_deephf
train_args:
energy_factor: 1.0
force_factor: 1.0
```

Required training changes:

- centralize force prediction in a reusable helper used by train and test;
- validate Jacobian shape, units, semantics, and metadata;
- report energy and force metrics separately instead of only the combined loss;
- support force-only validation metrics without requiring reference forces during pure inference;
- retain `force_mode: deepks`/legacy behavior for existing iterative DeePKS workflows;
- document that smooth twice-differentiable activations are required for force training because optimizing a force loss takes derivatives of `dE/dq` with respect to model parameters.

### 7. CLI and workflow integration

Add an explicit calculation mode rather than inferring perturbative versus self-consistent behavior from the presence of a model file:

```yaml
mode: deephf # base | deephf | deepks
model_file: model.pth
response_args:
method: zvector # zvector | direct
conv_tol: 1.0e-10
max_cycle: 50
dump_fields:
- e_base
- e_corr
- e_tot
- f_base
- f_corr
- f_tot
```

For training-data generation, allow:

```yaml
mode: base
dump_fields:
- e_base
- f_base
- dm_eig
- grad_vx_deephf
- l_e_delta
- l_f_delta
```

The default behavior of existing inputs must remain unchanged. A deprecation path can be introduced later, but this feature must not reinterpret `deepks scf -m model.pth` silently.

Update `deepks test` or add a dedicated DeePHF prediction command so energy-only and energy-plus-force inference are both available from saved data and directly from molecular geometries.

## Implementation phases

### Phase 0: establish a reproducible baseline

- Add a supported Python/PyTorch/PySCF version matrix and pin at least one CI combination.
- Add a minimal test framework; the current repository has no active force regression suite.
- Implement a central-difference reference evaluator for the complete DeePHF energy, used only by tests.
- Record SCF and CPHF convergence thresholds in test output.

### Phase 1: descriptor refactor and RHF direct response

- Extract reusable descriptor operations without changing DeePKS results.
- Implement RHF first-order density response using PySCF Hessian/CPHF helpers.
- Add `grad_vx_explicit`, `grad_vx_response`, and `grad_vx_deephf`.
- Validate `grad_vx_deephf` and its contraction with several smooth test models against energy finite differences.

### Phase 2: DeePHF method and RHF Z-vector force

- Add the non-self-consistent `DeePHF` method object.
- Implement reference, correction, and total energy/force fields.
- Implement direct and Z-vector RHF force backends.
- Add scanner and geometry-optimization integration.

### Phase 3: unrestricted and KS references

- Add UHF/UCPHF with independent alpha/beta occupations and response blocks.
- Add RKS/CPKS and UKS/CPKS through the corresponding PySCF response/Hessian adapters.
- Test open-shell systems and non-HF XC functionals.

### Phase 4: force training and evaluation

- Add strict data semantics and metadata validation.
- Refactor the force predictor shared by training and testing.
- Add exact DeePHF force loss, force metrics, and restart coverage.
- Add a small end-to-end energy-and-force training example.

### Phase 5: performance and documentation

- Coordinate batching and memory controls for direct response.
- Cache geometry-local response intermediates safely.
- Benchmark direct response versus Z-vector inference.
- Document formulas, supported references, units, convergence controls, and known limitations.

## Validation plan

### Unit tests

- `dE_corr/dP` from PyTorch versus finite differences of the AO density matrix.
- Explicit descriptor derivative versus finite differences at fixed AO density.
- First-order reference density from CP-HF/CP-KS versus finite differences of converged reference densities.
- Relaxed descriptor Jacobian versus finite differences of descriptors after rerunning the reference SCF at displaced geometries.
- Direct-response correction force versus contraction of the stored `grad_vx_deephf`.
- Z-vector force versus direct-response force.
- Force-loss parameter gradients versus finite differences of selected network parameters.

### Integration systems

- RHF: H2 and a bent H2O geometry.
- UHF: OH or another small, well-behaved open-shell molecule.
- RKS and UKS: one small system with a simple semilocal functional.
- Ghost/projector-center regression using the current ghost-atom conventions.
- A small water example using the repository's existing model/data format.

Use non-symmetric, non-equilibrium geometries to avoid accidentally passing because of zero forces.

### Numerical acceptance criteria

- Total analytic DeePHF gradient versus central finite differences of `E_ref + E_corr`:
- maximum absolute error <= `1e-5 Eh/Bohr` for well-converged small test systems;
- demonstrate the expected finite-difference plateau with at least two displacement sizes.
- Direct CP-HF/CP-KS and Z-vector forces agree to <= `1e-8 Eh/Bohr` under tight response convergence.
- Zero correction model reproduces the native PySCF reference gradient to numerical precision.
- Translational invariance: the total force sum is near zero for isolated molecules.
- Existing DeePKS energy/force examples remain unchanged within their current numerical tolerance.
- A force-trained DeePHF model can reload from checkpoint and reproduce identical energy and force predictions.

## Numerical and compatibility considerations

- Use double precision for descriptors, response intermediates, and force tests.
- Keep gradient (`dE/dR`) and force (`-dE/dR`) signs explicit at API boundaries.
- Preserve the repository's Bohr/Angstrom conversion behavior and test both units.
- Require a converged reference SCF before solving response equations.
- Check and report the CPHF/CPKS residual; do not return unconverged forces as valid data.
- Support response level shifting as an opt-in stabilization tool and record it in metadata.
- Handle atom subsets and ghost atoms consistently in explicit and response contributions.
- Invalidate all projector and response caches when scanner geometry changes.
- Descriptor eigenvalue degeneracies can make individual sorted-eigenvalue derivatives non-unique. Tests should initially use nondegenerate descriptors; the implementation should detect near-degenerate projected-density eigenvalues and warn or use a documented subspace-safe treatment before claiming support at exact degeneracy.
- Avoid relying on unstable private PySCF APIs without a compatibility adapter and version tests. Issue #82 and issue #89 already identify PySCF compatibility as a repository risk.

## Definition of done

- [ ] A dedicated non-self-consistent DeePHF Python API exists.
- [ ] RHF, UHF, RKS, and UKS exact analytic DeePHF forces are implemented for molecular PySCF calculations.
- [ ] Direct response and Z-vector backends agree and pass finite-difference validation.
- [ ] `grad_vx_deephf` contains the complete relaxed descriptor derivative and has explicit metadata.
- [ ] Exact DeePHF force training, validation, checkpoint restart, and inference are implemented.
- [ ] Existing `grad_vx` datasets and DeePKS workflows remain backward compatible and cannot be silently mistaken for exact DeePHF data.
- [ ] Geometry scanner/optimization works with the new method.
- [ ] CI includes analytic-force and force-training regression tests.
- [ ] User and developer documentation explain the theory, configuration, data fields, units, convergence, and limitations.

## Estimated effort

- RHF direct-response MVP with finite-difference validation: approximately 1-2 developer weeks.
- Complete RHF/UHF/RKS/UKS implementation, Z-vector optimization, training integration, compatibility work, tests, and documentation: approximately 3-5 developer weeks for a contributor familiar with PySCF response theory.

The work should be split into reviewable pull requests following the phases above; the direct-response implementation and numerical reference tests should land before the Z-vector optimization.

---

Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh