Skip to content

Commit

Permalink
Fix bug in build_dataset.from_model. And update e3 example
Browse files Browse the repository at this point in the history
  • Loading branch information
@QG-phy
QG-phy committed Aug 14, 2024
1 parent 2bb0bb5 commit 5dfced2
Show file tree
Hide file tree
Showing 5 changed files with 140 additions and 32 deletions.
7 changes: 5 additions & 2 deletions dptb/data/build.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from dptb.utils.tools import j_loader
from dptb.utils.argcheck import normalize_setinfo, normalize_lmdbsetinfo
from dptb.utils.argcheck import collect_cutoffs
from dptb.utils.argcheck import get_cutoffs_from_model_options
import logging
import torch
import copy
Expand Down Expand Up @@ -302,8 +303,10 @@ def from_model(self,
Returns:
dataset: The built dataset.
"""
cutoff_options = collect_cutoffs(model.model_options)

# cutoff_options = collect_cutoffs(model.model_options)
r_max, er_max, oer_max = get_cutoffs_from_model_options(model.model_options)
cutoff_options = {'r_max': r_max, 'er_max': er_max, 'oer_max': oer_max}

dataset = self(
root = root,
**cutoff_options,
Expand Down
34 changes: 8 additions & 26 deletions examples/e3/band_plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,44 +3,26 @@
from dptb.nn.build import build_model
from dptb.data import build_dataset

model = build_model(checkpoint="./e3_silicon/checkpoint/nnenv.best.pth") # buld the trained e3 hamiltonian and overlap model
# buld the trained e3 hamiltonian and overlap model
model = build_model(checkpoint="./e3_silicon/checkpoint/nnenv.best.pth")

dataset = build_dataset(
# build the dataset from the model
dataset = build_dataset.from_model(
model=model,
root="./data",
prefix="Si64"
) # build the dataset
)

jdata = j_loader("./band.json")
kpath_kwargs = jdata["task_options"] # read the kpath information from the json file
# "task": "band",
# "kline_type":"abacus",
# "kpath":[[0.0000000000, 0.0000000000, 0.0000000000, 20],
# [0.5000000000, 0.0000000000, 0.0000000000, 1],
# [0.0000000000, 0.5000000000, 0.0000000000, 20],
# [0.0000000000, 0.0000000000, 0.0000000000, 20],
# [0.0000000000, 0.0000000000, 0.5000000000, 1],
# [-0.5000000000, -0.5000000000, 0.5000000000, 20],
# [0.0000000000, 0.0000000000, 0.0000000000, 20],
# [0.0000000000, -0.5000000000, 0.5000000000, 1 ],
# [-0.5000000000, 0.0000000000, 0.5000000000, 20],
# [0.0000000000, 0.0000000000, 0.0000000000, 20],
# [0.5000000000, -0.500000000, 0.0000000000, 1]
# ],
# "klabels":["G","X","Y","G","Z","R_2","G","T_2","U_2","G","V_2"],
# "nel_atom":{"Si":4},
# "E_fermi":0.0,
# "emin":-7,
# "emax":18
kpath_kwargs = jdata["task_options"]

bcal = Band(model=model,
use_gui=False,
results_path="./",
device=model.device)
bcal.get_bands(data=dataset[0],
kpath_kwargs=kpath_kwargs) # compute band structure

bcal.band_plot(
E_fermi = kpath_kwargs["E_fermi"],
emin = kpath_kwargs["emin"],
emax = kpath_kwargs["emax"]
) # plot the band structure
emax = kpath_kwargs["emax"]) # plot the band structure
72 changes: 72 additions & 0 deletions examples/e3/data/Si64.vasp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
Si64
1.0
1.088740444183349609e+01 0.000000000000000000e+00 0.000000000000000000e+00
0.000000000000000000e+00 1.088740444183349609e+01 0.000000000000000000e+00
0.000000000000000000e+00 0.000000000000000000e+00 1.088740444183349609e+01
Si
64
Cart
1.085161685943603516e+01 1.280032098293304443e-01 1.078928661346435547e+01
1.841060072183609009e-02 1.989128626883029938e-02 5.386826038360595703e+00
1.083230304718017578e+01 5.443179607391357422e+00 1.084113216400146484e+01
1.073209476470947266e+01 5.354926109313964844e+00 5.515145301818847656e+00
5.484899997711181641e+00 1.084975528717041016e+01 1.489288061857223511e-01
5.495874404907226562e+00 5.350070446729660034e-02 5.522494316101074219e+00
5.428655624389648438e+00 5.469603061676025391e+00 5.730040743947029114e-02
5.486892223358154297e+00 5.412618637084960938e+00 5.512423515319824219e+00
5.753992870450019836e-02 2.735220670700073242e+00 2.737147808074951172e+00
9.110579639673233032e-02 2.743179321289062500e+00 8.217595100402832031e+00
5.919481441378593445e-02 8.242929458618164062e+00 2.665236473083496094e+00
1.213727816939353943e-01 8.150724411010742188e+00 8.206772804260253906e+00
5.396494388580322266e+00 2.671355247497558594e+00 2.717561483383178711e+00
5.411704063415527344e+00 2.765727281570434570e+00 8.161295890808105469e+00
5.338889122009277344e+00 8.191737174987792969e+00 2.803484678268432617e+00
5.485008716583251953e+00 8.147904396057128906e+00 8.188732147216796875e+00
2.697985887527465820e+00 2.655328989028930664e+00 1.086999511718750000e+01
2.803299665451049805e+00 2.753925323486328125e+00 5.408394336700439453e+00
2.731758594512939453e+00 8.291128158569335938e+00 1.087376213073730469e+01
2.743048667907714844e+00 8.140621185302734375e+00 5.500893592834472656e+00
8.194023132324218750e+00 2.727316379547119141e+00 1.087789916992187500e+01
8.145998954772949219e+00 2.603254556655883789e+00 5.454698085784912109e+00
7.996123313903808594e+00 8.126315116882324219e+00 5.489429086446762085e-02
8.157540321350097656e+00 8.178661346435546875e+00 5.396069526672363281e+00
2.724681615829467773e+00 1.085770320892333984e+01 2.624702692031860352e+00
2.642884731292724609e+00 1.709322445094585419e-02 8.155264854431152344e+00
2.728481531143188477e+00 5.489842891693115234e+00 2.687501192092895508e+00
2.719804286956787109e+00 5.382710933685302734e+00 8.202189445495605469e+00
8.232271194458007812e+00 1.088958159089088440e-01 2.695187807083129883e+00
8.190147399902343750e+00 1.076476860046386719e+01 8.095950126647949219e+00
8.187970161437988281e+00 5.468699455261230469e+00 2.664496183395385742e+00
8.083069801330566406e+00 5.405988216400146484e+00 8.226228713989257812e+00
4.130180358886718750e+00 1.358443140983581543e+00 4.137616157531738281e+00
4.051333904266357422e+00 1.297201514244079590e+00 9.625466346740722656e+00
4.133435726165771484e+00 6.705410480499267578e+00 4.101284980773925781e+00
4.076243877410888672e+00 6.794273376464843750e+00 9.562548637390136719e+00
9.573958396911621094e+00 1.423081636428833008e+00 4.038333892822265625e+00
9.454534530639648438e+00 1.340337395668029785e+00 9.410832405090332031e+00
9.585140228271484375e+00 6.824736595153808594e+00 3.959444046020507812e+00
9.489406585693359375e+00 6.853000164031982422e+00 9.511748313903808594e+00
1.335405468940734863e+00 1.307370305061340332e+00 1.314915299415588379e+00
1.366260290145874023e+00 1.351388096809387207e+00 6.778203964233398438e+00
1.336178421974182129e+00 6.824344635009765625e+00 1.330636739730834961e+00
1.347283601760864258e+00 6.734087944030761719e+00 6.762024879455566406e+00
6.878825187683105469e+00 1.403963446617126465e+00 1.293412685394287109e+00
6.797452926635742188e+00 1.434306621551513672e+00 6.866337299346923828e+00
6.644256114959716797e+00 6.731377124786376953e+00 1.499554872512817383e+00
6.844486236572265625e+00 6.848329544067382812e+00 6.809374332427978516e+00
1.377604961395263672e+00 4.136429786682128906e+00 4.032868862152099609e+00
1.389494061470031738e+00 4.109994888305664062e+00 9.576973915100097656e+00
1.447992086410522461e+00 9.560729980468750000e+00 4.123854637145996094e+00
1.328622579574584961e+00 9.550005912780761719e+00 9.473250389099121094e+00
6.836505889892578125e+00 4.088655471801757812e+00 4.096276760101318359e+00
6.809603214263916016e+00 4.098802566528320312e+00 9.570180892944335938e+00
6.812346458435058594e+00 9.609702110290527344e+00 3.982939004898071289e+00
6.827219009399414062e+00 9.470495223999023438e+00 9.475601196289062500e+00
4.046401500701904297e+00 4.024322032928466797e+00 1.342819690704345703e+00
4.066891670227050781e+00 4.043124675750732422e+00 6.784518241882324219e+00
4.096853733062744141e+00 9.517572402954101562e+00 1.429788351058959961e+00
4.126554965972900391e+00 9.538574218750000000e+00 6.838574409484863281e+00
9.460348129272460938e+00 4.077964305877685547e+00 1.332836031913757324e+00
9.455187797546386719e+00 3.956558942794799805e+00 6.875069141387939453e+00
9.541971206665039062e+00 9.649669647216796875e+00 1.294305443763732910e+00
9.570387840270996094e+00 9.492716789245605469e+00 6.793272018432617188e+00
5 changes: 1 addition & 4 deletions examples/e3/data/info.json
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,5 @@
{
"nframes": 1,
"pos_type": "cart",
"AtomicData_options": {
"r_max": 7.4,
"pbc": true
}
"pbc": true
}
54 changes: 54 additions & 0 deletions examples/e3/input_short_cpu.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
{
"common_options": {
"basis": {
"Si": "1s1p"
},
"device": "cpu",
"overlap": true
},
"model_options": {
"embedding": {
"method": "slem",
"r_max": {
"Si": 7.4
},
"irreps_hidden": "32x0e+32x1o+16x2e",
"n_layers": 3,
"avg_num_neighbors": 51,
"tp_radial_emb": true
},
"prediction":{
"method": "e3tb",
"neurons": [64,64]
}
},
"train_options": {
"num_epoch": 1500,
"batch_size": 1,
"optimizer": {
"lr": 0.005,
"type": "Adam"
},
"lr_scheduler": {
"type": "rop",
"factor": 0.8,
"patience": 50,
"min_lr": 1e-6
},
"loss_options":{
"train":{"method": "hamil_abs"}
},
"save_freq": 100,
"validation_freq": 10,
"display_freq": 1
},
"data_options": {
"train": {
"root": "./data",
"prefix": "Si64",
"get_Hamiltonian": true,
"get_overlap": true
}
}
}

0 comments on commit 5dfced2

Please sign in to comment.