Source code for gdt.core.plot.spectrum

# CONTAINS TECHNICAL DATA/COMPUTER SOFTWARE DELIVERED TO THE U.S. GOVERNMENT WITH UNLIMITED RIGHTS
#
# Contract No.: CA 80MSFC17M0022
# Contractor Name: Universities Space Research Association
# Contractor Address: 7178 Columbia Gateway Drive, Columbia, MD 21046
#
# Copyright 2017-2022 by Universities Space Research Association (USRA). All rights reserved.
#
# Developed by: William Cleveland and Adam Goldstein
#               Universities Space Research Association
#               Science and Technology Institute
#               https://sti.usra.edu
#
# Developed by: Daniel Kocevski
#               National Aeronautics and Space Administration (NASA)
#               Marshall Space Flight Center
#               Astrophysics Branch (ST-12)
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software distributed under the License
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing permissions and limitations under the
# License.
#
import numpy as np
from ..data_primitives import ChannelBins
from .defaults import *
from .lib import *
from .plot import GdtPlot, Histo, HistoErrorbars, HistoFilled, \
    SpectrumBackground

__all__ = ['Spectrum']

[docs]class Spectrum(GdtPlot): """Class for plotting count spectra and count spectra paraphernalia. This class can plot differential count spectra using energy information, or can plot a count spectrum based on raw energy channel. Parameters: data (:class:`~gdt.core.data_primitives.EnergyBins`, optional): The count spectrum data to plot background (:class:`~gdt.background.primitives.BackgroundSpectrum`, optional): The background spectrum to plot **kwargs: Options to pass to :class:`~gdt.core.plot.plot.GdtPlot` """ def __init__(self, data=None, background=None, canvas=None, **kwargs): super().__init__(canvas=canvas, **kwargs) self._spec = None self._errorbars = None self._bkgd = None self._selections = [] # initialize the plot axes, labels, ticks, and scales self._ax.set_xlabel('Energy (keV)', fontsize=PLOTFONTSIZE) self._ax.set_ylabel('Rate (count/s-keV)', fontsize=PLOTFONTSIZE) self._ax.xaxis.set_tick_params(labelsize=PLOTFONTSIZE) self._ax.yaxis.set_tick_params(labelsize=PLOTFONTSIZE) self._ax.set_xscale('log') self._ax.set_yscale('log') # plot data and/or background if set on init if data is not None: self.set_data(data) if background is not None: self.set_background(background) @property def background(self): """(:class:`~gdt.core.plot.plot.SpectrumBackground`): The count spectrum background plot element""" return self._bkgd @property def errorbars(self): """(:class:`~gdt.core.plot.plot.HistoErrorbars`): The error bars plot element """ return self._errorbars @property def selections(self): """(list of :class:`~gdt.core.plot.plot.HistoFilled`): The count spectrum selection plot element""" return self._selections @property def spectrum(self): """(:class:`~gdt.core.plot.plot.Histo`): The count spectrum plot element""" return self._spec # mark FIXME: store selections in a collection
[docs] def add_selection(self, data): """Add a selection to the plot. This adds a new selection to a list of existing selections. Args: data (:class:`~gdt.core.data_primitives.EnergyBins`): The count spectrum data selections to plot """ color, alpha, fill_alpha, kwargs = self._selection_settings() select = HistoFilled(data, self._ax, color=color, alpha=alpha, fill_alpha=fill_alpha, **kwargs) self._selections.append(select)
[docs] def set_background(self, background): """Set the background plotting data. If a background already exists, this triggers a replot of the background. Args: background (:class:`~gdt.background.primitives.BackgroundSpectrum`): The background spectrum to plot """ color, alpha, band_alpha, kwargs = self._bkgd_settings() self._bkgd = SpectrumBackground(background, self._ax, color=color, alpha=alpha, band_alpha=band_alpha, zorder=kwargs.pop("zorder", 1000), **kwargs)
[docs] def set_data(self, data): """Set the count spectrum plotting data. If a count spectrum already exists, this triggers a replot of the count spectrum. If an EnergyBins object is used, this will plot a differential energy spectrum (count/s/keV), and if a ChannelBins object is used (i.e. no energy information), the count spectrum per channel is plotted. Args: data (:class:`~gdt.core.data_primitives.EnergyBins` or :class:`~gdt.core.data_primitives.ChannelBins`): The data """ spec_color, spec_alpha, spec_kwargs = self._spec_settings() self._spec = Histo(data, self._ax, color=spec_color, alpha=spec_alpha, **spec_kwargs) eb_color, eb_alpha, eb_kwargs = self._eb_settings() self._errorbars = HistoErrorbars(data, self._ax, color=eb_color, alpha=eb_alpha, **eb_kwargs) mask = (data.rates > 0.0) if isinstance(data, ChannelBins): self._ax.set_xlabel('Channel Number', fontsize=PLOTFONTSIZE) self._ax.set_ylabel('Rate (count/s)', fontsize=PLOTFONTSIZE) self._ax.set_xscale('linear') self._ax.set_ylim(0.9 * data.rates[mask].min(), 1.1 * data.rates.max()) else: self._ax.set_ylim(0.9 * data.rates_per_kev[mask].min(), 1.1 * data.rates_per_kev.max()) self._ax.set_xlim(data.range)
[docs] def remove_background(self): """Remove the background from the plot. """ self._bkgd.remove() self._bkgd = None
[docs] def remove_data(self): """Remove the count spectrum from the plot. """ self._spec.remove() self._spec = None
[docs] def remove_errorbars(self): """Remove the count spectrum errorbars from the plot. """ self._errorbars.remove() self._errorbars = None
[docs] def remove_selections(self): """Remove the selections from the plot. """ [selection.remove() for selection in self._selections] self._selections = []
def _bkgd_settings(self): """The default settings for the background. If a background already exists, use its settings instead. """ if self._bkgd is None: color = BKGD_COLOR alpha = BKGD_ALPHA band_alpha = BKGD_ERROR_ALPHA kwargs = {'linewidth': BKGD_WIDTH} else: color = self._bkgd.color alpha = self._bkgd.alpha band_alpha = self._bkgd.band_alpha kwargs = self._bkgd._kwargs return color, alpha, band_alpha, kwargs def _eb_settings(self): """The default settings for the errorbars. If a lightcurve already exists, use its errorbars settings instead. """ if self._errorbars is None: eb_color = DATA_ERROR_COLOR eb_alpha = DATA_ERROR_ALPHA eb_kwargs = {} else: eb_color = self._errorbars.color eb_alpha = self._errorbars.alpha eb_kwargs = self._errorbars._kwargs return (eb_color, eb_alpha, eb_kwargs) def _selection_settings(self): """The default settings for a selection. If a selection already exists, use its settings instead. """ if len(self._selections) == 0: color = DATA_SELECTED_COLOR alpha = 1.0 fill_alpha = DATA_SELECTED_ALPHA kwargs = {} else: color = self._selections[0].color alpha = self._selections[0].alpha fill_alpha = self._selections[0].fill_alpha kwargs = self._selections[0]._kwargs return color, alpha, fill_alpha, kwargs def _spec_settings(self): """The default settings for the count spectrum. If a count spectrum already exists, use its settings instead. """ if self._spec is None: spec_color = DATA_COLOR spec_alpha = None spec_kwargs = {} else: spec_color = self._spec.color spec_alpha = self._spec.alpha spec_kwargs = self._spec._kwargs return (spec_color, spec_alpha, spec_kwargs)