示例#1
0
        def wrapped_f(context: strax.Context, run_id: str, **kwargs):
            # Validate arguments
            known_kwargs = (
                'time_range seconds_range time_within time_selection '
                'ignore_time_warning '
                'selection_str t_reference to_pe config').split()
            for k in kwargs:
                if k not in known_kwargs and k not in parameters:
                    # Python itself also raises TypeError for invalid kwargs
                    raise TypeError(f"Unknown argument {k} for {f.__name__}")

            if 'config' in kwargs:
                context = context.new_context(config=kwargs['config'])
            if 'config' in parameters:
                kwargs['config'] = context.config

            # Say magic words to enable holoviews
            if hv_bokeh:
                # Generally using globals is not great, but it would be
                # the same as doing a slow import on the top of this file
                # pylint: disable=global-statement
                global _hv_bokeh_initialized
                if not _hv_bokeh_initialized:
                    import holoviews
                    holoviews.extension('bokeh')
                    _hv_bokeh_initialized = True

            if 'to_pe' in parameters and 'to_pe' not in kwargs:
                kwargs['to_pe'] = straxen.get_correction_from_cmt(
                    run_id, context.config['gain_model'])

            # Prepare selection arguments
            kwargs['time_range'] = context.to_absolute_time_range(
                run_id,
                targets=requires,
                **{
                    k: kwargs.get(k)
                    for k in ('time_range seconds_range time_within'.split())
                })
            kwargs.setdefault('time_selection', default_time_selection)
            kwargs.setdefault('selection_str', None)

            kwargs['t_reference'], _ = context.estimate_run_start_and_end(
                run_id, requires)

            if warn_beyond_sec is not None and not kwargs.get(
                    'ignore_time_warning'):
                tr = kwargs['time_range']
                if tr is None:
                    sec_requested = float('inf')
                else:
                    sec_requested = (tr[1] - tr[0]) / int(1e9)
                if sec_requested > warn_beyond_sec:
                    tr_str = "the entire run" if tr is None else f"{sec_requested} seconds"
                    raise ValueError(
                        f"The author of this mini analysis recommends "
                        f"not requesting more than {warn_beyond_sec} seconds. "
                        f"You are requesting {tr_str}. If you wish to proceed, "
                        "pass ignore_time_warning = True.")

            # Load required data, if any
            if len(requires):
                deps_by_kind = strax.group_by_kind(requires, context=context)
                for dkind, dtypes in deps_by_kind.items():
                    if dkind in kwargs:
                        # Already have data, just apply cuts
                        kwargs[dkind] = strax.apply_selection(
                            kwargs[dkind],
                            selection_str=kwargs['selection_str'],
                            time_range=kwargs['time_range'],
                            time_selection=kwargs['time_selection'])
                    else:
                        kwargs[dkind] = context.get_array(
                            run_id,
                            dtypes,
                            selection_str=kwargs['selection_str'],
                            time_range=kwargs['time_range'],
                            time_selection=kwargs['time_selection'],
                            # Arguments for new context, if needed
                            config=kwargs.get('config'),
                            register=kwargs.get('register'),
                            storage=kwargs.get('storage', tuple()),
                            progress_bar=False,
                        )

                # If user did not give time kwargs, but the function expects
                # a time_range, try to add one based on the time range of the data
                base_dkind = list(deps_by_kind.keys())[0]
                x = kwargs[base_dkind]
                if len(x) and kwargs.get('time_range') is None:
                    x0 = x.iloc[0] if isinstance(x, pd.DataFrame) else x[0]
                    try:
                        kwargs.setdefault('time_range',
                                          (x0['time'], strax.endtime(x).max()))

                    except AttributeError:
                        # If x is a holoviews dataset, this will fail.
                        pass

            if 'seconds_range' in parameters:
                if kwargs.get('time_range') is None:
                    scr = None
                else:
                    scr = tuple([(t - kwargs['t_reference']) / int(1e9)
                                 for t in kwargs['time_range']])
                kwargs.setdefault('seconds_range', scr)

            kwargs.setdefault('run_id', run_id)
            kwargs.setdefault('context', context)

            if 'kwargs' in parameters:
                # Likely this will be passed to another mini-analysis
                to_pass = kwargs
                # Do not pass time_range and seconds_range both (unless explicitly requested)
                # strax does not like that
                if 'seconds_range' in to_pass and not 'seconds_range' in parameters:
                    del to_pass['seconds_range']
                if 'time_within' in to_pass and not 'time_within' in parameters:
                    del to_pass['time_within']
            else:
                # Pass only arguments the function wants
                to_pass = {k: v for k, v in kwargs.items() if k in parameters}
            return f(**to_pass)
示例#2
0
def plot_wf(st: strax.Context,
            containers,
            run_id, plot_log=True, plot_extension=0, hit_pattern=True,
            timestamp=True, time_fmt="%d-%b-%Y (%H:%M:%S)",
            **kwargs):
    """
    Combined waveform plot
    :param st: strax.Context
    :param containers: peaks/records/events where from we want to plot
        all the peaks that are within it's time range +- the
        plot_extension. For example, you can provide three adjacent
        peaks and plot them in a single figure.
    :param run_id: run_id of the containers
    :param plot_log: Plot the y-scale of the wf in log-space
    :param plot_extension: include this much nanoseconds around the
        containers (can be scalar or list of (-left_extension,
        right_extension).
    :param hit_pattern: include the hit-pattern in the wf
    :param timestamp: print the timestamp to the plot
    :param time_fmt: format fo the timestamp (datetime.strftime format)
    :param kwargs: kwargs for plot_peaks
    """

    if not isinstance(run_id, str):
        raise ValueError(f'Insert single run_id, not {run_id}')

    p = containers  # usually peaks
    run_start, _ = st.estimate_run_start_and_end(run_id)
    t_range = np.array([p['time'].min(), strax.endtime(p).max()])

    # Extend the time range if needed.
    if not np.iterable(plot_extension):
        t_range += np.array([-plot_extension, plot_extension])
    elif len(plot_extension) == 2:
        if not plot_extension[0] < 0:
            warnings.warn('Left extension is positive (i.e. later than start '
                          'of container).')
        t_range += plot_extension
    else:
        raise ValueError('Wrong dimensions for plot_extension. Use scalar or '
                         'object of len( ) == 2')
    t_range -= run_start
    t_range = t_range / 10 ** 9
    t_range = np.clip(t_range, 0, np.inf)


    if hit_pattern:
        plt.figure(figsize=(14, 11))
        plt.subplot(212)
    else:
        plt.figure(figsize=(14, 5))
    # Plot the wf
    plot_peaks(st, run_id, seconds_range=t_range, single_figure=False, **kwargs)

    if timestamp:
        _ax = plt.gca()
        t_stamp = datetime.datetime.fromtimestamp(
            containers['time'].min() / 10 ** 9).strftime(time_fmt)
        _ax.text(0.975, 0.925, t_stamp,
                 horizontalalignment='right',
                 verticalalignment='top',
                 transform=_ax.transAxes)
    # Select the additional two panels to show the top and bottom arrays
    if hit_pattern:
        axes = plt.subplot(221), plt.subplot(222)
        plot_hit_pattern(st, run_id,
                         seconds_range=t_range,
                         axes=axes,
                         vmin=1 if plot_log else None,
                         log_scale=plot_log,
                         label='Area per channel [PE]')