{'fraction_non_alphanumeric': 0.02328702241869277, 'fraction_numerical': 0.010341016735080517, 'mean_word_length': 5.061722488038278, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Transferable E(3) equivariant parameterization for Hamiltonian of molecules and solids

Machine learning, especially deep learning, can build a direct mapping from structure to properties with its huge parameter space, making it possible to perform high-throughput screening for the desired properties of materials. However, since the electronic Hamiltonian transforms non-trivially under rotation operations, it is challenging to accurately predict the electronic Hamiltonian while strictly satisfying this constraint. There is currently a lack of transferable machine learning models that can bypass the computationally demanding density functional theory (DFT) to obtain the ab initio Hamiltonian of molecules and materials by complete data-driven methods. In this work, we point out the necessity of explicitly considering the parity symmetry of the electronic Hamiltonian in addition to rotational equivariance. We propose a parameterized Hamiltonian that strictly satisfies rotational equivariance and parity symmetry simultaneously, based on which we develop an E(3) equivariant neural network called HamNet to predict the ab initio tight-binding Hamiltonian of various molecules and solids. The tests show that this model has similar transferability to that of machine learning potentials and can be applied to a class of materials with different configurations using the same set of trained network weights. The proposed framework provides a general transferable model for accelerating electronic structure calculations.


Introduction
Nowadays machine learning (ML) has a wide range of applications in molecular and materials science, including the direct prediction of various properties of materials [1][2][3] , the construction of neural network (NN) potentials with quantum mechanical precision [4][5][6][7][8] , the high-throughput generation of molecular and crystal structures [9][10][11] , and the construction of more precise exchange-correlation functionals 12,13 . However, the acquisition of the electronic structure of materials still relies almost exclusively on density functional theory (DFT) based calculations.
Unfortunately, these methods are very time-consuming to get the Hamiltonian of the systems through self-consistent iterations and scale poorly with system size. Semi-empirical tightbinding (TB) approximations 14 , such as the Slater-Koster method 15 , can reduce a lot of computation compared to the ab initio DFT methods. However, this approach often directly uses the existing or manually fine-tuned TB parameters and thus cannot accurately reproduce the electronic structure of general systems. Machine learning methods have powerful fitting and inference capabilities, and can accurately fit the ab initio tight-binding Hamiltonian of the system by a fully data-driven method. Fitting the ab initio tight-binding Hamiltonian with ML can reconcile the accuracy with speed and can be used to accelerate large-scale electronic structure calculations in the same way as machine learning potentials in molecular dynamics (MD).
Hegde and Brown first used kernel ridge regression (KRR) to learn semi-empirical tightbinding Hamiltonian matrices 16 . They successfully fitted the Hamiltonian of the Cu system containing only rotation invariant s orbitals and the diamond system consisting only of s and p orbitals. Similarly, Wang et al. designed a neural network model to obtain semi-empirical TB parameters by fitting the ab initio band structure 17 . An important feature of the Hamiltonian matrix is that its components transform equivariantly with the rotation of the coordinate system. However, neither of these approaches deals with the rotational equivariance of the Hamiltonian.
Moreover, the two methods can only fit the empirical model Hamiltonian rather than the true ab initio tight-binding Hamiltonian generated by the self-consistent iteration of ab initio tightbinding methods based on numerical atomic orbitals such as OpenMX 18,19 and Siesta 20, 21 .
Zhang et al. 22 and Nigam et al. 23 proposed a method to predict the ab-initial TB Hamiltonian of small molecules and simple solid systems by constructing an equivariant kernel (kernel) in Gaussian Process Regression (GPR) to parameterize the Hamiltonian. Since GPR uses a fixed kernel and representation [24][25][26][27] , its training accuracy and multi-element generalization ability are usually lower than those of deep neural networks such as the message passing neural networks (MPNNs) [28][29][30][31][32][33] when the number of training samples is sufficient. Therefore, developing graph neural networks (GNNs) capable of predicting the Hamiltonian of general periodic and aperiodic systems would be the best option.
However, traditional GNNs can only predict rotation-invariant scalars such as energy, band gap, etc. GNNs must encode the directional information of the system in an appropriate way to predict equivariant directional properties. In order to make the predicted Hamiltonian satisfy the rotational equivariance, Schütt et al. designed the SchNorb neural network architecture by embedding the direction information of the bonds into the message passing function 34 . This network constructs the ab initio Hamiltonian of molecules from the directional edge features of atom pairs. However, SchNorb needs to learn the rotational equivariance of the Hamiltonian through data augmentation, which greatly increases the amounts of training data and redundant parameters of the network. Unke et al proposed the PhisNet model 35 , which realized the SE (3) equivariant parameterization of the Hamiltonian matrices with GNN based on SO (3) representations and achieved state-of-the-art accuracy on the Hamiltonian of small molecules such as water and ethanol. However, it should be noted that the PhisNet model is not the most universal representation of the Hamiltonian as it ignores the parity of the Hamiltonian, which may present serious problems when predicting periodical systems with infinite sizes. The limitations of SE(3) representation will be discussed in detail later.
Recently, Li et al. proposed a GNN model called DeepH to predict the ab initio Hamiltonian by constructing a local coordinate system in a crystal 36 . DeepH successfully predicted the tightbinding Hamiltonian of some simple periodic systems such as graphene, carbon nanotubes, etc.
Their original intention of introducing the local coordinate system is to solve the rotation equivariance problem of the Hamiltonian, but DeepH still embeds the local directional information of interacting atom pairs in the invariant message passing function, which will undoubtedly increase the number of redundant parameters of the network but may require less data augmentation than a fully invariant model without local coordinate systems. In addition, the hopping distance between two interacting atoms far exceeds the lengths of the general chemical bonds. Taking the smallest hydrogen atom as an example, the cutoff radius of the numerical atomic orbital of the hydrogen atom used by OpenMX is 6 Bohr, so the furthest hopping between any two atomic bases used by OpenMX in periodic systems can exceed at least 12 Bohr (~6.4 Å), a distance that even exceeds the lattice parameters of some crystals. It is difficult to describe such long hopping in a well-defined local coordinate system.

E(3) equivariant parametrized Hamiltonian
The core of the electronic structure problem in DFT is to solve the Kohn-Sham equation for According to the angular momentum theory, is a reducible representation, which can be further decomposed into the direct sum of several irreducible Wigner D matrices: 8 Combining the parity of the Hamiltonian matrix block   , i i j j n l n l , we can get According to Eq. (7), , , The prediction of the Hamiltonian is transformed into the prediction of on i Ω and off ij Ω , which can be obtained by mapping from the equivariant features of the nodes and the pair interactions, respectively. The final parameterized Hamiltonian can be expressed as: 3 , , , The above formula is O(3) equivariant. Since GNN naturally has translational symmetry, the parameterized Hamiltonian represented by equivariance.

Network Implementation
As can be seen from the above discussion, the Hamiltonian matrix block satisfies the rotation equivariance and has a definite parity under the inversion operation, so we designed E  The distance between atom i and its neighboring atom j within the cutoff radius r c is expanded using the Bessel basis function: where c f is the cosine cutoff function, which guarantees physical continuity as atoms enter and exit the cutoff sphere. The directional information between atom i and atom j is embedded in a set of real spherical harmonics which is used to construct the rotationequivariant filter 37 in the equivariant message passing.

Limitations of SO(3) equivariant parametrized Hamiltonian
Hamiltonian on the test dataset of graphene structures. It can be seen from Fig. 2(b) that some of the Hamiltonian matrix elements predicted by PhiSNet are incorrect and are always approximately zero, thus a vertical band appears at the zero point of Fig. 2(b).
Since the number of neighbors of each atom in a small molecule is at most the total number of atoms − 1, the number of hopping terms for small molecules to be fitted is small. The successful prediction of PhiSNet on the Hamiltonian of small molecules shows that it can adjust the network weights by training to make its output parameterized Hamiltonian have a pseudo parity, but this cannot ensure the strict parity symmetry of the Hamiltonian matrix intrinsically.
In periodic crystal structures, the number of Hamiltonian matrices between each atom and its surrounding neighbors is much greater than that in small molecules.

Tests and applications
To test the accuracy and transferability of HamNet on periodic and aperiodic systems, we

Carbon allotropes.
We collected 426 structures with less than 60 atoms in the primitive cell from Samara Carbon Allotrope Database (SACADA) 41 Fig. 5(g). The structure of the Si allotrope labeled mp-1244971 is shown in Fig. 5(b). This structure is in an amorphous glassy state and quite different from the common four-coordinated Si allotrope. Each Si atom in this structure is not always tetrahedral coordinated but forms an irregular polyhedron with 3−5 coordinating atoms, so the prediction difficulty for this Si structure is expected to be relatively high. However, HamNet still accurately predicted its Hamiltonian matrix and reached an MAE  shows that it can be used to bypass the expensive DFT calculation to study the electronic structure of large amorphous SiO 2 glass.   Table 1. A two-layer MLP with 64 neurons is used to map the invariant edge embeddings to the weights of each tensor product path in Eqs. (14) and (16). Shifted softplus 30 function is used as the activation function in the MLP.

Discussion
To avoid overfitting, we include the error of the calculated energy bands as a regularization term in the loss function: where the variables marked with a tilde refer to the corresponding predictions and λ denotes the loss weight of the band energy error. When the training of the network has not converged, the error of the predicted Hamiltonian is large, resulting in poor or even divergent prediction values of the energy bands. Adding the band loss value at the beginning of training may cause the total loss value to diverge. Therefore, we train the network in two steps. First, only the mean absolute error of Hamiltonians is used as the loss value to train the network until the network weights converge. The parameters were optimized with AdamW 48, 49 optimizer using an initial learning rate of 10 −3 . Then the mean absolute error of each band calculated at N k random points in the reciprocal space is added to the loss function and starts the training at an initial learning rate of 10 −4 . When the accuracy of the model on the validation set is not improved after successive N patience epochs, the learning rate will be reduced by a factor of 0.5. When the accuracy of the model on the validation set is not improved after successive N stop epochs or the learning rate is lower than 10 −6 , the training will be stopped and the model that has the best accuracy on the validation set will be used on the test set. The values of some key network and training parameters on each dataset are listed in Table 1. All models were trained on a single NVIDIA A100 GPU.

Data and Code availability
The data and code that support the results of this work will be available online after the manuscript is accepted.