6D Phase Space Reconstruction: MENT-Flow Validation on Complex High-Dimensional Distributions

Table of Links
I. Introduction
II. Maximum Entropy Tomography
A. Ment
B. Ment-Flow
III. Numerical Experiments
A. 2D reconstructions from 1D projections
B. 6D reconstructions from 1D projections
IV. Conclusion and Extensions
V. Acknowledgments and References
B. 6D reconstructions from 1D projections
It is more difficult to design and evaluate high-dimensional numerical experiments. First, establishing reconstruction accuracy requires high-dimensional visualization or statistical distance metrics. We selected ground-truth distributions with clear high-dimensional structure and leveraged complete sets of pairwise projections and limited sets of partial projections (projections of slices) to aid the visualization. Second, we cannot determine the distance from the reconstructed distribution to the true maximum-entropy distribution without an analytic solution. We point to Fig. 2 as evidence that the entropy penalty can push the MENT-Flow solution close to the exact solution. We also continued to train an unregularized neural network on the same data to show that additional solutions can exist far from the prior. Third, for a given beamline and a fixed number of measurements, we do not yet know how to find the information-maximizing set of 6D phase space transformations. In 2:1 tomography, if the transformations are linear, the reconstruction quality is tied to a single parameter (the projection angle). There is no such connection in n:2 tomography when n > 3, as there is no obvious analog of the projection angle in these cases. Here, to demonstrate the method, we instead restrict our attention to 1D projections. A 1D projection axis can be specified by a point on the unit sphere; if the distribution is spherically symmetric, we hypothesize that the optimal projection axes are uniformly spaced on the sphere. In 2D, this leads to evenly spaced projection angles between 0 and π radians. In our numerical experiments, we approximated this condition by randomly sampling points from a uniform distribution on the sphere. The points will not be uniformly spaced, but in the limit of many projections, the reconstruction should converge to the true distribution [36]. Our first high-dimensional experiment, shown in Figs. 4-5, reconstructs a seven-mode Gaussian mixture distribution (a superposition of seven Gaussian distributions, each with a random mean and variance) from random 1D projections. Fig. 4 uses 25 projections and Fig. 5 uses 100 projections. This reconstruction used the same flow architecture as the 2D experiments. The NN architecture was changed to 2 layers of 50 units, still with tanh activation functions. We draw the following conclusions. (i) Normalizing flows can represent complicated 6D distributions far from the unimodal base distribution. All simulated measurements match the training data. Charged particle beams are often smooth and unimodal, so this example represents a challenging case. Therefore, flow-based models are likely sufficient for many applications in accelerator physics. (ii) MENT-Flow can simultaneously fit a large number of measurements. (iii) The entropy penalty works as intended. The entropy-regularized solution fits the data just as well as the NN solution but eliminates high-frequency terms in the distribution function. The MENT-Flow solution is much closer to the smooth prior. The Gaussian mixture distribution has little overlap between modes, so mismatch between the true and reconstructed distribution is obvious from low-dimensional views. Hollow structures in high-dimensional phase space are not always evident from low-dimensional views. As an example, measurements at the Spallation Neutron Source (SNS) Beam Test Facility (BTF) show spacecharge-driven hollowing in 3D and 5D projections of the 6D phase space distribution [1–3]. This motivates us to consider distributions with hidden internal structure. To this end, an n-dimensional “rings” distribution serves as the ground truth in Fig. 6-7; particles populate two concentric n-spheres with radii r2 = 2r1, and the radii are perturbed with Gaussian noise to generate a smooth density. The entropy-regularized solution maintains the spherical symmetry of the Gaussian prior, flattening and eventually inverting its radial density profile to fit the data.[6] The sliced views reveal an internal structure— a dense core surrounded by a low-density cloud—that MENT-Flow better approximates when measurements are scarce. In addition to injecting unnecessary correlations between planes, the unregularized solution ejects all particles from the core. Surprisingly, adding additional measurements does not solve the problem and generates two distinct modes in the reconstructed density. Using a different random seed to define the measurement axes can generate different patterns, but the hollowing and splitting just described are typical. Note that this internal structure is not obvious from the full 2D projections in the left column of Figs. 6-7.

[6] The one-dimensional projections in Figs. 6-7 are nearly Gaussian. Klartag [37] proved that almost all m-dimensional projections of an isotropic (no linear correlations) n-dimensional logconcave distribution function are nearly Gaussian when n ≫ m. Many distributions commonly used in accelerator modeling are log-concave, such as the n-dimensional Gaussian, Waterbag (uniformly filled ball), and KV (uniformly filled sphere) distributions. A practical implication of this theorem is that small fluctuations in the m-dimensional projections have a greater impact on the n-dimensional reconstructed distribution as n−m increases—for instance, completely inverting the density profile from peaked to hollow. Thus, we found that later training epochs can significantly change the distribution while only slightly decreasing the loss function. It follows that, for certain distributions, there may be some value of n − m for which n:m tomography is practically impossible.