GotFlow3D: Recurrent Graph Optimal Transport for Learning 3D Flow Motion in Particle Tracking

To ensure an accurate flow estimation, it is of great importance to capture reliable pointwise features, which are supposed to be discriminative enough for exact correspondences searching. Considering that the 3D coordinate is the only available input, the key point is to learn pointwise features representing the local geometry information. To achieve the effective feature exploration on the unstructured and irregular input format, we construct two kinds of geometric graphs on the particle sets and employ the static-dynamic fusion GNN (SDGNN) to derive the deep features of each particle. Specifically, we construct the static graph $\mathcal{G}^{s}$ on 3D spatial space and the dynamic graph $\mathcal{G}^{d}$ on high-dimensional feature space.

We first define the static graph $\mathcal{G}^{s}=(\mathcal{V}_{er},\mathcal{E}^{s})$, where the vertices $\mathcal{V}_{er}=\{1,2,...,n_{1}(n_{2})\}$ are composed of the entire given particle set $\mathcal{P}$ (or $\mathcal{Q}$). The edges $\mathcal{E}^{s}\subseteq\lvert\mathcal{V}_{er}\rvert\times\lvert\mathcal{V}_{er}\rvert$ are yielded by connecting each particle (vertex) $p_{i}$ to its $k$-nearest spatial neighbors $\forall p_{i_{k}}\in\mathcal{N}_{k}^{s}(p_{i})$ with the 3D Euclidean distance metric. The edges combining $p_{i}$ and its neighbors form local geometric structures and further generate the associated vertex feature $\tilde{f}_{i}=[p_{i_{k}}-p_{i};r_{i_{k}};\theta_{i_{k}};\varphi_{i_{k}}]$, $\tilde{f}_{i}\in\mathbb{R}^{k\times 6}$, where $p_{i_{k}}-p_{i}$ denotes the 3D directional vector. In addition, $r_{i_{k}},\theta_{i_{k}}$ and $\varphi_{i_{k}}$ represent the radial distance, azimuthal angle and polar angle in spherical coordinates, respectively, which supplement the relative spatial relationship of the neighboring vertices from another perspective. Such spherical representations further provide rotation invariant correspondences, which are advantageous in the estimation of complex motions and larger deformations [30]. The static graph for a particle set needs to be computed only once before training and is kept for the following retrieval.

As shown in Figure 1(d), given $\mathcal{G}^{s}$ and $\mathcal{P}$, we employ the SDGNN to encode the local geometric characteristics to generate the distinctive feature descriptor for each particle. We denote the extracted pointwise feature of vertex $p_{i}$ at the $\beta$-th layer as $g_{i}^{\beta}\in\mathbb{R}^{D_{\beta}}$. In the deep hierarchical SDGNN, graph convolution operations are performed on the local geometric structure of each vertex $p_{i}$ to map the vertex feature $g_{i}^{\beta-1}$ of the previous layer to a new feature $g_{i}^{\beta}$ in the subsequent layer. Inspired by the PointNet++ framework [31], we propose a GeoSetConv layer to further embed the vertex features of the static graph to a high-dimensional feature space to better characterize the local geometric structures:

here ${{}^{s}g_{i}^{\beta}}$ denotes the feature extracted from the static graph at the $\beta$-th layer, and $\mathcal{N}_{s}^{k}(p_{i})$ depicts the $k$-nearest neighbors of vertex $p_{i}$ in the static graph. In addition, $\phi^{\beta}$ is the $\beta$-th GeoSetConv layer, which is implemented as multi-layer perceptrons (MLPs). And $\max$ denotes the max-pooling operation, which aggregates features over the $k$ neighbors. The ${}^{s}g_{i}^{0}$ is initialized by the 3D coordinate of $p_{i}$ in the first GeoSetConv layer.

