A lightning wave generation method and system based on diffusion model

Abstract

The present invention relates to the field of lightning data processing technology, and in particular to a method and system for generating lightning waves based on a diffusion model. The present invention uses a diffusion model to generate lightning wave data, and the denoising neural network used in the reverse diffusion process is improved on the UNET model structure, and the improvement method is: a wavelet attention unit is added after each downsampling and before each upsampling. The present invention improves the denoising neural network, combines the attention technology wavelet transform, and captures multi-scale transformations at the same time, thereby improving the model's detail capture capability, especially capturing high-frequency components, and helping the model better retain the key features of lightning waves, thereby improving the quality of the diffusion model applied to lightning wave generation tasks. The present invention overcomes the defect of the above-mentioned prior art in which the diffusion model is directly applied to lightning waves.

Description

Technical Field

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

[0002] Narrow dipole events (NBEs) are short, intense electromagnetic waves caused by lightning discharges. This characteristic makes them difficult to detect in the very low frequency band, resulting in a limited number of samples, hindering comprehensive lightning research. Therefore, it is essential to use data generation methods to learn the underlying structure and distribution of data for sample augmentation.

[0003] Currently, mainstream data augmentation models include variational autoencoders (VAEs), generative adversarial networks (GANs), and diffusion models. VAE models often produce blurry and distorted results. GAN models require simultaneous training of both the generator and the discriminator, making it difficult to find a balance. Furthermore, generation quality is overly dependent on the quality of the training set, resulting in poor training results. Diffusion models employ a different approach when generating samples, gradually diffusing noise during the generation process. This approach results in clearer, more realistic, and more stable samples, while also maintaining good diversity.

[0004] The existing diffusion model has the following difficulties in generating lightning waveforms:

[0005] Lightning waveforms change instantaneously, making it difficult to capture the details of instantaneous jumps;

[0006] It is necessary to generate high-resolution samples. When there are many of them, the computational overhead of gathering all the information of the image is high.

[0007] Focusing on the entire image leads to the introduction of redundant information, and the interference model focuses on the key waveform area. Summary of the Invention

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

[0009] The present invention proposes a lightning wave generation method based on a diffusion model. The diffusion model is used to generate lightning wave data. The denoising neural network used in the reverse diffusion process is improved on the UNET model structure. The improvement is as follows: a wavelet attention unit is added after each downsampling and before each upsampling.

[0010] 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;

[0011] The input of the wavelet attention unit is processed by the attention network and decomposed into a low-frequency sub-band, a high-frequency sub-band in the horizontal direction, a high-frequency sub-band in the vertical direction, and a high-frequency sub-band in the diagonal direction by the feature decomposition function. 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, and then synthesized by the feature synthesis function. The synthesized data is normalized and used as the output of the wavelet attention unit.

[0012] Preferably, the denoising neural network comprises: an input convolution layer, a plurality of encoding modules, an intermediate module, a plurality of decoding modules and an output convolution layer connected in sequence; the number of encoding modules and decoding modules is equal;

[0013] The encoding module includes a downsampling unit, a wavelet attention unit, and a convolution unit connected sequentially. The output of the downsampling unit is also connected to the input of the convolution unit.

[0014] The decoding module consists of a wavelet attention unit, a convolution unit, and an upsampling unit connected sequentially; the input of the wavelet attention unit is also connected to the input of the convolution unit;

[0015] The encoding modules and decoding modules are arranged 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.

[0016] Preferably, the data of the encoding module and the decoding module in the denoising neural network are both 3.

[0017] Preferably, the loss function L(total) is used to back-propagate gradients to update parameters during the denoising neural network training process, and the loss function is:

[0018] L(total)=L(base)+λ(L(low)+L(high))

[0019] where L(base) represents the random noise ∈ calculated on the (x0,∈,t) distribution and the prediction noise ∈ of the denoising neural network θ (x t,t), L(low) is the expectation of the square of the two-norm 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 difference between the input and output of the high-frequency branch calculated on the (x0,t) distribution. The high-frequency branches include high-frequency horizontal branches, high-frequency vertical branches and high-frequency diagonal branches. x0 is a lightning wave data sample; ∈ is the random noise used for the forward diffusion of sample x0, and ∈ conforms to the normal distribution; t is the number of forward diffusion steps of sample x0, and θ is the parameter of the denoising neural network.

