Physics-constrained 3D Convolutional Neural Networks for Electrodynamics
We present a physics-constrained neural network (PCNN) approach to solving
Maxwell's equations for the electromagnetic fields of intense relativistic
charged particle beams. We create a 3D convolutional PCNN to map time-varying
current and charge densities J(r,t) and p(r,t) to vector and scalar potentials
A(r,t) and V(r,t) from which we generate electromagnetic fields according to
Maxwell's equations: B=curl(A), E=-div(V)-dA/dt. Our PCNNs satisfy hard
constraints, such as div(B)=0, by construction. Soft constraints push A and V
towards satisfying the Lorenz gauge.
Alexander Scheinker, Reeju Pokharel
## Physics-constrained 3D Convolutional Neural Networks for Electrodynamics

## Abstract

We present a physics-constrained neural network (PCNN) approach to solving Maxwell's equations for the electromagnetic fields of intense relativistic charged particle beams. We create a 3D convolutional PCNN to map time-varying current and charge densities \(\mathbf{J(r,t)}\) and \(\mathbf{\rho(r,t)}\) to vector and scalar potentials \(\mathbf{A(r,t)}\) and \(\mathbf{\phi(r,t)}\) from which we generate electromagnetic fields according to Maxwell's equations: \(\mathbf{B=\nabla\times A}\), \(\mathbf{E=-\nabla\phi-\partial A/\partial t}\). Our PCNNs satisfy hard constraints, such as \(\nabla\cdot\mathbf{B=0}\), by construction. Soft constraints push \(\mathbf{A}\) and \(\mathbf{\phi}\) towards satisfying the Lorenz gauge.

## Introduction

Electrodynamics is ubiquitous in describing physical processes governed by charged particle dynamics including everything from models of universe expansion, galactic disks forming cosmic ray halos, accelerator-based high energy X-ray light sources, achromatic metasurfaces, metasurfaces for dynamic holography and on-chip diffractive neural networks, down to the radiative damping of individual accelerated electrons [1-21].

Despite widely available high-performance computing, numerically calculating relativistic charged particle dynamics is still a challenge and an open area of research for large collections of particles undergoing collective effects in dynamics involving plasma turbulence [22], space charge forces [23,24], and coherent synchrotron radiation [25,26]. For example, the photo-injectors of modern X-ray free electron lasers such as the LCLS, SwissFEL, and EuXFEL and plasma wakefield accelerators such as FACET-II can produce high quality intense bunches with few picosecond rms lengths of up to 2 nC charge per bunch that are accelerated and squeezed down to lengths of tens to hundreds of femtoseconds [27-32]. At low energy near the injector, the 6D phase space (\(x\),\(y\),\(z\),\(px\),\(py\),\(p\),\(z\) ) dynamics of such bunches are strongly coupled through collective space charge (SC) forces. At higher energies, especially in bunch compressors where charged particle trajectories are curved through magnetic chicanes the dynamics are coupled through collective coherent synchrotron radiation (CSR).

A 2 nC bunch contains \(N\)\(\simeq\) 1.25 \(\times\) 10\({}^{10}\) electrons for which calculating exact individual particle to particle SC and CSR interactions is a computationally expensive O(\(N^{2}\)) process. For SC calculations, an O(\(N^{2}\)) process, such as the _SpaceCharge3D_ routine in the particle dynamics simulation code General Particle Tracer (GPT) may be necessary for intense low energy beams near the injector where the longitudinal (\(z\)) velocities of individual particles in a bunch have a large variation and are comparable to transverse (\(x\),\(y\)) velocities [33,34]. For relativistic particles, many conventional approaches for SC calculations greatly reduce the number of required calculations by utilizing particle-in-cell methods with macro-particles, such as the _SpaceCharge3DMesh_ routine in GPT. For CSR, relativistic state-of-the-art 3D CSR calculations still rely on a full set of point-to-point calculations [35].

A charged particle's electromagnetic Lagrangian is

\[L=-\frac{mc^{2}}{\gamma}+ev\cdot A-e\varphi,\qquad\gamma=\frac{1}{\sqrt{1-v^{ 2}/c^{2}}}, \tag{1}\]

