DeePKS: support collinear spin (nspin=2)#7433
Open
ErjieWu wants to merge 6 commits into
Open
Conversation
Add charge/magnetization data members for collinear nspin=2 DeePKS without changing the nspin=1 path: * dm_r_mag : (rho_up - rho_dn) real-space density matrix * pdm_mag : magnetization projected density matrix * gedm_mag : dE/d(pdm_mag) Allocated and freed only when nspin==2; for nspin==1 they stay null/empty and the existing dm_r/pdm/gedm behave exactly as before. Scaffolding only, no physics yet. tests/09_DeePKS: 284/284. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… model
Traditional (non-equivariant) nspin=2 collinear DeePKS now computes both the
charge and magnetization channels and evaluates a 2-channel model:
* update_dmr gains nspin/mag options; setup_deepks fills dm_r_mag = rho_up - rho_dn
* the operator builds pdm_mag and descriptor_mag, stacks charge+magnetization
descriptors into model input (1, nat, 2, des), and autograds in one pass to
fill gedm = dE/d(pdm) and gedm_mag = dE/d(pdm_mag)
The equivariant version stays a separate top-level branch; all magnetization
allocation/use is gated on nspin==2 && !deepks_equiv. The correction is not yet
applied per spin to the Hamiltonian (next step), so nspin=1 is unchanged.
tests/09_DeePKS: 284/284.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
SCF correction: calculate_HR takes the gedm to assemble from; for nspin=2 (traditional) the operator rebuilds V_delta_R per spin as |alpha>(gedm +/- gedm_mag)<alpha| and toggles current_spin like Veff. The label/output interface now feeds the 2-channel model for nspin=2 (previously crashed a real 2-channel model post-SCF). Forces/stress: the single live path is getForceStress -> cal_f_delta (integral_part/ftable is dead code). cal_f_delta gains an optional magnetization channel (dmr_mag, gedm_mag) and adds its contribution in the same pass, so F_delta = dm_r * nlm(gedm) + dm_r_mag * nlm(gedm_mag) with correct single symmetrize+weight on the stress. Validation (synthetic 2-channel model): * non-magnetic CH4: nspin=2 == nspin=1 total energy to 2.5e-9 eV, forces to ~1e-5 * polarized CH4 (nupdown=2): finite-difference force matches analytic to 0.05% nspin=1 tests/09_DeePKS: 284/284 throughout. nspin=1 path unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
For nspin=2 (traditional), the interface now builds the magnetization-channel descriptor before saving labels and reuses it for the 2-channel model. save_npy_d gains an optional descriptor_mag: when present it writes dm_eig as (nat, 2, des) - channel 0 = charge, channel 1 = magnetization - matching the 2-channel model input, so deepks-kit can train on nspin=2 descriptors. nspin=1 writes (nat, des) as before. Validated: nspin=2 deepks_out_labels=1 on CH4 produces deepks_dm_eig.npy of shape (5, 2, 18). nspin=1 tests/09_DeePKS: 284/284. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
For nspin=2 (traditional), every precalc/output label now carries the
charge/magnetization channel (channel 0 = charge, 1 = magnetization), inserted
before the descriptor dimension to match dm_eig (nat, 2, des):
* gevdm computed for both pdm and pdm_mag
* gradvx (force) and gvepsl (stress): gdmx/gdmepsl built from dm_r_mag, then
gvx/gvepsl per channel, stacked
* orbpre (bandgap): orbital_precalc per channel, stacked (cal_o_delta already
sums spin)
* vdpre / vdrpre (v_delta precalc) and gevdm export: per-channel + stacked
The precalc kernels are unchanged (channel-agnostic); only the interface calls
them per channel and stacks. deepks_spin2 hoisted to function scope.
Validated (synthetic 2-channel model, nupdown=2): label shapes gain the size-2
channel dim (dm_eig (5,2,18), gradvx (5,3,5,2,18), gvepsl (6,5,2,18),
orbpre (2,5,2,18), vdpre (2,33,33,5,2,18)); nspin=1 shapes unchanged.
tests/09_DeePKS (nspin=1): 284/284.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add spin-polarized DeePKS integration tests under tests/09_DeePKS, using a small synthetic 2-channel (charge + magnetization) demo model committed as Model_ProjOrb/model_nspin2_demo.ptg: * 29_NO_GO_deepks_scf_nspin2 : gamma SCF + force + stress (net moment) * 30_NO_KP_deepks_scf_nspin2 : multi-k SCF (net moment) * 31_NO_GO_deepks_bandgap_nspin2 : output labels (all precalc) + bandgap All use nupdown to give a real moment so the magnetization channel is active. tests/09_DeePKS now passes 313/313 (28 original + 3 new) under the standard np=4 harness. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
This PR adds collinear spin-polarized (nspin=2) support to the traditional DeePKS LCAO path by introducing a charge (n = ρ↑ + ρ↓) and magnetization (m = ρ↑ - ρ↓) channel decomposition, while leaving the equivariant DeePKS path unchanged. It extends DeePKS SCF correction, force/stress contributions, and output-label/precalc plumbing to operate on a 2-channel descriptor layout and adds integration tests covering the new behavior.
Changes:
- Add magnetization-channel DeePKS data structures (
dm_r_mag,pdm_mag,gedm_mag) and compute a 2-channel descriptor fornspin=2(traditional). - Apply spin-dependent Hamiltonian corrections by combining
gedm ± gedm_mag, and include magnetization-channel contributions in the force/stress path. - Extend DeePKS output labels/precalc tensors to carry a size-2 channel axis (charge/magnetization) and add new
nspin=2regression tests.
Reviewed changes
Copilot reviewed 31 out of 32 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/09_DeePKS/CASES_CPU.txt | Registers new nspin=2 DeePKS CPU integration cases. |
| tests/09_DeePKS/29_NO_GO_deepks_scf_nspin2/STRU | Adds structure input for gamma-only nspin=2 SCF test. |
| tests/09_DeePKS/29_NO_GO_deepks_scf_nspin2/KPT | Adds K-points file for gamma-only nspin=2 SCF test. |
| tests/09_DeePKS/29_NO_GO_deepks_scf_nspin2/INPUT | Adds DeePKS nspin=2 SCF + force + stress inputs (gamma-only). |
| tests/09_DeePKS/29_NO_GO_deepks_scf_nspin2/README | Documents the new test case intent. |
| tests/09_DeePKS/29_NO_GO_deepks_scf_nspin2/result.ref | Adds reference outputs for the new nspin=2 SCF test. |
| tests/09_DeePKS/30_NO_KP_deepks_scf_nspin2/STRU | Adds structure input for multi-k nspin=2 SCF test. |
| tests/09_DeePKS/30_NO_KP_deepks_scf_nspin2/KPT | Adds multi-k K-points file for nspin=2 SCF test. |
| tests/09_DeePKS/30_NO_KP_deepks_scf_nspin2/INPUT | Adds DeePKS nspin=2 SCF inputs for multi-k system. |
| tests/09_DeePKS/30_NO_KP_deepks_scf_nspin2/README | Documents the new multi-k nspin=2 test. |
| tests/09_DeePKS/30_NO_KP_deepks_scf_nspin2/result.ref | Adds reference outputs for the new multi-k test. |
| tests/09_DeePKS/31_NO_GO_deepks_bandgap_nspin2/STRU | Adds structure input for nspin=2 bandgap + label outputs test. |
| tests/09_DeePKS/31_NO_GO_deepks_bandgap_nspin2/KPT | Adds K-points file for gamma-only nspin=2 bandgap test. |
| tests/09_DeePKS/31_NO_GO_deepks_bandgap_nspin2/INPUT | Adds DeePKS nspin=2 bandgap + output-label inputs. |
| tests/09_DeePKS/31_NO_GO_deepks_bandgap_nspin2/README | Documents the nspin=2 bandgap/labels case. |
| tests/09_DeePKS/31_NO_GO_deepks_bandgap_nspin2/result.ref | Adds reference outputs for nspin=2 labels/bandgap. |
| source/source_lcao/setup_deepks.cpp | Updates DeePKS setup to build the magnetization-channel real-space DM when nspin=2. |
| source/source_lcao/module_operator_lcao/deepks_lcao.h | Updates DeePKS HR builder API to accept a gedm pointer for spin-dependent builds. |
| source/source_lcao/module_operator_lcao/deepks_lcao.cpp | Implements 2-channel (charge/mag) SCF correction and per-spin HR rebuild logic. |
| source/source_lcao/module_deepks/LCAO_deepks.h | Adds magnetization-channel DeePKS members (dm_r_mag, pdm_mag, gedm_mag). |
| source/source_lcao/module_deepks/LCAO_deepks.cpp | Allocates and initializes magnetization-channel PDM/DMR/gedm storage for nspin=2. |
| source/source_lcao/module_deepks/LCAO_deepks_io.h | Extends descriptor .npy writer interface to optionally include magnetization-channel descriptor. |
| source/source_lcao/module_deepks/LCAO_deepks_io.cpp | Writes dm_eig as (nat, 2, des) when magnetization-channel descriptor is provided. |
| source/source_lcao/module_deepks/LCAO_deepks_interface.cpp | Plumbs 2-channel labels/precalc outputs and bandgap-related tensors for nspin=2. |
| source/source_lcao/module_deepks/deepks_pdm.h | Extends update_dmr API with nspin/magnetization flags. |
| source/source_lcao/module_deepks/deepks_pdm.cpp | Implements magnetization-channel DMR via sign change for spin-down DMK blocks. |
| source/source_lcao/module_deepks/deepks_force.h | Extends DeePKS force/stress kernel API with optional magnetization-channel inputs. |
| source/source_lcao/module_deepks/deepks_force.cpp | Adds magnetization-channel contributions into force/stress accumulation. |
| source/source_lcao/module_deepks/deepks_basic.h | Extends cal_edelta_gedm API to support optional magnetization-channel descriptor/PDM and output gradients. |
| source/source_lcao/module_deepks/deepks_basic.cpp | Builds 2-channel model input and extracts gradients for both charge and magnetization channels. |
| source/source_lcao/FORCE_STRESS.cpp | Passes magnetization-channel DM/gedm into DeePKS force/stress evaluation for nspin=2. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
272
to
276
| std::vector<torch::Tensor> gedm_tensor = torch::autograd::grad(ec, | ||
| pdm, | ||
| grad_inputs, | ||
| gedm_shell, | ||
| /*retain_grad=*/true, | ||
| /*create_graph=*/false, |
| for (int inl = 0; inl < deepks_param.inlmax; ++inl) | ||
| { | ||
| int nm = 2 * deepks_param.inl2l[inl] + 1; | ||
| auto accessor = gedm_tensor[deepks_param.inlmax + inl].accessor<double, 2>(); |
Comment on lines
+554
to
+559
| if (deepks_spin2) | ||
| { | ||
| torch::Tensor gevdm_out_mag; | ||
| DeePKS_domain::prepare_gevdm(nat, deepks_param, orb, gevdm_mag, gevdm_out_mag); | ||
| gevdm_out = torch::stack({gevdm_out, gevdm_out_mag}, gevdm_out.dim() - 1); | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds collinear spin-polarized (
nspin=2) support to the traditional DeePKSmodule, using a charge / magnetization channel decomposition. The equivariant
version is intentionally left unchanged and stays a separate code path;
nspin=4(non-collinear) is future work. The
nspin=1path is unchanged (every spinbranch is gated on
nspin==2 && !deepks_equiv).What changed
dm_r_mag(rho_up - rho_dn),pdm_mag,gedm_mag, allocated only fornspin==2.both the charge channel
n = rho_up + rho_dnand the magnetization channelm = rho_up - rho_dn; the model is fed a 2-channel descriptor(1, nat, 2, des);one autograd pass fills
gedm = dE/d(pdm^n)andgedm_mag = dE/d(pdm^m). TheHamiltonian correction is applied per spin as
|alpha>(gedm +/- gedm_mag)<alpha|,toggling the operator's
current_spinlikeVeff.getForceStress -> cal_f_deltapath gains themagnetization contribution in a single pass (
F = dm_r*nlm(gedm) + dm_r_mag*nlm(gedm_mag)),preserving the single stress symmetrize+scale.
dm_eig,gradvx,gvepsl,orbpre,vdpre/vdrpre,gevdm) carries a size-2 channel axis before thedescriptor axis (channel 0 = charge, 1 = magnetization), matching the 2-channel
model input; band-gap
o_delta/orbpreare spin-resolved. Precalc kernels areunchanged (channel-agnostic); the interface calls them per channel and stacks
(on rank 0, where the rank-0-assembled tensors are valid).
Validation
tests/09_DeePKS: 313/313 pass under the standardnp=4harness (28 originalnspin=2cases).nspin=2 == nspin=1total energy to 2.5e-9 eV (the m->0reduction) and forces to ~1e-5.
nupdown=2): finite-difference force matches the analytic DeePKSforce to 0.05% (the magnetization force term is correct).
New tests (
tests/09_DeePKS, 29-31)Use a small synthetic 2-channel demo model committed at
Model_ProjOrb/model_nspin2_demo.ptg:29_NO_GO_deepks_scf_nspin2- gamma SCF + force + stress (net moment)30_NO_KP_deepks_scf_nspin2- multi-k SCF (net moment)31_NO_GO_deepks_bandgap_nspin2- output labels (all precalc) + band gapNotes
The
nspin=2descriptor / precalc label layout places the channel axis (size 2)immediately before the descriptor axis, channel 0 = charge, 1 = magnetization, so
deepks-kit can train a 2-channel model. A trained physical
nspin=2model is notincluded; the committed demo model is a small synthetic fixture for the
integration tests only.
This work was developed with AI assistance (commits carry a
Co-Authored-Bytrailer).