Enhancing Atmospheric Turbulence Phase Screen Generation with an Improved Diffusion Model and U-Net Noise Generation Network

Abstract

Simulating atmospheric turbulence phase screens is essential for optical system research and turbulence compensation. Traditional methods, such as multi-harmonic power spectrum inversion and Zernike polynomial fitting, often suffer from sampling errors and limited diversity. To overcome these challenges, this paper proposes an improved denoising diffusion probabilistic model (DDPM) for generating high-fidelity atmospheric turbulence phase screens. The model effectively captures the statistical distribution of turbulence phase screens using small training datasets. A refined loss function incorporating the structure function enhances accuracy. Additionally, a self-attention module strengthens the model’s ability to learn phase screen features. The experimental results demonstrate that the proposed approach significantly reduces the Fréchet Inception Distance (FID) from 154.45 to 59.80, with the mean loss stabilizing around 0.1 after 50,000 iterations. The generated phase screens exhibit high precision and diversity, providing an efficient and adaptable solution for atmospheric turbulence simulation.

1. Introduction

Atmospheric turbulence phase screens are essential for simulating optical distortions in turbulent media. Turbulence-induced refractive index fluctuations cause wavefront distortion, beam spreading [1], and imaging degradation [2], which are mathematically represented by phase screens through spatial phase shift distributions. The Fourier Transform (FT) method [3] is widely used to generate phase screens due to its computational efficiency. However, it insufficiently captures low spatial frequencies, thereby reducing accuracy. To overcome this limitation, Lane et al. [4] introduced subharmonic corrections to enhance low-frequency sampling. This method was later validated by Frehlich [5]. Zernike polynomials [6] also model phase distortions, but their reliance on low-order expansions often neglects higher-order turbulence effects [7,8].

Recent advancements in deep learning have enabled the application of neural networks to atmospheric turbulence research. Wang and Basu [9] applied artificial neural networks (ANNs) to optical turbulence estimation and marine refractive index structure constant simulation, demonstrating improved performance over traditional similarity theory-based methods. Later, Lohani et al. [10] introduced an optical feedback network to correct mode distortions caused by turbulence, resulting in near-ideal receiver profiles. More recently, Zeyang Wang et al. [11] used deep convolutional networks for phase screen generation, achieving an RMSE reduction of 8.76 while enhancing computational efficiency. However, these data-driven methods often require large training datasets, which can be a significant limitation.

In 2020, Ho et al. [12] introduced Denoising Diffusion Probabilistic Models (DDPMs), a novel unsupervised framework for image generation through iterative denoising. This method gradually adds noise to an image until it becomes Gaussian noise and then learns to reverse this process, thereby enabling phase screen generation by learning the underlying data distribution without relying on paired training samples.

This study proposes an enhanced DDPM architecture that integrates self-attention mechanisms to better capture fine-scale phase features. The modified model effectively synthesizes phase screens from limited training data by improving their ability to learn phase screen distributions. Additionally, a refined loss function is introduced, leading to improved accuracy and diversity in the generated phase screens compared to the baseline DDPM.

The paper is organized as follows: Section 2 discusses phase screen generation using the Fourier Transform method, Section 3 presents the experimental framework and diffusion model architecture, Section 4 analyzes experimental results, and Section 5 concludes with key findings.

2. Fourier Transform-Based Atmospheric Phase Screen Generation

The power spectrum inversion method is computationally efficient and widely adopted for simulating atmospheric turbulence phase screens. It applies the atmospheric phase power spectral density function to filter a zero-mean, unit-variance Hermitian complex Gaussian random matrix, generating the turbulence phase screen through an inverse Fourier transform [13]:

$$\phi (x,y)={\displaystyle {\int }_{f_y}{\displaystyle {\int }_{f_x}h}}(f_x,f_y)\sqrt{F_{\phi }(f_x,f_y)}\times \exp [\mathrm{j}2\mathrm{\pi }(f_{x}x+f_{y}y)]d{f}_{x}d{f}_y,$$

(1)

In this equation, $\phi (x,y)$ represents the generated atmospheric turbulence phase screen in the spatial domain. The variables $f_x$ and $f_y$ denote the spatial frequencies in the $x$- and $y$-directions, respectively, while $x$ and $y$ are the corresponding spatial coordinates. The term $h\left(f_x,f_y\right)$ is a Hermitian complex Gaussian random function with zero mean and unit variance. $F_{\phi }(f_x,f_y)$ denotes the phase power spectral density function, characterizing the statistical properties of atmospheric turbulence in the frequency domain. Under the assumptions of isotropy and local homogeneity, the Kolmogorov spectrum is employed as the atmospheric refractive index power spectral density function $F_n(\kappa )$, given by:

$$F_n(\kappa )=0.033C_n^2{\kappa }^{-11/3}(1/L_0\le \kappa \le 1/l_0),$$

(2)

where $\kappa =2\pi (f_x\widehat{i}+f_y\widehat{j})$ is the spatial angular frequency, $C_n^2$ is the refractive index structure constant, $L_0$ is the outer scale, and $l_0$ is the inner scale. Based on this, the power spectral density function $F_{\phi }(\kappa )$ of the phase can be calculated as follows:

$$F_{\phi }(\kappa )=2{\mathsf{\pi }}^2{k}^2\mathrm{\Delta }\mathrm{z}{F}_n(\kappa ),$$

(3)

where $\mathrm{\Delta }\mathrm{z}$ is the distance along the propagation path, and $k$ is the wavenumber. Substituting the Kolmogorov refractive index power spectral density (Equation (2)) gives [14]:

$$F_{\phi }(\kappa )=0.023r_0^{-5/3}{f}^{-11/3},$$

(4)

$f$ represents the spatial frequency, with units of cycles per meter. $r_0$ is the atmospheric coherence length, which characterizes the spatial scale over which atmospheric turbulence remains coherent, with units of meter. It is commonly expressed as:

$$r_0={\left[0.423{({\displaystyle \frac{2\pi }{\lambda }})}^2{\displaystyle {\int }_0^{\infty }{C}_n^2}(h)dh\right]}^{-3/5}$$

(5)

where $\lambda$ is the wavelength of light and $C_n^2$ is the refractive index structure constant.

Since computers can only process discrete signals, for a phase screen of size $L\times L$, the sampling interval is $\mathrm{\Delta }L\times \mathrm{\Delta }L$, and the number of samples is $N\times N$. The discrete spatial frequencies in the frequency domain yield the minimum and maximum spatial frequencies as:

$$\mathrm{\Delta }{f}_x=\mathrm{\Delta }{f}_y=f_{\min }=\mathrm{\Delta }f={\displaystyle \frac{1}{L}},$$

(6)

$$f_{\max }={\displaystyle \frac{N\times \mathrm{\Delta }f}{2}}={\displaystyle \frac{N}{2L}}={\displaystyle \frac{1}{2\mathrm{\Delta }L}},$$

(7)

Rewriting Equation (1) gives:

$$\phi (m,n)={\displaystyle \sum _{m^{\prime }=-\frac{N_x}{2}+1}^{\frac{N_x}{2}}{\displaystyle \sum _{n^{\prime }=-\frac{N_y}{2}+1}^{\frac{N_y}{2}}h(m^{\prime },n^{\prime })\sqrt{F_{\phi }(m^{\prime },n^{\prime })}}}\exp [\mathrm{j}2\mathrm{\pi }({\displaystyle \frac{m^{\prime }m}{N_x}}+{\displaystyle \frac{n^{\prime }n}{N_y}})]\mathrm{\Delta }{f}_x\mathrm{\Delta }{f}_y,$$

(8)

where $h(m^{\prime },n^{\prime })={\displaystyle \frac{g(m^{\prime },n^{\prime })}{\sqrt{\mathrm{\Delta }{f}_x\mathrm{\Delta }{f}_y}}},$$g(m^{\prime },n^{\prime })=g_{real}(m^{\prime },n^{\prime })+i{g}_{imag}(m^{\prime },n^{\prime })$, $g_{imag}(m^{\prime },n^{\prime })$, and $g_{real}(m^{\prime },n^{\prime })$ are Gaussian random matrices with a mean of 0 and a variance of 1, and they satisfy the condition $g(m^{\prime },n^{\prime })=g^{*}(-m^{\prime },-n^{\prime })$. Figure 1a shows the power spectral density corresponding to the Kolmogorov atmospheric turbulence model, while Figure 1b presents a phase screen generated using the power spectrum inversion method. The simulation is conducted with a Fried parameter of $r_0=0.1\mathrm{m}$, a phase screen size of $L=2\mathrm{m}$, and a resolution of $N=512$. The outer and inner scale parameters of the turbulence are set to $L_0=100\mathrm{m}$ and $l_0=0.01\mathrm{m}$, respectively. Although the traditional spectrum inversion method is fast and convenient, the power in the low-frequency range of the Kolmogorov atmospheric power spectral density is not captured. Specifically, the power in the range of $(-\mathrm{\Delta }{f}_x/2,\mathrm{\Delta }{f}_x/2)$ to $(-\mathrm{\Delta }{f}_y/2,\mathrm{\Delta }{f}_y/2)$ is missing, which leads to significant errors in the phase screen when the distance $\mathrm{\Delta }r$ is large.

