diff --git a/.gitignore b/.gitignore index 2c0ef0cc..89fbe255 100644 --- a/.gitignore +++ b/.gitignore @@ -21,10 +21,10 @@ asserts *.css band.png *.pth -*.npy -*.traj -*.dat -*.vasp +#*.npy +#*.traj +#*.dat +#*.vasp *log* checkpoint.pl # Byte-compiled / optimized / DLL files diff --git a/dptb/data/dataset/_default_dataset.py b/dptb/data/dataset/_default_dataset.py index 34f036d3..4a0d7f30 100644 --- a/dptb/data/dataset/_default_dataset.py +++ b/dptb/data/dataset/_default_dataset.py @@ -96,7 +96,7 @@ def __init__(self, kpoints = np.expand_dims(kpoints, axis=0) self.data["kpoints"] = np.broadcast_to(kpoints, (self.info["nframes"], kpoints.shape[1], 3)) - if kpoints.shape[0] == self.info["nframes"]: + elif kpoints.ndim == 3 and kpoints.shape[0] == self.info["nframes"]: # array of kpoints, (nframes, nkpoints, 3) self.data["kpoints"] = kpoints else: diff --git a/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/eigenvalues.npy b/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/eigenvalues.npy new file mode 100644 index 00000000..62abb802 Binary files /dev/null and b/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/eigenvalues.npy differ diff --git a/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/info.json b/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/info.json new file mode 100644 index 00000000..d2d7bd90 --- /dev/null +++ b/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/info.json @@ -0,0 +1,19 @@ +{ + "nframes": 1, + "natoms": 2, + "pos_type": "ase", + "AtomicData_options": { + "r_max": 5.0, + "er_max": 5.0, + "oer_max": 2.5, + "pbc": true + }, + "bandinfo": { + "nkpoints": 61, + "nbands": 14, + "band_min": 0, + "band_max": 6, + "emin": null, + "emax": null + } +} \ No newline at end of file diff --git a/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/kpoints.npy b/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/kpoints.npy new file mode 100644 index 00000000..df905f75 Binary files /dev/null and b/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/kpoints.npy differ diff --git a/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/xdat.traj b/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/xdat.traj new file mode 100644 index 00000000..d877e685 Binary files /dev/null and b/dptb/tests/data/test_data_nequip/dataset/kpath_spk.0/xdat.traj differ diff --git a/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/eigenvalues.npy b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/eigenvalues.npy new file mode 100644 index 00000000..d541deb3 Binary files /dev/null and b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/eigenvalues.npy differ diff --git a/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/eigenvalues_ref.npy b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/eigenvalues_ref.npy new file mode 100644 index 00000000..45c8622c Binary files /dev/null and b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/eigenvalues_ref.npy differ diff --git a/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/info.json b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/info.json new file mode 100644 index 00000000..6485de04 --- /dev/null +++ b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/info.json @@ -0,0 +1,17 @@ +{ + "nframes": 10, + "natoms": 8, + "pos_type": "ase", + "AtomicData_options": { + "r_max": 5.0, + "er_max": 5.0, + "oer_max": 2.5, + "pbc": true + }, + "bandinfo": { + "band_min": 0, + "band_max": 8, + "emin": null, + "emax": null + } +} \ No newline at end of file diff --git a/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/kpoints.npy b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/kpoints.npy new file mode 100644 index 00000000..77e2885a Binary files /dev/null and b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/kpoints.npy differ diff --git a/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/struct.vasp b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/struct.vasp new file mode 100644 index 00000000..4dad1b2b --- /dev/null +++ b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/struct.vasp @@ -0,0 +1,16 @@ +Si + 1.0000000000000000 + 5.4300150871000001 0.0000000000000000 0.0000000000000000 + 0.0000000000000000 5.4300150871000001 0.0000000000000000 + 0.0000000000000000 0.0000000000000000 5.4300150871000001 + Si + 8 +Cartesian + 0.0264071000000000 5.4365699999999997 5.4306599999999996 + 2.7157399999999998 2.7145600000000001 5.4341600000000003 + 2.6988300000000001 0.0021865800000000 2.7229999999999999 + -0.0085004700000000 2.6946599999999998 2.7118500000000001 + 1.3702900000000000 1.3672899999999999 1.3587400000000001 + 4.0686400000000003 4.0873600000000003 1.3605700000000001 + 4.0663099999999996 1.3515699999999999 4.0690799999999996 + 1.3523300000000000 4.0658700000000003 4.0620099999999999 diff --git a/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/xdat.traj b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/xdat.traj new file mode 100644 index 00000000..57dcbfe8 Binary files /dev/null and b/dptb/tests/data/test_data_nequip/dataset/kpathmd25.0/xdat.traj differ diff --git a/dptb/tests/data/test_data_nequip/input/input_dptb.json b/dptb/tests/data/test_data_nequip/input/input_dptb.json new file mode 100644 index 00000000..1bc42b9e --- /dev/null +++ b/dptb/tests/data/test_data_nequip/input/input_dptb.json @@ -0,0 +1,54 @@ +{ + "common_options": { + "basis": { + "Si": ["3s", "3p", "d*"] + }, + "device": "cpu", + "dtype": "float32", + "overlap": false + }, + "train_options": { + "num_epoch": 1500, + "batch_size": 1, + "optimizer": { + "lr": 0.001, + "type": "Adam" + }, + "lr_scheduler": { + "type": "exp", + "gamma": 0.999 + }, + "loss_options":{ + "train": {"method": "eigvals"} + }, + "save_freq": 500, + "validation_freq": 10, + "display_freq": 100 + }, + "model_options": { + "embedding":{ + "method": "se2", + "rs": 2.5, + "rc": 5.0, + "radial_net": { + "neurons": [10,20,30] + } + }, + "prediction":{ + "method": "sktb", + "neurons": [16,16,16] + }, + "nnsk": { + "onsite": {"method": "strain", "rs":2.5 ,"w":0.3}, + "hopping": {"method": "powerlaw", "rs":5.0, "w": 0.1}, + "freeze": true + } + }, + "data_options": { + "train": { + "root": "./data/", + "prefix": "kpath_spk", + "get_eigenvalues": true + } + } +} \ No newline at end of file diff --git a/dptb/tests/data/test_data_nequip/input/input_md.json b/dptb/tests/data/test_data_nequip/input/input_md.json new file mode 100644 index 00000000..bd0babcb --- /dev/null +++ b/dptb/tests/data/test_data_nequip/input/input_md.json @@ -0,0 +1,42 @@ +{ + "common_options": { + "basis": { + "Si": ["3s", "3p", "d*"] + }, + "device": "cpu", + "dtype": "float32", + "overlap": false + }, + "train_options": { + "num_epoch": 2, + "batch_size": 1, + "optimizer": { + "lr": 0.01, + "type": "Adam" + }, + "lr_scheduler": { + "type": "exp", + "gamma": 0.999 + }, + "loss_options":{ + "train": {"method": "eigvals"} + }, + "save_freq": 500, + "validation_freq": 10, + "display_freq": 100 + }, + "model_options": { + "nnsk": { + "onsite": {"method": "strain", "rs":2.5 ,"w":0.3}, + "hopping": {"method": "powerlaw", "rs":2.6, "w": 0.35}, + "freeze": false + } + }, + "data_options": { + "train": { + "root": "./dptb/tests/data/test_data_nequip/dataset/", + "prefix": "kpathmd25", + "get_eigenvalues": true + } + } +} \ No newline at end of file diff --git a/dptb/tests/data/test_data_nequip/input/input_push_rs.json b/dptb/tests/data/test_data_nequip/input/input_push_rs.json new file mode 100644 index 00000000..65e22c0d --- /dev/null +++ b/dptb/tests/data/test_data_nequip/input/input_push_rs.json @@ -0,0 +1,42 @@ +{ + "common_options": { + "basis": { + "Si": ["3s", "3p", "d*"] + }, + "device": "cpu", + "dtype": "float32", + "overlap": false + }, + "train_options": { + "num_epoch": 10, + "batch_size": 1, + "optimizer": { + "lr": 0.01, + "type": "Adam" + }, + "lr_scheduler": { + "type": "exp", + "gamma": 0.999 + }, + "loss_options":{ + "train": {"method": "eigvals"} + }, + "save_freq": 1, + "display_freq": 100 + }, + "model_options": { + "nnsk": { + "onsite": {"method": "strain", "rs":2.5 ,"w":0.3}, + "hopping": {"method": "powerlaw", "rs":2.6, "w": 0.3}, + "freeze": false, + "push": {"rs_thr": 0.01, "period": 1} + } + }, + "data_options": { + "train": { + "root": "./dptb/tests/data/test_data_nequip/dataset/", + "prefix": "kpath_spk", + "get_eigenvalues": true + } + } +} \ No newline at end of file diff --git a/dptb/tests/data/test_data_nequip/input/input_push_w.json b/dptb/tests/data/test_data_nequip/input/input_push_w.json new file mode 100644 index 00000000..2fc1d978 --- /dev/null +++ b/dptb/tests/data/test_data_nequip/input/input_push_w.json @@ -0,0 +1,42 @@ +{ + "common_options": { + "basis": { + "Si": ["3s", "3p", "d*"] + }, + "device": "cpu", + "dtype": "float32", + "overlap": false + }, + "train_options": { + "num_epoch": 10, + "batch_size": 1, + "optimizer": { + "lr": 0.05, + "type": "Adam" + }, + "lr_scheduler": { + "type": "exp", + "gamma": 0.999 + }, + "loss_options":{ + "train": {"method": "eigvals"} + }, + "save_freq": 1, + "display_freq": 100 + }, + "model_options": { + "nnsk": { + "onsite": {"method": "strain", "rs":2.5 ,"w":0.3}, + "hopping": {"method": "powerlaw", "rs":5.0, "w": 0.3}, + "freeze": false, + "push": {"w_thr": 0.01, "period": 1} + } + }, + "data_options": { + "train": { + "root": "./dptb/tests/data/test_data_nequip/dataset/", + "prefix": "kpath_spk", + "get_eigenvalues": true + } + } +} \ No newline at end of file diff --git a/dptb/tests/data/test_data_nequip/input/input_strain_polar.json b/dptb/tests/data/test_data_nequip/input/input_strain_polar.json new file mode 100644 index 00000000..d4ac8084 --- /dev/null +++ b/dptb/tests/data/test_data_nequip/input/input_strain_polar.json @@ -0,0 +1,41 @@ +{ + "common_options": { + "basis": { + "Si": ["3s", "3p", "d*"] + }, + "device": "cpu", + "dtype": "float32", + "overlap": false + }, + "train_options": { + "num_epoch": 5, + "batch_size": 1, + "optimizer": { + "lr": 0.05, + "type": "Adam" + }, + "lr_scheduler": { + "type": "exp", + "gamma": 0.999 + }, + "loss_options":{ + "train": {"method": "eigvals"} + }, + "save_freq": 1, + "display_freq": 100 + }, + "model_options": { + "nnsk": { + "onsite": {"method": "strain", "rs":2.5 ,"w":0.3}, + "hopping": {"method": "powerlaw", "rs":2.6, "w": 0.3}, + "freeze": false + } + }, + "data_options": { + "train": { + "root": "./dptb/tests/data/test_data_nequip/dataset/", + "prefix": "kpath_spk", + "get_eigenvalues": true + } + } +} \ No newline at end of file diff --git a/dptb/tests/data/test_data_nequip/input/input_valence.json b/dptb/tests/data/test_data_nequip/input/input_valence.json new file mode 100644 index 00000000..9a8778ca --- /dev/null +++ b/dptb/tests/data/test_data_nequip/input/input_valence.json @@ -0,0 +1,42 @@ +{ + "common_options": { + "basis": { + "Si": ["3s", "3p"] + }, + "device": "cpu", + "dtype": "float32", + "overlap": false, + "seed": 120468 + }, + "train_options": { + "num_epoch": 5, + "batch_size": 1, + "optimizer": { + "lr": 0.05, + "type": "Adam" + }, + "lr_scheduler": { + "type": "exp", + "gamma": 0.999 + }, + "loss_options":{ + "train": {"method": "eigvals"} + }, + "save_freq": 1, + "display_freq": 10 + }, + "model_options": { + "nnsk": { + "onsite": {"method": "none"}, + "hopping": {"method": "powerlaw", "rs":2.6, "w": 0.3}, + "freeze": false + } + }, + "data_options": { + "train": { + "root": "./dptb/tests/data/test_data_nequip/dataset/", + "prefix": "kpath_spk", + "get_eigenvalues": true + } + } +} \ No newline at end of file diff --git a/dptb/tests/test_new_nnsk_train.py b/dptb/tests/test_new_nnsk_train.py new file mode 100644 index 00000000..cf67f331 --- /dev/null +++ b/dptb/tests/test_new_nnsk_train.py @@ -0,0 +1,49 @@ +from dptb.entrypoints import train +import pytest +import torch +import numpy as np + +@pytest.fixture(scope='session', autouse=True) +def root_directory(request): + return str(request.config.rootdir) + +def test_nnsk_valence(): + INPUT_file = "./dptb/tests/data/test_data_nequip/input/input_valence.json" + output = "./dptb/tests/data/test_data_nequip/output" + + train(INPUT=INPUT_file, init_model=None, restart=None, train_soc=False,\ + output=output+"/test_valence", log_level=5, log_path=output+"/test_valence.log") + +def test_nnsk_strain_polar(): + INPUT_file = "./dptb/tests/data/test_data_nequip/input/input_strain_polar.json" + output = "./dptb/tests/data/test_data_nequip/output" + init_model = "./dptb/tests/data/test_data_nequip/output/test_valence/checkpoint/nnsk.iter6.pth" + + train(INPUT=INPUT_file, init_model=init_model, restart=None, train_soc=False,\ + output=output+"/test_strain_polar", log_level=5, log_path=output+"/test_strain_polar.log") + +def test_nnsk_push(): + INPUT_file_rs = "./dptb/tests/data/test_data_nequip/input/input_push_rs.json" + INPUT_file_w = "./dptb/tests/data/test_data_nequip/input/input_push_w.json" + output = "./dptb/tests/data/test_data_nequip/output" + init_model = "./dptb/tests/data/test_data_nequip/output/test_strain_polar/checkpoint/nnsk.iter6.pth" + + train(INPUT=INPUT_file_rs, init_model=init_model, restart=None, train_soc=False,\ + output=output+"/test_push_rs", log_level=5, log_path=output+"/test_push_rs.log") + train(INPUT=INPUT_file_w, init_model=init_model, restart=None, train_soc=False,\ + output=output+"/test_push_w", log_level=5, log_path=output+"/test_push_w.log") + + model_rs = torch.load("./dptb/tests/data/test_data_nequip/output/test_push_rs/checkpoint/nnsk.iter11.pth") + model_w = torch.load("./dptb/tests/data/test_data_nequip/output/test_push_w/checkpoint/nnsk.iter11.pth") + # test push limits + # 10 epoch, 0.01 step, 1 period -> 0.05 added. + assert np.isclose(model_rs["config"]["model_options"]["nnsk"]["hopping"]["rs"], 2.65) + assert np.isclose(model_w["config"]["model_options"]["nnsk"]["hopping"]["w"], 0.35) + +def test_md(): + INPUT_file = "./dptb/tests/data/test_data_nequip/input/input_md.json" + output = "./dptb/tests/data/test_data_nequip/output" + init_model = "./dptb/tests/data/test_data_nequip/output/test_push_w/checkpoint/nnsk.iter11.pth" + + train(INPUT=INPUT_file, init_model=init_model, restart=None, train_soc=False,\ + output=output+"/test_md", log_level=5, log_path=output+"/test_md.log") \ No newline at end of file