Lightning wave generation method and system based on diffusion model

Abstract

The invention relates to the technical field of lightning data processing, in particular to a lightning wave generation method and system based on a diffusion model. A diffusion model is adopted to generate lightning wave data, a denoising neural network adopted in the reverse diffusion process is obtained by improving a UNET model structure, and the improvement mode is that a wavelet attention unit is added after each down-sampling and before each up-sampling. According to the method, the denoising neural network is improved, attention technology wavelet transformation is combined, multi-scale transformation is captured at the same time, the detail capturing capacity of the model is improved, particularly, high-frequency components are captured, the model is helped to better reserve key features of lightning waves, and therefore the quality of the diffusion model applied to a lightning wave generation task is improved. The defect that a diffusion model is directly applied to lightning waves in the prior art is overcome.

Description

Technical Field

[0001] The present invention relates to the technical field of lightning data processing, and in particular to a lightning wave generation method and system based on a diffusion model. Background Art

[0002] Narrow bipolar events (NBEs) are short and intense electromagnetic wave phenomena caused by lightning discharges. Due to this characteristic, their detection in the very low frequency band is very difficult and the number of samples is small, which brings inconvenience to the comprehensive study of lightning. Therefore, it is very necessary to use data generation methods to learn the potential structure and distribution of data for sample enhancement.

[0003] The current mainstream data enhancement models mainly include variational autoencoders (VAEs), generative adversarial networks (GANs), and diffusion models. The generation results of VAE models are usually blurred and distorted. GAN models need to train both the generator and the discriminator simultaneously, and it is difficult to find a balance point. Moreover, the generation quality depends too much on the quality of the training set, resulting in poor training effects. Diffusion models adopt a different way to generate samples. By gradually diffusing noise during the generation process, the generated samples are clearer, more real, more stable, and can maintain good diversity.

[0004] The existing diffusion models for generating lightning waveforms mainly have the following difficulties: (1) Lightning waveforms change instantaneously, and it is not easy to capture the details of instantaneous jumps; (2) High-resolution samples need to be generated. When the quantity is large, the computational cost of aggregating all the information of the pictures is high; (3) Paying attention to the whole image leads to the introduction of redundant information, interfering with the model's focus on the key waveform regions. Summary of the Invention

[0005] In order to overcome the defects of directly applying the diffusion model to lightning waves in the above-mentioned prior art, the present invention proposes a lightning wave generation method based on a diffusion model, which improves the ability to capture complex waveform features by improving the denoising neural network, thereby greatly improving the quality of lightning wave generation.

[0006] A lightning wave generation method based on a diffusion model proposed by the present invention uses a diffusion model to generate lightning wave data. The denoising neural network used in the reverse diffusion process is improved on the basis of the UNET model structure. The improvement method is: adding wavelet attention units after each downsampling and before each upsampling; The wavelet attention unit includes an attention network, a feature decomposition function, a low-frequency branch, a high-frequency horizontal branch, a high-frequency vertical branch, a high-frequency diagonal branch, and a feature synthesis function; the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch are connected in parallel between the feature decomposition function and the feature synthesis function; After the input of the wavelet attention unit is processed by the attention network, it is decomposed into a low-frequency subband, a high-frequency subband in the horizontal direction, a high-frequency subband in the vertical direction, and a high-frequency subband in the diagonal direction by the feature decomposition function. After the four are processed by the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch respectively, they are synthesized by the feature synthesis function. The synthesized data is normalized and used as the output of the wavelet attention unit.

[0007] Preferably, the denoising neural network includes: an input convolutional layer, a plurality of encoding modules, an intermediate module, a plurality of decoding modules, and an output convolutional layer connected in sequence; the number of encoding modules and decoding modules is equal; The encoding module includes a downsampling unit, a wavelet attention unit, and a convolutional unit connected in sequence, and the output of the downsampling unit is also connected to the input of the convolutional unit; The decoding module includes a wavelet attention unit, a convolutional unit, and an upsampling unit connected in sequence; the input of the wavelet attention unit is also connected to the input of the convolutional unit; Let the encoding modules and decoding modules be sorted in sequence along the data transmission direction, and the output of the downsampling unit of the i-th encoding module is also connected to the input of the upsampling unit of the (N - i + 1)-th decoding module.

[0008] Preferably, the data of the encoding modules and decoding modules in the denoising neural network is 3.