[0020] Preferably, the training steps of the denoising neural network are as follows:

[0021] First, we extract the training sample x0 from the learning data set and add noise through the diffusion model to obtain the noise sample x t ; The noise sample x t Input denoising neural network and get prediction noise ∈ θ (x t ,t);

[0022] Then calculate the loss function L(total) of the model, and update the parameters to be learned of the denoising neural network through the back propagation gradient of L(total);

[0023] Repeat the above steps until the number of updates of the denoising network reaches the set value.

[0024] Preferably, the method for constructing the learning data set is: the original lightning wave data is formed into multiple window data through a sliding window, the mean and standard deviation of each window data are calculated, and the window data whose maximum value is greater than or equal to the sum of the mean value and the standard deviation multiple and the minimum value is less than or equal to the difference between the mean value and the standard deviation multiple are extracted and added to the learning data set.

[0025] Preferably, the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch all process the input sub-band by performing convolution processing on the input sub-band and then activating it, and performing convolution processing on the activation value and then performing weighted calculation with the set weight to obtain output data;

[0026] 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.

[0027] Preferably, group convolution is performed on the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch.

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

[0029] The present invention proposes a system for implementing the method for generating lightning waves based on a diffusion model, comprising a memory and a processor, wherein a computer program is stored in the memory, and the processor is connected to the memory, and the processor is used to execute the computer program to implement the method for generating lightning waves based on a diffusion model.

[0030] The advantages of the present invention are:

[0031] (1) Considering that lightning wave data contains rich local features and mutation points, such as peaks and steep edges, the present invention introduces wavelet attention units into the denoising neural network, combines the attention technology wavelet transform, and captures multi-scale transformations at the same time, thereby improving the model's ability to capture details, especially high-frequency components, helping the model to better retain the key features of lightning waves, thereby improving the quality of the diffusion model applied to lightning wave generation tasks.

