6353497555

Committed 29 Sep 2023 03:10PM UTC coverage: 93.724% (+0.05%) from 93.673%

Build # 6353497555

Build Type

Pull #1240

github

Committed by

web-flow

Commit Message

Merge 6854cb9dc into 9bf6192f0

Pull Request Pull Request #1240: Proposal to use pre-commit for continuous integration

Run Details

2930 of 2930 new or added lines in 106 files covered. (100.0%)

8751 of 9337 relevant lines covered (93.72%)

1.87 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

92.03

/straxen/analyses/waveform_plot.py

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import strax
import straxen
from mpl_toolkits.axes_grid1 import inset_locator

from .daq_waveforms import group_by_daq
from .records_matrix import DEFAULT_MAX_SAMPLES

export, __all__ = strax.exporter()
__all__.extend(["plot_wf"])


@straxen.mini_analysis()
def plot_waveform(
    context,
    deep=False,
    show_largest=100,
    figsize=None,
    max_samples=DEFAULT_MAX_SAMPLES,
    ignore_max_sample_warning=True,
    cbar_loc="lower right",
    lower_panel_height=2,
    **kwargs,
):
    """Plot the sum waveform and optionally per-PMT waveforms.

    :param deep: If True, show per-PMT waveform matrix under sum     waveform. If 'raw', use
    raw_records instead of records to do so. :param show_largest: Show only the largest show_largest
    peaks. :param figsize: Matplotlib figure size for the plot Additional     options for deep =
    True or raw: :param cbar_loc: location of the intensity color bar. Set to None to     omit it
    altogether. :param lower_panel_height: Height of the lower panel in terms of the     height of
    the upper panel.

    """
    if figsize is None:
        figsize = (10, 6 if deep else 4)

    if not deep:
        context.plot_peaks(**kwargs, show_largest=show_largest, figsize=figsize)

    else:
        f, axes = plt.subplots(
            2,
            1,
            constrained_layout=True,
            figsize=figsize,
            gridspec_kw={"height_ratios": [1, lower_panel_height]},
        )

        plt.sca(axes[0])
        context.plot_peaks(**kwargs, show_largest=show_largest, single_figure=False, xaxis=False)

        plt.sca(axes[1])
        context.plot_records_matrix(
            **kwargs,
            cbar_loc=cbar_loc,
            max_samples=max_samples,
            ignore_max_sample_warning=ignore_max_sample_warning,
            raw=deep == "raw",
            single_figure=False,
        )

        plt.subplots_adjust(hspace=0)


@straxen.mini_analysis(
    requires=("peaks", "peak_basics"), default_time_selection="touching", warn_beyond_sec=60
)
def plot_peaks(
    peaks,
    seconds_range,
    t_reference,
    show_largest=100,
    single_figure=True,
    figsize=(10, 4),
    xaxis=True,
):
    if single_figure:
        plt.figure(figsize=figsize)
    plt.axhline(0, c="k", alpha=0.2)

    peaks = peaks[np.argsort(-peaks["area"])[:show_largest]]
    peaks = strax.sort_by_time(peaks)

    for p in peaks:
        plot_peak(p, t0=t_reference, color={0: "gray", 1: "b", 2: "g"}[p["type"]])

    if xaxis == "since_start":
        seconds_range_xaxis(seconds_range, t0=seconds_range[0])
    elif xaxis:
        seconds_range_xaxis(seconds_range)
        plt.xlim(*seconds_range)
    else:
        plt.xticks([])
        plt.xlim(*seconds_range)
    plt.ylabel("Intensity [PE/ns]")
    if single_figure:
        plt.tight_layout()


