Stochastic Parameterization of Moist Physics Using Probabilistic Diffusion Model

1. Introduction

Convection plays a crucial role in atmospheric circulation, transferring heat from the earth surface to the upper atmosphere, creating vertical air movement to mix air at different altitudes, driving cloud formation and extreme precipitation, shaping the global mass, momentum, and energy budget. Meanwhile, current numerical models fail to explicitly resolve convection processes, due to resolution limits posed by computational constraints. This makes it instrumental to include parameterization schemes that empirically account for the effects of convection, based on the resolved variables of numerical model.

Traditional convection parameterization schemes, coming with error-prone empirical function forms and free parameters, may introduce significant errors in estimating heating, moistening, and precipitation rate at grid scale [1,2], which translates to errors in simulated large-scale circulation patterns, such as intertropical convergence zone, Madden-Julian oscillation, and mesoscale convective systems [3,4,5,6,7].

A growing interest is to apply high-fidelity data, including high-resolution simulations and observations, to calibrate existing convection schemes, or to develop data-driven simulators using deep neural network. The latter provides more flexibility and higher accuracy, promising a breaker to the deadlock of convection parameterization problem [8,9]. Up till now, deep neural networks have been tested thoroughly in numerical models under aquaplanet or realistic geography setup, showing its potential to be the next-generation convection scheme [10,11,12,13,14,15,16,17,18,19,20].

Yet, the lack of stochasticity by neural-network-based convection schemes will harm the performance of numerical model. Gentine et al. [9] found heating and moistening tendencies predicted by deep neural network lack variability below 700 hPa. Rasp et al. [16] showed that the standard deviation of convective heating tendency below 700 hPa in multi-year simulations is significantly lower after the same neural network is integrated in numerical model. Moreover, traditional stochastic parameterization schemes can improve the ensemble prediction skill and numerical simulation of Madden-Julian oscillation, climate mean, and variability [21,22,23]. It is expected that stochastic parameterization of moist physics by neural network may further improve its performance in numerical models.

Generative model is a kind of deep learning model based on probability theory. It establishes a mapping between a known prior distribution and the target distribution. It is naturally suited for stochastic parameterization. Classic generative models include variational auto encoder (VAE) [24] and generative adversarial network (GAN) [25]. GAN was explored in stochastic parameterization of convection, subgrid stress of ocean model, and stochastic tendencies in Lorenz-96 model [26,27,28,29,30,31,32]. Yet, VAE is hindered by its blurring generated samples and posterior collapse [33,34]. GAN suffers from difficulties in training and mode collapse [35,36]. Efforts have been made to stabilize the training and secure the diversity and quality of generated samples [33,35,36,37,38]. But it still requires strong case-specific experiences to train a good generative model based on GAN and VAE.

Recent years a new family of generative model, probabilistic diffusion model (PDM) [39], has received special attentions. PDM is the foundation model of the well-known text-to-image models like Stable Diffusion [40] and Dall-E [41]. PDM splits the generative task into a series of relatively simple denoising tasks. This paradigm shift makes PDM much easier to train and suffer from much less mode collapse [39,42,43]. PDM shows its strong dominance due to its high quality of generative samples [44].

In this study, a stochastic parameterization scheme for moist physics based on PDM (DIFF-MP) is developed. A cloud-resolving global simulation is coarse grained into resolved variables and subgrid contributions to form the training data of DIFF-MP. The fatal flaw of PDM is its slow inference speed. Adapted from Chen et al. [45], we let DIFF-MP train on a series of noised levels in a stochastic way to generalize to large denoising steps for acceleration. Cassifier-free guidance [46] is used for tuning the influence of conditional information by fusing the denoised latents of conditioned and unconditioned models at inference stage for further improvements of DIFF-MP. DIFF-MP is further compared with GAN and VAE on testing data for stochastic parameterization of moist physics. Finally, the physical interpretability of DIFF-MP is explored.

The remainder of this study is organized as follows. Section 2 presents the details of cloud-resolving data and scale-separation techniques for training data generation. Section 3 presents the training details of the DIFF-MP, and how to accelerate DIFF-MP and use classifier-free guidance to improve its performance. Section 4 presents the testing data performance of DIFF-MP and its comparisons with baselines, GAN and VAE. Section 4 also discusses the interpretability of DIFF-MP. Finally, section 5 concludes this work and presents future perspectives.

