예제 #1
0
def analyze_region(spectrogram: Spectrogram, roi: dict):
    """

    """
    time, velocity, intensity = spectrogram.slice(roi['time'], roi['velocity'])
    # Do we want to convert to power?
    power = spectrogram.power(intensity)
    # let's look at the peak power at each time
    peaks = np.max(power, axis=0)
    speaks = sorted(peaks)
    threshold = speaks[int(len(peaks) * 0.04)]
    amax = np.argmax(power, axis=0)
    power[power < threshold] = 0.0

    span = 30
    nmax = len(velocity) - 1
    times, centers, widths, amps = [], [], [], []

    def peg(x):
        if x < 0:
            return 0
        if x > nmax:
            return nmax
        return x

    for t in range(len(time) - 1):
        if power[amax[t], t] > threshold:
            acenter = amax[t]
            vfrom, vto = peg(acenter - span), peg(acenter + span)
            gus = Gaussian(velocity[vfrom:vto], power[vfrom:vto, t])
            if gus.valid:
                times.append(time[t])
                centers.append(gus.center)
                widths.append(gus.width)
                amps.append(gus.amplitude)
            else:
                print(f"[{t}] {gus.error}")

    #power[power>=threshold] = 1
    plt.pcolormesh(time * 1e6, velocity, power)
    plt.plot(np.array(times) * 1e6, centers, 'r.', alpha=0.5)
    plt.show()
예제 #2
0
 def __init__(self, spec: Spectrogram, time_chunk: int = 16, velocity_chunk: int = 16):
     self.spectrogram = spec
     self.time_chunk = time_chunk
     self.velocity_chunk = velocity_chunk
     powers = spec.power(spec.intensity)
     nrows, ncols = powers.shape
     nrows = nrows // velocity_chunk
     ncols = ncols // time_chunk
     self.intensity = ar = np.zeros((nrows, ncols))
     rows = np.linspace(0, nrows * velocity_chunk,
                        num=nrows + 1, dtype=np.uint32)
     cols = np.linspace(0, ncols * time_chunk,
                        num=ncols + 1, dtype=np.uint32)
     for row in range(nrows):
         rfrom, rto = rows[row], rows[row + 1]
         for col in range(ncols):
             ar[row, col] = np.mean(
                 powers[rfrom:rto, cols[col]:cols[col + 1]])
     self.time = spec.time[0:len(spec.time):time_chunk]
     self.velocity = spec.velocity[0:len(spec.velocity):velocity_chunk]
     vals = ar.flatten()
     vals.sort()
     self.spower = vals
     self.threshold = vals[int(len(vals) * 0.8)]
