A Crystal Structure Optimization Method Based on a Diffusion Model

Abstract

This invention discloses a crystal structure optimization method based on a diffusion model, belonging to the field of crystal structure optimization. The method includes: a training phase, in which a neural network model is trained using a diffusion model training method on a training set containing only the optimal crystal structure, enabling the neural network model to learn the mapping relationship from noisy structures to the optimal structure; the training set includes a sequence of atomic types, optimal fractional coordinates, and a unit cell basis vector matrix; during the training process, noise is applied to the optimal crystal structure at preset time steps to generate a noisy structure, and the noise intensity is used as the model input parameter; the output of the neural network model includes predicted atomic displacements and predicted unit cell strains; and an inference phase, in which the initial crystal structure to be optimized is input to the trained neural network model for iterative optimization, outputting the optimized crystal structure. This invention achieves low-step, stable, and controllable optimization of initial structures from any source, provided that only the optimal structure is used for training.

Description

Technical Field

[0001] This invention belongs to the field of crystal structure optimization technology, and particularly relates to a crystal structure optimization method based on a diffusion model. Background Technology

[0002] Crystal structure optimization is a core task in computational materials science, aiming to find the stable structure with the lowest energy by adjusting atomic coordinates and unit cell parameters. Currently, this field mainly relies on first-principles calculation methods based on physical principles, such as density functional theory (DFT). While these methods can achieve high optimization accuracy, they typically require multiple iterative calculations of energy and force, making the computation process extremely time-consuming and heavily reliant on high-performance computing resources, which is insufficient to meet the efficiency requirements of high-throughput material screening and design.

[0003] In recent years, machine learning techniques, particularly diffusion models and denoising score matching methods, have shown great potential in molecular and material structure generation. However, existing techniques largely focus on generating structures "from scratch" rather than being specifically designed for structure optimization tasks "from suboptimal to optimal." Furthermore, these machine learning methods typically rely on a supervised learning paradigm of one-to-one correspondence between "initial structure" and "optimal structure," while real-world material databases often only contain the final stable structure, lacking a suitable initial structure to pair with it, leading to difficulties in obtaining training data and limiting the applicability of the models. Simultaneously, most existing schemes only optimize atomic coordinates, neglecting the cooperative relaxation of unit cell parameters, making it difficult to achieve global optimization of crystal structures, and presenting challenges in consistency and continuity when dealing with periodic boundary conditions. Summary of the Invention

[0004] To address the aforementioned technical problems, this invention proposes a crystal structure optimization method based on a diffusion model, thereby resolving the issues present in the prior art.

[0005] To achieve the above objectives, in a first aspect, the present invention provides a crystal structure optimization method based on a diffusion model, comprising: During the training phase, a neural network model is trained using a diffusion model training method based on a training set containing only the optimal crystal structure. This allows the neural network model to learn the mapping relationship from noisy structures to the optimal structure. The training set includes a sequence of atomic types, fractional coordinates of the optimal crystal structure, and the unit cell basis vector matrix. The output of the neural network model includes predicted atomic displacements and predicted unit cell strains. The training phase includes: applying noise to the optimal crystal structure incrementally at preset time steps to construct a sequence of noisy crystal structures; using the noisy crystal structure, time step parameters, and noise intensity as inputs, and the optimal structure as the supervised target, training the neural network model to learn the reverse diffusion process. During the inference phase, the initial crystal structure to be optimized is input into the trained neural network model, iteratively optimized, and the optimized crystal structure is output.

[0006] Preferably, the process of applying noise perturbation to the optimal crystal structure in the training set includes: Gaussian noise is applied to the atomic coordinates of the optimal crystal structure to generate noisy atomic coordinates; A symmetric strain perturbation is applied to the cell basis vector matrix of the optimal crystal structure to generate a noisy cell.

[0007] Preferably, the process of training the neural network model further includes: Calculate the coordinate denoising loss based on the predicted atomic displacement and the actual atomic displacement; The cell loss is calculated based on the predicted cell strain value and the actual cell parameters. The parameters of the neural network model are updated by combining the coordinate denoising loss and the cell loss.

[0008] Preferably, the formula for calculating the coordinate denoising loss is: ; in, These are predicted atomic shift values. For coordinate noise amplitude, These are actual atomic shifts; The formula for calculating cell loss is: ; in, The unit cell parameters are obtained from neural network prediction, representing the lengths of the three sides and the included angles between the triangles. These are the corresponding optimal unit cell parameters.