@straxen.mini_analysis(
    requires=("peaks", "peak_basics"), default_time_selection="touching", warn_beyond_sec=60
)
def plot_hit_pattern(
    peaks,
    seconds_range,
    t_reference,
    axes=None,
    vmin=None,
    log_scale=False,
    label=None,
    single_figure=False,
    xenon1t=False,
    figsize=(10, 4),
):
    if single_figure:
        plt.figure(figsize=figsize)
    if len(peaks) > 1:
        print(f"warning showing total area of {len(peaks)} peaks")
    straxen.plot_pmts(
        np.sum(peaks["area_per_channel"], axis=0),
        axes=axes,
        vmin=vmin,
        log_scale=log_scale,
        label=label,
        xenon1t=xenon1t,
    )


@straxen.mini_analysis()
def plot_records_matrix(
    context,
    run_id,
    seconds_range,
    cbar_loc="upper right",
    raw=False,
    single_figure=True,
    figsize=(10, 4),
    group_by=None,
    max_samples=DEFAULT_MAX_SAMPLES,
    ignore_max_sample_warning=False,
    vmin=None,
    vmax=None,
    **kwargs,
):
    if seconds_range is None:
        raise ValueError(
            "You must pass a time selection (e.g. seconds_range) to plot_records_matrix."
        )

    if single_figure:
        plt.figure(figsize=figsize, constrained_layout=True)

    f = context.raw_records_matrix if raw else context.records_matrix

    wvm, ts, ys = f(
        run_id,
        max_samples=max_samples,
        ignore_max_sample_warning=ignore_max_sample_warning,
        **kwargs,
    )
    if group_by is not None:
        ylabs, wvm_mask = group_by_daq(run_id, group_by)
        wvm = wvm[:, wvm_mask]
        plt.ylabel(group_by)
    else:
        plt.ylabel("Channel number")

    # extract min and max from kwargs or set defaults
    if vmin is None:
        vmin = min(0.1 * wvm.max(), 1e-2)
    if vmax is None:
        vmax = wvm.max()

    plt.pcolormesh(
        ts,
        ys,
        wvm.T,
        norm=matplotlib.colors.LogNorm(
            vmin=vmin,
            vmax=vmax,
        ),
        cmap=plt.cm.inferno,
    )
    plt.xlim(*seconds_range)

    ax = plt.gca()
    if group_by is not None:
        # Do some magic to convert all the labels to an integer that
        # allows for remapping of the y labels to whatever is provided
        # in the "ylabs", otherwise matplotlib shows nchannels different
        # labels in the case of strings.
        # Make a dict that converts the label to an int
        int_labels = {h: i for i, h in enumerate(set(ylabs))}
        mask = np.ones(len(ylabs), dtype=np.bool_)
        # If the label (int) is different wrt. its neighbour, show it
        mask[1:] = np.abs(np.diff([int_labels[y] for y in ylabs])) > 0
        # Only label the selection
        ax.set_yticks(np.arange(len(ylabs))[mask])
        ax.set_yticklabels(ylabs[mask])
    plt.xlabel("Time [s]")

    if cbar_loc is not None:
        # Create a white box to place the color bar in
        # See https://stackoverflow.com/questions/18211967
        bbox = inset_locator.inset_axes(ax, width="20%", height="22%", loc=cbar_loc)
        _ = [bbox.spines[k].set_visible(False) for k in bbox.spines]
        bbox.patch.set_facecolor((1, 1, 1, 0.9))
        bbox.set_xticks([])
        bbox.set_yticks([])

        # Create the actual color bar
        cax = inset_locator.inset_axes(bbox, width="90%", height="20%", loc="upper center")
        plt.colorbar(cax=cax, label="Intensity [PE/ns]", orientation="horizontal")
        cax.xaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter("%g"))

    plt.sca(ax)