To further enrich the feature representations describing the detailed topological information, we explore the feature learning on the dynamic feature graph, as EdgeConv did in [32]. It is the first attempt at applying the dynamic graph CNN to extract features in the flow learning task. The following experiments also demonstrate that the introduction of the dynamic feature graph contributes to the improvement of the accuracy. The dynamic graph $\mathcal{G}^{d}$ is constructed and updated by aggregating the $k$-nearest neighbors on the high-dimensional feature space generated at each layer. Specifically, for each layer, we build the dynamic graph $\mathcal{G}^{d}=(\mathcal{V}_{er},\mathcal{E}^{d})$, where $\mathcal{G}^{d}$ shares the same vertices as that of $\mathcal{G}^{s}$ and the edges $\mathcal{E}^{d}$ link feature level neighbors to form local structures in the feature manifold. Then, we formulate the update of the vertex feature ${{}^{d}g_{i}^{\beta}}$ of $p_{i}$ on the dynamic graph at the $\beta$-th layer as

here $\mathcal{N}_{k}^{f}(p_{i})$ depicts the index of the $k$-nearest neighbors of $p_{i}$ in the feature space. And $\varphi^{\beta}$ denotes the $\beta$-th EdgeConv layer, which is also implemented as MLPs. Different from the static graph, whose receptive field is limited in the local neighborhood and hinders the acquirement of long-term dependencies, the dynamic graph enlarges the receptive field to the whole input set. The dynamic graph learns to construct a special receptive field by assigning more attention to the neighbors with similar feature (flow) properties.

Then, the proposed hierarchical SDGNN incorporates the extrinsic geometric topology and the intrinsic feature correspondences by fusing the learnt features from the two types of graphs. More precisely, we further concatenate the hierarchical features of these graphs from different layers of the GNN. Finally, several MLPs with a skip connection are employed to derive the final deep pointwise feature

where ${\psi}$ depicts the MLPs and $\mathop{\oplus}\limits_{\beta}(\cdot,\cdot)$ denotes the concatenation of features from different graphs and different layers. The extracted pointwise feature is adopted for the following optimal transport planning to infer the particle correspondences via the feature similarity. Furthermore, we employ another feature encoder to extract the context feature $\mathcal{F}_{\mathcal{C}}$ of $\mathcal{P}$, which provides additional context information to the flow update. Such context feature encoder simply consists of several SetConv layers from PointNet++, which only operates on the spatial static graph.

Optimal transport. Optimal transport theory [33] guides the transportation and allocation between two distributions. In this work, we regard the correspondences searching of two particle sets as the transportation and allocation of the particle mass, where each particle $p_{i}$ ($q_{j}$) is set to have a mass of $n_{1}^{-1}$ ($n_{2}^{-1}$). Then, the cost has to be paid to transport each unit of mass from $p_{i}$ to $q_{j}$. Therefore, the key issue of optimal transport is to seek a transport plan transforming a source particle distribution $\mathcal{P}$ to a target particle distribution $\mathcal{Q}$ whilst minimizing the transport cost [34]. In more detail, the computation of a transport plan $\mathcal{T}\in\mathbb{R}_{+}^{n_{1}\times n_{2}}$ between two particle sets can be formulated as

where $C_{ij}$ is the transport cost from particle $p_{i}\in\mathcal{P}$ to $q_{j}\in\mathcal{Q}$, and $U_{ij}$ denotes the assigned quantity of mass to be transported from $p_{i}$ to $q_{j}$. In addition, we define $\bm{\mu}_{p}=\frac{1}{n_{1}}\textbf{1}_{n_{1}}$ and $\bm{\mu}_{q}=\frac{1}{n_{2}}\textbf{1}_{n_{2}}$ in the equality constraints. The constraint manifests that the mass of each particle $p_{i}$ from the source distribution is supposed to be completely transported to the target and each particle $q_{j}$ of the target distribution is expected to be filled up with the transported mass.

However, due to the measurement noise in real experiments such as 3D reconstruction errors and object occlusions, the map reflecting the correspondences between two particle sets is neither injective nor surjective. Therefore, the mass constraints cannot hold ideally. To cope with this problem, it is preferred to apply the constraint relaxation to solve the optimal transport problem. We follow the formulation proposed by FLOT [26] to employ the Sinkhorn algorithm [35] to efficiently optimize the transport plan with an entropic regularization and a mass regularization:

here $KL$ represents the $KL$-divergence. In addition, $\epsilon$ and $\lambda$ denote the entropic and the mass regularization parameter, which are also learnt through the network training. A small $\epsilon$ leads to a sparse transport plan, while a high $\lambda$ tends to reinforce the effect of the mass constraint. Then, we define the transport cost between all pairs as the dissimilarity of the extracted pointwise feature via the cosine similarity:

where $\|\cdot\|_{2}$ denotes the $l_{2}$-norm. A small cost between two particles generally corresponds to a high plan, which implies a latent possibility of particle correspondences. The transport plan is constructed only once and is kept available for the following retrieval, which conduces to avoid numerous repeat matrix calculations.

Adaptive point-voxel optimal transport. In this paper, we propose a novel approach to adaptively retrieve the optimal transport plan for the iterative flow updating. For each iteration $t$, given the currently estimated flow $\mathcal{V}_{t}$, the source $\mathcal{P}$ is translated to $\mathcal{P}^{\prime}_{t}$, where $P^{\prime}_{t}=P+\mathcal{V}_{t}$. The estimated flow is initialized at $\mathcal{V}_{0}=0$ and accumulated by the previous flow updates as

where $l$ denotes the index of the iteration. Next, we can consult the transport plan between $\mathcal{P}^{\prime}_{t}$ and $\mathcal{Q}$ to guide the flow update $\Delta f_{t+1}$ of the next iteration. More specifically, the consultation is adopted by searching the neighbors of the warped source $\mathcal{P}^{\prime}_{t}$ in the target $\mathcal{Q}$ and retrieving the corresponding transport plan $\mathcal{T}_{t}$. This iterative transport plan consultation contributes to lowering the difficulty of long range motion estimation and promote the refinement of small scale motion.

An essential step in the above-mentioned algorithm is: how to efficiently retrieve abundance and reliable information ($\mathcal{T}_{t}$) from the full transport plan ($\mathcal{T}$) through the neighborhood relation reasoning of the unstructured particle set. To address this issue, the proposed approach further generates the adaptive point-voxel optimal transport (AdaPT) plan $\mathcal{T}^{*}_{t}$ during each consultation based on the original point-voxel operation [28]. More precisely, we integrate the captured point-based transport plan $\mathcal{T}^{*}_{p}$ and the voxel-based transport plan $\mathcal{T}^{*}_{v}$ to offer the long-range and small-scale correspondence of particle pairs in the consultation:

First, we present the generation of the point-based transport plan $\mathcal{T}^{*}_{p}$ using a dynamic $k_{t}$-nearest neighbors searching, where the number of the located neighbors are gradually reduced as the iteration progressed. Such construction adaptively narrows the search region to refine the fine-grained correspondence along with the increasing of iterations. From the deep point of view, the learning of the subsequent smaller residual flow amounts to cope with particles with relatively larger spatial distance. Accordingly, we dynamically reduce the searched neighboring particles to adapt to sparsely seeded regions, in which small number of neighbors perform better as noted in [30]. The computation of $\mathcal{T}^{*}_{p}$ and the attenuation of $k_{t}$ are formulated as:

here $\mathcal{N}_{Q}^{k_{t}}(P^{\prime}_{t})$ denotes the index set of the $k_{t}$-nearest neighbors of $P^{\prime}_{t}$ in $Q$, and $\mathcal{T}(Q_{j})$ is the corresponding retrieved transport plan between $Q_{j}$ and $P^{\prime}_{t}$. In addition, we further combine the retrieved $\mathcal{T}(Q_{j})$ and the relative spatial relation $Q_{j}-P^{\prime}_{t}$ of the corresponding neighbors to deep encode the final plan $\mathcal{T}^{*}_{p}$, where the local spatial relation plays a similar role as the geometric attention weight. The deep encoding is performed by the MLPs $\Phi_{1}$ with a $\max$ operation over the $k_{t}$ dimension. Furthermore, $k_{0}$ is the initial value of $k_{t}$, while $\eta$ is set to a constant value.