[0009] Preferably, the reasoning stage includes: Set the iteration parameters, including noise intensity, iteration step size, and maximum number of iterations; Based on the iteration parameters, the initial crystal structure is subjected to multiple rounds of deterministic iterative updates.

[0010] Preferably, the process of performing multiple rounds of deterministic iterative updates on the initial crystal structure includes: In a single iteration, the current atomic fractional coordinates and the current unit cell basis vector matrix are converted into atomic Cartesian coordinates; The atomic type sequence, the current atomic fraction coordinates, the current unit cell basis vector matrix, and a fixed small noise are input into the trained neural network model to obtain the current atomic displacement prediction value and the current unit cell strain prediction value. Based on the current predicted atomic displacement value and the iteration step size, update the atomic Cartesian coordinates, and convert the updated atomic Cartesian coordinates back to atomic fractional coordinates; Update the current cell basis matrix based on the current cell strain prediction value.

[0011] Preferably, the formula for updating the atomic Cartesian coordinates is: ; ; in, Let Cartesian coordinates be the atomic coordinates of the current iteration step. The iteration step size parameter, These are predicted atomic shift values. For a fixed small noise standard deviation, This is the atomic fractional coordinate matrix. Let be the unit cell basis vector matrix of the current iteration step; The formula for updating the current unit cell basis vector matrix is: ; in, To obtain the symmetrical part, This represents the predicted value of the unit cell strain.

[0012] Preferably, the reasoning stage further includes: During the iteration process, the atomic displacement and the unit cell change are monitored; The iteration terminates when the atomic displacement and the unit cell change are both below a preset threshold.

[0013] In a second aspect, the present invention also provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the steps of the method described in the first aspect.

[0014] Thirdly, the present invention also discloses a computer program product, including a computer program that, when executed by a processor, implements the steps of the method described in the first aspect.

[0015] Compared with the prior art, the present invention has the following advantages and technical effects: This invention provides a crystal structure optimization method based on a diffusion model. Through a "training phase based on a training set containing only the optimal crystal structure" technique, it directly achieves the technical effect of training without requiring paired 'initial structure-optimal structure' data. It fundamentally eliminates the dependence on scarce and expensive paired data, utilizing only common optimal structure sets found in materials databases, greatly reducing data acquisition costs and barriers.

[0016] This invention achieves a fast and streamlined optimization process by employing a technique where the initial crystal structure to be optimized is input into a trained neural network model during the inference phase for iterative optimization. It transforms the traditional model, which relies on physical calculations and requires multiple iterations of first-principles calculations, into a model of one-time training and multiple rapid inferences. Utilizing the trained neural network as the optimization operator, optimization can be completed within a small number of iterations, significantly reducing computation time and resource consumption. Attached Figure Description

[0017] The accompanying drawings, which form part of this application, are used to provide a further understanding of this application. The illustrative embodiments and descriptions of this application are used to explain this application and do not constitute an undue limitation of this application. In the drawings: Figure 1 This is a flowchart of a crystal structure optimization method based on a diffusion model according to an embodiment of the present invention; Figure 2 This is a flowchart of the training phase method according to an embodiment of the present invention; Figure 3 This is a flowchart of the reasoning stage method according to an embodiment of the present invention. Detailed Implementation

[0018] It should be noted that, unless otherwise specified, the embodiments and features described in this application can be combined with each other. This application will now be described in detail with reference to the accompanying drawings and embodiments.

[0019] It should be noted that the steps shown in the flowchart in the accompanying drawings can be executed in a computer system such as a set of computer-executable instructions, and although a logical order is shown in the flowchart, in some cases the steps shown or described may be executed in a different order than that shown here.

[0020] Example 1 like Figure 1 As shown, this embodiment provides a crystal structure optimization method based on a diffusion model, including: Step S01: Training phase. Based on a training set containing only the optimal crystal structure, a neural network model is trained using a diffusion model training method to enable the neural network model to learn the mapping relationship from noisy structures to the optimal structure. The training set includes a sequence of atomic types, fractional coordinates of the optimal crystal structure, and unit cell basis vector matrix. The output of the neural network model includes predicted atomic displacements and predicted unit cell strains. The training phase includes: applying noise to the optimal crystal structure incrementally at preset time steps to construct a noisy crystal structure sequence; using the noisy crystal structure, time step parameters, and noise intensity as inputs, and the optimal structure as the supervised target, training the neural network model to learn the reverse diffusion process. Furthermore, applying noise perturbation to the optimal crystal structure in the training set includes: Gaussian noise is applied to the atomic coordinates of the optimal crystal structure to generate noisy atomic coordinates; A symmetric strain perturbation is applied to the cell basis vector matrix of the optimal crystal structure to generate a noisy cell.