The “subharmonic method” can compensate for this shortcoming. Its principle involves using a combination of spatial frequencies with non-uniform sampling to introduce ultra-low spatial frequency components. Referring to the subharmonic method used by Lane [4] and incorporating the approach of Schmidt J.D. [15], the formula for generating the low-frequency phase screen can be expressed as follows:

$${\phi }_L(m,n)={\displaystyle \sum _{p=1}^{N_p}{\displaystyle \sum _{m=-1}^1{\displaystyle \sum }_{n=-1}^1}}[\exp ({\displaystyle \frac{2\mathsf{\pi }\mathrm{j}{3}^{-p}m{m}^{\prime }}{N}}+{\displaystyle \frac{2\mathsf{\pi }\mathrm{j}{3}^{-p}n{n}^{\prime }}{N}})\times f(m^{\prime },n^{\prime })\times R(m^{\prime },n^{\prime })],$$

(9)

In this equation, ${\phi }_L(m,n)$ denotes the low-frequency phase screen generated using the subharmonic method. The summations are performed over the subharmonic order $p$ and the frequency indices $m’$ and $n’$, with each ranging from −1 to 1. The parameter $N$ represents the size of the phase screen, and $3^{-p}$ determines the sampling interval for the $p$-th subharmonic layer. The exponential term introduces subharmonic spatial frequencies into the synthesis process. The function $f\left(m’,n’\right)$ defines the power spectrum weighting factor for each subharmonic frequency component and is given by $f(m^{\prime },n^{\prime })=C\times 3^{-2p}\times r_0^{-5/6}(f_{lx}^2+f_{ly}^2)$, where $C$ is a constant, and $r_0$ is the atmospheric coherence length (Fried parameter). The terms $f_{lx}=3^{-p}{m}^{\prime }\mathrm{\Delta }{f}_x$ and $f_{ly}=3^{-p}{n}^{\prime }\mathrm{\Delta }{f}_y$ represent the subharmonic spatial frequencies in the $x$- and $y$-directions, respectively. $R\left(m’,n’\right)$ is a complex Gaussian random variable with zero mean and unit variance. By selecting the subharmonic order as 3 and combining Equations (8) and (9), the phase screen with low-frequency compensation can be obtained, as shown in Figure 2:

At the same time, it has been theoretically proven that the wave structure function of Kolmogorov turbulence plane waves can be expressed as [16]:

$$D(|\mathrm{\Delta }r\left|\right)=6.88({\displaystyle \frac{\left|\mathrm{\Delta }r\right|}{r_0}}{)}^{5/3},$$

(10)

where $\mathrm{\Delta }r$ is the relative distance, with units of meter.

By generating 100 phase screens using the methods mentioned above and averaging them to obtain the phase screen structure function, the structure functions of the phase screens generated by the two methods are compared, as shown in Figure 3:

The yellow curve represents the theoretical structure function, the blue curve corresponds to the structure function obtained using the standard Fourier transform method, and the green curve represents the phase screen structure function obtained using the subharmonic compensation method. Although some deviation from the theoretical values remains, the low-frequency components are more accurately captured by the subharmonic approach.

3. Experimental Framework and Diffusion Model Architecture

3.1. Network Architecture

The DDPM is a generative model based on a two-stage process: a forward diffusion process and a reverse denoising process. Given the highly random and diverse nature of atmospheric turbulence phase screens, this study employs DDPM to model their distribution characteristics and generate phase screens.

During training, a sample $x_0$ is drawn from the dataset, and a random time step $t$ is selected. The forward diffusion process then adds random noise $ϵ$ to $x_0$, producing $x_{\mathrm{t}}$. The model takes $x_{\mathrm{t}}$ and $t$ as inputs to predict the noise ${ϵ}^{\prime }$. The loss between $ϵ$ and ${ϵ}^{\prime }$ is computed, and the network weights are updated using gradient descent. This iterative process trains the model to progressively refine noise prediction. Figure 4 illustrates the DDPM training workflow.

The neural network in this study adopts a U-Net architecture enhanced with self-attention mechanisms. As illustrated in Figure 5, it consists of three main modules: the encoder, bottleneck, and decoder. Key components include the DoubleConv, Down, self-attention, and Up layers. The DoubleConv module applies two consecutive convolution layers, each followed by normalization and a GELU activation function, enabling the extraction of complex features; the Down module performs down sampling while increasing the number of channels. It consists of a max-pooling layer, two DoubleConv layers, and a time embedding layer, progressively reducing the spatial resolution and enhancing feature representations; the Up module is responsible for up sampling the feature map, gradually restoring spatial resolution while reducing the number of channels. This architecture ensures effective feature extraction and enhances the model’s ability to generate accurate phase screens.

