QG-phy
diff --git a/‎README.md
Lines changed: 14 additions & 16 deletions b/‎README.md
Lines changed: 14 additions & 16 deletions
diff --git a/‎docs/quick_start/easy_install.md
Lines changed: 8 additions & 12 deletions b/‎docs/quick_start/easy_install.md
Lines changed: 8 additions & 12 deletions
diff --git a/‎dptb/data/transforms.py
Lines changed: 3 additions & 1 deletion b/‎dptb/data/transforms.py
Lines changed: 3 additions & 1 deletion
diff --git a/‎dptb/entrypoints/config.py
Lines changed: 11 additions & 11 deletions b/‎dptb/entrypoints/config.py
Lines changed: 11 additions & 11 deletions
diff --git a/‎dptb/entrypoints/emp_sk.py
Lines changed: 118 additions & 0 deletions b/‎dptb/entrypoints/emp_sk.py
Lines changed: 118 additions & 0 deletions
diff --git a/‎dptb/entrypoints/main.py
Lines changed: 38 additions & 1 deletion b/‎dptb/entrypoints/main.py
Lines changed: 38 additions & 1 deletion
@@ -66,42 +66,40 @@ Installing **DeePTB** is straightforward. We recommend using a virtual environme
 
     Highly recommended to install DeePTB from source to get the latest features and bug fixes.
   1. **Setup Python environment**:
-    
         Using conda (recommended, python >=3.9, <=3.12 ), e.g.,
         ```bash
         conda create -n dptb_venv python=3.10
         conda activate dptb_venv
         ```
         or using venv (make sure python >=3.9,<=3.12)
+    
         ```bash
         python -m venv dptb_venv
         source dptb_venv/bin/activate
+        ```
 
   2. **Clone DeePTB and  Navigate to the root directory**:
         ```bash
         git clone https://github.com/deepmodeling/DeePTB.git
         cd DeePTB
         ```
-  3. **Install `torch` and `torch-scatter`** (two ways):
-        - **Recommended**: Install torch and torch-scatter using the following commands:
 
+  3. **Install `torch`**:
+        ```bash
+        pip install "torch>=2.0.0,<=2.5.0"
+        ```
+  4. **Install `torch-scatter`** (two ways):
+        - **Recommended**: Install torch and torch-scatter using the following commands:
             ```bash
             python docs/auto_install_torch_scatter.py
             ```
-
         - **Manual**: Install torch and torch-scatter manually:
-          1. install torch:
-                ```bash
-                pip install "torch>=2.0.0,<=2.5.0"
-                ```
-
-          2. install torch-scatter:
-                ```bash
-                pip install torch-scatter -f https://data.pyg.org/whl/torch-${version}+${CUDA}.html
-                ```
-                where `${version}` is the version of torch, e.g., 2.5.0, and `${CUDA}` is the CUDA version, e.g., cpu, cu118, cu121, cu124. See [torch_scatter doc](https://github.com/rusty1s/pytorch_scatter) for more details.   
-
-  4. **Install DeePTB**:
+            ```bash
+            pip install torch-scatter -f https://data.pyg.org/whl/torch-${version}+${CUDA}.html
+            ```
+            where `${version}` is the version of torch, e.g., 2.5.0, and `${CUDA}` is the CUDA version, e.g., cpu, cu118, cu121, cu124. See [torch_scatter doc](https://github.com/rusty1s/pytorch_scatter) for more details.   
+
+  5. **Install DeePTB**:
         ```bash
         pip install .
         ```
 
@@ -29,32 +29,28 @@ Highly recommended to install DeePTB from source to get the latest features and
     ```bash
     python -m venv dptb_venv
     source dptb_venv/bin/activate
-
+    ```
 2. **Clone DeePTB and  Navigate to the root directory**:
     ```bash
     git clone https://github.com/deepmodeling/DeePTB.git
     cd DeePTB
     ```
-3. **Install `torch` and `torch-scatter`** (two ways):
-- **Recommended**: Install torch and torch-scatter using the following commands:
-
+3. **Install `torch`**:
     ```bash
-    python docs/auto_install_torch_scatter.py
+    pip install "torch>=2.0.0,<=2.5.0"
     ```
-
-- **Manual**: Install torch and torch-scatter manually:
-  1. install torch:
+4. **Install `torch-scatter`** (two ways):
+    - **Recommended**: Install torch and torch-scatter using the following commands:
         ```bash
-        pip install "torch>=2.0.0,<=2.5.0"
+         python docs/auto_install_torch_scatter.py
         ```
-
-  2. install torch-scatter:
+    - **Manual**: Install torch and torch-scatter manually:
         ```bash
         pip install torch-scatter -f https://data.pyg.org/whl/torch-${version}+${CUDA}.html
         ```
         where `${version}` is the version of torch, e.g., 2.5.0, and `${CUDA}` is the CUDA version, e.g., cpu, cu118, cu121, cu124. See [torch_scatter doc](https://github.com/rusty1s/pytorch_scatter) for more details.   
 
-4. **Install DeePTB**:   
+5. **Install DeePTB**:   
     ```bash
     pip install .
     ```
 
@@ -489,7 +489,9 @@ def __init__(
 
             for ko in orbtype_count.keys():
                 orbtype_count[ko] = max(orbtype_count[ko])
- 
+        else:
+            raise ValueError(f"Invalid basis {self.basis} found. now only support string or list basis.")
+            
         self.orbtype_count = orbtype_count
         full_basis_norb = 0
         for ko in orbtype_count.keys():
 
@@ -2,15 +2,13 @@
 import json
 from pathlib import Path
 import os
-from dptb.utils.config_sk import TrainFullConfigSK, TestFullConfigSK
-from dptb.utils.config_skenv import TrainFullConfigSKEnv, TestFullConfigSKEnv
-from dptb.utils.config_e3 import TrainFullConfigE3, TestFullConfigE3
 import logging
+from dptb.utils.gen_inputs import gen_inputs
 
 __all__ = ["get_full_config", "config"]
 log = logging.getLogger(__name__)
 
-def get_full_config(train, test, e3tb, sktb, sktbenv):
+def get_full_config(model, train, test, e3tb, sktb, sktbenv):
     """
     This function determines the appropriate full config based on the provided parameters.
 
@@ -31,16 +29,17 @@ def get_full_config(train, test, e3tb, sktb, sktbenv):
     name = ''
     if train:
         name += 'train'
+        
         # Use train configs based on e3tb, sktb, sktbenv
         if e3tb:
             name += '_E3'
-            full_config = TrainFullConfigE3
+            full_config = gen_inputs(mode='e3', task='train', model=model)
         elif sktb:
             name += '_SK'
-            full_config = TrainFullConfigSK
+            full_config = gen_inputs(mode='sk', task='train', model=model)
         elif sktbenv:
             name += '_SKEnv'
-            full_config = TrainFullConfigSKEnv
+            full_config = gen_inputs(mode='skenv', task='train', model=model)
         else:
             logging.error("Unknown config type in training mode")
             raise ValueError("Unknown config type in training mode")
@@ -49,13 +48,13 @@ def get_full_config(train, test, e3tb, sktb, sktbenv):
         name += 'test'
         if e3tb:
             name += '_E3'
-            full_config = TestFullConfigE3
+            full_config = gen_inputs(mode='e3', task='test', model=model)
         elif sktb:
             name += '_SK'
-            full_config = TestFullConfigSK
+            full_config = gen_inputs(mode='sk', task='test', model=model)
         elif sktbenv:
             name += '_SKEnv'
-            full_config = TestFullConfigSKEnv
+            full_config = gen_inputs(mode='skenv', task='test', model=model)
         else:
             logging.error("Unknown config type in testing mode")
             raise ValueError("Unknown config type in testing mode")
@@ -72,6 +71,7 @@ def config(
         e3tb: bool = False,
         sktb: bool = False,
         sktbenv: bool = False,
+        model: str = None,
         log_level: int = logging.INFO,
         log_path: Optional[str] = None,
         **kwargs
@@ -115,7 +115,7 @@ def config(
         train = True
 
     # Error handling and logic moved to get_full_config
-    name, full_config = get_full_config(train, test, e3tb, sktb, sktbenv)
+    name, full_config = get_full_config(model, train, test, e3tb, sktb, sktbenv)
     # Ensure PATH ends with .json
     if not PATH.endswith(".json"):
         PATH = os.path.join(PATH, "input_templete.json")
 
@@ -0,0 +1,118 @@
+import torch
+import numpy as np
+from dptb.nn.build import build_model
+import json
+import logging
+from dptb.nn.sktb.onsiteDB import onsite_energy_database
+import re
+import os
+from dptb.utils.gen_inputs import gen_inputs
+import json
+log = logging.getLogger(__name__)
+
+def to_empsk(
+    INPUT,
+    output='./', 
+    basemodel='poly2',
+    **kwargs):
+    """
+    Convert the model to empirical SK parameters.
+    """
+    if INPUT is None:
+        raise ValueError('INPUT is None.')
+    with open(INPUT, 'r') as f:
+        input = json.load(f)
+    common_options = input['common_options']
+    EmpSK(common_options, basemodel=basemodel).to_json(outdir=output)
+
+class EmpSK(object):
+    """
+    Empirical SK parameters.
+    """
+    def __init__(self, common_options, basemodel='poly2'):
+        """
+        Args:
+            common_options: common options for the model. especially contain the basis information.
+            basemodel: base model type for the empirical SK parameters  either 'poly2' or 'poly4'.
+        """
+        self.common_options,self.basisref = self.format_common_options(common_options)
+        if basemodel == 'poly2':
+            model_ckpt = os.path.join(os.path.dirname(__file__), '..', 'nn', 'dftb', "base_poly2.pth")
+        elif basemodel == 'poly4':
+            model_ckpt = os.path.join(os.path.dirname(__file__), '..', 'nn', 'dftb', "base_poly4.pth")
+        else:
+            raise ValueError(f'basemodel {basemodel} is not supported.')
+
+        self.model = build_model(model_ckpt, common_options=common_options, no_check=True)
+
+    def to_json(self, outdir='./'):
+        """
+        Convert the model to json format.
+        """
+        # 判断是否存在输出目录
+        if not os.path.exists(outdir):
+            os.makedirs(outdir, exist_ok=True)
+        json_dict = self.model.to_json(basisref=self.basisref)
+        with open(os.path.join(outdir,'sktb.json'), 'w') as f:
+            json.dump(json_dict, f, indent=4)
+        
+        # save input template
+        # input_template = gen_inputs(model=self.model, task='train', mode=mode)
+        
+        #with open(os.path.join(outdir,'input_template.json'), 'w') as f:
+        #    json.dump(input_template, f, indent=4)
+        log.info(f'Empirical SK parameters are saved in {os.path.join(outdir,"sktb.json")}')
+        log.info('If you want to further train the model, please use `dptb config` command to generate input template.')
+        return json_dict
+
+    def format_common_options(self, common_options):
+        """
+        Format the common options for the model. and construct the mapping between two kind of basis definition.
+        The two kind of basis definition are:
+            1. common_options = {'basis': {'C': ['s','p','d']}}
+            2. common_options = {'basis': {'C': ['2s','2p','d*']}}
+        
+        Args:
+            common_options: common options for the model. especially contain the basis information.
+            e.g. common_options = {'basis': {'C': ['s','p','d']}} or common_options = {'basis': {'C': ['2s','2p','d*']}}
+        
+        Returns:
+            common_options: common options for the model.
+            basisref: basis reference for the model.
+        """        
+        # check basis in common_options
+        if 'basis' not in common_options:
+            raise ValueError('basis information is not given in common_options.')
+        # check basis type
+        assert isinstance(common_options['basis'], dict), 'basis information is not a dictionary.'
+        basis = common_options['basis'] 
+        sys_ele =  "".join(list(basis.keys()))
+        log.info(f'Extracting empirical SK parameters for {sys_ele}')
+
+        use_basis_ref = False
+        basisref = {}
+        for ie in basis.keys():
+            basisref[ie] = {}
+            assert isinstance(basis[ie], list), f'basis information for {ie} is not a list.'
+            for ieorb in basis[ie]:
+                assert isinstance(ieorb, str), f'basis information for {ie} is not a string.'
+                if len(ieorb) == 1:
+                    assert use_basis_ref is False, 'Invalid basis setting: cannot mix s, p, d with ns, np, d*.'
+                    continue
+                else:
+                    use_basis_ref = True
+                    assert ieorb in onsite_energy_database[ie], f'basis information for {ie} is not in onsite_energy_database : {onsite_energy_database[ie].keys()}.'
+                    orbsymb = re.findall(r'[A-Za-z]', ieorb)[0]
+                    basisref[ie][orbsymb] = ieorb
+
+        if use_basis_ref:
+            std_basis = {}
+            for ie in basis.keys():
+                std_basis[ie] = []
+                for ieorb in basis[ie]:
+                    std_basis[ie].append(re.findall(r'[A-Za-z]', ieorb)[0])
+            common_options['basis'].update(std_basis)
+        
+            return common_options, basisref
+        else:
+            return common_options, None
@@ -13,6 +13,8 @@
 from dptb.utils.loggers import set_log_handles
 from dptb.utils.config_check import check_config_train
 from dptb.entrypoints.collectskf import skf2pth, skf2nnsk
+from dptb.entrypoints.emp_sk import to_empsk
+
 from dptb import __version__
 
 
@@ -86,6 +88,14 @@ def main_parser() -> argparse.ArgumentParser:
         default="./input_templete.json"
     )
 
+    parser_config.add_argument(
+        "-m", 
+        "--model",
+        type=str,
+        default=None,
+        help="load model to update input template."
+    )
+
     parser_config.add_argument(
         "-tr", 
         "--train",
@@ -394,7 +404,31 @@ def main_parser() -> argparse.ArgumentParser:
         help="The output files in training.",
     )
 
-
+    parser_esk = subparsers.add_parser(
+        "esk",
+        parents=[parser_log],
+        help="Generate initial empirical SK parameters.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser_esk.add_argument(
+        "INPUT", help="the input parameter file in json or yaml format",
+        type=str,
+        default=None
+    )
+    parser_esk.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        default="./",
+        help="The output files in training."
+    )
+    parser_esk.add_argument(
+        "-m",
+        "--basemodel",
+        type=str,
+        default="poly2",
+        help="The base model type can be poly2 or poly4."
+    )
     return parser
 
 def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
@@ -458,3 +492,6 @@ def main():
 
     elif args.command == 'skf2nn':
         skf2nnsk(**dict_args)
+
+    elif args.command == 'esk':
+        to_empsk(**dict_args)