[0021] Furthermore, training the neural network model includes: The atomic type sequence, noisy fractional coordinates, noisy unit cell basis vector matrix, and noise intensity are input into the neural network model. Obtain the predicted atomic displacement and predicted unit cell strain values ​​output by the neural network model.

[0022] Furthermore, training the neural network model also includes: Calculate the coordinate denoising loss based on the predicted atomic displacement and the actual atomic displacement; The cell loss is calculated based on the predicted cell strain value and the actual cell parameters. The parameters of the neural network model are updated by combining the coordinate denoising loss and the cell loss.

[0023] Specifically, the aim is to train parameters as network So that it is in a given Under these conditions, it is possible to simultaneously predict the atomic Cartesian displacement field. With cell strain (Symmetric, 6 degrees of freedom). Among them, It is implemented using a Graph Neural Network (GNN). Training samples are drawn only from the set of optimal structures, without relying on paired "initial structure - optimal structure" data. Represents a sequence of atom types (length is...) ); Represents the fractional coordinate matrix, and the unit cell Define periodic boundary conditions together; Represents the unit cell basis vector matrix (row vectors are) ),and Parameterization is equivalent and interchangeable; The conditional input representing the disturbance intensity must include at least the coordinate noise amplitude. It may also include cell perturbation / strain amplitude. , recorded as .

[0024] like Figure 2 As shown, in step S01, the network The training process is as follows: S0101 Data Standardization. Read the training set. The representation is unified using fractional coordinates and matrix cells, while retaining the element vocabulary / embedding definitions. Necessary graph construction parameters (such as cutoff radius and radial / angular basis functions) are prepared according to periodic boundary condition rules, but the specific architecture is not limited.

[0025] S0102 Synthetic Perturbation. Noise is added to the coordinates in Cartesian space and to the unit cell under strain parameterization to synthesize a "noisy initial" to replace expensive paired data.

[0026] Specifically, in step S0102, the synthesized perturbation includes coordinate perturbation and cell perturbation; The method for coordinate perturbation is: Let ,sampling (logarithmic uniformity) ),make And write back the fractional coordinates .

[0027] The method for cell perturbation (strain) is: sampling symmetrical strain. (range from) (Control), and then perform unit cell To perform a disturbance, i.e. ,right Amplitude limiting or spectral norm clipping is performed to ensure positive volume and non-degeneracy.

[0028] S0103 Condition Preparation and Input Packaging. As input for a single training sample; Perform scalar embedding and use it for feature modulation (such as FiLM).

[0029] S0104 Forward Prediction. Based on the input network. Output .

[0030] in Interpreted in Cartesian space; Used for unit cell renewal (During training, it can be used as a supervisory signal or simply as a regularization signal).

[0031] S0105 Target Construction and Loss. To enable the model to learn the orientation field "from noisy regression to the optimal" under periodic boundaries, this step first constructs the supervised target according to the minimum image rule, and then defines the coordinate denoising loss and cell correlation loss respectively, and adds the necessary regularization / constraints.

[0032] Specifically, in step S0105, the loss function includes coordinate denoising loss, cell loss, and regularization / constraint term; The form of coordinate denoising loss is: ; in, These are predicted atomic shift values. For coordinate noise amplitude, These represent the actual atomic displacements.

[0033] The form of unit cell loss is: ; in, The unit cell parameters are obtained from neural network prediction, representing the lengths of the three sides and the included angles between the triangles. These are the corresponding optimal unit cell parameters.

[0034] Regularization / constraint terms can include minimum interatomic spacing penalty, volume positivity penalty, angular range penalty, etc.

[0035] Finally, the total loss function is: ; in, For coordinate denoising loss, For unit cell strain loss, used to constrain the difference between the predicted strain tensor and the actual strain. For unit cell loss, This is a regularization term used to prevent overfitting. , , These are the weighting coefficients for each loss term, used to balance the contributions of different loss terms.

[0036] S0106 parameter update. Training was performed using optimizers such as AdamW; small improvements were made in the later stages of training. Sample weights or course learning may be used to reinforce small-step convergence; for Applying tanh limiting or spectral norm clipping can improve stability.

