Feature: Support nspin=4 PW partial-charge and wavefunction .cube output (Useful Information to know how to use out_pchg, out_wfc_norm, and out_wfc_re_im functions)

docs/advanced/input_files/input-main.md

@@ -2374,28 +2374,30 @@
 
 - **Type**: String
 - **Availability**: *[`basis_type`](#basis_type)==pw or ([`basis_type`](#basis_type)==lcao and [`calculation`](#calculation)==get_pchg)*
-- **Description**: Specifies the electronic states to calculate the charge densities with state index for, using a space-separated string of 0s and 1s. Each digit in the string corresponds to a state, starting from the first state. A 1 indicates that the charge density should be calculated for that state, while a 0 means the state will be ignored. The parameter allows a compact and flexible notation (similar to ocp_set), for example the syntax 1 4*0 5*1 0 is used to denote the selection of states: 1 means calculate for the first state, 4*0 skips the next four states, 5*1 means calculate for the following five states, and the final 0 skips the next state. It's essential that the total count of states does not exceed the total number of states (nbands); otherwise, it results in an error, and the process exits. The input string must contain only numbers and the asterisk (*) for repetition, ensuring correct format and intention of state selection. The outputs comprise multiple .cube files following the naming convention pchgi[state]s[spin]k[kpoint].cube.
+- **Description**: Selects electronic states for partial (band-decomposed) charge-density output using a space-separated string of `0`s and `1`s, where `1` selects a state and `0` skips it. Repetition follows the `ocp_set` syntax, for example `1 4*0 5*1 0`; the expanded list must not exceed `nbands`. Each output represents a complete one-particle state rather than its SCF occupation. The spin degeneracy is 2 for `nspin=1` and 1 for `nspin=2` or `nspin=4`. For `nspin=1`, `s1` contains the charge density. For `nspin=2`, `s1` and `s2` contain the spin-up and spin-down charge densities, respectively. For `nspin=4`, `s1`, `s2`, `s3`, and `s4` respectively contain $\rho_0$, $m_x$, $m_y$, and $m_z$. With `if_separate_k=true`, files are named `pchgi[state]s[component]k[kpoint].cube`; otherwise, the weighted k-point sum is named `pchgi[state]s[component].cube`.
+
+  > Note: Enabling symmetry may produce unintended partial charge densities because of reduced k-point weights and real-space symmetry operations. If the desired symmetry treatment is uncertain, set `symmetry = -1`. Use the same symmetry setting as in the SCF calculation.
 - **Default**: none
 
 ### out_wfc_norm
 
 - **Type**: String
 - **Availability**: *[`basis_type`](#basis_type)==pw or ([`basis_type`](#basis_type)==lcao and [`calculation`](#calculation)==get_wf)*
-- **Description**: Specifies the electronic states to calculate the real-space wave function modulus (norm, or known as the envelope function) with state index. The syntax and state selection rules are identical to out_pchg, but the output is the norm of the wave function. The outputs comprise multiple .cube files following the naming convention wfi[state]s[spin]k[kpoint].cube.
+- **Description**: Selects electronic states for real-space wavefunction-modulus output using the selection syntax and complete-state normalization of `out_pchg`. For `nspin=1`, `s1` contains the wavefunction modulus. For `nspin=2`, `s1` and `s2` contain the spin-up and spin-down wavefunction moduli, respectively. For `nspin=4`, `s1` contains the total spinor modulus. Files are named `wfi[state]s[spin]k[kpoint].cube`.
 - **Default**: none
 
 ### out_wfc_re_im
 
 - **Type**: String
 - **Availability**: *[`basis_type`](#basis_type)==pw or ([`basis_type`](#basis_type)==lcao and [`calculation`](#calculation)==get_wf)*
-- **Description**: Specifies the electronic states to calculate the real and imaginary parts of the wave function with state index. The syntax and state selection rules are identical to out_pchg, but the output contains both the real and imaginary components of the wave function. The outputs comprise multiple .cube files following the naming convention wfi[state]s[spin]k[kpoint][re/im].cube.
+- **Description**: Selects electronic states for real-space wavefunction real- and imaginary-part output using the selection syntax and complete-state normalization of `out_pchg`. For `nspin=1`, `s1` contains the wavefunction. For `nspin=2`, `s1` and `s2` contain the spin-up and spin-down wavefunctions, respectively. For `nspin=4`, `s1` and `s2` contain the upper and lower spinor components, respectively. Files are named `wfi[state]s[spin]k[kpoint][re/im].cube`.
 - **Default**: none
 
 ### if_separate_k
 
 - **Type**: Boolean
 - **Availability**: *([`basis_type`](#basis_type)==pw and [`out_pchg`](#out_pchg)!=none) or ([`basis_type`](#basis_type)==lcao and [`calculation`](#calculation)==get_pchg and [`gamma_only`](#gamma_only)==0)*
-- **Description**: Specifies whether to write the partial charge densities for all k-points to individual files or merge them. Warning: Enabling symmetry may produce unwanted results due to reduced k-point weights and symmetry operations in real space. Therefore when calculating partial charge densities, if you are not sure what you want exactly, it is strongly recommended to set symmetry = -1. It is noteworthy that your symmetry setting should remain the same as that in the SCF procedure.
+- **Description**: Specifies whether to write partial charge densities for individual k-points or merge them.
 - **Default**: false
 
 ### out_elf

docs/parameters.yaml

@@ -3624,31 +3624,33 @@ parameters:
     category: Output information
     type: String
     description: |
-      Specifies the electronic states to calculate the charge densities with state index for, using a space-separated string of 0s and 1s. Each digit in the string corresponds to a state, starting from the first state. A 1 indicates that the charge density should be calculated for that state, while a 0 means the state will be ignored. The parameter allows a compact and flexible notation (similar to ocp_set), for example the syntax 1 4*0 5*1 0 is used to denote the selection of states: 1 means calculate for the first state, 4*0 skips the next four states, 5*1 means calculate for the following five states, and the final 0 skips the next state. It's essential that the total count of states does not exceed the total number of states (nbands); otherwise, it results in an error, and the process exits. The input string must contain only numbers and the asterisk (*) for repetition, ensuring correct format and intention of state selection. The outputs comprise multiple .cube files following the naming convention pchgi[state]s[spin]k[kpoint].cube.
+      Selects electronic states for partial (band-decomposed) charge-density output using a space-separated string of `0`s and `1`s, where `1` selects a state and `0` skips it. Repetition follows the `ocp_set` syntax, for example `1 4*0 5*1 0`; the expanded list must not exceed `nbands`. Each output represents a complete one-particle state rather than its SCF occupation. The spin degeneracy is 2 for `nspin=1` and 1 for `nspin=2` or `nspin=4`. For `nspin=1`, `s1` contains the charge density. For `nspin=2`, `s1` and `s2` contain the spin-up and spin-down charge densities, respectively. For `nspin=4`, `s1`, `s2`, `s3`, and `s4` respectively contain $\rho_0$, $m_x$, $m_y$, and $m_z$. With `if_separate_k=true`, files are named `pchgi[state]s[component]k[kpoint].cube`; otherwise, the weighted k-point sum is named `pchgi[state]s[component].cube`.
+
+      [NOTE] Enabling symmetry may produce unintended partial charge densities because of reduced k-point weights and real-space symmetry operations. If the desired symmetry treatment is uncertain, set `symmetry = -1`. Use the same symmetry setting as in the SCF calculation.
     default_value: none
     unit: ""
     availability: basis_type==pw or (basis_type==lcao and calculation==get_pchg)
   - name: out_wfc_norm
     category: Output information
     type: String
     description: |
-      Specifies the electronic states to calculate the real-space wave function modulus (norm, or known as the envelope function) with state index. The syntax and state selection rules are identical to out_pchg, but the output is the norm of the wave function. The outputs comprise multiple .cube files following the naming convention wfi[state]s[spin]k[kpoint].cube.
+      Selects electronic states for real-space wavefunction-modulus output using the selection syntax and complete-state normalization of `out_pchg`. For `nspin=1`, `s1` contains the wavefunction modulus. For `nspin=2`, `s1` and `s2` contain the spin-up and spin-down wavefunction moduli, respectively. For `nspin=4`, `s1` contains the total spinor modulus. Files are named `wfi[state]s[spin]k[kpoint].cube`.
     default_value: none
     unit: ""
     availability: basis_type==pw or (basis_type==lcao and calculation==get_wf)
   - name: out_wfc_re_im
     category: Output information
     type: String
     description: |
-      Specifies the electronic states to calculate the real and imaginary parts of the wave function with state index. The syntax and state selection rules are identical to out_pchg, but the output contains both the real and imaginary components of the wave function. The outputs comprise multiple .cube files following the naming convention wfi[state]s[spin]k[kpoint][re/im].cube.
+      Selects electronic states for real-space wavefunction real- and imaginary-part output using the selection syntax and complete-state normalization of `out_pchg`. For `nspin=1`, `s1` contains the wavefunction. For `nspin=2`, `s1` and `s2` contain the spin-up and spin-down wavefunctions, respectively. For `nspin=4`, `s1` and `s2` contain the upper and lower spinor components, respectively. Files are named `wfi[state]s[spin]k[kpoint][re/im].cube`.
     default_value: none
     unit: ""
     availability: basis_type==pw or (basis_type==lcao and calculation==get_wf)
   - name: if_separate_k
     category: Output information
     type: Boolean
     description: |
-      Specifies whether to write the partial charge densities for all k-points to individual files or merge them. Warning: Enabling symmetry may produce unwanted results due to reduced k-point weights and symmetry operations in real space. Therefore when calculating partial charge densities, if you are not sure what you want exactly, it is strongly recommended to set symmetry = -1. It is noteworthy that your symmetry setting should remain the same as that in the SCF procedure.
+      Specifies whether to write partial charge densities for individual k-points or merge them.
     default_value: "false"
     unit: ""
     availability: "(basis_type==pw and out_pchg!=none) or (basis_type==lcao and calculation==get_pchg and gamma_only==0)"

source/source_base/parallel_grid.cpp

@@ -1,9 +1,11 @@
 #include "parallel_grid.h"
+
 #include "source_base/global_function.h"
 #include "source_base/global_variable.h"
 
 #ifdef __MPI
 #include "source_base/parallel_comm.h" // use POOL_WORLD
+
 #include <mpi.h>
 #endif
 
@@ -167,6 +169,95 @@ void Parallel_Grid::z_distribution()
     return;
 }
 
+void Parallel_Grid::reduce_across_pools(double* data) const
+{
+#ifdef __MPI
+    if (GlobalV::KPAR <= 1)
+    {
+        return;
+    }
+
+    assert(data != nullptr);
+    if (KP_WORLD != MPI_COMM_NULL)
+    {
+        // Equal-sized pools give corresponding ranks identical z-slab layouts,
+        // so their local buffers can be summed directly without redistribution.
+        MPI_Allreduce(MPI_IN_PLACE, data, this->nrxx, MPI_DOUBLE, MPI_SUM, KP_WORLD);
+        return;
+    }
+
+    // Uneven pool sizes have no KP_WORLD and may assign different z-slabs to
+    // corresponding ranks. Validate the local distribution before rebuilding
+    // a common global layout for the cross-pool reduction.
+    assert(!this->numz.empty());
+    assert(GlobalV::MY_POOL >= 0 && GlobalV::MY_POOL < static_cast<int>(this->numz.size()));
+    assert(GlobalV::RANK_IN_POOL >= 0 && GlobalV::RANK_IN_POOL < static_cast<int>(this->numz[GlobalV::MY_POOL].size()));
+    assert(this->nczp == this->numz[GlobalV::MY_POOL][GlobalV::RANK_IN_POOL]);
+    assert(this->nrxx == this->ncxy * this->nczp);
+
+    const int pool_size = this->nproc_in_pool[GlobalV::MY_POOL];
+    std::vector<int> receive_counts(pool_size);
+    std::vector<int> displacements(pool_size);
+    for (int ip = 0; ip < pool_size; ++ip)
+    {
+        receive_counts[ip] = this->numz[GlobalV::MY_POOL][ip] * this->ncxy;
+        displacements[ip] = this->startz[GlobalV::MY_POOL][ip] * this->ncxy;
+    }
+
+    std::vector<double> local_data(this->nrxx);
+    // The allgather below replicates one complete pool grid on every rank in
+    // that pool. INT_BGROUP then sums all of those replicas, so divide each
+    // local slab by the pool size to make each pool contribute exactly once.
+    const double pool_normalization = 1.0 / static_cast<double>(pool_size);
+    for (int ir = 0; ir < this->nrxx; ++ir)
+    {
+        local_data[ir] = data[ir] * pool_normalization;
+    }
+
+    std::vector<double> pool_data(this->ncxyz);
+    // Collect the rank-local [xy][local_z] slabs into rank-contiguous blocks.
+    MPI_Allgatherv(local_data.data(),
+                   this->nrxx,
+                   MPI_DOUBLE,
+                   pool_data.data(),
+                   receive_counts.data(),
+                   displacements.data(),
+                   MPI_DOUBLE,
+                   POOL_WORLD);
+
+    std::vector<double> global_layout(this->ncxyz);
+    // Convert the rank-contiguous allgather result to the canonical
+    // [xy][global_z] order required for element-wise reduction across pools.
+    for (int ip = 0; ip < pool_size; ++ip)
+    {
+        const int local_nz = this->numz[GlobalV::MY_POOL][ip];
+        const int global_z_start = this->startz[GlobalV::MY_POOL][ip];
+        const int gathered_start = global_z_start * this->ncxy;
+        for (int ixy = 0; ixy < this->ncxy; ++ixy)
+        {
+            for (int iz = 0; iz < local_nz; ++iz)
+            {
+                global_layout[ixy * this->ncz + global_z_start + iz] = pool_data[gathered_start + ixy * local_nz + iz];
+            }
+        }
+    }
+
+    MPI_Allreduce(MPI_IN_PLACE, global_layout.data(), this->ncxyz, MPI_DOUBLE, MPI_SUM, INT_BGROUP);
+
+    // Return only the z-slab owned by this rank under its pool's distribution.
+    const int local_z_start = this->startz[GlobalV::MY_POOL][GlobalV::RANK_IN_POOL];
+    for (int ixy = 0; ixy < this->ncxy; ++ixy)
+    {
+        for (int iz = 0; iz < this->nczp; ++iz)
+        {
+            data[ixy * this->nczp + iz] = global_layout[ixy * this->ncz + local_z_start + iz];
+        }
+    }
+#else
+    (void)data;
+#endif
+}
+
 #ifdef __MPI
 void Parallel_Grid::bcast(const double* const data_global, double* data_local, const int& rank, const bool is_sdft) const
 {

source/source_base/parallel_grid.h

@@ -21,6 +21,16 @@ class Parallel_Grid
         const int &nczp, const int &nrxx, const int &nbz, const int &bz,
         const int nprocgroup);
 
+    /**
+     * @brief Sum a distributed real-space grid across k-point pools.
+     *
+     * Uses a direct local-slab reduction for equal-sized pools. For uneven
+     * pools, reconstructs a common global layout before the cross-pool sum.
+     *
+     * @param data Local real-space grid data ordered as [x][y][z].
+     */
+    void reduce_across_pools(double* data) const;
+
 #ifdef __MPI
     /// @brief  Broadcast data from root to all processors. The index order is [x][y][z].
     void bcast(const double* const data_global, double* data_local, const int& rank, const bool is_sdft) const;

source/source_estate/module_charge/charge.cpp

@@ -42,10 +42,6 @@ Charge::Charge()
 Charge::~Charge()
 {
     this->destroy();
-#ifdef __MPI
-    delete[] rec;
-    delete[] dis;
-#endif
 }
 
 void Charge::set_rhopw(ModulePW::PW_Basis* rhopw_in)

source/source_estate/module_charge/charge.h

@@ -124,11 +124,6 @@ class Charge
     void init_final_scf(); //LiuXh add 20180619
 
 	public:
-    /**
-     * @brief init some arrays for mpi_inter_pools, rho_mpi
-     */
-    void init_chgmpi();
-
     /**
      * @brief Sum rho at different pools (k-point parallelism).
      *        Only used when GlobalV::KPAR > 1
@@ -169,12 +164,6 @@ class Charge
     bool allocate_rho;
 
     bool allocate_rho_final_scf; // LiuXh add 20180606
-
-#ifdef __MPI
-    int *rec = nullptr; //The number of elements each process should receive into the receive buffer.
-    int *dis = nullptr; //The displacement (relative to recvbuf) for each process in the receive buffer.
-#endif
-
 };
 
 #endif // charge

source/source_estate/module_charge/charge_init.cpp

@@ -249,9 +249,6 @@ void Charge::init_rho(const UnitCell& ucell,
         GlobalC::restart.info_load.load_charge_finish = true;
     }
 
-#ifdef __MPI
-    this->init_chgmpi();
-#endif
     if (PARAM.inp.init_chg == "wfc")
     {
         if (wfcpw == nullptr)