2. Training Data Preprocessing

The training data is diagnosed from high-fidelity data from cloud-resolving simulation based on Global-Regional Integrated forecast SysTem (GRIST) [47,48]. GRIST is formulated on primitive equations. The horizontal mesh adopts structured Delaunay-Voronoi grid [49]. This kind of grid mesh is appropriate for the global simulation due to its isotropic properties. High-fidelity data from high-resolution simulations are widely used for training data of machine-learning based parameterization schemes because most of atmospheric motions including convections are explicitly resolved.

The simulations consist of four periods: 1988.10.1-1988.10.20, 1998.1.1-1998.1.20, 2005.4.1-2005.4.20, and 2013.7.10-2013.7.29, including four seasons with ENSO and MJO events of various intensities. La Nina and El Nino are strong in 1988 and 1998 respectively. MJO are strong in 1988, 1998, and 2005. Those four periods are chosen for diversity of training data. The horizontal resolution of the simulation is 5 km, which is high enough to resolve deep convections. There are 30 levels under 20 km, with extra levels in boundary layer. The initial condition is interpolated from ECMWF Reanalysis v5 (ERA5) [50]. The boundary condition (sea surface temperature) is updated every 24 hours. GRIST adopts Yonsei University scheme (YSU) [51] for boundary layer parameterization, Noah-MP land surface model for surface-atmosphere flux parameterization, WRF single-moment 6-class scheme (WSM6) [52] for microphysics parameterization, and RRTMG shortwave and longwave schemes [53] for radiation parameterization. Model outputs are saved every modeling hour.

The high-resolution data need to be preprocessed into subgrid contributions and large-scale resolved variables to form the training data. Microphysics and subgrid vertical transports are regarded as the subgrid processes, which are the output of DIFF-MP. They include subgrid vertical transports of heat and water vapor, and four outputs from WSM6 scheme: temperature tendency (${Tend}_{T-mp}$), water vapor tendency (${Tend}_{q_v-mp}$), cloud water ($q_c$) and cloud ice mixing ratios ($q_i$). Rain, snow, graupel mixing ratios, and subgrid vertical transports of $q_c$ and $q_i$ are not considered because they are negligible. Six variables are selected as the conditional input of DIFF-MP: temperature (T), water vapor mixing ratio ($q_v$), surface pressure ($P_s$), sensible heat flux ($SHF$), latent heat flux ($LHF$), and short-wave radiation at surface (SOLIN). The procedure to process high-resolution data is shown as follows.

A random group of points on high-resolution Delaunay-Voronoi grid that are seamlessly connected are chosen for coarse graining and subgrid diagnosing. Those points are labeled as $P_1$, $P_2$, $P_3$, …, $P_n$. Those $n$ points must try their best to form a regular hexagon or pentagon. The coarse-grained variable is labeled as $\overline{a}$ for variable $a$ in high-resolution data, then,

$$\overline{a}={\displaystyle \frac{1}{n}}\sum _{P_i\in \left\{P_1, P_2, P_3, \dots , P_n\right\}}{a}_{P_i}$$

(1)