[0037] S0107 training output. A parameter-fixed output is obtained. , as an optimization operator in the reasoning stage.

[0038] Step S02: Inference stage, the initial crystal structure to be optimized is input into the trained neural network model, iterative optimization is performed, and the optimized crystal structure is output.

[0039] Furthermore, the reasoning phase includes: Set the iteration parameters, including noise intensity, iteration step size, and maximum number of iterations; Based on the iteration parameters, the initial crystal structure is subjected to multiple rounds of deterministic iterative updates.

[0040] Furthermore, performing multiple rounds of deterministic iterative updates on the initial crystal structure includes: In a single iteration, the current atomic fractional coordinates and the current unit cell basis vector matrix are converted into atomic Cartesian coordinates; The atomic type sequence, the current atomic fraction coordinates, the current unit cell basis vector matrix, and a fixed small noise are input into the trained neural network model to obtain the current atomic displacement prediction value and the current unit cell strain prediction value. Based on the current predicted atomic displacement value and the iteration step size, update the atomic Cartesian coordinates, and convert the updated atomic Cartesian coordinates back to atomic fractional coordinates; Update the current cell basis matrix based on the current cell strain prediction value.

[0041] Furthermore, the reasoning stage also includes: During the iteration process, the atomic displacement and the unit cell change are monitored; The iteration terminates when the atomic displacement and the unit cell change are both below a preset threshold.

[0042] Specifically, deterministic iteration with low noise is used to ensure that any given initial structure Converging to a low-energy structure. In this embodiment, a "low-energy structure" is defined as: the structure obtained at the termination of inference, used as input for single-point DFT evaluation, where the maximum atomic force does not exceed 0.05 eV / .

