24701692962

Committed 21 Apr 2026 02:59AM UTC coverage: 88.965% (+0.01%) from 88.954%

Build # 24701692962

Build Type

Pull #1650

github

Committed by

web-flow

Commit Message

Merge 38786e1d0 into 310bd239d

Pull Request Pull Request #1650: Fix NumPy 2.4+ / pandas 3 compatibility

Coverage Stats

21 of 24 new or added lines in 6 files covered. (87.5%)

272 existing lines in 29 files now uncovered.

8917 of 10023 relevant lines covered (88.97%)

0.89 hits per line

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

89.08

/straxen/matplotlib_utils.py

import numpy as np
import matplotlib
import matplotlib.pyplot as plt

import strax
import straxen

export, __all__ = strax.exporter()


@export
def plot_pmts(
    c,
    label="",
    figsize=None,
    show_tpc=True,
    extend="neither",
    vmin=None,
    vmax=None,
    **kwargs,
):
    """Plot the PMT arrays side-by-side, coloring the PMTS with c.

    :param c: Array of colors to use. Must have len() n_tpc_pmts
    :param label: Label for the color bar
    :param figsize: Figure size to use.
    :param extend: same as plt.colorbar(extend=...)
    :param vmin: Minimum of color scale
    :param vmax: maximum of color scale
    :param show_axis_labels: if True it will show x and y labels
    Other arguments are passed to plot_on_single_pmt_array.

    """
    if vmin is None:
        vmin = np.nanmin(c)
    if vmax is None:
        vmax = np.nanmax(c)
    if vmin == vmax:
        # Single-valued array passed
        vmax += 1
    if figsize is None:
        figsize = (13.25, 5.75)

    f, axes = plt.subplots(1, 2, figsize=figsize)
    plot_result = None
    for array_i, array_name in enumerate(["top", "bottom"]):
        ax = axes[array_i]
        plt.sca(ax)
        plt.title(array_name.capitalize())

        plot_result = plot_on_single_pmt_array(
            c,
            array_name=array_name,
            show_tpc=show_tpc,
            vmin=vmin,
            vmax=vmax,
            **kwargs,
        )

    axes[0].set_xlabel("x [cm]")
    axes[0].xaxis.set_label_coords(1.035, -0.075)
    axes[0].set_ylabel("y [cm]")

    axes[1].yaxis.tick_right()
    axes[1].yaxis.set_label_position("right")

    plt.tight_layout()
    plt.subplots_adjust(wspace=0)
    plt.colorbar(mappable=plot_result, ax=axes, extend=extend, label=label)


@export
def plot_on_single_pmt_array(
    c,
    array_name="top",
    r=straxen.tpc_r * 1.03,
    pmt_label_size=8,
    pmt_label_color="white",
    show_tpc=True,
    log_scale=False,
    vmin=None,
    vmax=None,
    dead_pmts=None,
    dead_pmt_color="gray",
    **kwargs,
):
    """Plot one of the PMT arrays and color it by c.

    :param c: Array of colors to use. Must be len() of the number of TPC PMTs
    :param label: Label for the color bar
    :param pmt_label_size: Fontsize for the PMT number labels.
    Set to 0 to disable.
    :param pmt_label_color: Text color of the PMT number labels.
    :param log_scale: If True, use a logarithmic color scale
    :param extend: same as plt.colorbar(extend=...)
    :param vmin: Minimum of color scale
    :param vmax: maximum of color scale
    Other arguments are passed to plt.scatter.

    """
    if vmin is None:
        vmin = c.min()
    if vmax is None:
        vmax = c.max()

    pmt_positions = straxen.pmt_positions().to_records()

    ax = plt.gca()
    ax.set_aspect("equal")
    plt.xlim(-r, r)
    plt.ylim(-r, r)

    mask = pmt_positions["array"] == array_name
    pos = pmt_positions[mask]

    kwargs.setdefault("s", 280)
    if log_scale:
        kwargs.setdefault("norm", matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax))
    else:
        kwargs.setdefault("vmin", vmin)
        kwargs.setdefault("vmax", vmax)
    result = plt.scatter(pos["x"], pos["y"], c=c[mask], **kwargs)

    if show_tpc:
        ax.set_facecolor("lightgrey")
        ax.add_artist(
            plt.Circle((0, 0), straxen.tpc_r, edgecolor="k", facecolor="w", zorder=-5, linewidth=1)
        )
    else:
        ax.set_axis_off()
    if dead_pmts is not None:
        _dead_mask = [pi in dead_pmts for pi in pos["i"]]
        plt.scatter(pos[_dead_mask]["x"], pos[_dead_mask]["y"], c=dead_pmt_color, **kwargs)

    if pmt_label_size:
        for p in pos:
            plt.text(
                p["x"],
                p["y"],
                str(p["i"]),
                horizontalalignment="center",
                verticalalignment="center",
                fontsize=pmt_label_size,
                color=pmt_label_color,
            )
    return result


