24898394147

Committed 24 Apr 2026 03:45PM UTC coverage: 88.974% (+0.02%) from 88.954%

Build # 24898394147

Build Type

Pull #1657

github

Committed by

web-flow

Commit Message

Merge 0d70d8ff8 into 310bd239d

Pull Request Pull Request #1657: docformatter 1.7.6 --> 1.7.8

Coverage Stats

8901 of 10004 relevant lines covered (88.97%)

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

XENONnT / straxen / 24898394147

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