The self-attention layer enhances the original U-Net architecture, introduced by Google’s machine translation team in 2017 [17]. It effectively captures spatial relationships in 2D images and consists of a normalization layer, multi-head attention module, and feedforward neural network, all integrated with residual connections to enhance model expressiveness (Figure 6).

The input consists of a 1 × 512 × 512 phase screen and a time step t. After an initial double convolution, the feature map expands to 64 × 512 × 512. Downsampling is then applied, reducing the spatial dimensions by half while doubling the number of channels, resulting in 128 × 256 × 256. A self-attention module follows without changing the feature map size. This process is repeated twice, further downsampling the feature map to 512 × 64 × 64. In the bottleneck layer, the feature map is adjusted to 256 × 64 × 64 before up sampling restores the resolution step by step, with additional self-attention modules incorporated to refine spatial details. The final convolution produces a 1 × 512 × 512 phase screen. The detailed network parameters are listed in

Table 1. Input and output parameters of each module.

Module	Output Size
Input	1 × 512 × 512
DoubleConv	64 × 512 × 512
Down + Self-Attention	128 × 256 × 256
Down + Self-Attention	256 × 128 × 128
Down + Self-Attention	512 × 64 × 64
DoubleConv	512 × 64 × 64
DoubleConv	512 × 64 × 64
DoubleConv	256 × 64 × 64
Up + Self-Attention	128 × 128 × 128
Up + Self-Attention	64 × 256 × 256
Up + Self-Attention	64 × 512 × 512
Conv2d	1 × 512 × 512

.

3.2. Loss Function Improvement

To improve the accuracy of the atmospheric phase screens generated by DDPM, this paper modifies the DDPM loss function. The loss function consists of two parts. The first part calculates the error between the predicted noise and the real noise using the mean squared error (MSE), and its expression is as follows:

$${\mathcal{L}}_{\mathrm{simple}}=\mathbb{E}\left[{‖ϵ-{ϵ}_{\theta }(x_t,t)‖}^2\right],$$

(11)

where $ϵ$ is the real noise, ${ϵ}_{\theta }$ is the noise predicted by the model, $t$ is the time step, and $x_t$ is the noisy image at step $t$. The second part calculates the mean squared error between the structure function of the image with real noise and the structure function of the image with model-predicted noise:

$${\mathcal{L}}_{\mathrm{s}tructure}=k_t\mathbb{E}\left[{‖D_0(x_{t_0})-D_t(x_t)‖}^2\right]$$

(12)

where $D_0$ represents the structure function of the image with real noise, and $D_t$ represents the structure function of the image with model-predicted noise. $k_t$ is a hyperparameter that is adjusted based on the time step $t$. As the time step $t$ decreases, the generated image becomes closer to the target image $x_0$. In this case, the value of $k_t$ needs to be larger to ensure that the structure function of the image with model-predicted noise is more accurate. Combining the two parts of the error, the total loss function is expressed as follows:

$${\mathcal{L}}_{\mathrm{total}}={\mathcal{L}}_{simple}+\lambda {\mathcal{L}}_{\mathrm{structure}}$$

(13)

$\lambda$ is a hyperparameter used to balance the magnitudes of ${\mathcal{L}}_{simple}$ and ${\mathcal{L}}_{\mathrm{structure}}$ to achieve the best training performance.

For a random field $g\left(r\right)$, the structure function is defined as follows:

$$D_g(\mathrm{\Delta }\mathit{r})=\langle {[g(\mathit{r}+\mathrm{\Delta }\mathit{r})-g(\mathit{r})]}^2\rangle$$

(14)

where $\mathrm{\Delta }r$ represents the displacement vector, and $r$ is the position vector. ⟨⋅⟩ denotes a spatial average over all positions $r$, under the assumption of ergodicity. The ensemble average can be approximated by a spatial average; as the statistical properties of the system are equivalent for both over a sufficiently large spatial domain [18], we use an ensemble average to present a spatial average.

