Make plotting of spectra more robust

Also works with mega-complexes
glotaran · Mar 28, 2021 · c46daad · c46daad
1 parent 33632b8
commit c46daad
Showing 1 changed file with 33 additions and 16 deletions.
diff --git a/pyglotaran_extras/plotting/plot_spectra.py b/pyglotaran_extras/plotting/plot_spectra.py
@@ -12,28 +12,45 @@ def plot_spectra(res, axes):
 
 
 def plot_sas(res, ax, title="SAS"):
-    sas = res.species_associated_spectra
-    sas.plot.line(x="spectral", ax=ax)
-    ax.set_title(title)
-    ax.get_legend().remove()
+    keys = [v for v in res.data_vars if v.startswith("species_associated_spectra")]
+    for key in keys:
+        sas = res[key]
+        sas.plot.line(x="spectral", ax=ax)
+        ax.set_title(title)
+        ax.get_legend().remove()
+        plt.draw()
+        plt.pause(0.001)
 
 
 def plot_norm_sas(res, ax, title="norm SAS"):
-    sas = res.species_associated_spectra
-    (sas / np.abs(sas).max(dim="spectral")).plot.line(x="spectral", ax=ax)
-    ax.set_title(title)
-    ax.get_legend().remove()
+    keys = [v for v in res.data_vars if v.startswith("species_associated_spectra")]
+    for key in keys:
+        sas = res[key]
+        sas = res.species_associated_spectra
+        (sas / np.abs(sas).max(dim="spectral")).plot.line(x="spectral", ax=ax)
+        ax.set_title(title)
+        ax.get_legend().remove()
+        plt.draw()
+        plt.pause(0.001)
 
 
 def plot_das(res, ax, title="DAS"):
-    das = res.decay_associated_spectra
-    das.plot.line(x="spectral", ax=ax)
-    ax.set_title(title)
-    ax.get_legend().remove()
+    keys = [v for v in res.data_vars if v.startswith("decay_associated_spectra")]
+    for key in keys:
+        das = res[key]
+        das.plot.line(x="spectral", ax=ax)
+        ax.set_title(title)
+        ax.get_legend().remove()
+        plt.draw()
+        plt.pause(0.001)
 
 
 def plot_norm_das(res, ax, title="norm DAS"):
-    das = res.decay_associated_spectra
-    (das / np.abs(das).max(dim="spectral")).plot.line(x="spectral", ax=ax)
-    ax.set_title(title)
-    ax.get_legend().remove()
+    keys = [v for v in res.data_vars if v.startswith("decay_associated_spectra")]
+    for key in keys:
+        das = res[key]
+        (das / np.abs(das).max(dim="spectral")).plot.line(x="spectral", ax=ax)
+        ax.set_title(title)
+        ax.get_legend().remove()
+        plt.draw()
+        plt.pause(0.001)