P$^2$C$^2$Net: PDE-Preserved Coarse Correction Network for efficient prediction of spatiotemporal dynamics

Thank you for achknowledging the data-efficiency and satisfactory performance on OOD tests of our work. We address your concerns as follows.

Weakness

Q1(a). Concerns on the larger size of the Poisson solver and neural correction modules compared to baselines.

The Poisson solver is essentially a numerical method (e.g., a spectral solver). The basic idea is to convert the original problem into the frequency domain, where the spatial differential operation becomes a multiplication operation via Fourier transform. The Poisson equation is then solved in the frequency domain, and the inverse Fourier transform is then used to obtain the solution in the original spatial domain. This entire process does not involve a Fourier neural network or any learnable parameters.

The correction block utilizes only two Fourier layers that retain a few number of modes, and is therefore not considered a large-scale neural network. Given the unique network architecture, our model achieves good generalization performance under a small and sparse data environment, highlighting its advantages over existing baseline models. For your reference, we show the number of model parameters in Table R1 in the 1-page PDF rebuttal file.

Q1(b). Training efficiency of the proposed method.

Please see our response to Q4 in General Reply Q4 for more details.

Questions

Q1. Training time on all benchmark examples compared to baselines.

Excellent comment! The training time for different mothods is shown in Table A below. It can be observed that the training time of our model is generally acceptable. Note that the training has been performed on a single 80GB Nvidia A100 GPU, as described in Appendix F in the paper.

Table A. Training time for different models.

Q2. Robustness to noisy or incomplete training data.

Great remark! Using the Burgers equation as an example, we introduced Gaussian noise of varying scales during the training process and observed its minimal impact on the results. The results are presented in Table B below, which indicate that our model is robust to certain noise and maintains the HCT (high correlation time) metric without reduction.

Table B. Impact of noise on P$^2$C$^2$Net performance.

Moreover, for the Burgers case, the time steps of the training data consist of only 28% of the inference steps in the test data. Namely, our training data itself is incomplete, and the missing data accounts for 72% of the test set. Based on this sparse dataset, we further randomly reduced the training data by 20% to simulate data incompleteness (e.g., randomly deleting data according to the rollout step size to make the trajectory incomplete, where the specific size can be found in Table 1 in the paper). The results of this experiment are shown in Table C below, where values are averaged over 10 test sets. Additionally, we conducted the same experiment on the GS dataset, with the results presented in Table D below. We can observe that, after making the data sparser, our model performance slightly decreases, which means that our model is capable of handling scenarios with incomplete data.

Table C. Impact of sparser Burgers dataset on P$^2$C$^2$Net performance.

Table D. Impact of sparser GS dataset on P$^2$C$^2$Net performance.

Q3. Applicability to different BCs beyond the periodic BCs.

Despite we showcase the eficacy of our model on datasets with Periodic BCs, it is applicable to handle other types of BCs, such as Dirichlet and Neumann BCs (see Table R2, Table R5 and Fig. R1 in the 1-page PDF rebuttal file). Please see our detailed response to Q1 in General Reply.

Remark: We sincerely thank you for yopur constructive comments. We hope the above responses are helpful to clarify your questions. We will be happy to hear your feedback and look forward to addressing any additional questions. Your consideration of improving the rating of our paper will be much appreciated!

Dear Reviewer bvnV,

We hope our point-by-point response in the rebuttal has addressed your concerns. We are very much looking forward to your feedback during the discussion period. We would be more than happy to answer any further questions you may have.

Best regards,

The Authors

Dear Reviewer bvnV:

As the author-reviewer discussion period will end soon, we would appreciate it if you could kindly review our responses at your earliest convenience. If there are any further questions or comments, we will do our best to address them before the discussion period ends.

Thank you very much for your time and efforts!

Sincerely,

The Authors

Thank you for acknowledging the novelty and the effectiveness of our proposed method. We address your concerns as follows.

Weakness

Q1. Clarifications of each block component, their connections and differences.

The clarifications of each block component and the connections can be found in our response to Q2 in General Reply.

It is noted that the correction block aims to address the information loss that occurs due to resolution reduction before calculating derivatives, helping the model adapt to the coarse grid. During training, it serves as a scaling factor for the derivative term. Achieving this with only a shallow network (2 layers with modes 12) is feasible. However, since the network is trained on a coarse grid with a high spatial coarsening ratio (up to $32^2=1024$ on NS), the reduced high-resolution values do not fully satisfy the governing equation. To mitigate error accumulation during the rollout prediction, compensating the PDE block is necessary. This requires a deeper network (e.g., 4 to 6 layers of FNO).

Q2. Difference between our method and PeRCNN except for the Poisson block.

The differences between these two methods are listed as follows.