To further cope with the large-scale motion embedded in the long-range correspondence of transport plan, we innovatively construct adaptively deformable multi-scale cuboids, which compartmentalize the geometric space containing particles into pyramid voxels to yield the voxel-based transport plan $\mathcal{T}^{*}_{v}$. Specifically, each constructed cuboid is centered at each particle of $P^{\prime}_{t}$ and collects the enveloped neighbors in $Q$. As depicted in Figure 1(e), the cuboid is composed of $d^{3}$ identical sub-cuboids, whose length ($\vec{e}_{1}$), width ($\vec{e}_{2}$) and height ($\vec{e}_{3}$) vectors are adaptively defined according to the several recent flow updates as:

where $\mathcal{S}(\lvert\Delta f_{l}\rvert)$ and $\mathcal{S}(\lvert\Delta f_{l-1}\rvert)$ denote the normalized length of the current updated flow and the last updated flow over all particles in $P^{\prime}_{t}$, respectively. Such normalization encourages a longer-range correspondence search when large variations appear on the flow. And $cos\left\langle\Delta f_{l},\Delta f_{l-1}\right\rangle$ calculates the cosine similarity between $\Delta f_{l}$ and $\Delta f_{l-1}$. Such similarity calculation also reinforces a longer-range correspondence search when recent flow updates keep consistent but prompts a smaller-range correspondence search when they remain inconsistent. Moreover, $\vec{l}_{1}$, $\vec{l}_{2}$ and $\vec{l}_{3}$ are unit orthogonal basis vectors of the three principal orientations used to guide the search direction of the correspondence, which are established as:

where $\left\langle\cdot,\cdot\right\rangle$ denotes the inner product of two vectors, while $\otimes$ represents the outer product; $\|\cdot\|_{n}$ means that the vector is normalized to unit length. Furthermore, we define a basic scale parameter $r$, which controls the size deformation of the cuboid:

where $d_{m}$ denotes the mean distance between each particle and its $k_{d}$-nearest neighbors in the input set, which is utilized to obtain the local particle concentration information. And $\alpha$ is a learnable scale parameter optimized with the network training, which contributes to an adaptive process of particle sets with various particle densities. In addition, $s$ is the index of the pyramid iteration, which gradually doubles the length size of the sub-cuboid to retrieve the multi-scale voxel-based plan. The adaptive variation of the multi-scale cuboid size $r$ makes it possible to cover farther desired particles.

As illustrated above, the constructed cuboid presents the powerful deformability to adapt to current flow updates and input particle densities, while showing puissant capabilities to capture long-range correspondences with the multi-scale strategy. Then, the cuboid partitions the enveloped neighboring particles in $Q$ into $d^{3}$ sub-cuboids, which indicate the relative direction of the neighbours to the central particle in $P^{\prime}_{t}$. Furthermore, the neighboring particles and the corresponding transport plan in each sub-cuboid are aggregated and averaged for the final calculation of $\mathcal{T}^{*}_{v}$:

here ${}^{s}\mathcal{N}_{Q}^{b}\left(P^{\prime}_{t}\right)$ denotes the neighboring particles of $P^{\prime}_{t}$ in $Q$ enveloped in the $b$-th sub-cuboid for the $s$-th pyramid iteration. And $\overline{\mathcal{T}}$ is the average over the transport plan of particles in the $b$-th sub-cuboid. Moreover, $\mathop{\uplus}_{s}$ and $\mathop{\uplus}_{b}$ denote the concatenation over all pyramid iterations and all sub-cuboids, respectively. There followed another MLPs $\Phi_{2}$, which is used to implement the similar deep encoding for the final plan $\mathcal{T}^{*}_{v}$.

