13791839586

Committed 11 Mar 2025 03:25PM UTC coverage: 89.67% (+0.02%) from 89.649%

Build # 13791839586

Build Type

Pull #1581

github

Committed by

web-flow

Commit Message

Merge 826703b83 into f1b430214

Pull Request Pull Request #1581: Added more colors+types to waveform_plot

Run Details

5 of 6 new or added lines in 1 file covered. (83.33%)

9 existing lines in 1 file now uncovered.

8620 of 9613 relevant lines covered (89.67%)

0.9 hits per line

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

91.97

/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()


@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.
    :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[strax.stable_argsort(-peaks["area"])[:show_largest]]
    peaks = strax.sort_by_time(peaks)

    for p in peaks:
        plot_peak(p, t0=t_reference, color={-1:'black', 0:'gray',1:'b',2:'g',3:'purple',20:'orange',22:'teal'}[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,
    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,
    )


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

XENONnT / straxen / 13791839586

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