-PeRCNN utilizes feature map multiplication to build polynomial combinations (with redundant terms) to approximate the underlying governing equation, which has great performance given fine mesh grids. In contrast, our proposed model directly incorporates the governing PDE, preserving its complete and precise form, to predict spatiotemporal dynamics on coarse mesh grids under a small and sparse data environemnt.

-Learning the diffusion term (e.g., $\Delta \mathbf{u}$) in the governing PDE is challenging. PeRCNN employs fixed finite difference Conv to compute the diffusion term, leading to significant errors. Our approach introduces a learnable Conv with symmetric constraint to compute the diffusion term, effectively mitigating errors on the coarse grid.

-PeRCNN is limited to testing simple cases (e.g., small domains and no external forces) and struggles with generalization to external forces and Reynolds number. Finally, PeRCNN aims to learn high-resolution dynamics with an initial state generator super-resolving from low-resolution grids.

Hence, our model is fundamentally different from PeRCNN.

Q3. Adaptations for NS equations and fair comparisons by adding NN blocks for baselines.

The NN block is a flexible component of our model that can be included or excluded during training based on demand. Since the PDE block alone is sufficient to achieve the state-of-the-art performance on Burgers, FN, and GS datasets, we chose to remove the NN block in these cases to establish a lighter model. However, for the NS equation, given the complexity of the flow ($Re=1000$), we need the NN block as a supplement to alleviate the error accumulation of long-term prediction and the instability issue on coarse grids.

To address the fairness comparison issue, we integrated the NNblock into both UNet and FNO models. The results, presented in the Table R1 in the 1-page PDF rebuttal file, show that while NNblock enhances the performance of both models, the improvements are still significantly lower compared to our proposed model.

Questions

Q1. End-to-end training or individual training?

Good question! Our model features end-to-end training, making the training process user-friendly and avoiding separate training of individual modules. This also can lead to better overall performance and coherence in the model's outputs, as the interdependencies between different parts of the model are explicitly accounted for during training. Additionally, the inclusion of prior knowledge facilitates the model convergence as shown in Fig. R4 in the 1-page PDF rebuttal file. We will clarify this in our revised paper.

Q2. More training for generalization tests, such as external force scenarios and Reynolds numbers.

No additional training was performed for the generalization tests. This aspect has been detailed in our paper, specifically in Section 4.2 (Generalization Test). The model was trained solely on five training sets (trajectories) with Reynolds number $Re=1000$ and external force $\mathbf{f} = \mathrm{sin}(4y)\mathbf{n}_{\mathit{x}} - 0.1\mathbf{u}$. Namely, once our model is trained, it can generalize to the different ICs, Reynolds numbers, and external force terms as shown in the paper.

Remark: We sincerely thank you for putting forward constructive comments/suggestions. We hope the above responses are helpful to clarify your questions. We will be happy to hear your feedback and look forward to addressing any additional questions. Your consideration of improving the rating of our paper will be much appreciated!

Dear Reviewer aBLf,

We hope our point-by-point response in the rebuttal has addressed your concerns. We are very much looking forward to your feedback during the discussion period. We would be more than happy to answer any further questions you may have.

Best regards,

The Authors

Dear Reviewer aBLf:

As the author-reviewer discussion period will end soon, we would appreciate it if you could kindly review our responses at your earliest convenience. If there are any further questions or comments, we will do our best to address them before the discussion period ends.

Thank you very much for your time and efforts!

Sincerely,

The Authors

Weakness

Q1(a). Generation of ICs of GS model for training and testing data.

First, to create ICs for the GS equation, we define a grid based on the spatiotemporal resolution and initialize the concentrations of chemicals A and B. Second, we set different random seeds to add random noise, changing the values of the chemicals at different positions. Finally, we obtain different ICs. Hence, the training and testing samples are independent. We will clarify this in our revised paper.

Q1(b). Concern on generalization with limited training data.

Our model is designed to adhere to PDEs. While different ICs lead to varied steady-state patterns, these variations are ultimately governed by the underlying PDEs. Our model encodes the PDEs within a learnable PDE block which helps reduce the model’s dependence on data (please see our reply to Q3 in General Reply). This module ensures the model to accurately capture the underlying dynamics even on a coarse grid and with limited data. Since PDE is embedded into the network with RK4 integrator, the model is able to generalize prediction of solution trajectories over different ICs.

Moreover, we employ different rollout timesteps during the training stage to improve the model's performance. That is, we treat a geneirc frame as an IC, and divide the trajectories into many samples based on the selected rollout timestep, which are then shuffled and grouped into batches for training (see Section 4.1 in our paper). As a result, the trained model is capable of extrapolating the prediction in the time horizon.

Finally, we also demonstrated that our model can generalize to different types of boundaries (as discussed in our reply to Q1 in General Reply).