[0009] Preferably, during the training process of the denoising neural network, the loss function L(total) is used to backpropagate the gradient to update the parameters, and its loss function is: L(total)=L(base)+λ(L(low)+L(high)); Among them, L(base) represents the expectation of the square of the two-norm of the difference between the random noise ∈ calculated on the (x0, ∈, t) distribution and the predicted noise ∈ of the denoising neural network θ (x t , t), L(low) is the expectation of the two-norm of the difference between the input and output of the low-frequency branch calculated on the (x0, t) distribution, and L(high) is the expectation of the two-norm of the difference between the input and output of the high-frequency branch calculated on the (x0, t) distribution. The high-frequency branch includes the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch; x0 is the lightning wave data sample; ∈ is the random noise used for the forward diffusion of the sample x0, and ∈ follows a normal distribution; t is the number of steps of forward diffusion of the sample x0, and θ is the parameter of the denoising neural network.

[0010] Preferably, the training steps of the denoising neural network are as follows: First, extract the training sample x0 from the learning dataset and add noise through the diffusion model to obtain the noise sample xt; input the noise sample x t into the denoising neural network to obtain the predicted noise ∈θ (x t , t); Then calculate the loss function L(total) of the model, and update the parameters to be learned of the denoising neural network by backpropagating the gradient through L(total); Repeat the above steps until the number of updates of the denoising network reaches the set value.

[0011] Preferably, the construction method of the learning dataset is as follows: form multiple window data from the original lightning wave data through a sliding window, calculate the mean and standard deviation of each window data, and extract the window data that satisfies that the maximum value is greater than or equal to the sum of the multiples of the mean and the standard deviation, and the minimum value is less than or equal to the difference between the mean and the multiples of the standard deviation and add it to the learning dataset.

[0012] Preferably, the processing methods of the low-frequency branch, high-frequency horizontal branch, high-frequency vertical branch, and high-frequency diagonal branch for the input sub-band are all: perform convolution processing on the input sub-band and then activate it, and perform convolution processing on the activation value and then perform weighted calculation with the set weight to obtain the output data; The set weight of the low-frequency branch is the low-frequency weight, and the set weights of the high-frequency horizontal branch, high-frequency vertical branch, and high-frequency diagonal branch are all high-frequency weights.

[0013] Preferably, grouped convolution is used for convolution processing on the high-frequency horizontal branch, high-frequency vertical branch, and high-frequency diagonal branch.

[0014] Preferably, the feature synthesis function uses the IWT function, and the normalization method of the synthesized data is: take half of the synthesized data and add 0.5.

[0015] A system for implementing the lightning wave generation method based on the diffusion model described above, including a memory and a processor, wherein a computer program is stored in the memory, the processor is connected to the memory, and the processor is used to execute the computer program to implement the lightning wave generation method based on the diffusion model.

[0016] The advantages of the present invention are as follows: (1) Considering that the lightning wave data contains rich local features and mutation points, such as peaks and steep edges, etc., the present invention introduces a wavelet attention unit into the denoising neural network, combines the attention technology wavelet transform, and simultaneously captures multi-scale transforms, improves the detail capture ability of the model, especially captures high-frequency components, helps the model better retain the key features of the lightning wave, and thus improves the quality of the diffusion model applied to the lightning wave generation task.

[0017] (2) The wavelet attention unit decomposes the data into four sub-bands: low-frequency, horizontal, vertical, and diagonal high-frequency; each sub-band has an independent processing branch and is finally synthesized. In this way, the model is allowed to process high-frequency information in different directions separately, avoiding the loss of details caused by mixed processing. In particular, the high-frequency branch uses grouped convolution, which can reduce the computational amount while maintaining the efficiency of feature extraction. In addition, different weights are set for the low-frequency and high-frequency branches, which facilitates the model to distinguish between the main contour and the detail part, and is beneficial to further improving the generation quality.

[0018] (3) In the present invention, the wavelet attention unit normalizes the synthesized data, which is beneficial to stabilizing the numerical range and preventing gradient problems.

[0019] (4) In the denoising neural network provided by the present invention, wavelet attention units are added to both the encoding module and the decoding module, and are added after each downsampling and before each upsampling. Thus, wavelet attention is applied multiple times during the entire feature extraction process, which can enhance the feature representation at different scales and further improve the effect of feature fusion.