GRU flow refinement. Inspired by RAFT [36], we finally employ a recurrent GRU to estimate the flow update $\Delta f_{t+1}$, which mimics the iterative strategy in the optimization algorithm. The applied GRU consumes the integrated AdaPT plan $\mathcal{T}^{*}_{t}$, the current flow estimation $\mathcal{V}_{t}$, the context feature $\mathcal{F}_{\mathcal{C}}$ and the hidden state $h_{t-1}$ transmitted from the last iteration, while yielding a new hidden state $h_{t}$ encoding the future flow information:

where $x_{t}$ denotes the concatenated feature of $\mathcal{T}^{*}_{t}$, $\mathcal{V}_{t}$ and $\mathcal{F}_{\mathcal{C}}$. And $\odot$ is the element-wise product. In addition, $W_{z(r,h)}$ denotes the weights of convolutional layers in different steps. Then, the $h_{t}$ is passed through a SetConv layer and several convolutional layers to predict the flow update $\Delta f_{t+1}$.

In this framework, there are some predefined parameters in the graph construction and the AdaPT plan retrieval process. The number of the searched neighboring particles used to construct the static graph and the dynamic graph is set as $k=32$. For the AdaPT plan retrieval, the dynamic $k_{t}$ is initialized as $k_{0}=32$, while $\eta$ is fixed to 2. And when the iteration exceeds 8, the $k_{t}$ is fixed to 16. The resolution of the cuboid $d$ and the number of the pyramid iteration $s$ are both set as 3. In addition, the $k_{d}$ is set as 3 to calculate the mean local particle distance. Moreover, the scale parameter $\alpha$ is initialized as 0.5 and has a special learning rate, which is 0.1 times the global learning rate, to enhance the training stability. The dimension of the final pointwise feature extracted by the hierarchical GNN is set as 128. Furthermore, the $\epsilon$ and $\lambda$ used to control the two regularizations are initialized as 0 and 1, respectively. And the flow is iteratively updated for 8 iterations.

The network training are performed with supervised loss using the entire estimated flow sequence $\{\mathcal{V}^{e}_{1},...\mathcal{V}^{e}_{T}\}$:

where $\mathcal{V}^{e}_{t}$ denotes the flow estimation of the $t$-th iteration, while $\mathcal{V}^{gt}$ is the ground truth flow. In addition, $\|\cdot\|$ denotes the $l_{1}$ distance, and $T$ is the total number of flow update iterations. Moreover, $N$ denotes the total number of the particles. Mini-batch training is employed with a batch size of 16. Then, we further impose an exponentially increased weight $\gamma_{t}$ to the estimated flow sequence:

where $\mu$ is set to 0.8. The network is trained on six TITIAN XP GPUs for 40 epoches using Adam optimizer [37] with a initial learning rate of 0.001.

Training data with ground-truth information plays an indispensable role in the network optimization with supervised learning strategies. However, the difficulty of recovering highly accurate ground-truth flow from real data enforces us turn to the simulated data from computational fluid dynamics (CFD), which is capable of generating various flow cases representing the true fluid mechanical characteristics by solving the Navier–Stokes equations [16].

To perform the flow supervision, we create the first synthetic 3D fluid flow dataset named FluidFlow3D, which is characterized by two consecutive particle sets $\mathcal{P}$, $\mathcal{Q}$ and the corresponding ground-truth flow set $\mathcal{V}^{gt}$. In particular, we first generate a particle set $\mathcal{P}$, which comprises more than two thousand randomly distributed particles in the 3D observation volume. The observation volume is also defined in the random location of the CFD computational domain. Next, we use the public CFD data from Johns Hopkins Turbulence Database (JHTDB) [38], which is proven to provide physically correct simulated flow structures, to generate the ground-truth flow $\mathcal{V}^{gt}$. This is achieved by applying the built-in functions to query the JHTDB database with the second order accurate Runge-Kutta integration scheme in time and the fourth order Lagrange interpolation in space, which are executed on the database nodes. The flow structures queried from JHTDB includes the incompressible isotropic turbulent flow, the incompressible magneto-hydrodynamic (MHD) turbulence, the fully developed turbulent channel flow and the transitional boundary layer flow. Therefore, we can get the next frame $\mathcal{Q}_{0}$ of the given particle set $\mathcal{P}$ using the integrated trajectory with a designated time interval, subsequently get the the corresponding ground-truth flow $\mathcal{V}^{gt}=\mathcal{Q}_{0}-\mathcal{P}$. The sketch of this generation procedure is illustrated in Figure 1(b). In addition to the data from JHTDB, several flow benchmarks, including the uniform flow and the Beltrami flow [39], are also generated to increase the diversity of the training dataset. Eventually, the synthetic training dataset consists of six categories of typical flow cases.

