Package: chgnet
CHGNetA pretrained universal neural network potential forcharge-informed atomistic modeling (see publication)CrystalHamiltonianGraph neuralNetwork is pretrained on the GGA/GGA+U static and relaxation trajectories from Materials Project,
a comprehensive dataset consisting of more than 1.5 Million structures from 146k compounds spanning the whole periodic table.CHGNet highlights its ability to study electron interactions and charge distribution
in atomistic modeling with near DFT accuracy. The charge inference is realized by regularizing the atom features with
DFT magnetic moments, which carry rich information about both local ionic environments and charge distribution.Pretrained CHGNet achieves excellent performance on materials stability prediction from unrelaxed structures according toMatbench Discovery[repo].Example notebooksNotebooksGoogle ColabDescriptionsCHGNet BasicsExamples for loading pre-trained CHGNet, predicting energy, force, stress, magmom as well as running structure optimization and MD.Tuning CHGNetExamples of fine tuning the pretrained CHGNet to your system of interest.Visualize RelaxationCrystal Toolkit app that visualizes convergence of atom positions, energies and forces of a structure during CHGNet relaxation.Phonon DOS + BandsUse CHGNet with theatomate2phonon workflowbased on finite displacements as implemented in Phonopy to calculate phonon density of states and band structure forSi(mp-149).Elastic tensor + bulk/shear modulusUse CHGNet with theatomate2elastic workflowbased on a stress-strain approach to calculate elastic tensor and derived bulk and shear modulus forSi(mp-149).Installationpipinstallchgnetif PyPI installation fails or you need the latestmainbranch commits, you can install from source:pipinstallgit+https://github.com/CederGroupHub/chgnetTutorials and DocsSee thesciML webinar tutorialon 2023-11-02 andAPI docs.UsageDirect Inference (Static Calculation)PretrainedCHGNetcan predict the energy (eV/atom), force (eV/A), stress (GPa) and
magmom ($\mu_B$) of a given structure.fromchgnet.model.modelimportCHGNetfrompymatgen.coreimportStructurechgnet=CHGNet.load()structure=Structure.from_file('examples/mp-18767-LiMnO2.cif')prediction=chgnet.predict_structure(structure)forkey,unitin[("energy","eV/atom"),("forces","eV/A"),("stress","GPa"),("magmom","mu_B"),]:print(f"CHGNet-predicted{key}({unit}):\n{prediction[key[0]]}\n")Molecular DynamicsCharge-informed molecular dynamics can be simulated with pretrainedCHGNetthroughASEpython interface (see below),
or throughLAMMPS.fromchgnet.model.modelimportCHGNetfromchgnet.model.dynamicsimportMolecularDynamicsfrompymatgen.coreimportStructureimportwarningswarnings.filterwarnings("ignore",module="pymatgen")warnings.filterwarnings("ignore",module="ase")structure=Structure.from_file("examples/mp-18767-LiMnO2.cif")chgnet=CHGNet.load()md=MolecularDynamics(atoms=structure,model=chgnet,ensemble="nvt",temperature=1000,# in Ktimestep=2,# in femto-secondstrajectory="md_out.traj",logfile="md_out.log",loginterval=100,)md.run(50)# run a 0.1 ps MD simulationThe MD defaults to CUDA if available, to manually set device to cpu or mps:MolecularDynamics(use_device='cpu').MD outputs are saved to the ASE trajectory file, to visualize the MD trajectory
and magnetic moments after the MD run:fromase.io.trajectoryimportTrajectoryfrompymatgen.io.aseimportAseAtomsAdaptorfromchgnet.utilsimportsolve_charge_by_magtraj=Trajectory("md_out.traj")mag=traj[-1].get_magnetic_moments()# get the non-charge-decorated structurestructure=AseAtomsAdaptor.get_structure(traj[-1])print(structure)# get the charge-decorated structurestruct_with_chg=solve_charge_by_mag(structure)print(struct_with_chg)To manipulate the MD trajectory, convert to other data formats, calculate mean square displacement, etc,
please refer toASE trajectory documentation.Structure OptimizationCHGNetcan perform fast structure optimization and provide site-wise magnetic moments. This makes it ideal for pre-relaxation andMAGMOMinitialization in spin-polarized DFT.fromchgnet.modelimportStructOptimizerrelaxer=StructOptimizer()result=relaxer.relax(structure)print("CHGNet relaxed structure",result["final_structure"])print("relaxed total energy in eV:",result['trajectory'].energies[-1])Available WeightsCHGNet 0.3.0 is released with new pretrained weights! (release date: 10/22/23)CHGNet.load()now loads0.3.0by default,
previous0.2.0version can be loaded withCHGNet.load('0.2.0')CHGNet_0.3.0CHGNet_0.2.0Model Training / Fine-tuneFine-tuning will help achieve better accuracy if a high-precision study is desired. To train/tune aCHGNet, you need to define your data in a
pytorchDatasetobject. The example datasets are provided indata/dataset.pyfromchgnet.data.datasetimportStructureData,get_train_val_test_loaderfromchgnet.trainerimportTrainerdataset=StructureData(structures=list_of_structures,energies=list_of_energies,forces=list_of_forces,stresses=list_of_stresses,magmoms=list_of_magmoms,)train_loader,val_loader,test_loader=get_train_val_test_loader(dataset,batch_size=32,train_ratio=0.9,val_ratio=0.05)trainer=Trainer(model=chgnet,targets="efsm",optimizer="Adam",criterion="MSE",learning_rate=1e-2,epochs=50,use_device="cuda",)trainer.train(train_loader,val_loader,test_loader)Notes for TrainingCheckfine-tuning example notebookThe target quantity used for training should be energy/atom (not total energy) if you're fine-tuning the pretrainedCHGNet.The pretrained dataset ofCHGNetcomes from GGA+U DFT withMaterialsProject2020Compatibilitycorrections applied.
The parameter for VASP is described inMPRelaxSet.
If you're fine-tuning withMPRelaxSet, it is recommended to apply theMP2020compatibility to your energy labels so that they're consistent with the pretrained dataset.If you're fine-tuning to functionals other than GGA, we recommend you refit theAtomRef.CHGNetstress is in units of GPa, and the unit conversion has already been included indataset.py. SoVASPstress can be directly fed toStructureDataTo save time from graph conversion step for each training, we recommend you useGraphDatadefined indataset.py, which reads graphs directly from saved directory. To create saved graphs,
seeexamples/make_graphs.py.MPtrj DatasetThe Materials Project trajectory (MPtrj) dataset used to pretrain CHGNet is available atfigshare.The MPtrj dataset consists of all the GGA/GGA+U DFT calculations from the September 2022Materials Project.
By using the MPtrj dataset, users agree to abide theMaterials Project terms of use.ReferenceIf you use CHGNet or MPtrj dataset, please citethis paper:@article{deng_2023_chgnet,title={CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling},DOI={10.1038/s42256-023-00716-3},journal={Nature Machine Intelligence},author={Deng, Bowen and Zhong, Peichen and Jun, KyuJung and Riebesell, Janosh and Han, Kevin and Bartel, Christopher J. and Ceder, Gerbrand},year={2023},pages={1–11}}Development & BugsCHGNetis under active development, if you encounter any bugs in installation and usage,
please open anissue. We appreciate your contributions!