Add new functions/methods to featurize module

src/lobsterpy/featurize/__init__.py

@@ -3,4 +3,4 @@
 
 """This package provides the modules for featurzing Lobster data ready for ML studies."""
 
-from .utils import get_file_paths, get_structure_path
+from .utils import get_electronegativities, get_file_paths, get_reduced_mass, get_structure_path, sort_dict_by_value

src/lobsterpy/featurize/core.py

@@ -24,7 +24,7 @@
 
 from lobsterpy.cohp.analyze import Analysis
 
-from . import get_file_paths
+from . import get_electronegativities, get_file_paths, get_reduced_mass
 
 warnings.filterwarnings("ignore")
 
@@ -318,6 +318,79 @@ def get_lobsterpy_cba_dict(path_to_lobster_calc: str | Path, bonds: str, orbital
 
         return data
 
+    @staticmethod
+    def get_unique_bonds_df(
+        path_to_lobster_calc: str | Path,
+        bonds: str,
+        summed_icohps: bool,
+        rm_weighted_icohps: bool,
+        ids: str | None = None,
+    ) -> pd.DataFrame:
+        """
+        Generate a Pandas dataframe with icohps mean / summed icohps for unique bonds.
+
+        :param path_to_lobster_calc: path to lobsterpy lightweight json file
+        :param bonds: "all" or "cation-anion" bonds
+        :param summed_icohps:  bool indicating whether to sum the icohps for the unique bonds
+        :param rm_weighted_icohps: bool indicating whether to use reduced mass as weights for icohps
+        :param ids: set index name in the pandas dataframe. Default is None.
+
+        Returns:
+            Returns a pandas dataframe with icohps for unique bonds as columns
+
+        """
+        file_paths = get_file_paths(
+            path_to_lobster_calc=path_to_lobster_calc, requested_files=["structure", "cohpcar", "icohplist", "charge"]
+        )
+
+        try:
+            analyse = Analysis(
+                path_to_poscar=str(file_paths.get("structure")),
+                path_to_icohplist=str(file_paths.get("icohplist")),
+                path_to_cohpcar=str(file_paths.get("cohpcar")),
+                path_to_charge=str(file_paths.get("charge")),
+                cutoff_icohp=0.10,
+                which_bonds=bonds,
+                orbital_resolved=False,
+            )
+
+        except ValueError:
+            analyse = None
+
+        if not ids:
+            ids = Path(path_to_lobster_calc).name
+
+        # define a pandas dataframe
+        df = pd.DataFrame(index=[ids])
+
+        pair_icohps = {}
+        rm_pairs = {}
+        for plot_label in analyse.get_site_bond_resolved_labels():
+            atom_pair = plot_label.split(":")[-1].strip().split("-")
+            pair_rm = get_reduced_mass(atom_pair)
+            pair_en = get_electronegativities(atom_pair)
+            # sort atom names based on electronegativity
+            en_sorted_atom_pair = [x for _, x in sorted(zip(pair_en, atom_pair))]
+            rm_pairs["-".join(en_sorted_atom_pair)] = pair_rm
+            for lab in analyse.get_site_bond_resolved_labels()[plot_label]:
+                val = analyse.chemenv.Icohpcollection.get_icohp_by_label(lab)
+                if "-".join(en_sorted_atom_pair) not in pair_icohps:
+                    pair_icohps["-".join(en_sorted_atom_pair)] = [val]
+                else:
+                    pair_icohps["-".join(en_sorted_atom_pair)].append(val)
+
+        for atom_pair in pair_icohps:
+            if not summed_icohps and not rm_weighted_icohps:
+                df.loc[ids, f"{atom_pair}_icohp_mean"] = np.mean(pair_icohps[atom_pair])
+            elif summed_icohps and not rm_weighted_icohps:
+                df.loc[ids, f"{atom_pair}_icohp_sum"] = np.sum(pair_icohps[atom_pair])
+            elif not summed_icohps and rm_weighted_icohps:
+                df.loc[ids, f"{atom_pair}_rm_icohp_mean"] = np.mean(pair_icohps[atom_pair]) * rm_pairs[atom_pair]
+            else:
+                df.loc[ids, f"{atom_pair}_rm_icohp_sum"] = np.sum(pair_icohps[atom_pair]) * rm_pairs[atom_pair]
+
+        return df
+
 
 class CoxxFingerprint(NamedTuple):
     """

src/lobsterpy/featurize/utils.py

@@ -7,6 +7,7 @@
 
 from pathlib import Path
 
+from mendeleev import element
 from monty.os.path import zpath
 
 
@@ -91,3 +92,40 @@ def get_structure_path(lobster_path: Path) -> Path:
             return gz_file_path
 
     raise Exception
+
+
+def get_reduced_mass(atom_pair: list[str]) -> float:
+    """
+    Compute reduced mass between a pair of atoms.
+
+    :param atom_pair: list of atomic species symbols in string
+
+    :return: reduced mass
+    """
+    atom1 = element(atom_pair[0])
+    atom2 = element(atom_pair[1])
+    return (atom1.atomic_weight * atom2.atomic_weight) / (atom1.atomic_weight + atom2.atomic_weight)
+
+
+def get_electronegativities(atom_pair: list[str]) -> list[float]:
+    """
+    Get allen electronegativities for a pair of atoms.
+
+    :param atom_pair: list of atomic species symbols in string
+
+    :return: list of allen electronegativities
+    """
+    atom1 = element(atom_pair[0])
+    atom2 = element(atom_pair[1])
+    return [atom1.electronegativity_allen(), atom2.electronegativity_allen()]
+
+
+def sort_dict_by_value(input_dict: dict[str, float]) -> dict:
+    """
+    Sort dictionary by values.
+
+    :param input_dict: input dictionary
+
+    :return: sorted dictionary
+    """
+    return dict(sorted(input_dict.items(), key=lambda item: item[1]))