Charge density prediction method and device based on plane wave fix and equivariant neural network, and medium
By employing a hybrid architecture of plane wave-attached equivariant neural networks, combined with spherical harmonic functions and radial basis functions, high-precision charge density prediction for complex material systems is achieved. This solves the operational complexity and error problems of existing models in large-scale system calculations, and improves the model's generalization ability and computational efficiency.
Patent Information
- Application Number
- CN202511225993.7
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-08-29
- Publication Date
- 2025-12-05
AI Technical Summary
Existing charge density prediction models suffer from operational complexity, large result errors, and insufficient generalization ability in large-scale system calculations. They are particularly prone to significant errors in charge transfer excitation and long-range correlation modeling, making them difficult to adapt to complex material systems.
A hybrid architecture based on a plane wave-attached equivariant neural network is adopted. Atomic features are initialized by acquiring crystal structure parameters, spherical harmonic function expansion is performed by combining structural neighborhood edges to generate directional features, and equivariant neural network is generated by combining radial basis functions to perform message aggregation. A probe model is used for hierarchical interaction and updating to output target prediction results.
It improves the accuracy and efficiency of charge density prediction, reduces the dependence on the amount of training data, adapts to the calculation of ground state charge density of multi-scale material systems, and enhances the model's generalization ability and computational accuracy.
Smart Images
Figure CN121075486A_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of data processing technology, and in particular to a method, device and medium for predicting charge density based on a plane wave augmented equivariant neural network. Background Technology
[0002] Charge density, as one of the fundamental characteristics of electronic structure, plays a crucial role in atomic, molecular, solid, and even biological macromolecular systems. Charge density prediction, a core problem in quantum mechanics and computational materials science, traces its theoretical foundation back to the foundational research of density functional theory (DFT). The Hohenberg-Kohn theorem established the unique mapping of electron density to ground-state properties in 1964, while the Kohn-Sham equation further transformed the complexity of multi-electron systems into a solvable framework under the single-electron approximation. However, traditional DFT methods are limited by cubic computational complexity (O(N)). 3 This leads to significant performance bottlenecks in systems containing thousands of atoms (such as protein lysozymes or metal-organic frameworks), with computation times typically lasting weeks or even months in typical applications.
[0003] Charge density plays a crucial role in various fields, including chemistry, materials science, biochemistry, and drug design. With the development of computational materials science and artificial intelligence, charge density has gradually become an important input parameter in high-throughput material screening, machine learning model training, and first-principles simulations. Research shows that feature vectors constructed using charge density information can significantly improve the accuracy of prediction models and also help accelerate the calculation of electronic structures of large-scale systems. To overcome the efficiency limitations of traditional methods, researchers have turned to data-driven machine learning strategies, aiming to directly construct spatial mapping relationships of charge density using atomic coordinates and elemental information. For example, the Gaussian-2 network and SchNet architecture combined with graph neural network technology have achieved high-precision predictions of organic molecules and periodic crystal systems. These methods have shown application potential in enzyme catalytic reaction mechanism analysis and high-throughput material screening, and by coupling with self-consistent DFT iterations ("cooled DFT" technology), the convergence cycle of traditional methods has been significantly shortened.
[0004] Despite some progress in efficiency and accuracy, existing technologies still face the following technical challenges: In charge transfer excitation and long-range correlation modeling, existing models have significant errors in describing charge redistribution between molecules or at heterogeneous interfaces; In terms of generalization ability in heterogeneous systems, models trained on a single material class are difficult to transfer to metal-organic composite systems or multi-element systems; There are bottlenecks in ultra-large-scale system computation, and for systems with tens of billions of atoms (such as biomacromolecules or nanoporous materials), a performance breakthrough needs to be achieved by combining linear scaling algorithms with distributed computing architectures. Summary of the Invention
[0005] In order to overcome the above-mentioned technical defects, the purpose of this invention is to provide a charge density prediction method, device and medium based on a plane wave-added equivariant neural network, so as to solve the problems of complex operation and large error in the existing models used for charge density prediction.
[0006] This invention discloses a charge density prediction method based on a plane wave augmented equivariant neural network, comprising: Obtain the crystal structure parameters of the object to be predicted; Atomic features are initialized based on the crystal structure parameters to obtain basic atomic features; Construct a structural neighborhood edge for the predicted object, restricting the distance between all atoms to within the cutoff radius; Based on the structural domain edges, a spherical harmonic function expansion is performed to obtain directional features containing local symmetry; Radial basis functions are generated based on the aforementioned directional features to establish radial sensitivity features; Geometric feature fusion is performed based on the radial sensitivity and the basic atomic features to generate an equilateral edge message vector; Message aggregation is performed based on the equivariant edge message vector, and target atomic features are generated by gated residual update. Based on the target atomic features, hierarchical features are generated, and a probe model is used to perform hierarchical interaction and updates based on the hierarchical features to output the target prediction results.
[0007] Preferably, the atomic feature initialization includes: For each atom in the predicted object, a 118-dimensional embedding vector is encoded using a two-dimensional discrete function based on the atom number.
[0008] Preferably, constructing structural neighborhood edges for the predicted object, restricting the distance between all atoms to within the cutoff radius, includes: Based on atomic characteristics, the lattice basis vector matrix, fractional coordinates of each atom, and the number of periodic mirror displacements of each atom are obtained to calculate the relative position vector of any two atoms, thereby restricting the distance between all atoms to within the cutoff radius. The relative position vector between any two atoms is calculated according to the following formula: ;in For atoms arrive The relative position vector; For the lattice basis vector matrix; , For atoms , Fractional coordinates; For atoms The number of periodic mirror displacements.
[0009] Preferably, the step of performing spherical harmonic function expansion based on the structural neighborhood edges to obtain directional features containing local symmetry includes: The edge of the structural domain is expanded using a spherical harmonic function, converted into spherical coordinates, and direction-related features are generated. The directionally related features are split into atomic bond types to construct a directionally dependent domain mask matrix, and the spherical harmonic coefficients are superimposed and projected onto the angular vector space to obtain directional features containing local symmetry.
[0010] Preferably, the step of generating radial basis functions based on the directional features to establish radial sensitivity features includes: The radial basis function generation is performed according to the following formula: ;in Index of basis functions; The distance between atoms; Center of the Gaussian function; This is the width of the Gaussian function.
[0011] Preferably, the step of fusing geometric features based on the radial sensitivity and the basic atomic features to generate an equilateral edge message vector includes: Irreducible representation coupling is implemented through tensor product based on fundamental atomic features, spherical harmonic features, and radial basis functions, and equivariant edge message vectors are generated through Klebsch-Gordon coefficient constraints; wherein the radial basis functions are generated by the radial basis functions, and the spherical harmonic features are obtained by expanding the spherical harmonic functions.
[0012] Preferably, message aggregation is performed according to the following formula: ,in, For atoms The neighborhood group, For atoms The number of neighbors, For the first The edge message vector of the layer is obtained from the equivariant edge message vector.
[0013] Preferably, in the probe model, the features of surrounding atoms are aggregated layer by layer according to the hierarchical features, and updated through a gating mechanism, thereby combining each layer through a linear layer and correcting it through a KAN network to output the target prediction result.
[0014] The present invention also discloses a computer device, the computer device including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the steps of the above-described charge density prediction method.
[0015] The present invention also discloses a computer-readable storage medium having a computer program stored thereon, wherein the computer program, when executed by a processor, implements the steps of the above-described charge density prediction method.
[0016] Compared with existing technologies, the above technical solution has the following advantages: The charge density prediction method provided in this application combines a hybrid architecture of plane wave basis functions and equivariant neural networks to construct a multi-scale simulation framework for calculating the ground-state charge density of material systems. Based on high-precision first-principles calculations (plane wave basis functions), it constructs structural neighborhood edges and then performs spherical harmonic function expansion to obtain directional and radial features. Geometric features are then fused to adapt to the equivariant neural network. Finally, a probe model is used to save hierarchical features based on the equivariant edge message vectors obtained after network processing, enabling hierarchical interaction and updates. This reduces the model's dependence on the amount of training data and improves the accuracy of charge density prediction, solving the problems of complex operation and large error in existing models used for charge density prediction. Attached Figure Description
[0017] Figure 1 This is a flowchart illustrating an embodiment of the charge density prediction method, device, and medium based on a plane wave augmented equivariant neural network according to the present invention. Figure 2 This is a reference figure showing the predicted and actual charge density visualization results of BaTiO3 in an embodiment of the charge density prediction method, device and medium based on plane wave appended equivariant neural network described in this invention. Figure 3 This is a reference plot for evaluating the scatter plot level in an embodiment of the charge density prediction method, device and medium based on plane wave embellished equivariant neural network described in this invention; Figure 4 This is a reference diagram for evaluating gradient accuracy in an embodiment of the charge density prediction method, device, and medium based on a plane wave-added equivariant neural network described in this invention. Detailed Implementation
[0018] The advantages of the present invention will be further illustrated below with reference to the accompanying drawings and specific embodiments.
[0019] Exemplary embodiments will now be described in detail, examples of which are illustrated in the accompanying drawings. When the following description relates to the drawings, unless otherwise indicated, the same numerals in different drawings denote the same or similar elements. The embodiments described in the following exemplary embodiments do not represent all embodiments consistent with this disclosure. Rather, they are merely examples of apparatuses and methods consistent with some aspects of this disclosure as detailed in the appended claims.
[0020] The terminology used in this disclosure is for the purpose of describing particular embodiments only and is not intended to be limiting of the disclosure. The singular forms “a,” “the,” and “the” as used in this disclosure and the appended claims are also intended to include the plural forms unless the context clearly indicates otherwise. It should also be understood that the term “and / or” as used herein refers to and includes any and all possible combinations of one or more of the associated listed items.
[0021] It should be understood that although the terms first, second, third, etc., may be used in this disclosure to describe various information, such information should not be limited to these terms. Depending on the context, the word "if" as used herein can be interpreted as "when," "when," or "in response to determination."
[0022] In the following description, suffixes such as "module," "part," or "unit" used to denote elements are used only for the convenience of the description of the invention and have no specific meaning in themselves. Therefore, "module" and "part" can be used interchangeably.
[0023] Example: This example discloses a charge density prediction method based on a plane wave-attached equivariant neural network. As an illustration, the model on which the charge density prediction depends applies a plane wave-attached equivariant neural network, which is a hybrid architecture combining plane wave basis functions and equivariant neural networks, designed to efficiently handle physical systems with periodic structures (such as crystals and quasicrystals) or long-range interactions.
[0024] Specifically, it can be implemented using E3NN, an open-source library developed based on the PyTorch framework. E3NN focuses on building neural network models with E(3) equivariance, suitable for data modeling in three-dimensional Euclidean space. Its core design is based on Group Representation Theory, ensuring mathematical consistency of network outputs with geometric transformations of the input by explicitly encoding rotation, translation, and inversion symmetries. Specifically, the core computation module of E3NN includes the combined application of TensorProduct and Spherical Harmonics, achieving efficient operations on geometric tensors through irreducible representations. For example, in atomic-scale modeling, input features (such as occupancy matrices) can be decomposed into irreducible representations of O(3) groups, thus constructing a rotation-invariant feature space. Furthermore, the framework supports hierarchical reconstruction strategies, such as progressively predicting inter-atomic distance vectors through equivariant graph neural networks, achieving multi-scale modeling from coarse-grained to atomic-level structures.
[0025] Therefore, compared with other existing deep learning methods, E3NN reduces the model's dependence on the amount of training data through strict mathematical symmetry constraints, while improving the ability to interpret physical laws. This embodiment can rely on it to build a multi-scale simulation framework to achieve ground state charge density prediction in simple or complex material systems.
[0026] Specifically, for the prediction method provided in this embodiment, please refer to... Figure 1 ,include: S10: Obtain the crystal structure parameters of the object to be predicted; initialize atomic features based on the crystal structure parameters to obtain basic atomic features; Specifically, the atomic feature initialization includes: encoding each atom in the prediction object into a 118-dimensional embedding vector using a two-dimensional discrete function based on its atomic number. This can be represented as: ,in For atoms The atomic number (1-118). It is the Kronecker delta function, a two-dimensional discrete function, when... It is 1 if it is true, otherwise it is 0.
[0027] In this embodiment, the object of prediction can be a single system, a composite system, or a multi-element system of three-dimensional materials. In this embodiment, the BaTiO3 perovskite structure is used as an example: Input the BaTiO3 perovskite structure from the MP database. The structure information is cubic crystal system, space group Pm-3m (#221), lattice constant a = 3.99 Å, Ba 2+ Located at vertex (0,0,0), Ti 4+ Located at the body center (0.5a, 0.5a, 0.5a), O 2- Located at each face center (0.5a, 0.5a, 0, etc.). Ba 2+ Equivalent to twelve Os 2- Atomic bonding forms BaO 12 The dodecahedron, which is equivalent to twelve other BaO atoms. 12 The dodecahedron shares a vertex and is associated with six equivalent BaO vertices. 12 The dodecahedrons share a common face, and it also shares a face with eight equivalent TiO6 octahedra. All Ba-O bond lengths are 2.83 Å. Ti 4+ Equivalent to six O 2- Atomic bonds form a TiO6 octahedron, which shares vertices with six other equivalent TiO6 octahedrons and with eight equivalent BaO atoms. 12 The dodecahedrons share a common face. These octahedrons with shared vertices are not tilted. All Ti-O bond lengths are 2.00 Å. 2- It exhibits a distorted linear geometric configuration, with four equivalent Ba... 2+ Atoms and two equivalent Ti 4+ Atomic bonding. Align atomic features for initialization, and encode the type of each atom as a 118-dimensional one-hot vector (anchored to the atomic number corresponding to the outer electron configuration), resulting in an 118-dimensional embedding vector representation of Ba (56-dimensional) / Ti (22-dimensional) / O (8-dimensional), which is the basic atomic feature mentioned above.
[0028] In practice, it is understandable that these basic atomic features can provide the atomic-related parameters required for geometric feature fusion, such as constructing structural domain edges.
[0029] S20: Construct a structural domain edge for the predicted object, restricting the distance between all atoms to within the cutoff radius; Specifically, constructing structural neighborhood edges for the predicted object, and restricting the distances between all atoms to within the cutoff radius, includes: Based on atomic characteristics, the lattice basis vector matrix, fractional coordinates of each atom, and the number of periodic mirror displacements of each atom are obtained to calculate the relative position vector of any two atoms, thereby restricting the distance between all atoms to within the cutoff radius. The relative position vector between any two atoms is calculated according to the following formula: ;in For atoms arrive The relative position vector; For the lattice basis vector matrix; , For atoms , Fractional coordinates; For atoms The number of periodic mirror displacements.
[0030] Understandably, the purpose of constructing structural domain edges is to effectively divide the computational region, simplify problem complexity, and improve computational efficiency.
[0031] The above BaTiO3 perovskite structure is constructed by building the structural neighborhood edges of the S10 input. The relative position vectors after considering periodic mirroring are calculated to ensure that the distance between all atoms is within the cutoff radius (specifically, the cutoff radius is set to 5.0 Å to eliminate PBC boundary effects). After complete screening, duplicate edges (identification of mirror equivalent bonds) and edges with abnormal displacement amplitudes (>5.5 Å) are removed, resulting in an edge feature set G(E) = (V, E) containing three chemical bond types (Ba-O, Ti-O, OO) and corresponding lattice period extensions.
[0032] S30: Perform spherical harmonic function expansion based on the edge of the structural domain to obtain directional features containing local symmetry; Specifically, the step of performing spherical harmonic function expansion based on the structural neighborhood edges to obtain directional features containing local symmetry includes: performing spherical harmonic function expansion on the structural neighborhood edges, converting them into spherical coordinates to generate directional features; splitting the directional features into atomic bond types to construct a directional dependency domain mask matrix, superimposing spherical harmonic coefficients and projecting them onto the angular vector space to obtain directional features containing local symmetry.
[0033] The expansion of the spherical harmonic function can be seen in the following formula: ;in It is the quantum number of angular momentum. It is the magnetic quantum number; This is associated Legendre polynomials. The edge vectors are converted to spherical coordinates, and the values of the spherical harmonics of each order are calculated to generate the corresponding features in the aforementioned directions.
[0034] For the BaTiO3 perovskite structure mentioned above, based on the results of S20 above, spherical harmonic function expansion is performed, the edge vectors are converted into spherical coordinates, the values of spherical harmonic functions of each order are calculated, direction-related features are generated, and the features are split according to atomic bond type to construct a direction-dependent domain mask matrix (specifically: Ti-O bond activates l=4 features, OO bond only retains l=0), and the spherical harmonic coefficients are superimposed and projected onto the angular vector space to form a directional representation containing local symmetry.
[0035] S40: Generate radial basis functions based on the directional features to establish radial sensitivity features; Specifically, the step of generating radial basis functions based on the directional features to establish radial sensitivity features includes: The radial basis function generation is performed according to the following formula: ;in Index of basis functions; The distance between atoms; Center of the Gaussian function; This is the width of the Gaussian function.
[0036] The above-mentioned S30 determines the symmetrical directional characteristics, and S40 determines the radial characteristics; For the BaTiO3 perovskite structure described above, the output of S30 can be used to generate radial basis functions, resulting in 20 different... Gaussian function with spacing of 0.1 to 5.0 Å , =2.5, truncate the smooth edges of the sigmoid decay function and perform Z-score normalization to ensure that these basis functions have zero mean and unit variance, and establish the radial sensitivity matrix R(r).
[0037] S50: Based on the radial sensitivity and the basic atomic features, perform geometric feature fusion to generate an equilateral edge message vector; Specifically, the step of fusing geometric features based on the radial sensitivity and the basic atomic features to generate an equilateral edge message vector includes: implementing irreducible representation coupling based on the basic atomic features, spherical harmonic features, and radial basis through tensor product ((atomic features × spherical harmonic features × radial basis), and generating an equilateral edge message vector through Klebsch-Gordon coefficient constraints; wherein the radial basis is obtained by generating the radial basis function, and the spherical harmonic features are obtained by expanding the spherical harmonic function.
[0038] The above geometric features are fused using the following formula: ;in These are the Clebsch-Gordan coefficients; Weights generated for the radial network; For atoms The Among the layer features, belonging to The value of the component.
[0039] For the BaTiO3 perovskite structure mentioned above, the results of S40 are geometrically fused, and irreducible representations of L=0 to 4th order coupling are implemented through tensor product. The variable edge message vectors such as SO(3) are generated by Clebsch-Gordan coefficient constraints.
[0040] S60: Message aggregation is performed based on the equivariant edge message vector, and target atomic features are generated by gated residual update; Specifically, message aggregation is performed according to the following formula: ,in, For atoms The neighborhood group, For atoms The number of neighbors, For the first The edge message vector of the layer is obtained from the equivariant edge message vector.
[0041] Perform gated residual updates based on the following expression: ,in For the first Layer Atom eigenvectors, For the first The message vector after layer aggregation For dynamic rotation angle, Generate weights for the rotation angle. For gated linear transformation weights, This is the bias for the gated linear transformation. In special cases... When processing scalar features ( When ), the spherical harmonic function degenerates into a constant: At this point, the tensor product simplifies to scalar multiplication: After aggregation, updates are directly performed through the fully connected layer, without the need for direction-dependent processing.
[0042] Therefore, the BaTiO3 perovskite structure described above is used for message aggregation based on the S50 results, and based on attention weights. For each atom, truncate the message vector within its neighborhood by 5.0 Å. Vector addition is performed, and the dimension is raised to the spherical tensor domain through τ-lifting operation for further gated residual updates.
[0043] Further gated residual update: Updating atomic features via a bilinear gate function (LSTM variant): Residual connections are preserved to ensure smooth information transfer, especially for orbitals containing d orbitals (Ti). 4+ ) and inert gas closed shell (Ba 2+ / O 2- In the case of spin-orbit coupling and LDA+U initial state, the two algorithms (spin-orbit coupling and LDA+U initial state) will automatically classify the relevant scenarios and call Wannier90 (localized orbit and tight-binding model) or Hubbard U adaptive solver (strong correlation correction) to automatically process and obtain the results.
[0044] S70: Generate hierarchical features based on the target atomic features, and use the probe model to perform hierarchical interaction and updates based on the hierarchical features to output the target prediction result.
[0045] Specifically, for the atomic features updated by the gated residuals, i.e., the results of S60 above, hierarchical features are saved, as shown in the following formula: ;in For the first Atomic features output by the layer The feature dimensions are dynamically determined.
[0046] Specifically, in the probe model, the features of surrounding atoms are aggregated layer by layer according to the hierarchical features, updated through a gating mechanism, thereby combining each layer through a linear layer and correcting it through a KAN network to output the target prediction result.
[0047] Specifically, the probe hierarchy interaction reference in the probe model: ;in probe Charge density at that location, The hierarchy weights are implicit in the parameters of the Linear layer. For linear transformation weights, For the first Layer probe eigenvectors, Corrections generated for the Kolmogorov-Arnold network For a multilayer perceptron, features are transferred from... Mapped to 24 dimensions, for The weight, These are the learnable basis functions of KAN.
[0048] The features of each layer of the probe are linearly combined to generate the main predicted charge density: Residual correction is achieved by generating correction terms through a KAN network to enhance local accuracy. The final prediction result is obtained as follows: .
[0049] Therefore, it can be understood that in the model implementing S10-S70 described above, this embodiment is as follows: Changes in atomic characteristics: ; Changes in predicted charge density: Based on the above, in this embodiment, the example BaTiO3 perovskite structure is used to save the S60 results hierarchically, and the output features of each convolutional layer are saved to a list for use by the probe model layer by layer; the results of the hierarchical feature saving are used for probe hierarchical interaction, for each probe... The probe aggregates features from surrounding atoms layer by layer and updates these features through a gating mechanism; multi-level prediction output is then performed, autonomously combining the two regions for linear readout. The features of each layer of the probe are combined through linear layers to generate the master predicted charge density. Residual correction is then performed, i.e., the correction term is generated by the aforementioned KAN network. Finally, the predicted and true charge densities of BaTiO3 are visualized as follows: Figure 2 (in Figure 2 (a) represents the true structure. Figure 2 (b) is the prediction structure corresponding to the target prediction results; for evaluating the scatter plot level, please refer to [reference needed]. Figure 3 For an assessment of gradient accuracy, please refer to [link / reference]. Figure 4 .
[0050] Understandably, based on the above steps, this embodiment provides a charge density prediction method based on a plane wave-attached equivariant neural network, which effectively improves prediction accuracy and realizes the calculation of ground state charge density of simple / complex material systems by constructing a multi-scale simulation framework.
[0051] It is worth noting that the dataset is derived from high-precision first-principles calculations, ensuring the reliability of the training data. The material structure information is standardized (i.e., the atomic feature initialization mentioned above) to adapt to the input requirements of the neural network (equivariant neural network).
[0052] Experimental verification showed that the prediction method achieved a coefficient of determination of 0.99 for the predicted charge density in a three-dimensional material system. Figure 3 The coefficient of determination for the gradient is 0.97 (e.g., Figure 4 This demonstrates its extremely high accuracy and reliability in predicting charge density and its gradient.
[0053] This embodiment relates to the field of large-scale system charge density prediction technology. Specifically, it provides a charge density prediction method based on first-principles plane wave ultrasoft pseudopotential technology. By constructing a multi-scale simulation framework / model to perform the above steps, the ground state charge density of different material systems can be calculated.
[0054] Based on the above, the model implemented in this embodiment can not only serve as a new and efficient prediction tool in multiple fields such as materials science and chemistry, but also provide strong support for exploring the development of new functional materials. It can be applied in different fields and has high reliability and accuracy.
[0055] The above steps can be implemented by multiple modules / units, or integrated into the same module / unit.
[0056] To achieve the above objectives, this embodiment also provides a computer device, which may be one or more. The modules / components controlling the implementation of the above steps may be integrated or distributed in different computer devices. The computer device may be a laptop computer, desktop computer, large data server, etc., that executes programs. The computer device in this embodiment includes, but is not limited to, a memory and a processor that can be interconnected via a system bus.
[0057] In this embodiment, the memory may include a program storage area and a data storage area. The program storage area may store the program required by the operating system and at least one function; the data storage area may store user operation data on the computer device, such as the crystal structure parameters mentioned above or the results of each processing step. Furthermore, the memory may include high-speed random access memory, and may also include non-volatile memory, etc.
[0058] The processor can be a central processing unit, a controller, or other data processing chip. This processor is typically used to control the overall operation of a computer device. In this embodiment, the processor is used to run program code stored in memory or process data to implement the above-described method.
[0059] To achieve the above objectives, the present invention also provides a computer-readable storage medium, which may include multiple storage media, such as read-only memory, electrically erasable programmable read-only memory, programmable read-only memory, etc., on which a computer program is stored, and the program, when executed by a processor, performs a corresponding function. In this embodiment, the data stored in the computer-readable storage medium is executed by a processor to implement the above method.
[0060] It should be noted that the embodiments of the present invention have better implementability and are not intended to limit the present invention in any way. Any person skilled in the art may use the above-disclosed technical content to change or modify it into equivalent effective embodiments. However, any modifications or equivalent changes and modifications made to the above embodiments based on the technical essence of the present invention without departing from the content of the technical solution of the present invention shall still fall within the scope of the technical solution of the present invention.
Claims
1. A charge density prediction method based on a plane wave augmented equivariant neural network, characterized in that, The method comprises the following steps: obtaining crystal structure parameters of a prediction object; initializing atomic features according to the crystal structure parameters to obtain basic atomic features; constructing structure field edges of the prediction object to limit the distance between all atoms within a cutoff radius; performing spherical harmonic function expansion based on the structure field edges to obtain directional features containing local symmetry; generating radial basis functions based on the directional features to establish radial sensitivity features; fusing geometric features based on the radial sensitivity and the basic atomic features to generate isometric edge message vectors; performing message aggregation based on the isometric edge message vectors to generate target atomic features through gated residual update; generating hierarchical features based on the target atomic features, and performing hierarchical interaction and update based on the hierarchical features by using a probe model to output a target prediction result.
2. The charge density prediction method of claim 1, wherein, The initializing atomic features comprises the following steps: for each atom in the prediction object, a two-dimensional discrete function is used to encode the atomic number into a 118-dimensional embedding vector.
3. The charge density prediction method of claim 1, wherein, The constructing structure field edges of the prediction object to limit the distance between all atoms within a cutoff radius comprises the following steps: obtaining a lattice basis vector matrix, fractional coordinates of each atom, and periodic mirror displacement times of each atom according to atomic features to calculate the relative position vectors of any two atoms to limit the distance between all atoms within a cutoff radius; The relative position vectors of any two atoms are calculated according to the following formula: ; wherein is the atomic to the relative position vector of is the lattice basis matrix; , is the atomic , fractional coordinates of is the atomic periodic image displacement number of 4. The charge density prediction method of claim 1, wherein, The performing spherical harmonic function expansion based on the structure field edges to obtain directional features containing local symmetry comprises the following steps: The spherical harmonic function expansion is converted into spherical coordinates to generate direction-dependent features; The direction-dependent features are split according to atomic bond types to construct a direction-dependent domain mask matrix, and the spherical harmonic coefficient projection is superimposed into an angular vector space to obtain directional features containing local symmetry.
5. The charge density prediction method of claim 1, wherein, The generating radial basis functions based on the directional features to establish radial sensitivity features comprises the following steps: The radial basis function generation is performed according to the following formula: ; where is a basis function index; is an interatomic distance; is a Gaussian function center; is a Gaussian function width.
6. The charge density prediction method of claim 1, wherein, The fusing geometric features based on the radial sensitivity and the basic atomic features to generate isometric edge message vectors comprises the following steps: Irreducible representation coupling is implemented by tensor product according to the basic atomic features, spherical harmonic features, and radial basis, and isometric edge message vectors are generated by Kleibsh-Gordon coefficient constraint; wherein the radial basis is obtained by the radial basis function generation, and the spherical harmonic features are obtained by the spherical harmonic function expansion.
7. The charge density prediction method according to claim 1, wherein: message aggregation is performed according to the following formula: wherein, is a set of neighbors of atom , is a number of neighbors of atom , is an edge message vector of the th layer, obtained from the equivariant edge message vector.
8. The charge density prediction method according to claim 1, wherein: In the probe model, the features of surrounding atoms are aggregated layer by layer according to hierarchical features, and are updated through a gating mechanism, so that each layer is combined through a linear layer, and is corrected through a KAN network to output a target prediction result.
9. A computer device, comprising: The computer device comprises a memory, a processor, and a computer program stored on the memory and executable on the processor, and the processor executes the computer program to implement the steps of the charge density prediction method according to any one of claims 1 to 8.
10. A computer-readable storage medium having stored thereon a computer program, characterized in that, The computer program, when executed by a processor, implements the steps of the charge density prediction method of any one of claims 1 to 8.