refactor: plotting code in dftb2nnsk.py fix push_decay method and update push options

dptb/nn/dftb2nnsk.py

@@ -96,17 +96,19 @@ def viz(self, r_min=1.5, r_max=5.0):
             r = torch.linspace(r_min,r_max, steps=100)
             hops = vmap(self.step)(r.reshape(-1,1))
 
+
             dftb_hopping = self.dftb(r, mode="hopping").permute(1,0,2)
             dftb_overlap = self.dftb(r, mode="overlap").permute(1,0,2)
 
+            r = r.numpy()
             fig = plt.figure(figsize=(6,4))
             # hops[0] shape - [n_r, n_edge, n_skintegrals]
 
             for i in range(hops[0].shape[1]):
-                plt.plot(r, hops[0][:,i, :-1], c="C"+str(i))
-                plt.plot(r, hops[0][:,i, -1], c="C"+str(i))
-                plt.plot(r, dftb_hopping[:,i, :-1], c="C"+str(i), linestyle="--")
-                plt.plot(r, dftb_hopping[:,i, -1], c="C"+str(i), linestyle="--")
+                plt.plot(r, hops[0][:,i, :-1].detach().numpy(), c="C"+str(i))
+                plt.plot(r, hops[0][:,i, -1].detach().numpy(), c="C"+str(i))
+                plt.plot(r, dftb_hopping[:,i, :-1].numpy(), c="C"+str(i), linestyle="--")
+                plt.plot(r, dftb_hopping[:,i, -1].numpy(), c="C"+str(i), linestyle="--")
             plt.title("hoppings")
             plt.xlabel("r(angstrom)")
             plt.tight_layout()
@@ -115,10 +117,10 @@ def viz(self, r_min=1.5, r_max=5.0):
 
             fig = plt.figure(figsize=(6,4))
             for i in range(hops[1].shape[1]):
-                plt.plot(r, hops[1][:,i, :-1], c="C"+str(i))
-                plt.plot(r, hops[1][:,i, -1], c="C"+str(i))
-                plt.plot(r, dftb_overlap[:,i, :-1], c="C"+str(i), linestyle="--")
-                plt.plot(r, dftb_overlap[:,i, -1], c="C"+str(i), linestyle="--")
+                plt.plot(r, hops[1][:,i, :-1].detach().numpy(), c="C"+str(i))
+                plt.plot(r, hops[1][:,i, -1].detach().numpy(), c="C"+str(i))
+                plt.plot(r, dftb_overlap[:,i, :-1].numpy(), c="C"+str(i), linestyle="--")
+                plt.plot(r, dftb_overlap[:,i, -1].numpy(), c="C"+str(i), linestyle="--")
             plt.title("overlaps")
             plt.xlabel("r(angstrom)")
             plt.tight_layout()

dptb/nn/nnsk.py