[0020] (5) In the loss function of model training of the present invention, loss terms for the low-frequency and high-frequency sub-bands are added. The low-frequency part can ensure the accuracy of the overall structure, and the high-frequency part strengthens the detail reconstruction, which is beneficial to improving the fidelity of the generated data, especially for data sensitive to details such as lightning waves.

[0021] (6) During the construction of the data set of the present invention, window data with high volatility is screened, which is beneficial to capturing extreme situations, making the trained model more robust when generating complex waveforms, more suitable for the characteristics of lightning wave data, and avoiding the generated lightning wave data from being too smooth. Description of the Drawings

[0022] Figure 1 is the structural diagram of the denoising neural network proposed by the present invention; Figure 2 is the structural diagram of the wavelet attention unit proposed by the present invention; Figure 3 is the flow chart of the data set construction method proposed by the present invention; Figure 4 is the flow chart of the denoising neural network training method proposed by the present invention; Figure 5 is the flow chart of the lightning wave generation method based on the diffusion model proposed by the present invention; Figure 6 is the display of the effect of Experiment 1; Figure 7 is the display of the effect of Experiment 2. Detailed Embodiments

[0023] Next, the technical solutions in the embodiments of the present invention will be clearly and completely described in conjunction with the accompanying drawings in the embodiments of the present invention. Obviously, the described embodiments are only a part of the embodiments of the present invention, rather than all the embodiments. All other embodiments obtained by those of ordinary skill in the art based on the embodiments of the present invention without creative efforts shall fall within the protection scope of the present invention.

[0024] In this embodiment, a denoising neural network is constructed to predict the noise contained in the input lightning wave data. Refer to Figure 1 , the denoising neural network includes: an input convolutional layer, a plurality of encoding modules, an intermediate module, a plurality of decoding modules, and an output convolutional layer connected in sequence.

[0025] The encoding module includes a downsampling unit, a wavelet attention unit, and a convolutional unit connected in sequence; the input end of the downsampling unit is the input end of the encoding module, the output of the downsampling unit is also connected to the input of the convolutional unit, and the output of the convolutional unit is used as the output of the encoding module. In the encoding module, the input data is processed by the downsampling unit and then by the wavelet attention unit. The data processed by the wavelet attention unit and the data processed by the downsampling unit are dimensionally concatenated and enter the convolutional unit for convolutional processing to obtain the output of the encoding module.

[0026] The decoding module includes a wavelet attention unit, a convolutional unit, and an upsampling unit connected in sequence; the input end of the wavelet attention unit is the input end of the decoding module, and the output of the upsampling unit is used as the output of the decoding module; the input of the convolutional unit is also connected to the input of the wavelet attention unit.

[0027] The data of the encoding module and the decoding module are both N, and are incrementally numbered according to the data transmission order. The output of the downsampling unit of the i-th encoding module is also connected to the input of the upsampling unit of the (N - i + 1)-th decoding module.

[0028] In this way, the input data of the decoding module is processed by the wavelet attention unit and then dimensionally concatenated with the input data and enters the convolutional unit for processing. The data processed by the convolutional unit is dimensionally concatenated with the data output by the downsampling unit in the corresponding decoding module and enters the upsampling unit. The upsampling unit upsamples the input data to obtain the output data of the decoding module.

[0029] Specifically in implementation, the number of both the encoding module and the decoding module is set to 3. Then, the output of the downsampling unit in the first encoding module is also connected to the input of the upsampling unit of the third decoding module, the output of the downsampling unit in the second encoding module is also connected to the input of the upsampling unit of the second decoding module, and the output of the downsampling unit in the third encoding module is also connected to the input of the upsampling unit of the first decoding module.

[0030] The intermediate module may specifically be composed of two serially connected convolutional units.

[0031] Referring to Figure 2 , the wavelet attention unit in this solution includes an attention network, a feature decomposition function, a low-frequency branch, a high-frequency horizontal branch, a high-frequency vertical branch, a high-frequency diagonal branch, and a feature synthesis function; The input of the attention network is the input of the wavelet attention unit. The output of the attention network is connected to the input of the feature decomposition function. The low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch are connected in parallel between the feature decomposition function and the feature synthesis function; The feature decomposition function uses the DWT function to decompose the input data into a low-frequency subband LL, a high-frequency subband LH in the horizontal direction, a high-frequency subband HL in the vertical direction, and a high-frequency subband HH in the diagonal direction; After LL, LH, HL, and HH are processed by the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch respectively, they are synthesized by the feature synthesis function IWT. The synthesized data XI is normalized and used as the output Xinorm of the wavelet attention unit; The formula is expressed as: Xinorm = [XI + 1] / 2; XI = IWT[LL′, HL′, LH′, HH′]; Among them, LL′, HL′, LH′, and HH′ respectively represent the outputs of the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch.

