10286081415

Committed 07 Aug 2024 02:17PM UTC coverage: 91.122% (-0.01%) from 91.135%

Build # 10286081415

Build Type

Pull #1404

github

Committed by

web-flow

Commit Message

Merge 68e283066 into 4d6c2b7ba

Pull Request Pull Request #1404: Plugins for position reconstruction with conditional normalizing flow

Run Details

170 of 190 new or added lines in 6 files covered. (89.47%)

3 existing lines in 1 file now uncovered.

9145 of 10036 relevant lines covered (91.12%)

1.82 hits per line

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

83.78

/straxen/plugins/peaks/peak_positions_flow.py

import numpy as np
import strax
import straxen

export, __all__ = strax.exporter()


@export
class PeakPositionsFlow(strax.Plugin):
    """Conditional Normalizing Flow for position reconstruction.

    This plugin reconstructs the position of S2 peaks using a conditional normalizing flow model.
    It provides x and y coordinates of the reconstructed position, along with uncertainty contours
    and uncertainty estimates in r and theta.

    Depends on: 'peaks'
    Provides: 'peak_positions_flow'

    Configuration options:
    - min_reconstruction_area: Minimum area (PE) required for reconstruction
    - n_poly: Size of the uncertainty contour
    - sig: Confidence level of the contour
    - log_area_scale: Scaling parameter for log area
    - n_top_pmts: Number of top PMTs
    - pred_function: Path to the compiled JAX function for predictions

    """

    depends_on = "peaks"
    provides = "peak_positions_flow"

    __version__ = "0.0.3"

    compressor = "zstd"
    parallel = True  # can set to "process" after #82

    # Configuration options
    min_reconstruction_area = straxen.URLConfig(
        help="Skip reconstruction if area (PE) is less than this",
        default=10,
        infer_type=False,
    )

    n_poly = straxen.URLConfig(
        default=16,
        infer_type=False,
        help="Size of uncertainty contour",
    )

    sig = straxen.URLConfig(
        default=0.393,
        infer_type=False,
        help="Confidence level of contour",
    )

    log_area_scale = straxen.URLConfig(
        default=10,
        infer_type=False,
        help="Scaling parameter for log area",
    )

    n_top_pmts = straxen.URLConfig(
        default=straxen.n_top_pmts, infer_type=False, help="Number of top PMTs"
    )

    pred_function = straxen.URLConfig(
        default=(
            "jax://resource://flow_20240730.tar.gz?"
            "n_poly=plugin.n_poly&sig=plugin.sig&fmt=abs_path"
        ),
        help="Compiled JAX function",
    )

    def vectorized_prediction_chunk(self, flow_condition, N_chunk_max=4096):
        """Compute predictions for a chunk of data.

        Args:
            flow_condition: Input data for the flow model
            N_chunk_max: Maximum chunk size (default: 4096)

        Returns:
            xy: Predicted x and y coordinates
            contour: Uncertainty contours

        """
        N_entries = flow_condition.shape[0]
        if N_entries > N_chunk_max:
            raise ValueError("Chunk greater than max size")
        else:
            inputs = np.zeros((N_chunk_max, self.n_top_pmts + 1))
            inputs[:N_entries] = flow_condition
            xy, contour = self.pred_function(inputs)
            return xy[:N_entries], contour[:N_entries]

    def prediction_loop(self, flow_condition, N_chunk_max=4096):
        """Compute predictions for arbitrary-size inputs using a loop.

        Args:
            flow_condition: Input data for the flow model
            N_chunk_max: Maximum chunk size (default: 4096)

        Returns:
            xy: Predicted x and y coordinates
            contour: Uncertainty contours

        """
        N_entries = flow_condition.shape[0]
        if N_entries <= N_chunk_max:
            return self.vectorized_prediction_chunk(flow_condition, N_chunk_max=N_chunk_max)
        N_chunks = N_entries // N_chunk_max

        xy_list = []
        contour_list = []
        for i in range(N_chunks):
            xy, contour = self.vectorized_prediction_chunk(
                flow_condition[i * N_chunk_max : (i + 1) * N_chunk_max]
            )
            xy_list.append(xy)
            contour_list.append(contour)

        if N_chunks * N_chunk_max < N_entries:
            xy, contour = self.vectorized_prediction_chunk(flow_condition[(i + 1) * N_chunk_max :])
            xy_list.append(xy)
            contour_list.append(contour)
        return np.concatenate(xy_list, axis=0), np.concatenate(contour_list, axis=0)

    def infer_dtype(self):
        """Define the data type for the output.

        Returns:
            dtype: Numpy dtype for the output array

        """
        dtype = [
            (("Reconstructed flow S2 X position (cm), uncorrected", "x_flow"), np.float32),
            (("Reconstructed flow S2 Y position (cm), uncorrected", "y_flow"), np.float32),
            (
                ("Position uncertainty contour", "position_contours_flow"),
                np.float32,
                (self.n_poly + 1, 2),
            ),
            (("Position uncertainty in r (cm)", "r_uncertainty_flow"), np.float32),
            (("Position uncertainty in theta (rad)", "theta_uncertainty_flow"), np.float32),
        ]
        dtype += strax.time_fields
        return dtype

    def compute(self, peaks):
        """Compute the position reconstruction for the given peaks.

        Args:
            peaks: Input peak data

        Returns:
            result: Array with reconstructed positions and uncertainties

        """
        # Initialize result array
        result = np.ones(len(peaks), dtype=self.dtype)
        result["time"], result["endtime"] = peaks["time"], strax.endtime(peaks)

        # Set default values to NaN
        result["x_flow"] *= float("nan")
        result["y_flow"] *= float("nan")
        result["position_contours_flow"] *= float("nan")
        result["r_uncertainty_flow"] *= np.nan
        result["theta_uncertainty_flow"] *= np.nan

        # Keep large peaks only
        peak_mask = peaks["area"] > self.min_reconstruction_area
        if not np.sum(peak_mask):
            # Nothing to do, and .predict crashes on empty arrays
            return result

        # Prepare input data for the flow model
        area_per_channel_top = peaks["area_per_channel"][peak_mask, 0 : self.n_top_pmts]
        total_top_areas = np.sum(area_per_channel_top, axis=1)
        with np.errstate(divide="ignore", invalid="ignore"):
            flow_data = np.concatenate(
                [
                    area_per_channel_top / total_top_areas[..., np.newaxis],
                    np.log(total_top_areas[..., np.newaxis]) / self.log_area_scale,
                ],
                axis=1,
            )

        # Get position reconstruction
        xy, contours = self.prediction_loop(flow_data)

        # Write output to the result array
        result["x_flow"][peak_mask] = xy[:, 0]
        result["y_flow"][peak_mask] = xy[:, 1]
        result["position_contours_flow"][peak_mask] = contours

        # Calculate uncertainties in r and theta
        r_array = np.linalg.norm(contours, axis=2)
        r_min = np.min(r_array, axis=1)
        r_max = np.max(r_array, axis=1)

        theta_array = np.arctan2(contours[..., 1], contours[..., 0])
        theta_min = np.min(theta_array, axis=1)
        theta_max = np.max(theta_array, axis=1)
        theta_diff = theta_max - theta_min
        theta_diff[theta_diff > np.pi] -= 2 * np.pi

        result["r_uncertainty_flow"][peak_mask] = (r_max - r_min) / 2
        result["theta_uncertainty_flow"][peak_mask] = np.abs(theta_diff) / 2

        return result