Q2. Clarifications on $\tilde{x}$ and $x$, and $u$ in Eq. (1) and use $\hat{u}$ in Eq. (3).

We downsample the data from the fine grid to the coarse grid, denoted by $x$ and $\tilde{x}$ respectively, as mentioned in Sections 3.1 and 3.2.1 in our paper. We denote $\mathbf{u}$ as the ground truth of spatiotemporal dynamics and $\hat{\mathbf{u}}$ as the solution after passing through the correction block, resulting in a corrected coarse solution. This process can be represented as $\hat{\mathbf{u}}_k=`NN` (\mathbf{u}_k)$, where the correction block is a trainable network defined in Section 3.2.2. Hope this clarifies your question.

Q3(a). Fair comparisons between our method and baselines.

Our baselines include both purely data-driven (UNet and FNO) and physics-aware (PeRCNN and LI) methods. The reason why we consider UNet and FNO as baselines is to test their extreme performance on small datasets. Initially, we also considered DeepONet, but it was removed from the NS case due to its poor performance in the first three experiments.

We have also added temporal stencil modeling (TSM) [Sun, et al. ICML 2023] and learned correction (LC) [Kochkov, et al. PNAS 2021] as additional baselines, which use a history of solution as input and incorporate physics knowledge. The correlation curve is ploted in Fig. R2 in the 1-page PDF rebuttal file. We believe the additional experiments would improve the fairness in the baseline comparisons.

Q3(b). Details of settings and configurations of runtine comparisons in Figure 6.

Please see our detailed response to Q4 in General Reply. We will clarify this in our revised paper.

Questions

Q1. The purpose and necessity of the lower path in Figure 1(a).

We present the role of each module in our response to Q2 in General Reply and discuss their importance in our response to Q3 in General Reply. Given the potential instability and accumulated errors in the learnable PDE block, especially in the NS example, the lower path is indispensable. It serves to correct the solution generated by the PDE block. Moreover, since the input of the NN block includes pressure ${p}$ to provide more information, the lower path also includes an optinal Poisson block. Adding such a block help improve the model's performance as shown in Table R6 in the 1-page PDF rebuttal file.

Q2. Initialization of conv filters and the correction block, and learning curves in training.

We use random initialization for the symmetric Conv filters and scale the parameters to a low magnitude with an empirical value of 0.001. For the correction module, the Conv filters use the default Kaiming Uniform initialization, while other components use random initialization. We also plot the learning curves (averaged over 5 independent training trails for the NS dataset) as an example, shown in Fig. R4 in the 1-page PDF rebuttal file.

Q3. Sampling process for datasets and fairness of the testing datasets.

For the Burgers, FN, and GS datasets, which utilize the finite difference method as solver, the solution is calculated at mesh grids. To establish a coarse grid ground truth, we apply uniform downsampling [Rao, et al. Nature Machine Intelligence 2023] in both temporal and spatial dimensions. In contrast, the NS dataset, which uses the finite volume method, computes solutions at cell faces in a staggered manner. Therefore, we use staggered average downsampling [Kochkov, et al. PNAS 2021] (shown in Fig. R5 in the 1-page PDF rebuttal file) in the spatial dimension and uniform downsampling in the temporal dimension.

For fairness of testing datasets, we randomly select 10 seeds from the range of [1, 10,000], which are used to generate different Gaussian noise disturbances, which are then applied to the initial velocity field to create varying ICs. We also perform a warm-up phase during data generation to ensure that the variance and mean of the trajectories are closely aligned, thereby maintaining fairness. We wiil clarify this in the revised paper.

Remark: Thanks for your constructive comments. Your consideration of improving the rating of our paper will be much appreciated!

Dear Reviewer WfL9,

We hope our point-by-point response in the rebuttal has addressed your concerns. We are very much looking forward to your feedback during the discussion period. We would be more than happy to answer any further questions you may have.

Best regards,

The Authors

Dear Reviewer WfL9:

As the author-reviewer discussion period will end soon, we would appreciate it if you could kindly review our responses at your earliest convenience. If there are any further questions or comments, we will do our best to address them before the discussion period ends.

Thank you very much for your time and efforts!

Sincerely,

The Authors

Weakness

W1. Different BC encoding and derivative calculations.

Despite we showcase the eficacy of our model on datasets with Periodic BCs, it is applicable to handle other types of BCs (see Table R2, Table R5 and Fig. R1 in the 1-page PDF rebuttal file). Please see our detailed reply to Q1 in General Reply.

W2. Challenges for unstructured grids.

Great remark! We acknowledge the challenges on extending the current pipeline to unstructured meshes. Nevertheless, the overarching framework is extensible, e.g., integrating graph neural networks and differentiable finite element/volume methods to learn on unstructured meshes. We will discuss the challenges in the revised paper.

Questions

Q1. Re-organization of PDE-block, Poisson block, & NN block in Fig. 1.