@export
def log_y(a=None, b=None, scalar_ticks=True, tick_at=None):
    """Make the y axis use a log scale from a to b."""
    plt.yscale("log")
    if a is not None:
        if b is None:
            a, b = a[0], a[-1]
        ax = plt.gca()
        plt.ylim(a, b)
        if scalar_ticks:
            ax.yaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter("%g"))
            ax.set_yticks(logticks(a, b, tick_at))


@export
def log_x(a=None, b=None, scalar_ticks=True, tick_at=None):
    """Make the x axis use a log scale from a to b."""
    plt.xscale("log")
    if a is not None:
        if b is None:
            a, b = a[0], a[-1]
        plt.xlim(a, b)
        ax = plt.gca()
        if scalar_ticks:
            ax.xaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter("%g"))
            ax.set_xticks(logticks(a, b, tick_at))


def logticks(tmin, tmax=None, tick_at=None):
    if tick_at is None:
        tick_at = (1, 2, 5, 10)
    a, b = np.log10([tmin, tmax])
    a = np.floor(a)
    b = np.ceil(b)
    ticks = strax.stable_sort(
        np.unique(np.outer(np.array(tick_at), 10.0 ** np.arange(a, b)).ravel())
    )
    ticks = ticks[(tmin <= ticks) & (ticks <= tmax)]
    return ticks


@export
def draw_box(x, y, **kwargs):
    """Draw rectangle, given x-y boundary tuples."""
    plt.gca().add_patch(
        matplotlib.patches.Rectangle(
            (x[0], y[0]), x[1] - x[0], y[1] - y[0], facecolor="none", **kwargs
        )
    )


@export
def plot_single_pulse(records, run_id, pulse_i=""):
    """Function which plots a single pulse.

    :param records: Records which belong to the pulse.
    :param run_id: Id of the run.
    :param pulse_i: Index of the pulse to be plotted.
    :return: fig, axes objects.

    """
    pulse = _make_pulse(records)

    fig, axes = plt.subplots()
    sec, ns = _split_off_ns(records[0]["time"])
    date = np.datetime_as_string(sec.astype("<M8[ns]"), unit="s")
    plt.title(f"Pulse {pulse_i} from {run_id}\nRecorded at {date[:10]}, {date[10:]} UTC {ns} ns")

    plt.step(np.arange(len(pulse)), pulse, where="post", label=f'Ch: {records[0]["channel"]}')
    plt.legend()
    plt.xlabel(f'Sample [{records[0]["dt"]} ns]')
    plt.ylabel("Height [ADC counts]")
    plt.grid()
    return fig, axes


def _make_pulse(records):
    """Helper to make a pulse based on fragements."""
    pulse = np.zeros(records[0]["pulse_length"], dtype=np.float32)

    offset = 0
    for r in records:
        pulse[offset : offset + r["length"]] = r["data"][: r["length"]]
        offset += r["length"]

    return pulse


