Skip to content

Feature: Support nspin=4 PW partial-charge and wavefunction .cube output - #7841

Open
AsTonyshment wants to merge 2 commits into
deepmodeling:developfrom
AsTonyshment:fix/pchg_wfc
Open

Feature: Support nspin=4 PW partial-charge and wavefunction .cube output#7841
AsTonyshment wants to merge 2 commits into
deepmodeling:developfrom
AsTonyshment:fix/pchg_wfc

Conversation

@AsTonyshment

Copy link
Copy Markdown
Collaborator

What's changed?

This PR implements complete PW nspin=4 real-space state output through out_pchg, out_wfc_norm, and out_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.

  • Implement the four physical nspin=4 partial-charge components: charge $\rho_0$ and magnetization $m_x$, $m_y$, and $m_z$.
  • Output the total spinor norm through out_wfc_norm, and output the real and imaginary parts of both $\psi_{\uparrow}$ and $\psi_{\downarrow}$ through out_wfc_re_im.
  • Apply scalar symmetry restoration to $\rho_0$ and the coupled SOC vector transformation to $(m_x,m_y,m_z)$, reusing the same convention and implementation as the total charge-density path.
  • Use the actual local smooth-grid and dense-grid sizes instead of mixing global nxyz, local nrxx, 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 .cube output.
  • Replace three manually allocated PW inverse-FFT buffers with local ABACUS Tensor objects. Their CPU/GPU memory now follows RAII ownership and is released when the output routine returns.
  • Simplify the output interfaces by deriving band counts and buffer sizes from Psi and the relevant PW basis objects, and by removing redundant ctx, kpar, my_pool, grid-size, and Charge arguments.
  • Replace the old standalone .cube integrator with cube_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

$\rho_0 = |\psi_{\uparrow}|^2 + |\psi_{\downarrow}|^2$,

$m_x = 2\mathrm{Re}(\psi_{\uparrow}^{*}\psi_{\downarrow})$,

$m_y = 2\mathrm{Im}(\psi_{\uparrow}^{*}\psi_{\downarrow})$,

$m_z = |\psi_{\uparrow}|^2 - |\psi_{\downarrow}|^2$.

For nspin = 4, out_pchg components s1 through s4 are respectively $\rho_0$, $m_x$, $m_y$, and $m_z$. out_wfc_re_im components s1 and s2 are respectively $\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.

Feature validation: pointwise nspin=4 spinor identities

image

This figure uses the GPU scf_out_wf_spinor case on an NVIDIA GeForce RTX 2080 Ti with two MPI ranks, OMP_NUM_THREADS=1, KPAR=2, nspin=4, symmetry=0, a $1\times1\times2$ Gamma-centered k-point mesh, pw_seed=1, ecutwfc=20 Ry, and ecutrho=160 Ry. The ratio ecutrho / ecutwfc = 8 forces 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 .cube files for $\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.

The 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 9.875e-13 for the norm squared, 9.213e-13 for $\rho_0$, 3.890e-13 for $m_x$, 7.925e-13 for $m_y$, and 9.452e-13 for $m_z$, all far below the 1e-8 acceptance threshold.

Feature validation: nspin = 4 scalar/SOC symmetry and KPAR

image

This figure uses a controlled single-atom simple-cubic Fe operator probe with the fully relativistic Fe.upf, noncolin=1, lspinorb=1, nspin=4, a three-dimensional initial moment along the cubic [111] direction, a $2\times2\times2$ Gamma-centered mesh, band 1, ecutwfc=20 Ry, and ecutrho=160 Ry. Four independent calculations cover the unreduced full BZ with symmetry=0, KPAR=1, and four MPI ranks (8 k-points); the S_6 IBZ with symmetry=1, KPAR=1, and four MPI ranks (4 k-points); the same IBZ with KPAR=2 and four MPI ranks (equal 2+2 process pools); and the same IBZ with KPAR=2 and five MPI ranks (uneven 3+2 process pools, exercising the fallback path without KP_WORLD). All runs use OMP_NUM_THREADS=1 and 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 S_6 magnetic group directly from cubic signed-permutation matrices. It transforms $\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 1.000e-15; the worst residual among the coupled vector components is 1.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-9 for the equal-pool layout and 1.467e-9 for the uneven-pool fallback, and every integral error is below 7.723e-10. All symmetry, pointwise, and integral checks pass the 1e-7 acceptance 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_Grid now owns the common cross-pool real-space reduction operation. Charge::reduce_diff_pools delegates to this method, and the old Charge::init_chgmpi receive-count/displacement arrays are removed.
  • Symmetry_rho adds a raw-array SOC overload so state-output buffers can reuse the established coupled magnetization transformation without pretending to be a Charge object.

@AsTonyshment
AsTonyshment requested review from mohanchen and a lite review from Copilot August 20, 2026 12:17
@AsTonyshment
AsTonyshment removed the request for review from Copilot August 20, 2026 12:19
@AsTonyshment
AsTonyshment requested a lite review from Copilot and removed request for Copilot August 20, 2026 15:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant