Interactive plotter

lobsterpy/cohp/describe.py

@@ -7,6 +7,7 @@
 
 from pathlib import Path
 from lobsterpy.plotting import PlainCohpPlotter
+from LobsterPy.lobsterpy.plotting import InteractiveCohpPlotter
 
 
 class Description:
@@ -266,6 +267,66 @@ def plot_cohps(
         else:
             plot.close()
 
+    def plot_interactive_cohps(
+        self,
+        save_as_html=False,
+        filename=None,
+        ylim=None,
+        xlim=None,
+        integrated=False,
+        title="",
+        sigma=None,
+    ):
+        """
+        Will automatically generate interactive plots of the most relevant COHP
+
+        Args:
+            save (bool): will save the plot to a file
+            filename (str/Path):
+            ylim (list of float): energy scale that is shown in plot (eV)
+            xlim(list of float): energy range for COHPs in eV
+            integrated (bool): if True, integrated COHPs will be shown
+            sigma: Standard deviation of Gaussian broadening applied to
+                population data. If None, no broadening will be added.
+            skip_show (bool): if True, the plot will not be shown.
+
+        Returns:
+            A plotly.graph_objects.Figure object.
+
+        """
+        set_cohps = self.analysis_object.set_cohps
+        if self.analysis_object.whichbonds == "cation-anion":
+            set_inequivalent_cations = self.analysis_object.set_inequivalent_ions
+        elif self.analysis_object.whichbonds == "all":
+            set_inequivalent_cations = self.analysis_object.set_inequivalent_ions
+        set_labels_cohps = self.analysis_object.set_labels_cohps
+        structure = self.analysis_object.structure
+
+        for iplot, (ication, labels, cohps) in enumerate(
+            zip(set_inequivalent_cations, set_labels_cohps, set_cohps)
+        ):
+
+            namecation = str(structure[ication].specie)
+
+            ia = InteractiveCohpPlotter()
+            for label, cohp in zip(labels, cohps):
+                if label is not None:
+                    ia.add_cohp(namecation + str(ication + 1) + ": " + label, cohp)
+            plot = ia.get_plot(integrated=integrated, xlim=xlim, ylim=ylim)
+
+            plot.update_layout(title_text=title)
+            if save_as_html:
+                if len(set_inequivalent_cations) > 1:
+                    if isinstance(filename, str):
+                        filename = Path(filename)
+                    filename_new = (
+                        filename.parent / f"{filename.stem}-{iplot}{filename.suffix}"
+                    )
+                else:
+                    filename_new = filename
+                plot.write_html(filename_new)
+            plot.show()
+
     @staticmethod
     def _coordination_environment_to_text(ce):
         """

lobsterpy/plotting/__init__.py

@@ -228,26 +228,11 @@ def _broaden(energies: np.ndarray, population: np.ndarray, sigma=None, cutoff=4.
         return convolve(population, kernel, mode="same") / kernel.sum()
 
 
-class InteractiveCohpPlotter:
+class InteractiveCohpPlotter(CohpPlotter):
     """
     Interactive COHP plotter to view all relevant / multiple COHPs in one figure
     """
 
-    def __init__(self, zero_at_efermi=True, are_coops=False, are_cobis=False):
-        """
-        Args:
-            zero_at_efermi: Whether to shift all populations to have zero
-                energy at the Fermi level. Defaults to True.
-            are_coops: Switch to indicate that these are COOPs, not COHPs.
-                Defaults to False for COHPs.
-            are_cobis: Switch to indicate that these are COBIs, not COHPs/COOPs.
-                Defaults to False for COHPs
-        """
-        self.zero_at_efermi = zero_at_efermi
-        self.are_coops = are_coops
-        self.are_cobis = are_cobis
-        self._cohps = {}
-
     def add_all_relevant_cohps(self, complete_cohp, analyse):
         """
         Adds all relevant COHPs from lobsterpy analyse object.
@@ -355,13 +340,12 @@ def get_plot(
         else:
             energy_label = "$E$ (eV)"
 
-
         # Setting up repeating color scheme (same as for matplotlib plots in .mplstyle)
         palette = ['#e41a1c', '#377eb8', '#4daf4a', '#984ea3', '#ff7f00', '#ffff33', '#a65628', '#f781bf', '#999999']
         pal_iter = cycle(palette)
 
         traces = []
-        for item in self._cohps.values():
+        for label, item in self._cohps.items():
             population_key = item["ICOHP"] if integrated else item["COHP"]
             band_color = next(pal_iter)
             for spin in [Spin.up, Spin.down]:
@@ -370,7 +354,7 @@ def get_plot(
                     x = population if invert_axes else item["energies"]
                     y = item["energies"] if invert_axes else population
                     if spin == Spin.up:
-                        trace = go.Scatter(x=x, y=y, name=item["plot_label"])
+                        trace = go.Scatter(x=x, y=y, name=label)
                         trace.update(ld.spin_up_trace_style_dict)
                     else:
                         trace = go.Scatter(x=x, y=y, name="")
@@ -399,7 +383,7 @@ def get_plot(
         #TODO:
         # improve display of legend
         # somehow y axis scaling inside image?
-        # inherit CohpPlotter fr. pymatgen?
+        # add sigma arg
         # maybe dashed line at Ef
 
         return fig