update h-BN example, ckpt and the related doc. (#127)

* update mix.ep50.pth, mix.iter500.pth and train_config.json

* update hBN example

* docs: update hands_on.md

docs/quick_start/hands_on.md

@@ -65,7 +65,9 @@ Having the data file and input parameter, we can start training our first **DeeP
     "nnsk": {
             "onsite": {"method": "none"},
             "hopping": {"method": "powerlaw", "rs":1.6, "w":0.3},
-            "freeze": false
+            "soc":{},
+            "freeze": false,
+            "push":false
         }
     }
 ```
@@ -77,16 +79,16 @@ Since we are using only the valence orbitals at this stage, we can limit the ene
     "bandinfo": {
         "band_min": 0,
         "band_max": 6,
-        "emin": -0.1,
-        "emax": 20.0
+        "emin": null,
+        "emax": null
     }
 ```
 
 Using the follwing command and we can train the first model:
 
 ```bash
 cd deeptb/examples/hBN
-dptb train input_short.json -o ./first
+dptb train ./input/input_first.json -o ./first
 ```
 
 Here ``-o`` indicate the output directory. During the fitting procedure, we can see the loss curve of hBN is decrease consistently. When finished, we get the fitting results in folders ```first```.
@@ -95,6 +97,10 @@ By modify the checkpoint path in the script `plot_band.py` and running it, the b
 ```bash
 python plot_band.py
 ```
+or just using the command line 
+```bash
+dptb run ./run/band.json -i ./first/checkpoint/nnsk.best.pth -o ./band_plot
+```
 > Note: the ```basis``` setting in the plotting script must be the same as in the input.
 
 ![band_first](../img/hBN_band_first.png)
@@ -121,7 +127,7 @@ To train the conduction band, the energy window we previously set in `info.json`
 
 We can then start the training using the previous model and modified input:
 ```bash
-dptb train input_short.json -i ./first/checkpoint/nnsk.ep1000.pth -o ./condband
+dptb train input/input_condband.json -i ./first/checkpoint/nnsk.ep500.pth -o ./condband
 ```
 ``-i`` states initialize the model from the checkpoint file, where the previous model is provided.
 
@@ -147,7 +153,7 @@ We can further improve the accuracy by incorporating more features of our code,
 After setting we can run the training for strain model:
 
 ```bash
-dptb train input_short.json  -i ./condband/checkpoint/nnsk.ep500.pth -o ./strain
+dptb train input/input_strain.json  -i ./condband/checkpoint/nnsk.ep500.pth -o ./strain
 ```
 
 We can also plot the band structure of the strain model:
@@ -161,30 +167,74 @@ It looks ok, we can further improve the accuracy by adding more neighbours, and
         "nnsk": {
             "onsite": {"method": "strain", "rs":1.6, "w":0.3},
             "hopping": {"method": "powerlaw", "rs":1.6, "w": 0.3},
+            "soc":{},
             "push": {"rs_thr": 0.02, "period": 10},
             "freeze": false
         }
     }
 ```
-This means that we gradually add up the `rs` in decay function, pushing up to 3rd nearest neighbour for considering in calculating bonding. see the input file `hBN/input/3_varycutoff.json` for detail. Then we can run the training again:
+This means that we gradually add up the `rs` in decay function, pushing up to 3rd nearest neighbour for considering in calculating bonding. see the input file `hBN/input/input_push_rs.json` for detail. Then we can run the training again:
 
 ```bash