[0032] The synthesized data output by the IWT function is a value in the numerical interval (-1, 1). According to the above normalization formula, Xinorm can be limited to the interval (0, 1). In this way, the extreme value of the feature peak is reduced, which facilitates subsequent data processing.

[0033] The low-frequency branch consists of two serially connected convolutional layers. LL is activated after the first convolutional processing, and then after the second convolutional processing, it is weighted and calculated with the set low-frequency weight α L to obtain the output LL' of the low-frequency branch; The high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch all consist of two serially connected grouped convolutional layers. LH is activated after the first grouped convolutional processing, and then after the second grouped convolutional processing, it is weighted and calculated with the set high-frequency weight α H to obtain the output LH' of the high-frequency horizontal branch; HL is activated after the first grouped convolutional processing, and then after the second grouped convolutional processing, it is weighted and calculated with the high-frequency weight α H to obtain the output HL' of the high-frequency vertical branch; HH is activated after the first grouped convolutional processing, and then after the second grouped convolutional processing, it is weighted and calculated with the high-frequency weight α HThrough weighted calculation, the output HH' of the high-frequency diagonal branch is obtained; The formula is expressed as: LL' = α L ×Conv(Silu(Conv(LL))); [LH', HL', HH'] T = α H ×GroupConv(Silu(GroupConv[LH, HL, HH] T )); For the convenience of subsequent processing, the output of the wavelet attention unit needs to be normalized, that is, the synthesized data output by the feature synthesis function is normalized to obtain the output Xinorm of the wavelet attention unit.

[0034] Temporal encoding embedding is performed during the downsampling process of the encoding module and during the upsampling process of the decoding module.

[0035] Let the input data of the i-th encoding module be denoted as Xin_i, and the output data be denoted as Xout_i; the output of the downsampling unit in the i-th encoding module is denoted as Xd_i; The processing process of the i-th encoding module for the input data Xin_i is as follows: Xin_i is sampled by the downsampling unit to obtain Xd_i, and the processing result of Xd_i by the wavelet attention unit is concatenated with Xd_i in dimension and then processed by the convolutional unit to obtain Xout_i; the formula is expressed as: Xout_i = Conv[Concat(Xinorm _i, Xd_i)]; Xinorm _i represents the processing result of the wavelet attention unit on Xd_i in the i-th encoding module, and Concat represents concatenation in dimension.

[0036] The input data of the i-th decoding module is denoted as X'in_i, and the output data is denoted as X'out_i; The data obtained by processing the input data X'in_i of the i-th decoding module by the wavelet attention unit is concatenated with X'in_i and then enters the convolutional unit for processing. The output data of the convolutional unit is concatenated with Xd_(N - i + 1) in dimension and then input into the upsampling unit for sampling to generate X'out_i; Xd_(N - i + 1) is the output of the downsampling unit in the (N - i + 1)-th encoding module; The processing process formula of the i-th decoding module is expressed as: X'out_i = UP{Concat(Conv(Concat(X'in_i, X'inorm _i)), Xd_(N - i + 1))}; X'inorm _i represents the processing result of the X'in_i by the wavelet attention unit in the i-th decoding module, and UP represents upsampling.

[0037] Referring to Figure 3 , the construction process of the learning dataset of the denoising neural network is as follows: SA1. Obtain the original lightning wave data S and form multiple window data S(g) by sliding the window, where 1 ≤ g ≤ G, G is the number of window data divided from the original lightning wave data S, and g is the ordinal number; S(g) = {s((g - 1)×z), s((g - 1)×z + 1), s((g - 1)×z + 2), ……, s((g - 1)×z + K - 1)}; Among them, K is the window length and z is the step size. In specific implementation, z = 1 can be set.

[0038] SA2. Calculate the mean μ and standard deviation σ of each window data, and denote the mean and standard deviation of the window data S(g) as μ(g) and σ(g) respectively; SA3. Extract the window data that satisfies the condition that the maximum value is greater than or equal to the sum of the mean and the multiple of the standard deviation, and the minimum value is less than or equal to the difference between the mean and the multiple of the standard deviation from all window data and add it to the learning dataset.

