Merge 7dbfd81 into adb541c

pymatgen/command_line/bader_caller.py

@@ -154,7 +154,7 @@ def __init__(
             self.is_vasp = False
             self.cube = Cube(fpath)
             self.structure = self.cube.structure
-            self.nelects = self.structure.site_properties.get("nelect", [])  # For cube, see if struc has nelects
+            self.nelects = None
 
         tmpfile = "CHGCAR" if chgcar_filename else "CUBE"
         with ScratchDir("."):
@@ -275,25 +275,27 @@ def get_charge(self, atom_index):
         """
         return self.data[atom_index]["charge"]
 
-    def get_charge_transfer(self, atom_index):
+    def get_charge_transfer(self, atom_index, nelect=None):
         """
-        Returns the charge transferred for a particular atom. Requires POTCAR
-        to be supplied.
+        Returns the charge transferred for a particular atom. If the arg nelect
+        is not supplied, then POTCAR must be supplied to determine nelectrons.
 
         Args:
             atom_index:
                 Index of atom.
+            nelect:
+                number of electrons associated with an isolated atom at this index.
+                For most DFT codes this corresponds to the number of valence electrons
+                associated with the pseudopotential
 
         Returns:
             Charge transfer associated with atom from the Bader analysis.
-            Given by final charge on atom - nelectrons in POTCAR for
+            Given by final charge on atom - nelectrons for
             associated atom.
         """
-        if not self.nelects:
-            raise ValueError(
-                "No NELECT info! Need POTCAR for VASP, or a structure object" "with nelect as a site property."
-            )
-        return self.data[atom_index]["charge"] - self.nelects[atom_index]
+        if not self.nelects and nelect is None:
+            raise ValueError("No NELECT info! Need POTCAR for VASP or nelect argument" "for cube file")
+        return self.data[atom_index]["charge"] - (nelect if nelect is not None else self.nelects[atom_index])
 
     def get_charge_decorated_structure(self):
         """
@@ -308,28 +310,54 @@ def get_charge_decorated_structure(self):
         struc.add_site_property("charge", charges)
         return struc
 
-    def get_oxidation_state_decorated_structure(self):
+    def get_oxidation_state_decorated_structure(self, nelects=None):
         """
         Returns an oxidation state decorated structure based on bader analysis results.
 
         Note, this assumes that the Bader analysis was correctly performed on a file
         with electron densities
         """
-        charges = [-self.get_charge_transfer(i) for i in range(len(self.structure))]
+        charges = [
+            -self.get_charge_transfer(i, None if not nelects else nelects[i]) for i in range(len(self.structure))
+        ]
         struc = self.structure.copy()
         struc.add_oxidation_state_by_site(charges)
         return struc
 
-    def get_spin_state_decorated_structure(self):
+    def get_decorated_structure(self, property_name, average=False):
         """
-        Returns a structure decorated with spins.
+        Get a property-decorated structure from the Bader analysis.
 
-        Note, this assumes that the Bader analysis was correctly performed on a file
-        with spin densities
+        This is distinct from getting charge decorated structure, which assumes
+        the "standard" Bader analysis of electron densities followed by converting
+        electron count to charge. The expected way to use this is to call Bader on
+        a non-charge density file such as a spin density file, electrostatic potential
+        file, etc., while using the charge density file as the reference (chgref_filename)
+        so that the partitioning is determined via the charge, but averaging or integrating
+        is done for another property.
+
+        User warning: Bader analysis cannot automatically determine what property is
+        inside of the file. So if you want to use this for a non-conventional property
+        like spin, you must ensure that you have the file is for the appropriate
+        property and you have an appropriate reference file.
+
+        Args:
+            property_name: name of the property to assign to the structure, note that 
+                if name is "spin" this is handled as a special case, and the appropriate 
+                spin properties are set on the species in the structure
+            average: whether or not to return the average of this property, rather
+                than the total, by dividing by the atomic volume.
+
+        Returns:
+            structure with site properties assigned via Bader Analysis
         """
-        spins = [self.get_charge(i) for i in range(len(self.structure))]
+        vals = [self.get_charge(i) for i in range(len(self.structure))]
         struc = self.structure.copy()
-        struc.add_spin_by_site(spins)
+        if average:
+            vals = np.divide(vals, [d["atomic_vol"] for d in self.data])
+        struc.add_site_property(property_name, vals)
+        if property_name == "spin":
+            struc.add_spin_by_site(vals)
         return struc
 
     @property