1	import numpy as np	2✔
2	import strax	2✔
3	import straxen	2✔
4
5	export, __all__ = strax.exporter()	2✔
6
7
8	@export	2✔
9	class PeakPositionsFlow(strax.Plugin):	2✔
10	"""Conditional Normalizing Flow for position reconstruction.
11
12	This plugin reconstructs the position of S2 peaks using a conditional normalizing flow model.
13	It provides x and y coordinates of the reconstructed position, along with uncertainty contours
14	and uncertainty estimates in r and theta.
15
16	Depends on: 'peaks'
17	Provides: 'peak_positions_flow'
18
19	Configuration options:
20	- min_reconstruction_area: Minimum area (PE) required for reconstruction
21	- n_poly: Size of the uncertainty contour
22	- sig: Confidence level of the contour
23	- log_area_scale: Scaling parameter for log area
24	- n_top_pmts: Number of top PMTs
25	- pred_function: Path to the compiled JAX function for predictions
26
27	"""
28
29	depends_on = "peaks"	2✔
30	provides = "peak_positions_flow"	2✔
31
32	__version__ = "0.0.3"	2✔
33
34	compressor = "zstd"	2✔
35	parallel = True # can set to "process" after #82	2✔
36
37	# Configuration options
38	min_reconstruction_area = straxen.URLConfig(	2✔
39	help="Skip reconstruction if area (PE) is less than this",
40	default=10,
41	infer_type=False,
42	)
43
44	n_poly = straxen.URLConfig(	2✔
45	default=16,
46	infer_type=False,
47	help="Size of uncertainty contour",
48	)
49
50	sig = straxen.URLConfig(	2✔
51	default=0.393,
52	infer_type=False,
53	help="Confidence level of contour",
54	)
55
56	log_area_scale = straxen.URLConfig(	2✔
57	default=10,
58	infer_type=False,
59	help="Scaling parameter for log area",
60	)
61
62	n_top_pmts = straxen.URLConfig(	2✔
63	default=straxen.n_top_pmts, infer_type=False, help="Number of top PMTs"
64	)
65
66	pred_function = straxen.URLConfig(	2✔
67	default=(
68	"jax://resource://flow_20240730.tar.gz?"
69	"n_poly=plugin.n_poly&sig=plugin.sig&fmt=abs_path"
70	),
71	help="Compiled JAX function",
72	)
73
74	def vectorized_prediction_chunk(self, flow_condition, N_chunk_max=4096):	2✔
75	"""Compute predictions for a chunk of data.
76
77	Args:
78	flow_condition: Input data for the flow model
79	N_chunk_max: Maximum chunk size (default: 4096)
80
81	Returns:
82	xy: Predicted x and y coordinates
83	contour: Uncertainty contours
84
85	"""
86	N_entries = flow_condition.shape[0]	2✔
87	if N_entries > N_chunk_max:	2✔
88	raise ValueError("Chunk greater than max size")
89	else:
90	inputs = np.zeros((N_chunk_max, self.n_top_pmts + 1))	2✔
91	inputs[:N_entries] = flow_condition	2✔
92	xy, contour = self.pred_function(inputs)	2✔
93	return xy[:N_entries], contour[:N_entries]	2✔
94
95	def prediction_loop(self, flow_condition, N_chunk_max=4096):	2✔
96	"""Compute predictions for arbitrary-size inputs using a loop.
97
98	Args:
99	flow_condition: Input data for the flow model
100	N_chunk_max: Maximum chunk size (default: 4096)
101
102	Returns:
103	xy: Predicted x and y coordinates
104	contour: Uncertainty contours
105
106	"""
107	N_entries = flow_condition.shape[0]	2✔
108	if N_entries <= N_chunk_max:	2✔
109	return self.vectorized_prediction_chunk(flow_condition, N_chunk_max=N_chunk_max)	2✔
NEW 110	N_chunks = N_entries // N_chunk_max	×
111
NEW 112	xy_list = []	×
NEW 113	contour_list = []	×
NEW 114	for i in range(N_chunks):	×
NEW 115	xy, contour = self.vectorized_prediction_chunk(	×
116	flow_condition[i * N_chunk_max : (i + 1) * N_chunk_max]
117	)
NEW 118	xy_list.append(xy)	×
NEW 119	contour_list.append(contour)	×
120
NEW 121	if N_chunks * N_chunk_max < N_entries:	×
NEW 122	xy, contour = self.vectorized_prediction_chunk(flow_condition[(i + 1) * N_chunk_max :])	×
NEW 123	xy_list.append(xy)	×
NEW 124	contour_list.append(contour)	×
NEW 125	return np.concatenate(xy_list, axis=0), np.concatenate(contour_list, axis=0)	×
126
127	def infer_dtype(self):	2✔
128	"""Define the data type for the output.
129
130	Returns:
131	dtype: Numpy dtype for the output array
132
133	"""
134	dtype = [	2✔
135	(("Reconstructed flow S2 X position (cm), uncorrected", "x_flow"), np.float32),
136	(("Reconstructed flow S2 Y position (cm), uncorrected", "y_flow"), np.float32),
137	(
138	("Position uncertainty contour", "position_contours_flow"),
139	np.float32,
140	(self.n_poly + 1, 2),
141	),
142	(("Position uncertainty in r (cm)", "r_uncertainty_flow"), np.float32),
143	(("Position uncertainty in theta (rad)", "theta_uncertainty_flow"), np.float32),
144	]
145	dtype += strax.time_fields	2✔
146	return dtype	2✔
147
148	def compute(self, peaks):	2✔
149	"""Compute the position reconstruction for the given peaks.
150
151	Args:
152	peaks: Input peak data
153
154	Returns:
155	result: Array with reconstructed positions and uncertainties
156
157	"""
158	# Initialize result array
159	result = np.ones(len(peaks), dtype=self.dtype)	2✔
160	result["time"], result["endtime"] = peaks["time"], strax.endtime(peaks)	2✔
161
162	# Set default values to NaN
163	result["x_flow"] *= float("nan")	2✔
164	result["y_flow"] *= float("nan")	2✔
165	result["position_contours_flow"] *= float("nan")	2✔
166	result["r_uncertainty_flow"] *= np.nan	2✔
167	result["theta_uncertainty_flow"] *= np.nan	2✔
168
169	# Keep large peaks only
170	peak_mask = peaks["area"] > self.min_reconstruction_area	2✔
171	if not np.sum(peak_mask):	2✔
172	# Nothing to do, and .predict crashes on empty arrays
173	return result	2✔
174
175	# Prepare input data for the flow model
176	area_per_channel_top = peaks["area_per_channel"][peak_mask, 0 : self.n_top_pmts]	2✔
177	total_top_areas = np.sum(area_per_channel_top, axis=1)	2✔
178	with np.errstate(divide="ignore", invalid="ignore"):	2✔
179	flow_data = np.concatenate(	2✔
180	[
181	area_per_channel_top / total_top_areas[..., np.newaxis],
182	np.log(total_top_areas[..., np.newaxis]) / self.log_area_scale,
183	],
184	axis=1,
185	)
186
187	# Get position reconstruction
188	xy, contours = self.prediction_loop(flow_data)	2✔
189
190	# Write output to the result array
191	result["x_flow"][peak_mask] = xy[:, 0]	2✔
192	result["y_flow"][peak_mask] = xy[:, 1]	2✔
193	result["position_contours_flow"][peak_mask] = contours	2✔
194
195	# Calculate uncertainties in r and theta
196	r_array = np.linalg.norm(contours, axis=2)	2✔
197	r_min = np.min(r_array, axis=1)	2✔
198	r_max = np.max(r_array, axis=1)	2✔
199
200	theta_array = np.arctan2(contours[..., 1], contours[..., 0])	2✔
201	theta_min = np.min(theta_array, axis=1)	2✔
202	theta_max = np.max(theta_array, axis=1)	2✔
203	theta_diff = theta_max - theta_min	2✔
204	theta_diff[theta_diff > np.pi] -= 2 * np.pi	2✔
205
206	result["r_uncertainty_flow"][peak_mask] = (r_max - r_min) / 2	2✔
207	result["theta_uncertainty_flow"][peak_mask] = np.abs(theta_diff) / 2	2✔
208
209	return result	2✔

XENONnT / straxen / 10286081415

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