[feat] allow spinave in calc_density_correction also for vasp and qe

Author: the-hampel

python/triqs_dft_tools/sumk_dft.py

@@ -2150,8 +2150,10 @@ def calc_density_correction(self, filename=None, dm_type=None, spinave=False, kp
                    DFT code to write the density correction for. Options:
                    'vasp', 'wien2k', 'elk' or 'qe'. Needs to be set for 'qe'
         spinave : logical
-                   Elk specific and for magnetic calculations in DMFT only.
-                   It averages the spin to keep the DFT part non-magnetic.
+                   For magnetic calculations in DMFT only.
+                   If True, the density matrix is averaged over spin up/down to
+                   keep the DFT part non-magnetic (supported for 'elk', 'vasp'
+                   and 'qe').
         kpts_to_write : iterable of int
                    Indices of k points that are written to file. If None (default),
                    all k points are written. Only implemented for dm_type 'vasp'
@@ -2170,11 +2172,11 @@ def calc_density_correction(self, filename=None, dm_type=None, spinave=False, kp
                          the corresponing total charge `dens`.
 
         """
-        #automatically set dm_type if required
+        # automatically set dm_type if required
         if dm_type is None:
             dm_type = self.dft_code
 
-        assert dm_type in ('vasp', 'wien2k', 'elk', 'qe'), "'dm_type' must be either 'vasp', 'wienk', 'elk' or 'qe'"
+        assert dm_type in ('vasp', 'wien2k', 'elk', 'qe'), "'dm_type' must be either 'vasp', 'wien2k', 'elk' or 'qe'"
         # default file names
         if filename is None:
             if dm_type == 'wien2k':
@@ -2334,7 +2336,10 @@ def calc_density_correction(self, filename=None, dm_type=None, spinave=False, kp
             else:
                 assert np.min(kpts_to_write) >= 0 and np.max(kpts_to_write) < self.n_k
 
-            assert self.SP == 0, "Spin-polarized density matrix is not implemented"
+            if not spinave:
+                assert self.SP == 0, "Spin-polarized density matrix is not implemented"
+            elif self.SP != 0:
+                mpi.report("SumK calc_density_correction: WARNING! Averaging out spin-polarized correction in the density channel")
 
             if mpi.is_master_node():
                 if filename == 'vaspgamma.h5':
@@ -2343,8 +2348,13 @@ def calc_density_correction(self, filename=None, dm_type=None, spinave=False, kp
                         bnd_win_towrite = [band_window[0][:n_k_ibz,:]]
                         vasp_h5['band_window'] = bnd_win_towrite
                         vasp_h5.create_group('deltaN')
-                        vasp_h5['deltaN']['up'] = deltaN['up'][:n_k_ibz]
-                        vasp_h5['deltaN']['down'] = deltaN['down'][:n_k_ibz]
+                        if spinave and self.SP != 0 and self.SO == 0:
+                            deltaN_ave = [(u + d) / 2.0 for (u, d) in zip(deltaN['up'][:n_k_ibz], deltaN['down'][:n_k_ibz])]
+                            vasp_h5['deltaN']['up'] = deltaN_ave
+                            vasp_h5['deltaN']['down'] = deltaN_ave
+                        else:
+                            vasp_h5['deltaN']['up'] = deltaN['up'][:n_k_ibz]
+                            vasp_h5['deltaN']['down'] = deltaN['down'][:n_k_ibz]
                 else:
                     with open(filename, 'w') as f:
                         f.write(" -1  -1  ! Number of k-points, default number of bands\n")  # % len(kpts_to_write))
@@ -2410,7 +2420,9 @@ def calc_density_correction(self, filename=None, dm_type=None, spinave=False, kp
                             f.write("\n")
 
         elif dm_type == 'qe':
-            if self.SP == 0:
+            if not spinave:
+                assert self.SP == 0, "Spin-polarized density matrix is not implemented"
+            elif self.SP != 0:
                 mpi.report("SUMK calc_density_correction: WARNING! Averaging out spin-polarized correction in the density channel")
 
             subgrp = 'dft_update'