[0032] (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, which is then synthesized. This allows the model 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 amount of computation while maintaining the efficiency of feature extraction. In addition, different weights are set for the low and high frequency branches, which makes it easier for the model to distinguish between the main contours and details, which is conducive to further improving the generation quality.

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

[0034] (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, so that wavelet attention is applied multiple times in the entire feature extraction process, which can enhance feature representation at different scales and further improve the effect of feature fusion.

[0035] (5) The present invention adds loss terms for low-frequency and high-frequency subbands to the loss function of model training. The low-frequency part can ensure the accuracy of the overall structure, while the high-frequency part enhances the reconstruction of details, which is beneficial to improving the fidelity of the generated data, especially for detail-sensitive data such as lightning waves.

[0036] (6) In the process of constructing the data set, the present invention screens window data with high volatility, which is conducive to capturing extreme situations, making the trained model more robust when generating complex waveforms and more suitable for the characteristics of lightning wave data, thus avoiding the generated lightning wave data being too smooth. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] Figure 1This is the structural diagram of the denoising neural network proposed in the present invention;

[0038] Figure 2 This is the structural diagram of the wavelet attention unit proposed in the present invention;

[0039] Figure 3 This is a flow chart of the dataset construction method proposed in the present invention;

[0040] Figure 4 This is a flow chart of the denoising neural network training method proposed in the present invention;

[0041] Figure 5 This is a flow chart of the lightning wave generation method based on the diffusion model proposed in the present invention;

[0042] Figure 6 This is the effect demonstration of Experiment 1;

[0043] Figure 7 This is the effect demonstration of Experiment 2. DETAILED DESCRIPTION

[0044] The following will provide a clear and complete description of the technical solutions in the embodiments of the present invention in conjunction with the accompanying drawings. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. All other embodiments obtained by ordinary technicians in this field based on the embodiments of the present invention without making any creative efforts shall fall within the scope of protection of the present invention.

[0045] In this embodiment, a denoising neural network is constructed to predict the noise contained in the input lightning wave data. Figure 1 ,The denoising neural network includes: a sequentially connected input convolution layer, multiple encoding modules, an intermediate module, multiple decoding modules and an output convolution layer.

[0046] The encoding module consists of a sequentially connected downsampling unit, a wavelet attention unit, and a convolution unit. The input of the downsampling unit serves as the input of the encoding module, and the output of the downsampling unit is also connected to the input of the convolution unit. The output of the convolution unit serves 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 concatenated and then fed into the convolution unit for convolution processing to obtain the output of the encoding module.

[0047] The decoding module includes a wavelet attention unit, a convolution unit and an upsampling unit connected in sequence; the input of the wavelet attention unit is the input of the decoding module, and the output of the upsampling unit is the output of the decoding module; the input of the convolution unit is also connected to the input of the wavelet attention unit.

[0048] The data of the encoding module and the decoding module are both N, and are numbered in ascending order according to the data transmission sequence. 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.

[0049] In this way, the input data of the decoding module is processed by the wavelet attention unit and then spliced ​​with the input data dimension to enter the convolution unit for processing. The data processed by the convolution unit is then spliced ​​with the data dimension output by the downsampling unit in the corresponding decoding module to enter the upsampling unit. The upsampling unit upsamples the input data to obtain the output data of the decoding module.

[0050] In a specific implementation, the number of encoding modules and decoding modules is set to 3, then the output of the downsampling unit in the first encoding module is also connected to the upsampling unit input of the third decoding module, the output of the downsampling unit in the second encoding module is also connected to the upsampling unit input of the second decoding module, and the output of the downsampling unit in the third encoding module is also connected to the upsampling unit input of the first decoding module.

[0051] The intermediate module can be specifically composed of two convolutional units connected in series.

[0052] Reference Figure 2 ,The wavelet attention unit in this scheme 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;

[0053] The input of the attention network is the input of the wavelet attention unit, and the output of the attention network is connected to the input of the eigendecomposition function. The low-frequency branch, high-frequency horizontal branch, high-frequency vertical branch, and high-frequency diagonal branch are connected in parallel between the eigendecomposition function and the feature synthesis function.

[0054] The eigendecomposition function uses the DWT function to decompose the input data into a low-frequency sub-band LL, a horizontal high-frequency sub-band LH, a vertical high-frequency sub-band HL, and a diagonal high-frequency sub-band HH;

[0055] LL, LH, HL and HH are processed by the low-frequency branch, high-frequency horizontal branch, high-frequency vertical branch and high-frequency diagonal branch respectively, and then synthesized by the feature synthesis function IWT. The synthesized data XI is normalized and used as the output Xinorm of the wavelet attention unit;

[0056] The formula is: Xinorm=[XI+1] / 2

[0057] XI= IWT[LL′,HL′,LH′,HH′]

[0058] Wherein, LL′, HL′, LH′ and HH′ represent the outputs of the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch and the high-frequency diagonal branch, respectively.

[0059] The synthetic 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 characteristic peak is reduced, which facilitates subsequent data processing.

[0060] The low-frequency branch is composed of two convolutional layers in series. LL is activated after the first convolution process and then after the second convolution process with the set low-frequency weight α L Weighted calculation is performed to obtain the output LL' of the low-frequency branch;

[0061] The high-frequency horizontal branch, high-frequency vertical branch and high-frequency diagonal branch are composed of two group convolution layers connected in series. LH is activated after the first group convolution process and then after the second group convolution process with the set high-frequency weight α H Weighted calculation is performed to obtain the output LH' of the high-frequency horizontal branch;

[0062] HL is activated after the first group convolution process, and then after the second group convolution process, it is combined with the high-frequency weight α H Weighted calculation is performed to obtain the output HL' of the high-frequency vertical branch;

[0063] HH is activated after the first group convolution process, and then after the second group convolution process, it is combined with the high-frequency weight α H Weighted calculation is performed to obtain the output HH' of the high-frequency diagonal branch;

[0064] The formula is:

[0065] LL'=α L ×Conv(Silu(Conv(LL)))

[0066] [LH',HL',HH'] T =α H ×GroupConv(Silu(GroupConv[LH,HL,HH] T ))

[0067] In order to facilitate subsequent processing, the output of the wavelet attention unit needs to be normalized, that is, the synthetic data output by the feature synthesis function is normalized to obtain the output Xinorm of the wavelet attention unit.

[0068] Temporal code embedding is performed during the downsampling process of the encoding module, and temporal code embedding is performed during the upsampling process of the decoding module.

[0069] Let the input data of the i-th encoding module be recorded as Xin_i, and the output data be recorded as Xout_i; the output of the downsampling unit in the i-th encoding module be recorded as Xd_i;

[0070] The processing process of the input data Xin_i in the i-th encoding module is as follows: Xin_i is sampled by the downsampling unit to obtain Xd_i, the result of Xd_i processing by the wavelet attention unit is concatenated with the dimension of Xd_i and then processed by the convolution unit to obtain Xout_i; the formula is expressed as:

[0071] Xout_i=Conv[Concat(Xinorm _i,Xd_i)]

[0072] Xinorm_i represents the processing result of Xd_i by the wavelet attention unit in the i-th encoding module, and Concat represents dimension splicing.

[0073] The input data of the i-th decoding module is recorded as X'in_i, and the output data is recorded as X'out_i;

[0074] The input data X'in_i of the i-th decoding module is processed by the wavelet attention unit, and then concatenated with X'in_i and entered into the convolution unit for processing. The output data of the convolution unit is concatenated with the Xd_(N-i+1) 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;

[0075] The processing formula of the i-th decoding module is expressed as:

[0076] X'out_i=UP{Concat(Conv(Concat(X'in_i,X'inorm _i)),Xd_(N-i+1))}

[0077] X'inorm_i represents the processing result of X'in_i by the wavelet attention unit in the i-th decoding module, and UP represents upsampling.

[0078] Reference Figure 3 , the learning dataset construction process of the denoising neural network is as follows:

[0079] SA1. Obtain the original lightning wave data S and use a sliding window to form multiple window data S(g), 1≤g≤G, G is the number of window data divided by the original lightning wave data S, and g is an ordinal number;

[0080] S(g)={s((g-1)×z),s((g-1)×z+1),s((g-1)×z+2),…,s((g-1)×z+K-1)}

[0081] Among them, K is the window length, z is the step size, and in specific implementation, z=1 can be set.

[0082] SA2. Calculate the mean μ and standard deviation σ of each window data, and let the mean and standard deviation of the window data S(g) be denoted as μ(g) and σ(g) respectively;

[0083] SA3. Extract window data from all window data whose maximum value is greater than or equal to the sum of the mean value and the standard deviation multiple, and whose minimum value is less than or equal to the difference between the mean value and the standard deviation multiple, and add them to the learning dataset.

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

[0085] In specific implementation, step SA3 can be divided into the following sub-steps:

[0086] SA31. Determine 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;

[0087] If yes, then add the window data S(g) to the learning data set and then execute step SA32;

[0088] If not, proceed to step SA32;

[0089] SA32, determine whether g is less than G;

[0090] If yes, update g to g+1 and return to step SA31;

[0091] If not, the original lightning wave data S is updated, and the process returns to step SA1.

[0092] Reference Figure 4 , the training method of the denoising neural network includes the following steps:

[0093] St1, initialize the parameters to be learned of the denoising neural network, including the node weights of the model;

[0094] St2, extract the training sample x0 from the learning data set and add noise through the diffusion model to obtain the noise sample x t ;

[0095]

[0096] α t =1-β t

[0097]

[0098] Among them, α t 、 and β t are all hyperparameters, β t is a random value in the interval (0,1), ∈ is a random noise that conforms to the normal distribution, that is, ∈~N(0,1), and t is the set number of diffusion steps;

[0099] St3, the noise sample x t Input denoising neural network, denoising neural network output prediction noise ∈ θ (x t ,t), θ represents the parameters of the denoising neural network;

[0100] St4. Calculate the model's loss function L(total) and backpropagate the gradient of L(total) to update the parameters to be learned in the denoising neural network.

[0101] L(total)=L(base)+λ(L(low)+L(high));

[0102] Where L(base) represents the random noise ∈ and prediction noise ∈ calculated on the (x0,∈,t) distribution. θ (x t ,t), L(low) is the expectation of the square of the two-norm 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 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.

[0103] Right now:

[0104] L(low)= [||LL'-LL||2]

[0105] L(high)= [||LH'-LH||2]+ [||HL'-HL||2]+ [||HH'-HH||2]

[0106] St5. Repeat the above steps St1-St4 until the number of updates of the denoising network reaches the set value.

[0107] The lightning wave generation method based on the diffusion model proposed in this embodiment is used to learn existing lightning wave data samples and obtain more similar lightning wave data samples, such as Figure 5 As shown, the generation method includes the following steps:

[0108] S1. Obtain the lightning wave data sample X0 to be learned, and add noise to X0 using the diffusion model to obtain the noise sample XT;

[0109]

[0110] α T =1-β T

[0111]

[0112] Among them, α T 、 and β T are all hyperparameters, β T is a random value in the interval (0,1), ∈ is a random noise that conforms to the normal distribution, that is, ∈~N(0,1), and T is the set number of diffusion steps.

[0113] S2. Predict the sample noise through the denoising neural network, iterate the diffusion step number t from T to 0, realize the reverse diffusion denoising of the noise sample XT, and obtain the generated lightning wave data sample X'0. X'0 is the new lightning wave data obtained by learning X0.

[0114] Step S2 is specifically as follows:

[0115] S21, input the noise sample Xt into the denoising neural network, and the denoising neural network outputs the predicted noise ∈ θ (Xt, t); using the inverse diffusion formula, the noise sample Xt is denoised and the noise sample X(t-1) is obtained; the initial value of t is T;

[0116] The reverse diffusion formula is:

[0117]

[0118]

[0119]

[0120] S22. Determine whether t is equal to 1; if yes, obtain X'0=X(t-1); if not, update t to t-1, and then return to step S21.

[0121] Thus, by adopting the present 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 needs.

[0122] The lightning wave generation method based on the diffusion model is verified below with reference to specific embodiments.

[0123] In this embodiment, three contrast generation methods are provided.

[0124] Comparative generation method 1: VAE for short, directly uses the variational encoder (VAE) to encode and decode the input lightning wave data to generate new lightning wave data;

[0125] Comparative generation method 2: GAN for short, based on the generative adversarial network (GAN), reconstructs the input lightning wave data to obtain new lightning wave data;

[0126] Comparative generation method 3: DDPM for short, uses a diffusion model to generate new lightning wave data, and uses the Unet network to predict noise in the reverse diffusion process;

[0127] The method of the present invention, referred to as WAB, uses a diffusion model to generate new lightning wave data, and the noise prediction in the reverse diffusion process uses a denoising neural network with added wavelet attention units proposed by the present invention.

[0128] The comparison generation method three is the baseline model of the method of the present invention, and the generation steps of the two refer to the above steps S1-S2.

[0129] In this embodiment, a learning dataset is first constructed using the dataset construction method proposed in the present invention, with specific reference to steps SA1-SA3 above.

[0130] In this embodiment, the learning data set is divided into a training set and a test set. The training set is used to train the model, including VAE, GAN network, Unet network and the denoising neural network proposed in the present invention. Then, the trained model is used to generate the same amount of lightning wave data samples as the test set to form a generated data set, and then the spatial distribution differences between the generated data set and the test set are compared.

[0131] Specifically, in the contrast generation method 1 and the contrast generation method 2, the trained VAE network and GAN network are directly used to generate lightning wave data to construct the corresponding generated datasets;

[0132] Comparison of generation method 3 and the method of the present invention uses the above steps S1-S2 to generate lightning wave data to construct a corresponding generation data set.

[0133] In this example, data generation methods are compared in two experiments.

[0134] In both experiments, a learning dataset was constructed based on lightning wave data from different historical time periods, and then divided into a training set containing 900 samples and a test set containing 300 samples.

[0135] In both experiments, the training and test set parameters were set to 64*64, that is, the dataset samples were adjusted to image data with a length and width of 64 pixels. The batch size of the model training was set to 16, the initial learning rate was 0.002, and the training rounds were 1000 rounds.

[0136] In both experiments, 300 images were generated using the trained model and compared with the test set. In experiment 1, the spatial distribution difference between the generated dataset and the test set was measured using the FID (Frechet Inception Distance) metric. The experimental results are as follows: Figure 6 As shown in the second experiment, the spatial distribution difference between the generated dataset and the test set is measured by the SSIM (structural similarity) indicator. The experimental results are shown in Figure 7 shown.

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

[0138] Of course, it will be apparent to those skilled in the art that the present invention is not limited to the details of the exemplary embodiments described above, but also encompasses the same or similar structures that can be implemented in other specific forms without departing from the spirit or essential characteristics of the invention. Therefore, the embodiments should be considered in all respects as illustrative and non-restrictive, and the scope of the invention is defined by the appended claims, not the foregoing description, and it is intended that all variations that fall within the meaning and range of equivalents of the claims be encompassed within the present invention. Any reference signs in the claims should not be construed as limiting the claim to which they relate.

[0139] In addition, it should be understood that although this specification is described in terms of implementation methods, not every implementation method contains only one independent technical solution. This narrative method of the specification is only for the sake of clarity. Those skilled in the art should regard the specification as a whole. The technical solutions in each embodiment can also be appropriately combined to form other implementation methods that can be understood by those skilled in the art.

[0140] The technology, shape, and structure not described in detail in the present invention are all well-known technologies.

Claims

1. A lightning wave generation method based on a diffusion model, characterized in that: A diffusion model is used to generate lightning wave data. The denoising neural network used in the reverse diffusion process is improved on the UNET model structure by adding a wavelet attention unit 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 decomposed into a low-frequency sub-band, a high-frequency sub-band in the horizontal direction, a high-frequency sub-band in the vertical direction, and a high-frequency sub-band in the diagonal direction by the feature decomposition function. The four sub-bands are processed by the low-frequency branch, the high-frequency horizontal branch, the high-frequency vertical branch, and the high-frequency diagonal branch respectively, and then synthesized by the feature synthesis function. The synthesized data is normalized and used as the output of the wavelet attention unit. The denoising neural network includes: a sequentially connected input convolution layer, multiple encoding modules, an intermediate module, multiple decoding modules and an output convolution layer; the number of encoding modules and decoding modules is equal; The encoding module includes a downsampling unit, a wavelet attention unit, and a convolution unit connected sequentially. The output of the downsampling unit is also connected to the input of the convolution unit. The decoding module consists of a wavelet attention unit, a convolution unit, and an upsampling unit connected sequentially; the input of the wavelet attention unit is also connected to the input of the convolution unit; The encoding modules and decoding modules are arranged 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.

2. The method for generating lightning waves based on a diffusion model according to claim 1, wherein: The data of the encoding module and decoding module in the denoising neural network are both 3.

3. The lightning wave generation method based on the diffusion model according to claim 1, characterized in that: During the denoising neural network training process, the loss function L(total) is used to back-propagate the gradient to update the parameters. The loss function is: L(total)=L(base)+λ(L(low)+L(high)); where L(base) represents the random noise ∈ calculated on the (x0,∈,t) distribution and the prediction noise ∈ of the denoising neural network θ (x t ,t), L(low) is the expectation of the square of the two-norm 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 difference between the input and output of the high-frequency branch calculated on the (x0,t) distribution. The high-frequency branches include high-frequency horizontal branches, high-frequency vertical branches and high-frequency diagonal branches. x0 is a lightning wave data sample; ∈ is the random noise used for the forward diffusion of sample x0, and ∈ conforms to the normal distribution; t is the number of forward diffusion steps of sample x0, and θ is the parameter of the denoising neural network.

4. The method for generating lightning waves based on a diffusion model according to claim 3, wherein: The training steps of the denoising neural network are as follows: First, we extract the training sample x0 from the learning data set and add noise through the diffusion model to obtain the noise sample x t ; The noise sample x t Input denoising neural network and get prediction 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 through the back propagation gradient of L(total); Repeat the above steps until the number of updates of the denoising network reaches the set value.

5. The method for generating lightning waves based on a diffusion model according to claim 4, characterized in that: The method for constructing the learning dataset is as follows: the original lightning wave data is formed into multiple window data through a sliding window, the mean and standard deviation of each window data are calculated, and the window data whose maximum value is greater than or equal to the sum of the mean value and the standard deviation multiple and the minimum value is less than or equal to the difference between the mean value and the standard deviation multiple are extracted and added to the learning dataset.

6. The method for generating lightning waves based on a diffusion model according to claim 1, wherein: The low-frequency branch, high-frequency horizontal branch, high-frequency vertical branch, and high-frequency diagonal branch all process the input subband by performing convolution on the input subband and then activating it. The activation value is then convolved with the set weights 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.

7. The method for generating lightning waves based on a diffusion model according to claim 1, wherein: Group convolution is used for convolution processing on the high-frequency horizontal branch, high-frequency vertical branch and high-frequency diagonal branch.

8. The method for generating lightning waves based on a diffusion model according to claim 1, wherein: The feature synthesis function uses the IWT function, and the normalization method of the synthetic data is: take half of the synthetic data and add 0.

5.

9. A system for implementing the method for generating lightning waves based on a diffusion model according to any one of claims 1 to 8, characterized in that: The invention comprises 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 according to any one of claims 1 to 8.

Images