where \(e\) is the particle's charge, \(c\) is the speed of light, \(v=|\mathbf{v}|\), and A and \(\mathbf{\phi}\) are the vector and scalar potentials, respectively, which define the magnetic (\(\mathbf{B}\)) and electric (\(\mathbf{E}\)) fields as

\[\mathbf{B=\nabla\times A},\qquad\mathbf{E=-\nabla\varphi-\frac{\partial A}{ \partial t}}, \tag{2}\]

For which the relativistic Lorentz force law is

\[\frac{d\mathbf{p}}{dt}=e(\mathbf{E+\nu\times B}),\qquad\mathbf{p=\gamma mv}. \tag{3}\]

The \(\mathbf{E}\) and \(\mathbf{B}\) dynamics are coupled and depend on current and charge densities \(\mathbf{J}\) and \(\mathbf{\rho}\) as described by Maxwell's equations

\[\nabla\cdot\mathbf{E=\frac{\rho}{\varepsilon_{0}},\qquad\nabla\cdot B=0}, \tag{4}\]

\[\nabla\times\mathbf{E=-\frac{\partial B}{\partial t},\qquad\nabla\times B= \mu_{0}\left(J+\varepsilon_{0}\frac{\partial\mathbf{E}}{\partial t}\right)}. \tag{5}\]

A typical approach to numerically solving Eq. 1-5 starts with initial charge \(\mathbf{\rho(x,y,z,t=0)}\) and current profiles \(\mathbf{J(x,\,y,\,z,\,t=0)}\) and their rates of change as well as any external electric and magnetic fields \(\mathbf{E}\)_ext_(\(x\),\(y\),\(z\),\(t=0\)), \(\mathbf{E}\)_ext_(\(x\),\(y\),\(z\),\(t=0\)) and their rates of change, which may be produced by mag- nets and radio-frequency resonant acceleration cavities as is typical in high intensity charged particle accelerators. The total electromagnetic fields are then calculated as the sum of the external fields and the self-fields produced by the current and charge densities themselves according to Eq. 4, 5. The initial fields apply a force on the particles causing a change in momentum and position, as defined by Eq. 3. The most computationally expensive part of the process is the calculation of the self-fields generated by the particle distribution.

In this work, we introduce a physics-constrained neural net- work (PCNN) approach to solving Maxwell's equations for the self-fields generated by relativistic charged particle beams. For example, for the problem of mapping current density \(\mathbf{J}\) to an estimate \(\mathbf{B^{{}^{\ast}}}\) of the associated magnetic field \(\mathbf{B}\) we build Eq. 2 into the structure of our NN and generate the vector potential \(\mathbf{A^{{}^{\ast}}}\), which defines the magnetic field as

\[\mathbf{\bar{B}=\nabla\times\bar{A}\Rightarrow\nabla\cdot\bar{B}=\nabla\cdot \left(\nabla\times\bar{A}\right)=0}, \tag{6}\]

which satisfies the physics constraint by construction.
Neural networks (NN) are powerful machine learning (ML) tools which can extract complex physical relationships directly from data and have been used for speeding up the studies of complex physical systems [36-43]. Incredibly powerful and flexible physics-informed neural networks (PINNs), which include soft constraints in the NN's cost function, have been developed and have shown great capabilities for complex fluid dynamics simulations [44], material science [45], for symplectic single particle tracking [46], for learning molecular force fields [47], and for large classes of partial differential equations [48-50].

For the problem of mapping current density \(\lambda\) to an estimate B\({}^{\wedge}\), the PINN approach is to train a neural network with a cost function defined as

\[C =w_{B}\left\|\iint\left[B-\beta\right]^{2}dV+w_{\nabla}\iint\left[ \nabla\cdot\beta\right]^{2}dV\right.\] \[=w_{B}\left\|B-\hat{B}\right\|_{2}+w_{\nabla}\left\|\nabla\cdot \hat{B}\right\|_{2}, \tag{7}\]