[0039] In this step, by performing the above steps SA1 - SA3 on different original lightning wave data S, sufficient learning samples can be obtained to form a learning dataset. The original lightning wave data S can be directly read from the lightning detection station.

[0040] In specific implementation, step SA3 can be divided into the following sub - steps: SA31. Judge whether the window data S(g) satisfies: MAX(S(g)) ≥ μ(g) + 2σ(g) and MIN(S(g)) ≤ μ(g) - 2σ(g); the initial value of g is 1; If yes, add the window data S(g) to the learning dataset, and then perform step SA32; If no, perform step SA32; SA32. Judge whether g is less than G; If yes, update g to g + 1, and then return to step SA31; If no, update the original lightning wave data S, and then return to step SA1.

[0041] Referring to Figure 4 , the training method of the denoising neural network includes the following steps: St1. Initialize the parameters to be learned of the denoising neural network, and the parameters to be learned include the node weights of the model; St2. Extract the training sample \(x_0\) from the learning dataset and add noise through the diffusion model to obtain the noisy sample \(x_t\); ; \(\alpha\) t = 1 - \(\beta\) t ; ; where \(\alpha\) t and \(\beta\) t are both hyperparameters, \(\beta\) t is a random value in the interval (0, 1), \(\epsilon\) is a random noise conforming to the normal distribution, i.e., \(\epsilon\sim N(0, 1)\), and \(t\) is the set number of diffusion steps; St3. Input the noisy sample \(x\) t into the denoising neural network, and the denoising neural network outputs the predicted noise \(\epsilon\) θ \((x\) t , \(t)\), and \(\theta\) represents the parameters of the denoising neural network; St4. Calculate the loss function \(L(total)\) of the model, and update the parameters to be learned of the denoising neural network by backpropagating the gradient through \(L(total)\); \(L(total)=L(base)+\lambda(L(low)+L(high))\); where \(L(base)\) represents the expectation of the square of the two-norm of the difference between the random noise \(\epsilon\) and the predicted noise \(\epsilon\) θ \((x\) t , \(t)\) calculated on the \((x_0, \epsilon, t)\) distribution, \(L(low)\) is the expectation of the two-norm of the difference between the input and output of the low-frequency branch calculated on the \((x_0, t)\) distribution, and \(L(high)\) is the expectation of the two-norm of the difference between the input and output of the high-frequency branch calculated on the \((x_0, t)\) distribution. The high-frequency branch includes the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch; That is: ; \(L(low)=\) [||LL'-LL||2]; \(L(high)=\) [||LH'-LH||2]+ [||HL'-HL||2]+ [||HH'-HH||2]; St5. Repeat the above steps St1 - St4 until the number of updates of the denoising network reaches the set value.

[0042] The lightning wave generation method based on the diffusion model proposed in this embodiment is used to learn the existing lightning wave data samples to obtain more lightning wave data samples of the same type, such as Figure 5As shown, the generation method includes the following steps: S1. Obtain the to-be-learned lightning wave data sample X0, and add noise to X0 through a diffusion model to obtain a noise sample XT; ; α T = 1 - β T ; ; where α T 、 and β T are both hyperparameters, β T is a random value in the interval (0, 1), ∈ is a random noise conforming to a normal distribution, i.e., ∈ ~ N(0, 1), and T is the set number of diffusion steps.

[0043] S2. Predict the sample noise through a denoising neural network, and iteratively decrease the diffusion step t from T to 0 to achieve reverse diffusion denoising of the noise sample XT, and obtain the generated lightning wave data sample X'0, and X'0 is the new lightning wave data learned from X0.

[0044] Specifically, step S2 is as follows: S21. Input the noise sample Xt into the denoising neural network, and the denoising neural network outputs the predicted noise ∈ θ (Xt, t); adopt the reverse diffusion formula to obtain the noise sample X(t - 1) after denoising the noise sample Xt; the initial value of t is T; The reverse diffusion formula is: ; ; ; S22. Judge whether t is equal to 1; if so, obtain X'0 = X(t - 1); if not, let t be updated to t - 1, and then return to step S21.

[0045] In this way, by adopting this application, after determining the denoising neural network, the above steps S1 - S2 can be combined to expand the limited lightning wave data samples to obtain more lightning wave data samples to meet more experimental requirements.

