\section{Introduction}
\label{sec:intro}
The dynamics of aerodynamically coupled particles in turbulent fluids are of interest for astrophysical, atmospheric physics, experimental, and industrial applications \citep{toschi09}. The coupling between the fluid and particle phases modifies the turbulent cascade, and leads to clustering, effective diffusion, and collisions among the particles. The generic problem has numerous physically significant variations, depending upon the regime of aerodynamic drag (Stokes or Epstein), the compressibility of the turbulence, the relative mass in particles versus that in the fluid, and the amount of fluid volume excluded by the particles.

In the astrophysical domain, the response of solid particles to turbulence is important for multiple problems in protoplanetary disk physics and early phase planet formation. On large spatial scales, the ubiquitous substructure seen in dust continuum images of protoplanetary disks \citep{vdm13,alma14,andrews18,andrews20} is interpreted as concentrations of solids within gaseous vortices \citep{barge95} and axisymmetric pressure maxima \citep{whipple72}, structures which may themselves have formed as a consequence of disk turbulence \citep{johansen09,manger18}. On small scales, the collision velocity of small (micron to cm-sized) particles is typically determined by how the particles interact with turbulence in the gas \citep{weidenschilling93,johansen14}. For small particles the collision velocity induced by turbulence increases with size. Because dust particles fragment (or bounce) above some threshold velocity \citep{blum08}, this behavior leads to the prediction of a {\em maximum particle size} that can result from coagulation. On intermediate scales, the two-way momentum transfer between particles and gas in disks results in a linear instability, the {\em streaming instability} \citep{youdin05,krapp19,paardekooper20,yang21}, that may be responsible for the formation of planetesimals via gravitational collapse \citep{johansen15,simon16,li21,yang17,schafer17}.

The various manifestations of particle-gas interactions in protoplanetary disks can be studied using analytic methods \citep[e.g.][]{volk80,ormel07,squire18,krapp19,lin21}, and by simulations that treat the particle phase as either discrete Lagrangian particles or as a second fluid \citep[e.g.][]{pan11,ishihara18,bhatnagar18,krapp20,sakurai21}. Although powerful, these methods by no means suffice to answer all of the physically interesting questions that arise in planet formation studies. Numerical simulations can only model a small subset of the full range of spatial scales, leaving open questions as to how interactions on different scales combine \citep{chambers10,hartlep20}. Even when only a limited range of scales need to be studied, it is computationally expensive to run models that span the range of plausible size distributions for particles in protoplanetary disks. Machine Learning (ML) approaches may provide new theoretical tools to study these problems. In cosmology, ML architectures have been able to generate realistic predictions of gas properties from dark matter only simulations, and to interpolate reliably from grids of simulations that sparsely cover the full parameter space \citep{vn21}. By analogy, in planet formation it might be possible to extend the range of explicit simulations with an effective ML sub-grid model, and to make predictions for the response of arbitrary particle size distributions using a fairly small set of input simulations. Multiple ML architectures would likely be needed to realize these gains.

In this paper we present an initial investigation into the utility of established ML methods for predicting the spatial and statistical properties of particles in turbulent fluids. We consider the response of a single particle size to isotropic forced turbulence, in a periodic domain, with aerodynamic drag in the Epstein regime. This setup excludes the forces that lead radial and vertical epicyclic oscillations of free particles, which are physically important \citep{youdin07}. It is, however,  the simplest model problem that is relevant to studies of the rate and velocity of particle collisions on small scales in protoplanetary disks. Clustering in this regime is well-studied, and is understood to result from the interaction of particles with small-scale fluid vortices \citep[e.g.][]{johansen14}. This physical understanding does not, however, lead to a predictive model for the particle distribution. The collision velocity, on the other hand, is amenable to an analytic model \citep[several exist, the most common in disk applications being][]{ormel07}. We generate training data using a recently developed Lagrangian particle module within the {\tt Athena++} code (Yang et. al., in prep.), and use a U-Net network architecture \citep{ronneberger2015u} to learn the mapping between the three-dimensional fluid fields and gridded representations of the particle data. We then assess the fidelity of the predictions for both statistical properties of the particles (the strength of clustering and the relative velocity as a function of scale), and for their instantaneous spatial distribution.

The plan of this paper is as follows. In \S\ref{sec:simulations} we describe the {\tt Athena++} code, the implementation of Lagrangian particles within the code, and the setup of the forced turbulence simulations. The Machine Learning methods are described in \S\ref{sec:ML}. Sections \S\ref{sec:density} and  \S\ref{sec:normal_particle_velc_prediction} quantify the network performance for the particle density and velocity fields, and for statistical properties that would be relevant to the calculation of the rate and mutual velocity of particle collisions. We present some generalization tests in \S\ref{sec:generalization} and conclude in \S\ref{sec:conclusion}.

\section{Numerical simulations}
\label{sec:simulations}

\subsection{Physics of turbulent clustering}
\label{sec:phy_of_clustering}
To aid in the discussion of clustering statistics in later sections, we briefly introduce some relevant physics of turbulence and turbulence-driven dust dynamics.

In the standard description of fluid turbulence \citep{kolmogorov41,kolmogorov1991local}, there exist a range of scales, known as the inertial range, within which turbulent statistics depend only on the mean energy transfer rate $\langle \epsilon \rangle$. In this limit, the two point relative velocity $\langle \delta \vec{u} \cdot \delta\vec{u}\rangle$ at a separation $\ell$ scales as $\langle \epsilon \rangle^{2/3} \ell^{2/3}$. In the inertial range, dissipation due to viscosity is negligible and energy mainly cascades to smaller scales. The scale at which viscosity becomes important can be found by equating $\langle \epsilon \rangle $ to the energy dissipation rate $\nu \left[\delta u(\ell_d)/\ell_d\right]^2 \approx \langle \epsilon \rangle$ \citep{ching2014statistics}, from which one can obtain the
Kolmogorov dissipative length scale, timescale and velocity scale as $\ell_\eta = (\nu^3/\langle \epsilon \rangle)^{1/4}, \tau_\eta = (\nu/\langle \epsilon \rangle)^{1/2}$, and $u_\eta = (\nu \epsilon)^{1/4}$ respectively.

Now consider dust particles coupled to such a turbulent flow. Assuming that friction is the dominant force driving the dust, the equation of motion of a dust particle takes the form,
\begin{equation}
    \frac{d\vec{v}_p}{dt} = \frac{\vec{v}_g(p) - \vec{v}_p}{\tau_s},
    \label{eq:particle_eom}
\end{equation}
where $\vec{v}_p$ is the velocity of particles, $\vec{v}_g$ is the velocity of gas, and $\tau_s$ is the friction time characterizing the strength of coupling between the dust and gas. For dust particles in the Epstein regime\footnote{In this regime the drag force is {\em linear} in the velocity difference between the particle and the fluid, rather than quadratic as in the Stokes regime typically relevant to terrestrial fluids}, that is, particles with size $a_p$ much smaller than the mean free path $\lambda_g$ of the gas molecules $(a_p\ll \lambda_g)$, the friction time can be estimated as \citep{weidenschilling93},

\begin{equation}
    \tau_s = \left(\frac{\rho_d}{\rho_g} \right) \left(\frac{a_p}{c_s}\right),
    \label{eq:epstein_drag}
\end{equation}
where $\rho_d$ is the material density of the dust, $\rho_g$ is the fluid density, and $c_s$ is the local speed of sound of the gas.