where the first term depends on magnetic field prediction accuracy and the second term penalizes violation of the physics constraint \(\nabla\cdot\text{B}^{\wedge}=0\), as shown in Figure 1. However, with soft PINN-type constraints there is no guarantee that the constraints are always satisfied, which is in contrast to the hard constraints implemented in our approach, which guarantee that constraints are not violated within numerical and finite discretization limits. Furthermore, when utilizing PINN-type soft constraints there is a tradeoff between the minimization of the two terms in Eq. 7 based on the choice of weights \(w_{\text{B}}\) and \(w_{\Delta}\). Intuitively this tradeoff can be understood by the fact that the easiest way for a neural network to satisfy \(\nabla\cdot\text{B}^{\wedge}=0\) is \(\text{B}^{\wedge}\equiv C\) for any constant \(C\). For hard constraints there is no such tradeoff, the cost function only penalizes field accuracy and the constraint itself is built into how the field is constructed. In our PCNN approach, our cost function is simply

\[C=\left\|B-\hat{B}\right\|_{2}, \tag{8}\]

and there is no tradeoff between reconstruction accuracy and physics constraint enforcement.

This is important because when simulating charged particle dynamics, great care must be taken to satisfy the physics constraints as defined by Eq. 1-5. It is very important to enforce well known beam properties such as phase space volume-preserving symplectic maps that satisfy Liouville's theorem so that the beam dynamics are self-consistent [51-56]. Results on physics informed NNs with hard constraints have mostly focused on fluid dynamics and climate modeling and are much more limited than PINN approaches [57-61].

## 3 Physics-constrained neural networks

In Figure 1 we summarize three NN approaches: 1) a NN approach without physics constraints, 2) a PINN approach with soft constraints, and 3) our PCNN approach. We demonstrate our PCNN method with numerical studies of relativistic (5 MeV), short (\(\alpha\)t = 800 fs), high charge (2 nc) electron bunches represented by \(N=50\) million macro particles. We utilize the charged particle dynamics simulation code General Parcticle Tracer (GPT) with 3D space charge forces [33,34]. The charged particle distributions were simulated for 1.2 ns with all of the data saved to a file at each \(\Delta\)t = 12 ps interval so that the beam was displaced 0.36 m over 100 saved steps.

The Figure 2 (A) shows the \(x\) and \(y\) trajectories of 10000 random particles sampled from the bunch distribution over the entire 100 saved steps as the beam is compressed by a 0.5 T solenoid magnet whose \(B_{2}\) field is shown in green. Only the first 75 steps, shown in black, were used for training and the final 25 steps were used for testing, shown in red. Figure 2 (B) shows the {\(x\),\(y\)} and {\(x\),\(z\)} projections of the electron bunch density at steps 0 and 74.

The training beam we have created is designed to have multiple length scales in order to help the trained PCNN generalize to new unseen distributions. We have created several closely spaced Gaussian bunches of varying \(\alpha\)x and \(\alpha\)y as seen in the (\(x\),\(y\)) projection of step 0. Furthermore as seen in the (\(x\),\(z\)) projection the beam has an overall bunch length of \(\alpha_{Z}=800\)\(\mu\)m with density fluctuations of various \(\alpha_{Z}\) along the length of the beam. By step 75 the beam has been over-compressed in the (\(x\),\(y\)) plane as seen by the (\(x\),\(y\)) projection and the beam density has started to spread in the \(z\) direction due to space charge forces.

At each time step we generate discrete versions of J, \(\rho\), E, and B by breaking up the 2.4 mm x 2.4 mm x 4.4 mm volume which is co-moving with the center of the beam into a 128x128x128 pixel cube with sides of length \(\Delta\)y = 18.9 \(\mu\)m, \(\Delta\)y = 18.9 \(\mu\)m, \(\Delta\)z = 34.6 \(\mu\)m, and averaging over all of the macroparticles in each cube. We compare the three neural network approaches to map J to B, as shown in Figure 1: 1) A standard NN using (8) as the cost function for training, 2) A PINN using (7) as the cost function, and 3) A PCNN us - in (8) as the cost function with the physics constraint built into the structure of the ML approach. The NN, PINN, and PCNN are able to achieve similar errors on the training data as they all use a similar 3D convolutional neural network (CNN) encoder-decoder architecture, as shown in Figure 3.

There is however an important distinction in terms of neural network size when comparing the NN, PINN, and PCNN approaches. The PCNN is actually smaller while achieving

Figure 1: Various NN approaches to generate \(a\) magnetic field \(B^{\wedge}\) (on estimate of B) from current density \(\lambda\) are shown.
better test data results and a much smaller violation of the physics constraint. For the NN and PINN all three components (\(\{\textit{{{\rm{}}}}_{x},\textit{{{\rm{}}}}_{y},\textit{{{\rm{}}}}_{z}\}\)) of J must be used as inputs to generate all three components (\(\{\textit{{{\rm{}}}}_{x}^{\prime},\textit{{{\rm{}}}}_{x}^{\prime},\textit{{{\rm{}}} }_{z}\}\)) of B\({}^{\textbf{{{\rm{}}}}_{\textbf{{{\rm{}}}}}}\). Therefore, the inputs and outputs of this 3D CNN are objects of size 128x128x128x3 and the input-output mapping of this 3D CNN, \(\textit{{{\rm{}}}}_{N_{b}}\), is given by

\[\left\{{{{\textit{{\rm{}}}}_{x}},{{\textit{{\rm{}}}}_{y}},{{\textit{{\rm{}}}}_{z }}}\right\}\to{{{\textit{{\rm{}}}}}_{B}}\rightarrow\left\{{{{\textit{{\rm{}}} }_{x}},{{\textit{{\rm{}}}}_{y}},{{\textit{{\rm{}}}}_{z}}}\right\}. \tag{9}\]

However, for the PCNN we are generating an estimate of **A**, which satisfies

\[\textit{{A}}(\textit{{r}},\textit{{t}})=\frac{{{{\textit{{\rm{}}}}_{0}}}}{{{ \textit{{4}}}{\textit{{\pi}}}}}\left[\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!
to a volume gives the partial derivative \(V\!\rightarrow\!\!V\!\times\!W_{\Delta}\!\!\times\!\!W\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial\! \times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial\!\times\!\!V \!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\!V \!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times \!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial\!\times\!\!V\! \times\!\!\partial\!\times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\! \partial\!\times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial \!\times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial \!\times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!V\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!V\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial\! \times\!\!V\!\times\!\!\partial\!\times\!\!V\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \partial\!\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \partial\!\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \partial\!\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\!\times\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \partial\!\times\!\!\partial\!\!\times\!\partial\!\times\!\!\partial\!\! \times\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\!\times\!\! \partial\!\times\!\partial\!\!\times\!\partial\!\!\times\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\partial\!\! \times\!\!\partial\!\times\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\! \partial\!\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\partial\! \times\!\!\partial\!\times\!\!\times\!\partial\!\!\times\!\!\partial\!\times\! \partial\!\!\times\!\!\partial\!\times\!\!\times\!\partial\!\!\times\!\! \partial\!\times\!\!\partial\!\times\!\partial\!\!\times\!\!\times\!\partial\! \times\!\!\partial\!\times\!\!\times\!\partial\!\!\times\!\partial\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\times\!\! \partial\!\times\!\!\partial\!\times\!\!\times\!\partial\!\times\!\! \partial\!\times\!\!\times\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\times\!\partial\!\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\times\!\partial\!\times\!\!\partial\!\times\!\! \times\!\partial\!\!\times\!\!\partial\!\times\!\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\times\!\!\partial\!\times\!\!\times\!\partial\!\! \times\!\!\partial\!\times\!\!\times\!\partial\!\times\!\!\partial\!\times\!\! \partial\!\times\!\!\times\!\partial\!\!\times\!\partial\!\times\!\!\! \times\!\partial\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\times\!\partial\!\! \times\!\partial\!\!\times\!\partial\!\times\!\!\!\times\!\partial\!\!\times\!\! \partial\!\times\!\!\times\!\!\partial\!\times\!\!\partial\!\times\!\!\! \times\!\partial\!\times\!\!\partial\!\!\times\!\!\partial\!\times\!\!\!\times\!\! \partial\!\times\!\!\partial\!\times\
Therefore our initial particle distribution can be thought of as an intense beam surrounded by a cube-shaped halo of diminishing density. This is due to the fact that we defined our initial beam in terms of Gaussian distributions without any hard cut- offs. Once this cube shaped region begins to travel through the solenoid it is rotated and squeezed, resulting in regions of non-zero and zero density that have sharp straight contours, as can be seen in the bottom part of Figure 5 and cause numerical problems for calculating derivatives. The most obvious mitigation for this would be to create a mask that cuts off all field calculations related to the beam beyond some minimal cutoff density. Despite this limitation, which each 3D CNN-based approach will suffer, the PCNN approach can be seen to be more accurate than NN without constraints and also than the PINN approach. In Figure 6 we compare PINN and PCNN predictions for two states of the beam, one within the training data set for which both are highly accurate and one beyond the training set where the accuracy quickly drops off. The next step is to add a prediction E\({}^{\alpha}\) of the electric field E. We generate \(\phi\) from \(\rho\) via a second neural network N\({}_{\phi}\) which gives the mapping

\[\rho\to N_{\varphi}\rightarrow\hat{\phi}. \tag{16}\]

As above, we approximate \(\partial/\partial t\) as

\[\frac{\partial A}{\partial t}=\frac{A(t+\Delta_{t})-A(t-\Delta_{t})}{2\Delta _{t}}+O\big{(}\Delta_{t}{}^{2}\big{)}, \tag{17}\]

where \(\Delta_{t}\)=1.2\(\times\)10-11. After \(\partial\Delta^{\alpha}/\partial t\) is calculated, a single forward pass for E\({}^{\alpha}\) is given by

\[\{\rho,J\}\rightarrow\big{\{}N_{\varphi},N_{A}\}\rightarrow\big{\{}\tilde{ \varphi},\tilde{A}\big{\}}\rightarrow\mathbf{\tilde{E}}=-\nabla\tilde{\varphi}- \frac{\partial\mathbf{\tilde{A}}}{\partial t}, \tag{18}\]

_as shown in Figure 7._

Because uncountably many non-unique choices of A and \(\phi\) generate the same E and B fields, we add the Lorenz gauge as a PINN-type soft constraint to the training cost function

\[w_{B}\big{\|}B-\hat{B}\big{\|}_{2}+w_{E}\big{\|}E-\hat{E}\big{\|}_{2}+w_{L} \big{\|}\nabla\cdot\mathbf{\tilde{A}}+\frac{1}{c^{2}}\frac{\partial\hat{\varphi} }{\partial t}\big{\|}_{2}, \tag{19}\]

which has the additional benefit that it introduces more data for the magnetic field calculation as the magnetic field is now informed by the Lorenz condition. Predictions for the entire 3D beam at step 1 by the Lorenz PCNN are shown in Figure 8. In Figure 9 we show the Lorenz PCNN-generated (B\({}^{\alpha}\), E\({}^{\alpha}\)) fields at just a single 2D (x, y) slice at various steps including those beyond the training data.

## 4 Discussion

Our final demonstration of the strength of building in hard physics constraints in the 3D CNN is a demonstration of its non-catastrophic failure when predicting the electromagnetic fields of two additional 2 nc beams that are very different from both the test and training data shown so

Figure 6: The top row is the initial beam state with the (x,y) projection of B\({}^{\alpha}\) as generated by the PINN and PCNN shown along with B and the differences plotted over the (x,y) projection of the beam’s charge density \(\rho\). For training data the methods perform equally well in field reconstruction accuracy. In the middle row we see the first step (76) beyond the training data set and an immediate drop in the accuracy of the PINN. In the third row the (x,z) projections of step (76) are shown, the roughness of the (x,z) projection shows intuitively how the PINN matches the B field in a mean squared error sense, but violates the constraint \(\nabla\cdot\textbf{B}=0\).
The first additional beam is three parallel electron beams, each of whose length is \(\alpha_{z}\approx 3mm\) which is similar to the overall length of the various bunches used in the training data. The three parallel beams are different from the training data in having empty space between the individual bunches. The second additional beam is a hollow tube of electrons with the same length as the three parallel beams, but whose topology is entirely different from anything the PCNN has seen so far. Figure 10 shows results of predicting the (**E**, **B**) fields for the three parallel and the hollow beams. As expected, the PCNN performs worse than previously, but is qualitatively very accurate in both E and B field prediction for the three parallel beams. The hollow beam is a much bigger challenge and much larger field errors are seen, but crucially, the predicted fields are still qualitatively correct in terms of direction and flow, with most of the error due to the wrong amplitude being predicted. We believe that this final result honestly shows the generality and strength of the PCNN approach and its limitations. We should not expect a trained CNN to predict well on an entirely unseen data set, this is a well-known problem known as distribution shift in the ML community in which NNs must be re-trained for inputs different from the training data set distribution. The fact that the PCNN produces reasonable outputs for inputs wildly different from that of the training data is a major strength of the approach. As Maxwell's equations are important for describing an extremely wide range of physical phenomenon the applications of such a method for electrodynamics applications are many. Here we will briefly touch on charged particle dynamics in high energy accelerators. There is a growing literature on utilizing ML-based surrogate models as virtual diagnostics in the particle accelerator community. For these approaches the NNs are typically trained as input-output maps utilizing experimental input data together with computationally expensive output data such as measuring a charged particle current density at one location of an accelerator and then running physics codes to map that to another location [41], or mapping input accelerator data directly to beam characteristics at other accelerator locations [39]. For such applications, the PCNN method can enable the development of much more robust real-time virtual diagnostics that satisfy physics constraints.

Another large family of applications is for accelerator design. For example, given a fixed input beam charge and current density distributions a beam line may be designed with various electromagnetic magnet and resonant acceleration components. For each design choice, such as distance between magnets or magnetic field strengths, high-fidelity physics-based models must be used to track the charged particle dynamics. With our approach, once a PCNN is trained for a family of input beam distributions, we have demonstrated that we can make accurate field predictions that respect physics constraints even as the beam is significantly changed by the application of external fields based on the accelerator's design. The next step of this work, which is beyond the scope of this paper and an ongoing effort, is to utilize our PCNN approach to quickly push particles and to confirm that the field predictions are accurate enough such that the particle dynamics are physically consistent. As we have already seen some slight numerical limitations as discussed above, this might push us to utilize even higher resolution discretization, such as 512\({}^{3}\) or 1024\({}^{3}\) pixel volumes, which remains to be determined. If this approach is able to provide physically consistent beam dynamics, even if they slightly violate constraints, this will be a fast and powerful way to zoom in on an optimal design estimate, after which more accurate slower physics-based simulations can be used for detailed studies.

## Conclusions

A robust PCNN method has been developed in order to explicitly take Maxwell's equations into account in the structure of generative 3D convolutional neural networks in order to more accurately satisfy physics constraints. Although this method is less general than the incredibly flexible PINN approach, in which any partial differential equation can be easily introduced as a soft constraint, the resulting physics constraints are more accurately respected. Furthermore, we have shown how to combine this PCNN approach with the PINN approach in our Lorenz CNN in which hard physics con- straints are enforced in the generation of the E and B fields and the soft penalty on violation of the Lorenz guage is added to the cost function in Equation 19.

Figure 7: The Lorenz PCNN generates the vector and scalar potentials and their associated electromagnetic fields.
Figure 8: (A) Electromagnetic fields are shown for all positions within the \(128^{3}\) pixel volume for normalized charge density \(\rho>0.0025\) for the first state of the beam. (B) We zoom in on only the part of the electron bunch which has the largest \(\sigma_{2}\) profile and show it from two angles. (C) Fields from only a single \((x,y)\) slice of the 3D volume are shown at two different angles.
## References

* [1] A. A.
## Acknowledgments

This work was supported by the U.S. Department of En- ergy (DOE), Office of Science, Office of High Energy Physics contract number 89233218CNA000001 and the Los Alamos National Laboratory LDRD Program Directed Research (DR) project 20220074DR. National Laboratory LDRD Program Directed Research (DR) project 20220074DR.