def _split_off_ns(time):
    """Mini helper to divide time into seconds and ns."""
    sec = (time // 10**9) * 10**9
    ns = time - sec
    return sec, ns

1	import numpy as np	1✔
2	import matplotlib	1✔
3	import matplotlib.pyplot as plt	1✔
4
5	import strax	1✔
6	import straxen	1✔
7
8	export, __all__ = strax.exporter()	1✔
9
10
11	@export	1✔
12	def plot_pmts(	1✔
13	c,
14	label="",
15	figsize=None,
16	show_tpc=True,
17	extend="neither",
18	vmin=None,
19	vmax=None,
20	**kwargs,
21	):
22	"""Plot the PMT arrays side-by-side, coloring the PMTS with c.
23
24	:param c: Array of colors to use. Must have len() n_tpc_pmts
25	:param label: Label for the color bar
26	:param figsize: Figure size to use.
27	:param extend: same as plt.colorbar(extend=...)
28	:param vmin: Minimum of color scale
29	:param vmax: maximum of color scale
30	:param show_axis_labels: if True it will show x and y labels
31	Other arguments are passed to plot_on_single_pmt_array.
32
33	"""
34	if vmin is None:	1✔
35	vmin = np.nanmin(c)	1✔
36	if vmax is None:	1✔
37	vmax = np.nanmax(c)	1✔
38	if vmin == vmax:	1✔
39	# Single-valued array passed
40	vmax += 1	1✔
41	if figsize is None:	1✔
42	figsize = (13.25, 5.75)	1✔
43
44	f, axes = plt.subplots(1, 2, figsize=figsize)	1✔
45	plot_result = None	1✔
46	for array_i, array_name in enumerate(["top", "bottom"]):	1✔
47	ax = axes[array_i]	1✔
48	plt.sca(ax)	1✔
49	plt.title(array_name.capitalize())	1✔
50
51	plot_result = plot_on_single_pmt_array(	1✔
52	c,
53	array_name=array_name,
54	show_tpc=show_tpc,
55	vmin=vmin,
56	vmax=vmax,
57	**kwargs,
58	)
59
60	axes[0].set_xlabel("x [cm]")	1✔
61	axes[0].xaxis.set_label_coords(1.035, -0.075)	1✔
62	axes[0].set_ylabel("y [cm]")	1✔
63
64	axes[1].yaxis.tick_right()	1✔
65	axes[1].yaxis.set_label_position("right")	1✔
66
67	plt.tight_layout()	1✔
68	plt.subplots_adjust(wspace=0)	1✔
69	plt.colorbar(mappable=plot_result, ax=axes, extend=extend, label=label)	1✔
70
71
72	@export	1✔
73	def plot_on_single_pmt_array(	1✔
74	c,
75	array_name="top",
76	r=straxen.tpc_r * 1.03,
77	pmt_label_size=8,
78	pmt_label_color="white",
79	show_tpc=True,
80	log_scale=False,
81	vmin=None,
82	vmax=None,
83	dead_pmts=None,
84	dead_pmt_color="gray",
85	**kwargs,
86	):
87	"""Plot one of the PMT arrays and color it by c.
88
89	:param c: Array of colors to use. Must be len() of the number of TPC PMTs
90	:param label: Label for the color bar
91	:param pmt_label_size: Fontsize for the PMT number labels.
92	Set to 0 to disable.
93	:param pmt_label_color: Text color of the PMT number labels.
94	:param log_scale: If True, use a logarithmic color scale
95	:param extend: same as plt.colorbar(extend=...)
96	:param vmin: Minimum of color scale
97	:param vmax: maximum of color scale
98	Other arguments are passed to plt.scatter.
99
100	"""
101	if vmin is None:	1✔
UNCOV 102	vmin = c.min()	×
103	if vmax is None:	1✔
104	vmax = c.max()	1✔
105
106	pmt_positions = straxen.pmt_positions().to_records()	1✔
107
108	ax = plt.gca()	1✔
109	ax.set_aspect("equal")	1✔
110	plt.xlim(-r, r)	1✔
111	plt.ylim(-r, r)	1✔
112
113	mask = pmt_positions["array"] == array_name	1✔
114	pos = pmt_positions[mask]	1✔
115
116	kwargs.setdefault("s", 280)	1✔
117	if log_scale:	1✔
118	kwargs.setdefault("norm", matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax))	1✔
119	else:
120	kwargs.setdefault("vmin", vmin)	1✔
121	kwargs.setdefault("vmax", vmax)	1✔
122	result = plt.scatter(pos["x"], pos["y"], c=c[mask], **kwargs)	1✔
123
124	if show_tpc:	1✔
125	ax.set_facecolor("lightgrey")	1✔
126	ax.add_artist(	1✔
127	plt.Circle((0, 0), straxen.tpc_r, edgecolor="k", facecolor="w", zorder=-5, linewidth=1)
128	)
129	else:
UNCOV 130	ax.set_axis_off()	×
131	if dead_pmts is not None:	1✔
132	_dead_mask = [pi in dead_pmts for pi in pos["i"]]	1✔
133	plt.scatter(pos[_dead_mask]["x"], pos[_dead_mask]["y"], c=dead_pmt_color, **kwargs)	1✔
134
135	if pmt_label_size:	1✔
136	for p in pos:	1✔
137	plt.text(	1✔
138	p["x"],
139	p["y"],
140	str(p["i"]),
141	horizontalalignment="center",
142	verticalalignment="center",
143	fontsize=pmt_label_size,
144	color=pmt_label_color,
145	)
146	return result	1✔
147
148
149	@export	1✔
150	def log_y(a=None, b=None, scalar_ticks=True, tick_at=None):	1✔
151	"""Make the y axis use a log scale from a to b."""
152	plt.yscale("log")	×
153	if a is not None:	×
154	if b is None:	×
155	a, b = a[0], a[-1]	×
156	ax = plt.gca()	×
157	plt.ylim(a, b)	×
158	if scalar_ticks:	×
UNCOV 159	ax.yaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter("%g"))	×
UNCOV 160	ax.set_yticks(logticks(a, b, tick_at))	×
161
162
163	@export	1✔
164	def log_x(a=None, b=None, scalar_ticks=True, tick_at=None):	1✔
165	"""Make the x axis use a log scale from a to b."""
166	plt.xscale("log")	1✔
167	if a is not None:	1✔
168	if b is None:	1✔
UNCOV 169	a, b = a[0], a[-1]	×
170	plt.xlim(a, b)	1✔
171	ax = plt.gca()	1✔
172	if scalar_ticks:	1✔
173	ax.xaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter("%g"))	1✔
174	ax.set_xticks(logticks(a, b, tick_at))	1✔
175
176
177	def logticks(tmin, tmax=None, tick_at=None):	1✔
178	if tick_at is None:	1✔
179	tick_at = (1, 2, 5, 10)	1✔
180	a, b = np.log10([tmin, tmax])	1✔
181	a = np.floor(a)	1✔
182	b = np.ceil(b)	1✔
183	ticks = strax.stable_sort(	1✔
184	np.unique(np.outer(np.array(tick_at), 10.0 ** np.arange(a, b)).ravel())
185	)
186	ticks = ticks[(tmin <= ticks) & (ticks <= tmax)]	1✔
187	return ticks	1✔
188
189
190	@export	1✔
191	def draw_box(x, y, **kwargs):	1✔
192	"""Draw rectangle, given x-y boundary tuples."""
UNCOV 193	plt.gca().add_patch(	×
194	matplotlib.patches.Rectangle(
195	(x[0], y[0]), x[1] - x[0], y[1] - y[0], facecolor="none", **kwargs
196	)
197	)
198
199
200	@export	1✔
201	def plot_single_pulse(records, run_id, pulse_i=""):	1✔
202	"""Function which plots a single pulse.
203
204	:param records: Records which belong to the pulse.
205	:param run_id: Id of the run.
206	:param pulse_i: Index of the pulse to be plotted.
207	:return: fig, axes objects.
208
209	"""
210	pulse = _make_pulse(records)	1✔
211
212	fig, axes = plt.subplots()	1✔
213	sec, ns = _split_off_ns(records[0]["time"])	1✔
214	date = np.datetime_as_string(sec.astype("<M8[ns]"), unit="s")	1✔
215	plt.title(f"Pulse {pulse_i} from {run_id}\nRecorded at {date[:10]}, {date[10:]} UTC {ns} ns")	1✔
216
217	plt.step(np.arange(len(pulse)), pulse, where="post", label=f'Ch: {records[0]["channel"]}')	1✔
218	plt.legend()	1✔
219	plt.xlabel(f'Sample [{records[0]["dt"]} ns]')	1✔
220	plt.ylabel("Height [ADC counts]")	1✔
221	plt.grid()	1✔
222	return fig, axes	1✔
223
224
225	def _make_pulse(records):	1✔
226	"""Helper to make a pulse based on fragements."""
227	pulse = np.zeros(records[0]["pulse_length"], dtype=np.float32)	1✔
228
229	offset = 0	1✔
230	for r in records:	1✔
231	pulse[offset : offset + r["length"]] = r["data"][: r["length"]]	1✔
232	offset += r["length"]	1✔
233
234	return pulse	1✔
235
236
237	def _split_off_ns(time):	1✔
238	"""Mini helper to divide time into seconds and ns."""
239	sec = (time // 10*9) 10**9	1✔
240	ns = time - sec	1✔
241	return sec, ns	1✔

XENONnT / straxen / 24701692962

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