-dptb train input_short.json  -i ./strain/checkpoint/nnsk.ep500.pth -o ./varycutoff
+dptb train input/input_push_rs.json  -i ./strain/checkpoint/nnsk.ep500.pth -o ./push_rs
 ```
 
 We finally get the model with more neighbors. We can plot the result again:
 
-![band_varycutoff](../img/hBN_band_varycutoff.png)
+![band_varycutoff](../img/hBN_band_pushrs.png)
+
+
+we can further push the decay w to 0.2 and train the model again. modify the model options:
+```json
+    "model_options": {
+        "nnsk": {
+            "onsite": {"method": "strain", "rs":1.6, "w":0.3},
+            "hopping": {"method": "powerlaw", "rs":3.4, "w": 0.3},
+            "soc":{},
+            "push": {"w_thr": -0.001, "period": 10},
+            "freeze": false
+        }
+    }
+```
+note:  we change the hopping cutoff `rs` to 3.4, and the push w_thr to -0.001.
+
+see the input file `hBN/input/input_push_w.json` and run the training:
+
+```bash
+dptb train input/input_push_w.json  -i ./push_rs/checkpoint/nnsk.iter_rs3.400_w0.300.pth -o ./push_w
+```
+
+We can the plot the band structure again:
+
+![band_varycutoff](../img/hBN_band_pushw.png)
+
+
+We can again increase more training epochs, using the pushed parameters and turn off push tag. see the input file `hBN/input/input_final.json` and run the training:
+
+```json
+    "model_options": {
+        "nnsk": {
+            "onsite": {"method": "strain", "rs":1.6, "w":0.3},
+            "hopping": {"method": "powerlaw", "rs":3.4, "w": 0.2},
+            "soc":{},
+            "push": false,
+            "freeze": false
+        }
+    }
+```
+
+```bash
+dptb train input/input_final.json  -i ./push_w/checkpoint/nnsk.iter_rs3.400_w0.210.pth -o ./final
+```
 
-We can again increase more training epochs, using the larger cutoff checkpoint. This can be done simply by assigning a large `num_epoch` in `train_options`. 
 
 And we can get a fairly good fitting result:
 
 ![band_longtrain](../img/hBN_band_longtrain.png)
 
 Now you have learned the basis use of **DeePTB**, however, the advanced functions still need to be explored for accurate and flexible electron structure representation, such as:
-- atomic orbitals
 - environmental correction
+- spin-orbital interaction
 - ...
 
 Altogether, we can simulate the electronic structure of a crystal system in a dynamic trajectory. **DeePTB** is capable of handling atom movement, volume change under stress, SOC effect and can use DFT eigenvalues with different orbitals and xc functionals as training targets.

examples/hBN/band_plot.py

@@ -1,10 +1,11 @@
 from dptb.postprocess.bandstructure.band import Band
 from dptb.nn.nnsk import NNSK
 from dptb.utils.tools import j_loader
+from dptb.nn.build import build_model
 import os
 
-model = NNSK.from_reference(checkpoint="./ref_ckpts/condband/checkpoint/nnsk.ep500.pth",
-                            basis={'B': ['2s', '2p', 'd*'], 'N': ['2s', '2p', 'd*']})
+model = build_model(checkpoint="./ref_ckpts/condband/checkpoint/nnsk.ep500.pth")
+
 jdata = j_loader("./run/band.json")
 results_path = "./band_plot"
 kpath_kwargs = jdata["task_options"]

examples/hBN/data/kpath.0/info.json

@@ -11,7 +11,7 @@
     "bandinfo": {
         "band_min": 0,
         "band_max": 6,
-        "emin": -0.1,
-        "emax": 40.0
+        "emin": null,
+        "emax": null
     }
-}
+}

examples/hBN/input/2_condband.json → examples/hBN/input/input_condband.json

@@ -1,44 +1,46 @@
-{
-    "common_options": {
-            "basis": {
-                "B": ["2s", "2p", "d*"],
-                "N": ["2s", "2p", "d*"]
-            },
-            "device": "cpu",
-            "dtype": "float32",
-            "overlap": false,
-            "seed": 114514
-    },
-    "train_options": {
-        "num_epoch": 500,
-        "batch_size": 1,
-        "optimizer": {
-            "lr": 0.05,
-            "type": "Adam"
-        },
-        "lr_scheduler": {
-            "type": "exp",
-            "gamma": 0.999
-        },
-        "loss_options":{
-            "train": {"method": "eigvals"}
-        },
-        "save_freq": 100,
-        "validation_freq": 10,
-        "display_freq": 10
-    },
-    "model_options": {
-        "nnsk": {
-            "onsite": {"method": "none"},
-            "hopping": {"method": "powerlaw", "rs":1.6, "w": 0.3},
-            "freeze": false
-        }
-    },
-    "data_options": {
-        "train": {
-            "root": "./data/",
-            "prefix": "kpath",
-            "get_eigenvalues": true
-        }
-    }
+{
+    "common_options": {
+            "basis": {
+                "B": ["2s", "2p","d*"],
+                "N": ["2s", "2p","d*"]
+            },
+            "device": "cpu",
+            "dtype": "float32",
+            "overlap": false,
+            "seed": 42
+    },
+    "train_options": {
+        "num_epoch": 500,
+        "batch_size": 1,
+        "optimizer": {
+            "lr": 0.05,
+            "type": "Adam"
+        },
+        "lr_scheduler": {
+            "type": "exp",
+            "gamma": 0.999
+        },
+        "loss_options":{
+            "train": {"method": "eigvals"}
+        },
+        "save_freq": 50,
+        "validation_freq": 10,
+        "display_freq": 10
+    },
+    "model_options": {
+        "nnsk": {
+            "onsite": {"method": "none"},
+            "hopping": {"method": "powerlaw", "rs":1.6, "w": 0.3},
+            "soc":{},
+            "freeze": false,
+            "push":false
+        }
+    },
+    "data_options": {
+        "train": {
+            "root": "./data/",
+            "prefix": "kpath",
+            "get_eigenvalues": true
+        }
+    }
 }

examples/hBN/input/5_longtrain.json → examples/hBN/input/input_final.json

@@ -1,44 +1,46 @@
-{
-    "common_options": {
-            "basis": {
-                "B": ["2s", "2p", "d*"],
-                "N": ["2s", "2p", "d*"]
-            },
-            "device": "cpu",
-            "dtype": "float32",
-            "overlap": false,
-            "seed": 42
-    },
-    "train_options": {
-        "num_epoch": 10000,
-        "batch_size": 1,
-        "optimizer": {
-            "lr": 0.001,
-            "type": "Adam"
-        },
-        "lr_scheduler": {
-            "type": "exp",
-            "gamma": 0.999
-        },
-        "loss_options":{
-            "train": {"method": "eigvals"}
-        },
-        "save_freq": 1000,
-        "validation_freq": 10,
-        "display_freq": 10
-    },
-    "model_options": {
-        "nnsk": {
-            "onsite": {"method": "strain", "rs":1.6, "w":0.3},
-            "hopping": {"method": "powerlaw", "rs":3.6, "w": 0.3},
-            "freeze": false
-        }
-    },
-    "data_options": {
-        "train": {
-            "root": "./data/",
-            "prefix": "kpath",
-            "get_eigenvalues": true
-        }
-    }
+{
+    "common_options": {
+            "basis": {
+                "B": ["2s", "2p","d*"],
+                "N": ["2s", "2p","d*"]
+            },
+            "device": "cpu",
+            "dtype": "float32",
+            "overlap": false,
+            "seed": 42
+    },
+    "train_options": {
+        "num_epoch": 1000,
+        "batch_size": 1,
+        "optimizer": {
+            "lr": 0.05,
+            "type": "Adam"
+        },
+        "lr_scheduler": {
+            "type": "exp",
+            "gamma": 0.999
+        },
+        "loss_options":{
+            "train": {"method": "eigvals"}
+        },
+        "save_freq": 10,
+        "validation_freq": 10,
+        "display_freq": 10
+    },
+    "model_options": {
+        "nnsk": {
+            "onsite": {"method": "strain", "rs": 1.6, "w": 0.3},
+            "hopping": {"method": "powerlaw", "rs":3.4, "w": 0.2},
+            "soc":{},
+            "freeze": false,
+            "push": false
+        }
+    },
+    "data_options": {
+        "train": {
+            "root": "./data/",
+            "prefix": "kpath",
+            "get_eigenvalues": true
+        }
+    }
 }