[0046] The following combines specific embodiments to verify the above lightning wave generation method based on a diffusion model.

[0047] In this embodiment, three comparative generation methods are provided.

[0048] Comparative generation method one: Abbreviated as VAE, directly use a variational autoencoder (VAE) to encode and decode the input lightning wave data to generate new lightning wave data; Comparison generation method two: Abbreviated as GAN, based on the generative adversarial network (GAN), the input lightning wave data is reconstructed to obtain new lightning wave data; Comparison generation method three: Abbreviated as DDPM, a diffusion model is used to generate new lightning wave data, and the Unet network is used for noise prediction in the reverse diffusion process; The method of the present invention, abbreviated as WAB, uses a diffusion model to generate new lightning wave data, and the denoising neural network with a wavelet attention unit added proposed by the present invention is used for noise prediction in the reverse diffusion process.

[0049] Comparison generation method three is the baseline model of the method of the present invention, and the generation steps of both refer to the above steps S1 - S2.

[0050] In this embodiment, first, the learning dataset is constructed by using the dataset construction method proposed by the present invention, specifically referring to the above steps SA1 - SA3.

[0051] In this embodiment, the learning dataset is divided into a training set and a test set. The models including VAE, GAN network, Unet network, and the denoising neural network proposed by the present invention are trained using the training set. Then, the trained models are used to generate lightning wave data samples equal in quantity to the test set to form a generated dataset, and then the spatial distribution differences between the generated dataset and the test set are compared.

[0052] Specifically, in comparison generation method one and comparison generation method two, the trained VAE network and GAN network are directly used to generate lightning wave data to construct the corresponding generated datasets; Comparison generation method three and the method of the present invention use the above steps S1 - S2 to generate lightning wave data to construct the corresponding generated datasets.

[0053] In this embodiment, the data generation methods are compared in two experiments.

[0054] In the two experiments, learning datasets are constructed based on lightning wave data from different historical time periods, and then divided into a training set with a sample size of 900 and a test set with a sample size of 300.

[0055] In the two experiments, the parameters of the training set and the test set are set to 64 * 64, that is, the dataset samples are adjusted to image data with a length and width of 64 pixels; the batch size batch_Size for model training is set to 16, the initial learning rate is 0.002, and the number of training epochs is 1000 rounds.

[0056] In the two experiments, 300 pictures are respectively generated using the trained models and compared with the test set. In Experiment 1, the spatial distribution differences between the generated dataset and the test set are measured by the FID (Frechet Inception Distance) metric, and the experimental results are as Figure 6As shown; in Experiment 2, the spatial distribution difference between the generated dataset and the test set was measured by the SSIM (structural similarity) index, and the experimental results are as Figure 7 shown.

[0057] It can be seen that in both experiments, the lightning wave data generated by the method of the present invention has the smallest difference from the test set, and the index gap is obvious, which proves the superiority of the present invention.

[0058] Of course, for those skilled in the art, the present invention is not limited to the details of the above exemplary embodiments, but also includes the same or similar structures that can be implemented in other specific forms without departing from the spirit or basic characteristics of the present invention. Therefore, from any point of view, the embodiments should be regarded as exemplary and non-limiting. The scope of the present invention is defined by the appended claims rather than the above description. Therefore, all changes falling within the meaning and scope of the equivalent elements of the claims are intended to be embraced within the present invention. Any reference signs in the claims should not be construed as limiting the claims involved.

[0059] In addition, it should be understood that although this specification is described according to embodiments, not every embodiment only contains an independent technical solution. This narrative way of the specification is only for clarity. Those skilled in the art should regard the specification as a whole, and the technical solutions in each embodiment can also be appropriately combined to form other embodiments that can be understood by those skilled in the art.

[0060] The technologies, shapes, and structures not described in detail in the present invention are all well-known technologies.

Claims

1. A method for generating lightning waves based on a diffusion model, characterized in that, The diffusion model is adopted to generate lightning wave data. The denoising neural network used in the reverse diffusion process is improved based on the UNET model structure. The improvement method is as follows: wavelet attention units are added after each downsampling and before each upsampling; The wavelet attention unit includes an attention network, a feature decomposition function, a low-frequency branch, a high-frequency horizontal branch, a high-frequency vertical branch, a high-frequency diagonal branch, and a feature synthesis function; the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch are connected in parallel between the feature decomposition function and the feature synthesis function; The input of the wavelet attention unit is processed by the attention network and then decomposed by the feature decomposition function into a low-frequency subband, a high-frequency subband in the horizontal direction, a high-frequency subband in the vertical direction, and a high-frequency subband in the diagonal direction. After being processed by the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch respectively, they are synthesized by the feature synthesis function. The synthesized data is normalized and used as the output of the wavelet attention unit.

