Feature: Support nspin=4 PW partial-charge and wavefunction .cube output - #7841
Open
AsTonyshment wants to merge 2 commits into
Open
Feature: Support nspin=4 PW partial-charge and wavefunction .cube output#7841AsTonyshment wants to merge 2 commits into
.cube output#7841AsTonyshment wants to merge 2 commits into
Conversation
AsTonyshment
force-pushed
the
fix/pchg_wfc
branch
from
August 20, 2026 15:36
c2dbf62 to
2ca3d6c
Compare
AsTonyshment
requested
a lite review from Copilot
and removed request for
Copilot
August 20, 2026 15:36
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.
What's changed?
This PR implements complete PW
nspin=4real-space state output throughout_pchg,out_wfc_norm, andout_wfc_re_im. The new feature works on CPU and GPU, supports smooth and double FFT grids, preserves separate-k and k-summed output modes, and handles both symmetry reduction and equal/uneven KPAR layouts.nspin=4partial-charge components: chargeout_wfc_norm, and output the real and imaginary parts of bothout_wfc_re_im.nxyz, localnrxx, and incompatible FFT layouts. When double-grid mode is active, transfer the smooth-grid spinor wavefunctions to the dense charge grid by reciprocal-space zero padding before.cubeoutput.Tensorobjects. Their CPU/GPU memory now follows RAII ownership and is released when the output routine returns.Psiand the relevant PW basis objects, and by removing redundantctx,kpar,my_pool, grid-size, andChargearguments..cubeintegrator withcube_tool.py, which provides scalar integration, phase-aligned complex-wavefunction comparison, compact U(1)-invariant wavefunction fingerprints, and pointwise spinor-identity checks.The PW spinor fields use
For$\rho_0$ , $m_x$ , $m_y$ , and $m_z$ . $\psi_{\uparrow}$ and $\psi_{\downarrow}$ . As for the existing collinear path, these files represent complete one-particle states and do not include the SCF occupation factor.
nspin = 4,out_pchgcomponentss1throughs4are respectivelyout_wfc_re_imcomponentss1ands2are respectivelyFeature validation: pointwise nspin=4 spinor identities
This figure uses the GPU$1\times1\times2$ Gamma-centered k-point mesh,
scf_out_wf_spinorcase on an NVIDIA GeForce RTX 2080 Ti with two MPI ranks,OMP_NUM_THREADS=1,KPAR=2,nspin=4,symmetry=0, apw_seed=1,ecutwfc=20 Ry, andecutrho=160 Ry. The ratioecutrho / ecutwfc = 8forces the wavefunctions to be output on a 25 × 25 × 25 dense grid rather than their smooth FFT grid. Band 1 and k-point 1 are shown.The five columns verify, from left to right, the total wavefunction norm squared,$\rho_0$ , $m_x$ , $m_y$ , and $m_z$ . In the upper row, the translucent filled surface is the value written directly by ABACUS. The black mesh is reconstructed independently from the four Re/Im $\psi_{\uparrow}$ and $\psi_{\downarrow}$ . Nonnegative quantities use the isovalue $0.2\max(f)$ ; signed quantities use every existing surface at $\pm0.2\max(|f|)$ . Exact overlap of the filled and mesh surfaces is the visual pointwise check.
.cubefiles forThe lower row shows$\log_{10}(|f_{\mathrm{ABACUS}}-f_{\mathrm{reconstructed}}|+10^{-12})$ on the z slice containing each identity's global maximum error, with a shared range from -12 to -8. The cyan circle marks that maximum-error grid point. The maximum absolute errors are $\rho_0$ , $m_x$ , $m_y$ , and $m_z$ , all far below the
9.875e-13for the norm squared,9.213e-13for3.890e-13for7.925e-13for9.452e-13for1e-8acceptance threshold.Feature validation:
nspin = 4scalar/SOC symmetry and KPARThis figure uses a controlled single-atom simple-cubic Fe operator probe with the fully relativistic$2\times2\times2$ Gamma-centered mesh, band 1,
Fe.upf,noncolin=1,lspinorb=1,nspin=4, a three-dimensional initial moment along the cubic [111] direction, aecutwfc=20 Ry, andecutrho=160 Ry. Four independent calculations cover the unreduced full BZ withsymmetry=0,KPAR=1, and four MPI ranks (8 k-points); theS_6IBZ withsymmetry=1,KPAR=1, and four MPI ranks (4 k-points); the same IBZ withKPAR=2and four MPI ranks (equal 2+2 process pools); and the same IBZ withKPAR=2and five MPI ranks (uneven 3+2 process pools, exercising the fallback path withoutKP_WORLD). All runs useOMP_NUM_THREADS=1and deliberately stop after one SCF step, thereby holding the operator-level input fixed while testing symmetry restoration and parallel output assembly. This is not presented as a converged Fe ground-state calculation.The first row shows the four IBZ partial-charge components: the scalar charge$\rho_0$ and the three nontrivial magnetization components $m_x$ , $m_y$ , and $m_z$ . Blue and red isosurfaces denote positive and negative values for signed components. Each title gives the common relative isovalue and the field integral.
The second row checks the output against an independent symmetry oracle. The script enumerates all 12 unitary and antiunitary operations of the$\rho_0$ as a scalar and transforms $\boldsymbol{m}=(m_x,m_y,m_z)$ as one coupled axial SOC vector, including the time-reversal sign and the ABACUS Pauli-$y$ handedness convention. The panels show $\log_{10}(|\Delta f|+10^{-16})$ on the z slice containing the worst residual over all 12 operations. The worst scalar residual is
S_6magnetic group directly from cubic signed-permutation matrices. It transforms1.000e-15; the worst residual among the coupled vector components is1.000e-16.The third and fourth rows compare all four components from the equal 2+2 and uneven 3+2 KPAR layouts with the single-pool IBZ result. They show$\log_{10}(|\Delta f|+10^{-12})$ on each component's worst-residual z slice; cyan circles mark the corresponding grid points. The largest pointwise errors are
1.278e-9for the equal-pool layout and1.467e-9for the uneven-pool fallback, and every integral error is below7.723e-10. All symmetry, pointwise, and integral checks pass the1e-7acceptance threshold.The full-BZ run supplies the expected unreduced 8-point reference branch and the IBZ runs confirm reduction to 4 points. The SOC vector field is validated through its full coupled transformation law rather than by incorrectly requiring each Cartesian magnetization component to be separately invariant under full-BZ-to-IBZ reduction.
Core-module impact
Parallel_Gridnow owns the common cross-pool real-space reduction operation.Charge::reduce_diff_poolsdelegates to this method, and the oldCharge::init_chgmpireceive-count/displacement arrays are removed.Symmetry_rhoadds a raw-array SOC overload so state-output buffers can reuse the established coupled magnetization transformation without pretending to be aChargeobject.