class SpectrogramWidget:
    """
    A Jupyter notebook widget to represent a spectrogram, along with
    numerous controls to adjust the appearance of the spectrogram.
    For the widget to behave in a Jupyter notebook, place::

        %matplotlib widget

    at the top of the notebook. This requires that the package
    ipympl is installed, which can be done either with pip3
    or conda install ipympl.

    I also recommend editing ~/.jupyter/custom/custom.css to
    modify the definition of .container::

        .container {
            width: 100% !important;
            margin-right: 40px;
            margin-left: 40px;
            }

    **Inputs**

    - digfile: either a string or DigFile
    - kwargs: optional keyword arguments. These are passed to
      the Spectrogram constructor and to the routine that
      creates the control widgets.

    **Data members**

    - digfile: the source data DigFile
    - title: the title displayed above the spectrogram
    - baselines: a list of baseline velocities
    - spectrogram: a Spectrogram object deriving from digfile
    - fig:
    - axSpectrum:
    - axSpectrogram:
    - image:
    - colorbar:
    - individual_controls: dictionary of widgets
    - controls:
    """
    _gspec = {
        'width_ratios': [6, 1.25],
        'height_ratios': [1],
        'wspace': 0.05,
        'left': 0.075,
        'right': 0.975,
    }

    def __init__(self, *args, **kwargs):
        """
        If one passes in a single unnamed arg, it can either be a digfile,
        a string pointing to a digfile, or a two-dimensional ndarray.
        If we are founded on a dig file, it is possible to recompute
        things. Only a subset of operations are possible when we're
        based on a two-dimensional array, but perhaps that is sometimes
        desirable.
        """
        self.digfile = None
        if len(args) == 1:
            arg = args[0]
            if isinstance(arg, str):
                self.digfile = DigFile(arg)
            elif isinstance(arg, DigFile):
                self.digfile = arg
        if self.digfile == None and len(args):
            # Let's see if we have enough information to display a spectrogram
            # That means we have a two-dimensional ndarray, and possibly corresponding
            # time and velocity arrays. The signature would be (intensity, [times, velocities]).
            arg = args[0]
            if isinstance(arg, np.ndarray) and len(arg.shape) == 2:
                self._static = {
                    'intensity': arg,
                }

                if len(args) == 3:
                    self._static['time'] = np.asarray(args[1])
                    self._static['velocity'] = np.asarray(args[2])
                else:
                    self._static['time'] = np.arange(0, -1 + arg.shape[1])
                    self._static['velocity'] = np.arange(0, -1 + arg.shape[0])

            assert hasattr(
                self, '_static'
            ), "Inappropriate arguments passed to the SpectrogramWidget constructor"

        # If LaTeX is enabled in matplotlib, underscores in the title
        # cause problems in displaying the histogram. However, we can
        # solve this by not using latex in displaying the title, so there
        # is no need to alter characters here.
        self.title = "" if self.static else self.digfile.filename.split(
            '/')[-1]
        self.baselines = []

        # Compute the base spectrogram (do we really need this?)
        self.spectrogram = None
        if self.dig:
            self.spectrogram = Spectrogram(self.digfile, None, None, **kwargs)
            self.spectrogram_fresh = True  # flag for the first pass

            self.spectrogram.overlap = 0.875

        self.fig, axes = plt.subplots(nrows=1,
                                      ncols=2,
                                      sharey=True,
                                      squeeze=True,
                                      gridspec_kw=self._gspec)
        self.axSpectrogram, self.axSpectrum = axes

        self.subfig = None
        self.axTrack = None
        self.axSpare = None

        # At the moment, clicking on the image updates the spectrum
        # shown on the left axes. It would be nice to be more sophisticated and
        # allow for more sophisticated interactions, including the ability
        # to display more than one spectrum.
        self.fig.canvas.mpl_connect('button_press_event',
                                    lambda x: self.handle_click(x))
        self.fig.canvas.mpl_connect('key_press_event',
                                    lambda x: self.handle_key(x))

        self.spectrum(None, "")

        self.image = None  # we will set in update_spectrogram
        self.colorbar = None  # we will set this on updating, based on the

        self.peak_followers = []  # will hold any PeakFollowers
        self.spectra = []  # will hold spectra displayed at right
        self.spectra_in_db = True  # should spectra be displayed in db?

        self.controls = dict()  # widgets stored by name
        self.layout = None  # how the controls get laid out
        self.selecting = False  # we are not currently selecting a ROI
        self.roi = []  # and we have no regions of interest
        self.threshold = None
        self.make_controls(**kwargs)

        display(self.layout)
        self.update_spectrogram()

    @property
    def static(self):
        "If we are not associated with a dig file, return True"
        return hasattr(self, '_static')

    @property
    def dig(self):
        "True if we are associated with a dig file and can recompute the spectrogram"
        return isinstance(self.digfile, DigFile)

    @property
    def intensity(self):
        "Return the two-dimensional array of intensities"
        if self.dig:
            return self.spectrogram.intensity
        else:
            return self._static['intensity']

    def make_controls(self, **kwargs):
        """
        Create the controls for this widget and store them in self.controls.
        """
        cd = self.controls  # the dictionary of controls
        t_range = kwargs.get('t_range', (0, 100))

        # If we are associated with a dig file, include controls that
        # allow recomputation of the spectrogram
        if self.dig:
            df = self.digfile

            # FFT size  ###########################################
            # Set the size of each spectrum
            pps = self.spectrogram.points_per_spectrum
            # pps = kwargs.get('points_per_spectrum', 8192)
            val = int(np.log2(pps))
            cd['spectrum_size'] = slide = widgets.IntSlider(
                value=val, min=8, max=18, step=1, description="FFT 2^n")
            slide.continuous_update = False
            slide.observe(
                lambda x: self.overhaul(points_per_spectrum=2**x['new']),
                names="value")

            # Set the overlap percentage of successive time intervals
            cd['overlap'] = slide = widgets.FloatSlider(
                description='Overlap %',
                value=100.0 * self.spectrogram.overlap,
                min=0,
                max=100)
            slide.continuous_update = False
            slide.observe(lambda x: self.overhaul(overlap=x['new'] * 0.01),
                          names="value")

            # Thumbnail  ###########################################
            # Display a thumbnail of the raw signal
            cd['raw_signal'] = widgets.ToggleButton(value=False,
                                                    description="Show V(t)",
                                                    button_style='',
                                                    icon='check')
            cd['raw_signal'].observe(lambda b: self.show_raw_signal(b),
                                     names="value")

            # Time range ###########################################

            cd['t_range'] = slide = ValueSlider("Time (µs)",
                                                t_range, (df.t0, df.t_final),
                                                1e6,
                                                readout_format=".1f")
            slide.continuous_update = False

            # Velocity range ###########################################
            cd['velocity_range'] = slide = ValueSlider(
                "Velocity (km/s)", (0, 50), (0.0, self.spectrogram.v_max),
                1e-3,
                readout_format=".1f",
                continuous_update=False)
        else:
            d = self._static
            cd['t_range'] = ValueSlider("Time",
                                        t_range, (d['time'][0], d['time'][-1]),
                                        1e-3,
                                        readout_format=".1f",
                                        continuous_update=False)
            cd['velocity_range'] = ValueSlider("Velocity (km/s)", (0, 100),
                                               (0.0, d['velocity'][-1]),
                                               1e-3,
                                               readout_format=".1f",
                                               continuous_update=False)

        cd['t_range'].observe(lambda x: self.do_update(x), names="value")
        cd['velocity_range'].observe(lambda x: self.update_velocity_range(x),
                                     names="value")

        # Color range ###########################################

        imax = self.intensity.max()
        if self.dig:
            hl = self.spectrogram.histogram_levels
            imin = hl['tens'][3]

            # Let's figure out the range likely to produce a clear
            # image. Put the 50% point at the bottom end and the 95%
            # at the top?

            def scale(x):
                return int(100.0 * (x - imin) / (imax - imin))

            start_range = (scale(hl['tens'][8]), scale(hl['tenths'][9]))
        else:
            imin = imax - 150  # ?? this is bad and assumes db
            start_range = (40, 70)

        cd['intensity_range'] = slide = ValueSlider("Color",
                                                    start_range, (imin, imax),
                                                    multiplier=1,
                                                    readout_format=".0f",
                                                    continuous_update=False)
        slide.observe(lambda x: self.update_color_range(), names="value")

        # Threshold percentage #####################################
        cd['threshold'] = slide = widgets.FloatSlider(
            description='Noise floor %',
            value=0,
            min=0,
            max=100.0,
            continuous_update=False)
        slide.observe(lambda x: self.update_threshold(x['new']), names="value")

        # Color map selector ###########################################
        the_maps = sorted(COLORMAPS.keys())
        the_maps.append('Computed')
        cd['color_map'] = widgets.Dropdown(
            options=the_maps,
            value=DEFMAP,
            description='Color Map',
            disabled=False,
        )
        cd['color_map'].observe(lambda x: self.update_cmap(), names="value")

        # Click selector  ###########################################
        # What to do when registering a click in the spectrogram
        cd['clicker'] = widgets.Select(options=("Spectrum (dB)", "Spectrum",
                                                "Peak", "Gauss",
                                                "Template_Matching"),
                                       value='Spectrum (dB)',
                                       description="Click:",
                                       disabled=False)

        cd['marquee'] = mwidgets.RectangleSelector(
            self.axSpectrogram,
            lambda eclick, erelease: self.RSelect(eclick, erelease),
            interactive=True,
            useblit=True,
            rectprops=dict(facecolor='yellow',
                           edgecolor='red',
                           alpha=0.2,
                           fill=True),
            drawtype='box')
        cd['marquee'].set_active(False)

        # Clear spectra ###########################################
        cd['clear_spectra'] = widgets.Button(description="Clear Spectra")
        cd['clear_spectra'].on_click(lambda b: self.clear_spectra())

        # Clear peak_followers ###########################################
        cd['clear_followers'] = widgets.Button(
            description="Clear Peak Followers")
        cd['clear_followers'].on_click(lambda b: self.clear_followers())

        # Computing baselines ###########################################
        cd['baselines'] = widgets.Dropdown(options=('_None_', 'Squash', 'FFT'),
                                           value='_None_',
                                           description='Baselines',
                                           disabled=False)
        cd['baselines'].observe(lambda x: self.update_baselines(x["new"]),
                                names="value")

        cd['squash'] = widgets.Button(description="Squash in Vertical")
        cd['squash'].on_click(lambda b: self.squash_vertical())

        columns = [
            'color_map;t_range;velocity_range;intensity_range;threshold',
            'clicker;clear_spectra;clear_followers',
        ]
        if self.dig:
            columns.append('spectrum_size;overlap;raw_signal;baselines;squash')

        vboxes = []
        for col in columns:
            vboxes.append(widgets.VBox([cd[x] for x in col.split(';')]))
        self.layout = widgets.HBox(vboxes)

    def range(self, var):
        "Return the range of the named control, or None if not found."
        if var in self.controls:
            return self.controls[var].range
        return None

    def RSelect(self, eclick, erelease):
        "Called when self.selecting is True and the marquee is active"
        if self.selecting:
            t0, t1 = eclick.xdata, erelease.xdata
            v0, v1 = eclick.ydata, erelease.ydata
            # make sure they are in the right order
            if t1 < t0:
                t0, t1 = t1, t0
            if v1 < v0:
                v0, v1 = v1, v0
            self.roi.append(dict(time=(t0, t1), velocity=(v0, v1)))

    def do_update(self, what):
        self.update_spectrogram()

    def show_raw_signal(self, box):
        """
        Display or remove the thumbnail of the time series data
        at the top of the spectrogram window.
        """
        if box.new:
            # display the thumbnail
            t_range = self.range('t_range')
            thumb = self.digfile.thumbnail(*t_range)
            # we have to superpose the thumbnail on the
            # existing velocity axis, so we need to rescale
            # the vertical.
            tvals = thumb['times'] * 1e6  # convert to µs
            yvals = thumb['peak_to_peak']
            ylims = self.axSpectrum.get_ylim()
            # Map the thumbnail to the top 20%
            ymax = yvals.max()
            yrange = ymax - yvals.min()
            yscale = 0.2 * (ylims[1] - ylims[0]) / yrange
            vvals = ylims[1] - yscale * (ymax - yvals)
            self.raw = self.axSpectrogram.plot(tvals, vvals, 'r-',
                                               alpha=0.5)[0]
        else:
            try:
                self.axSpectrogram.lines.remove(self.raw)
                self.raw = None
                self.fig.canvas.draw()
                self.fig.canvas.flush_events()
            except:
                pass

    def overhaul(self, **kwargs):
        """
        A parameter affecting the base spectrogram has been changed, so
        we need to recompute everything.
        """
        if self.dig:
            if self.spectrogram_fresh:
                self.spectrogram_fresh = False
            else:
                self.spectrogram.set(**kwargs)
        self.update_spectrogram()

    def update_spectrogram(self):
        """
        Recompute and display everything
        """

        intense = self.spectrogram.intensity if self.dig else self._static[
            'intensity']

        # Having recomputed the spectrum, we need to set the yrange

        # of the color map slider
        cmin = intense.min()
        cmax = intense.max()
        self.controls['intensity_range'].range = (cmin, cmax)
        self.display_spectrogram()

    def display_spectrogram(self):
        """

        """
        trange = self.range('t_range')
        vrange = self.range('velocity_range')

        if self.dig:
            # extract the requisite portions
            times, velocities, intensities = self.spectrogram.slice(
                trange, vrange)
        else:
            d = self._static
            times, velocities, intensities = d['time'], d['velocity'], d[
                'intensity']

        # if we have already displayed an image, remove it
        if self.colorbar:
            self.colorbar.remove()
            self.colorbar = None
        if self.image:
            self.image.remove()
        self.image = None

        if self.threshold:
            intensities[intensities < self.threshold] = self.threshold

        self.image = self.axSpectrogram.pcolormesh(times * 1e6, velocities,
                                                   intensities)

        self.colorbar = self.fig.colorbar(self.image,
                                          ax=self.axSpectrogram,
                                          fraction=0.08)

        self.axSpectrogram.set_title(self.title, usetex=False)
        self.axSpectrogram.set_xlabel('Time ($\mu$s)')
        self.axSpectrogram.set_xlim(*(np.array(trange) * 1e6))
        self.axSpectrogram.set_ylabel('Velocity (m/s)')
        self.update_velocity_range()
        self.update_color_range()
        self.update_cmap()

    def update_threshold(self, x):
        sg = self.spectrogram
        if int(x) == 0:
            self.threshold = None
        else:
            if x < 90:
                threshold = sg.histogram_levels['tens'][x // 10]
            elif x < 99:
                threshold = sg.histogram_levels['ones'][int(x - 90)]
            else:
                threshold = sg.histogram_levels['tenths'][int(10 * (x - 99))]
            self.threshold = self.spectrogram.transform(threshold)
        self.display_spectrogram()

    def update_cmap(self):
        """
        Update the color map used to display the spectrogram
        """
        mapname = self.controls['color_map'].value
        if mapname == 'Computed':
            from UI_Elements.generate_color_map import make_spectrogram_color_map
            mapinfo = make_spectrogram_color_map(self.spectrogram, 4, mapname)
            maprange = (mapinfo['centroids'][1], mapinfo['centroids'][-2])
            self.controls['intensity_range'].value = maprange
        self.image.set_cmap(COLORMAPS[mapname])

    def update_velocity_range(self, info=None):
        """
        Update the displayed velocity range using values obtained
        from the 'velocity_range' slider.
        """
        if info:
            old_vmin, old_vmax = info['old']
            vmin, vmax = info['new']
            if vmax > old_vmax or vmin < old_vmin:
                return self.update_spectrogram()
        vmin, vmax = self.range('velocity_range')
        self.axSpectrogram.set_ylim(vmin, vmax)
        self.axSpectrum.set_ylim(vmin, vmax)

    def update_color_range(self):
        self.image.set_clim(self.range('intensity_range'))

    def handle_click(self, event):
        try:
            # convert time to seconds
            t, v = event.xdata * 1e-6, event.ydata
        except:
            return 0
        if self.selecting:
            return 0
        # Look up what we should do with the click
        action = self.controls['clicker'].value
        # print(f"The action I am attempting to do is {action}")
        try:
            if 'Spectrum' in action:
                self.spectrum(t, action)
            elif 'Template_Matching' in action:
                ti = self.spectrogram._time_to_index(t)
                vi = self.spectrogram._velocity_to_index(v)
                self.match_templates(ti, vi)
            else:
                # print("I am handling a click that should be a peak follower")
                self.follow(t, v, action)

        except Exception as eeps:
            pass

    def handle_key(self, event):
        try:
            # convert time to seconds
            _, v = event.xdata * 1e-6, event.ydata
        except:
            pass
        char = event.key
        if char == 'x':
            # remove the all spectra
            self.clear_spectra()
        if char in ('f', 'b', 'F', 'B'):
            # We'd like to go exploring
            if not hasattr(self, 'explorer_mark'):
                self.explorer_mark = 2
            else:
                shifts = dict(f=4, F=20, b=-4, B=-40)
                self.explorer_mark += shifts[char]

            self.gaussian_explorer(self.explorer_mark)
        if char in ('m', 'M'):
            self.selecting = not self.selecting
            self.controls['marquee'].set_active(self.selecting)
        if char in "0123456789":
            n = int(char)
            # self.fan_out(int(char))
            self.gauss_out(n)
        if char in ('a', 'A') and self.roi:
            self.analyze_roi()

    def clear_spectra(self):
        """Remove all spectra from axSpectrum and the corresponding
        markers from axSpectrogram
        """
        for x in self.spectra:
            self.axSpectrogram.lines.remove(x['marker'])
            self.axSpectrum.lines.remove(x['line'])
        self.spectra = []
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()

    def clear_followers(self):
        """Remove all followers"""
        for x in self.peak_followers:
            self.axSpectrogram.lines.remove(x.line)
        self.peak_followers = []
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()

    def follow(self, t, v, action):
        """Attempt to follow the path starting with the clicked
        point."""
        print(
            f"Let's follow something starting at {t, v} using action {action}."
        )
        if action == "Gauss":
            fitter = GaussianFitter(self.spectrogram, (t, v))
            self.gauss = fitter
        elif action == "Peak":
            follower = PeakFollower(self.spectrogram, (t, v))
            # self.peak = follower
            self.peak_followers.append(follower)
            follower.run()
            tsec, v = follower.v_of_t
            follower.line = self.axSpectrogram.plot(tsec * 1e6,
                                                    v,
                                                    'r.',
                                                    alpha=0.4,
                                                    markersize=2)[0]
        # print("Create a figure and axes, then call self.gauss.show_fit(axes)")

    def gauss_out(self, n: int):
        """
        Show center velocity, width, and amplitude for gaussian
        fits to the data in follower n.
        """
        if n >= len(self.peak_followers):
            return 0
        WRITEOUT, fnum = False, 0
        pf = self.peak_followers[n]
        times, centers, widths, amps = [], [], [], []
        vind = pf.frame['vi_span'].to_numpy()
        tind = pf.frame['t_index'].to_numpy()
        sp = self.spectrogram
        for j in range(len(tind)):
            t = sp.time[tind[j]] * 1e6
            vfrom, vto = vind[j]
            powers = sp.power(sp.intensity[vfrom:vto, tind[j]])
            speeds = sp.velocity[vfrom:vto]

            gus = Gaussian(speeds, powers)
            if gus.valid:
                times.append(t)
                centers.append(gus.center)
                widths.append(gus.width)
                amps.append(gus.amplitude)
                if WRITEOUT:
                    fname = f"{os.path.splitext(self.digfile.filename)[0]}_{fnum:04d}.csv"
                    gus.write_csv(fname)
                    fnum += 1
        if WRITEOUT:
            # write out the times, too
            fname = f"{os.path.splitext(self.digfile.filename)[0]}_t.csv"
            v = np.asarray(times)
            np.savetxt(v, fname)

        fig, axes = plt.subplots(nrows=1, ncols=3, squeeze=True)
        ax1, ax2, ax3 = axes
        ax1.errorbar(times, centers, fmt='b-', yerr=widths)
        ax1.set_xlabel(r'$t$ ($\mu$s)')
        ax1.set_ylabel(r'$v$ (m/s)')

        ax2.plot(times, widths, 'r-')
        ax2.set_xlabel(r'$t$ ($\mu$s)')
        ax2.set_ylabel(r'$\delta v$ (m/s)')

        ax3.plot(times, amps, 'g-')
        ax3.set_xlabel(r'$t$ ($\mu$s)')
        ax3.set_ylabel('Amplitude')

        # Store the values for later access
        if not hasattr(self, "gauss_outs"):
            self.gauss_outs = [None for x in range(len(self.peak_followers))]
        else:
            while len(self.gauss_outs) < len(self.peak_followers):
                self.gauss_outs.append(None)
        self.gauss_outs[n] = dict(time=np.array(times),
                                  center=np.array(centers),
                                  width=np.array(widths),
                                  amplitude=np.array(amps))

    def gaussian_explorer(self, follower_pt: int):
        """
        Show center velocity, width, and amplitude for gaussian
        fits to the data in follower n.
        """
        if len(self.peak_followers) == 0:
            return 0
        pf = self.peak_followers[0]
        res = pf.results
        points = len(res['t_index'])

        def bound(x):
            return x % points

        follower_pt = bound(follower_pt)

        # We'd like to show data for this index, the previous one,
        #  and the next one, along with the gaussian fit
        hoods = [pf.hood(n=bound(x + follower_pt)) for x in (-2, -1, 0, 1, 2)]

        # Check that we have the requisite figure, and make it if we don't
        if not hasattr(self, 'explorer_fig'):
            self.explorer_fig, self.explorer_axes = plt.subplots(1,
                                                                 5,
                                                                 sharey=True)
            # also add a marker to the follower's representation on the
            # spectrogram to make it easier to see where we are
            self.explorer_marker = self.axSpectrogram.plot([], [], 'k*')[0]

        # show where we are
        tsec, v = pf.v_of_t
        self.explorer_marker.set_data([
            tsec[follower_pt] * 1e6,
        ], [
            v[follower_pt],
        ])

        min_v, max_v, max_peak = 1e10, 0, 0
        for ax, hood in zip(self.explorer_axes, hoods):
            ax.clear()
            # plot the data
            ax.plot(hood.velocity, hood.intensity, 'ko', alpha=0.5)

            # plot the background level used for the moment calculation
            bg = hood.moment['background']
            ax.plot([hood.velocity[0], hood.velocity[-1]], [bg, bg], 'r-')
            # show the center and widths from the moment calculation
            tallest = np.max(hood.intensity)
            max_peak = max(tallest, max_peak)
            ax.plot([
                hood.moment['center'] + x * hood.moment['std_err']
                for x in (-1, 0, 1)
            ], 0.5 * tallest * np.ones(3), 'r.')
            # show the gaussian
            hood.plot_gaussian(ax)
            vcenter, width = hood.peak_v, hood.moment['std_dev']
            min_v = min(min_v, vcenter - 12 * width)
            max_v = max(max_v, vcenter + 12 * width)
            # ax.set_xlim(vcenter - 12 * width, vcenter + 12 * width)
            ax.set_xlabel(f"$v$ (m/s)")
            ax.set_title(f"{hood.time*1e6:.2f}" + " $\\mu$s")
            txt = f"m = {hood.moment['center']:.1f} ± {hood.moment['std_err']:.1f}"
            txt = f"{txt}\ng = {hood.gaussian.center:.1f} ± {hood.gaussian.width:.1f}"
            ax.annotate(txt,
                        xy=(0.05, 0.95),
                        xycoords='axes fraction',
                        horizontalalignment='left',
                        verticalalignment='top')

        for ax in self.explorer_axes:
            ax.set_xlim(min_v, max_v)

        # label the common velocity axis
        ax = self.explorer_axes[0]
        ax.set_ylim(-0.05 * max_peak, 1.2 * max_peak)
        ax.set_ylabel("Intensity")

    def analyze_roi(self):
        """
        Extract the region(s) of interest and process them
        """
        for roi in self.roi:
            analyze_region(self.spectrogram, roi['time'])

    def fan_out(self, n: int):
        """Produce a zoomed in version of this trace, showing
        the neighborhood around the determined peak.
        """
        if n >= len(self.peak_followers):
            return 0
        pf = self.peak_followers[n]
        vind = pf.frame['vi_span'].to_numpy()
        tind = pf.frame['t_index'].to_numpy()
        self.subfig, axes = plt.subplots(nrows=1,
                                         ncols=2,
                                         sharey=True,
                                         squeeze=True,
                                         gridspec_kw=self._gspec)
        self.axSpare, self.axTrack = axes

        # We will create a "waterfall" of curves surrounding
        # the peaks, each offset by a bit. The x axis will
        # represent intensity, with subsequent time traces offset
        # by an amount I need to determine. The y axis
        # is velocity.

        spans = []
        vvec = self.spectrogram.velocity  # shortcut to velocity vector
        tvec = self.spectrogram.time
        ivec = self.spectrogram.intensity

        # pre-extract a bunch of one-dimensional curves
        # and be sure to convert to power
        for n in range(len(tind)):
            vfrom, vto = vind[n]
            spans.append({
                'v':
                vvec[vfrom:vto],
                'power':
                self.spectrogram.power(ivec[vfrom:vto, tind[n]]),
                't':
                tvec[tind[n]] * 1e6,
            })

        maxima = np.array([np.max(x['power']) for x in spans])
        maxpower = maxima.max()
        # Let's set the offset between times to be one tenth of
        # the maxpower
        offset = 0.025 * maxpower

        for n in reversed(list(range(len(spans)))):
            span = spans[n]
            self.axTrack.plot(span['power'] + n * offset,
                              span['v'],
                              'b-',
                              alpha=0.33)
        self.axTrack.set_ylabel('$v$')

    def squash_vertical(self):
        normy = self.spectrogram.squash(dB=False) * 2000
        self.axSpectrogram.plot(self.spectrogram.time * 1e6,
                                normy,
                                'b-',
                                alpha=0.75)

    def update_baselines(self, method):
        """
        Handle the baselines popup menu
        """
        from ProcessingAlgorithms.SignalExtraction.baselines import baselines_by_squash
        blines = []
        self.baselines = []  # remove any existing baselines
        if method == "Squash":
            peaks, sigs, heights = baselines_by_squash(self.spectrogram)
            blines.extend(peaks)

            # for n in range(len(heights)):
            # if heights[n] > 0.1:
            # blines.append(peaks[n])

        # Now show the baselines in blines or remove any
        # if blines is empty

        if not blines:
            for b in self.baselines:
                self.axSpectrum.lines.remove(b['line'])
            self.baselines = []  # remove them
        else:
            edges = (self.spectrogram.intensity.min(),
                     self.spectrogram.intensity.max())
            for v in blines:
                bline = self.axSpectrum.plot([edges[0], edges[1]], [v, v],
                                             'k-',
                                             alpha=0.4)
                self.baselines.append(dict(v=v, line=bline))

    def baseline_intensity(self):
        self.update_baselines("Squash")
        figgy = plt.figure()
        ax = figgy.add_subplot(1, 1, 1)
        sg = self.spectrogram
        for bline in self.baselines:
            index = sg._velocity_to_index(bline['v'])
            if index > 2:
                ax.semilogy(sg.time * 1e6, sg.power(sg.intensity[index, :]))
        ax.set_xlabel(r"$t$ ($\mu$ s)")
        ax.set_ylabel(r"Power")

    def Spectrum(self, the_time: float):
        """
        Return the column from the spectrogram in power form
        """
        sg = self.spectrogram
        t_index = sg._time_to_index(the_time)
        vals = sg.intensity[:, t_index]
        return sg.power(vals)

    def spectrum(self, the_time: float, form: str):
        """
        Display a spectrum in the left axes corresponding to the
        passed value of the_time (which is in seconds).
        """
        _colors = ["r", "g", "b", "y"]
        if the_time is None:
            # Initialize the axes
            # self.axSpectrum.plot([0, 1], [0, 1], 'r-')
            self.axSpectrum.grid(axis='x', which='both', color='b', alpha=0.4)
        else:
            if True:
                delta_t = self.spectrogram.points_per_spectrum / 2 * \
                    self.digfile.dt
                the_spectrum = Spectrum(self.digfile.values(
                    the_time - delta_t, the_time + delta_t),
                                        self.digfile.dt,
                                        remove_dc=True)
                # compute the level of the 90th percentile
                spec = dict(spectrum=the_spectrum)
                vals = the_spectrum.db
                ordering = np.argsort(vals)
                if self.baselines:
                    blines = [x['v'] for x in self.baselines]
                    n = -1
                    while the_spectrum.velocities[ordering[n]] in blines:
                        n -= 1
                else:
                    n = -1
                spec['max'] = vals[ordering[n]]
                noise_floor = int(n - 0.1 * len(vals))
                spec['90'] = vals[ordering[noise_floor]]
            else:
                t_index = self.spectrogram._time_to_index(the_time)
                vals = self.spectrogram.intensity[:, t_index]

            # We need to worry about the format of the spectrum
            db = ('dB' in form)
            field = 'db' if db else 'power'
            the_line = self.axSpectrum.plot(getattr(the_spectrum, field),
                                            the_spectrum.velocities,
                                            _colors[len(self.spectra)],
                                            alpha=0.33)
            spec['line'] = the_line[0]

            tval = the_time * 1e6  # convert to microseconds
            marker = self.axSpectrogram.plot([tval, tval],
                                             [0, self.spectrogram.v_max],
                                             _colors[len(self.spectra)],
                                             alpha=0.33)
            spec['marker'] = marker[0]

            self.spectra.append(spec)

            if db != self.spectra_in_db:
                self.spectra_in_db = db  # switch our mode
                # and replot all the spectra
                for spec in self.spectra:
                    li = spec['line']
                    sp = spec['spectrum']
                    li.set(xdata=getattr(sp, field), ydata=sp.velocities)

            self.axSpectrum.set_xlabel("Power (dB)" if db else "Power")
            if db:
                # we should order the values and set a limit at something
                # like the strongest decile
                ninety = max([x['90'] for x in self.spectra])
                peak = max([x['max'] for x in self.spectra])
                self.axSpectrum.set_xlim(ninety, peak)
            return 0
            line = self.axSpectrum.lines[0]
            intensities = the_spectrum.db
            line.set(xdata=intensities, ydata=the_spectrum.velocities)

            # We should also add a line to the spectrogram showing where
            # the spectrum came from.
            if not self.axSpectrogram.lines:
                self.axSpectrogram.plot([0, 0], [0, 1], 'r-', alpha=0.33)
            # this won't scale when we add baselines
            line = self.axSpectrogram.lines[0]

            line.set(xdata=[tval, tval], ydata=[0, self.spectrogram.v_max])

    def match_templates(self, time, velocity):

        methodsToUse = [
            'cv2.TM_CCORR_NORMED', 'cv2.TM_SQDIFF', 'cv2.TM_SQDIFF_NORMED'
        ]

        matcher = TemplateMatcher(
            self.spectrogram,
            template=Templates.opencv_long_start_pattern5.value,
            span=200,
            k=20,
            methods=methodsToUse)

        times, velos, scores, methodsUsed = matcher.match()

        time_offset = 2.0 + (2 * abs(
            matcher.spectrogram.time[matcher.template_time_offset_index] * 1e6)
                             )
        times = np.array(times) + time_offset

        matcher.add_to_plot(self.axSpectrogram,
                            times,
                            velos,
                            scores,
                            methodsUsed,
                            show_points=True,
                            show_medoids=True,
                            verbose=False,
                            visualize_opacity=False,
                            show_bounds=True)
예제 #4
0
def baselines_by_squash(
    spectrogram: Spectrogram,
    segments = 8,
    min_height = 0.0005,  # minimum fraction of the tallest peak for consideration
    min_separation = 400,  # minimum separation of peaks in m/s
    peak_width = 20,
    min_percent = 80,
    max_side_lobe = 0.5
):
    """
    Return a list of baseline velocities and their uncertainties.
    Inputs:
      -  spectrogram: an instance of Spectrogram
    Optional inputs:
      -  segments: an integer specifying the number of segments into which to
         divide the time dimension to look for a consistent peak
      -  min_height: the minimum fraction of the strongest squashed peak
         to consider when looking for baselines
      -  min_separation: the minimum separation of baselines, in m/s
      -  peak_width: to be considered a peak, look at points at this remove
         (in m/s) from a candidate peak to see that their values are lower than
         max_side_lobe times the peak's value
      -  min_percent: minimum percentage of segments satisfying the peak criterion
         described above under peak_width
      -  max_side_lobe: the segment's squashed value at ± peak_width must be
         less than or equal to max_side_lobe times the peak's value
    Outputs:
      -  peaks: an array of peaks, in descending order of strength
      -  widths: an array of uncertainty values for the peaks
      -  heights: corresponding peak heights, normalized to
            the greatest height
    """

    # Collapse along the time axis, making sure to use power,
    # not dB
    dv = spectrogram.velocity[1] - spectrogram.velocity[0]

    powers = spectrogram.power(spectrogram.intensity)

    # We'd like to squash into about 8 distinct segments
    # along axis 1 and then make sure that we get consistent
    # values before calling a peak a baseline

    boundaries = list(range(0, len(spectrogram.time),
                            len(spectrogram.time) // segments))
    boundaries.append(len(spectrogram.time))
    squash = []
    for n in range(segments):
        squash.append(
            powers[:, boundaries[n]:boundaries[n + 1]].mean(axis=1))

    combined_spectrum = spectrogram.power(
        spectrogram.intensity).mean(axis=1)
    tallest = combined_spectrum.max()
    # this should be improved
    peaks, properties = find_peaks(
        combined_spectrum,
        height=min_height * tallest,  # peaks must be this tall to count
        distance=min_separation / dv  # peaks must be separated by this much at minimum
    )
    heights = properties['peak_heights']
    peak_ht = heights.max()
    # produce an ordering from tallest to shortest peak
    ordering = np.flip(np.argsort(heights))
    peak_positions = peaks[ordering]
    peaks = spectrogram.velocity[peak_positions]
    heights = heights[ordering] / peak_ht

    # Now we want to filter out the peaks that don't correspond to
    # consistent signal across the segments

    strong_peaks = []
    delta = int(peak_width / dv)
    for n, pk_pos in enumerate(peak_positions):
        if pk_pos >= delta:
            totes = 0
            for sq in squash:
                try:
                    if sq[pk_pos + delta] < max_side_lobe * sq[pk_pos] > sq[pk_pos - delta]:
                        totes += 1
                except IndexError:
                    pass
            strong_peaks.append(totes >= segments * min_percent * 0.01)
            # print(f"[{pk_pos}] -> {spectrogram.velocity[pk_pos]} ({heights[n]})) got {totes}")
        else:
            strong_peaks.append(False)
    keepers = np.array(strong_peaks)

    return peaks[keepers], (np.ones(len(peaks)) * dv)[keepers], heights[keepers]