2. The lightning wave generation method based on a diffusion model according to claim 1, wherein The denoising neural network includes: an input convolutional layer, multiple encoding modules, an intermediate module, multiple decoding modules, and an output convolutional layer connected in sequence; the number of encoding modules is equal to the number of decoding modules; The encoding module includes a downsampling unit, a wavelet attention unit, and a convolutional unit connected in sequence. The output of the downsampling unit is also connected to the input of the convolutional unit; The decoding module includes a wavelet attention unit, a convolutional unit, and an upsampling unit connected in sequence; the input of the wavelet attention unit is also connected to the input of the convolutional unit; Let the encoding modules and the decoding modules be sorted in sequence along the data transmission direction. The output of the downsampling unit of the i-th encoding module is also connected to the input of the upsampling unit of the (N - i + 1)-th decoding module.

3. The lightning wave generation method based on the diffusion model according to claim 2, wherein The number of data in the encoding modules and the decoding modules in the denoising neural network is both 3.

4. The lightning wave generation method based on a diffusion model according to claim 2, wherein During the training process of the denoising neural network, the loss function L(total) is used to backpropagate the gradient to update the parameters. Its loss function is: L(total)=L(base)+λ(L(low)+L(high)); Among them, L(base) represents the expected value of the square of the two-norm of the difference between the random noise ∈ calculated on the (x0, ∈, t) distribution and the predicted noise ∈ of the denoising neural network. θ (x t , t), L(low) is the expected value of the two-norm of the difference between the input and output of the low-frequency branch calculated on the (x0, t) distribution, L(high) is the expected value of the two-norm of the difference between the input and output of the high-frequency branch calculated on the (x0, t) distribution, and the high-frequency branch includes high-frequency horizontal branches, high-frequency vertical branches, and high-frequency diagonal branches; x0 is the lightning wave data sample; ∈ is the random noise used for the forward diffusion of the sample x0, and ∈ follows a normal distribution; t is the number of steps for the forward diffusion of the sample x0, and θ is the parameter of the denoising neural network.

5. The lightning wave generation method based on a diffusion model according to claim 4, wherein The training steps of the denoising neural network are as follows: First, extract the training sample x0 from the learning dataset and add noise through the diffusion model to obtain the noisy sample xt; input the noisy sample x t into the denoising neural network to obtain the predicted noise ∈ θ (x t , t); Then calculate the loss function L(total) of the model, and update the parameters to be learned of the denoising neural network by backpropagating the gradient through L(total); Repeat the above steps until the number of updates of the denoising network reaches the set value.

6. The lightning wave generation method based on a diffusion model according to claim 5, wherein The construction method of the learning dataset is as follows: the original lightning wave data is formed into multiple window data through a sliding window. Calculate the mean and standard deviation of each window data, and extract the window data that satisfies the condition that the maximum value is greater than or equal to the sum of the multiples of the mean and the standard deviation, and the minimum value is less than or equal to the difference between the mean and the multiples of the standard deviation, and add it to the learning dataset.

7. The lightning wave generation method based on the diffusion model according to claim 1, characterized in that, The processing methods of the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch for the input subband are all: perform convolutional processing on the input subband and then activate it. After the activation value is subjected to convolutional processing, it is weighted and calculated with the set weight to obtain the output data; The set weight of the low-frequency branch is the low-frequency weight, and the set weights of the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch are all the high-frequency weights.

8. The lightning wave generation method based on a diffusion model according to claim 1, wherein Group convolution is used for convolutional processing on the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch.

9. The lightning wave generation method based on a diffusion model according to claim 1, wherein, The feature synthesis function uses the IWT function, and the normalization method of the synthesized data is: take half of the synthesized data and add 0.

5.

10. A system for implementing the lightning wave generation method based on a diffusion model according to any one of claims 1-9, characterized in that, It includes a memory and a processor. A computer program is stored in the memory. The processor is connected to the memory and is used to execute the computer program to implement the lightning wave generation method based on the diffusion model according to any one of claims 1-9.

Images