ENH: add spectral demo

Author: tacaswell

spectral/germ.h5

@@ -0,0 +1,76 @@
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from matplotlib.widgets import SpanSelector
+from matplotlib.colors import LogNorm
+import numpy as np
+import h5py
+
+def plot_all_chan_spectrum(spectrum, bins, *, ax=None, **kwargs):
+
+    def integrate_to_angles(spectrum, bins, lo, hi):
+        lo_ind, hi_ind = bins.searchsorted([lo, hi])
+        return spectrum[lo_ind:hi_ind].sum(axis=0)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(13.5, 9.5))
+    else:
+        fig = ax.figure
+
+    div = make_axes_locatable(ax)
+    ax_r = div.append_axes('right', 2, pad=0.1, sharey=ax)
+    ax_t = div.append_axes('top', 2, pad=0.1, sharex=ax)
+
+    ax_r.yaxis.tick_right()
+    ax_r.yaxis.set_label_position("right")
+    ax_t.xaxis.tick_top()
+    ax_t.xaxis.set_label_position("top")
+
+    im = ax.imshow(spectrum, origin='lower', aspect='auto',
+                   extent=(-.5, 383.5,
+                           bins[0], bins[-1]),
+                   norm=LogNorm())
+
+    e_line, = ax_r.plot(spectrum.sum(axis=1), bins[:-1] + np.diff(bins))
+    p_line, = ax_t.plot(spectrum.sum(axis=0))
+    label = ax_t.annotate('[0, 70] kEv', (0, 1), (10, -10),
+                          xycoords='axes fraction',
+                          textcoords='offset pixels',
+                          va='top', ha='left')
+
+    def update(lo, hi):
+        p_data = integrate_to_angles(spectrum, bins, lo, hi)
+        p_line.set_ydata(p_data)
+        ax_t.relim()
+        ax_t.autoscale(axis='y')
+
+        label.set_text(f'[{lo:.1f}, {hi:.1f}] keV')
+        fig.canvas.draw_idle()
+
+    span = SpanSelector(ax_r, update, 'vertical', useblit=True,
+                        rectprops={'alpha': .5, 'facecolor': 'red'},
+                        span_stays=True)
+
+    ax.set_xlabel('channel [#]')
+    ax.set_ylabel('E [keV]')
+
+    ax_t.set_xlabel('channel [#]')
+    ax_t.set_ylabel('total counts')
+
+    ax_r.set_ylabel('E [keV]')
+    ax_r.set_xlabel('total counts')
+    ax.set_xlim(-.5, 383.5)
+    ax.set_ylim(bins[0], bins[-1])
+    ax_r.set_xlim(xmin=0)
+
+    return spectrum, bins, {'center': {'ax': ax, 'im': im},
+                            'top': {'ax': ax_t, 'p_line': p_line},
+                            'right': {'ax': ax_r, 'e_line': e_line,
+                                      'span': span}}
+
+
+with h5py.File('germ.h5', 'r') as fin:
+    spectrum = fin['spectrum'][:]
+    bins = fin['bins'][:]
+
+plot_all_chan_spectrum(spectrum, bins)
+plt.show()