Bugfix: save density when spin unrestricted (pyscf#205)

* customize radii * fixed a bug when saving electron density with spin unrestricted
bytedance · Aug 30, 2024 · 081d7bb · 081d7bb
1 parent a065d7d
commit 081d7bb
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/gpu4pyscf/drivers/dft_driver.py b/gpu4pyscf/drivers/dft_driver.py
@@ -190,7 +190,8 @@ def run_dft(mol_name, config, charge=None, spin=0):
         if save_density and xc.lower() != 'hf':
             weights = mf.grids.weights
             coords = mf.grids.coords
-            rho = mf._numint.get_rho(mf.mol, dm, mf.grids)
+            dm0 = dm[0] + dm[1] if dm.ndim == 3 else dm
+            rho = mf._numint.get_rho(mf.mol, dm0, mf.grids)
 
             if isinstance(weights, cupy.ndarray): weights = weights.get()
             if isinstance(coords, cupy.ndarray):  coords  = coords.get()