Furthermore, to enhance the generalization capabilities of the trained network for realistic experimental configurations, we augment the above flow cases with various flow conditions. The measurement performance in PTV experiments is highly limited by large dynamic velocity ranges and high particle densities. These two factors can be characterized by the particle displacement ratio $\rho$, which has been widely used in the PTV community [40] and can comprehensively depict the flow conditions:

here $d_{a}$ and $u_{\max}\Delta t$ denote the average distance between particles in the whole set (e.g., $\mathcal{P}$) and the maximum displacement between two input sets, respectively. It is reported that the classical flow estimators can easily achieve satisfactory performance for flow cases with $\rho\gg 1$, while facing more challenges when $\rho$ gets smaller [40]. In our training dataset, we change the characteristic observation volume to vary $d_{a}$ (the number of particles remains constant) and alter the time interval $\Delta t$ to obtain different maximum displacements. As a result, this important parameter $\rho$ in the synthetic dataset ranges from $0.35$ to $6.58$, much beyond the normal range of real PTV experiments. Moreover, to mimic the real-world applications, we add some noises to create more realistic cases, which mainly arise from the physical motion of particles in and out of the field of view, the measurement noise and the occlusion. To this end, we shuffle the particles in $\mathcal{Q}_{0}$ to obtain the second set $\mathcal{Q}$, which directly disorders the correspondence. Then, we randomly replace $n\%$ of the particles in $\mathcal{Q}$ with new random particles. Such noises eliminates $n\%$ of the correspondence of $\mathcal{P}$, while introducing new isolated particles in $\mathcal{Q}$ with no correspondence. Here we set $n$ with $1\%$ level of noise for this dataset. Finally, we generate a synthetic dataset with 16K training samples and 1.6K validation samples. A table summarizing the synthetic dataset is given in Table III.

A family of synthetic FluidFlow3D-derived datasets, a public turbulent cylinder wake flow dataset (denoted by CylinderFlow) [41] and a real-world dataset (denoted by SphericalIndentationFlow) across different domains are adopted to conduct the experimental evaluations.

FluidFlow3D-family. Apart from the aforementioned FluidFlow3D, we further transform FluidFlow3D to a new dataset with normalized spatial scale (denoted by FluidFlow3D-norm) to compute the unified evaluation metric among flow cases with different observation volumes. The FluidFlow3D-norm is used for the comparison with the state-of-the-art methods and ablation studies. To analyze the robustness, we create a synthetic subset FluidFlow3D-noise with the different noise level. Furthermore, another subset FluidFlow3D-ratio regarding various $\rho$ is also generated.

CylinderFlow. This public dataset contains Eulerian velocities and Lagrangian particle trajectories of the cylinder wake flow with a Reynolds number of 3900 [41]. Such flow data is also calculated using CFD. In addition, the fourth-order Runge-Kutta scheme in time and trilinear interpolations in space are utilized to transport about 200,000 particles to obtain the Lagrangian trajectories. Trajectories of two 3D observation volumes are available for the assessment of tracking algorithms.

DeformationFlow. This experiment is conducted by measuring the volumetric deformation of a soft polyacrylamide (PA) hydrogel exerted by the spherical indentation [30]. A stainless steel sphere is placed on the surface of the PA hydrogel, which leads to the indentation deformation due to the force of gravity. Then the 3D volumetric image of the hydrogel material as well as the embedded fluorescent particles are scanned both before and after the deformation with multiphoton microscopy.