Added methods to compute and compare DOS fingerprints

pymatgen/electronic_structure/dos.py

@@ -9,6 +9,7 @@
 
 import functools
 import warnings
+from collections import namedtuple
 from typing import Mapping
 
 import numpy as np
@@ -1165,6 +1166,150 @@ def get_upper_band_edge(
         upper_band_edge = energies[np.argmax(densities)]
         return upper_band_edge
 
+    def get_dos_fp(self, type="summed_pdos", binning=True, min_e=None, max_e=None, nbins=256, normalize=True):
+        """
+        Generates the DOS fingerprint based on work of
+        F. Knoop, T. A. r Purcell, M. Scheffler, C. Carbogno, J. Open Source Softw. 2020, 5, 2671.
+        Source - https://gitlab.com/vibes-developers/vibes/-/tree/master/vibes/materials_fp
+        Copyright (c) 2020 Florian Knoop, Thomas A.R.Purcell, Matthias Scheffler, Christian Carbogno
+
+
+        Args:
+            type (str): Specify fingerprint type needed can accept 's/p/d/f/summed_pdos/tdos'
+            min_e (float): The minimum mode energy to include in the fingerprint
+            max_e (float): The maximum mode energy to include in the fingerprint
+            nbins (int): Number of bins to be used in the fingerprint
+            normalize (bool): If true, normalizes the area under fp to equal to 1
+
+        Returns:
+            Fingerprint(namedtuple) : The electronic density of states fingerprint
+            of format (energies, densities, type, nbins)
+        """
+        fp_tup = namedtuple("fingerprint", "energies densities type nbins binwidth")
+        energies = self.energies - self.efermi
+
+        if max_e is None:
+            max_e = np.max(energies)
+
+        if min_e is None:
+            min_e = np.min(energies)
+
+        pdos_obj = self.get_spd_dos()
+
+        pdos = {}
+        for k, _v in pdos_obj.items():
+            dens = pdos_obj[k].get_densities()
+
+            pdos.update({k.name: dens})
+
+        sum_pdos_array = []
+        for dos in pdos.values():
+            sum_pdos_array.append(np.array(dos))
+
+        pdos.update({"summed_pdos": np.sum(sum_pdos_array, axis=0)})
+        pdos.update({"tdos": self.get_densities()})
+
+        try:
+
+            densities = pdos[type]
+            if len(energies) < nbins:
+                inds = np.where((energies >= min_e) & (energies <= max_e))
+                return fp_tup(energies[inds], densities[inds], type, len(energies), np.diff(energies)[0])
+
+            if binning:
+                ener_bounds = np.linspace(min_e, max_e, nbins + 1)
+                ener = ener_bounds[:-1] + (ener_bounds[1] - ener_bounds[0]) / 2.0
+                bin_width = np.diff(ener)[0]
+            else:
+                ener_bounds = np.array(energies)
+                ener = np.append(energies, [energies[-1] + np.abs(energies[-1]) / 10])
+                nbins = len(energies)
+                bin_width = np.diff(energies)[0]
+
+            dos_rebin = np.zeros(ener.shape)
+
+            for ii, e1, e2 in zip(range(len(ener)), ener_bounds[0:-1], ener_bounds[1:]):
+                inds = np.where((energies >= e1) & (energies < e2))[0]
+                dos_rebin[ii] = np.sum(densities[inds])
+            if normalize:  # scale dos bins to make area under histogram equal 1
+                area = np.sum(dos_rebin * bin_width)
+                dos_rebin_sc = dos_rebin / area
+            else:
+                dos_rebin_sc = dos_rebin
+
+            return fp_tup(np.array([ener]), dos_rebin_sc, type, nbins, bin_width)
+
+        except KeyError:
+            raise ValueError(
+                "Please recheck type requested, either the orbital projections unavailable in input dos or "
+                "some there exist some mistake in the spelling."
+            )
+
+    @staticmethod
+    def fp_to_dict(fp):
+        """Converts a fingerprint into a dictionary
+
+        Args:
+            fp: The DOS fingerprint to be converted into a dictionary
+
+        Returns:
+            dict: A dict of the fingerprint Keys=type, Values=np.ndarray(energies, states)
+        """
+        fp_dict = {}
+        fp_dict[fp[2]] = np.array([fp[0], fp[1]]).T
+
+        return fp_dict
+
+    @staticmethod
+    def get_dos_fp_similarity(fp1, fp2, col=1, pt="All", normalize=False, tanimoto=False):
+        """Calculates the similarity index (dot product) of two fingerprints
+
+        Args:
+            fp1 (get_dos_fp): The 1st dos fingerprint
+            fp2 (get_dos_fp):  The 2nd dos fingerprint
+            col (int): The item in the fingerprints to take the dot product of (0=Energies 1=Densities)
+            pt (int or 'All') : The index of the point that the dot product is to be taken
+            normalize (bool): If True normalize the scalar product to 1
+            tanimoto (bool): If True will compute tanimoto index
+
+        Returns:
+        Similarity index (float): The value of dot product
+        """
+        if not isinstance(fp1, dict):
+            fp1_dict = CompleteDos.fp_to_dict(fp1)
+        else:
+            fp1_dict = fp1
+
+        if not isinstance(fp2, dict):
+            fp2_dict = CompleteDos.fp_to_dict(fp2)
+        else:
+            fp2_dict = fp2
+
+        if pt == "All":
+            vec1 = np.array([pt[col] for pt in fp1_dict.values()]).flatten()
+            vec2 = np.array([pt[col] for pt in fp2_dict.values()]).flatten()
+        else:
+            vec1 = fp1_dict[fp1[2][pt]][col]
+            vec2 = fp2_dict[fp2[2][pt]][col]
+
+        if not normalize and tanimoto:
+            rescale = np.linalg.norm(vec1) ** 2 + np.linalg.norm(vec2) ** 2 - np.dot(vec1, vec2)
+            return np.dot(vec1, vec2) / rescale
+
+        elif not tanimoto and normalize:
+            rescale = np.linalg.norm(vec1) * np.linalg.norm(vec2)
+            return np.dot(vec1, vec2) / rescale
+
+        elif not tanimoto and not normalize:
+            rescale = 1.0
+            return np.dot(vec1, vec2) / rescale
+
+        else:
+            raise ValueError(
+                "Cannot compute similarity index, Please set either one of normalize/tanimoto arg to true "
+                "or set both to false"
+            )
+
     @classmethod
     def from_dict(cls, d) -> CompleteDos:
         """

pymatgen/electronic_structure/tests/test_dos.py

@@ -257,6 +257,54 @@ def test_kurtosis(self):
         kurtosis = dos.get_band_kurtosis()
         self.assertAlmostEqual(kurtosis, 7.764506941340621)
 
+    def test_get_dos_fp(self):
+        # normalize=True
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, nbins=56, normalize=True)
+        bin_width = np.diff(dos_fp.energies)[0][0]
+        self.assertLessEqual(max(dos_fp.energies[0]), 0)
+        self.assertGreaterEqual(min(dos_fp.energies[0]), -10)
+        self.assertEqual(len(dos_fp.energies[0]), 56)
+        self.assertEqual(dos_fp.type, "s")
+        self.assertAlmostEqual(sum(dos_fp.densities * bin_width), 1, delta=0.001)
+        # normalize=False
+        dos_fp2 = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, nbins=56, normalize=False)
+        bin_width2 = np.diff(dos_fp2.energies)[0][0]
+        self.assertAlmostEqual(sum(dos_fp2.densities * bin_width2), 7.279303571428509, delta=0.001)
+        self.assertAlmostEqual(dos_fp2.binwidth, bin_width2, delta=0.001)
+        # binning=False
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=None, max_e=None, nbins=56, normalize=True, binning=False)
+        self.assertEqual(dos_fp.nbins, len(self.dos.energies))
+
+    def test_get_dos_fp_similarity(self):
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, nbins=56, normalize=True)
+        dos_fp2 = self.dos.get_dos_fp(type="tdos", min_e=-10, max_e=0, nbins=56, normalize=True)
+        similarity_index = self.dos.get_dos_fp_similarity(dos_fp, dos_fp2, col=1, tanimoto=True)
+        self.assertAlmostEqual(similarity_index, 0.3342481451042263, delta=0.0001)
+
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, nbins=56, normalize=True)
+        dos_fp2 = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, nbins=56, normalize=True)
+        similarity_index = self.dos.get_dos_fp_similarity(dos_fp, dos_fp2, col=1, tanimoto=True)
+        self.assertEqual(similarity_index, 1)
+
+    def test_dos_fp_exceptions(self):
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, nbins=56, normalize=True)
+        dos_fp2 = self.dos.get_dos_fp(type="tdos", min_e=-10, max_e=0, nbins=56, normalize=True)
+        # test exceptions
+        with self.assertRaises(ValueError) as err:
+
+            self.dos.get_dos_fp_similarity(dos_fp, dos_fp2, col=1, tanimoto=True, normalize=True)
+        self.assertEqual(
+            err.exception.__str__(),
+            "Cannot compute similarity index, Please set either one of normalize/tanimoto arg to true or set both to false",
+        )
+        with self.assertRaises(ValueError) as err:
+
+            self.dos.get_dos_fp(type="k", min_e=-10, max_e=0, nbins=56, normalize=True)
+        self.assertEqual(
+            err.exception.__str__(),
+            "Please recheck type requested, either the orbital projections unavailable in input dos or some there exist some mistake in the spelling.",
+        )
+
 
 class DOSTest(PymatgenTest):
     def setUp(self):