13791848185

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

Build # 13791848185

Build Type

Pull #1581

github

Committed by

web-flow

Commit Message

Merge c3e3de69f into f1b430214

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

Run Details

6 of 6 new or added lines in 4 files covered. (100.0%)

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

XENONnT / straxen / 13791848185

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