To efficiently calculate the structure function of a random field, Sadhukhan et al. [19] proposed using parallel computing to accelerate the calculation of the structure function. In this paper, we leverage GPU-based parallel computing to parallelize the direct structure function calculation algorithm. Experiments were conducted on a computer with an Intel Core i7-13790F CPU (Intel Corporation, Santa Clara, CA, USA) and an NVIDIA GeForce RTX 4060 GPU (NVIDIA Corporation, Santa Clara, CA, USA). The input image size was 512 × 512. Directly calculating the structure function for one image took approximately 32 s. After applying the parallel computing method, the speed improved to around 0.3 s, an improvement of approximately 100 times, meeting the speed requirements for integration as a loss function during training.

4. Results

This study uses a dataset of 5000 atmospheric turbulence phase screens, each with a physical size of $2\mathrm{m}\times 2\mathrm{m}$ and a resolution of $512\times 512$ sampling points. The atmospheric coherence length $r_0$ is randomly sampled in the range of 0.03 m to 0.15 m. The outer and inner scale parameters are set to $L_0=100\mathrm{m}$ and $l_0=0.01\mathrm{m}$, respectively. The number of subharmonic compensation layers is set to $p=3$. The model was trained with a batch size of 10, using AdamW as the optimizer. The learning rate was 3 × 10⁻⁴, and training was conducted for 50 k iterations on an NVIDIA A100 GPU, completed in 2 h.

Figure 7a shows the loss trend over iterations, and Figure 7b presents its smoothed version. The loss gradually decreases, stabilizing below 0.1 after 50 k iterations, indicating model convergence.

Figure 8 presents atmospheric turbulence phase screens generated by DDPM’s reverse process at different stages, with eight images per batch. At Epoch = 2, the outputs appear disorganized, resembling random noise. By Epoch = 50, some images still exhibit noise-like patterns, but certain phase screen structures begin to emerge. By Epoch = 80, the generated phase screens closely match those in the original dataset, displaying finer details and richer features.

Figure 9 presents a quantitative evaluation of the enhanced DDPM’s ability to generate atmospheric phase screens. A comparison of the average structure functions from 100 model-generated phase screens and those from the original dataset shows close agreement. While minor deviations exist in individual realizations, the overall trend and spatial scaling characteristics remain consistent, with an RMS error of less than 5%.

The Fréchet Inception Distance (FID) [20] is more suitable than the Inception Score (IS) [21] for evaluating phase screens. The IS measures image quality based on classifiability, while FID compares feature distributions to assess both quality and diversity, making it more effective for validating generated phase screens. This paper conducts an ablation study on three DDPM variants: the Baseline DDPM, which uses a standard U-Net; an Enhanced DDPM, which incorporates a self-attention module into the architecture while keeping the loss function unchanged; and Ours (Enhanced DDPM + Loss), which further improves the loss function. Each experiment generates 100 images per epoch, with the results summarized in the

Table 2. FID values of improved DDPM at different Iterations.

Models	Iterations = 0	Iterations = 10 k	Iterations = 20 k	Iterations = 30 k	Iterations = 40 k	Iterations = 50 k
DDPM	145.26	86.14	82.46	77.62	73.53	74.16
Enhanced DDPM	125.17	76.95	80.78	78.24	67.31	65.86
Ours	154.45	71.19	63.61	68.15	60.45	59.80

below.

Table 2 shows a significant reduction in FID during training, indicating the effective learning of the phase screen distribution. A comparison among the Baseline DDPM, Enhanced DDPM, and Ours (Enhanced DDPM + Loss) reveals a 15% improvement in the latter, demonstrating higher generative accuracy. FID values stabilize between 40,000 and 50,000 iterations, with differences of less than 0.8%, suggesting statistical convergence. This is likely due to (a) the saturation of loss function minimization with oscillatory gradients at higher iterations and (b) the inherent sensitivity limitations of FID when evaluating high-dimensional turbulence distributions, particularly when the test dataset is limited [21].

5. Conclusions

This paper proposes an improved DDPM diffusion model for generating atmospheric turbulence phase screens. To address the challenge of limited phase screen data, the model enables training on small datasets while effectively capturing phase screen distribution characteristics, reducing the FID from 154.45 to 59.80. The loss function was refined by incorporating the structure function, and parallel computing was employed to accelerate its calculation, achieving a 100 × speedup for 512 × 512 images. Additionally, a self-attention module was integrated into the network, enhancing its ability to learn phase screen distributions. After 50,000 training iterations, the mean loss stabilizes around 0.1.

Despite these improvements, some limitations remain. The inherent sequential nature of DDPM leads to relatively slow generation speed, rendering it unsuitable for real-time applications. While the generated phase screens exhibit improved diversity, variability in output quality is still observed.

Overall, the proposed model demonstrates flexibility and can be further optimized through parameter tuning or architectural enhancements, offering a promising approach for atmospheric turbulence phase screen generation.