To quantify the clustering behavior, $\tau_s$ should be compared with a characteristic time-scale of the turbulent flow. A commonly used time scale relevant for our purposes is the Kolmogorov time-scale $\tau_\eta$ \citep{sundaram_collins_1997,pan11,ishihara18}. It represents the turnover time of the smallest eddies. The ratio of these time-scales is the dimensionless Stokes number $\text{St} = \tau_s/\tau_\eta$\footnote{In protoplanetary disk studies, the ``Stokes number" is often defined instead using as a characteristic time-scale the dynamical time, $\Omega_{\rm K}^{-1}$, where $\Omega_{\rm K}$ is the Keplerian orbital frequency. Because our simulations do not include shear, we use the definition appropriate for fluid turbulence.}. For $\text{St} \ll 1$, the particle is well coupled to the gas at all scales and behave like a trace particle. With very little memory of its previous motion, the clustering will be mainly determined by the velocity divergence of the surrounding flow at one point in time. It has been shown that the particle velocity divergence $\sum_i \partial_i v_{p}^i \approx  \text{St}/\tau_\eta$ \citep{pan11}, where $i$ is the spatial index. This implies that for $\text{St} < 1$ the clustering intensifies with the Stokes number.

Conversely, when $\text{St} > 1$, turbulent fluctuate too rapidly for the particle to respond. This has two effects. First, as the friction time $\tau_s$ increases, the history of the particle has an increasingly large influence on the particle relative velocity. As shown in the formal solution of equation~\ref{eq:particle_eom} \citep{pan2013turbulence},
\begin{equation}
    \vec{v}_p(t) = \frac{1}{\tau_s} \int_{t_0}^{t} \vec{v}_g(\mathbf{X}(\tau),\tau) \exp(-\frac{t-\tau}{\tau_s}) d\tau,
\end{equation}
where $\mathbf{X}(t)$ is the position of the particle at some time $t$ and $t_0$ is a reference time chosen to be long enough (i.e. $t-t_0\gg \tau_s$) so that the velocity at time $t_0$ has long be "forgotten" by the particle. From the equation, we see that the particle's current velocity depends on the past history of gas velocity within a time window of $\approx \tau_s$. Hence, one should expect that for $\text{St} > 1$ the local gas velocity divergence, and also the $\text{St}/\tau_\eta$ scaling discussed above, should no longer be adequate for describing particle clustering.

Furthermore, particles with $\text{St} > 1$ respond to eddies of different turnover time selectively, giving rise to a characteristic clustering scale $\ell_{\tau_s}$ \citep{pan11}. Since $\tau_s > \tau_\eta$, some of the eddies have turnover time smaller than $\tau_s$. Flow on these scale is decoupled from the particle and instead give it random kicks, which inhibits clustering at those scales. On the other hand, the particles are well coupled to eddies with turnover time much larger than $\tau_s$, and obtain relative velocities similar to that of the flow. This depresses  clustering at large scale. Hence, the most significant clustering is expected at $\tau_{eddy} \approx \tau_{s}$. In the inertial range, we estimate $\tau_{eddy} \approx \ell^{2/3} \langle \epsilon \rangle^{-1/3}$ and obtain the  characteristic scale $\ell_{\tau_s} \approx \langle \epsilon \rangle^{1/2} \tau_s^{3/2} = \text{St}^{3/2} \ell_\eta$, where $\ell_\eta$ is the Kolmogorov length.

\subsection{Forced turbulence simulations}
\label{sec:forced_turbulence_simulation}
We use the \texttt{ATHENA++} magnetohydrodynamics (MHD) code \citep{stone20} to calculate the training data for our study.
All simulations are conducted using a uniform 3D Cartesian grid to represent the gas variables with a resolution of $256$ cells in each spatial direction.
We use the Harten--Lax--van Leer--Einfeld (HLLE) Riemann solver combined with a second-order accurate van Leer predictor-corrector scheme for time integration. Spatial reconstruction is done using the 3rd order piecewise parabolic method (PPM). We use periodic boundary conditions in all spatial directions.

The gas is initialized with a constant density of $\rho_\mathrm{gas} = 1$ and with all velocity components $v_i =0$. We use an isothermal equation of state,
\begin{equation}
    P_\mathrm{gas} = c_\mathrm{s}^2 \rho_\mathrm{gas},
\end{equation}
with the constant speed of sound $c_\mathrm{s}=1$.
To generate a turbulent velocity spectrum for the gas, we use the FFT turbulence driver module of \texttt{ATHENA++}, which continually feeds energy into the domain at the largest scales by adding large scale velocity perturbations to the flow. The scale of these perturbations is controlled by the parameters $k_\mathrm{min}$ and $k_\mathrm{max}$, denoting the smallest and largest wavenumber to be perturbed. (Here wavenumber is defined as $k=1/\lambda$, without a factor of $2\pi$ in front). Unless otherwise specified, we use $k_\mathrm{min}=1$ and $k_\mathrm{max}=5$ with simulation box size $L_x=L_y=L_z=1$, which allows the turbulent energy spectrum to develop naturally into a Kolmogorov type spectrum, where numerical dissipation occurs near the Nyquist frequency. We exclude data from an initial transient state from our ML dataset. The amount of energy injected per perturbation interval (which we set to be equal to one simulation time-step) is given by the energy injection rate $\epsilon_\mathrm{inj}$ and is set to $\epsilon_\mathrm{inj} = 0.005$ (except where specified). The RMS gas velocity calculated from the simulation is $v_{g}^\text{rms} \approx 0.104$, which is $\approx 10\%$ the speed of sound. The turbulence, although computed with a compressible code, is thus for practical purposes almost incompressible. As we are investigating subsonic turbulence, we enforce solenoidal turbubulence driving by suppressing the compressive component of the velocity perturbation.

\subsection{Lagrangian dust particles}
\label{sec:lagrangian_dust_particles}
We use the Lagrangian dust particle module of the \texttt{ATHENA++} code (Yang et. al., in prep.) to simulate the particle component in the simulations. In all our simulations, $256^3$ dust particles are used.
For simplicity, we only treat the drag force enacted onto the dust particles by the gas and neglect the back-reaction drag force enacted by the particles onto the gas, which is a good approximation in cases with low dust-to-gas mass ratios. To interpolate between the gridded values of the gas component and the particle positions, the module uses the triangular shaped cloud (TSC) algorithm.
Unless otherwise specified, we use a particle friction time of $\tau_s = 0.1$.

From the simulated output, we estimate the turbulent scales of the fiducial test case to be: $\langle \epsilon\rangle\approx 3.17\times 10^{-3}$, $\nu \approx 1.95 \times 10^{-5}$, $\ell_\eta \approx 1.24\times 10^{-3}$, $\tau_\eta \approx 7.84\times 10^{-2}$ and $u_\eta \approx 1.57\times 10^{-2}$, all in the relevant simulation units. Similarly, the particle clustering length is estimated to be $\ell_{\tau_s} \approx 1.58\times 10^{-3}$. The energy injection rate $\langle \epsilon\rangle$ and eddy turnover time $\tau_\eta$ are estimated from the third-order structure function and strain-rate tensor respectively; The particle friction time is taken directly from simulation setup, and the rest are obtained from their relationship with these three variables.
\section{Machine learning methods}
\label{sec:ML}

\subsection{Dataset}
Snapshots of the simulation are output at intervals chosen such that the temporal correlation between adjacent frames is negligible. The gas density field $\rho_g$, gas velocity field $\vec{v}_g$, particle density field $\rho_p$ and particle velocity field $\vec{v}_p$ are used to train and validate the network. All of the fields are gridded on a $256^3$ mesh. (The gas fields are natively gridded; the particle data is mapped to the same grid.) The first two serve as inputs and the later as outputs of the network.
\subsection{Data preprocessing and Normalization}
The following data preprocessing and normalization steps are applied to the input and output fields to aid the training.
\subsubsection{Input fields}
The gas density $\rho_g$ is log normalized,
\begin{equation}
    \rho_g^\ast = \log(\rho_g)
    \label{eq:rhog_norm},
\end{equation}
where $\rho_g^\ast$ is the normalized gas density and $\rho_g$ is the gas density
The log normalized density field helps separate gas density of various scales, allowing the network to better perceive the input field. No preprocessing is applied to the gas velocity $v_g$.
\subsubsection{Output fields}
The particle density field $\rho_p$ and velocity field $v_p$ are known to be non-gaussian \citep[e.g.][]{ishihara2018dust}. Therefore, data normalization and preprocessing techniques are employed to aid the training of the neural network.

For the particle density field $\rho_p$, after experimentation, the following normalization is used,
\begin{equation}
    \rho^\ast_p = \ln(\exp(\rho_p/\rho_0)-1).
    \label{eq:rhop_norm}
\end{equation}
Here, $\rho^\ast_p$ is the normalized particle density and $\rho_0=0.14$ is an empirically chosen density scale that results in a  distribution of $\rho^\ast_p$ that is most nearly symmetric. This value corresponds roughly to the average particle density. The mapping between $\rho_p^\ast$ and $\rho_0$ is one-to-one, and it can be easily inverted using a softplus-like function:
\begin{equation}
    \rho_p = \rho_0\ln(\exp(\rho_p^\ast)+1).
    \label{eq:rhop_inv_norm}
\end{equation}
Similarly, to prevent the bulk gas velocity $v_g$ from masking the fine details of the particle velocity $v_p$, we choose to train our network (referred to later as UNET-R) on the relative velocity $v_{rel}$ instead of the particle velocity $v_p$,
\begin{equation}
    \vec{v}_{rel} = \vec{v}_p - \vec{v}_g.
    \label{eq:vp_norm}
\end{equation}
This is intended to help the network to focus on the physically interesting deviation of the particle velocity $v_p$ from the gas velocity $v_g$, instead of learning primarily just the background flow. To evaluate the effectiveness of this strategy, we also trained the network directly on $v_p$ (later, UNET-V) and compared it against that on $v_{rel}$. Details of this comparison can be found in \S\ref{sec:normal_particle_velc_prediction}.

\subsection{Model and training}

\subsubsection{Model}
\begin{figure*}
    \centering
    \includegraphics[width=0.9\textwidth]{Figures/dustnet-architecture.pdf}
    \caption{The architecture of the U-Net. The thickness of individual layers corresponds roughly to the number of channels in each layer and the width and height corresponds to the spatial size. The color scheme corresponds to: convolution layers of kernel size 3 and stride 1 \textbf{(orange)}, Leaky ReLU activation with batch normalization \textbf{(deep orange)}, downsampling layers using convolution of kernel size 2 and stride 2 \textbf{(red)} and upsampling layers using transposed convolution of kernel size 2 and stride 2 \textbf{(purple)}. Residue connection is implemented for all convolution layers except those used in up and down sampling. Similar to other U-Net architectures, large scale information is incorporated by narrow down and concatenating compression layer outputs with the upsampled outputs.}
    \label{fig:unet_architecture}
\end{figure*}

The network architecture U-Net is used. First developed by \cite{ronneberger2015u} for biomedical image segmentation and later generalized by \cite{milletari2016v} to handle volumetric data, U-Net is a "fully convolutional network" that is capable of taking a multichannel, 3D input block and producing an output block of similar spatial size but with a different number of channels. It exhibits translational equivariance, such that its predictions depend only on the relative positions of the grid points \citep{long2015fully}. This is desirable because it is consistent with the physical picture of homogeneous turbulence \citep{ching2014statistics}, and because it allows the network to be applied to input fields of different spatial sizes. U-Net has recently been successfuly employed in cosmology for tasks including super-resolution (generating high resolution predictions from lower resolution simulations) \citep{schaurecker2021super}, and prediction of galaxy distributions from the dark matter density field \citep{yip2019dark}. This motivates us to apply the architecture to our problem.

The network is implemented using the \textbf{map2map} neural network emulator using the architecture specified in Fig.~\ref{fig:unet_architecture}. The network consists of two major parts, the "compression part" and the "decompression part". The former consists of successive convolutional and downsampling layers that compress information of the input field into a smaller size latent representation; while the latter consists of successive levels of convolutional and up-sampling layers that expand the latent information into the output block. Information extracted in different levels of the network is incorporated by concatenating outputs of the compression part and those of the expanding part on the same level.

\subsubsection{Training}

The simulation frames, 252 in total, are randomly divided into training and validation sets in a 4:1 ratio. Training is done using only the frames in the training set.

The entire training set is iterated through during each epoch of training. For each frame in the set, we divide the input fields into eight octants of size $128^3$ so that they can fit onto one GPU.  Then, we pad them periodically by $5$ grids on all sides so that it is $148^3$ in size. The padding, as discussed by \cite{schaurecker2021super}, allows the network to preserve full translational symmetry and generate outputs without boundary artifacts. This results in an input field of dimension $(4,148,148,148)$, with three of its channels coming from the spatial components of the gas velocity $v_g^{(x/y/z)}$ and one from the gas density $\rho_g$. Discrete rotation and parity transforms on the data are also applied to the input block using map2map. This helps the network to capture the isotropic nature of the turbulent flow. The input enters the U-Net and produces an output field of size $(4,108,108,108)$: One of the channels represents particle density $\rho_p$, and the rest represent particle relative velocity $\vec{v}_{rel}$. Finally, the whole $(4,256,256,256)$ field is reconstructed from the output patches produced by map2map.

For the loss function, we choose to minimize the mean square error (MSE) loss over all output channels. Weight decay is also used. The loss function, which is minimized using the Adam optimizer \citep{kingma2014adam}, takes the following form,
\begin{equation}
    \mathcal{L} = \left[\frac{1}{\mathcal{B}}\sum_{i=1}^{\mathcal{B}}(\mathbf{y}_\text{tgt} - \mathbf{y}_\text{out}^{(i)})^2\right] + \lambda \norm{\mathbf{w}}^2,
\end{equation}
where $\mathbf{y}_\text{tgt},\mathbf{y}^{(i)}_\text{out}$ represents the target and output fields, $\mathcal{B}$ is the batch size, and $\lambda$ is the weight decay and $\norm{\mathbf{w}}$ is the square norm of weights. Unless otherwise specified, all training runs in this paper use learning rate $10^{-4}$, weight decay $\lambda= 10^{-4}$, and batch size $\mathcal{B}=1$, for around 400 epochs until the results are stable.

After training, we evaluate the model by comparing the output with the target over a variety of statistics. The results are summarized in \S\ref{sec:density} and  \ref{sec:normal_particle_velc_prediction}.

\section{Prediction of particle density field}
\label{sec:density}

\subsection{Visual Comparison}
\label{sec:rhop_visual_comp}
Using the trained network, we output the predicted particle density field $\rho_p$ and compare it to the target on a frame-by-frame basis. Snapshots of selected frames are presented in Fig.~\ref{fig:rhop_slice_comp_normal}. Because the densities span a wide range, a $\log_{10}$ scale color map is chosen to aid visualization. We also show the normalized cross correlations\footnote{The normalized cross correlation of two function $f,g$ is defined to be $\left[(f\star g)(r)\right]/\sqrt{\left[(f\star f)(0)\right]\left[(g\star g)(0)\right]}$, which is a normalized inner product of $f,g$. For volumetric data, the result is radially binned.} of these frames, defined as in Fig.~\ref{fig:rhop_cross_corr}.

As shown in Fig.~\ref{fig:rhop_slice_comp_normal}, both machine learning outputs (UNET-V, UNET-R) produce density distributions that are similar to the ground truth (ATHENA++). Locations with voids in the density distribution align, and the filamentary structures of the dust are captured. However, both predicted outputs appear to be less sharp than the target field, indicating the algorithm has smoothed out the densities. The density at the locations of peaks is also less pronounced, indicating an under-prediction in those areas.
These observations are borne out in Fig.~\ref{fig:rhop_cross_corr}, where we see that the peak values of normalized cross correlations for both UNET-V and UNET-R are only $\approx 0.6$. The larger peak values for UNET-R show that UNET-R, trained on the relative velocity, is slightly better at predicting $\rho_p$ than UNET-V.

\begin{figure*}
    \centering
    \includegraphics[width=\textwidth]{Figures/rhop_slice_comp_normal.pdf}
    \vskip-0.75cm
    \caption{Particle density field slice $\rho_p$ from ATHENA++ (left), UNET-V (center, a network trained on the relative velocity between particles and gas) and UNET-R (right, a network trained on the absolute particle velocity), for a typical frame. Qualitatively, there is good agreement between the predicted structures and the ground truth, though the network outputs are visually less sharp than the target simulation data.}
    \label{fig:rhop_slice_comp_normal}
\end{figure*}

\begin{figure}
    \centering
    \includegraphics[width=0.4\textwidth]{Figures/rhop_cross_corr.pdf}
    \caption{Normalized cross correlation of $\rho_p$ of target field (ATHENA++) with itself, UNET-R and UNET-V. The peak values of the statistics are: 1.0, 0.622 and 0.613 respectively.}
    \label{fig:rhop_cross_corr}
\end{figure}

\subsection{Radial Distribution Function (RDF)}
\label{sec:normal_stat_rdf}
We compute the particle radial distribution function (RDF) $g(r)$ of the simulated and U-Net produced outputs to study their difference in spatial distributions more carefully. Let $\bar{n}$ be the average dust density. The RDF is defined so that the average number of dust particles $P(r)$ inside a volume element $dV$ at a distance $r$ relative to a reference particle is given by,
\begin{equation}
    P(r) = \bar{n} g(r) dV.
    \label{eq:radial_distribution_func}
\end{equation}
Although this definition can be applied directly to compute the RDF, a faster way of obtaining it on a grid is available. First, we compute the particle density correlation function (DCF) $\xi(r)$, which is defined as,
\begin{equation}
    \xi(r) = \frac{1}{\bar{n}^2} \langle \left(n(\vec{x})-\bar{n}\right) \left(n(\vec{x}+\vec{r})-\bar{n}\right)\rangle .
    \label{eq:density_correlation_func}
\end{equation}
Then we use the relation $g(r) = 1 + \xi(r)$ between the two quantities to obtain the RDF \citep{shaw2003particle}. This allows for an efficient computation, as  Eq.~\ref{eq:density_correlation_func} can be computed quickly using the fast fourier transform (FFT). Unless otherwise specified, all RDFs computed in this paper are based on this method. The use of grid-based metrics, rather than metrics computed directly from the particle distribution, differs from most prior work. Differences in results that arise between the method based on Eq.~\ref{eq:radial_distribution_func}, and that based on Eq.~\ref{eq:density_correlation_func} for the simulated particle density, are discussed in Appendix~\ref{sec:grid_stats_vs_part_stats}.

In the left-most panel of Fig.~\ref{fig:rdf_grouped}, we show the RDFs of $\rho_p$ produced by the simulation (ATHENA++), and by the U-Nets (UNET-V, UNET-R). Recall, UNET-V is the U-Net trained directly on $v_p$ and UNET-R on $v_{rel}$. The features of the RDFs reflect what we see in the slice plots. The simulation and U-Nets agree well for $r \gtrsim 0.01$, indicating the large scale features are captured. For $r\approx 7\times10^{-3}$, a dip of roughly $5\%$ is observed, indicating the underestimation of $\rho_p$ at small scales. The over-prediction for $r\gtrsim 7\times 10^{-3}$ might be due to the smoothed density around the peaks. Overall, the error in RDF statistics is bounded by $\pm 5\%$. Although we have picked only one frame for discussion, the features observed are general. Discussion on the consistency across frames can be found in Appendix~\ref{sec:consistency_across_frames}.

\begin{figure*}
    \centering
    \includegraphics[width=\textwidth]{Figures/rdf_compare_all.pdf}
    \caption{Predicted versus simulated strength of particle clustering as a function of spatial scale, quantified via the radial distribution function statistic (RDF). From left to right, the RDF plots show different simulation setups. The left-most panels show the fiducial setup, while the remaining panels show the results of generalization tests. (See \S\ref{sec:forced_turbulence_simulation},\ref{sec:lagrangian_dust_particles} and \S\ref{sec:generalization} for details.)  {ATHENA++} shows the RDF computed directly from the simulation, {UNET-R} is the U-Net output trained with relative velocity between gas and dust, and {UNET-V} is the U-Net output trained directly on dust velocity. The first row compares the RDF of the  simulation output and the U-Net prediction. The second row shows the ratio between the outputs and the target to better quantify the error. The zoom-in plot in the lower-right panel shows the peak of the error at small scales.}
    \label{fig:rdf_grouped}
\end{figure*}

\section{Prediction of particle velocity field}
\label{sec:normal_particle_velc_prediction}

\subsection{Visual Comparison}
In Fig.~\ref{fig:vrel_slice_comp}, we display slices of the three spatial components of the relative velocities field $v_{rel}$ produced by simulation and the U-Nets. Recall that the particle velocity $v_p$ can be thought of as having two contributions: (i) a gas velocity term $v_g$, which represents the velocity of the underlying flow driving the dust, and (ii) a relative velocity term $v_{rel} = v_p - v_g$, which measures the deviation of the particle velocity from the gas velocity due to the imperfect coupling between them. To prevent $v_g$ from overshadowing $v_{rel}$ in visual comparisons, we show $v_{rel}$.

From Fig.~\ref{fig:vrel_slice_comp}, we see that both UNET-V and UNET-R produce $v_{rel}$ fields that are qualitatively similar to that of the ATHENA++ ground truth. It is difficult to tell which is better solely by visual comparison. For example, at around $(v_{rel}^{3},0.5,0.6)$, UNET-V produce a sharper output for places where that component peaks, whereas at around $(v_{rel}^{3},0.1-0.3, 0.6-0.9)$, UNET-R captures the zeros of the field better. This motivates us to use more sophisticated statistics in \S\ref{sec:rel_velc_normal} and \S\ref{sec:rel_velc_rad_normal}, to assess which of them performs better.
\begin{figure*}
    \centering
    \includegraphics[width=0.90\textwidth]{Figures/velc_slice_comp.pdf}
    \caption{Prediction of the spatial distribution of particle relative velocity ($v_g-v_p$), in domains of isotropic forced turbulence. The figure compares the ground truth, obtained with {\sc ATHENA++}, with predictions from two deep learning models: UNET-V trained directly with particle velocity $v_p$, and UNET-R with relative velocity $v_{rel}$. The columns represent different setups (ATHENA++, UNET-V, UNET-R) and the rows show the 3 components of the velocity field. Each panel is a slice through the three-dimensional domain.}
    \label{fig:vrel_slice_comp}
\end{figure*}

\subsection{Relative velocity $\langle w^2 \rangle^{1/2}$}
\label{sec:rel_velc_normal}
The relative particle velocity $\langle w^2 \rangle^{1/2}$ is a commonly studied statistic of particle clustering in turbulence. Following \citet{ishihara2018dust}, we define it as,
\begin{equation}
    \langle w^2 \rangle  = \frac{1}{N_p}\sum_{|\vec{r_1} - \vec{r_2}| = r} |\vec{V}_1 - \vec{V}_2|^2,
    \label{eq:rel_vel}
\end{equation}
where $r$ is the separation between the particle pairs, $\vec{V}_1$, $\vec{V}_2$ are the velocities of the particles located at $\vec{r}_1,\vec{r}_2$ and $N_p$ is the number of pairs of particles.

To compute Eq.~\ref{eq:rel_vel} on the grid, we convert the average into a weighted one using the relevant particle density $\rho_p$ as weight. More precisely, let $p$ be a grid cell and $\text{C}_p(r_1,r_2)$ be the set of all grid cells whose center point lies inside the spherical shell with inner and outer radius $r_1$ and $r_2$ that is centered at $p$. Then a grid average of some quantity $Q$, weighted using particle density $\rho_p$, is defined as,
\begin{equation}
    \langle Q\rangle^{\rho_p}_\text{grid} = \frac{\sum_{q\in C_p(r,r+\Delta r) }\rho_p(q) Q(q)}{\sum_{q\in C_p(r,r+\Delta r)} \rho_p(q)},
\end{equation}
where $\rho_p(q)$ and $Q(q)$ are understood as the grid values at cell $q$. Using this definition, we write $\langle w^2_\text{tgt}\rangle^{1/2} = \left[\langle w^2_\text{tgt}\rangle^{\rho^\text{tgt}}_\text{grid}\right]^{1/2}$ for the ground truth relative velocity and $\langle w^2\rangle^{1/2}=\langle w^2_\text{out}\rangle^{1/2} = \left[\langle w^2_\text{out}\rangle^{\rho^\text{out}}_\text{grid}\right]^{1/2}$ for the output velocity field.

Unlike grid RDFs, computing these quantities on the grid are relatively expensive. To see that, notice if we were to compute the statistics for the whole grid $r\in [256^{-1},1]$, we would need to check every grid point against every other grid point. If $N$ is the number of grid points along one axis, this would imply a complexity $\approx (N^3)^2=N^6$.
To reduce the computational cost, we restrict $r\in[256^{-1},10^{-1}]$ for all relative velocity statistics. This remains practical because for many applications it is the accuracy of relative velocities at small scales that is the most important. Furthermore, from the estimated clustering scale in \S\ref{sec:lagrangian_dust_particles}, we see that the upper bound of the radii $r=10^{-1} \gg \ell_\eta \sim 10^{-3}$, so the relative velocity curve for $r>10^{-1}$ is likely not relevant for our application. Again, possible errors introduced by grid averaging are explored in Appendix~\ref{sec:grid_stats_vs_part_stats}.

We plot $\langle w^2\rangle^{1/2}$ and its error as a function of $r$ in Fig.~\ref{fig:vp_rel_grouped}. The upper left panel shows that the particle velocity fields produced by UNET-V and UNET-R have a similar shape to ground truth (ATHENA++), and the lower left panel shows that both models have their error bounded by approximately $15\%$ above and $7\%$ below. For $r \gtrsim 2.5 \times 10^{-2}$, both models under-predict the relative velocity, while for $r\lesssim 2.5 \times 10^{-2}$ both over-predict. However, UNET-R's absolute error is always less than UNET-V's at all length scales. In particular, the errors in UNET-R's prediction are $\approx 2.5\%$ less than UNET-V's at the smallest and largest scales. Furthermore, the region around the point where UNET-R's error curve crosses unity, especially that toward the tail of the curve, is much flatter than that of UNET-V's. This indicates the error varies less with separation in UNET-R.

The results above support our hypothesis that targeting the relative velocity aids in the network training. UNET-R can capture the radial dependence of relative velocity in the predicted particle velocity field within a maximum error of approximately 10\%. Although this error is systematic and scale-dependent rather than random, we consider this to be good performance.

\begin{figure*}
    \centering
    \includegraphics[width=\textwidth]{Figures/vrel_compare_all.pdf}
    \caption{Comparison of relative velocity between dust particles $\langle w^2 \rangle^{1/2}$ for different simulation setups (see \S\ref{sec:forced_turbulence_simulation},\ref{sec:lagrangian_dust_particles} and \S\ref{sec:generalization} for details). {ATHENA++} represents the relative velocity computed from the simulation, {UNET-R} is the U-Net output trained with relative dust velocity, {UNET-V} is the U-Net output trained directly on dust velocity, and {VG} is the average gas velocity weighted by particle density (i.e. the relative velocity curve one would have obtained if the particles follow perfectly the flow). The left-most panels show the fiducial results, the other panels show the results of generalization tests.}
    \label{fig:vp_rel_grouped}
\end{figure*}

\subsection{Relative radial velocity $\langle w^2_\text{r} \rangle^{1/2}$}
\label{sec:rel_velc_rad_normal}
Similarly, the root-mean-square relative radial velocity $\langle w_\text{r}^2\rangle^{1/2}$ is defined as,
\begin{equation}
    \langle w_\text{r}^2 \rangle  = \frac{1}{N_p}\sum_{|\vec{r_1} - \vec{r_2}| = r} \left(\left(\vec{V}_1 - \vec{V}_2\right)\cdot \hat{\eta}\right)^2,
\end{equation}
where all the symbols retain the same meaning as in Eq.~\ref{eq:rel_vel}, and $\hat{\eta} = (\vec{r}_1 - \vec{r}_2)/|\vec{r}_1 - \vec{r}_2|$. Using the same weighting procedure as in \S\ref{sec:rel_velc_normal}, we again computed $\langle w^2_\text{r}\rangle$ on ATHENA++, UNET-V and UNET-R data. The results are shown in Fig.~\ref{fig:vp_rel_rad_grouped}.

In the figure, we see that both the relative particle velocity $\langle w_{r}^2\rangle^{1/2}$ predicted by UNET-R and UNET-V have similar shape to the target. However, the performance of the U-Nets varies at different scales. For $r\gtrsim 2.5\times 10^{-2}$ both models perform similarly and predict a radial velocity around $\approx 2.5\%$ above and $\approx 5\%$ below the target respectively. The two models over-predict at scales $r\lesssim 4\times 10^{-2}$, being $\approx 10\%$ and $\approx 12.5$ larger respectively. Unlike $\langle w^2\rangle^{1/2}$, the error curves of the models do not intersect at unity. However, UNET-R's error remains  smaller than UNET-V's across most scales, except for a small region around $r\approx 3\times 10^{-2}$.

Also, while both UNET-R and UNET-V overpredict the radial velocity at small scales, that of UNET-V is much closer to the curve labeled VG---the average gas velocity weighted by particle density---than UNET-R. This indicates that UNET-V is essentially outputting the gas velocity $v_g$ as the particle velocity at small scales, instead of learning the particles' deviation from the gas flow. This further demonstrates that training on $v_{rel}$ is effective in helping the network to single out how the particle velocity differs from the background flow.

Finally, unlike the relative velocity, which shows a flat part at intermediate scales, there are no obvious flat region on UNET-V's error curve. This suggests the network in general predicts the direction of $v_g$ less accurately than its amplitude.

\begin{figure*}
    \centering
    \includegraphics[width=\textwidth]{Figures/vrel_rad_compare_all.pdf}
    \caption{Comparison of relative radial velocity $\langle w^2_r \rangle^{1/2}$ for different simulation setups (see \S\ref{sec:forced_turbulence_simulation},\ref{sec:lagrangian_dust_particles} and \S\ref{sec:generalization} for details). Refer to Fig.~\ref{fig:vp_rel_grouped} for the description on the abbreviated labels.}
    \label{fig:vp_rel_rad_grouped}
\end{figure*}

\subsection{Error estimation}
The predicted particle density and velocity fields, and their associated statistics, show good but not perfect agreement with the ground truth derived from the simulation data. The errors are systematic and scale-dependent. It is of interest to ask: how much of this error is due to uncertainty in the deep learning model? Addressing this question is difficult, and various approximate methods have been used in the machine learning community. We have experimented with two popular frameworks, (i) Monte-Carlo Dropout \citep{gal2016dropout,zhu2017deep} and (ii) SWA-Gaussian (SWAG) \citep{maddox2019simple}. Details are given in Appendix~\ref{sec_errors}. In summary, we find evidence using the SWAG approach that model uncertainty has a larger effect on the velocity predictions than on the density predictions. However, the magnitude of the estimated model uncertainty is significantly smaller than the size of the discrepancy between the model predictions and the ground truth. The source of the discrepancy thus remains unclear, but it is clearly of a systematic nature.

\section{Generalization Tests}
\label{sec:generalization}

Our deep learning model, trained on a single simulation of forced turbulence with a specified set of parameters, cannot be expected to generalize to distinctly different regimes of particle-gas coupling. However, it {\em may} have learned enough about how particles interact with turbulent fluids to behave robustly when confronted by modestly different turbulent conditions. To test our model's ability to generalize, we have designed three test cases with physical parameters different from the fiducial case discussed so far. The test cases are: (1) turbulence driven with a different range of cut off wavenumbers $k_{min},k_{max}$; (2) turbulence driven using a different energy injection rate of $\langle\dot{e}\rangle = 5\times 10^{-4}$; (3) particles with stopping time $\tau_s$ 100 times smaller than the original simulation. ATHENA++ simulations with these different parameters, but otherwise similar to the fiducial run, were performed for each test case.

For each of the test cases, we feed their simulation frames into the UNET-R trained from the fiducial run without any additional training, and produce field predictions for $\rho_p$ and $v_p$. Then, using the same statistics code in \S\ref{sec:normal_stat_rdf}, \S\ref{sec:rel_velc_normal} and \S\ref{sec:rel_velc_rad_normal}, we compute the relevant statistics from the field outputs and compare them with the ground truth.

\subsection{Different wavenumber cut-off}
\label{sec:generalize_k_cutoff}
In this generalization set, different minimum and maximum cutoff of driving wavenumbers $k$ are used. Instead of the original $(k_{min},k_{max})=(1,5)$ combination, we remove the largest scale driving and use the combination $(3,8)$. The resulting gas velocity power spectrum is illustrated in Fig.~\ref{fig:vg_power_spectrum}. The dip in power at around $k=1-3$ confirms that the new flow has less power in large scales compared to the original flow.

\begin{figure}
    \centering
    \includegraphics[width=0.45\textwidth]{Figures/vg_power_spectrum.pdf}
    \caption{Energy spectrum of driving flow $v_g$ using $(k_{min},k_{max})=(1,5)$ (fiducial) and $(k_{min},k_{max})=(3,8)$ (generalization test). The color bands illustrate the range of $k$ used in driving the flow for each case and the line shows the resulting power spectrum of the gas kinetic energy}
    \label{fig:vg_power_spectrum}
\end{figure}

The RDF is computed and plotted in the second column in Fig.~\ref{fig:rdf_grouped}.In the figure, we notice that the U-Net predicted RDF resembles closely that of ATHENA++. Furthermore, from the ratio plot, we see the discrepancy between the two densities takes a similar shape as that produced on the original validation set.

As mentioned in \S\ref{sec:phy_of_clustering}, the particle clustering is most responsive to eddies with size $\ell \approx \text{St}^{3/2} \ell_\eta$. Therefore, the absence  of large-scale driving should not affect the resulting clustering. This test indicates the network can pick up the correct scales in the power spectrum responsible for the clustering.

The relative velocity $\langle w^2\rangle^{1/2}$ and the radial relative velocity $\langle w_r^2\rangle^{1/2}$ behave similarly to the density. Due to the energy cascade inherent to Kolmogorov turbulence, it is expected that the removal of large-scale driving in this test case should not affect the small-scale relative velocity. Therefore a similar target relative velocity profile is expected. This is verified in the second plot in Fig.~\ref{fig:vp_rel_grouped} and Fig.~\ref{fig:vp_rel_rad_grouped}. Furthermore, the network also outputs a similar relative velocity as the target. The relative velocity of the predicted $v_p$ is relatively flat for $r\gtrsim 2\times 10^{-2}$ and raises for scales smaller than $r\lesssim 2\times 10^{-2}$. The "flat" part is at most $\approx 3\%$ lower than the ground truth and the "peak" is at most $\approx 12\%$ above the target just like that in the normal validation set. The predicted relative velocity also differ from VG by $\approx 5\%$, which indicates the network can discriminate gas velocities from dust velocities. This shows that the network correctly ignores the large-scale driving that is irrelevant to the relative velocity on small-scales and identifies the wave-numbers that are relevant for the problem. As a result the network is able to predict the relative velocity field with similar quality as the one it is trained on.

The radial component shows very similar character to the relative velocity. From Fig.~\ref{fig:vp_rel_rad_grouped}, we see that at the largest scale the prediction is $\approx 1\%$ below the ground truth and at the smallest scale it is $\approx 10\%$ larger. However, unlike the $\langle w^2\rangle^{1/2}$ predictions, the difference between the predicted radial velocity and VG is smaller in this case. Still, the predictions made by the network have similar quality as that in the fiducial case. This shows that for all statistics above, the network is indeed able to find the correct wave-numbers that aree responsible for the relative velocity.

\subsection{Different energy injection rate}
\label{sec:generalize_dedt}
In this generalization test, we drive the turbulence with an energy injection rate that is 10 times smaller than that of the fiducial case ($\langle \epsilon \rangle =5\times 10^{-3}$). As shown in \S\ref{sec:phy_of_clustering}, the length scale at which the strongest clustering occurs is given as $\ell_\tau \approx \langle \epsilon \rangle^{1/2} \tau_s^{3/2}$. With $\langle \epsilon \rangle  \to 0.1\langle \epsilon \rangle $, the clustering length scale should be $0.1^{1/2} \approx 0.32$ times the original. Hence, the dust clustering structure becomes finer and the RDF should rise at smaller scales. Since our statistics are limited to grid resolution $r\approx 256^{-1}$, the RDF instead appears flatter. This is reflected in the RDF of the simulated ground truth in the third subplot of Fig.~\ref{fig:rdf_grouped}, where the peak $g(r)$ value is reduced from $2.25$ to $1.8$.

However, the network output does not show similar scaling. As shown in the figure, starting from $r\lesssim 10^{-1}$, the network consistently underpredicts the RDF. From $10^{-2}\lesssim r\lesssim 10^{-1}$, the prediction is $\lesssim 3\%$ from the target. However, the deviation quickly grows for scales $r\lesssim 10^{-2}$, reaching an error of $\approx 15\%$ at a length of $\approx 5\times 10^{-3}$. This shows that while the network is able to partially learn the dependency of clustering on $\langle \epsilon \rangle $, it is unable to account for the senhanced clustering at smaller scales. We note that all predicted RDFs go to 1 as $r\gtrsim 10^{-1}$. This shows that the predicted dust density indeed goes to the average density at large scale, as expected from the definition of the RDF.

Similarly, reducing the energy injection rate will affect the gas velocity scale. By Kolmogorov scaling, the gas velocity in the inertial range should scale as $v_g \propto \langle \epsilon \rangle^{1/3}$. Since the dust is driven by the gas via aerodynamic drag, we expect the relative velocity of the particles to scale partly as $\propto \langle \epsilon \rangle^{1/3}$ too. Therefore, reducing $\langle \epsilon \rangle  \to 0.1\langle \epsilon \rangle $ should scale down the velocity by a factor of $10^{1/3} \approx 2.15$. This scaling is roughly reflected in the third columns of Figs.~\ref{fig:vp_rel_grouped},\ref{fig:vp_rel_rad_grouped}. The minimum and maximum target relative velocity scale of the fiducial average relative velocity and radial velocity are $\approx 5\times 10^{-2} - 2\times 10^{-1}$ and $\approx 1.3\times 10^{-1} - 3\times10^{-1}$ respectively, while that in this generalization test is $\approx 2\times 10^{-2} - 8\times 10^{-2}$ and $\approx 9\times 10^{-2} - 2 \times 10^{-1}$ respectively. Therefore, the target relative velocities are all smaller than the fiducial case by approximately the expected factor, and for small scale the difference roughly corresponds to the scaling we derived above. However, the network does not display similar scaling. As shown in Fig.~\ref{fig:vp_rel_grouped}, the network always over-predicts the relative velocity by a factor of at least $\approx 10\%$. Furthermore, the deviation grows steadily as the scale gets smaller, peaking at $\approx 40\%$ error at the smallest scale.

Analogous results is shown in radial velocity statistics in Fig.~\ref{fig:vp_rel_rad_grouped}. At all scales, the network over-predicts the statistics by at least $6\%$. This deviation continues to grow for smaller $r$ and an $\approx 20\%$ overshoot in radial velocity at grid-scale is observed.

Furthermore, comparing the VG line and the network output in Fig.~\ref{fig:vp_rel_grouped} and Fig.~\ref{fig:vp_rel_rad_grouped}, we see that the predicted particle velocities are much higher than what would have been obtained if all dust was perfectly following the flow (VG). Hence, the network is likely predicting particle relative velocities that correspond to simulations with a higher velocity scale. This reflects the network's inability to extract correct injection energy scales from the gas velocity field. While this is a failure of the model, it is expected because the network is trained only with one energy injection rate.

\subsection{Different particle stopping time $\tau_s$}
\label{sec:generalize_diff_stop_time}
Finally, we briefly discuss the results of a more challenging generalization test, in which we changed the dust stopping time $\tau_s$ for the particles to 0.001, a hundred times smaller than that of the original simulation. Since $\ell_\tau \approx \tau_s^{3/2}\epsilon^{1/2}$, this means the particle clustering scale will be scaled down to around $0.1 \%$ of the original flow. Again, by the same reasoning as \S\ref{sec:generalize_dedt}, a flat ground-truth RDF is expected. This is reflected in the rightmost plot in Fig.~\ref{fig:rdf_grouped}. The violet line, which represents the RDF of the target flow, is nearly flat and only deviates slightly from $1$ at $r\lesssim 10^{-2}$. The strong coupling means that the velocity of the particles in the simulation follows that of the gas closely at all scales.

In contrast to the less challenging generalization tests presented above, the large change in stopping time leads to a catastrophic failure in the prediction of the particle RDF. The results for the different velocity statistics, while not wrong by a similarly large factor, are systematically lower than the ground truth at all scales (See Fig.~\ref{fig:vp_rel_grouped} \& \ref{fig:vp_rel_rad_grouped}). These failures are to be expected because $\tau_s$ is a physically important parameter that our model has no knowledge of. They illustrate the importance of including simulations with different stopping times as part of the training set for future models.

\section{Discussion}
\label{sec:conclusion}
In this work, we have quantified the ability of the U-Net machine learning architecture to predict the clustering of particles that are aerodynamically coupled to turbulent gas. Training data was generated using a recently developed particle module for the {\sc ATHENA++} code, with the included physics and parameters chosen to be appropriate for the problem of small-scale turbulent clustering of dust in protoplanetary disks. After training, using a loss function that targets the three-dimensional particle density and velocity fields, we find that the U-Net model yields satisfactory performance for gridded representations of the {\em statistical} metrics that are of greatest interest for planet formation---the Radial Distribution Function and particle relative velocity. Errors on these quantities are scale-dependent, but typically at the percent level. Smaller maximum errors are found for density rather than velocity statistics. The model also generates, of course, a fast approximate mapping from the fluid to the particle fields in three dimensions. This is a novel capability that is not, to our knowledge, possible using analytic methods.

The use of a U-Net architecture requires a gridded representation of the particle density, and the definition of a gridded (single-valued) velocity field. These characteristics lead directly to some of the limitations observed in our study. The Lagrangian representation of particles in the training simulations yields useful information on particle statistics on scales {\em smaller} than those of the fluid grid, which is discarded as a pre-processing step in our approach. For applications where the sole interest is the statistical properties of the particles, better results extending to smaller scales would likely be obtained by using a convolutional neural network to target directly the radial distribution and relative velocity functions.

The model fails to generalize to different friction times and energy injection rates, as expected given that these are key physical parameters and the network was only exposed to a single value of them. There are no obvious obstacles to training the network on a suite of simulations, and thereby learning the dependence of particle clustering on physical input parameters \citep[see, e.g. analogous work in cosmology;][]{vn21}. The combination of three-dimensional spatial prediction of particle positions, together with prediction of the particles' velocity relative to the gas, would then allow calculation of particle collision rates and outcomes for arbitrary size distributions. Different approaches would be needed to incorporate the back-reaction of particles on the gas, an aspect ignored in our work but which is central to models of planetesimal formation \citep{youdin05,simon16}. Learned simulators for turbulent flows show considerable promise \citep{dang22}, and it would be interesting to apply similar methods to multi-phase fluid turbulence. The existence of new  instabilities that rely on the two-way coupling between gas and particles \citep{youdin05,squire18} suggests that multi-phase turbulence would be a stringent test of such methods.

The work presented here explores arguably the simplest problem in two-fluid turbulence that has relevance to planet formation. Related problems are present on multiple other scales in protoplanetary disks, and in other domains including atmospheric science and experimental studies of turbulence. Our results suggest that deep learning offers a complementary approach to direct numerical simulations, that can predict particle clustering at a lower computational cost.

\acknowledgments
We thank Jim Stone for helpful discussions during the development of this project. PJA acknowledges support from NASA TCAN award 80NSSC19K0639. CCY and ZZ acknowledge support from NASA via the Emerging Worlds program (award 80NSSC20K0347) and the Astrophysics Theory Program (award 80NSSC21K0141). CCY acknowledges support from NASA TCAN award 80NSSC21K0497.

\bibliographystyle{aasjournal}

\section{Introduction}
\label{sec:intro}
The dynamics of aerodynamically coupled particles in turbulent fluids are of interest for astrophysical, atmospheric physics, experimental, and industrial applications \citep{toschi09}. The coupling between the fluid and particle phases modifies the turbulent cascade, and leads to clustering, effective diffusion, and collisions among the particles. The generic problem has numerous physically significant variations, depending upon the regime of aerodynamic drag (Stokes or Epstein), the compressibility of the turbulence, the relative mass in particles versus that in the fluid, and the amount of fluid volume excluded by the particles.

In the astrophysical domain, the response of solid particles to turbulence is important for multiple problems in protoplanetary disk physics and early phase planet formation. On large spatial scales, the ubiquitous substructure seen in dust continuum images of protoplanetary disks \citep{vdm13,alma14,andrews18,andrews20} is interpreted as concentrations of solids within gaseous vortices \citep{barge95} and axisymmetric pressure maxima \citep{whipple72}, structures which may themselves have formed as a consequence of disk turbulence \citep{johansen09,manger18}. On small scales, the collision velocity of small (micron to cm-sized) particles is typically determined by how the particles interact with turbulence in the gas \citep{weidenschilling93,johansen14}. For small particles the collision velocity induced by turbulence increases with size. Because dust particles fragment (or bounce) above some threshold velocity \citep{blum08}, this behavior leads to the prediction of a {\em maximum particle size} that can result from coagulation. On intermediate scales, the two-way momentum transfer between particles and gas in disks results in a linear instability, the {\em streaming instability} \citep{youdin05,krapp19,paardekooper20,yang21}, that may be responsible for the formation of planetesimals via gravitational collapse \citep{johansen15,simon16,li21,yang17,schafer17}.

The various manifestations of particle-gas interactions in protoplanetary disks can be studied using analytic methods \citep[e.g.][]{volk80,ormel07,squire18,krapp19,lin21}, and by simulations that treat the particle phase as either discrete Lagrangian particles or as a second fluid \citep[e.g.][]{pan11,ishihara18,bhatnagar18,krapp20,sakurai21}. Although powerful, these methods by no means suffice to answer all of the physically interesting questions that arise in planet formation studies. Numerical simulations can only model a small subset of the full range of spatial scales, leaving open questions as to how interactions on different scales combine \citep{chambers10,hartlep20}. Even when only a limited range of scales need to be studied, it is computationally expensive to run models that span the range of plausible size distributions for particles in protoplanetary disks. Machine Learning (ML) approaches may provide new theoretical tools to study these problems. In cosmology, ML architectures have been able to generate realistic predictions of gas properties from dark matter only simulations, and to interpolate reliably from grids of simulations that sparsely cover the full parameter space \citep{vn21}. By analogy, in planet formation it might be possible to extend the range of explicit simulations with an effective ML sub-grid model, and to make predictions for the response of arbitrary particle size distributions using a fairly small set of input simulations. Multiple ML architectures would likely be needed to realize these gains.

In this paper we present an initial investigation into the utility of established ML methods for predicting the spatial and statistical properties of particles in turbulent fluids. We consider the response of a single particle size to isotropic forced turbulence, in a periodic domain, with aerodynamic drag in the Epstein regime. This setup excludes the forces that lead radial and vertical epicyclic oscillations of free particles, which are physically important \citep{youdin07}. It is, however,  the simplest model problem that is relevant to studies of the rate and velocity of particle collisions on small scales in protoplanetary disks. Clustering in this regime is well-studied, and is understood to result from the interaction of particles with small-scale fluid vortices \citep[e.g.][]{johansen14}. This physical understanding does not, however, lead to a predictive model for the particle distribution. The collision velocity, on the other hand, is amenable to an analytic model \citep[several exist, the most common in disk applications being][]{ormel07}. We generate training data using a recently developed Lagrangian particle module within the {\tt Athena++} code (Yang et. al., in prep.), and use a U-Net network architecture \citep{ronneberger2015u} to learn the mapping between the three-dimensional fluid fields and gridded representations of the particle data. We then assess the fidelity of the predictions for both statistical properties of the particles (the strength of clustering and the relative velocity as a function of scale), and for their instantaneous spatial distribution.

The plan of this paper is as follows. In \S\ref{sec:simulations} we describe the {\tt Athena++} code, the implementation of Lagrangian particles within the code, and the setup of the forced turbulence simulations. The Machine Learning methods are described in \S\ref{sec:ML}. Sections \S\ref{sec:density} and  \S\ref{sec:normal_particle_velc_prediction} quantify the network performance for the particle density and velocity fields, and for statistical properties that would be relevant to the calculation of the rate and mutual velocity of particle collisions. We present some generalization tests in \S\ref{sec:generalization} and conclude in \S\ref{sec:conclusion}.

\section{Numerical simulations}
\label{sec:simulations}

\subsection{Physics of turbulent clustering}
\label{sec:phy_of_clustering}
To aid in the discussion of clustering statistics in later sections, we briefly introduce some relevant physics of turbulence and turbulence-driven dust dynamics.

In the standard description of fluid turbulence \citep{kolmogorov41,kolmogorov1991local}, there exist a range of scales, known as the inertial range, within which turbulent statistics depend only on the mean energy transfer rate $\langle \epsilon \rangle$. In this limit, the two point relative velocity $\langle \delta \vec{u} \cdot \delta\vec{u}\rangle$ at a separation $\ell$ scales as $\langle \epsilon \rangle^{2/3} \ell^{2/3}$. In the inertial range, dissipation due to viscosity is negligible and energy mainly cascades to smaller scales. The scale at which viscosity becomes important can be found by equating $\langle \epsilon \rangle $ to the energy dissipation rate $\nu \left[\delta u(\ell_d)/\ell_d\right]^2 \approx \langle \epsilon \rangle$ \citep{ching2014statistics}, from which one can obtain the
Kolmogorov dissipative length scale, timescale and velocity scale as $\ell_\eta = (\nu^3/\langle \epsilon \rangle)^{1/4}, \tau_\eta = (\nu/\langle \epsilon \rangle)^{1/2}$, and $u_\eta = (\nu \epsilon)^{1/4}$ respectively.

Now consider dust particles coupled to such a turbulent flow. Assuming that friction is the dominant force driving the dust, the equation of motion of a dust particle takes the form,
\begin{equation}
    \frac{d\vec{v}_p}{dt} = \frac{\vec{v}_g(p) - \vec{v}_p}{\tau_s},
    \label{eq:particle_eom}
\end{equation}
where $\vec{v}_p$ is the velocity of particles, $\vec{v}_g$ is the velocity of gas, and $\tau_s$ is the friction time characterizing the strength of coupling between the dust and gas. For dust particles in the Epstein regime\footnote{In this regime the drag force is {\em linear} in the velocity difference between the particle and the fluid, rather than quadratic as in the Stokes regime typically relevant to terrestrial fluids}, that is, particles with size $a_p$ much smaller than the mean free path $\lambda_g$ of the gas molecules $(a_p\ll \lambda_g)$, the friction time can be estimated as \citep{weidenschilling93},

\begin{equation}
    \tau_s = \left(\frac{\rho_d}{\rho_g} \right) \left(\frac{a_p}{c_s}\right),
    \label{eq:epstein_drag}
\end{equation}
where $\rho_d$ is the material density of the dust, $\rho_g$ is the fluid density, and $c_s$ is the local speed of sound of the gas.

To quantify the clustering behavior, $\tau_s$ should be compared with a characteristic time-scale of the turbulent flow. A commonly used time scale relevant for our purposes is the Kolmogorov time-scale $\tau_\eta$ \citep{sundaram_collins_1997,pan11,ishihara18}. It represents the turnover time of the smallest eddies. The ratio of these time-scales is the dimensionless Stokes number $\text{St} = \tau_s/\tau_\eta$\footnote{In protoplanetary disk studies, the ``Stokes number" is often defined instead using as a characteristic time-scale the dynamical time, $\Omega_{\rm K}^{-1}$, where $\Omega_{\rm K}$ is the Keplerian orbital frequency. Because our simulations do not include shear, we use the definition appropriate for fluid turbulence.}. For $\text{St} \ll 1$, the particle is well coupled to the gas at all scales and behave like a trace particle. With very little memory of its previous motion, the clustering will be mainly determined by the velocity divergence of the surrounding flow at one point in time. It has been shown that the particle velocity divergence $\sum_i \partial_i v_{p}^i \approx  \text{St}/\tau_\eta$ \citep{pan11}, where $i$ is the spatial index. This implies that for $\text{St} < 1$ the clustering intensifies with the Stokes number.

Conversely, when $\text{St} > 1$, turbulent fluctuate too rapidly for the particle to respond. This has two effects. First, as the friction time $\tau_s$ increases, the history of the particle has an increasingly large influence on the particle relative velocity. As shown in the formal solution of equation~\ref{eq:particle_eom} \citep{pan2013turbulence},
\begin{equation}
    \vec{v}_p(t) = \frac{1}{\tau_s} \int_{t_0}^{t} \vec{v}_g(\mathbf{X}(\tau),\tau) \exp(-\frac{t-\tau}{\tau_s}) d\tau,
\end{equation}
where $\mathbf{X}(t)$ is the position of the particle at some time $t$ and $t_0$ is a reference time chosen to be long enough (i.e. $t-t_0\gg \tau_s$) so that the velocity at time $t_0$ has long be "forgotten" by the particle. From the equation, we see that the particle's current velocity depends on the past history of gas velocity within a time window of $\approx \tau_s$. Hence, one should expect that for $\text{St} > 1$ the local gas velocity divergence, and also the $\text{St}/\tau_\eta$ scaling discussed above, should no longer be adequate for describing particle clustering.

Furthermore, particles with $\text{St} > 1$ respond to eddies of different turnover time selectively, giving rise to a characteristic clustering scale $\ell_{\tau_s}$ \citep{pan11}. Since $\tau_s > \tau_\eta$, some of the eddies have turnover time smaller than $\tau_s$. Flow on these scale is decoupled from the particle and instead give it random kicks, which inhibits clustering at those scales. On the other hand, the particles are well coupled to eddies with turnover time much larger than $\tau_s$, and obtain relative velocities similar to that of the flow. This depresses  clustering at large scale. Hence, the most significant clustering is expected at $\tau_{eddy} \approx \tau_{s}$. In the inertial range, we estimate $\tau_{eddy} \approx \ell^{2/3} \langle \epsilon \rangle^{-1/3}$ and obtain the  characteristic scale $\ell_{\tau_s} \approx \langle \epsilon \rangle^{1/2} \tau_s^{3/2} = \text{St}^{3/2} \ell_\eta$, where $\ell_\eta$ is the Kolmogorov length.

\subsection{Forced turbulence simulations}
\label{sec:forced_turbulence_simulation}
We use the \texttt{ATHENA++} magnetohydrodynamics (MHD) code \citep{stone20} to calculate the training data for our study.
All simulations are conducted using a uniform 3D Cartesian grid to represent the gas variables with a resolution of $256$ cells in each spatial direction.
We use the Harten--Lax--van Leer--Einfeld (HLLE) Riemann solver combined with a second-order accurate van Leer predictor-corrector scheme for time integration. Spatial reconstruction is done using the 3rd order piecewise parabolic method (PPM). We use periodic boundary conditions in all spatial directions.

The gas is initialized with a constant density of $\rho_\mathrm{gas} = 1$ and with all velocity components $v_i =0$. We use an isothermal equation of state,
\begin{equation}
    P_\mathrm{gas} = c_\mathrm{s}^2 \rho_\mathrm{gas},
\end{equation}
with the constant speed of sound $c_\mathrm{s}=1$.
To generate a turbulent velocity spectrum for the gas, we use the FFT turbulence driver module of \texttt{ATHENA++}, which continually feeds energy into the domain at the largest scales by adding large scale velocity perturbations to the flow. The scale of these perturbations is controlled by the parameters $k_\mathrm{min}$ and $k_\mathrm{max}$, denoting the smallest and largest wavenumber to be perturbed. (Here wavenumber is defined as $k=1/\lambda$, without a factor of $2\pi$ in front). Unless otherwise specified, we use $k_\mathrm{min}=1$ and $k_\mathrm{max}=5$ with simulation box size $L_x=L_y=L_z=1$, which allows the turbulent energy spectrum to develop naturally into a Kolmogorov type spectrum, where numerical dissipation occurs near the Nyquist frequency. We exclude data from an initial transient state from our ML dataset. The amount of energy injected per perturbation interval (which we set to be equal to one simulation time-step) is given by the energy injection rate $\epsilon_\mathrm{inj}$ and is set to $\epsilon_\mathrm{inj} = 0.005$ (except where specified). The RMS gas velocity calculated from the simulation is $v_{g}^\text{rms} \approx 0.104$, which is $\approx 10\%$ the speed of sound. The turbulence, although computed with a compressible code, is thus for practical purposes almost incompressible. As we are investigating subsonic turbulence, we enforce solenoidal turbubulence driving by suppressing the compressive component of the velocity perturbation.

\subsection{Lagrangian dust particles}
\label{sec:lagrangian_dust_particles}
We use the Lagrangian dust particle module of the \texttt{ATHENA++} code (Yang et. al., in prep.) to simulate the particle component in the simulations. In all our simulations, $256^3$ dust particles are used.
For simplicity, we only treat the drag force enacted onto the dust particles by the gas and neglect the back-reaction drag force enacted by the particles onto the gas, which is a good approximation in cases with low dust-to-gas mass ratios. To interpolate between the gridded values of the gas component and the particle positions, the module uses the triangular shaped cloud (TSC) algorithm.
Unless otherwise specified, we use a particle friction time of $\tau_s = 0.1$.

From the simulated output, we estimate the turbulent scales of the fiducial test case to be: $\langle \epsilon\rangle\approx 3.17\times 10^{-3}$, $\nu \approx 1.95 \times 10^{-5}$, $\ell_\eta \approx 1.24\times 10^{-3}$, $\tau_\eta \approx 7.84\times 10^{-2}$ and $u_\eta \approx 1.57\times 10^{-2}$, all in the relevant simulation units. Similarly, the particle clustering length is estimated to be $\ell_{\tau_s} \approx 1.58\times 10^{-3}$. The energy injection rate $\langle \epsilon\rangle$ and eddy turnover time $\tau_\eta$ are estimated from the third-order structure function and strain-rate tensor respectively; The particle friction time is taken directly from simulation setup, and the rest are obtained from their relationship with these three variables.
\section{Machine learning methods}
\label{sec:ML}

\subsection{Dataset}
Snapshots of the simulation are output at intervals chosen such that the temporal correlation between adjacent frames is negligible. The gas density field $\rho_g$, gas velocity field $\vec{v}_g$, particle density field $\rho_p$ and particle velocity field $\vec{v}_p$ are used to train and validate the network. All of the fields are gridded on a $256^3$ mesh. (The gas fields are natively gridded; the particle data is mapped to the same grid.) The first two serve as inputs and the later as outputs of the network.
\subsection{Data preprocessing and Normalization}
The following data preprocessing and normalization steps are applied to the input and output fields to aid the training.
\subsubsection{Input fields}
The gas density $\rho_g$ is log normalized,
\begin{equation}
    \rho_g^\ast = \log(\rho_g)
    \label{eq:rhog_norm},
\end{equation}
where $\rho_g^\ast$ is the normalized gas density and $\rho_g$ is the gas density
The log normalized density field helps separate gas density of various scales, allowing the network to better perceive the input field. No preprocessing is applied to the gas velocity $v_g$.
\subsubsection{Output fields}
The particle density field $\rho_p$ and velocity field $v_p$ are known to be non-gaussian \citep[e.g.][]{ishihara2018dust}. Therefore, data normalization and preprocessing techniques are employed to aid the training of the neural network.

For the particle density field $\rho_p$, after experimentation, the following normalization is used,
\begin{equation}
    \rho^\ast_p = \ln(\exp(\rho_p/\rho_0)-1).
    \label{eq:rhop_norm}
\end{equation}
Here, $\rho^\ast_p$ is the normalized particle density and $\rho_0=0.14$ is an empirically chosen density scale that results in a  distribution of $\rho^\ast_p$ that is most nearly symmetric. This value corresponds roughly to the average particle density. The mapping between $\rho_p^\ast$ and $\rho_0$ is one-to-one, and it can be easily inverted using a softplus-like function:
\begin{equation}
    \rho_p = \rho_0\ln(\exp(\rho_p^\ast)+1).
    \label{eq:rhop_inv_norm}
\end{equation}
Similarly, to prevent the bulk gas velocity $v_g$ from masking the fine details of the particle velocity $v_p$, we choose to train our network (referred to later as UNET-R) on the relative velocity $v_{rel}$ instead of the particle velocity $v_p$,
\begin{equation}
    \vec{v}_{rel} = \vec{v}_p - \vec{v}_g.
    \label{eq:vp_norm}
\end{equation}
This is intended to help the network to focus on the physically interesting deviation of the particle velocity $v_p$ from the gas velocity $v_g$, instead of learning primarily just the background flow. To evaluate the effectiveness of this strategy, we also trained the network directly on $v_p$ (later, UNET-V) and compared it against that on $v_{rel}$. Details of this comparison can be found in \S\ref{sec:normal_particle_velc_prediction}.

\subsection{Model and training}

\subsubsection{Model}
\begin{figure*}
    \centering
    \includegraphics[width=0.9\textwidth]{Figures/dustnet-architecture.pdf}
    \caption{The architecture of the U-Net. The thickness of individual layers corresponds roughly to the number of channels in each layer and the width and height corresponds to the spatial size. The color scheme corresponds to: convolution layers of kernel size 3 and stride 1 \textbf{(orange)}, Leaky ReLU activation with batch normalization \textbf{(deep orange)}, downsampling layers using convolution of kernel size 2 and stride 2 \textbf{(red)} and upsampling layers using transposed convolution of kernel size 2 and stride 2 \textbf{(purple)}. Residue connection is implemented for all convolution layers except those used in up and down sampling. Similar to other U-Net architectures, large scale information is incorporated by narrow down and concatenating compression layer outputs with the upsampled outputs.}
    \label{fig:unet_architecture}
\end{figure*}

The network architecture U-Net is used. First developed by \cite{ronneberger2015u} for biomedical image segmentation and later generalized by \cite{milletari2016v} to handle volumetric data, U-Net is a "fully convolutional network" that is capable of taking a multichannel, 3D input block and producing an output block of similar spatial size but with a different number of channels. It exhibits translational equivariance, such that its predictions depend only on the relative positions of the grid points \citep{long2015fully}. This is desirable because it is consistent with the physical picture of homogeneous turbulence \citep{ching2014statistics}, and because it allows the network to be applied to input fields of different spatial sizes. U-Net has recently been successfuly employed in cosmology for tasks including super-resolution (generating high resolution predictions from lower resolution simulations) \citep{schaurecker2021super}, and prediction of galaxy distributions from the dark matter density field \citep{yip2019dark}. This motivates us to apply the architecture to our problem.

The network is implemented using the \textbf{map2map} neural network emulator using the architecture specified in Fig.~\ref{fig:unet_architecture}. The network consists of two major parts, the "compression part" and the "decompression part". The former consists of successive convolutional and downsampling layers that compress information of the input field into a smaller size latent representation; while the latter consists of successive levels of convolutional and up-sampling layers that expand the latent information into the output block. Information extracted in different levels of the network is incorporated by concatenating outputs of the compression part and those of the expanding part on the same level.

\subsubsection{Training}

The simulation frames, 252 in total, are randomly divided into training and validation sets in a 4:1 ratio. Training is done using only the frames in the training set.

The entire training set is iterated through during each epoch of training. For each frame in the set, we divide the input fields into eight octants of size $128^3$ so that they can fit onto one GPU.  Then, we pad them periodically by $5$ grids on all sides so that it is $148^3$ in size. The padding, as discussed by \cite{schaurecker2021super}, allows the network to preserve full translational symmetry and generate outputs without boundary artifacts. This results in an input field of dimension $(4,148,148,148)$, with three of its channels coming from the spatial components of the gas velocity $v_g^{(x/y/z)}$ and one from the gas density $\rho_g$. Discrete rotation and parity transforms on the data are also applied to the input block using map2map. This helps the network to capture the isotropic nature of the turbulent flow. The input enters the U-Net and produces an output field of size $(4,108,108,108)$: One of the channels represents particle density $\rho_p$, and the rest represent particle relative velocity $\vec{v}_{rel}$. Finally, the whole $(4,256,256,256)$ field is reconstructed from the output patches produced by map2map.

For the loss function, we choose to minimize the mean square error (MSE) loss over all output channels. Weight decay is also used. The loss function, which is minimized using the Adam optimizer \citep{kingma2014adam}, takes the following form,
\begin{equation}
    \mathcal{L} = \left[\frac{1}{\mathcal{B}}\sum_{i=1}^{\mathcal{B}}(\mathbf{y}_\text{tgt} - \mathbf{y}_\text{out}^{(i)})^2\right] + \lambda \norm{\mathbf{w}}^2,
\end{equation}
where $\mathbf{y}_\text{tgt},\mathbf{y}^{(i)}_\text{out}$ represents the target and output fields, $\mathcal{B}$ is the batch size, and $\lambda$ is the weight decay and $\norm{\mathbf{w}}$ is the square norm of weights. Unless otherwise specified, all training runs in this paper use learning rate $10^{-4}$, weight decay $\lambda= 10^{-4}$, and batch size $\mathcal{B}=1$, for around 400 epochs until the results are stable.

After training, we evaluate the model by comparing the output with the target over a variety of statistics. The results are summarized in \S\ref{sec:density} and  \ref{sec:normal_particle_velc_prediction}.

\section{Prediction of particle density field}
\label{sec:density}

\subsection{Visual Comparison}
\label{sec:rhop_visual_comp}
Using the trained network, we output the predicted particle density field $\rho_p$ and compare it to the target on a frame-by-frame basis. Snapshots of selected frames are presented in Fig.~\ref{fig:rhop_slice_comp_normal}. Because the densities span a wide range, a $\log_{10}$ scale color map is chosen to aid visualization. We also show the normalized cross correlations\footnote{The normalized cross correlation of two function $f,g$ is defined to be $\left[(f\star g)(r)\right]/\sqrt{\left[(f\star f)(0)\right]\left[(g\star g)(0)\right]}$, which is a normalized inner product of $f,g$. For volumetric data, the result is radially binned.} of these frames, defined as in Fig.~\ref{fig:rhop_cross_corr}.

As shown in Fig.~\ref{fig:rhop_slice_comp_normal}, both machine learning outputs (UNET-V, UNET-R) produce density distributions that are similar to the ground truth (ATHENA++). Locations with voids in the density distribution align, and the filamentary structures of the dust are captured. However, both predicted outputs appear to be less sharp than the target field, indicating the algorithm has smoothed out the densities. The density at the locations of peaks is also less pronounced, indicating an under-prediction in those areas.
These observations are borne out in Fig.~\ref{fig:rhop_cross_corr}, where we see that the peak values of normalized cross correlations for both UNET-V and UNET-R are only $\approx 0.6$. The larger peak values for UNET-R show that UNET-R, trained on the relative velocity, is slightly better at predicting $\rho_p$ than UNET-V.

\begin{figure*}
    \centering
    \includegraphics[width=\textwidth]{Figures/rhop_slice_comp_normal.pdf}
    \vskip-0.75cm
    \caption{Particle density field slice $\rho_p$ from ATHENA++ (left), UNET-V (center, a network trained on the relative velocity between particles and gas) and UNET-R (right, a network trained on the absolute particle velocity), for a typical frame. Qualitatively, there is good agreement between the predicted structures and the ground truth, though the network outputs are visually less sharp than the target simulation data.}
    \label{fig:rhop_slice_comp_normal}
\end{figure*}

\begin{figure}
    \centering
    \includegraphics[width=0.4\textwidth]{Figures/rhop_cross_corr.pdf}
    \caption{Normalized cross correlation of $\rho_p$ of target field (ATHENA++) with itself, UNET-R and UNET-V. The peak values of the statistics are: 1.0, 0.622 and 0.613 respectively.}
    \label{fig:rhop_cross_corr}
\end{figure}

\subsection{Radial Distribution Function (RDF)}
\label{sec:normal_stat_rdf}
We compute the particle radial distribution function (RDF) $g(r)$ of the simulated and U-Net produced outputs to study their difference in spatial distributions more carefully. Let $\bar{n}$ be the average dust density. The RDF is defined so that the average number of dust particles $P(r)$ inside a volume element $dV$ at a distance $r$ relative to a reference particle is given by,
\begin{equation}
    P(r) = \bar{n} g(r) dV.
    \label{eq:radial_distribution_func}
\end{equation}
Although this definition can be applied directly to compute the RDF, a faster way of obtaining it on a grid is available. First, we compute the particle density correlation function (DCF) $\xi(r)$, which is defined as,
\begin{equation}
    \xi(r) = \frac{1}{\bar{n}^2} \langle \left(n(\vec{x})-\bar{n}\right) \left(n(\vec{x}+\vec{r})-\bar{n}\right)\rangle .
    \label{eq:density_correlation_func}
\end{equation}
Then we use the relation $g(r) = 1 + \xi(r)$ between the two quantities to obtain the RDF \citep{shaw2003particle}. This allows for an efficient computation, as  Eq.~\ref{eq:density_correlation_func} can be computed quickly using the fast fourier transform (FFT). Unless otherwise specified, all RDFs computed in this paper are based on this method. The use of grid-based metrics, rather than metrics computed directly from the particle distribution, differs from most prior work. Differences in results that arise between the method based on Eq.~\ref{eq:radial_distribution_func}, and that based on Eq.~\ref{eq:density_correlation_func} for the simulated particle density, are discussed in Appendix~\ref{sec:grid_stats_vs_part_stats}.

In the left-most panel of Fig.~\ref{fig:rdf_grouped}, we show the RDFs of $\rho_p$ produced by the simulation (ATHENA++), and by the U-Nets (UNET-V, UNET-R). Recall, UNET-V is the U-Net trained directly on $v_p$ and UNET-R on $v_{rel}$. The features of the RDFs reflect what we see in the slice plots. The simulation and U-Nets agree well for $r \gtrsim 0.01$, indicating the large scale features are captured. For $r\approx 7\times10^{-3}$, a dip of roughly $5\%$ is observed, indicating the underestimation of $\rho_p$ at small scales. The over-prediction for $r\gtrsim 7\times 10^{-3}$ might be due to the smoothed density around the peaks. Overall, the error in RDF statistics is bounded by $\pm 5\%$. Although we have picked only one frame for discussion, the features observed are general. Discussion on the consistency across frames can be found in Appendix~\ref{sec:consistency_across_frames}.

\begin{figure*}
    \centering
    \includegraphics[width=\textwidth]{Figures/rdf_compare_all.pdf}
    \caption{Predicted versus simulated strength of particle clustering as a function of spatial scale, quantified via the radial distribution function statistic (RDF). From left to right, the RDF plots show different simulation setups. The left-most panels show the fiducial setup, while the remaining panels show the results of generalization tests. (See \S\ref{sec:forced_turbulence_simulation},\ref{sec:lagrangian_dust_particles} and \S\ref{sec:generalization} for details.)  {ATHENA++} shows the RDF computed directly from the simulation, {UNET-R} is the U-Net output trained with relative velocity between gas and dust, and {UNET-V} is the U-Net output trained directly on dust velocity. The first row compares the RDF of the  simulation output and the U-Net prediction. The second row shows the ratio between the outputs and the target to better quantify the error. The zoom-in plot in the lower-right panel shows the peak of the error at small scales.}
    \label{fig:rdf_grouped}
\end{figure*}

\section{Prediction of particle velocity field}
\label{sec:normal_particle_velc_prediction}

\subsection{Visual Comparison}
In Fig.~\ref{fig:vrel_slice_comp}, we display slices of the three spatial components of the relative velocities field $v_{rel}$ produced by simulation and the U-Nets. Recall that the particle velocity $v_p$ can be thought of as having two contributions: (i) a gas velocity term $v_g$, which represents the velocity of the underlying flow driving the dust, and (ii) a relative velocity term $v_{rel} = v_p - v_g$, which measures the deviation of the particle velocity from the gas velocity due to the imperfect coupling between them. To prevent $v_g$ from overshadowing $v_{rel}$ in visual comparisons, we show $v_{rel}$.

From Fig.~\ref{fig:vrel_slice_comp}, we see that both UNET-V and UNET-R produce $v_{rel}$ fields that are qualitatively similar to that of the ATHENA++ ground truth. It is difficult to tell which is better solely by visual comparison. For example, at around $(v_{rel}^{3},0.5,0.6)$, UNET-V produce a sharper output for places where that component peaks, whereas at around $(v_{rel}^{3},0.1-0.3, 0.6-0.9)$, UNET-R captures the zeros of the field better. This motivates us to use more sophisticated statistics in \S\ref{sec:rel_velc_normal} and \S\ref{sec:rel_velc_rad_normal}, to assess which of them performs better.
\begin{figure*}
    \centering
    \includegraphics[width=0.90\textwidth]{Figures/velc_slice_comp.pdf}
    \caption{Prediction of the spatial distribution of particle relative velocity ($v_g-v_p$), in domains of isotropic forced turbulence. The figure compares the ground truth, obtained with {\sc ATHENA++}, with predictions from two deep learning models: UNET-V trained directly with particle velocity $v_p$, and UNET-R with relative velocity $v_{rel}$. The columns represent different setups (ATHENA++, UNET-V, UNET-R) and the rows show the 3 components of the velocity field. Each panel is a slice through the three-dimensional domain.}
    \label{fig:vrel_slice_comp}
\end{figure*}

\subsection{Relative velocity $\langle w^2 \rangle^{1/2}$}
\label{sec:rel_velc_normal}
The relative particle velocity $\langle w^2 \rangle^{1/2}$ is a commonly studied statistic of particle clustering in turbulence. Following \citet{ishihara2018dust}, we define it as,
\begin{equation}
    \langle w^2 \rangle  = \frac{1}{N_p}\sum_{|\vec{r_1} - \vec{r_2}| = r} |\vec{V}_1 - \vec{V}_2|^2,
    \label{eq:rel_vel}
\end{equation}
where $r$ is the separation between the particle pairs, $\vec{V}_1$, $\vec{V}_2$ are the velocities of the particles located at $\vec{r}_1,\vec{r}_2$ and $N_p$ is the number of pairs of particles.

To compute Eq.~\ref{eq:rel_vel} on the grid, we convert the average into a weighted one using the relevant particle density $\rho_p$ as weight. More precisely, let $p$ be a grid cell and $\text{C}_p(r_1,r_2)$ be the set of all grid cells whose center point lies inside the spherical shell with inner and outer radius $r_1$ and $r_2$ that is centered at $p$. Then a grid average of some quantity $Q$, weighted using particle density $\rho_p$, is defined as,
\begin{equation}
    \langle Q\rangle^{\rho_p}_\text{grid} = \frac{\sum_{q\in C_p(r,r+\Delta r) }\rho_p(q) Q(q)}{\sum_{q\in C_p(r,r+\Delta r)} \rho_p(q)},
\end{equation}
where $\rho_p(q)$ and $Q(q)$ are understood as the grid values at cell $q$. Using this definition, we write $\langle w^2_\text{tgt}\rangle^{1/2} = \left[\langle w^2_\text{tgt}\rangle^{\rho^\text{tgt}}_\text{grid}\right]^{1/2}$ for the ground truth relative velocity and $\langle w^2\rangle^{1/2}=\langle w^2_\text{out}\rangle^{1/2} = \left[\langle w^2_\text{out}\rangle^{\rho^\text{out}}_\text{grid}\right]^{1/2}$ for the output velocity field.

Unlike grid RDFs, computing these quantities on the grid are relatively expensive. To see that, notice if we were to compute the statistics for the whole grid $r\in [256^{-1},1]$, we would need to check every grid point against every other grid point. If $N$ is the number of grid points along one axis, this would imply a complexity $\approx (N^3)^2=N^6$.
To reduce the computational cost, we restrict $r\in[256^{-1},10^{-1}]$ for all relative velocity statistics. This remains practical because for many applications it is the accuracy of relative velocities at small scales that is the most important. Furthermore, from the estimated clustering scale in \S\ref{sec:lagrangian_dust_particles}, we see that the upper bound of the radii $r=10^{-1} \gg \ell_\eta \sim 10^{-3}$, so the relative velocity curve for $r>10^{-1}$ is likely not relevant for our application. Again, possible errors introduced by grid averaging are explored in Appendix~\ref{sec:grid_stats_vs_part_stats}.

We plot $\langle w^2\rangle^{1/2}$ and its error as a function of $r$ in Fig.~\ref{fig:vp_rel_grouped}. The upper left panel shows that the particle velocity fields produced by UNET-V and UNET-R have a similar shape to ground truth (ATHENA++), and the lower left panel shows that both models have their error bounded by approximately $15\%$ above and $7\%$ below. For $r \gtrsim 2.5 \times 10^{-2}$, both models under-predict the relative velocity, while for $r\lesssim 2.5 \times 10^{-2}$ both over-predict. However, UNET-R's absolute error is always less than UNET-V's at all length scales. In particular, the errors in UNET-R's prediction are $\approx 2.5\%$ less than UNET-V's at the smallest and largest scales. Furthermore, the region around the point where UNET-R's error curve crosses unity, especially that toward the tail of the curve, is much flatter than that of UNET-V's. This indicates the error varies less with separation in UNET-R.

The results above support our hypothesis that targeting the relative velocity aids in the network training. UNET-R can capture the radial dependence of relative velocity in the predicted particle velocity field within a maximum error of approximately 10\%. Although this error is systematic and scale-dependent rather than random, we consider this to be good performance.

\begin{figure*}
    \centering
    \includegraphics[width=\textwidth]{Figures/vrel_compare_all.pdf}
    \caption{Comparison of relative velocity between dust particles $\langle w^2 \rangle^{1/2}$ for different simulation setups (see \S\ref{sec:forced_turbulence_simulation},\ref{sec:lagrangian_dust_particles} and \S\ref{sec:generalization} for details). {ATHENA++} represents the relative velocity computed from the simulation, {UNET-R} is the U-Net output trained with relative dust velocity, {UNET-V} is the U-Net output trained directly on dust velocity, and {VG} is the average gas velocity weighted by particle density (i.e. the relative velocity curve one would have obtained if the particles follow perfectly the flow). The left-most panels show the fiducial results, the other panels show the results of generalization tests.}
    \label{fig:vp_rel_grouped}
\end{figure*}

\subsection{Relative radial velocity $\langle w^2_\text{r} \rangle^{1/2}$}
\label{sec:rel_velc_rad_normal}
Similarly, the root-mean-square relative radial velocity $\langle w_\text{r}^2\rangle^{1/2}$ is defined as,
\begin{equation}
    \langle w_\text{r}^2 \rangle  = \frac{1}{N_p}\sum_{|\vec{r_1} - \vec{r_2}| = r} \left(\left(\vec{V}_1 - \vec{V}_2\right)\cdot \hat{\eta}\right)^2,
\end{equation}
where all the symbols retain the same meaning as in Eq.~\ref{eq:rel_vel}, and $\hat{\eta} = (\vec{r}_1 - \vec{r}_2)/|\vec{r}_1 - \vec{r}_2|$. Using the same weighting procedure as in \S\ref{sec:rel_velc_normal}, we again computed $\langle w^2_\text{r}\rangle$ on ATHENA++, UNET-V and UNET-R data. The results are shown in Fig.~\ref{fig:vp_rel_rad_grouped}.

In the figure, we see that both the relative particle velocity $\langle w_{r}^2\rangle^{1/2}$ predicted by UNET-R and UNET-V have similar shape to the target. However, the performance of the U-Nets varies at different scales. For $r\gtrsim 2.5\times 10^{-2}$ both models perform similarly and predict a radial velocity around $\approx 2.5\%$ above and $\approx 5\%$ below the target respectively. The two models over-predict at scales $r\lesssim 4\times 10^{-2}$, being $\approx 10\%$ and $\approx 12.5$ larger respectively. Unlike $\langle w^2\rangle^{1/2}$, the error curves of the models do not intersect at unity. However, UNET-R's error remains  smaller than UNET-V's across most scales, except for a small region around $r\approx 3\times 10^{-2}$.

Also, while both UNET-R and UNET-V overpredict the radial velocity at small scales, that of UNET-V is much closer to the curve labeled VG---the average gas velocity weighted by particle density---than UNET-R. This indicates that UNET-V is essentially outputting the gas velocity $v_g$ as the particle velocity at small scales, instead of learning the particles' deviation from the gas flow. This further demonstrates that training on $v_{rel}$ is effective in helping the network to single out how the particle velocity differs from the background flow.

Finally, unlike the relative velocity, which shows a flat part at intermediate scales, there are no obvious flat region on UNET-V's error curve. This suggests the network in general predicts the direction of $v_g$ less accurately than its amplitude.

\begin{figure*}
    \centering
    \includegraphics[width=\textwidth]{Figures/vrel_rad_compare_all.pdf}
    \caption{Comparison of relative radial velocity $\langle w^2_r \rangle^{1/2}$ for different simulation setups (see \S\ref{sec:forced_turbulence_simulation},\ref{sec:lagrangian_dust_particles} and \S\ref{sec:generalization} for details). Refer to Fig.~\ref{fig:vp_rel_grouped} for the description on the abbreviated labels.}
    \label{fig:vp_rel_rad_grouped}
\end{figure*}

\subsection{Error estimation}
The predicted particle density and velocity fields, and their associated statistics, show good but not perfect agreement with the ground truth derived from the simulation data. The errors are systematic and scale-dependent. It is of interest to ask: how much of this error is due to uncertainty in the deep learning model? Addressing this question is difficult, and various approximate methods have been used in the machine learning community. We have experimented with two popular frameworks, (i) Monte-Carlo Dropout \citep{gal2016dropout,zhu2017deep} and (ii) SWA-Gaussian (SWAG) \citep{maddox2019simple}. Details are given in Appendix~\ref{sec_errors}. In summary, we find evidence using the SWAG approach that model uncertainty has a larger effect on the velocity predictions than on the density predictions. However, the magnitude of the estimated model uncertainty is significantly smaller than the size of the discrepancy between the model predictions and the ground truth. The source of the discrepancy thus remains unclear, but it is clearly of a systematic nature.

\section{Generalization Tests}
\label{sec:generalization}

Our deep learning model, trained on a single simulation of forced turbulence with a specified set of parameters, cannot be expected to generalize to distinctly different regimes of particle-gas coupling. However, it {\em may} have learned enough about how particles interact with turbulent fluids to behave robustly when confronted by modestly different turbulent conditions. To test our model's ability to generalize, we have designed three test cases with physical parameters different from the fiducial case discussed so far. The test cases are: (1) turbulence driven with a different range of cut off wavenumbers $k_{min},k_{max}$; (2) turbulence driven using a different energy injection rate of $\langle\dot{e}\rangle = 5\times 10^{-4}$; (3) particles with stopping time $\tau_s$ 100 times smaller than the original simulation. ATHENA++ simulations with these different parameters, but otherwise similar to the fiducial run, were performed for each test case.

For each of the test cases, we feed their simulation frames into the UNET-R trained from the fiducial run without any additional training, and produce field predictions for $\rho_p$ and $v_p$. Then, using the same statistics code in \S\ref{sec:normal_stat_rdf}, \S\ref{sec:rel_velc_normal} and \S\ref{sec:rel_velc_rad_normal}, we compute the relevant statistics from the field outputs and compare them with the ground truth.

\subsection{Different wavenumber cut-off}
\label{sec:generalize_k_cutoff}
In this generalization set, different minimum and maximum cutoff of driving wavenumbers $k$ are used. Instead of the original $(k_{min},k_{max})=(1,5)$ combination, we remove the largest scale driving and use the combination $(3,8)$. The resulting gas velocity power spectrum is illustrated in Fig.~\ref{fig:vg_power_spectrum}. The dip in power at around $k=1-3$ confirms that the new flow has less power in large scales compared to the original flow.

\begin{figure}
    \centering
    \includegraphics[width=0.45\textwidth]{Figures/vg_power_spectrum.pdf}
    \caption{Energy spectrum of driving flow $v_g$ using $(k_{min},k_{max})=(1,5)$ (fiducial) and $(k_{min},k_{max})=(3,8)$ (generalization test). The color bands illustrate the range of $k$ used in driving the flow for each case and the line shows the resulting power spectrum of the gas kinetic energy}
    \label{fig:vg_power_spectrum}
\end{figure}

The RDF is computed and plotted in the second column in Fig.~\ref{fig:rdf_grouped}.In the figure, we notice that the U-Net predicted RDF resembles closely that of ATHENA++. Furthermore, from the ratio plot, we see the discrepancy between the two densities takes a similar shape as that produced on the original validation set.

As mentioned in \S\ref{sec:phy_of_clustering}, the particle clustering is most responsive to eddies with size $\ell \approx \text{St}^{3/2} \ell_\eta$. Therefore, the absence  of large-scale driving should not affect the resulting clustering. This test indicates the network can pick up the correct scales in the power spectrum responsible for the clustering.

The relative velocity $\langle w^2\rangle^{1/2}$ and the radial relative velocity $\langle w_r^2\rangle^{1/2}$ behave similarly to the density. Due to the energy cascade inherent to Kolmogorov turbulence, it is expected that the removal of large-scale driving in this test case should not affect the small-scale relative velocity. Therefore a similar target relative velocity profile is expected. This is verified in the second plot in Fig.~\ref{fig:vp_rel_grouped} and Fig.~\ref{fig:vp_rel_rad_grouped}. Furthermore, the network also outputs a similar relative velocity as the target. The relative velocity of the predicted $v_p$ is relatively flat for $r\gtrsim 2\times 10^{-2}$ and raises for scales smaller than $r\lesssim 2\times 10^{-2}$. The "flat" part is at most $\approx 3\%$ lower than the ground truth and the "peak" is at most $\approx 12\%$ above the target just like that in the normal validation set. The predicted relative velocity also differ from VG by $\approx 5\%$, which indicates the network can discriminate gas velocities from dust velocities. This shows that the network correctly ignores the large-scale driving that is irrelevant to the relative velocity on small-scales and identifies the wave-numbers that are relevant for the problem. As a result the network is able to predict the relative velocity field with similar quality as the one it is trained on.

The radial component shows very similar character to the relative velocity. From Fig.~\ref{fig:vp_rel_rad_grouped}, we see that at the largest scale the prediction is $\approx 1\%$ below the ground truth and at the smallest scale it is $\approx 10\%$ larger. However, unlike the $\langle w^2\rangle^{1/2}$ predictions, the difference between the predicted radial velocity and VG is smaller in this case. Still, the predictions made by the network have similar quality as that in the fiducial case. This shows that for all statistics above, the network is indeed able to find the correct wave-numbers that aree responsible for the relative velocity.

\subsection{Different energy injection rate}
\label{sec:generalize_dedt}
In this generalization test, we drive the turbulence with an energy injection rate that is 10 times smaller than that of the fiducial case ($\langle \epsilon \rangle =5\times 10^{-3}$). As shown in \S\ref{sec:phy_of_clustering}, the length scale at which the strongest clustering occurs is given as $\ell_\tau \approx \langle \epsilon \rangle^{1/2} \tau_s^{3/2}$. With $\langle \epsilon \rangle  \to 0.1\langle \epsilon \rangle $, the clustering length scale should be $0.1^{1/2} \approx 0.32$ times the original. Hence, the dust clustering structure becomes finer and the RDF should rise at smaller scales. Since our statistics are limited to grid resolution $r\approx 256^{-1}$, the RDF instead appears flatter. This is reflected in the RDF of the simulated ground truth in the third subplot of Fig.~\ref{fig:rdf_grouped}, where the peak $g(r)$ value is reduced from $2.25$ to $1.8$.

However, the network output does not show similar scaling. As shown in the figure, starting from $r\lesssim 10^{-1}$, the network consistently underpredicts the RDF. From $10^{-2}\lesssim r\lesssim 10^{-1}$, the prediction is $\lesssim 3\%$ from the target. However, the deviation quickly grows for scales $r\lesssim 10^{-2}$, reaching an error of $\approx 15\%$ at a length of $\approx 5\times 10^{-3}$. This shows that while the network is able to partially learn the dependency of clustering on $\langle \epsilon \rangle $, it is unable to account for the senhanced clustering at smaller scales. We note that all predicted RDFs go to 1 as $r\gtrsim 10^{-1}$. This shows that the predicted dust density indeed goes to the average density at large scale, as expected from the definition of the RDF.

Similarly, reducing the energy injection rate will affect the gas velocity scale. By Kolmogorov scaling, the gas velocity in the inertial range should scale as $v_g \propto \langle \epsilon \rangle^{1/3}$. Since the dust is driven by the gas via aerodynamic drag, we expect the relative velocity of the particles to scale partly as $\propto \langle \epsilon \rangle^{1/3}$ too. Therefore, reducing $\langle \epsilon \rangle  \to 0.1\langle \epsilon \rangle $ should scale down the velocity by a factor of $10^{1/3} \approx 2.15$. This scaling is roughly reflected in the third columns of Figs.~\ref{fig:vp_rel_grouped},\ref{fig:vp_rel_rad_grouped}. The minimum and maximum target relative velocity scale of the fiducial average relative velocity and radial velocity are $\approx 5\times 10^{-2} - 2\times 10^{-1}$ and $\approx 1.3\times 10^{-1} - 3\times10^{-1}$ respectively, while that in this generalization test is $\approx 2\times 10^{-2} - 8\times 10^{-2}$ and $\approx 9\times 10^{-2} - 2 \times 10^{-1}$ respectively. Therefore, the target relative velocities are all smaller than the fiducial case by approximately the expected factor, and for small scale the difference roughly corresponds to the scaling we derived above. However, the network does not display similar scaling. As shown in Fig.~\ref{fig:vp_rel_grouped}, the network always over-predicts the relative velocity by a factor of at least $\approx 10\%$. Furthermore, the deviation grows steadily as the scale gets smaller, peaking at $\approx 40\%$ error at the smallest scale.

Analogous results is shown in radial velocity statistics in Fig.~\ref{fig:vp_rel_rad_grouped}. At all scales, the network over-predicts the statistics by at least $6\%$. This deviation continues to grow for smaller $r$ and an $\approx 20\%$ overshoot in radial velocity at grid-scale is observed.

Furthermore, comparing the VG line and the network output in Fig.~\ref{fig:vp_rel_grouped} and Fig.~\ref{fig:vp_rel_rad_grouped}, we see that the predicted particle velocities are much higher than what would have been obtained if all dust was perfectly following the flow (VG). Hence, the network is likely predicting particle relative velocities that correspond to simulations with a higher velocity scale. This reflects the network's inability to extract correct injection energy scales from the gas velocity field. While this is a failure of the model, it is expected because the network is trained only with one energy injection rate.

\subsection{Different particle stopping time $\tau_s$}
\label{sec:generalize_diff_stop_time}
Finally, we briefly discuss the results of a more challenging generalization test, in which we changed the dust stopping time $\tau_s$ for the particles to 0.001, a hundred times smaller than that of the original simulation. Since $\ell_\tau \approx \tau_s^{3/2}\epsilon^{1/2}$, this means the particle clustering scale will be scaled down to around $0.1 \%$ of the original flow. Again, by the same reasoning as \S\ref{sec:generalize_dedt}, a flat ground-truth RDF is expected. This is reflected in the rightmost plot in Fig.~\ref{fig:rdf_grouped}. The violet line, which represents the RDF of the target flow, is nearly flat and only deviates slightly from $1$ at $r\lesssim 10^{-2}$. The strong coupling means that the velocity of the particles in the simulation follows that of the gas closely at all scales.

In contrast to the less challenging generalization tests presented above, the large change in stopping time leads to a catastrophic failure in the prediction of the particle RDF. The results for the different velocity statistics, while not wrong by a similarly large factor, are systematically lower than the ground truth at all scales (See Fig.~\ref{fig:vp_rel_grouped} \& \ref{fig:vp_rel_rad_grouped}). These failures are to be expected because $\tau_s$ is a physically important parameter that our model has no knowledge of. They illustrate the importance of including simulations with different stopping times as part of the training set for future models.

\section{Discussion}
\label{sec:conclusion}
In this work, we have quantified the ability of the U-Net machine learning architecture to predict the clustering of particles that are aerodynamically coupled to turbulent gas. Training data was generated using a recently developed particle module for the {\sc ATHENA++} code, with the included physics and parameters chosen to be appropriate for the problem of small-scale turbulent clustering of dust in protoplanetary disks. After training, using a loss function that targets the three-dimensional particle density and velocity fields, we find that the U-Net model yields satisfactory performance for gridded representations of the {\em statistical} metrics that are of greatest interest for planet formation---the Radial Distribution Function and particle relative velocity. Errors on these quantities are scale-dependent, but typically at the percent level. Smaller maximum errors are found for density rather than velocity statistics. The model also generates, of course, a fast approximate mapping from the fluid to the particle fields in three dimensions. This is a novel capability that is not, to our knowledge, possible using analytic methods.

The use of a U-Net architecture requires a gridded representation of the particle density, and the definition of a gridded (single-valued) velocity field. These characteristics lead directly to some of the limitations observed in our study. The Lagrangian representation of particles in the training simulations yields useful information on particle statistics on scales {\em smaller} than those of the fluid grid, which is discarded as a pre-processing step in our approach. For applications where the sole interest is the statistical properties of the particles, better results extending to smaller scales would likely be obtained by using a convolutional neural network to target directly the radial distribution and relative velocity functions.

The model fails to generalize to different friction times and energy injection rates, as expected given that these are key physical parameters and the network was only exposed to a single value of them. There are no obvious obstacles to training the network on a suite of simulations, and thereby learning the dependence of particle clustering on physical input parameters \citep[see, e.g. analogous work in cosmology;][]{vn21}. The combination of three-dimensional spatial prediction of particle positions, together with prediction of the particles' velocity relative to the gas, would then allow calculation of particle collision rates and outcomes for arbitrary size distributions. Different approaches would be needed to incorporate the back-reaction of particles on the gas, an aspect ignored in our work but which is central to models of planetesimal formation \citep{youdin05,simon16}. Learned simulators for turbulent flows show considerable promise \citep{dang22}, and it would be interesting to apply similar methods to multi-phase fluid turbulence. The existence of new  instabilities that rely on the two-way coupling between gas and particles \citep{youdin05,squire18} suggests that multi-phase turbulence would be a stringent test of such methods.

The work presented here explores arguably the simplest problem in two-fluid turbulence that has relevance to planet formation. Related problems are present on multiple other scales in protoplanetary disks, and in other domains including atmospheric science and experimental studies of turbulence. Our results suggest that deep learning offers a complementary approach to direct numerical simulations, that can predict particle clustering at a lower computational cost.

\acknowledgments
We thank Jim Stone for helpful discussions during the development of this project. PJA acknowledges support from NASA TCAN award 80NSSC19K0639. CCY and ZZ acknowledge support from NASA via the Emerging Worlds program (award 80NSSC20K0347) and the Astrophysics Theory Program (award 80NSSC21K0141). CCY acknowledges support from NASA TCAN award 80NSSC21K0497.

\bibliographystyle{aasjournal}