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

89.17

/straxen/matplotlib_utils.py

import warnings

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