fix estimate_nbands function (#3149)

* modify formula to be correct, update tests

* remove pytest approx calls for integers

* pre-commit auto-fixes

* add ability to correctly account for NPAR and noncollinear magnetism in estimating nbands

* woohoo I just love linting

* snake_case vars

* dedup test comments

* snake_case

---------

Co-authored-by: Janosh Riebesell <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

pymatgen/io/vasp/sets.py

@@ -677,21 +677,30 @@ def estimate_nbands(self) -> int:
         Estimate the number of bands that VASP will initialize a
         calculation with by default. Note that in practice this
         can depend on # of cores (if not set explicitly).
+        Note that this formula is slightly different than the formula on the VASP wiki
+        (as of July 2023). This is because the formula in the source code (`main.F`) is
+        slightly different than what is on the wiki.
         """
-        nions = len(self.structure)
+        n_ions = len(self.structure)
 
-        # from VASP's point of view, the number of magnetic atoms are
-        # the number of atoms with non-zero magmoms, so use Incar as
-        # source of truth
-        nmag = len([m for m in self.incar["MAGMOM"] if not np.allclose(m, 0)])
+        if self.incar["ISPIN"] == 1:  # per the VASP source, if non-spin polarized ignore n_mag
+            n_mag = 0
+        else:  # otherwise set equal to sum of total magmoms
+            n_mag = sum(self.incar["MAGMOM"])
+            n_mag = np.floor((n_mag + 1) / 2)
 
-        # by definition, if non-spin polarized ignore nmag
-        if (not nmag) or (self.incar["ISPIN"] == 1):
-            nbands = np.ceil(self.nelect / 2 + nions / 2)
-        else:
-            nbands = np.ceil(0.6 * self.nelect + nmag)
+        possible_val_1 = np.floor((self.nelect + 2) / 2) + max(np.floor(n_ions / 2), 3)
+        possible_val_2 = np.floor(self.nelect * 0.6)
+
+        n_bands = max(possible_val_1, possible_val_2) + n_mag
+
+        if self.incar.get("LNONCOLLINEAR") is True:
+            n_bands = n_bands * 2
+
+        if n_par := self.incar.get("NPAR"):
+            n_bands = (np.floor((n_bands + n_par - 1) / n_par)) * n_par
 
-        return int(nbands)
+        return int(n_bands)
 
     def __str__(self):
         return type(self).__name__

pymatgen/io/vasp/tests/test_sets.py

@@ -236,15 +236,6 @@ def test_nelect(self):
         struct = Structure(lattice, ["Si", "Si", "Fe"], coords)
         assert MITRelaxSet(struct).nelect == 16
 
-        # Test estimate of number of bands (function of nelect) with nmag>0
-        assert MITRelaxSet(struct).estimate_nbands() == 13
-        assert MPRelaxSet(struct).estimate_nbands() == 17
-
-        # Test estimate of number of bands (function of nelect) with nmag==0
-        struct = Structure(lattice, ["Si", "Si", "Si"], coords)
-        assert MITRelaxSet(struct).estimate_nbands() == 11
-        assert MPRelaxSet(struct).estimate_nbands() == 11
-
         # Check that it works even when oxidation states are present. Was a bug
         # previously.
         struct = Structure(lattice, ["Si4+", "Si4+", "Fe2+"], coords)
@@ -256,6 +247,32 @@ def test_nelect(self):
         assert MITRelaxSet(struct).nelect == 16
         assert MPRelaxSet(struct).nelect == 22
 
+    @skip_if_no_psp_dir
+    def test_estimate_nbands(self):
+        # estimate_nbands is a function of n_elect, n_ions, magmom, noncollinearity of magnetism, and n_par
+        coords = [[0] * 3, [0.5] * 3, [0.75] * 3]
+        lattice = Lattice.cubic(4)
+
+        # pure Si
+        struct = Structure(lattice, ["Si", "Si", "Si"], coords)
+        assert MITRelaxSet(struct).estimate_nbands() == 11
+        assert MPRelaxSet(struct).estimate_nbands() == 11
+
+        # Si + Fe
+        struct = Structure(lattice, ["Si", "Si", "Fe"], coords)
+        assert MITRelaxSet(struct).estimate_nbands() == 15
+        assert MPRelaxSet(struct).estimate_nbands() == 18
+
+        # Si + Fe with NPAR = 4
+        uis = {"NPAR": 4}
+        assert MITRelaxSet(struct, user_incar_settings=uis).estimate_nbands() == approx(16)
+        assert MPRelaxSet(struct, user_incar_settings=uis).estimate_nbands() == approx(20)
+
+        # Si + Fe with noncollinear magnetism turned on
+        uis = {"LNONCOLLINEAR": True}
+        assert MITRelaxSet(struct, user_incar_settings=uis).estimate_nbands() == approx(30)
+        assert MPRelaxSet(struct, user_incar_settings=uis).estimate_nbands() == approx(36)
+
     @skip_if_no_psp_dir
     def test_get_incar(self):
         incar = self.mp_set.incar