Through coarse graining (equation 1), we can get averaged variables for DIFF-MP’s conditional input and output except for subgrid contributions of $T$ and $q_v$. If the difference between $a$ and $\overline{a}$ is $a\text{'}$, then the subgrid vertical flux of $a$ is,

$$\overline{a’w\text{'}}={\displaystyle \frac{1}{n}}\sum _{P_i\in \left\{P_1, P_2, P_3, \dots , P_n\right\}}{a\text{'}}_{P_i}\bullet {w\text{'}}_{P_i}.$$

(2)

The tendency due to subgrid vertical flux of $a$ is,

$${Tend}_{a-flux}=-{\displaystyle \frac{\partial \overline{a’w\text{'}}}{\partial z}}.$$

(3)

According to the definition, the subgrid contributions for $T$ and $q_v$ are,

$$Ten{d}_{q_v-sgs}=\overline{Ten{d}_{q_v-mp}}+{Tend}_{q_v-flux}=\overline{Ten{d}_{q_v-mp}}-{\displaystyle \frac{\partial \overline{q_v’w^{\text{'}}}}{\partial z}},$$

(4)

$$Ten{d}_{T-sgs}=\overline{Ten{d}_{T-mp}}+{Tend}_{T-flux}=\overline{Ten{d}_{T-mp}}-{\displaystyle \frac{\partial \overline{T’w^{\text{'}}}}{\partial z}}.$$

(5)

The reason why the subgrid contributions are formulated as presented above is explained in supplementary material. The conditional input and output variables for DIFF-MP are shown in Table 1.

Four resolutions are considered in this study: 120 km, 60 km, 30 km, and 15 km, corresponding to 576, 144, 36, and 9 points in equations 1 and 2. The number of training data for all the four resolutions are the same. The conditional input of DIFF-MP has four variables at surface layer. Those variables are each duplicated vertically for 30 times to align with vertical profiles of $T$ and $q_v$ to form the input matrix. $T$ and $P_s$ are normalized by subtracting 0.05-quantile and dividing by the difference between 0.95-quantile ($a_{0.95}$) and 0.05-quantile ($a_{0.05}$),

$$a_{norm}={\displaystyle \frac{a-a_{0.05}}{a_{0.95}-a_{0.05}}}.$$

(6)

The other conditional input and output variables are normalized by dividing by $a_{0.95}$,

$$a_{norm}={\displaystyle \frac{a}{a_{0.95}}}.$$

(7)

After normalization, the samples that contain output variable larger than 3.0 are all replaced by the neighboring samples within 3.0 to exclude the abnormal high values. The four simulations have 20 days each. The first day is discarded due to the model spin up. The following 13 days are chosen as training data, 4 days as validation data and final 2 days as testing data. Data from the four periods are randomly mixed, and there are 51,118,080, 15,728,640 and 7,864,320 samples for training, validation and testing data. To accelerate the validation and testing of DIFF-MP, only random 10,000 and 1,600,000 samples of validation and testing data are used.

3. DIFF-MP, Inference Acceleration and Classifier-Free Guidance

3.1. DIFF-MP and Its Training

Figure 1 shows how high-resolution data is preprocessed and how DIFF-MP stochastically parameterize moist physics. DIFF-MP can be divided into forward and reverse diffusion processes. In forward diffusion process, target data was fused by gaussian noise step by step until complete noise (blue arrows in Figure 1). DIFF-MP is trained to reverse forward process stepwise and generate samples in reverse process (red arrows in Figure 1). A brief introduction to PDM’s mathematic derivations, training and sampling algorithms are presented in supplementary material.

DIFF-MP is a hierarchical one-dimensional PDM, whose structure is depicted in Figure A1. The number of filters and layers of DIFF-MP are 128 and 6 based on the hyper-parameter experiments. DIFF-MP is inspired by Wavegrad [45] for network structure, GAN-TTS [54] for Ublock structure. The feature-wise linear modulation module (FiLM) [55] is used to combine information from noise level $\sqrt{{\overline{\alpha }}_t}$, denoised ${\mathit{x}}_t$, and conditional input $\mathit{c}$.

Adam algorithm and cyclical learning rates ranging from 1e-4 to 1e-3 are used for DIFF-MP training [56,57]. Loss function is mean squared error. Batch size is 8,000. DIFF-MP is trained for five epochs. Model weights are saved at the end of last epoch. DIFF-MP is trained in Python package Keras 3.0 [58] on Nvidia 4090 GPU. Different DIFF-MPs are trained on different resolutions for thorough validation of its performance.

3.2. Inference Acceleration of DIFF-MP

The training and sampling algorithms of DIFF-MP with inference acceleration in detail is presented in Figure 2. Typically, PDM is only trained on a fixed time-step schedule. When fewer time steps are applied, PDM must follow a quite different denoising route that it is not trained on, leading to degraded performances. To improve the DIFF-MP’s adaptability to fewer denoising steps, it is conditioned on the noise level $\sqrt{{\overline{\alpha }}_t}$ directly. This is also adopted by Song and Ermon [59,60] in their score matching framework. Moreover, ${\overline{\alpha }}_t$ at time step $t$ is sampled from a uniform distribution $U({\overline{\alpha }}_{t-1}^{\mathrm{\text{'}}},{\overline{\alpha }}_t^{\mathrm{\text{'}}})$ during training. This distribution is defined from a predefined fixed noise schedule $\left\{{\overline{\alpha }}_t^{\mathrm{\text{'}}}\right\}$, in which ${\alpha }_t^{\mathrm{\text{'}}}$ linearly decreases from 9.9999 to 9.94 with 100 steps. During testing, denoising schedule $\left\{{\overline{\alpha }}_t\right\}$ is interpolated from noise level function $f\left(t\right)$ (Figure S1), which is constructed from the predefined fixed noise schedule $\left\{{\overline{\alpha }}_t^{\mathrm{\text{'}}}\right\}$. Then sampling of DIFF-MP is the same as ordinary PDM. This algorithm is adapted from Chen et al. [45] with modifications in inference stage.

To show the adaptability of DIFF-MP with inference acceleration to different denoising steps, it is compared with DIFF-MPs trained with other fixed steps only. The other DIFF-MPs are trained on fixed noise schedules: ${\alpha }_t$ linearly decreases from 9.9999 to 9.94, with total time steps of $n$. $n$ is set to be 100, 50, 20, 10, 5 and 2. All the models are trained for five epochs until converged. Model weights are saved at the end of last epoch.

Validation criteria are the mean squared error, Pearson correlation coefficient, coverage ratio, spread-skill correlation, standard deviation, kurtosis, and skewness. Because DIFF-MP is a stochastic parameterization scheme, the validation must be conducted on an ensemble of output. For a single conditional input, 32 different outputs are generated. Mean squared error and correlation coefficient measure the error of stochastic parameterization. Coverage ratio is the ratio of validation data that are in the value range of DIFF-MP output ensemble. A good output ensemble should cover the validation data, so DIFF-MP should have high coverage ratio. Spread-skill correlation is the correlation coefficient between the error of DIFF-MP and the spread of DIFF-MP. DIFF-MP should have large spread when it has large prediction error, so DIFF-MP should have high spread-skill correlation. Standard deviation, kurtosis, and skewness measure the high-order statistics of samples generated by DIFF-MP. Those statistics should be close to the corresponding statistics of validation data. How those criteria are calculated are presented in supplementary material.

Figure 3 and Figure 4 show the validation data performance of different DIFF-MPs. Those models are trained on data of resolution 120 km. Results on the other resolutions are similar. For DIFF-MPs trained on fixed steps, their performances degrade quickly as number of steps decreases, except for coverage ratio (Figure 3). Coverage ratio increases because the sample spread is unreasonably large when denoising steps are too few (Figure 3). But for DIFF-MP trained with inference acceleration, correlation coefficient, mean squared error, and spread-skill correlation merely change with different steps (Figure 4). Standard deviation, skewness and kurtosis only deviate from validation data limitedly (Figure 4). DIFF-MP trained with inference acceleration is significantly better than DIFF-MPs trained on fixed steps. Training on a series of noise levels will effectively alleviate the overfitting of DIFF-MP to limited fixed steps and help its performance generalize to different steps. Consider the balance between the acceleration and sample quality, DIFF-MP in the rest of the study uses five steps to denoise.

3.3. Classifier-Free Guidance of DIFF-MP

Standard deviation, skewness, and kurtosis of samples generated by DIFF-MP deviate from validation data on $Ten{d}_{T-sgs}$ and $Ten{d}_{q_v-sgs}$ when DIFF-MP uses five steps to denoise (Figures 4a, b). Standard deviation is also smaller for $q_c$ and $q_i$ (Figure 4c, d). Cassifier-free guidance [46] can alleviate the deviation of the statistics (standard deviation, skewness, and kurtosis). This method is originally proposed to improve the diversity of generated samples from conditional PDM. When conditional information of PDM is too strong (weak), the samples generated tend to be deterministic (diverse) and the statistics of samples will deviate from original data. Classifier-free guidance fuses the denoised latent from the conditional PDM with the denoised latent from another unconditional PDM. In this way, the statistics of the samples generated by conditional PDM will be drawn closer to the original data.

During training stage, DIFF-MP’s condition input $\mathit{c}$ will be replaced by denoised latent ${\mathit{x}}_t$ with a probability of 0.1, through which an unconditional DIFF-MP is trained in the meantime. The final output will be the combination of the denoised latents from the conditional and unconditional DIFF-MPs by mixing ratio $\omega$. For each denoising steps,

$$\widetilde{P_{\theta }}\left({\mathit{x}}_{t-1}|{\mathit{x}}_t,\mathit{c}\right)=\left(1+\omega \right)P_{\theta }\left({\mathit{x}}_{t-1}|{\mathit{x}}_t,\mathit{c}\right)-\omega P_{\theta }\left({\mathit{x}}_{t-1}|{\mathit{x}}_t,{\mathit{x}}_t\right),$$

(8)

where $P_{\theta }$ is original DIFF-MP, $\widetilde{P_{\theta }}$ is the combined DIFF-MP. The training and sampling algorithms of DIFF-MP using classifier-free guidance and inference acceleration are presented in Figure 5.

Figure 6 shows the influence of $\omega$ on DIFF-MP performance on validation data at the resolution of 120 km. As $\omega$ increases, kurtosis and skewness decrease evidently, but standard deviation increases. Classifier free guidance successfully pushes the statistics of samples generated by DIFF-MP to the validation data. In the meantime, the other four criteria do not change much. The best $\omega$ for $Ten{d}_{T-sgs}$ and $Ten{d}_{q_v-sgs}$ is 0.5. But for $q_c$ and $q_i$, no mixing is better. It is also found that different $\omega$ for different output variables will not interfere with each other during the inference stage. The same validation of $\omega$ is also conducted on the other resolutions. The best $\omega$ for all the four resolutions are in Table 2, which will be followed by DIFF-MP in this study.

4. Results

4.1. Baseline Models and Their Trainings

Conditional VAE (CVAE-MP) and conditional GAN (CGAN-MP) are selected as baselines. VAE maps each sample from training data into a known distribution within the latent space to form a random coding [24], which enables the random generation of the training data. GAN is composed of two competing networks, in which generator can produce data that cannot be distinguished from discriminator [25]. Both models are widely used in generative learning and appropriate for stochastic parameterization of moist physics. Variables in Table 1 are also used for conditional inputs of CVAE-MP and CGAN-MP. To make a fair comparison between different generative models, DIFF-MP, CVAE-MP, and CGAN-MP have the same model size. Model structures of CGAN-MP and CVAE-MP are presented in Figures S2 and S3.

Training settings of CVAE-MP and CGAN-MP are the same as DIFF-MP. CGAN-MP output is the sum of a pretrained neural network and the generator. The pretrained neural network predicts the output profiles deterministically, and its weight is frozen during training of CGAN-MP. Wasserstein GAN technique is also applied to stabilize CGAN-MP training [35]. CVAE-MP and CGAN-MP are trained for 2 and 22 epochs until performance converged. The best models on validation data are saved for further comparison with DIFF-MP on testing data. Validation criterion is the average of correlation coefficient, spread-skill correlation, and coverage ratio. Different models for CVAE-MP and CGAN-MP are also trained for different resolutions.

4.2. Performance Comparison between Models

The testing data performance of CGAN-MP, CVAE-MP, and DIFF-MP at different resolutions are presented in Figure 7. For mean squared error, DIFF-MP is significantly better than the others on $q_c$ and $q_i$, and nearly the best on $Ten{d}_{T-sgs}$ and $Ten{d}_{q_v-sgs}$. DIFF-MP is consistently the best on correlation coefficient and spread-skill correlation. It also beats the other two models on coverage ratio except for $Ten{d}_{T-sgs}$. In terms of the statistics of the samples generated, DIFF-MP is closer to testing data in general than the other two models. It also achieves a much more consistent performance among different resolutions. It is concluded that DIFF-MP performs robustly better than CGAN-MP and CVAE-MP on testing data.

It is noticed that performances of mean squared error and correlation coefficient degrade as resolution increases. This is because as grid spacing increases, coarse graining involves more averaging over turbulences and cloud processes, which makes the subgrid processes more predictable. The same degradation is also reported in other studies [19,61].

The global distribution of mean squared error by DIFF-MP, and the mean-squared-error differences between models at 120 km resolution are depicted in Figure 8. DIFF-MP’s mean squared error is mostly distributed in midlatitude where extratropical cyclones are active. Besides, mean squared error of $Ten{d}_{T-sgs}$ and $Ten{d}_{q_v-sgs}$ is also distributed along the large-scale terrain and tropical zone where shallow convections are active (Figures 8a, d). Except for $Ten{d}_{q_v-sgs}$, DIFF-MP has smaller mean squared error than the other two models globally. The red color at background indicates the lower systematic error by DIFF-MP. Figure 9 shows the global distribution of whether testing data is in the range of model ensemble at 120 km resolution. DIFF-MP covers almost 90% of testing data for all the variables. The uncovered locations spread along the regions with large mean squared error in Figure 8. Almost all the testing data of $Ten{d}_{T-sgs}$ are not covered by CVAE-MP (Figure 9b). So are $q_c$ and $q_i$ from CGAN-MP (Figures 9i, l). Figure 8 and Figure 9 furtherly confirm the robustness of DIFF-MP’s better performance over the other models.

Figure 10 shows the per-level distribution of model generated samples at 120 km resolution. Testing data distributions in Figure 10 all stop abruptly at -3.0 and 3.0 because abnormal extreme data are excluded. We can see two distinct signals of shallow and deep convections in Figure 10a. $Ten{d}_{T-sgs}$ is mostly positive on upper levels of boundary layer (about level 7), but negative on lower levels. This is due to the condensation of water vapor in shallow convections, where the latent heat is released near the boundary layer top but collected on the lower levels. In Figure 10e, $Ten{d}_{q_v-sgs}$ is positive near the boundary layer top but negative below, indicating the transport of water vapor by shallow convections. Shallow convections are also shown in lower-level extremes of $q_c$ and $q_i$ (Figures 10i, m). Above boundary layer, there is another peak of heating for $Ten{d}_{T-sgs}$ on levels 10-15 (Figure 10a). This is where the deep convections form. Deep convections are also seen in upper levels of $q_i$ (Figure 10m).

Signals of shallow and deep convections are generally captured by all three generative models. But the predicted distributions of extreme data are different. DIFF-MP almost perfectly reproduces the data distribution of testing data, even the artificial abrupt at ±3.0. Only the distributions of $Ten{d}_{T-sgs}$ above level 15 are underestimated (Figure 10b). CVAE-MP is overly conservative on extreme values. Only the probability density of large value is captured (Figures 10d, h, l, p). While CGAN-MP is excessively radical, leading to overly large area of probability distributions (Figures 10c, g, k, o). CGAN-MP also predicts unreasonable negative values for $q_c$ and $q_i$. Figure A2 shows the results from 30 km resolution, which are consistent with Figure 10.

Figure 11 shows the ensemble output profiles from different models and corresponding output profiles from testing data at 120 km resolution. All the models’ ensemble can capture the overall vertical variability of profiles from testing data. But for $Ten{d}_{T-sgs}$ and $Ten{d}_{q_v-sgs}$, CGAN-MP’s and CVAE-MP’s ensembles concentrate about zero value and deviate from testing data on levels 10-15 (Figures 11a, b, d, e). On the contrary, DIFF-MP’s ensemble shares the similar vertical variability as testing data (Figures 11c, f). As for $q_c$ and $q_i$, CGAN-MP’s ensemble extends to 2.0-2.5 (Figures 11g, h). This is also confirmed by Figure 10 where CGAN-MP predicts excessive extreme values. Value ranges of CVAE-MP’s and DIFF-MP’s ensembles match testing data better than CGAN-MP. Results on 30 km resolution are similar with Figure 11, which are presented in Figure A3.

DIFF-MP can recover the data distribution of subgrid moistening, heating, and cloud processes better than CGAN-MP and CVAE-MP. It is also obvious that compared to the other two models, DIFF-MP can produce reasonable ensemble that can cover the output profiles from testing data. DIFF-MP is better than the other two models on the stochastic parameterization of the moist physics.

4.3. Interpretability of DIFF-MP

It is of vital importance to test the interpretability of DIFF-MP to secure the physical robustness. Figure 12 shows the influence on DIFF-MP’s outputs due to the change of stratification in boundary layer. The original boundary layer has unstable stratification. After $T$ profile is neutralized, $Ten{d}_{T-sgs}$ and $Ten{d}_{q_v-sgs}$ are almost zero (Figures 12c, d). $q_i$ on higher levels also decreases, indicating the termination of shallow and deep convections (Figure 12f). $q_c$ and $q_i$ increase in boundary layer (Figures 12e, f). This is because $q_v$ accumulates in boundary layer due to suppressed convections, which cause excessive condensation. While after $q_v$ is neutralized, $Ten{d}_{q_v-sgs}$ becomes positive and $q_c$ disappear in boundary layer (Figures 12d, e). This is because $q_v$ is not saturated after the neutralization. Therefore, shallow clouds evaporate and $q_v$ in boundary layer is restored by DIFF-MP. Figure 13 shows the influence of surface flux. As surface flux decreases, $Ten{d}_{T-sgs}$, $Ten{d}_{q_v-sgs}$, and $q_i$ on higher levels decrease significantly (Figures 13a, b, d). This is caused by the suppressed deep convections due to surface flux reduction. $q_c$ in boundary layer increases first and then decreases as surface flux is further reduced. Similar to Figure 12e, $q_c$ is first accumulated in boundary layer when convections are constrained. But $q_c$ is then reduced when convections are severely constrained. Figure 12 and Figure 13 confirm that DIFF-MP’s response to the input variation is physically reasonable, which secure the future implementation of DIFF-MP into GRIST model.

5. Conclusions and Discussions

This study develops a stochastic moist physics parameterization scheme named DIFF-MP based on one-dimensional PDM. DIFF-MP is trained on a series of noise levels which boosts its generalizability to large denoising step to accelerate. DIFF-MP achieves 20 times acceleration without significant degradation of performance. Classifier-free guidance is furtherly adopted to eliminate the deviations of the statistic of DIFF-MP-generated samples from the validation data.

DIFF-MP’s ability to stochastically parameterize the subgrid contributions of moist physics is compared against CVAE-MP and CGAN-MP on testing data. DIFF-MP demonstrates a consistent improvement compared to the other two models in terms of prediction error, spread-skill correlation, coverage ratio, and the statistics of the reproduced subgrid contributions, including standard deviation, kurtosis, and skewness. DIFF-MP’s performance is also consistent among the four different resolutions. The improvement of DIFF-MP on prediction error can be up to 40% at most compared to the other models.

DIFF-MP’s prediction error and testing data that are not included in the predicted ensemble mainly distribute along the large-scale terrain and midlatitudes where extratropical cyclones are active. DIFF-MP’s prediction error is consistently smaller than the other two models globally. In terms of coverage ratio, nearly 90% of testing data are included in the predicted ensemble of DIFF-MP. DIFF-MP’s predicted ensemble profiles are also more reasonable in vertical variability and value range than the other two models.

DIFF-MP can reproduce the per-level distributions of different variables almost perfectly. CGAN-MP tend to predict too much extreme data while CVAE-MP too little. When unstable stratification in boundary layer is neutralized or surface flux is reduced, deep convections will be significantly suppressed while low clouds accumulate in boundary layer due to constrained shallow convections. The interpretability experiment shows that DIFF-MP’s prediction is physically reasonable.

Only the moist physics are parameterized in this study for proof of concept of PDM in stochastic parameterization. DIFF-MP can include more physical processes like boundary layer turbulences, longwave radiation, and shortwave radiation to achieve a unified parameterization of all the physical processes. Future work considers the implementation of DIFF-MP into GRIST, and studies the influence of DIFF-MP on numerical simulation of Madden-Julian oscillation, intertropical convergence zone, climate mean, and variability.

The implementation of Python-based machine learning models into Fortran-based earth system model is still difficult. Early works tend to hard code the Python-based models into Fortran through self-developed tools, which are troublesome and time consuming [11,16,17]. It is strongly appealed that the development team of earth system models should develop official tools to simplify the implementation procedure. One possible solution is splitting the earth system model into different modules and wrap them with Python interfaces that allow easy implementation and the heterogeneous computing based on GPU and CPU [62]. The machine learning models will be run on GPU while the numerical integration on CPU to enable the utmost efficiency. It is also possible that with the help of large language model, the Fortran-based earth system model will be translated into Python and enable the running of the whole system on GPUs only [63,64]. This will terminate the implementation problem and benefit the community for further development of machine learning parameterization schemes.