@@ -200,23 +200,19 @@ def push_decay(self, rs_thr: float=0., rc_thr: float=0., w_thr: float=0., ovp_th
             the threshold step to push the w
         """
 
-
-        if self.count_push // period > 0:
+        self.count_push += 1
+        if self.count_push % period == 0:
             if abs(rs_thr) > 0:
                 self.hopping_options["rs"] += rs_thr
             if abs(w_thr) > 0:
                 self.hopping_options["w"] += w_thr
             if abs(rc_thr) > 0:
                 self.hopping_options["rc"] += rc_thr
-            if abs(ovp_thr) > 0 and self.ovp_factor >=ovp_thr:
-                self.ovp_factor -= ovp_thr
-
+            if abs(ovp_thr) > 0 and self.ovp_factor >= abs(ovp_thr):
+                self.ovp_factor += ovp_thr
+                log.info(f"ovp_factor is decreased to {self.ovp_factor}")
             self.model_options["nnsk"]["hopping"] = self.hopping_options
 
-            self.count_push = 0
-        else:
-            self.count_push += 1
-
     def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         # get the env and bond from the data
         # calculate the sk integrals
@@ -227,7 +223,7 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         # map the parameters to the edge/node/env features
         # compute integrals from parameters using hopping and onsite clas
         if self.push is not None and self.push is not False:
-            if abs(self.push.get("rs_thr")) + abs(self.push.get("rc_thr")) + abs(self.push.get("w_thr")) > 0:
+            if abs(self.push.get("rs_thr")) + abs(self.push.get("rc_thr")) + abs(self.push.get("w_thr")) + abs(self.push.get("ovp_thr")) > 0:
                 self.push_decay(**self.push)
 
         reflective_bonds = np.array([self.idp_sk.bond_to_type["-".join(self.idp_sk.type_to_bond[i].split("-")[::-1])] for i  in range(len(self.idp_sk.bond_types))])
@@ -515,8 +511,12 @@ def from_reference(
             model = cls(**common_options, **nnsk, transform=transform)
 
             if f["config"]["common_options"]["basis"] == common_options["basis"] and \
-                f["config"]["model_options"] == model.model_options:
+                f["config"]["model_options"]["nnsk"]["onsite"] == model.model_options["nnsk"]["onsite"] and \
+                f["config"]["model_options"]["nnsk"]["hopping"] == model.model_options["nnsk"]["hopping"] and \
+                    f["config"]["model_options"]["nnsk"]["soc"] == model.model_options["nnsk"]["soc"]:
+
                 model.load_state_dict(f["model_state_dict"])
+
             else:
                 #TODO: handle the situation when ref_model config is not the same as the current model
                 # load hopping

dptb/plugins/saver.py

@@ -28,7 +28,17 @@ def register(self, trainer, checkpoint_path):
                 # 计算所有阈值之和
                 thrs = sum(abs(val) for key, val in push_option.items() if "thr" in key)
                 # 如果阈值之和不为 0, 则 push 为 True
-                push = abs(thrs) != 0.0
+                if abs(push_option['rs_thr'])  + abs(push_option['w_thr']) != 0.0 and abs(push_option['ovp_thr']) != 0.0:
+                    log.error("rs_thr, w_thr and ovp_thr cannot be pushed at the same time.")
+                    raise ValueError("rs_thr, w_thr and ovp_thr cannot be pushed at the same time.")
+
+                if abs(push_option['rs_thr'])  + abs(push_option['w_thr']) != 0.0:
+                    push = 'rs_w'
+                # push = abs(thrs) != 0.0
+                elif abs(push_option['ovp_thr']) != 0.0:
+                    push = 'overlap'
+                else:
+                    push = False            
             else:
                 push = False
         else:
@@ -39,10 +49,13 @@ def register(self, trainer, checkpoint_path):
     def iteration(self, **kwargs):
         # suffix = "_b"+"%.3f"%self.trainer.common_options["bond_cutoff"]+"_c"+"%.3f"%self.trainer.onsite_options["skfunction"]["sk_cutoff"]+"_w"+\
         #         "%.3f"%self.trainer.model_options["skfunction"]["sk_decay_w"]
-        if self.push:
+        if self.push == 'rs_w':
             suffix = ".iter_rs" + "%.3f"%self.trainer.model.hopping_options["rs"]+"_w"+"%.3f"%self.trainer.model.hopping_options["w"]
             # By default, the maximum number of saved checkpoints is 100 for pushing rs and w.
-            max_ckpt = 100
+            max_ckpt = self.trainer.train_options["max_ckpt"]
+        elif self.push == 'overlap':
+            suffix= ".iter_ovp" + "%.3f"%self.trainer.model.ovp_factor
+            max_ckpt = self.trainer.train_options["max_ckpt"]
         else:
             suffix = ".iter{}".format(self.trainer.iter)
             max_ckpt = self.trainer.train_options["max_ckpt"]

dptb/tests/test_sktb.py

@@ -48,19 +48,19 @@ def test_nnsk_push(root_directory):
     train(INPUT=INPUT_file_w, init_model=init_model, restart=None,\
           output=output+"/test_push_w", log_level=5, log_path=output+"/test_push_w.log")
 
-    model_rs = torch.load(f"{root_directory}/dptb/tests/data/test_sktb/output/test_push_rs/checkpoint/nnsk.iter_rs2.650_w0.300.pth")
-    model_w = torch.load(f"{root_directory}/dptb/tests/data/test_sktb/output/test_push_w/checkpoint/nnsk.iter_rs5.000_w0.350.pth")
+    model_rs = torch.load(f"{root_directory}/dptb/tests/data/test_sktb/output/test_push_rs/checkpoint/nnsk.iter_rs2.700_w0.300.pth")
+    model_w = torch.load(f"{root_directory}/dptb/tests/data/test_sktb/output/test_push_w/checkpoint/nnsk.iter_rs5.000_w0.400.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)
+    assert np.isclose(model_rs["config"]["model_options"]["nnsk"]["hopping"]["rs"], 2.700)
+    assert np.isclose(model_w["config"]["model_options"]["nnsk"]["hopping"]["w"], 0.40)
 
 # train on md structures.
 @pytest.mark.order(4)
 def test_md(root_directory):
     INPUT_file =root_directory + "/dptb/tests/data/test_sktb/input/input_md.json"
     output = root_directory + "/dptb/tests/data/test_sktb/output"
-    init_model = root_directory + "/dptb/tests/data/test_sktb/output/test_push_w/checkpoint/nnsk.iter_rs5.000_w0.350.pth"
+    init_model = root_directory + "/dptb/tests/data/test_sktb/output/test_push_w/checkpoint/nnsk.iter_rs5.000_w0.400.pth"
 
     check_config_train(INPUT=INPUT_file, init_model=init_model, restart=None)
     train(INPUT=INPUT_file, init_model=init_model, restart=None,\

dptb/utils/argcheck.py

@@ -655,7 +655,22 @@ def hopping():
         Argument("rs", float, optional=True, default=6.0, doc=doc_rs),
         Argument("w", float, optional=True, default=0.1, doc=doc_w),
     ]
-
+    poly1pow = [
+        Argument("rs", float, optional=True, default=6.0, doc=doc_rs),
+        Argument("w", float, optional=True, default=0.1, doc=doc_w),
+    ]
+    poly2pow = [
+        Argument("rs", float, optional=True, default=6.0, doc=doc_rs),
+        Argument("w", float, optional=True, default=0.1, doc=doc_w),
+    ]
+    poly3pow = [
+        Argument("rs", float, optional=True, default=6.0, doc=doc_rs),
+        Argument("w", float, optional=True, default=0.1, doc=doc_w),
+    ]    
+    poly2exp = [
+        Argument("rs", float, optional=True, default=6.0, doc=doc_rs),
+        Argument("w", float, optional=True, default=0.1, doc=doc_w),
+    ]
     varTang96 = [
         Argument("rs", float, optional=True, default=6.0, doc=doc_rs),
         Argument("w", float, optional=True, default=0.1, doc=doc_w),
@@ -669,6 +684,10 @@ def hopping():
 
     return Variant("method", [
                     Argument("powerlaw", dict, powerlaw),
+                    Argument("poly1pow", dict, poly1pow),
+                    Argument("poly2pow", dict, poly2pow),
+                    Argument("poly3pow", dict, poly3pow),
+                    Argument("poly2exp", dict, poly2exp),
                     Argument("varTang96", dict, varTang96),
                     Argument("NRL0", dict, NRL),
                     Argument("NRL1", dict, NRL),

examples/mos2/data/set.0/info.json

@@ -9,8 +9,8 @@
     },
     "bandinfo": {
         "band_min": 0,
-        "band_max": 12,
-        "emin": null,
-        "emax": null
+        "band_max": 14,
+        "emin": 0.0,
+        "emax": 20.0
     }
 }

examples/mos2/input_nnsk.json

@@ -0,0 +1,46 @@
+{
+    "common_options": {
+            "basis": {
+                "Mo":["5s","5p","4d"],
+                "S":["3s","3p","d*"]
+            },
+            "device": "cpu",
+            "dtype": "float32",
+            "overlap": false,
+            "seed": 42
+    },
+    "train_options": {
+        "num_epoch": 500,
+        "batch_size": 1,
+        "optimizer": {
+            "lr": 1e-3,
+            "type": "Adam"
+        },
+        "lr_scheduler": {
+            "type": "exp",
+            "gamma": 0.9990
+        },
+        "loss_options":{
+            "train": {"method": "eigvals"}
+        },
+        "save_freq": 50,
+        "validation_freq": 10,
+        "display_freq": 1
+    },
+    "model_options": {
+        "nnsk": {
+            "onsite": {"method": "uniform"},
+            "hopping": {"method": "poly2pow", "rs":5.3, "w": 0.2},
+            "soc":{},
+            "freeze": false,
+            "push":false
+        }
+    },
+    "data_options": {
+        "train": {
+            "root": "./data/",
+            "prefix": "set",
+            "get_eigenvalues": true
+        }
+    }
+}

examples/mos2/input_nnsk_push.json

@@ -0,0 +1,46 @@
+{
+    "common_options": {
+            "basis": {
+                "Mo":["5s","5p","4d"],
+                "S":["3s","3p","d*"]
+            },
+            "device": "cpu",
+            "dtype": "float32",
+            "overlap": true,
+            "seed": 42
+    },
+    "train_options": {
+        "num_epoch": 100,
+        "batch_size": 1,
+        "optimizer": {
+            "lr": 1e-3,
+            "type": "Adam"
+        },
+        "lr_scheduler": {
+            "type": "exp",
+            "gamma": 0.9999
+        },
+        "loss_options":{
+            "train": {"method": "eigvals"}
+        },
+        "save_freq": 10,
+        "validation_freq": 10,
+        "display_freq": 1
+    },
+    "model_options": {
+        "nnsk": {
+            "onsite": {"method": "uniform"},
+            "hopping": {"method": "poly2pow", "rs":5.3, "w": 0.2},
+            "soc":{},
+            "freeze": ["overlap"],
+            "push":{"ovp_thr":-0.01, "period":1}
+        }
+    },
+    "data_options": {
+        "train": {
+            "root": "./data/",
+            "prefix": "set",
+            "get_eigenvalues": true
+        }
+    }
+}