Considering the sequence of actions, placing the PDE block in Fig. 1.b first, followed by the Poisson block, would be better. We will modify it in the revised paper.

Q2. Clarification of three input components for time integration in the PDE block.

Please see our reponse to Q1 in General Reply.

Q3. Stability of Poisson block in early stage of training.

The Poisson block mainly relies on a numerical solver without any trainable parameters. We scale the initialized weights of the symmetric filter to ensure that the Laplacian term has a low magnitude. Thanks to the correction block, the solution does not reach a large magnitude, either. We added a loss curve to prove this point (averaged over 5 independent training trails), shown in Fig. R4 of the 1-page PDF rebuttal file. The consistent performance across these trials demonstrates that our model does not encounter instability during early training. We also explored relaxing the Poisson block by adding trainanle parameters (e.g., scaling coefficients). However, no improvement is observed.

Q4. Equation residue block is the RK-4 integration.

You are right -- the PDE block is indeed designed based on the RK4 integrator. However, the PDE block is trainable, consisting of the correction block and a trainable filter bank, which differs from the standard RK4-based finite diference (FD) solver.

Q5. Comparing derivatives from kernels on coarse grid with those from FD on fine grid.

The derivatives of these two strategies are different. We learn an approximate surrogate of derivatives on the coarse grid, which are different from the ground truth derivative values. Our goal is to minimize the overall PDE residual, instead of making each derivative equal to the corresponding ground truth.

Q6. More baselines that learn corrections and take history of solutions as inputs.

We further considered 3 correction learning baselines, including learned correction (LC) [PNAS 2021, 118(21):e2101784118], two variants of UNet and FNO by adding the NN Block for correction. The results in Table R1, in the 1-page PDF rebuttal file, showed that our model has the best performance.

We have also trained UNet and FNO that use 12 history snapshots as inputs. However, their performance didn't improve. Note that LI and TSM in our baselines also use historical information (32 steps) as inputs. These baselines cannot capture the underlying dynamics due to the limited training samples.

Q7. Filter size and its dependence on grid size.

Following your suggestion, we conducted experiments on different filter sizes, e.g., 3x3, 5x5, 7x7. As shown in Table R4 in the 1-page PDF rebuttal file, the model with the 5x5 filter outperforms the other two (the 3x3 filter has low accuracy while the 7x7 filter has instability issue on coarse grids). Hence, the 5x5 filter is empirically suggested.

Q8. Ablation on removing the NN block.

For simple systems, removing the NN block does not affect the results. For complex systems, e.g., NS, it causes the long rollout prediction diverged. Please see our reply to Q3 in General Reply.

Q9. Ablation to replace trainable symmetric filter with FD kernel in the PDE block.

We conducted an ablation study on the Burgers dataset, e.g., replacing our trainable conv filter with FD kernel in the PDE block (NN block excluded in the model). However, this led to NaN values in the test results due to large error accumulation. This indicates that FD kernels fails to approximate the derivatives on coarse grids.

Q10. Ablation on the Poisson block.

We conducted an ablation study on the NS dataset to evaluate the significance of the Poisson block, as shown in Table R6 in the 1-page PDF rebuttal file. We can see that the model's performance deteriorates without the Poisson block.

Q11. Generalization to all ICs; discretization error accumulation.

Since PDE is embedded into the network with RK4 integrator, the model can in theory generalize prediction of solution trajectories over different ICs. However, we agree with you that large error accumulation over time might occur given a complex discretized IC on coarse grids. We will clarify this in our revised paper.

Q12. Difference between the training and testing ICs?

We use random seeds to generate different ICs. In the NS example, different ICs are created by generating random noise for each component of the velocity field and filtering it to create a divergence-free field with desired properties. We randomly selected 10 ICs and drew histograms to show the differences, shown in Fig. R3 in the 1-page PDF rebuttal file.

Q13. Correlation plots over time steps.

The correlation curve over time steps for the NS dataset is shown in Fig. R2 in the 1-page PDF rebuttal file. To provide a more comprehensive comparison, we also included DNS results for various grid resolutions in the figure.

Q14. Computational cost.

Please see our reply to Q4 in General Reply.

Remark: Thanks for your constructive comments. Your consideration of improving the rating of our paper will be much appreciated!

Dear Reviewer xwvp,

We hope our point-by-point response in the rebuttal has addressed your concerns. We are very much looking forward to your feedback during the discussion period. We would be more than happy to answer any further questions you may have.

Best regards, The Authors

Dear Reviewer xwvp:

As the author-reviewer discussion period will end soon, we would appreciate it if you could kindly review our responses at your earliest convenience. If there are any further questions or comments, we will do our best to address them before the discussion period ends.

Thank you very much for your time and efforts!

Sincerely,

The Authors