Extend lobsterenv for coop/cobi

pymatgen/io/lobster/lobsterenv.py

@@ -44,13 +44,15 @@ class LobsterNeighbors(NearNeighbors):
     def __init__(
         self,
         are_coops=False,
+        are_cobis=False,
         filename_ICOHP=None,
         valences=None,
         limits=None,
         structure=None,
         additional_condition=0,
         only_bonds_to=None,
         perc_strength_ICOHP=0.15,
+        noise_cutoff=0.1,
         valences_from_charges=False,
         filename_CHARGE=None,
         which_charge="Mulliken",
@@ -64,11 +66,11 @@ def __init__(
         """
 
         Args:
-            are_coops: (bool) if True, the file is a ICOOPLIST.lobster and not a ICOHPLIST.lobster; only tested for
-            ICOHPLIST.lobster so far
-            filename_ICOHP: (str) Path to ICOHPLIST.lobster
+            are_coops: (bool) if True, the file is a ICOOPLIST.lobster and not a ICOHPLIST.lobster;
+            are_cobis: (Bool) if True, the file is a ICOBILIST.lobster and not a ICOHPLIST.lobster
+            filename_ICOHP: (str) Path to ICOHPLIST.lobster/ICOOPLIST.lobster/ICOBILIST.lobster
             valences: (list of integers/floats) gives valence/charge for each element
-            limits: limit to decide which ICOHPs should be considered
+            limits: list of limits [lower, upper] to decide which ICOHPs/ICOBI/ICOOP should be considered
             structure: (Structure Object) typically constructed by: Structure.from_file("POSCAR") (Structure object
             from pymatgen.core.structure)
             additional_condition:   Additional condition that decides which kind of bonds will be considered
@@ -82,7 +84,7 @@ def __init__(
             only_bonds_to: (list of str) will only consider bonds to certain elements (e.g. ["O"] for oxygen)
             perc_strength_ICOHP: if no limits are given, this will decide which icohps will still be considered (
             relative to
-            the strongest ICOHP)
+            the strongest ICOHP/ICOOP/ICOBI)
             valences_from_charges: if True and path to CHARGE.lobster is provided, will use Lobster charges (
             Mulliken) instead of valences
             filename_CHARGE: (str) Path to Charge.lobster
@@ -97,16 +99,18 @@ def __init__(
             id_blist_sg2: (str) Identity of data in filename_blist_sg2,
                 e.g., "icoop" or "icobi"
         """
-        self.ICOHP = Icohplist(are_coops=are_coops, filename=filename_ICOHP)
+        self.ICOHP = Icohplist(are_coops=are_coops, are_cobis=are_cobis, filename=filename_ICOHP)
         self.Icohpcollection = self.ICOHP.icohpcollection
         self.structure = structure
         self.limits = limits
         self.only_bonds_to = only_bonds_to
         self.adapt_extremum_to_add_cond = adapt_extremum_to_add_cond
         self.are_coops = are_coops
+        self.are_cobis = are_cobis
         self.add_additional_data_sg = add_additional_data_sg
         self.filename_blist_sg1 = filename_blist_sg1
         self.filename_blist_sg2 = filename_blist_sg2
+        self.noise_cutoff = noise_cutoff
 
         allowed_arguments = ["icoop", "icobi"]
         if id_blist_sg1.lower() not in allowed_arguments or id_blist_sg2.lower() not in allowed_arguments:
@@ -141,9 +145,6 @@ def __init__(
                 are_cobis=are_cobis_id2,
             )
 
-        if are_coops:
-            raise ValueError("Algorithm only works correctly for ICOHPLIST.lobster")
-
         # will check if the additional condition is correctly delivered
         if additional_condition in range(0, 7):
             self.additional_condition = additional_condition
@@ -355,16 +356,18 @@ def get_light_structure_environment(self, only_cation_environments=False, only_i
 
     def get_info_icohps_to_neighbors(self, isites=None, onlycation_isites=True):
         """
-        this method Return information of cohps of neighbors
+        this method Return information of cohps/coops/icobis of neighbors depending on input bond list file
+        (could be ICOOPLIST.lobster/ICOHPLIST.lobster/ICOBILIST.lobster)
 
         Args:
             isites: list of site ids, if isite==None, all isites will be used to add the icohps of the neighbors
             onlycation_isites: will only use cations, if isite==[]
 
 
         Returns:
-            sum of icohps of neighbors to certain sites [given by the id in structure], number of bonds to this site,
-            labels (from ICOHPLIST) for
+            sum of icohps/icoops/icobis of neighbors to certain sites [given by the id in structure],
+            number of bonds to this site,
+            labels (from ICOHPLIST/ICOOPLIST/ICOBILIST) for
             these bonds
             [the latter is useful for plotting summed COHP plots],
             list of the central isite for each label
@@ -431,7 +434,7 @@ def plot_cohps_of_neighbors(
         # include COHPPlotter and plot a sum of these COHPs
         # might include option to add Spin channels
         # implement only_bonds_to
-        cp = CohpPlotter()
+        cp = CohpPlotter(are_cobis=self.are_cobis, are_coops=self.are_coops)
 
         plotlabel, summed_cohp = self.get_info_cohps_to_neighbors(
             path_to_COHPCAR,
@@ -493,7 +496,13 @@ def get_info_cohps_to_neighbors(
             self.structure.to(filename=path, fmt="POSCAR")
 
             if not hasattr(self, "completecohp"):
-                self.completecohp = CompleteCohp.from_file(fmt="LOBSTER", filename=path_to_COHPCAR, structure_file=path)
+                self.completecohp = CompleteCohp.from_file(
+                    fmt="LOBSTER",
+                    filename=path_to_COHPCAR,
+                    structure_file=path,
+                    are_coops=self.are_coops,
+                    are_cobis=self.are_cobis,
+                )
 
         # will check that the number of bonds in ICOHPLIST and COHPCAR are identical
         # further checks could be implemented
@@ -1134,10 +1143,10 @@ def _get_limit_from_extremum(
 
         Args:
             icohpcollection: icohpcollection object
-            percentage: will determine which ICOHPs will be considered (only 0.15 from the maximum value)
+            percentage: will determine which ICOHPs/ICOOP/ICOBI will be considered (only 0.15 from the maximum value)
             adapt_extremum_to_add_cond: should the extrumum be adapted to the additional condition
             additional_condition: additional condition to determine which bonds are relevant
-        Returns: [-100000, min(max_icohp*0.15,-0.1)]
+        Returns: [-inf, min(max_icohp*0.15,noise_cutoff)] / [max(max_icohp*0.15, noise_cutoff),inf]
         """
         # TODO: make it work for COOPs
         if not adapt_extremum_to_add_cond or additional_condition == 0:
@@ -1154,15 +1163,22 @@ def _get_limit_from_extremum(
                 if (val1 < 0.0 < val2) or (val2 < 0.0 < val1):
                     list_icohps.append(value.summed_icohp)
 
-            extremum_based = min(list_icohps) * percentage
+            if not self.are_coops and not self.are_cobis:
+                extremum_based = min(list_icohps) * percentage
+            else:
+                extremum_based = max(list_icohps) * percentage
 
         elif additional_condition == 2:
             # NO_ELEMENT_TO_SAME_ELEMENT_BONDS
             list_icohps = []
             for value in icohpcollection._icohplist.values():
                 if value._atom1.rstrip("0123456789") != value._atom2.rstrip("0123456789"):
                     list_icohps.append(value.summed_icohp)
-            extremum_based = min(list_icohps) * percentage
+
+            if not self.are_coops and not self.are_cobis:
+                extremum_based = min(list_icohps) * percentage
+            else:
+                extremum_based = max(list_icohps) * percentage
 
         elif additional_condition == 3:
             # ONLY_ANION_CATION_BONDS_AND_NO_ELEMENT_TO_SAME_ELEMENT_BONDS = 3
@@ -1179,13 +1195,23 @@ def _get_limit_from_extremum(
                     and value._atom1.rstrip("0123456789") != value._atom2.rstrip("0123456789")
                 ):
                     list_icohps.append(value.summed_icohp)
-            extremum_based = min(list_icohps) * percentage
+
+            if not self.are_coops and not self.are_cobis:
+                extremum_based = min(list_icohps) * percentage
+            else:
+                extremum_based = max(list_icohps) * percentage
+
         elif additional_condition == 4:
             list_icohps = []
             for value in icohpcollection._icohplist.values():
                 if value._atom1.rstrip("0123456789") == "O" or value._atom2.rstrip("0123456789") == "O":
                     list_icohps.append(value.summed_icohp)
-            extremum_based = min(list_icohps) * percentage
+
+            if not self.are_coops and not self.are_cobis:
+                extremum_based = min(list_icohps) * percentage
+            else:
+                extremum_based = max(list_icohps) * percentage
+
         elif additional_condition == 5:
             # DO_NOT_CONSIDER_ANION_CATION_BONDS=5
             list_icohps = []
@@ -1197,7 +1223,11 @@ def _get_limit_from_extremum(
 
                 if (val1 > 0.0 and val2 > 0.0) or (val1 < 0.0 and val2 < 0.0):
                     list_icohps.append(value.summed_icohp)
-            extremum_based = min(list_icohps) * percentage
+
+            if not self.are_coops and not self.are_cobis:
+                extremum_based = min(list_icohps) * percentage
+            else:
+                extremum_based = max(list_icohps) * percentage
 
         elif additional_condition == 6:
             # ONLY_CATION_CATION_BONDS=6
@@ -1210,13 +1240,20 @@ def _get_limit_from_extremum(
 
                 if val1 > 0.0 and val2 > 0.0:
                     list_icohps.append(value.summed_icohp)
-            extremum_based = min(list_icohps) * percentage
 
-        # if not self.are_coops:
-        max_here = min(extremum_based, -0.1)
-        return -100000, max_here
-        # else:
-        #    return extremum_based, 100000
+            if not self.are_coops and not self.are_cobis:
+                extremum_based = min(list_icohps) * percentage
+            else:
+                extremum_based = max(list_icohps) * percentage
+
+        if not self.are_coops and not self.are_cobis:
+            max_here = min(extremum_based, -self.noise_cutoff) if self.noise_cutoff is not None else extremum_based
+            return -float("inf"), max_here
+        if self.are_coops or self.are_cobis:
+            min_here = max(extremum_based, self.noise_cutoff) if self.noise_cutoff is not None else extremum_based
+            return min_here, float("inf")
+
+        return None
 
 
 class LobsterLightStructureEnvironments(LightStructureEnvironments):

pymatgen/io/lobster/tests/test_lobsterenv.py

@@ -124,6 +124,21 @@ def setUp(self):
             structure=Structure.from_file(os.path.join(test_dir_env, "POSCAR.mp_353.gz")),
             additional_condition=6,
         )
+        # coop / cobi
+        self.chemenvlobster1_coop_NaCl = LobsterNeighbors(
+            are_coops=True,
+            filename_ICOHP=os.path.join(test_dir_env, "ICOOPLIST.lobster.NaCl.gz"),
+            structure=Structure.from_file(os.path.join(test_dir_env, "POSCAR.NaCl.gz")),
+            additional_condition=1,
+            noise_cutoff=None,
+        )
+
+        self.chemenvlobster1_cobi_mp470 = LobsterNeighbors(
+            are_coops=True,
+            filename_ICOHP=os.path.join(test_dir_env, "ICOBILIST.lobster.mp_470.gz"),
+            structure=Structure.from_file(os.path.join(test_dir_env, "POSCAR.mp_470.gz")),
+            additional_condition=1,
+        )
 
         # TODO: use charge instead of valence
         self.chemenvlobster1_charges = LobsterNeighbors(
@@ -217,17 +232,6 @@ def setUp(self):
             adapt_extremum_to_add_cond=True,
         )
 
-    def test_use_of_coop(self):
-        with pytest.raises(ValueError):
-            _ = LobsterNeighbors(
-                are_coops=True,
-                filename_ICOHP=os.path.join(test_dir_env, "ICOHPLIST.lobster.mp_353.gz"),
-                structure=Structure.from_file(os.path.join(test_dir_env, "POSCAR.mp_353.gz")),
-                valences_from_charges=True,
-                filename_CHARGE=os.path.join(test_dir_env, "CHARGE.lobster.mp-353.gz"),
-                additional_condition=1,
-            )
-
     def test_cation_anion_mode_without_ions(self):
         with pytest.raises(ValueError) as exc:
             _ = LobsterNeighbors(
@@ -444,6 +448,26 @@ def test_get_nn_info(self):
             == 2
         )
 
+        assert (
+            len(
+                self.chemenvlobster1_coop_NaCl.get_nn(
+                    structure=Structure.from_file(os.path.join(test_dir_env, "POSCAR.NaCl.gz")),
+                    n=0,
+                )
+            )
+            == 6
+        )
+
+        assert (
+            len(
+                self.chemenvlobster1_cobi_mp470.get_nn(
+                    structure=Structure.from_file(os.path.join(test_dir_env, "POSCAR.mp_470.gz")),
+                    n=3,
+                )
+            )
+            == 3
+        )
+
         # NO_ELEMENT_TO_SAME_ELEMENT_BONDS = 2
         assert (
             len(