def seconds_range_xaxis(seconds_range, t0=None):
    """Make a pretty time axis given seconds_range."""
    plt.xlim(*seconds_range)
    ax = plt.gca()
    ax.ticklabel_format(useOffset=False)
    xticks = plt.xticks()[0]
    if not len(xticks):
        return

    # Format the labels
    # I am not very proud of this code...
    def chop(x):
        return np.floor(x).astype(np.int64)

    if t0 is None:
        xticks_ns = np.round(xticks * int(1e9)).astype(np.int64)
    else:
        xticks_ns = np.round((xticks - xticks[0]) * int(1e9)).astype(np.int64)
    sec = chop(xticks_ns // int(1e9))
    ms = chop((xticks_ns % int(1e9)) // int(1e6))
    us = chop((xticks_ns % int(1e6)) // int(1e3))
    samples = chop((xticks_ns % int(1e3)) // 10)

    labels = [str(sec[i]) for i in range(len(xticks))]
    print_ns = np.any(samples != samples[0])
    print_us = print_ns | np.any(us != us[0])
    print_ms = print_us | np.any(ms != ms[0])
    if print_ms and t0 is None:
        labels = [l + f".{ms[i]:03}" for i, l in enumerate(labels)]
        if print_us:
            labels = [l + r" $\bf{" + f"{us[i]:03}" + "}$" for i, l in enumerate(labels)]
            if print_ns:
                labels = [l + f" {samples[i]:02}0" for i, l in enumerate(labels)]
        plt.xticks(ticks=xticks, labels=labels, rotation=90)
    else:
        labels = list(chop((xticks_ns // 10) * 10))
        labels[-1] = ""
        plt.xticks(ticks=xticks, labels=labels, rotation=0)
    if t0 is None:
        plt.xlabel("Time since run start [sec]")
    else:
        plt.xlabel("Time [ns]")


def plot_peak(p, t0=None, center_time=True, **kwargs):
    x, y = time_and_samples(p, t0=t0)
    kwargs.setdefault("linewidth", 1)

    # Plot waveform
    plt.plot(x, y, drawstyle="steps-pre", **kwargs)
    if "linewidth" in kwargs:
        del kwargs["linewidth"]
    kwargs["alpha"] = kwargs.get("alpha", 1) * 0.2
    plt.fill_between(x, 0, y, step="pre", linewidth=0, **kwargs)

    # Mark extent with thin black line
    plt.plot([x[0], x[-1]], [y.max(), y.max()], c="k", alpha=0.3, linewidth=1)

    # Mark center time with thin black line
    if center_time:
        if t0 is None:
            t0 = p["time"]
        ct = (p["center_time"] - t0) / int(1e9)
        plt.axvline(ct, c="k", alpha=0.4, linewidth=1, linestyle="--")


def time_and_samples(p, t0=None):
    """Return (x, y) numpy arrays for plotting the waveform data in p using 'steps-pre'.

    Where x is the time since t0 in seconds (or another time_scale), and y is intensity in PE / ns.
    :param p: Peak or other similar strax data type :param t0: Zero of time in ns since unix epoch

    """
    n = p["length"]
    if t0 is None:
        t0 = p["time"]
    x = ((p["time"] - t0) + np.arange(n + 1) * p["dt"]) / int(1e9)
    y = p["data"][:n] / p["dt"]
    return x, np.concatenate([[y[0]], y])

1	import matplotlib	2✔
2	import matplotlib.pyplot as plt	2✔
3	import numpy as np	2✔
4	import strax	2✔
5	import straxen	2✔
6	from mpl_toolkits.axes_grid1 import inset_locator	2✔
7
8	from .daq_waveforms import group_by_daq	2✔
9	from .records_matrix import DEFAULT_MAX_SAMPLES	2✔
10
11	export, __all__ = strax.exporter()	2✔
12	__all__.extend(["plot_wf"])	2✔
13
14
15	@straxen.mini_analysis()	2✔
16	def plot_waveform(	2✔
17	context,
18	deep=False,
19	show_largest=100,
20	figsize=None,
21	max_samples=DEFAULT_MAX_SAMPLES,
22	ignore_max_sample_warning=True,
23	cbar_loc="lower right",
24	lower_panel_height=2,
25	**kwargs,
26	):
27	"""Plot the sum waveform and optionally per-PMT waveforms.
28
29	:param deep: If True, show per-PMT waveform matrix under sum waveform. If 'raw', use
30	raw_records instead of records to do so. :param show_largest: Show only the largest show_largest
31	peaks. :param figsize: Matplotlib figure size for the plot Additional options for deep =
32	True or raw: :param cbar_loc: location of the intensity color bar. Set to None to omit it
33	altogether. :param lower_panel_height: Height of the lower panel in terms of the height of
34	the upper panel.
35
36	"""
37	if figsize is None:	2✔
38	figsize = (10, 6 if deep else 4)	2✔
39
40	if not deep:	2✔
41	context.plot_peaks(**kwargs, show_largest=show_largest, figsize=figsize)	2✔
42
43	else:
44	f, axes = plt.subplots(	2✔
45	2,
46	1,
47	constrained_layout=True,
48	figsize=figsize,
49	gridspec_kw={"height_ratios": [1, lower_panel_height]},
50	)
51
52	plt.sca(axes[0])	2✔
53	context.plot_peaks(**kwargs, show_largest=show_largest, single_figure=False, xaxis=False)	2✔
54
55	plt.sca(axes[1])	2✔
56	context.plot_records_matrix(	2✔
57	**kwargs,
58	cbar_loc=cbar_loc,
59	max_samples=max_samples,
60	ignore_max_sample_warning=ignore_max_sample_warning,
61	raw=deep == "raw",
62	single_figure=False,
63	)
64
65	plt.subplots_adjust(hspace=0)	2✔
66
67
68	@straxen.mini_analysis(	2✔
69	requires=("peaks", "peak_basics"), default_time_selection="touching", warn_beyond_sec=60
70	)
71	def plot_peaks(	2✔
72	peaks,
73	seconds_range,
74	t_reference,
75	show_largest=100,
76	single_figure=True,
77	figsize=(10, 4),
78	xaxis=True,
79	):
80	if single_figure:	2✔
81	plt.figure(figsize=figsize)	2✔
82	plt.axhline(0, c="k", alpha=0.2)	2✔
83
84	peaks = peaks[np.argsort(-peaks["area"])[:show_largest]]	2✔
85	peaks = strax.sort_by_time(peaks)	2✔
86
87	for p in peaks:	2✔
88	plot_peak(p, t0=t_reference, color={0: "gray", 1: "b", 2: "g"}[p["type"]])	2✔
89
90	if xaxis == "since_start":	2✔
91	seconds_range_xaxis(seconds_range, t0=seconds_range[0])	×
92	elif xaxis:	2✔
93	seconds_range_xaxis(seconds_range)	2✔
94	plt.xlim(*seconds_range)	2✔
95	else:
96	plt.xticks([])	2✔
97	plt.xlim(*seconds_range)	2✔
98	plt.ylabel("Intensity [PE/ns]")	2✔
99	if single_figure:	2✔
100	plt.tight_layout()	2✔
101
102
103	@straxen.mini_analysis(	2✔
104	requires=("peaks", "peak_basics"), default_time_selection="touching", warn_beyond_sec=60
105	)
106	def plot_hit_pattern(	2✔
107	peaks,
108	seconds_range,
109	t_reference,
110	axes=None,
111	vmin=None,
112	log_scale=False,
113	label=None,
114	single_figure=False,
115	xenon1t=False,
116	figsize=(10, 4),
117	):
118	if single_figure:	2✔
119	plt.figure(figsize=figsize)	×
120	if len(peaks) > 1:	2✔
121	print(f"warning showing total area of {len(peaks)} peaks")	×
122	straxen.plot_pmts(	2✔
123	np.sum(peaks["area_per_channel"], axis=0),
124	axes=axes,
125	vmin=vmin,
126	log_scale=log_scale,
127	label=label,
128	xenon1t=xenon1t,
129	)
130
131
132	@straxen.mini_analysis()	2✔
133	def plot_records_matrix(	2✔
134	context,
135	run_id,
136	seconds_range,
137	cbar_loc="upper right",
138	raw=False,
139	single_figure=True,
140	figsize=(10, 4),
141	group_by=None,
142	max_samples=DEFAULT_MAX_SAMPLES,
143	ignore_max_sample_warning=False,
144	vmin=None,
145	vmax=None,
146	**kwargs,
147	):
148	if seconds_range is None:	2✔
149	raise ValueError(
150	"You must pass a time selection (e.g. seconds_range) to plot_records_matrix."
151	)
152
153	if single_figure:	2✔
154	plt.figure(figsize=figsize, constrained_layout=True)	2✔
155
156	f = context.raw_records_matrix if raw else context.records_matrix	2✔
157
158	wvm, ts, ys = f(	2✔
159	run_id,
160	max_samples=max_samples,
161	ignore_max_sample_warning=ignore_max_sample_warning,
162	**kwargs,
163	)
164	if group_by is not None:	2✔
165	ylabs, wvm_mask = group_by_daq(run_id, group_by)	2✔
166	wvm = wvm[:, wvm_mask]	2✔
167	plt.ylabel(group_by)	2✔
168	else:
169	plt.ylabel("Channel number")	2✔
170
171	# extract min and max from kwargs or set defaults
172	if vmin is None:	2✔
173	vmin = min(0.1 * wvm.max(), 1e-2)	2✔
174	if vmax is None:	2✔
175	vmax = wvm.max()	2✔
176
177	plt.pcolormesh(	2✔
178	ts,
179	ys,
180	wvm.T,
181	norm=matplotlib.colors.LogNorm(
182	vmin=vmin,
183	vmax=vmax,
184	),
185	cmap=plt.cm.inferno,
186	)
187	plt.xlim(*seconds_range)	2✔
188
189	ax = plt.gca()	2✔
190	if group_by is not None:	2✔
191	# Do some magic to convert all the labels to an integer that
192	# allows for remapping of the y labels to whatever is provided
193	# in the "ylabs", otherwise matplotlib shows nchannels different
194	# labels in the case of strings.
195	# Make a dict that converts the label to an int
196	int_labels = {h: i for i, h in enumerate(set(ylabs))}	2✔
197	mask = np.ones(len(ylabs), dtype=np.bool_)	2✔
198	# If the label (int) is different wrt. its neighbour, show it
199	mask[1:] = np.abs(np.diff([int_labels[y] for y in ylabs])) > 0	2✔
200	# Only label the selection
201	ax.set_yticks(np.arange(len(ylabs))[mask])	2✔
202	ax.set_yticklabels(ylabs[mask])	2✔
203	plt.xlabel("Time [s]")	2✔
204
205	if cbar_loc is not None:	2✔
206	# Create a white box to place the color bar in
207	# See https://stackoverflow.com/questions/18211967
208	bbox = inset_locator.inset_axes(ax, width="20%", height="22%", loc=cbar_loc)	2✔
209	_ = [bbox.spines[k].set_visible(False) for k in bbox.spines]	2✔
210	bbox.patch.set_facecolor((1, 1, 1, 0.9))	2✔
211	bbox.set_xticks([])	2✔
212	bbox.set_yticks([])	2✔
213
214	# Create the actual color bar
215	cax = inset_locator.inset_axes(bbox, width="90%", height="20%", loc="upper center")	2✔
216	plt.colorbar(cax=cax, label="Intensity [PE/ns]", orientation="horizontal")	2✔
217	cax.xaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter("%g"))	2✔
218
219	plt.sca(ax)	2✔
220
221
222	def seconds_range_xaxis(seconds_range, t0=None):	2✔
223	"""Make a pretty time axis given seconds_range."""
224	plt.xlim(*seconds_range)	2✔
225	ax = plt.gca()	2✔
226	ax.ticklabel_format(useOffset=False)	2✔
227	xticks = plt.xticks()[0]	2✔
228	if not len(xticks):	2✔
229	return	×
230
231	# Format the labels
232	# I am not very proud of this code...
233	def chop(x):	2✔
234	return np.floor(x).astype(np.int64)	2✔
235
236	if t0 is None:	2✔
237	xticks_ns = np.round(xticks * int(1e9)).astype(np.int64)	2✔
238	else:
239	xticks_ns = np.round((xticks - xticks[0]) * int(1e9)).astype(np.int64)	×
240	sec = chop(xticks_ns // int(1e9))	2✔
241	ms = chop((xticks_ns % int(1e9)) // int(1e6))	2✔
242	us = chop((xticks_ns % int(1e6)) // int(1e3))	2✔
243	samples = chop((xticks_ns % int(1e3)) // 10)	2✔
244
245	labels = [str(sec[i]) for i in range(len(xticks))]	2✔
246	print_ns = np.any(samples != samples[0])	2✔
247	print_us = print_ns \| np.any(us != us[0])	2✔
248	print_ms = print_us \| np.any(ms != ms[0])	2✔
249	if print_ms and t0 is None:	2✔
250	labels = [l + f".{ms[i]:03}" for i, l in enumerate(labels)]	2✔
251	if print_us:	2✔
252	labels = [l + r" $\bf{" + f"{us[i]:03}" + "}$" for i, l in enumerate(labels)]	2✔
253	if print_ns:	2✔
254	labels = [l + f" {samples[i]:02}0" for i, l in enumerate(labels)]	2✔
255	plt.xticks(ticks=xticks, labels=labels, rotation=90)	2✔
256	else:
257	labels = list(chop((xticks_ns // 10) * 10))	×
258	labels[-1] = ""	×
259	plt.xticks(ticks=xticks, labels=labels, rotation=0)	×
260	if t0 is None:	2✔
261	plt.xlabel("Time since run start [sec]")	2✔
262	else:
263	plt.xlabel("Time [ns]")	×
264
265
266	def plot_peak(p, t0=None, center_time=True, **kwargs):	2✔
267	x, y = time_and_samples(p, t0=t0)	2✔
268	kwargs.setdefault("linewidth", 1)	2✔
269
270	# Plot waveform
271	plt.plot(x, y, drawstyle="steps-pre", **kwargs)	2✔
272	if "linewidth" in kwargs:	2✔
273	del kwargs["linewidth"]	2✔
274	kwargs["alpha"] = kwargs.get("alpha", 1) * 0.2	2✔
275	plt.fill_between(x, 0, y, step="pre", linewidth=0, **kwargs)	2✔
276
277	# Mark extent with thin black line
278	plt.plot([x[0], x[-1]], [y.max(), y.max()], c="k", alpha=0.3, linewidth=1)	2✔
279
280	# Mark center time with thin black line
281	if center_time:	2✔
282	if t0 is None:	2✔
283	t0 = p["time"]	×
284	ct = (p["center_time"] - t0) / int(1e9)	2✔
285	plt.axvline(ct, c="k", alpha=0.4, linewidth=1, linestyle="--")	2✔
286
287
288	def time_and_samples(p, t0=None):	2✔
289	"""Return (x, y) numpy arrays for plotting the waveform data in p using 'steps-pre'.
290
291	Where x is the time since t0 in seconds (or another time_scale), and y is intensity in PE / ns.
292	:param p: Peak or other similar strax data type :param t0: Zero of time in ns since unix epoch
293
294	"""
295	n = p["length"]	2✔
296	if t0 is None:	2✔
297	t0 = p["time"]	×
298	x = ((p["time"] - t0) + np.arange(n + 1) * p["dt"]) / int(1e9)	2✔
299	y = p["data"][:n] / p["dt"]	2✔
300	return x, np.concatenate([[y[0]], y])	2✔

XENONnT / straxen / 6353497555

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous