Simulation-Based Inference for Direction Reconstruction of Ultra-High-Energy Cosmic Rays with Radio Arrays

TL;DR

This work introduces a simulation-based inference pipeline that couples a physics-informed graph neural network (GNN) to a normalizing-flow posterior within the Learning the Universe Implicit Likelihood Inference framework, making it ideally suited for upcoming experiments targeting highly inclined events.

Abstract

Ultra-high-energy cosmic-ray (UHECR) observatories require unbiased direction reconstruction to enable multi-messenger astronomy with sparse, nanosecond-scale radio pulses. Explicit likelihood methods often rely on simplified models, which may bias results and understate uncertainties. We introduce a simulation-based inference pipeline that couples a physics-informed graph neural network (GNN) to a normalizing-flow posterior within the Learning the Universe Implicit Likelihood Inference framework. Each event is seeded by an analytic plane-wavefront fit; the GNN refines this estimate by learning spatiotemporal correlations among antenna signals, and its frozen embedding conditions an eight-block autoregressive flow that returns the full Bayesian posterior. Trained on about $8,000$ realistic UHECR air-shower simulations generated with the ZHAireS code, the posteriors are temperature-calibrated to meet empirical coverage targets. We demonstrate a sub-degree median angular resolution on test UHECR events, and find that the nominal 68% highest-posterior-density contours capture $71\% \pm 2\%$ of true arrival directions, indicating a mildly conservative uncertainty calibration. This approach provides physically interpretable reconstructions, well-calibrated uncertainties, and rapid inference, making it ideally suited for upcoming experiments targeting highly inclined events, such as GRAND, AugerPrime Radio, and BEACON.

Figures (11)

• Figure 1: End-to-end SBI pipeline within the LtU-ILI framework: Antenna signals simulated with ZHAireS are compressed by a physics-informed Graph Convolutional Network (GCN) and plane-wavefront (PWF) model into a residual offset $\Delta\mathbf{k}$. Adding this residual to an analytic fit yields the three-dimensional arrival direction $\mathbf{k}$, forming an information bottleneck focused on geometry. A masked autoregressive flow transforms a Gaussian base distribution into the posterior $p(\mathbf{k}\mid\text{data})$. Each flow layer employs a Masked Autoencoder Distribution Estimator (MADE), with triangular masks enforcing causality, resulting in a strictly lower-triangular Jacobian that enables efficient autoregressive posterior density evaluation without explicit likelihood computations.
• Figure 2: Constructed graph representations of two simulated air-shower events using the $k$-nearest-neighbor (kNN) method on the GRANDProto300 geometry. Triggered antennas are shown at their $(x,y)$ locations; marker color encodes relative peak time (ns) and marker size scales inversely with time to highlight early arrivals. Edges connect temporally nearest neighbors: for each node we link to $k=\sqrt{N}$ antennas (clipped to $3 \le k \le 12$) with the smallest $|\Delta t|$, where $N$ is the number of triggered nodes. For illustration purposes, only the highest-weight subset of edges is rendered; thus some spatially close stations may appear unconnected if they are not close in time-of-arrival. This emphasizes causal (wavefront) connectivity rather than Euclidean proximity and can produce edge bundles roughly orthogonal to the arrival direction. All edges are used in training; the sparse rendering is for visualization only.
• Figure 3: Simulated GRANDproto300 ultra-high-energy cosmic ray event drawn from the held-out test set. Grey open circles represent antenna stations that did not trigger, while colored filled circles denote triggered stations. The size of filled circles decreases with increasing trigger order, and their color encodes the arrival-time delay $(\Delta t)$ relative to the earliest triggered antenna. This synthetic event is extracted from the held-out validation dataset and is used to evaluate the accuracy of our direction reconstruction procedure.
• Figure 4: Top: Joint posterior distribution for the mock GRANDproto300-like event shown in Fig. \ref{['fig:hittpattern']}. Shaded regions represent the 68% and 95% credible intervals for zenith and azimuth, with the true arrival direction marked by a red cross. Corresponding one-dimensional marginal distributions for zenith (top axis) and azimuth (right axis) are also shown. Bottom: Percentile-percentile (P–P) plots for zenith (left) and azimuth (right), constructed from $2,048$ test events. Empirical percentiles (blue) lie mostly above the diagonal line and its 95% bootstrap uncertainty band, indicating mildly conservative posteriors.
• Figure 5: TARP (Tests of Accuracy with Random Points) coverage plot for $1,560$ validation events (see Ref. lemos:2023). The diagnostic was performed directly on the full direction vector $\mathbf{k}$, rather than separately on zenith and azimuth angles. The blue curve represents empirical expected coverage calculated from $2,048$ posterior samples per event, the shaded region indicates a bootstrap-derived $1\,\sigma$ uncertainty band, and the dashed diagonal marks perfect calibration. The observed near-diagonal alignment confirms well-calibrated posteriors, exhibiting mildly conservative uncertainty estimates at credibility levels above $40\%$.