[0043] like Figure 3 As shown, in step S02, the reasoning stage includes the following steps: S0201 Initialization and Threshold Setting. Settings (like Step size (like Step limit (like ), as well as minimum interatomic spacing, volume / angle variation threshold and median displacement early stopping threshold.

[0044] S0202 is a two-stage design, from coarse to fine. If the initial deviation is large, start with the medium stage. Stack First, perform a rough pullback, then switch to S0203 for fine-tuning with low noise.

[0045] S0203 Deterministic small-step iteration. For implement:

[0046] in, Let Cartesian coordinates be the atomic coordinates of the current iteration step. The iteration step size parameter, These are predicted atomic shift values. For a fixed small noise standard deviation, This is the atomic fractional coordinate matrix. Let be the unit cell basis vector matrix of the current iteration step; Take the symmetrical part, This represents the predicted value of the unit cell strain.

[0047] Perform constraint checks at each step (minimum spacing, volume, angle), and halve the steps if necessary. Alternatively, it can be reversed one step; early stopping occurs when both the median displacement and the cell change are below the threshold.

[0048] S0204 Inference Output. Optimized structure obtained. If necessary, Convert to Output.

[0049] The beneficial effects of this embodiment: Compared with existing technologies, the crystal structure optimization method based on a diffusion model proposed in this embodiment achieves stable and controllable optimization of initial structures from arbitrary sources in a few steps, while only using the optimal structure for training. It has the following advantages: First, it is data-friendly and independent of pairing: it does not require paired samples of "initial structure-optimal structure", and can complete training with only the optimal structure, which significantly reduces the cost of data collection and cleaning.

[0050] Second, rapid convergence with fewer steps: The inference stage adopts deterministic iteration with low noise, which can push the atomic coordinates and cell parameters to a low-energy reasonable solution within ten to tens of steps, reducing computation time and resource consumption.

[0051] Third, cell-coordinate co-optimization: cell updates are predicted by strain parameterization, avoiding problems such as negative volume and degradation caused by direct regression to the absolute cell, thus ensuring geometric and numerical stability.

[0052] Fourth, strict periodic boundary condition consistency: the displacement supervision and loss are constructed using the minimum image rule, the intercellular displacement is unique and continuous, and the convergence process is physically consistent.

[0053] Fifth, robust and controllable: It incorporates soft constraints such as volume positivity, angle range, and minimum atomic spacing, as well as early stopping criteria, making the optimization process interpretable and monitorable.

[0054] Sixth, cross-distribution adaptation: Through "noise condition-denoised direction field" learning, the robustness to different initial structure distributions is enhanced, which facilitates cross-data source migration and batch processing.

[0055] Seventh, easy to integrate and expand in engineering: The network, as a "callable optimization operator", can be seamlessly integrated with existing material processes and can be further enhanced with target property guidance or energy / force proxy distillation.

[0056] Example 2 This embodiment also discloses a computer-readable storage medium storing a computer program thereon, which, when executed by a processor, implements the steps of the method described in Embodiment 1.

[0057] Example 3 This embodiment also discloses a computer program product, including a computer program that, when executed by a processor, implements the steps of the method described in Embodiment 1.

[0058] The above are merely preferred embodiments of this application, but the scope of protection of this application is not limited thereto. Any variations or substitutions that can be easily conceived by those skilled in the art within the scope of the technology disclosed in this application should be included within the scope of protection of this application. Therefore, the scope of protection of this application should be determined by the scope of the claims.

Claims

1. A crystal structure optimization method based on a diffusion model, characterized in that, Includes the following steps: During the training phase, a neural network model is trained using a diffusion model training method based on a training set containing only the optimal crystal structure. This allows the neural network model to learn the mapping relationship from noisy structures to the optimal structure. The training set includes a sequence of atomic types, fractional coordinates of the optimal crystal structure, and the unit cell basis vector matrix. The output of the neural network model includes predicted atomic displacements and predicted unit cell strains. The training phase includes: applying noise to the optimal crystal structure incrementally at preset time steps to construct a sequence of noisy crystal structures; using the noisy crystal structure, time step parameters, and noise intensity as inputs, and the optimal structure as the supervised target, training the neural network model to learn the reverse diffusion process. The process of applying noise perturbation to the optimal crystal structure in the training set includes: applying Gaussian noise to the atomic coordinates of the optimal crystal structure to generate noisy atomic coordinates; and applying symmetric strain perturbation to the unit cell basis vector matrix of the optimal crystal structure to generate noisy unit cells. During the inference phase, the initial crystal structure to be optimized is input into the trained neural network model, iteratively optimized, and the optimized crystal structure is output. The inference phase includes: setting iteration parameters, including noise intensity, iteration step size, and maximum number of iteration steps; and performing multiple rounds of deterministic iterative updates on the initial crystal structure based on the iteration parameters. The process of performing multiple rounds of deterministic iterative updates on the initial crystal structure includes: in a single iteration, converting the current atomic fractional coordinates and the current unit cell basis matrix into atomic Cartesian coordinates; inputting the atomic type sequence, the current atomic fractional coordinates, the current unit cell basis matrix, and fixed small noise into the trained neural network model to obtain the current atomic displacement prediction value and the current unit cell strain prediction value; updating the atomic Cartesian coordinates according to the current atomic displacement prediction value and the iteration step size, and converting the updated atomic Cartesian coordinates back to atomic fractional coordinates; updating the current unit cell basis matrix according to the current unit cell strain prediction value. The formula for updating the atomic Cartesian coordinates is: ； ； in, Let Cartesian coordinates be the atomic coordinates of the current iteration step. The iteration step size parameter, These are predicted atomic shift values. For a fixed small noise standard deviation, This is the atomic fractional coordinate matrix. Let be the unit cell basis vector matrix of the current iteration step; The formula for updating the current unit cell basis vector matrix is: ； in, To obtain the symmetrical part, This represents the predicted value of the unit cell strain. The reasoning phase also includes: monitoring atomic displacement and unit cell change during the iteration process; terminating the iteration when the atomic displacement and unit cell change are lower than a preset threshold.

2. The method according to claim 1, characterized in that, The process of training the neural network model also includes: Calculate the coordinate denoising loss based on the predicted atomic displacement and the actual atomic displacement; The cell loss is calculated based on the predicted cell strain value and the actual cell parameters. The parameters of the neural network model are updated by combining the coordinate denoising loss and the cell loss.

3. The method according to claim 2, characterized in that, The formula for calculating the coordinate denoising loss is: ； in, These are predicted atomic shift values. For coordinate noise amplitude, These are actual atomic shifts; The formula for calculating cell loss is: ； in, The unit cell parameters are obtained from neural network prediction, representing the lengths of the three sides and the included angles between the triangles. These are the corresponding optimal unit cell parameters.

4. A computer-readable storage medium having a computer program stored thereon, characterized in that, When executed by a processor, the computer program implements the steps of the method according to any one of claims 1-3.

5. A computer program product, comprising a computer program, characterized in that, When executed by a processor, the computer program implements the steps of the method according to any one of claims 1-3.

Images