fg bg count plot working

examples/visualization/plot_fg_bg_counts.py

@@ -0,0 +1,30 @@
+# Copyright 2021 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
+# Under the terms of Contract DE-NA0003525 with NTESS,
+# the U.S. Government retains certain rights in this software.
+"""This example demonstrates how to plot gamma spectra."""
+import sys
+
+from riid.data.synthetic import get_dummy_seeds
+from riid.visualize import plot_fg_and_bg_spectra
+from riid.data.synthetic.seed import SeedMixer
+from riid.data.synthetic.static import StaticSynthesizer
+
+if len(sys.argv) == 2:
+    import matplotlib
+    matplotlib.use("Agg")
+
+fg_seeds_ss, bg_seeds_ss = get_dummy_seeds().split_fg_and_bg()
+mixed_bg_seed_ss = SeedMixer(bg_seeds_ss, mixture_size=3).generate(10)
+
+static_synth = StaticSynthesizer(
+    samples_per_seed=100,
+    snr_function="log10",
+    snr_function_args=(50, 100),
+    return_fg=True,
+    return_gross=False,
+    return_bg=True
+)
+fg_ss, bg_ss, _ = static_synth.generate(fg_seeds_ss, mixed_bg_seed_ss)
+
+plot_fg_and_bg_spectra(fg_ss, bg_ss, index=3000,
+                       xscale='log', yscale='log', xlim=(1, None))

riid/visualize.py

@@ -371,9 +371,62 @@ def plot_spectra(ss: SampleSet, in_energy: bool = False,
 
 
 @save_or_show_plot
-def plot_fg_and_bg_spectra(fg_ss, bg_ss):
-    """Plots the first foreground spectrum alongside the first background spectrum and shows the error bars per channel for the background spectrum."""
-    pass
+def plot_fg_and_bg_spectra(fg_ss, bg_ss, index=0,
+                           figsize: tuple = (12.8, 7.2),
+                           xscale: str = "linear", yscale: str = "log",
+                           xlim: tuple = (None, None), ylim: tuple = (None, None),
+                           ylabel: str = None, title: str = None, legend_loc: str = None):
+    """Plots the first foreground spectrum alongside the first background spectrum and shows the
+       error bars per channel for the background spectrum.
+    """
+
+    labels = []
+    colors = ['grey', 'steelblue']
+    bg = bg_ss.spectra.iloc[index]
+    labels.append("BG")
+    fg = fg_ss.spectra.iloc[index]
+    labels.append("FG")
+
+    bins = fg_ss.get_channel_energies(index)
+
+    fig, ax = plt.subplots(figsize=figsize)
+    ax.grid()
+    y, bin_edges, _ = ax.hist([bg, fg], bins, label=labels, color=colors,
+                              stacked=True, density=False)
+    y_bg = y[0]
+    bin_centers = .5 * (bin_edges[1:] + bin_edges[:-1])
+    ax.errorbar(bin_centers, y_bg, yerr=np.sqrt(y_bg), ecolor='black',
+                elinewidth=.5, fmt="k,", label="BG STD")
+
+    ax.set_xscale(xscale)
+    ax.set_yscale(yscale)
+    ax.set_xlim(xlim)
+    ax.set_ylim(ylim)
+
+    if xscale == "log":
+        ax.set_xlabel("log(Energy (keV))")
+    else:
+        ax.set_xlabel("Energy (keV)")
+    ax.tick_params(axis="x", rotation=45)
+
+    if ylabel:
+        ax.set_ylabel(ylabel)
+    elif yscale == "log":
+        ax.set_ylabel("log(Counts)")
+    else:
+        ax.set_ylabel("Counts")
+
+    if title:
+        ax.set_title(title)
+    else:
+        ax.set_title("FG/BG Counts")
+
+    if legend_loc:
+        ax.legend(loc=legend_loc)
+    else:
+        ax.legend()
+
+    return fig, ax
 
 
 @save_or_show_plot