Stacked histograms for spectral plots

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

@@ -364,6 +364,65 @@ def plot_spectra(ss: SampleSet, in_energy: bool = False,
  return fig, ax
 
 
+@save_or_show_plot
+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
 def plot_learning_curve(train_loss: list, validation_loss: list,
  xscale: str = "linear", yscale: str = "linear",