From cded772645abdc79eaf3add1e8326e196d245ee5 Mon Sep 17 00:00:00 2001
From: Qiangqiang Gu <98570179+QG-phy@users.noreply.github.com>
Date: Thu, 11 Apr 2024 14:14:17 +0800
Subject: [PATCH] update DPTB to NNENV  (#115)

* update argcheck.py

* rename DPTB to  NNENV

* update argcheck.py

* chore: update info.json

* docs: update 1 file and create 1 file
---
 docs/index.rst                                |    5 +
 docs/input_params/input_tags.rst              | 2314 +++++++++++++++++
 dptb/nn/__init__.py                           |    4 +-
 dptb/nn/build.py                              |   18 +-
 dptb/nn/deeptb.py                             |   16 +-
 .../test_sktb/dataset/kpath_spk.0/info.json   |    4 +-
 dptb/tests/test_build_model.py                |    8 +-
 dptb/utils/argcheck.py                        |  316 ++-
 8 files changed, 2517 insertions(+), 168 deletions(-)
 create mode 100644 docs/input_params/input_tags.rst

diff --git a/docs/index.rst b/docs/index.rst
index bad55735..f0eb8e5e 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -32,6 +32,11 @@ These capabilities make **DeePTB** adaptable to various research scenarios, exte
    advanced/interface/index
    advanced/acceleration/index
 
+.. toctree::
+   :maxdepth: 2
+   :caption: INPUT TAG
+   
+   input_params/input_tags
 
 .. toctree::
    :maxdepth: 2
diff --git a/docs/input_params/input_tags.rst b/docs/input_params/input_tags.rst
new file mode 100644
index 00000000..3201deaf
--- /dev/null
+++ b/docs/input_params/input_tags.rst
@@ -0,0 +1,2314 @@
+.. _`common_options`: 
+
+common_options: 
+    | type: ``dict``
+    | argument path: ``common_options``
+
+    .. _`common_options/basis`: 
+
+    basis: 
+        | type: ``dict``
+        | argument path: ``common_options/basis``
+
+        The atomic orbitals used to construct the basis. e.p. {'A':['2s','2p','s*'],'B':'[3s','3p']}
+
+    .. _`common_options/overlap`: 
+
+    overlap: 
+        | type: ``bool``, optional, default: ``False``
+        | argument path: ``common_options/overlap``
+
+        Whether to calculate the overlap matrix. Default: False
+
+    .. _`common_options/device`: 
+
+    device: 
+        | type: ``str``, optional, default: ``cpu``
+        | argument path: ``common_options/device``
+
+        The device to run the calculation, choose among `cpu` and `cuda[:int]`, Default: `cpu`
+
+    .. _`common_options/dtype`: 
+
+    dtype: 
+        | type: ``str``, optional, default: ``float32``
+        | argument path: ``common_options/dtype``
+
+        The digital number's precison, choose among: 
+                            Default: `float32`
+                                - `float32`: indicating torch.float32
+                                - `float64`: indicating torch.float64
+                
+
+    .. _`common_options/seed`: 
+
+    seed: 
+        | type: ``int``, optional, default: ``3982377700``
+        | argument path: ``common_options/seed``
+
+        The random seed used to initialize the parameters and determine the shuffling order of datasets. Default: `3982377700`
+
+
+.. _`train_options`: 
+
+train_options: 
+    | type: ``dict``, optional
+    | argument path: ``train_options``
+
+    Options that defines the training behaviour of DeePTB.
+
+    .. _`train_options/num_epoch`: 
+
+    num_epoch: 
+        | type: ``int``
+        | argument path: ``train_options/num_epoch``
+
+        Total number of training epochs. It is worth noted, if the model is reloaded with `-r` or `--restart` option, epoch which have been trained will counted from the time that the checkpoint is saved.
+
+    .. _`train_options/batch_size`: 
+
+    batch_size: 
+        | type: ``int``, optional, default: ``1``
+        | argument path: ``train_options/batch_size``
+
+        The batch size used in training, Default: 1
+
+    .. _`train_options/ref_batch_size`: 
+
+    ref_batch_size: 
+        | type: ``int``, optional, default: ``1``
+        | argument path: ``train_options/ref_batch_size``
+
+        The batch size used in reference data, Default: 1
+
+    .. _`train_options/val_batch_size`: 
+
+    val_batch_size: 
+        | type: ``int``, optional, default: ``1``
+        | argument path: ``train_options/val_batch_size``
+
+        The batch size used in validation data, Default: 1
+
+    .. _`train_options/optimizer`: 
+
+    optimizer: 
+        | type: ``dict``, optional, default: ``{}``
+        | argument path: ``train_options/optimizer``
+
+                The optimizer setting for selecting the gradient optimizer of model training. Optimizer supported includes `Adam`, `SGD` and `LBFGS` 
+
+                For more information about these optmization algorithm, we refer to:
+
+                - `Adam`: [Adam: A Method for Stochastic Optimization.](https://arxiv.org/abs/1412.6980)
+
+                - `SGD`: [Stochastic Gradient Descent.](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html)
+
+                - `LBFGS`: [On the limited memory BFGS method for large scale optimization.](http://users.iems.northwestern.edu/~nocedal/PDFfiles/limited-memory.pdf) 
+
+    
+
+
+        Depending on the value of *type*, different sub args are accepted. 
+
+        .. _`train_options/optimizer/type`: 
+
+        type:
+            | type: ``str`` (flag key), default: ``Adam``
+            | argument path: ``train_options/optimizer/type`` 
+            | possible choices: |code:train_options/optimizer[Adam]|_, |code:train_options/optimizer[SGD]|_
+
+            select type of optimizer, support type includes: `Adam`, `SGD` and `LBFGS`. Default: `Adam`
+
+            .. |code:train_options/optimizer[Adam]| replace:: ``Adam``
+            .. _`code:train_options/optimizer[Adam]`: `train_options/optimizer[Adam]`_
+            .. |code:train_options/optimizer[SGD]| replace:: ``SGD``
+            .. _`code:train_options/optimizer[SGD]`: `train_options/optimizer[SGD]`_
+
+        .. |flag:train_options/optimizer/type| replace:: *type*
+        .. _`flag:train_options/optimizer/type`: `train_options/optimizer/type`_
+
+
+        .. _`train_options/optimizer[Adam]`: 
+
+        When |flag:train_options/optimizer/type|_ is set to ``Adam``: 
+
+        .. _`train_options/optimizer[Adam]/lr`: 
+
+        lr: 
+            | type: ``float``, optional, default: ``0.001``
+            | argument path: ``train_options/optimizer[Adam]/lr``
+
+            learning rate. Default: 1e-3
+
+        .. _`train_options/optimizer[Adam]/betas`: 
+
+        betas: 
+            | type: ``list``, optional, default: ``[0.9, 0.999]``
+            | argument path: ``train_options/optimizer[Adam]/betas``
+
+            coefficients used for computing running averages of gradient and its square Default: (0.9, 0.999)
+
+        .. _`train_options/optimizer[Adam]/eps`: 
+
+        eps: 
+            | type: ``float``, optional, default: ``1e-08``
+            | argument path: ``train_options/optimizer[Adam]/eps``
+
+            term added to the denominator to improve numerical stability, Default: 1e-8
+
+        .. _`train_options/optimizer[Adam]/weight_decay`: 
+
+        weight_decay: 
+            | type: ``float``, optional, default: ``0``
+            | argument path: ``train_options/optimizer[Adam]/weight_decay``
+
+            weight decay (L2 penalty), Default: 0
+
+        .. _`train_options/optimizer[Adam]/amsgrad`: 
+
+        amsgrad: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``train_options/optimizer[Adam]/amsgrad``
+
+            whether to use the AMSGrad variant of this algorithm from the paper On the [Convergence of Adam and Beyond](https://openreview.net/forum?id=ryQu7f-RZ) ,Default: False
+
+
+        .. _`train_options/optimizer[SGD]`: 
+
+        When |flag:train_options/optimizer/type|_ is set to ``SGD``: 
+
+        .. _`train_options/optimizer[SGD]/lr`: 
+
+        lr: 
+            | type: ``float``, optional, default: ``0.001``
+            | argument path: ``train_options/optimizer[SGD]/lr``
+
+            learning rate. Default: 1e-3
+
+        .. _`train_options/optimizer[SGD]/momentum`: 
+
+        momentum: 
+            | type: ``float``, optional, default: ``0.0``
+            | argument path: ``train_options/optimizer[SGD]/momentum``
+
+            momentum factor Default: 0
+
+        .. _`train_options/optimizer[SGD]/weight_decay`: 
+
+        weight_decay: 
+            | type: ``float``, optional, default: ``0.0``
+            | argument path: ``train_options/optimizer[SGD]/weight_decay``
+
+            weight decay (L2 penalty), Default: 0
+
+        .. _`train_options/optimizer[SGD]/dampening`: 
+
+        dampening: 
+            | type: ``float``, optional, default: ``0.0``
+            | argument path: ``train_options/optimizer[SGD]/dampening``
+
+            dampening for momentum, Default: 0
+
+        .. _`train_options/optimizer[SGD]/nesterov`: 
+
+        nesterov: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``train_options/optimizer[SGD]/nesterov``
+
+            enables Nesterov momentum, Default: False
+
+    .. _`train_options/lr_scheduler`: 
+
+    lr_scheduler: 
+        | type: ``dict``, optional, default: ``{}``
+        | argument path: ``train_options/lr_scheduler``
+
+        The learning rate scheduler tools settings, the lr scheduler is used to scales down the learning rate during the training process. Proper setting can make the training more stable and efficient. The supported lr schedular includes: `Exponential Decaying (exp)`, `Linear multiplication (linear)`
+
+
+        Depending on the value of *type*, different sub args are accepted. 
+
+        .. _`train_options/lr_scheduler/type`: 
+
+        type:
+            | type: ``str`` (flag key), default: ``exp``
+            | argument path: ``train_options/lr_scheduler/type`` 
+            | possible choices: |code:train_options/lr_scheduler[exp]|_, |code:train_options/lr_scheduler[linear]|_, |code:train_options/lr_scheduler[rop]|_
+
+            select type of lr_scheduler, support type includes `exp`, `linear`
+
+            .. |code:train_options/lr_scheduler[exp]| replace:: ``exp``
+            .. _`code:train_options/lr_scheduler[exp]`: `train_options/lr_scheduler[exp]`_
+            .. |code:train_options/lr_scheduler[linear]| replace:: ``linear``
+            .. _`code:train_options/lr_scheduler[linear]`: `train_options/lr_scheduler[linear]`_
+            .. |code:train_options/lr_scheduler[rop]| replace:: ``rop``
+            .. _`code:train_options/lr_scheduler[rop]`: `train_options/lr_scheduler[rop]`_
+
+        .. |flag:train_options/lr_scheduler/type| replace:: *type*
+        .. _`flag:train_options/lr_scheduler/type`: `train_options/lr_scheduler/type`_
+
+
+        .. _`train_options/lr_scheduler[exp]`: 
+
+        When |flag:train_options/lr_scheduler/type|_ is set to ``exp``: 
+
+        .. _`train_options/lr_scheduler[exp]/gamma`: 
+
+        gamma: 
+            | type: ``float``, optional, default: ``0.999``
+            | argument path: ``train_options/lr_scheduler[exp]/gamma``
+
+            Multiplicative factor of learning rate decay.
+
+
+        .. _`train_options/lr_scheduler[linear]`: 
+
+        When |flag:train_options/lr_scheduler/type|_ is set to ``linear``: 
+
+        .. _`train_options/lr_scheduler[linear]/start_factor`: 
+
+        start_factor: 
+            | type: ``float``, optional, default: ``0.3333333``
+            | argument path: ``train_options/lr_scheduler[linear]/start_factor``
+
+            The number we multiply learning rate in the first epoch.         The multiplication factor changes towards end_factor in the following epochs. Default: 1./3.
+
+        .. _`train_options/lr_scheduler[linear]/end_factor`: 
+
+        end_factor: 
+            | type: ``float``, optional, default: ``0.3333333``
+            | argument path: ``train_options/lr_scheduler[linear]/end_factor``
+
+            The number we multiply learning rate in the first epoch.     The multiplication factor changes towards end_factor in the following epochs. Default: 1./3.
+
+        .. _`train_options/lr_scheduler[linear]/total_iters`: 
+
+        total_iters: 
+            | type: ``int``, optional, default: ``5``
+            | argument path: ``train_options/lr_scheduler[linear]/total_iters``
+
+            The number of iterations that multiplicative factor reaches to 1. Default: 5.
+
+
+        .. _`train_options/lr_scheduler[rop]`: 
+
+        When |flag:train_options/lr_scheduler/type|_ is set to ``rop``: 
+
+        rop: reduce on plateau
+
+        .. _`train_options/lr_scheduler[rop]/mode`: 
+
+        mode: 
+            | type: ``str``, optional, default: ``min``
+            | argument path: ``train_options/lr_scheduler[rop]/mode``
+
+            One of min, max. In min mode, lr will be reduced when the quantity monitored has stopped decreasing;         in max mode it will be reduced when the quantity monitored has stopped increasing. Default: 'min'.
+
+        .. _`train_options/lr_scheduler[rop]/factor`: 
+
+        factor: 
+            | type: ``float``, optional, default: ``0.1``
+            | argument path: ``train_options/lr_scheduler[rop]/factor``
+
+            Factor by which the learning rate will be reduced. new_lr = lr * factor. Default: 0.1.
+
+        .. _`train_options/lr_scheduler[rop]/patience`: 
+
+        patience: 
+            | type: ``int``, optional, default: ``10``
+            | argument path: ``train_options/lr_scheduler[rop]/patience``
+
+            Number of epochs with no improvement after which learning rate will be reduced. For example,         if patience = 2, then we will ignore the first 2 epochs with no improvement,         and will only decrease the LR after the 3rd epoch if the loss still hasn't improved then. Default: 10.
+
+        .. _`train_options/lr_scheduler[rop]/threshold`: 
+
+        threshold: 
+            | type: ``float``, optional, default: ``0.0001``
+            | argument path: ``train_options/lr_scheduler[rop]/threshold``
+
+            Threshold for measuring the new optimum, to only focus on significant changes. Default: 1e-4.
+
+        .. _`train_options/lr_scheduler[rop]/threshold_mode`: 
+
+        threshold_mode: 
+            | type: ``str``, optional, default: ``rel``
+            | argument path: ``train_options/lr_scheduler[rop]/threshold_mode``
+
+            One of rel, abs. In rel mode, dynamic_threshold = best * ( 1 + threshold ) in 'max' mode or         best * ( 1 - threshold ) in min mode. In abs mode,         dynamic_threshold = best + threshold in max mode or best - threshold in min mode. Default: 'rel'.
+
+        .. _`train_options/lr_scheduler[rop]/cooldown`: 
+
+        cooldown: 
+            | type: ``int``, optional, default: ``0``
+            | argument path: ``train_options/lr_scheduler[rop]/cooldown``
+
+            Number of epochs to wait before resuming normal operation after lr has been reduced. Default: 0.
+
+        .. _`train_options/lr_scheduler[rop]/min_lr`: 
+
+        min_lr: 
+            | type: ``list`` | ``float``, optional, default: ``0``
+            | argument path: ``train_options/lr_scheduler[rop]/min_lr``
+
+            A scalar or a list of scalars.         A lower bound on the learning rate of all param groups or each group respectively. Default: 0.
+
+        .. _`train_options/lr_scheduler[rop]/eps`: 
+
+        eps: 
+            | type: ``float``, optional, default: ``1e-08``
+            | argument path: ``train_options/lr_scheduler[rop]/eps``
+
+            Minimal decay applied to lr.         If the difference between new and old lr is smaller than eps, the update is ignored. Default: 1e-8.
+
+    .. _`train_options/save_freq`: 
+
+    save_freq: 
+        | type: ``int``, optional, default: ``10``
+        | argument path: ``train_options/save_freq``
+
+        Frequency, or every how many iteration to saved the current model into checkpoints, The name of checkpoint is formulated as `latest|best_dptb|nnsk_b<bond_cutoff>_c<sk_cutoff>_w<sk_decay_w>`. Default: `10`
+
+    .. _`train_options/validation_freq`: 
+
+    validation_freq: 
+        | type: ``int``, optional, default: ``10``
+        | argument path: ``train_options/validation_freq``
+
+        Frequency or every how many iteration to do model validation on validation datasets. Default: `10`
+
+    .. _`train_options/display_freq`: 
+
+    display_freq: 
+        | type: ``int``, optional, default: ``1``
+        | argument path: ``train_options/display_freq``
+
+        Frequency, or every how many iteration to display the training log to screem. Default: `1`
+
+    .. _`train_options/max_ckpt`: 
+
+    max_ckpt: 
+        | type: ``int``, optional, default: ``4``
+        | argument path: ``train_options/max_ckpt``
+
+        The maximum number of saved checkpoints, Default: 4
+
+    .. _`train_options/loss_options`: 
+
+    loss_options: 
+        | type: ``dict``
+        | argument path: ``train_options/loss_options``
+
+        .. _`train_options/loss_options/train`: 
+
+        train: 
+            | type: ``dict``
+            | argument path: ``train_options/loss_options/train``
+
+            Loss options for training.
+
+
+            Depending on the value of *method*, different sub args are accepted. 
+
+            .. _`train_options/loss_options/train/method`: 
+
+            method:
+                | type: ``str`` (flag key)
+                | argument path: ``train_options/loss_options/train/method`` 
+                | possible choices: |code:train_options/loss_options/train[hamil]|_, |code:train_options/loss_options/train[eigvals]|_, |code:train_options/loss_options/train[hamil_abs]|_, |code:train_options/loss_options/train[hamil_blas]|_
+
+                The loss function type, defined by a string like `<fitting target>_<loss type>`, Default: `eigs_l2dsf`. supported loss functions includes:
+
+                                    - `eigvals`: The mse loss predicted and labeled eigenvalues and Delta eigenvalues between different k.
+                                    - `hamil`: 
+                                    - `hamil_abs`:
+                                    - `hamil_blas`:
+                
+
+                .. |code:train_options/loss_options/train[hamil]| replace:: ``hamil``
+                .. _`code:train_options/loss_options/train[hamil]`: `train_options/loss_options/train[hamil]`_
+                .. |code:train_options/loss_options/train[eigvals]| replace:: ``eigvals``
+                .. _`code:train_options/loss_options/train[eigvals]`: `train_options/loss_options/train[eigvals]`_
+                .. |code:train_options/loss_options/train[hamil_abs]| replace:: ``hamil_abs``
+                .. _`code:train_options/loss_options/train[hamil_abs]`: `train_options/loss_options/train[hamil_abs]`_
+                .. |code:train_options/loss_options/train[hamil_blas]| replace:: ``hamil_blas``
+                .. _`code:train_options/loss_options/train[hamil_blas]`: `train_options/loss_options/train[hamil_blas]`_
+
+            .. |flag:train_options/loss_options/train/method| replace:: *method*
+            .. _`flag:train_options/loss_options/train/method`: `train_options/loss_options/train/method`_
+
+
+            .. _`train_options/loss_options/train[hamil]`: 
+
+            When |flag:train_options/loss_options/train/method|_ is set to ``hamil``: 
+
+            .. _`train_options/loss_options/train[hamil]/onsite_shift`: 
+
+            onsite_shift: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/train[hamil]/onsite_shift``
+
+                Whether to use onsite shift in loss function. Default: False
+
+
+            .. _`train_options/loss_options/train[eigvals]`: 
+
+            When |flag:train_options/loss_options/train/method|_ is set to ``eigvals``: 
+
+            .. _`train_options/loss_options/train[eigvals]/diff_on`: 
+
+            diff_on: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/train[eigvals]/diff_on``
+
+                Whether to use random differences in loss function. Default: False
+
+            .. _`train_options/loss_options/train[eigvals]/eout_weight`: 
+
+            eout_weight: 
+                | type: ``float``, optional, default: ``0.01``
+                | argument path: ``train_options/loss_options/train[eigvals]/eout_weight``
+
+                The weight of eigenvalue out of range. Default: 0.01
+
+            .. _`train_options/loss_options/train[eigvals]/diff_weight`: 
+
+            diff_weight: 
+                | type: ``float``, optional, default: ``0.01``
+                | argument path: ``train_options/loss_options/train[eigvals]/diff_weight``
+
+                The weight of eigenvalue difference. Default: 0.01
+
+
+            .. _`train_options/loss_options/train[hamil_abs]`: 
+
+            When |flag:train_options/loss_options/train/method|_ is set to ``hamil_abs``: 
+
+            .. _`train_options/loss_options/train[hamil_abs]/onsite_shift`: 
+
+            onsite_shift: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/train[hamil_abs]/onsite_shift``
+
+                Whether to use onsite shift in loss function. Default: False
+
+
+            .. _`train_options/loss_options/train[hamil_blas]`: 
+
+            When |flag:train_options/loss_options/train/method|_ is set to ``hamil_blas``: 
+
+            .. _`train_options/loss_options/train[hamil_blas]/onsite_shift`: 
+
+            onsite_shift: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/train[hamil_blas]/onsite_shift``
+
+                Whether to use onsite shift in loss function. Default: False
+
+        .. _`train_options/loss_options/validation`: 
+
+        validation: 
+            | type: ``dict``, optional
+            | argument path: ``train_options/loss_options/validation``
+
+            Loss options for validation.
+
+
+            Depending on the value of *method*, different sub args are accepted. 
+
+            .. _`train_options/loss_options/validation/method`: 
+
+            method:
+                | type: ``str`` (flag key)
+                | argument path: ``train_options/loss_options/validation/method`` 
+                | possible choices: |code:train_options/loss_options/validation[hamil]|_, |code:train_options/loss_options/validation[eigvals]|_, |code:train_options/loss_options/validation[hamil_abs]|_, |code:train_options/loss_options/validation[hamil_blas]|_
+
+                The loss function type, defined by a string like `<fitting target>_<loss type>`, Default: `eigs_l2dsf`. supported loss functions includes:
+
+                                    - `eigvals`: The mse loss predicted and labeled eigenvalues and Delta eigenvalues between different k.
+                                    - `hamil`: 
+                                    - `hamil_abs`:
+                                    - `hamil_blas`:
+                
+
+                .. |code:train_options/loss_options/validation[hamil]| replace:: ``hamil``
+                .. _`code:train_options/loss_options/validation[hamil]`: `train_options/loss_options/validation[hamil]`_
+                .. |code:train_options/loss_options/validation[eigvals]| replace:: ``eigvals``
+                .. _`code:train_options/loss_options/validation[eigvals]`: `train_options/loss_options/validation[eigvals]`_
+                .. |code:train_options/loss_options/validation[hamil_abs]| replace:: ``hamil_abs``
+                .. _`code:train_options/loss_options/validation[hamil_abs]`: `train_options/loss_options/validation[hamil_abs]`_
+                .. |code:train_options/loss_options/validation[hamil_blas]| replace:: ``hamil_blas``
+                .. _`code:train_options/loss_options/validation[hamil_blas]`: `train_options/loss_options/validation[hamil_blas]`_
+
+            .. |flag:train_options/loss_options/validation/method| replace:: *method*
+            .. _`flag:train_options/loss_options/validation/method`: `train_options/loss_options/validation/method`_
+
+
+            .. _`train_options/loss_options/validation[hamil]`: 
+
+            When |flag:train_options/loss_options/validation/method|_ is set to ``hamil``: 
+
+            .. _`train_options/loss_options/validation[hamil]/onsite_shift`: 
+
+            onsite_shift: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/validation[hamil]/onsite_shift``
+
+                Whether to use onsite shift in loss function. Default: False
+
+
+            .. _`train_options/loss_options/validation[eigvals]`: 
+
+            When |flag:train_options/loss_options/validation/method|_ is set to ``eigvals``: 
+
+            .. _`train_options/loss_options/validation[eigvals]/diff_on`: 
+
+            diff_on: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/validation[eigvals]/diff_on``
+
+                Whether to use random differences in loss function. Default: False
+
+            .. _`train_options/loss_options/validation[eigvals]/eout_weight`: 
+
+            eout_weight: 
+                | type: ``float``, optional, default: ``0.01``
+                | argument path: ``train_options/loss_options/validation[eigvals]/eout_weight``
+
+                The weight of eigenvalue out of range. Default: 0.01
+
+            .. _`train_options/loss_options/validation[eigvals]/diff_weight`: 
+
+            diff_weight: 
+                | type: ``float``, optional, default: ``0.01``
+                | argument path: ``train_options/loss_options/validation[eigvals]/diff_weight``
+
+                The weight of eigenvalue difference. Default: 0.01
+
+
+            .. _`train_options/loss_options/validation[hamil_abs]`: 
+
+            When |flag:train_options/loss_options/validation/method|_ is set to ``hamil_abs``: 
+
+            .. _`train_options/loss_options/validation[hamil_abs]/onsite_shift`: 
+
+            onsite_shift: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/validation[hamil_abs]/onsite_shift``
+
+                Whether to use onsite shift in loss function. Default: False
+
+
+            .. _`train_options/loss_options/validation[hamil_blas]`: 
+
+            When |flag:train_options/loss_options/validation/method|_ is set to ``hamil_blas``: 
+
+            .. _`train_options/loss_options/validation[hamil_blas]/onsite_shift`: 
+
+            onsite_shift: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/validation[hamil_blas]/onsite_shift``
+
+                Whether to use onsite shift in loss function. Default: False
+
+        .. _`train_options/loss_options/reference`: 
+
+        reference: 
+            | type: ``dict``, optional
+            | argument path: ``train_options/loss_options/reference``
+
+            Loss options for reference data in training.
+
+
+            Depending on the value of *method*, different sub args are accepted. 
+
+            .. _`train_options/loss_options/reference/method`: 
+
+            method:
+                | type: ``str`` (flag key)
+                | argument path: ``train_options/loss_options/reference/method`` 
+                | possible choices: |code:train_options/loss_options/reference[hamil]|_, |code:train_options/loss_options/reference[eigvals]|_, |code:train_options/loss_options/reference[hamil_abs]|_, |code:train_options/loss_options/reference[hamil_blas]|_
+
+                The loss function type, defined by a string like `<fitting target>_<loss type>`, Default: `eigs_l2dsf`. supported loss functions includes:
+
+                                    - `eigvals`: The mse loss predicted and labeled eigenvalues and Delta eigenvalues between different k.
+                                    - `hamil`: 
+                                    - `hamil_abs`:
+                                    - `hamil_blas`:
+                
+
+                .. |code:train_options/loss_options/reference[hamil]| replace:: ``hamil``
+                .. _`code:train_options/loss_options/reference[hamil]`: `train_options/loss_options/reference[hamil]`_
+                .. |code:train_options/loss_options/reference[eigvals]| replace:: ``eigvals``
+                .. _`code:train_options/loss_options/reference[eigvals]`: `train_options/loss_options/reference[eigvals]`_
+                .. |code:train_options/loss_options/reference[hamil_abs]| replace:: ``hamil_abs``
+                .. _`code:train_options/loss_options/reference[hamil_abs]`: `train_options/loss_options/reference[hamil_abs]`_
+                .. |code:train_options/loss_options/reference[hamil_blas]| replace:: ``hamil_blas``
+                .. _`code:train_options/loss_options/reference[hamil_blas]`: `train_options/loss_options/reference[hamil_blas]`_
+
+            .. |flag:train_options/loss_options/reference/method| replace:: *method*
+            .. _`flag:train_options/loss_options/reference/method`: `train_options/loss_options/reference/method`_
+
+
+            .. _`train_options/loss_options/reference[hamil]`: 
+
+            When |flag:train_options/loss_options/reference/method|_ is set to ``hamil``: 
+
+            .. _`train_options/loss_options/reference[hamil]/onsite_shift`: 
+
+            onsite_shift: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/reference[hamil]/onsite_shift``
+
+                Whether to use onsite shift in loss function. Default: False
+
+
+            .. _`train_options/loss_options/reference[eigvals]`: 
+
+            When |flag:train_options/loss_options/reference/method|_ is set to ``eigvals``: 
+
+            .. _`train_options/loss_options/reference[eigvals]/diff_on`: 
+
+            diff_on: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/reference[eigvals]/diff_on``
+
+                Whether to use random differences in loss function. Default: False
+
+            .. _`train_options/loss_options/reference[eigvals]/eout_weight`: 
+
+            eout_weight: 
+                | type: ``float``, optional, default: ``0.01``
+                | argument path: ``train_options/loss_options/reference[eigvals]/eout_weight``
+
+                The weight of eigenvalue out of range. Default: 0.01
+
+            .. _`train_options/loss_options/reference[eigvals]/diff_weight`: 
+
+            diff_weight: 
+                | type: ``float``, optional, default: ``0.01``
+                | argument path: ``train_options/loss_options/reference[eigvals]/diff_weight``
+
+                The weight of eigenvalue difference. Default: 0.01
+
+
+            .. _`train_options/loss_options/reference[hamil_abs]`: 
+
+            When |flag:train_options/loss_options/reference/method|_ is set to ``hamil_abs``: 
+
+            .. _`train_options/loss_options/reference[hamil_abs]/onsite_shift`: 
+
+            onsite_shift: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/reference[hamil_abs]/onsite_shift``
+
+                Whether to use onsite shift in loss function. Default: False
+
+
+            .. _`train_options/loss_options/reference[hamil_blas]`: 
+
+            When |flag:train_options/loss_options/reference/method|_ is set to ``hamil_blas``: 
+
+            .. _`train_options/loss_options/reference[hamil_blas]/onsite_shift`: 
+
+            onsite_shift: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``train_options/loss_options/reference[hamil_blas]/onsite_shift``
+
+                Whether to use onsite shift in loss function. Default: False
+
+
+.. _`data_options`: 
+
+data_options: 
+    | type: ``dict``
+    | argument path: ``data_options``
+
+    The options for dataset settings in training.
+
+    .. _`data_options/train`: 
+
+    train: 
+        | type: ``dict``
+        | argument path: ``data_options/train``
+
+        The dataset settings for training.
+
+        .. _`data_options/train/type`: 
+
+        type: 
+            | type: ``str``, optional, default: ``DefaultDataset``
+            | argument path: ``data_options/train/type``
+
+            The type of dataset.
+
+        .. _`data_options/train/root`: 
+
+        root: 
+            | type: ``str``
+            | argument path: ``data_options/train/root``
+
+            This is where the dataset stores data files.
+
+        .. _`data_options/train/prefix`: 
+
+        prefix: 
+            | type: ``str`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``data_options/train/prefix``
+
+            The prefix of the folders under root, which will be loaded in dataset.
+
+        .. _`data_options/train/separator`: 
+
+        separator: 
+            | type: ``str``, optional, default: ``.``
+            | argument path: ``data_options/train/separator``
+
+            the sepatator used to separate the prefix and suffix in the dataset directory. Default: '.'
+
+        .. _`data_options/train/get_Hamiltonian`: 
+
+        get_Hamiltonian: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/train/get_Hamiltonian``
+
+            Choose whether the Hamiltonian blocks (and overlap blocks, if provided) are loaded when building dataset.
+
+        .. _`data_options/train/get_overlap`: 
+
+        get_overlap: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/train/get_overlap``
+
+            Choose whether the overlap blocks are loaded when building dataset.
+
+        .. _`data_options/train/get_DM`: 
+
+        get_DM: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/train/get_DM``
+
+            Choose whether the density matrix is loaded when building dataset.
+
+        .. _`data_options/train/get_eigenvalues`: 
+
+        get_eigenvalues: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/train/get_eigenvalues``
+
+            Choose whether the eigenvalues and k-points are loaded when building dataset.
+
+    .. _`data_options/validation`: 
+
+    validation: 
+        | type: ``dict``, optional
+        | argument path: ``data_options/validation``
+
+        The dataset settings for validation.
+
+        .. _`data_options/validation/type`: 
+
+        type: 
+            | type: ``str``, optional, default: ``DefaultDataset``
+            | argument path: ``data_options/validation/type``
+
+            The type of dataset.
+
+        .. _`data_options/validation/root`: 
+
+        root: 
+            | type: ``str``
+            | argument path: ``data_options/validation/root``
+
+            This is where the dataset stores data files.
+
+        .. _`data_options/validation/prefix`: 
+
+        prefix: 
+            | type: ``str`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``data_options/validation/prefix``
+
+            The prefix of the folders under root, which will be loaded in dataset.
+
+        .. _`data_options/validation/separator`: 
+
+        separator: 
+            | type: ``str``, optional, default: ``.``
+            | argument path: ``data_options/validation/separator``
+
+            the sepatator used to separate the prefix and suffix in the dataset directory. Default: '.'
+
+        .. _`data_options/validation/get_Hamiltonian`: 
+
+        get_Hamiltonian: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/validation/get_Hamiltonian``
+
+            Choose whether the Hamiltonian blocks (and overlap blocks, if provided) are loaded when building dataset.
+
+        .. _`data_options/validation/get_overlap`: 
+
+        get_overlap: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/validation/get_overlap``
+
+            Choose whether the overlap blocks are loaded when building dataset.
+
+        .. _`data_options/validation/get_DM`: 
+
+        get_DM: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/validation/get_DM``
+
+            Choose whether the density matrix is loaded when building dataset.
+
+        .. _`data_options/validation/get_eigenvalues`: 
+
+        get_eigenvalues: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/validation/get_eigenvalues``
+
+            Choose whether the eigenvalues and k-points are loaded when building dataset.
+
+    .. _`data_options/reference`: 
+
+    reference: 
+        | type: ``dict``, optional
+        | argument path: ``data_options/reference``
+
+        The dataset settings for reference.
+
+        .. _`data_options/reference/type`: 
+
+        type: 
+            | type: ``str``, optional, default: ``DefaultDataset``
+            | argument path: ``data_options/reference/type``
+
+            The type of dataset.
+
+        .. _`data_options/reference/root`: 
+
+        root: 
+            | type: ``str``
+            | argument path: ``data_options/reference/root``
+
+            This is where the dataset stores data files.
+
+        .. _`data_options/reference/prefix`: 
+
+        prefix: 
+            | type: ``str`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``data_options/reference/prefix``
+
+            The prefix of the folders under root, which will be loaded in dataset.
+
+        .. _`data_options/reference/separator`: 
+
+        separator: 
+            | type: ``str``, optional, default: ``.``
+            | argument path: ``data_options/reference/separator``
+
+            the sepatator used to separate the prefix and suffix in the dataset directory. Default: '.'
+
+        .. _`data_options/reference/get_Hamiltonian`: 
+
+        get_Hamiltonian: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/reference/get_Hamiltonian``
+
+            Choose whether the Hamiltonian blocks (and overlap blocks, if provided) are loaded when building dataset.
+
+        .. _`data_options/reference/get_overlap`: 
+
+        get_overlap: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/reference/get_overlap``
+
+            Choose whether the overlap blocks are loaded when building dataset.
+
+        .. _`data_options/reference/get_DM`: 
+
+        get_DM: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/reference/get_DM``
+
+            Choose whether the density matrix is loaded when building dataset.
+
+        .. _`data_options/reference/get_eigenvalues`: 
+
+        get_eigenvalues: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``data_options/reference/get_eigenvalues``
+
+            Choose whether the eigenvalues and k-points are loaded when building dataset.
+
+
+.. _`model_options`: 
+
+model_options: 
+    | type: ``dict``, optional
+    | argument path: ``model_options``
+
+    The parameters to define the `nnsk`,`mix` and `dptb` model.
+
+    .. _`model_options/embedding`: 
+
+    embedding: 
+        | type: ``dict``, optional
+        | argument path: ``model_options/embedding``
+
+        The parameters to define the embedding model.
+
+
+        Depending on the value of *method*, different sub args are accepted. 
+
+        .. _`model_options/embedding/method`: 
+
+        method:
+            | type: ``str`` (flag key), default: ``se2``
+            | argument path: ``model_options/embedding/method`` 
+            | possible choices: |code:model_options/embedding[se2]|_, |code:model_options/embedding[baseline]|_, |code:model_options/embedding[deeph-e3]|_, |code:model_options/embedding[e3baseline_0]|_, |code:model_options/embedding[e3baseline_1]|_, |code:model_options/embedding[e3baseline_2]|_, |code:model_options/embedding[e3baseline_3]|_, |code:model_options/embedding[e3baseline_4]|_, |code:model_options/embedding[e3baseline_5]|_, |code:model_options/embedding[e3baseline_6]|_, |code:model_options/embedding[e3baseline_nonlocal]|_
+
+            The parameters to define the embedding model.
+
+            .. |code:model_options/embedding[se2]| replace:: ``se2``
+            .. _`code:model_options/embedding[se2]`: `model_options/embedding[se2]`_
+            .. |code:model_options/embedding[baseline]| replace:: ``baseline``
+            .. _`code:model_options/embedding[baseline]`: `model_options/embedding[baseline]`_
+            .. |code:model_options/embedding[deeph-e3]| replace:: ``deeph-e3``
+            .. _`code:model_options/embedding[deeph-e3]`: `model_options/embedding[deeph-e3]`_
+            .. |code:model_options/embedding[e3baseline_0]| replace:: ``e3baseline_0``
+            .. _`code:model_options/embedding[e3baseline_0]`: `model_options/embedding[e3baseline_0]`_
+            .. |code:model_options/embedding[e3baseline_1]| replace:: ``e3baseline_1``
+            .. _`code:model_options/embedding[e3baseline_1]`: `model_options/embedding[e3baseline_1]`_
+            .. |code:model_options/embedding[e3baseline_2]| replace:: ``e3baseline_2``
+            .. _`code:model_options/embedding[e3baseline_2]`: `model_options/embedding[e3baseline_2]`_
+            .. |code:model_options/embedding[e3baseline_3]| replace:: ``e3baseline_3``
+            .. _`code:model_options/embedding[e3baseline_3]`: `model_options/embedding[e3baseline_3]`_
+            .. |code:model_options/embedding[e3baseline_4]| replace:: ``e3baseline_4``
+            .. _`code:model_options/embedding[e3baseline_4]`: `model_options/embedding[e3baseline_4]`_
+            .. |code:model_options/embedding[e3baseline_5]| replace:: ``e3baseline_5``
+            .. _`code:model_options/embedding[e3baseline_5]`: `model_options/embedding[e3baseline_5]`_
+            .. |code:model_options/embedding[e3baseline_6]| replace:: ``e3baseline_6``
+            .. _`code:model_options/embedding[e3baseline_6]`: `model_options/embedding[e3baseline_6]`_
+            .. |code:model_options/embedding[e3baseline_nonlocal]| replace:: ``e3baseline_nonlocal``
+            .. _`code:model_options/embedding[e3baseline_nonlocal]`: `model_options/embedding[e3baseline_nonlocal]`_
+
+        .. |flag:model_options/embedding/method| replace:: *method*
+        .. _`flag:model_options/embedding/method`: `model_options/embedding/method`_
+
+
+        .. _`model_options/embedding[se2]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``se2``: 
+
+        .. _`model_options/embedding[se2]/rs`: 
+
+        rs: 
+            | type: ``float`` | ``int``
+            | argument path: ``model_options/embedding[se2]/rs``
+
+            The soft cutoff where the smooth function starts.
+
+        .. _`model_options/embedding[se2]/rc`: 
+
+        rc: 
+            | type: ``float`` | ``int``
+            | argument path: ``model_options/embedding[se2]/rc``
+
+            The hard cutoff where the smooth function value ~0.0
+
+        .. _`model_options/embedding[se2]/radial_net`: 
+
+        radial_net: 
+            | type: ``dict``
+            | argument path: ``model_options/embedding[se2]/radial_net``
+
+            network to build the descriptors.
+
+            .. _`model_options/embedding[se2]/radial_net/neurons`: 
+
+            neurons: 
+                | type: ``list``
+                | argument path: ``model_options/embedding[se2]/radial_net/neurons``
+
+                the size of nn for descriptor
+
+            .. _`model_options/embedding[se2]/radial_net/activation`: 
+
+            activation: 
+                | type: ``str``, optional, default: ``tanh``
+                | argument path: ``model_options/embedding[se2]/radial_net/activation``
+
+                activation
+
+            .. _`model_options/embedding[se2]/radial_net/if_batch_normalized`: 
+
+            if_batch_normalized: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``model_options/embedding[se2]/radial_net/if_batch_normalized``
+
+                whether to turn on the batch normalization.
+
+        .. _`model_options/embedding[se2]/n_axis`: 
+
+        n_axis: 
+            | type: ``int`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``model_options/embedding[se2]/n_axis``
+
+            the out axis shape of the deepmd-se2 descriptor.
+
+
+        .. _`model_options/embedding[baseline]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``baseline``: 
+
+        .. _`model_options/embedding[baseline]/p`: 
+
+        p: 
+            | type: ``float`` | ``int``
+            | argument path: ``model_options/embedding[baseline]/p``
+
+        .. _`model_options/embedding[baseline]/rc`: 
+
+        rc: 
+            | type: ``float`` | ``int``
+            | argument path: ``model_options/embedding[baseline]/rc``
+
+        .. _`model_options/embedding[baseline]/n_basis`: 
+
+        n_basis: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[baseline]/n_basis``
+
+        .. _`model_options/embedding[baseline]/n_radial`: 
+
+        n_radial: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[baseline]/n_radial``
+
+        .. _`model_options/embedding[baseline]/n_sqrt_radial`: 
+
+        n_sqrt_radial: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[baseline]/n_sqrt_radial``
+
+        .. _`model_options/embedding[baseline]/n_layer`: 
+
+        n_layer: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[baseline]/n_layer``
+
+        .. _`model_options/embedding[baseline]/radial_net`: 
+
+        radial_net: 
+            | type: ``dict``
+            | argument path: ``model_options/embedding[baseline]/radial_net``
+
+            .. _`model_options/embedding[baseline]/radial_net/neurons`: 
+
+            neurons: 
+                | type: ``list``
+                | argument path: ``model_options/embedding[baseline]/radial_net/neurons``
+
+            .. _`model_options/embedding[baseline]/radial_net/activation`: 
+
+            activation: 
+                | type: ``str``, optional, default: ``tanh``
+                | argument path: ``model_options/embedding[baseline]/radial_net/activation``
+
+            .. _`model_options/embedding[baseline]/radial_net/if_batch_normalized`: 
+
+            if_batch_normalized: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``model_options/embedding[baseline]/radial_net/if_batch_normalized``
+
+        .. _`model_options/embedding[baseline]/hidden_net`: 
+
+        hidden_net: 
+            | type: ``dict``
+            | argument path: ``model_options/embedding[baseline]/hidden_net``
+
+            .. _`model_options/embedding[baseline]/hidden_net/neurons`: 
+
+            neurons: 
+                | type: ``list``
+                | argument path: ``model_options/embedding[baseline]/hidden_net/neurons``
+
+            .. _`model_options/embedding[baseline]/hidden_net/activation`: 
+
+            activation: 
+                | type: ``str``, optional, default: ``tanh``
+                | argument path: ``model_options/embedding[baseline]/hidden_net/activation``
+
+            .. _`model_options/embedding[baseline]/hidden_net/if_batch_normalized`: 
+
+            if_batch_normalized: 
+                | type: ``bool``, optional, default: ``False``
+                | argument path: ``model_options/embedding[baseline]/hidden_net/if_batch_normalized``
+
+        .. _`model_options/embedding[baseline]/n_axis`: 
+
+        n_axis: 
+            | type: ``int`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``model_options/embedding[baseline]/n_axis``
+
+
+        .. _`model_options/embedding[deeph-e3]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``deeph-e3``: 
+
+        .. _`model_options/embedding[deeph-e3]/irreps_embed`: 
+
+        irreps_embed: 
+            | type: ``str``, optional, default: ``64x0e``
+            | argument path: ``model_options/embedding[deeph-e3]/irreps_embed``
+
+        .. _`model_options/embedding[deeph-e3]/irreps_mid`: 
+
+        irreps_mid: 
+            | type: ``str``, optional, default: ``64x0e+32x1o+16x2e+8x3o+8x4e+4x5o``
+            | argument path: ``model_options/embedding[deeph-e3]/irreps_mid``
+
+        .. _`model_options/embedding[deeph-e3]/lmax`: 
+
+        lmax: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[deeph-e3]/lmax``
+
+        .. _`model_options/embedding[deeph-e3]/n_basis`: 
+
+        n_basis: 
+            | type: ``int``, optional, default: ``128``
+            | argument path: ``model_options/embedding[deeph-e3]/n_basis``
+
+        .. _`model_options/embedding[deeph-e3]/rc`: 
+
+        rc: 
+            | type: ``float``
+            | argument path: ``model_options/embedding[deeph-e3]/rc``
+
+        .. _`model_options/embedding[deeph-e3]/n_layer`: 
+
+        n_layer: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[deeph-e3]/n_layer``
+
+
+        .. _`model_options/embedding[e3baseline_0]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``e3baseline_0``: 
+
+        .. _`model_options/embedding[e3baseline_0]/irreps_hidden`: 
+
+        irreps_hidden: 
+            | type: ``str``, optional, default: ``64x0e+32x1o+16x2e+8x3o+8x4e+4x5o``
+            | argument path: ``model_options/embedding[e3baseline_0]/irreps_hidden``
+
+        .. _`model_options/embedding[e3baseline_0]/lmax`: 
+
+        lmax: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_0]/lmax``
+
+        .. _`model_options/embedding[e3baseline_0]/avg_num_neighbors`: 
+
+        avg_num_neighbors: 
+            | type: ``float`` | ``int``, optional, default: ``50``
+            | argument path: ``model_options/embedding[e3baseline_0]/avg_num_neighbors``
+
+        .. _`model_options/embedding[e3baseline_0]/r_max`: 
+
+        r_max: 
+            | type: ``dict`` | ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_0]/r_max``
+
+        .. _`model_options/embedding[e3baseline_0]/n_layers`: 
+
+        n_layers: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_0]/n_layers``
+
+        .. _`model_options/embedding[e3baseline_0]/n_radial_basis`: 
+
+        n_radial_basis: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_0]/n_radial_basis``
+
+        .. _`model_options/embedding[e3baseline_0]/PolynomialCutoff_p`: 
+
+        PolynomialCutoff_p: 
+            | type: ``int``, optional, default: ``6``
+            | argument path: ``model_options/embedding[e3baseline_0]/PolynomialCutoff_p``
+
+            The order of polynomial cutoff function. Default: 6
+
+        .. _`model_options/embedding[e3baseline_0]/latent_kwargs`: 
+
+        latent_kwargs: 
+            | type: ``dict`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``model_options/embedding[e3baseline_0]/latent_kwargs``
+
+        .. _`model_options/embedding[e3baseline_0]/env_embed_multiplicity`: 
+
+        env_embed_multiplicity: 
+            | type: ``int``, optional, default: ``1``
+            | argument path: ``model_options/embedding[e3baseline_0]/env_embed_multiplicity``
+
+        .. _`model_options/embedding[e3baseline_0]/linear_after_env_embed`: 
+
+        linear_after_env_embed: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_0]/linear_after_env_embed``
+
+        .. _`model_options/embedding[e3baseline_0]/latent_resnet_update_ratios_learnable`: 
+
+        latent_resnet_update_ratios_learnable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_0]/latent_resnet_update_ratios_learnable``
+
+
+        .. _`model_options/embedding[e3baseline_1]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``e3baseline_1``: 
+
+        .. _`model_options/embedding[e3baseline_1]/irreps_hidden`: 
+
+        irreps_hidden: 
+            | type: ``str``, optional, default: ``64x0e+32x1o+16x2e+8x3o+8x4e+4x5o``
+            | argument path: ``model_options/embedding[e3baseline_1]/irreps_hidden``
+
+        .. _`model_options/embedding[e3baseline_1]/lmax`: 
+
+        lmax: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_1]/lmax``
+
+        .. _`model_options/embedding[e3baseline_1]/avg_num_neighbors`: 
+
+        avg_num_neighbors: 
+            | type: ``float`` | ``int``, optional, default: ``50``
+            | argument path: ``model_options/embedding[e3baseline_1]/avg_num_neighbors``
+
+        .. _`model_options/embedding[e3baseline_1]/r_max`: 
+
+        r_max: 
+            | type: ``dict`` | ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_1]/r_max``
+
+        .. _`model_options/embedding[e3baseline_1]/n_layers`: 
+
+        n_layers: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_1]/n_layers``
+
+        .. _`model_options/embedding[e3baseline_1]/n_radial_basis`: 
+
+        n_radial_basis: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_1]/n_radial_basis``
+
+        .. _`model_options/embedding[e3baseline_1]/PolynomialCutoff_p`: 
+
+        PolynomialCutoff_p: 
+            | type: ``int``, optional, default: ``6``
+            | argument path: ``model_options/embedding[e3baseline_1]/PolynomialCutoff_p``
+
+            The order of polynomial cutoff function. Default: 6
+
+        .. _`model_options/embedding[e3baseline_1]/latent_kwargs`: 
+
+        latent_kwargs: 
+            | type: ``dict`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``model_options/embedding[e3baseline_1]/latent_kwargs``
+
+        .. _`model_options/embedding[e3baseline_1]/env_embed_multiplicity`: 
+
+        env_embed_multiplicity: 
+            | type: ``int``, optional, default: ``1``
+            | argument path: ``model_options/embedding[e3baseline_1]/env_embed_multiplicity``
+
+        .. _`model_options/embedding[e3baseline_1]/linear_after_env_embed`: 
+
+        linear_after_env_embed: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_1]/linear_after_env_embed``
+
+        .. _`model_options/embedding[e3baseline_1]/latent_resnet_update_ratios_learnable`: 
+
+        latent_resnet_update_ratios_learnable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_1]/latent_resnet_update_ratios_learnable``
+
+
+        .. _`model_options/embedding[e3baseline_2]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``e3baseline_2``: 
+
+        .. _`model_options/embedding[e3baseline_2]/irreps_hidden`: 
+
+        irreps_hidden: 
+            | type: ``str``, optional, default: ``64x0e+32x1o+16x2e+8x3o+8x4e+4x5o``
+            | argument path: ``model_options/embedding[e3baseline_2]/irreps_hidden``
+
+        .. _`model_options/embedding[e3baseline_2]/lmax`: 
+
+        lmax: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_2]/lmax``
+
+        .. _`model_options/embedding[e3baseline_2]/avg_num_neighbors`: 
+
+        avg_num_neighbors: 
+            | type: ``float`` | ``int``, optional, default: ``50``
+            | argument path: ``model_options/embedding[e3baseline_2]/avg_num_neighbors``
+
+        .. _`model_options/embedding[e3baseline_2]/r_max`: 
+
+        r_max: 
+            | type: ``dict`` | ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_2]/r_max``
+
+        .. _`model_options/embedding[e3baseline_2]/n_layers`: 
+
+        n_layers: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_2]/n_layers``
+
+        .. _`model_options/embedding[e3baseline_2]/n_radial_basis`: 
+
+        n_radial_basis: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_2]/n_radial_basis``
+
+        .. _`model_options/embedding[e3baseline_2]/PolynomialCutoff_p`: 
+
+        PolynomialCutoff_p: 
+            | type: ``int``, optional, default: ``6``
+            | argument path: ``model_options/embedding[e3baseline_2]/PolynomialCutoff_p``
+
+            The order of polynomial cutoff function. Default: 6
+
+        .. _`model_options/embedding[e3baseline_2]/latent_kwargs`: 
+
+        latent_kwargs: 
+            | type: ``dict`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``model_options/embedding[e3baseline_2]/latent_kwargs``
+
+        .. _`model_options/embedding[e3baseline_2]/env_embed_multiplicity`: 
+
+        env_embed_multiplicity: 
+            | type: ``int``, optional, default: ``1``
+            | argument path: ``model_options/embedding[e3baseline_2]/env_embed_multiplicity``
+
+        .. _`model_options/embedding[e3baseline_2]/linear_after_env_embed`: 
+
+        linear_after_env_embed: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_2]/linear_after_env_embed``
+
+        .. _`model_options/embedding[e3baseline_2]/latent_resnet_update_ratios_learnable`: 
+
+        latent_resnet_update_ratios_learnable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_2]/latent_resnet_update_ratios_learnable``
+
+
+        .. _`model_options/embedding[e3baseline_3]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``e3baseline_3``: 
+
+        .. _`model_options/embedding[e3baseline_3]/irreps_hidden`: 
+
+        irreps_hidden: 
+            | type: ``str``, optional, default: ``64x0e+32x1o+16x2e+8x3o+8x4e+4x5o``
+            | argument path: ``model_options/embedding[e3baseline_3]/irreps_hidden``
+
+        .. _`model_options/embedding[e3baseline_3]/lmax`: 
+
+        lmax: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_3]/lmax``
+
+        .. _`model_options/embedding[e3baseline_3]/avg_num_neighbors`: 
+
+        avg_num_neighbors: 
+            | type: ``float`` | ``int``, optional, default: ``50``
+            | argument path: ``model_options/embedding[e3baseline_3]/avg_num_neighbors``
+
+        .. _`model_options/embedding[e3baseline_3]/r_max`: 
+
+        r_max: 
+            | type: ``dict`` | ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_3]/r_max``
+
+        .. _`model_options/embedding[e3baseline_3]/n_layers`: 
+
+        n_layers: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_3]/n_layers``
+
+        .. _`model_options/embedding[e3baseline_3]/n_radial_basis`: 
+
+        n_radial_basis: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_3]/n_radial_basis``
+
+        .. _`model_options/embedding[e3baseline_3]/PolynomialCutoff_p`: 
+
+        PolynomialCutoff_p: 
+            | type: ``int``, optional, default: ``6``
+            | argument path: ``model_options/embedding[e3baseline_3]/PolynomialCutoff_p``
+
+            The order of polynomial cutoff function. Default: 6
+
+        .. _`model_options/embedding[e3baseline_3]/latent_kwargs`: 
+
+        latent_kwargs: 
+            | type: ``dict`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``model_options/embedding[e3baseline_3]/latent_kwargs``
+
+        .. _`model_options/embedding[e3baseline_3]/env_embed_multiplicity`: 
+
+        env_embed_multiplicity: 
+            | type: ``int``, optional, default: ``1``
+            | argument path: ``model_options/embedding[e3baseline_3]/env_embed_multiplicity``
+
+        .. _`model_options/embedding[e3baseline_3]/linear_after_env_embed`: 
+
+        linear_after_env_embed: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_3]/linear_after_env_embed``
+
+        .. _`model_options/embedding[e3baseline_3]/latent_resnet_update_ratios_learnable`: 
+
+        latent_resnet_update_ratios_learnable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_3]/latent_resnet_update_ratios_learnable``
+
+
+        .. _`model_options/embedding[e3baseline_4]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``e3baseline_4``: 
+
+        .. _`model_options/embedding[e3baseline_4]/irreps_hidden`: 
+
+        irreps_hidden: 
+            | type: ``str``, optional, default: ``64x0e+32x1o+16x2e+8x3o+8x4e+4x5o``
+            | argument path: ``model_options/embedding[e3baseline_4]/irreps_hidden``
+
+        .. _`model_options/embedding[e3baseline_4]/lmax`: 
+
+        lmax: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_4]/lmax``
+
+        .. _`model_options/embedding[e3baseline_4]/avg_num_neighbors`: 
+
+        avg_num_neighbors: 
+            | type: ``float`` | ``int``, optional, default: ``50``
+            | argument path: ``model_options/embedding[e3baseline_4]/avg_num_neighbors``
+
+        .. _`model_options/embedding[e3baseline_4]/r_max`: 
+
+        r_max: 
+            | type: ``dict`` | ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_4]/r_max``
+
+        .. _`model_options/embedding[e3baseline_4]/n_layers`: 
+
+        n_layers: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_4]/n_layers``
+
+        .. _`model_options/embedding[e3baseline_4]/n_radial_basis`: 
+
+        n_radial_basis: 
+            | type: ``int``, optional, default: ``3``
+            | argument path: ``model_options/embedding[e3baseline_4]/n_radial_basis``
+
+        .. _`model_options/embedding[e3baseline_4]/PolynomialCutoff_p`: 
+
+        PolynomialCutoff_p: 
+            | type: ``int``, optional, default: ``6``
+            | argument path: ``model_options/embedding[e3baseline_4]/PolynomialCutoff_p``
+
+            The order of polynomial cutoff function. Default: 6
+
+        .. _`model_options/embedding[e3baseline_4]/latent_kwargs`: 
+
+        latent_kwargs: 
+            | type: ``dict`` | ``NoneType``, optional, default: ``None``
+            | argument path: ``model_options/embedding[e3baseline_4]/latent_kwargs``
+
+        .. _`model_options/embedding[e3baseline_4]/env_embed_multiplicity`: 
+
+        env_embed_multiplicity: 
+            | type: ``int``, optional, default: ``1``
+            | argument path: ``model_options/embedding[e3baseline_4]/env_embed_multiplicity``
+
+        .. _`model_options/embedding[e3baseline_4]/linear_after_env_embed`: 
+
+        linear_after_env_embed: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_4]/linear_after_env_embed``
+
+        .. _`model_options/embedding[e3baseline_4]/latent_resnet_update_ratios_learnable`: 
+
+        latent_resnet_update_ratios_learnable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_4]/latent_resnet_update_ratios_learnable``
+
+
+        .. _`model_options/embedding[e3baseline_5]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``e3baseline_5``: 
+
+        .. _`model_options/embedding[e3baseline_5]/irreps_hidden`: 
+
+        irreps_hidden: 
+            | type: ``str``
+            | argument path: ``model_options/embedding[e3baseline_5]/irreps_hidden``
+
+        .. _`model_options/embedding[e3baseline_5]/lmax`: 
+
+        lmax: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[e3baseline_5]/lmax``
+
+        .. _`model_options/embedding[e3baseline_5]/avg_num_neighbors`: 
+
+        avg_num_neighbors: 
+            | type: ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_5]/avg_num_neighbors``
+
+        .. _`model_options/embedding[e3baseline_5]/r_max`: 
+
+        r_max: 
+            | type: ``dict`` | ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_5]/r_max``
+
+        .. _`model_options/embedding[e3baseline_5]/n_layers`: 
+
+        n_layers: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[e3baseline_5]/n_layers``
+
+        .. _`model_options/embedding[e3baseline_5]/n_radial_basis`: 
+
+        n_radial_basis: 
+            | type: ``int``, optional, default: ``10``
+            | argument path: ``model_options/embedding[e3baseline_5]/n_radial_basis``
+
+        .. _`model_options/embedding[e3baseline_5]/PolynomialCutoff_p`: 
+
+        PolynomialCutoff_p: 
+            | type: ``int``, optional, default: ``6``
+            | argument path: ``model_options/embedding[e3baseline_5]/PolynomialCutoff_p``
+
+            The order of polynomial cutoff function. Default: 6
+
+        .. _`model_options/embedding[e3baseline_5]/cutoff_type`: 
+
+        cutoff_type: 
+            | type: ``str``, optional, default: ``polynomial``
+            | argument path: ``model_options/embedding[e3baseline_5]/cutoff_type``
+
+            The type of cutoff function. Default: polynomial
+
+        .. _`model_options/embedding[e3baseline_5]/env_embed_multiplicity`: 
+
+        env_embed_multiplicity: 
+            | type: ``int``, optional, default: ``1``
+            | argument path: ``model_options/embedding[e3baseline_5]/env_embed_multiplicity``
+
+        .. _`model_options/embedding[e3baseline_5]/tp_radial_emb`: 
+
+        tp_radial_emb: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_5]/tp_radial_emb``
+
+            Whether to use tensor product radial embedding.
+
+        .. _`model_options/embedding[e3baseline_5]/tp_radial_channels`: 
+
+        tp_radial_channels: 
+            | type: ``list``, optional, default: ``[128, 128]``
+            | argument path: ``model_options/embedding[e3baseline_5]/tp_radial_channels``
+
+            The number of channels in tensor product radial embedding.
+
+        .. _`model_options/embedding[e3baseline_5]/latent_channels`: 
+
+        latent_channels: 
+            | type: ``list``, optional, default: ``[128, 128]``
+            | argument path: ``model_options/embedding[e3baseline_5]/latent_channels``
+
+            The number of channels in latent embedding.
+
+        .. _`model_options/embedding[e3baseline_5]/latent_dim`: 
+
+        latent_dim: 
+            | type: ``int``, optional, default: ``256``
+            | argument path: ``model_options/embedding[e3baseline_5]/latent_dim``
+
+            The dimension of latent embedding.
+
+        .. _`model_options/embedding[e3baseline_5]/res_update`: 
+
+        res_update: 
+            | type: ``bool``, optional, default: ``True``
+            | argument path: ``model_options/embedding[e3baseline_5]/res_update``
+
+            Whether to use residual update.
+
+        .. _`model_options/embedding[e3baseline_5]/res_update_ratios`: 
+
+        res_update_ratios: 
+            | type: ``float``, optional, default: ``0.5``
+            | argument path: ``model_options/embedding[e3baseline_5]/res_update_ratios``
+
+            The ratios of residual update, should in (0,1).
+
+        .. _`model_options/embedding[e3baseline_5]/res_update_ratios_learnable`: 
+
+        res_update_ratios_learnable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_5]/res_update_ratios_learnable``
+
+            Whether to make the ratios of residual update learnable.
+
+
+        .. _`model_options/embedding[e3baseline_6]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``e3baseline_6``: 
+
+        .. _`model_options/embedding[e3baseline_6]/irreps_hidden`: 
+
+        irreps_hidden: 
+            | type: ``str``
+            | argument path: ``model_options/embedding[e3baseline_6]/irreps_hidden``
+
+        .. _`model_options/embedding[e3baseline_6]/lmax`: 
+
+        lmax: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[e3baseline_6]/lmax``
+
+        .. _`model_options/embedding[e3baseline_6]/avg_num_neighbors`: 
+
+        avg_num_neighbors: 
+            | type: ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_6]/avg_num_neighbors``
+
+        .. _`model_options/embedding[e3baseline_6]/r_max`: 
+
+        r_max: 
+            | type: ``dict`` | ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_6]/r_max``
+
+        .. _`model_options/embedding[e3baseline_6]/n_layers`: 
+
+        n_layers: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[e3baseline_6]/n_layers``
+
+        .. _`model_options/embedding[e3baseline_6]/n_radial_basis`: 
+
+        n_radial_basis: 
+            | type: ``int``, optional, default: ``10``
+            | argument path: ``model_options/embedding[e3baseline_6]/n_radial_basis``
+
+        .. _`model_options/embedding[e3baseline_6]/PolynomialCutoff_p`: 
+
+        PolynomialCutoff_p: 
+            | type: ``int``, optional, default: ``6``
+            | argument path: ``model_options/embedding[e3baseline_6]/PolynomialCutoff_p``
+
+            The order of polynomial cutoff function. Default: 6
+
+        .. _`model_options/embedding[e3baseline_6]/cutoff_type`: 
+
+        cutoff_type: 
+            | type: ``str``, optional, default: ``polynomial``
+            | argument path: ``model_options/embedding[e3baseline_6]/cutoff_type``
+
+            The type of cutoff function. Default: polynomial
+
+        .. _`model_options/embedding[e3baseline_6]/env_embed_multiplicity`: 
+
+        env_embed_multiplicity: 
+            | type: ``int``, optional, default: ``1``
+            | argument path: ``model_options/embedding[e3baseline_6]/env_embed_multiplicity``
+
+        .. _`model_options/embedding[e3baseline_6]/tp_radial_emb`: 
+
+        tp_radial_emb: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_6]/tp_radial_emb``
+
+            Whether to use tensor product radial embedding.
+
+        .. _`model_options/embedding[e3baseline_6]/tp_radial_channels`: 
+
+        tp_radial_channels: 
+            | type: ``list``, optional, default: ``[128, 128]``
+            | argument path: ``model_options/embedding[e3baseline_6]/tp_radial_channels``
+
+            The number of channels in tensor product radial embedding.
+
+        .. _`model_options/embedding[e3baseline_6]/latent_channels`: 
+
+        latent_channels: 
+            | type: ``list``, optional, default: ``[128, 128]``
+            | argument path: ``model_options/embedding[e3baseline_6]/latent_channels``
+
+            The number of channels in latent embedding.
+
+        .. _`model_options/embedding[e3baseline_6]/latent_dim`: 
+
+        latent_dim: 
+            | type: ``int``, optional, default: ``256``
+            | argument path: ``model_options/embedding[e3baseline_6]/latent_dim``
+
+            The dimension of latent embedding.
+
+        .. _`model_options/embedding[e3baseline_6]/res_update`: 
+
+        res_update: 
+            | type: ``bool``, optional, default: ``True``
+            | argument path: ``model_options/embedding[e3baseline_6]/res_update``
+
+            Whether to use residual update.
+
+        .. _`model_options/embedding[e3baseline_6]/res_update_ratios`: 
+
+        res_update_ratios: 
+            | type: ``float``, optional, default: ``0.5``
+            | argument path: ``model_options/embedding[e3baseline_6]/res_update_ratios``
+
+            The ratios of residual update, should in (0,1).
+
+        .. _`model_options/embedding[e3baseline_6]/res_update_ratios_learnable`: 
+
+        res_update_ratios_learnable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_6]/res_update_ratios_learnable``
+
+            Whether to make the ratios of residual update learnable.
+
+
+        .. _`model_options/embedding[e3baseline_nonlocal]`: 
+
+        When |flag:model_options/embedding/method|_ is set to ``e3baseline_nonlocal``: 
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/irreps_hidden`: 
+
+        irreps_hidden: 
+            | type: ``str``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/irreps_hidden``
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/lmax`: 
+
+        lmax: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/lmax``
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/avg_num_neighbors`: 
+
+        avg_num_neighbors: 
+            | type: ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/avg_num_neighbors``
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/r_max`: 
+
+        r_max: 
+            | type: ``dict`` | ``float`` | ``int``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/r_max``
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/n_layers`: 
+
+        n_layers: 
+            | type: ``int``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/n_layers``
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/n_radial_basis`: 
+
+        n_radial_basis: 
+            | type: ``int``, optional, default: ``10``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/n_radial_basis``
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/PolynomialCutoff_p`: 
+
+        PolynomialCutoff_p: 
+            | type: ``int``, optional, default: ``6``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/PolynomialCutoff_p``
+
+            The order of polynomial cutoff function. Default: 6
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/cutoff_type`: 
+
+        cutoff_type: 
+            | type: ``str``, optional, default: ``polynomial``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/cutoff_type``
+
+            The type of cutoff function. Default: polynomial
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/env_embed_multiplicity`: 
+
+        env_embed_multiplicity: 
+            | type: ``int``, optional, default: ``1``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/env_embed_multiplicity``
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/tp_radial_emb`: 
+
+        tp_radial_emb: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/tp_radial_emb``
+
+            Whether to use tensor product radial embedding.
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/tp_radial_channels`: 
+
+        tp_radial_channels: 
+            | type: ``list``, optional, default: ``[128, 128]``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/tp_radial_channels``
+
+            The number of channels in tensor product radial embedding.
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/latent_channels`: 
+
+        latent_channels: 
+            | type: ``list``, optional, default: ``[128, 128]``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/latent_channels``
+
+            The number of channels in latent embedding.
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/latent_dim`: 
+
+        latent_dim: 
+            | type: ``int``, optional, default: ``256``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/latent_dim``
+
+            The dimension of latent embedding.
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/res_update`: 
+
+        res_update: 
+            | type: ``bool``, optional, default: ``True``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/res_update``
+
+            Whether to use residual update.
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/res_update_ratios`: 
+
+        res_update_ratios: 
+            | type: ``float``, optional, default: ``0.5``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/res_update_ratios``
+
+            The ratios of residual update, should in (0,1).
+
+        .. _`model_options/embedding[e3baseline_nonlocal]/res_update_ratios_learnable`: 
+
+        res_update_ratios_learnable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/embedding[e3baseline_nonlocal]/res_update_ratios_learnable``
+
+            Whether to make the ratios of residual update learnable.
+
+    .. _`model_options/prediction`: 
+
+    prediction: 
+        | type: ``dict``, optional
+        | argument path: ``model_options/prediction``
+
+        The parameters to define the prediction model
+
+
+        Depending on the value of *method*, different sub args are accepted. 
+
+        .. _`model_options/prediction/method`: 
+
+        method:
+            | type: ``str`` (flag key)
+            | argument path: ``model_options/prediction/method`` 
+            | possible choices: |code:model_options/prediction[sktb]|_, |code:model_options/prediction[e3tb]|_
+
+            The options to indicate the prediction model. Can be sktb or e3tb.
+
+            .. |code:model_options/prediction[sktb]| replace:: ``sktb``
+            .. _`code:model_options/prediction[sktb]`: `model_options/prediction[sktb]`_
+            .. |code:model_options/prediction[e3tb]| replace:: ``e3tb``
+            .. _`code:model_options/prediction[e3tb]`: `model_options/prediction[e3tb]`_
+
+        .. |flag:model_options/prediction/method| replace:: *method*
+        .. _`flag:model_options/prediction/method`: `model_options/prediction/method`_
+
+
+        .. _`model_options/prediction[sktb]`: 
+
+        When |flag:model_options/prediction/method|_ is set to ``sktb``: 
+
+        neural network options for prediction model.
+
+        .. _`model_options/prediction[sktb]/neurons`: 
+
+        neurons: 
+            | type: ``list``
+            | argument path: ``model_options/prediction[sktb]/neurons``
+
+            neurons in the neural network.
+
+        .. _`model_options/prediction[sktb]/activation`: 
+
+        activation: 
+            | type: ``str``, optional, default: ``tanh``
+            | argument path: ``model_options/prediction[sktb]/activation``
+
+            activation function.
+
+        .. _`model_options/prediction[sktb]/if_batch_normalized`: 
+
+        if_batch_normalized: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/prediction[sktb]/if_batch_normalized``
+
+            if to turn on batch normalization
+
+
+        .. _`model_options/prediction[e3tb]`: 
+
+        When |flag:model_options/prediction/method|_ is set to ``e3tb``: 
+
+        neural network options for prediction model.
+
+        .. _`model_options/prediction[e3tb]/scales_trainable`: 
+
+        scales_trainable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/prediction[e3tb]/scales_trainable``
+
+            whether to scale the trianing target.
+
+        .. _`model_options/prediction[e3tb]/shifts_trainable`: 
+
+        shifts_trainable: 
+            | type: ``bool``, optional, default: ``False``
+            | argument path: ``model_options/prediction[e3tb]/shifts_trainable``
+
+            whether to shift the training target.
+
+    .. _`model_options/nnsk`: 
+
+    nnsk: 
+        | type: ``dict``, optional
+        | argument path: ``model_options/nnsk``
+
+        The parameters to define the nnsk model.
+
+        .. _`model_options/nnsk/onsite`: 
+
+        onsite: 
+            | type: ``dict``
+            | argument path: ``model_options/nnsk/onsite``
+
+            The onsite options to define the onsite of nnsk model.
+
+
+            Depending on the value of *method*, different sub args are accepted. 
+
+            .. _`model_options/nnsk/onsite/method`: 
+
+            method:
+                | type: ``str`` (flag key)
+                | argument path: ``model_options/nnsk/onsite/method`` 
+                | possible choices: |code:model_options/nnsk/onsite[strain]|_, |code:model_options/nnsk/onsite[uniform]|_, |code:model_options/nnsk/onsite[NRL]|_, |code:model_options/nnsk/onsite[none]|_
+
+                The onsite correction mode, the onsite energy is expressed as the energy of isolated atoms plus the model correction, the correction mode are:
+                                    Default: `none`: use the database onsite energy value.
+                                    - `strain`: The strain mode correct the onsite matrix densly by $$H_{i,i}^{lm,l^\prime m^\prime} = \epsilon_l^0 \delta_{ll^\prime}\delta_{mm^\prime} + \sum_p \sum_{\zeta} \Big[ \mathcal{U}_{\zeta}(\hat{r}_{ip}) \ \epsilon_{ll^\prime \zeta} \Big]_{mm^\prime}$$ which is also parameterized as a set of Slater-Koster like integrals.
+
+                                    - `uniform`: The correction is a energy shift respect of orbital of each atom. Which is formally written as: 
+                                                $$H_{i,i}^{lm,l^\prime m^\prime} = (\epsilon_l^0+\epsilon_l^\prime) \delta_{ll^\prime}\delta_{mm^\prime}$$ Where $\epsilon_l^0$ is the isolated energy level from the DeePTB onsite database, and $\epsilon_l^\prime$ is the parameters to fit.
+                                    - `NRL`: use the NRL-TB formula.
+                
+
+                .. |code:model_options/nnsk/onsite[strain]| replace:: ``strain``
+                .. _`code:model_options/nnsk/onsite[strain]`: `model_options/nnsk/onsite[strain]`_
+                .. |code:model_options/nnsk/onsite[uniform]| replace:: ``uniform``
+                .. _`code:model_options/nnsk/onsite[uniform]`: `model_options/nnsk/onsite[uniform]`_
+                .. |code:model_options/nnsk/onsite[NRL]| replace:: ``NRL``
+                .. _`code:model_options/nnsk/onsite[NRL]`: `model_options/nnsk/onsite[NRL]`_
+                .. |code:model_options/nnsk/onsite[none]| replace:: ``none``
+                .. _`code:model_options/nnsk/onsite[none]`: `model_options/nnsk/onsite[none]`_
+
+            .. |flag:model_options/nnsk/onsite/method| replace:: *method*
+            .. _`flag:model_options/nnsk/onsite/method`: `model_options/nnsk/onsite/method`_
+
+
+            .. _`model_options/nnsk/onsite[strain]`: 
+
+            When |flag:model_options/nnsk/onsite/method|_ is set to ``strain``: 
+
+            .. _`model_options/nnsk/onsite[strain]/rs`: 
+
+            rs: 
+                | type: ``float``, optional, default: ``6.0``
+                | argument path: ``model_options/nnsk/onsite[strain]/rs``
+
+                The smooth cutoff `fc` for strain model. rs is where fc = 0.5
+
+            .. _`model_options/nnsk/onsite[strain]/w`: 
+
+            w: 
+                | type: ``float``, optional, default: ``0.1``
+                | argument path: ``model_options/nnsk/onsite[strain]/w``
+
+                The decay factor of `fc` for strain and nrl model.
+
+
+            .. _`model_options/nnsk/onsite[uniform]`: 
+
+            When |flag:model_options/nnsk/onsite/method|_ is set to ``uniform``: 
+
+
+
+            .. _`model_options/nnsk/onsite[NRL]`: 
+
+            When |flag:model_options/nnsk/onsite/method|_ is set to ``NRL``: 
+
+            .. _`model_options/nnsk/onsite[NRL]/rc`: 
+
+            rc: 
+                | type: ``float``, optional, default: ``6.0``
+                | argument path: ``model_options/nnsk/onsite[NRL]/rc``
+
+                The smooth cutoff of `fc` for nrl model, rc is where fc ~ 0.0
+
+            .. _`model_options/nnsk/onsite[NRL]/w`: 
+
+            w: 
+                | type: ``float``, optional, default: ``0.1``
+                | argument path: ``model_options/nnsk/onsite[NRL]/w``
+
+                The decay factor of `fc` for strain and nrl model.
+
+            .. _`model_options/nnsk/onsite[NRL]/lda`: 
+
+            lda: 
+                | type: ``float``, optional, default: ``1.0``
+                | argument path: ``model_options/nnsk/onsite[NRL]/lda``
+
+                The lambda type encoding value in nrl model. now only support elementary substance
+
+
+            .. _`model_options/nnsk/onsite[none]`: 
+
+            When |flag:model_options/nnsk/onsite/method|_ is set to ``none``: 
+
+
+        .. _`model_options/nnsk/hopping`: 
+
+        hopping: 
+            | type: ``dict``
+            | argument path: ``model_options/nnsk/hopping``
+
+            The hopping options to define the hopping of nnsk model.
+
+
+            Depending on the value of *method*, different sub args are accepted. 
+
+            .. _`model_options/nnsk/hopping/method`: 
+
+            method:
+                | type: ``str`` (flag key)
+                | argument path: ``model_options/nnsk/hopping/method`` 
+                | possible choices: |code:model_options/nnsk/hopping[powerlaw]|_, |code:model_options/nnsk/hopping[varTang96]|_, |code:model_options/nnsk/hopping[NRL0]|_, |code:model_options/nnsk/hopping[NRL1]|_, |code:model_options/nnsk/hopping[custom]|_
+
+                The hopping formula. 
+                                    -  `powerlaw`: the powerlaw formula for bond length dependence for sk integrals.
+                                    -  `varTang96`: a variational formula based on Tang96 formula.
+                                    -  `NRL0`: the old version of NRL formula for overlap, we set overlap and hopping share same options.
+                                    -  `NRL1`: the new version of NRL formula for overlap. 
+                    
+
+                .. |code:model_options/nnsk/hopping[powerlaw]| replace:: ``powerlaw``
+                .. _`code:model_options/nnsk/hopping[powerlaw]`: `model_options/nnsk/hopping[powerlaw]`_
+                .. |code:model_options/nnsk/hopping[varTang96]| replace:: ``varTang96``
+                .. _`code:model_options/nnsk/hopping[varTang96]`: `model_options/nnsk/hopping[varTang96]`_
+                .. |code:model_options/nnsk/hopping[NRL0]| replace:: ``NRL0``
+                .. _`code:model_options/nnsk/hopping[NRL0]`: `model_options/nnsk/hopping[NRL0]`_
+                .. |code:model_options/nnsk/hopping[NRL1]| replace:: ``NRL1``
+                .. _`code:model_options/nnsk/hopping[NRL1]`: `model_options/nnsk/hopping[NRL1]`_
+                .. |code:model_options/nnsk/hopping[custom]| replace:: ``custom``
+                .. _`code:model_options/nnsk/hopping[custom]`: `model_options/nnsk/hopping[custom]`_
+
+            .. |flag:model_options/nnsk/hopping/method| replace:: *method*
+            .. _`flag:model_options/nnsk/hopping/method`: `model_options/nnsk/hopping/method`_
+
+
+            .. _`model_options/nnsk/hopping[powerlaw]`: 
+
+            When |flag:model_options/nnsk/hopping/method|_ is set to ``powerlaw``: 
+
+            .. _`model_options/nnsk/hopping[powerlaw]/rs`: 
+
+            rs: 
+                | type: ``float``, optional, default: ``6.0``
+                | argument path: ``model_options/nnsk/hopping[powerlaw]/rs``
+
+                The cut-off for smooth function fc for powerlaw and varTang96, fc(rs)=0.5
+
+            .. _`model_options/nnsk/hopping[powerlaw]/w`: 
+
+            w: 
+                | type: ``float``, optional, default: ``0.1``
+                | argument path: ``model_options/nnsk/hopping[powerlaw]/w``
+
+                 The decay w in fc
+
+
+            .. _`model_options/nnsk/hopping[varTang96]`: 
+
+            When |flag:model_options/nnsk/hopping/method|_ is set to ``varTang96``: 
+
+            .. _`model_options/nnsk/hopping[varTang96]/rs`: 
+
+            rs: 
+                | type: ``float``, optional, default: ``6.0``
+                | argument path: ``model_options/nnsk/hopping[varTang96]/rs``
+
+                The cut-off for smooth function fc for powerlaw and varTang96, fc(rs)=0.5
+
+            .. _`model_options/nnsk/hopping[varTang96]/w`: 
+
+            w: 
+                | type: ``float``, optional, default: ``0.1``
+                | argument path: ``model_options/nnsk/hopping[varTang96]/w``
+
+                 The decay w in fc
+
+
+            .. _`model_options/nnsk/hopping[NRL0]`: 
+
+            When |flag:model_options/nnsk/hopping/method|_ is set to ``NRL0``: 
+
+            .. _`model_options/nnsk/hopping[NRL0]/rc`: 
+
+            rc: 
+                | type: ``float``, optional, default: ``6.0``
+                | argument path: ``model_options/nnsk/hopping[NRL0]/rc``
+
+                The cut-off for smooth function fc for NRL, fc(rc) = 0.
+
+            .. _`model_options/nnsk/hopping[NRL0]/w`: 
+
+            w: 
+                | type: ``float``, optional, default: ``0.1``
+                | argument path: ``model_options/nnsk/hopping[NRL0]/w``
+
+                 The decay w in fc
+
+
+            .. _`model_options/nnsk/hopping[NRL1]`: 
+
+            When |flag:model_options/nnsk/hopping/method|_ is set to ``NRL1``: 
+
+            .. _`model_options/nnsk/hopping[NRL1]/rc`: 
+
+            rc: 
+                | type: ``float``, optional, default: ``6.0``
+                | argument path: ``model_options/nnsk/hopping[NRL1]/rc``
+
+                The cut-off for smooth function fc for NRL, fc(rc) = 0.
+
+            .. _`model_options/nnsk/hopping[NRL1]/w`: 
+
+            w: 
+                | type: ``float``, optional, default: ``0.1``
+                | argument path: ``model_options/nnsk/hopping[NRL1]/w``
+
+                 The decay w in fc
+
+
+            .. _`model_options/nnsk/hopping[custom]`: 
+
+            When |flag:model_options/nnsk/hopping/method|_ is set to ``custom``: 
+
+
+        .. _`model_options/nnsk/soc`: 
+
+        soc: 
+            | type: ``dict``, optional, default: ``{}``
+            | argument path: ``model_options/nnsk/soc``
+
+            The soc options to define the soc of nnsk model,
+                            Default: {} # empty dict
+
+                            - {'method':'none'} : use database soc value. 
+                            - {'method':uniform} : set lambda_il; assign a soc lambda value for each orbital -l on each atomtype i; l=0,1,2 for s p d.
+
+        .. _`model_options/nnsk/freeze`: 
+
+        freeze: 
+            | type: ``str`` | ``bool`` | ``list``, optional, default: ``False``
+            | argument path: ``model_options/nnsk/freeze``
+
+            The parameters to define the freeze of nnsk model can be bool and string and list.
+
+                                Default: False
+
+                                 - True: freeze all the nnsk parameters
+
+                                 - False: train all the nnsk parameters
+ 
+                                 - 'hopping','onsite','overlap' and 'soc' to freeze the corresponding parameters.
+                                 - list of the strings e.g. ['overlap','soc'] to freeze both overlap and soc parameters.
+
+        .. _`model_options/nnsk/std`: 
+
+        std: 
+            | type: ``float``, optional, default: ``0.01``
+            | argument path: ``model_options/nnsk/std``
+
+            The std value to initialize the nnsk parameters. Default: 0.01
+
+        .. _`model_options/nnsk/push`: 
+
+        push: 
+            | type: ``dict`` | ``bool``, optional, default: ``False``
+            | argument path: ``model_options/nnsk/push``
+
+            The parameters to define the push the soft cutoff of nnsk model.
+
+            .. _`model_options/nnsk/push/rs_thr`: 
+
+            rs_thr: 
+                | type: ``float`` | ``int``, optional, default: ``0.0``
+                | argument path: ``model_options/nnsk/push/rs_thr``
+
+                The step size for cutoff value for smooth function in the nnsk anlytical formula.
+
+            .. _`model_options/nnsk/push/rc_thr`: 
+
+            rc_thr: 
+                | type: ``float`` | ``int``, optional, default: ``0.0``
+                | argument path: ``model_options/nnsk/push/rc_thr``
+
+                The step size for cutoff value for smooth function in the nnsk anlytical formula.
+
+            .. _`model_options/nnsk/push/w_thr`: 
+
+            w_thr: 
+                | type: ``float`` | ``int``, optional, default: ``0.0``
+                | argument path: ``model_options/nnsk/push/w_thr``
+
+                The step size for decay factor w.
+
+            .. _`model_options/nnsk/push/period`: 
+
+            period: 
+                | type: ``int``, optional, default: ``100``
+                | argument path: ``model_options/nnsk/push/period``
+
+                the interval of iterations to modify the rs w values.
+
diff --git a/dptb/nn/__init__.py b/dptb/nn/__init__.py
index 4aa0820a..b3b32799 100644
--- a/dptb/nn/__init__.py
+++ b/dptb/nn/__init__.py
@@ -1,10 +1,10 @@
 from .build import build_model
-from .deeptb import DPTB
+from .deeptb import NNENV
 from .nnsk import NNSK
 
 __all__ = [
     build_model,
-    DPTB,
+    NNENV,
     NNSK,
 ]
 """
diff --git a/dptb/nn/build.py b/dptb/nn/build.py
index ed1c057c..8463a9fc 100644
--- a/dptb/nn/build.py
+++ b/dptb/nn/build.py
@@ -1,4 +1,4 @@
-from dptb.nn.deeptb import DPTB, MIX
+from dptb.nn.deeptb import NNENV, MIX
 import logging
 from dptb.nn.nnsk import NNSK
 import torch
@@ -35,7 +35,7 @@ def build_model(
             raise ValueError("You need to provide model_options and common_options when you are initializing a model from scratch.")
 
     # decide whether to initialize a mixed model, or a deeptb model, or a nnsk model
-    init_deeptb = False
+    init_nnenv = False
     init_nnsk = False
     init_mixed = False
 
@@ -78,7 +78,7 @@ def build_model(
             raise ValueError("Model_options are not set correctly!")
     else:
         if all((model_options.get("embedding"), model_options.get("prediction"))):
-            init_deeptb = True
+            init_nnenv = True
             if model_options["prediction"]['method'] == 'sktb':
                 log.warning("The prediction method is sktb, but the nnsk option is not set. this is highly not recommand.\n"+
                             "We recommand to train nnsk then train mix model for sktb. \n"+
@@ -103,15 +103,15 @@ def build_model(
             raise ValueError("Model_options are not set correctly!")
     
     
-    assert int(init_mixed) + int(init_deeptb) + int(init_nnsk) == 1, "You can only choose one of the mixed, deeptb, and nnsk options."
+    assert int(init_mixed) + int(init_nnenv) + int(init_nnsk) == 1, "You can only choose one of the mixed, deeptb, and nnsk options."
     # check if the model is deeptb or nnsk
 
     # init deeptb
     if from_scratch:
-        if init_deeptb:
-            model = DPTB(**model_options, **common_options)
+        if init_nnenv:
+            model = NNENV(**model_options, **common_options)
 
-            # do initialization from statistics if DPTB is e3tb and statistics is provided
+            # do initialization from statistics if NNENV is e3tb and statistics is provided
             if model.method == "e3tb" and statistics is not None:
                 scalar_mask = torch.BoolTensor([ir.dim==1 for ir in model.idp.orbpair_irreps])
                 node_shifts = statistics["node"]["scalar_ave"]
@@ -132,8 +132,8 @@ def build_model(
             
     else:
         # load the model from the checkpoint
-        if init_deeptb:
-            model = DPTB.from_reference(checkpoint, **model_options, **common_options)
+        if init_nnenv:
+            model = NNENV.from_reference(checkpoint, **model_options, **common_options)
         if init_nnsk:
             model = NNSK.from_reference(checkpoint, **model_options["nnsk"], **common_options)
         if init_mixed:
diff --git a/dptb/nn/deeptb.py b/dptb/nn/deeptb.py
index 46bda4a3..ef0c17c4 100644
--- a/dptb/nn/deeptb.py
+++ b/dptb/nn/deeptb.py
@@ -54,9 +54,9 @@ def get_neuron_config(nl):
 
     return config
 
-class DPTB(nn.Module):
+class NNENV(nn.Module):
     quantities = ["hamiltonian", "energy"]
-    name = "dptb"
+    name = "nnenv"
     def __init__(
             self,
             embedding: dict,
@@ -94,7 +94,7 @@ def __init__(
         NotImplementedError
             _description_
         """
-        super(DPTB, self).__init__()
+        super(NNENV, self).__init__()
 
         if isinstance(dtype, str):
             dtype = getattr(torch, dtype)
@@ -308,7 +308,7 @@ def __init__(
 
         self.dtype = dtype
         self.device = device
-        self.dptb = DPTB(
+        self.nnenv = NNENV(
             embedding=embedding, 
             prediction=prediction, 
             basis=basis, 
@@ -332,7 +332,7 @@ def __init__(
         assert not self.nnsk.push, "The push option is not supported in the mixed model. The push option is only supported in the nnsk model."
         
         self.model_options = self.nnsk.model_options
-        self.model_options.update(self.dptb.model_options)
+        self.model_options.update(self.nnenv.model_options)
         
         self.hamiltonian = self.nnsk.hamiltonian
         if overlap:
@@ -341,11 +341,11 @@ def __init__(
 
 
     def forward(self, data: AtomicDataDict.Type):
-        data_dptb = self.dptb(data)
+        data_nnenv = self.nnenv(data)
         data_nnsk = self.nnsk(data)
 
-        data_nnsk[AtomicDataDict.EDGE_FEATURES_KEY] = data_nnsk[AtomicDataDict.EDGE_FEATURES_KEY] * (1 + data_dptb[AtomicDataDict.EDGE_FEATURES_KEY])
-        data_nnsk[AtomicDataDict.NODE_FEATURES_KEY] = data_nnsk[AtomicDataDict.NODE_FEATURES_KEY] * (1 + data_dptb[AtomicDataDict.NODE_FEATURES_KEY])
+        data_nnsk[AtomicDataDict.EDGE_FEATURES_KEY] = data_nnsk[AtomicDataDict.EDGE_FEATURES_KEY] * (1 + data_nnenv[AtomicDataDict.EDGE_FEATURES_KEY])
+        data_nnsk[AtomicDataDict.NODE_FEATURES_KEY] = data_nnsk[AtomicDataDict.NODE_FEATURES_KEY] * (1 + data_nnenv[AtomicDataDict.NODE_FEATURES_KEY])
 
         data_nnsk = self.hamiltonian(data_nnsk)
         if hasattr(self, "overlap"):
diff --git a/dptb/tests/data/test_sktb/dataset/kpath_spk.0/info.json b/dptb/tests/data/test_sktb/dataset/kpath_spk.0/info.json
index d2d7bd90..4304669a 100644
--- a/dptb/tests/data/test_sktb/dataset/kpath_spk.0/info.json
+++ b/dptb/tests/data/test_sktb/dataset/kpath_spk.0/info.json
@@ -9,11 +9,9 @@
         "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/test_build_model.py b/dptb/tests/test_build_model.py
index 475b3fa2..1b81eb91 100644
--- a/dptb/tests/test_build_model.py
+++ b/dptb/tests/test_build_model.py
@@ -1,7 +1,7 @@
 import pytest
 import torch
 from dptb.nn.nnsk import NNSK
-from dptb.nn.deeptb import DPTB, MIX
+from dptb.nn.deeptb import NNENV, MIX
 from dptb.utils.tools import j_must_have
 from dptb.nn.build import build_model
 import os
@@ -117,12 +117,12 @@ def test_build_model_MIX_from_scratch():
 
     assert isinstance(model, MIX)
     assert model.name == "mix"
-    assert model.dptb.method == 'sktb'
-    assert model.dptb.name == 'dptb'
+    assert model.nnenv.method == 'sktb'
+    assert model.nnenv.name == 'nnenv'
     assert model.nnsk.name == 'nnsk'
 
     assert model.nnsk.transform == False
-    assert model.dptb.transform == False
+    assert model.nnenv.transform == False
 
 def test_build_model_failure():
     run_options = {
diff --git a/dptb/utils/argcheck.py b/dptb/utils/argcheck.py
index a39f0e54..34a0b680 100644
--- a/dptb/utils/argcheck.py
+++ b/dptb/utils/argcheck.py
@@ -10,29 +10,80 @@
                         'dptb-hopping_net_neuron', 'dptb-env_net_neuron', 'dptb-soc_net_neuron', 'dptb-onsite_net_neuron', 'dptb-axis_neuron', 'skfunction-skformula', 'sknetwork-sk_onsite_nhidden', 
                         'sknetwork-sk_hop_nhidden']
 
+
+def gen_doc_train(*, make_anchor=True, make_link=True, **kwargs):
+    if make_link:
+        make_anchor = True
+    co = common_options()
+    tr = train_options()
+    da = data_options()
+    mo = model_options()
+    ptr = []
+    ptr.append(co.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs))
+    ptr.append(tr.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs))
+    ptr.append(da.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs))
+    ptr.append(mo.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs))
+
+    key_words = []
+    for ii in "\n\n".join(ptr).split("\n"):
+        if "argument path" in ii:
+            key_words.append(ii.split(":")[1].replace("`", "").strip())
+    # ptr.insert(0, make_index(key_words))
+
+    return "\n\n".join(ptr)
+
+
+def gen_doc_run(*, make_anchor=True, make_link=True, **kwargs):
+    if make_link:
+        make_anchor = True
+    rop = run_options()
+
+    ptr = []
+    ptr.append(rop.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs))
+
+    key_words = []
+    for ii in "\n\n".join(ptr).split("\n"):
+        if "argument path" in ii:
+            key_words.append(ii.split(":")[1].replace("`", "").strip())
+    # ptr.insert(0, make_index(key_words))
+
+    return "\n\n".join(ptr)
+
+
+def gen_doc_setinfo(*, make_anchor=True, make_link=True, **kwargs):
+    if make_link:
+        make_anchor = True
+    sio = set_info_options()
+    ptr = []
+    ptr.append(sio.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs))
+
+    key_words = []
+    for ii in "\n\n".join(ptr).split("\n"):
+        if "argument path" in ii:
+            key_words.append(ii.split(":")[1].replace("`", "").strip())
+    # ptr.insert(0, make_index(key_words))
+
+    return "\n\n".join(ptr)
+
+
 def common_options():
     doc_device = "The device to run the calculation, choose among `cpu` and `cuda[:int]`, Default: `cpu`"
-    doc_dtype = "The digital number's precison, choose among: \n\n\
-            - `float32`: indicating torch.float32\n\n\
-            - `float64`: indicating torch.float64\n\n\
-            Default: `float32`\n\n"
+    doc_dtype = """The digital number's precison, choose among: 
+                    Default: `float32`
+                        - `float32`: indicating torch.float32
+                        - `float64`: indicating torch.float64
+                """
     
     doc_seed = "The random seed used to initialize the parameters and determine the shuffling order of datasets. Default: `3982377700`"
-    #doc_onsite_cutoff = "The cutoff-range considered when using strain mode correction. Out of which the atom are assume to have no effect on current atom's onsite energy."
-    #doc_bond_cutoff = "The cutoff-range of bond hoppings, beyond which it assume the atom pairs have 0 hopping integrals."
-    #doc_env_cutoff = "The cutoff-range of DeePTB environmental correction, recommand range is: (0.5*bond_cutoff, bond_cutoff)"
-    doc_basis = "The atomic orbitals used to construct the basis. E.p. {'A':'2s','2p','s*','B':'3s','3p' }"
-    doc_overlap = ""
+    doc_basis = "The atomic orbitals used to construct the basis. e.p. {'A':['2s','2p','s*'],'B':'[3s','3p']}"
+    doc_overlap = "Whether to calculate the overlap matrix. Default: False"
 
     args = [
-        #Argument("onsite_cutoff", float, optional = True, doc = doc_onsite_cutoff),
-        #Argument("bond_cutoff", float, optional = False, doc = doc_bond_cutoff),
-        #Argument("env_cutoff", float, optional = True, doc = doc_env_cutoff),
         Argument("basis", dict, optional=False, doc=doc_basis),
-        Argument("seed", int, optional=True, default=3982377700, doc=doc_seed),
         Argument("overlap", bool, optional=True, default=False, doc=doc_overlap),
         Argument("device", str, optional = True, default="cpu", doc = doc_device),
         Argument("dtype", str, optional = True, default="float32", doc = doc_dtype),
+        Argument("seed", int, optional=True, default=3982377700, doc=doc_seed),
     ]
 
     doc_common_options = ""
@@ -54,13 +105,15 @@ def train_options():
     "
     doc_lr_scheduler = "The learning rate scheduler tools settings, the lr scheduler is used to scales down the learning rate during the training process. Proper setting can make the training more stable and efficient. The supported lr schedular includes: `Exponential Decaying (exp)`, `Linear multiplication (linear)`"
     doc_batch_size = "The batch size used in training, Default: 1"
+    doc_ref_batch_size = "The batch size used in reference data, Default: 1"
+    doc_val_batch_size = "The batch size used in validation data, Default: 1"
     doc_max_ckpt = "The maximum number of saved checkpoints, Default: 4"
     
     args = [
         Argument("num_epoch", int, optional=False, doc=doc_num_epoch),
         Argument("batch_size", int, optional=True, default=1, doc=doc_batch_size),
-        Argument("ref_batch_size", int, optional=True, default=1, doc=doc_batch_size),
-        Argument("val_batch_size", int, optional=True, default=1, doc=doc_batch_size),
+        Argument("ref_batch_size", int, optional=True, default=1, doc=doc_ref_batch_size),
+        Argument("val_batch_size", int, optional=True, default=1, doc=doc_val_batch_size),
         Argument("optimizer", dict, sub_fields=[], optional=True, default={}, sub_variants=[optimizer()], doc = doc_optimizer),
         Argument("lr_scheduler", dict, sub_fields=[], optional=True, default={}, sub_variants=[lr_scheduler()], doc = doc_lr_scheduler),
         Argument("save_freq", int, optional=True, default=10, doc=doc_save_freq),
@@ -76,7 +129,7 @@ def train_options():
 
 def test_options():
     doc_display_freq = "Frequency, or every how many iteration to display the training log to screem. Default: `1`"
-    doc_batch_size = ""
+    doc_batch_size = "The batch size used in testing, Default: 1"
     
     args = [
         Argument("batch_size", int, optional=True, default=1, doc=doc_batch_size),
@@ -206,7 +259,7 @@ def train_data_sub():
         Argument("get_eigenvalues", bool, optional=True, default=False, doc=doc_eig)
     ]
 
-    doc_train = ""
+    doc_train = "The dataset settings for training."
 
     return Argument("train", dict, optional=False, sub_fields=args, sub_variants=[], doc=doc_train)
 
@@ -217,18 +270,20 @@ def validation_data_sub():
     doc_eig = "Choose whether the eigenvalues and k-points are loaded when building dataset."
     doc_vlp = "Choose whether the overlap blocks are loaded when building dataset."
     doc_DM = "Choose whether the density matrix is loaded when building dataset."
+    doc_separator = "the sepatator used to separate the prefix and suffix in the dataset directory. Default: '.'"
 
     args = [
         Argument("type", str, optional=True, default="DefaultDataset", doc="The type of dataset."),
         Argument("root", str, optional=False, doc=doc_root),
         Argument("prefix", str, optional=True, default=None, doc=doc_prefix),
+        Argument("separator", str, optional=True, default='.', doc=doc_separator),
         Argument("get_Hamiltonian", bool, optional=True, default=False, doc=doc_ham),
         Argument("get_overlap", bool, optional=True, default=False, doc=doc_vlp),
         Argument("get_DM", bool, optional=True, default=False, doc=doc_DM),
         Argument("get_eigenvalues", bool, optional=True, default=False, doc=doc_eig)
     ]
 
-    doc_validation = ""
+    doc_validation = "The dataset settings for validation."
 
     return Argument("validation", dict, optional=True, sub_fields=args, sub_variants=[], doc=doc_validation)
 
@@ -239,18 +294,20 @@ def reference_data_sub():
     doc_eig = "Choose whether the eigenvalues and k-points are loaded when building dataset."
     doc_vlp = "Choose whether the overlap blocks are loaded when building dataset."
     doc_DM = "Choose whether the density matrix is loaded when building dataset."
+    doc_separator = "the sepatator used to separate the prefix and suffix in the dataset directory. Default: '.'"
 
     args = [
         Argument("type", str, optional=True, default="DefaultDataset", doc="The type of dataset."),
         Argument("root", str, optional=False, doc=doc_root),
         Argument("prefix", str, optional=True, default=None, doc=doc_prefix),
+        Argument("separator", str, optional=True, default='.', doc=doc_separator),
         Argument("get_Hamiltonian", bool, optional=True, default=False, doc=doc_ham),
         Argument("get_overlap", bool, optional=True, default=False, doc=doc_vlp),
         Argument("get_DM", bool, optional=True, default=False, doc=doc_DM),
         Argument("get_eigenvalues", bool, optional=True, default=False, doc=doc_eig)
     ]
 
-    doc_reference = ""
+    doc_reference = "The dataset settings for reference."
     
     return Argument("reference", dict, optional=True, sub_fields=args, sub_variants=[], doc=doc_reference)
 
@@ -259,30 +316,34 @@ def test_data_sub():
     doc_prefix = "The prefix of the folders under root, which will be loaded in dataset."
     doc_ham = "Choose whether the Hamiltonian blocks (and overlap blocks, if provided) are loaded when building dataset."
     doc_eig = "Choose whether the eigenvalues and k-points are loaded when building dataset."
+    doc_vlp = "Choose whether the overlap blocks are loaded when building dataset."
+    doc_DM = "Choose whether the density matrix is loaded when building dataset."
+    doc_separator = "the sepatator used to separate the prefix and suffix in the dataset directory. Default: '.'"
 
     args = [
         Argument("type", str, optional=True, default="DefaultDataset", doc="The type of dataset."),
         Argument("root", str, optional=False, doc=doc_root),
         Argument("prefix", str, optional=True, default=None, doc=doc_prefix),
         Argument("get_Hamiltonian", bool, optional=True, default=False, doc=doc_ham),
-        Argument("get_eigenvalues", bool, optional=True, default=False, doc=doc_eig)
+        Argument("get_eigenvalues", bool, optional=True, default=False, doc=doc_eig),
+        Argument("get_overlap", bool, optional=True, default=False, doc=doc_vlp),
+        Argument("get_DM", bool, optional=True, default=False, doc=doc_DM),
+        Argument("separator", str, optional=True, default='.', doc=doc_separator)
     ]
 
-    doc_reference = ""
+    doc_test = "The dataset settings for testing."
     
-    return Argument("test", dict, optional=False, sub_fields=args, default={}, sub_variants=[], doc=doc_reference)
+    return Argument("test", dict, optional=False, sub_fields=args, default={}, sub_variants=[], doc=doc_test)
 
 
 def data_options():
-    doc_use_reference = "Whether to use a reference dataset that jointly train the model. It acting as a constraint or normalization to make sure the model won't deviate too much from the reference data."
-
     args = [
             train_data_sub(),
             validation_data_sub(),
             reference_data_sub()
             ]
 
-    doc_data_options = ""
+    doc_data_options = "The options for dataset settings in training."
 
     return Argument("data_options", dict, sub_fields=args, sub_variants=[], optional=False, doc=doc_data_options)
 
@@ -292,53 +353,13 @@ def test_data_options():
         test_data_sub()
     ]
 
-    doc_test_data_options = "parameters for dataset settings in testing"
+    doc_test_data_options = "The options for dataset settings in testing"
 
     return Argument("data_options", dict, sub_fields=args, sub_variants=[], optional=False, doc=doc_test_data_options)
 
-# def dptb():
-#     doc_soc_env = "button that allow environmental correction for soc parameters, used only when soc is open, Default: False"
-#     doc_axis_neuron = "The axis_neuron specifies the size of the submatrix of the embedding matrix, the axis matrix as explained in the [DeepPot-SE paper](https://arxiv.org/abs/1805.09003)."
-#     doc_onsite_net_neuron = r"The number of hidden neurons in the network for onsites $\langle i|H|i\rangle$ of `dptb` model. Default: `[128, 128, 256, 256]`"
-#     doc_env_net_neuron = "The number of hidden neurons in the environment embedding network of `dptb` model. Default: `[128, 128, 256, 256]`"
-#     doc_hopping_net_neuron = r"The number of hidden neurons in the network for hoppings $\langle i|H|j\rangle$ of `dptb` model. Default: `[128, 128, 256, 256]`"
-#     doc_onsite_net_activation = "The activation function for onsite networks. Default: `tanh`"
-#     doc_env_net_activation = "The activation function for environment embedding networks. Default: `tanh`"
-#     doc_hopping_net_activation = "The activation function for hopping networks. Default: `tanh`"
-#     doc_soc_net_activation = "The activation function for soc networks. Default: `tanh`"
-#     doc_soc_net_neuron = r"The number of hidden neurons in the network for soc $\lambda$ of `dptb` model. Default: `[128, 128, 256, 256]`"
-#     doc_soc_net_type = "The network type for soc, the value can be:\n\n\
-#         - `res`: for feedforward Network with residual connections, for more information about residual network, we refer to [Deep Residual Learning for Image Recognition](https://openaccess.thecvf.com/content_cvpr_2016/papers/He_Deep_Residual_Learning_CVPR_2016_paper.pdf). \n\n\
-#         - `ffn`: for feedforward Network."
-#     doc_onsite_net_type = "The network type for onsites."
-#     doc_env_net_type = "The network type for environment embeddings."
-#     doc_hopping_net_type = "The network type for hoppings."
-#     doc_if_batch_normalized = "Whether to use batch normalization after each layer in neural network. The batch normalization normalize the itermidiate values in neural network with the mean and variance estimated in the batch dimension. The batch here means the batch of onsite or hopping embeddings or position vectors that processed by neural network at one time computation, which is different from the batch defined in `data_options`. Default: False."
-
-#     args = [
-#         Argument("soc_env", bool, optional=True, default=False, doc=doc_soc_env),
-#         Argument("axis_neuron", int, optional=True, default=10, doc=doc_axis_neuron),
-#         Argument("onsite_net_neuron", list, optional=True, default=[128, 128, 256, 256], doc=doc_onsite_net_neuron),
-#         Argument("soc_net_neuron", list, optional=True, default=[128, 128, 256, 256], doc=doc_soc_net_neuron),
-#         Argument("env_net_neuron", list, optional=True, default=[128, 128, 256, 256], doc=doc_env_net_neuron),
-#         Argument("hopping_net_neuron", list, optional=True, default=[128, 128, 256, 256], doc=doc_hopping_net_neuron),
-#         Argument("onsite_net_activation", str, optional=True, default="tanh", doc=doc_onsite_net_activation),
-#         Argument("env_net_activation", str, optional=True, default="tanh", doc=doc_env_net_activation),
-#         Argument("hopping_net_activation", str, optional=True, default="tanh", doc=doc_hopping_net_activation),
-#         Argument("soc_net_activation", str, optional=True, default="tanh", doc=doc_soc_net_activation),
-#         Argument("onsite_net_type", str, optional=True, default="res", doc=doc_onsite_net_type),
-#         Argument("env_net_type", str, optional=True, default="res", doc=doc_env_net_type),
-#         Argument("hopping_net_type", str, optional=True, default="res", doc=doc_hopping_net_type),
-#         Argument("soc_net_type", str, optional=True, default="res", doc=doc_soc_net_type),
-#         Argument("if_batch_normalized", bool, optional=True, default=False, doc=doc_if_batch_normalized)
-#     ]
-
-#     doc_dptb = "The parameters for `dptb` model, which maps the environmental information to Tight-Binding parameters."
-
-#     return Argument("dptb", dict, optional=True, sub_fields=args, sub_variants=[], default={}, doc=doc_dptb)
 
 def embedding():
-    doc_method = ""
+    doc_method = "The parameters to define the embedding model."
 
     return Variant("method", [
             Argument("se2", dict, se2()),
@@ -356,14 +377,14 @@ def embedding():
 
 def se2():
 
-    doc_rs = ""
-    doc_rc = ""
-    doc_n_axis = ""
-    doc_radial_net = ""
+    doc_rs = "The soft cutoff where the smooth function starts."
+    doc_rc = "The hard cutoff where the smooth function value ~0.0"
+    doc_n_axis = "the out axis shape of the deepmd-se2 descriptor."
+    doc_radial_net = "network to build the descriptors."
 
-    doc_neurons = ""
-    doc_activation = ""
-    doc_if_batch_normalized = ""
+    doc_neurons = "the size of nn for descriptor"
+    doc_activation = "activation"
+    doc_if_batch_normalized = "whether to turn on the batch normalization."
 
     radial_net = [
         Argument("neurons", list, optional=False, doc=doc_neurons),
@@ -468,9 +489,6 @@ def e3baselinev5():
     doc_r_max = ""
     doc_n_layers = ""
     doc_env_embed_multiplicity = ""
-    doc_linear_after_env_embed = ""
-    doc_latent_resnet_update_ratios_learnable = ""
-    doc_latent_kwargs = ""
 
     return [
             Argument("irreps_hidden", str, optional=False, doc=doc_irreps_hidden),
@@ -493,9 +511,8 @@ def e3baselinev5():
 
 
 def prediction():
-    doc_method = ""
-    doc_nn = ""
-    doc_linear = ""
+    doc_method = "The options to indicate the prediction model. Can be sktb or e3tb."
+    doc_nn = "neural network options for prediction model."
 
     return Variant("method", [
             Argument("sktb", dict, sktb_prediction(), doc=doc_nn),
@@ -503,12 +520,9 @@ def prediction():
         ], optional=False, doc=doc_method)
 
 def sktb_prediction():
-    doc_neurons = ""
-    doc_activation = ""
-    doc_if_batch_normalized = ""
-    doc_quantities = ""
-    doc_hamiltonian = ""
-    doc_precision = ""
+    doc_neurons = "neurons in the neural network."
+    doc_activation = "activation function."
+    doc_if_batch_normalized = "if to turn on batch normalization"
 
     nn = [
         Argument("neurons", list, optional=False, doc=doc_neurons),
@@ -520,8 +534,8 @@ def sktb_prediction():
 
 
 def e3tb_prediction():
-    doc_scales_trainable = ""
-    doc_shifts_trainable = ""
+    doc_scales_trainable = "whether to scale the trianing target."
+    doc_shifts_trainable = "whether to shift the training target."
 
     nn = [
         Argument("scales_trainable", bool, optional=True, default=False, doc=doc_scales_trainable),
@@ -534,9 +548,9 @@ def e3tb_prediction():
 
 def model_options():
 
-    doc_model_options = "The parameters to define the `nnsk` and `dptb` model."
-    doc_embedding = ""
-    doc_prediction = ""
+    doc_model_options = "The parameters to define the `nnsk`,`mix` and `dptb` model."
+    doc_embedding = "The parameters to define the embedding model."
+    doc_prediction = "The parameters to define the prediction model"
 
     return Argument("model_options", dict, sub_fields=[
         Argument("embedding", dict, optional=True, sub_fields=[], sub_variants=[embedding()], doc=doc_embedding),
@@ -574,8 +588,8 @@ def nnsk():
 def push():
     doc_rs_thr = "The step size for cutoff value for smooth function in the nnsk anlytical formula."
     doc_rc_thr = "The step size for cutoff value for smooth function in the nnsk anlytical formula."
-    doc_w_thr = ""
-    doc_period = ""
+    doc_w_thr = "The step size for decay factor w."
+    doc_period = "the interval of iterations to modify the rs w values."
 
     return Argument("push", [bool,dict], sub_fields=[
         Argument("rs_thr", [int,float], optional=True, default=0., doc=doc_rs_thr),
@@ -585,18 +599,18 @@ def push():
     ], sub_variants=[], optional=True, default=False, doc="The parameters to define the push the soft cutoff of nnsk model.")
 
 def onsite():
-    doc_method = r"The onsite correction mode, the onsite energy is expressed as the energy of isolated atoms plus the model correction, the correction mode are:\n\n\
-            - `strain`: The strain mode correct the onsite matrix densly by $$H_{i,i}^{lm,l^\prime m^\prime} = \epsilon_l^0 \delta_{ll^\prime}\delta_{mm^\prime} + \sum_p \sum_{\zeta} \Big[ \mathcal{U}_{\zeta}(\hat{\br}_{ip}) \ \epsilon_{ll^\prime \zeta} \Big]_{mm^\prime}$$ which is also parameterized as a set of Slater-Koster like integrals.\n\n\
-            - `uniform`: The correction is a energy shift respect of orbital of each atom. Which is formally written as: \n\n\
-                  $$H_{i,i}^{lm,l^\prime m^\prime} = (\epsilon_l^0+\epsilon_l^\prime) \delta_{ll^\prime}\delta_{mm^\prime}$$ Where $\epsilon_l^0$ is the isolated energy level from the DeePTB onsite database, and $\epsilon_l^\prime$ is the parameters to fit. E.p. \n\n\
-            - `split`: (not recommanded) The split onsite mode correct onsite hamiltonian with a magnetic quantum number dependent form, which violate the rotation equivariace, but some times can be effective. The formula is: \
-                $$H_{i,i}^{lm,l^\prime m^\prime} = (\epsilon_l^0+\epsilon_{lm}^\prime) \delta_{ll^\prime}\delta_{mm^\prime}$$ \n\n\
-            Default: `none`"
+    doc_method = """The onsite correction mode, the onsite energy is expressed as the energy of isolated atoms plus the model correction, the correction mode are:
+                    Default: `none`: use the database onsite energy value.
+                    - `strain`: The strain mode correct the onsite matrix densly by $$H_{i,i}^{lm,l^\prime m^\prime} = \epsilon_l^0 \delta_{ll^\prime}\delta_{mm^\prime} + \sum_p \sum_{\zeta} \Big[ \mathcal{U}_{\zeta}(\hat{\br}_{ip}) \ \epsilon_{ll^\prime \zeta} \Big]_{mm^\prime}$$ which is also parameterized as a set of Slater-Koster like integrals.\n\n\
+                    - `uniform`: The correction is a energy shift respect of orbital of each atom. Which is formally written as: 
+                                $$H_{i,i}^{lm,l^\prime m^\prime} = (\epsilon_l^0+\epsilon_l^\prime) \delta_{ll^\prime}\delta_{mm^\prime}$$ Where $\epsilon_l^0$ is the isolated energy level from the DeePTB onsite database, and $\epsilon_l^\prime$ is the parameters to fit.
+                    - `NRL`: use the NRL-TB formula.
+                """
 
-    doc_rs = ""
-    doc_w = ""
-    doc_rc = ""
-    doc_lda = ""
+    doc_rs = "The smooth cutoff `fc` for strain model. rs is where fc = 0.5"
+    doc_w = "The decay factor of `fc` for strain and nrl model."
+    doc_rc = "The smooth cutoff of `fc` for nrl model, rc is where fc ~ 0.0"
+    doc_lda = "The lambda type encoding value in nrl model. now only support elementary substance"
 
     strain = [
         Argument("rs", float, optional=True, default=6.0, doc=doc_rs),
@@ -617,10 +631,15 @@ def onsite():
                 ],optional=False, doc=doc_method)
 
 def hopping():
-    doc_method = ""
-    doc_rs = ""
-    doc_w = ""
-    doc_rc = ""
+    doc_method = """The hopping formula. 
+                    -  `powerlaw`: the powerlaw formula for bond length dependence for sk integrals.
+                    -  `varTang96`: a variational formula based on Tang96 formula.
+                    -  `NRL0`: the old version of NRL formula for overlap, we set overlap and hopping share same options.
+                    -  `NRL1`: the new version of NRL formula for overlap. 
+                    """
+    doc_rs = "The cut-off for smooth function fc for powerlaw and varTang96, fc(rs)=0.5"
+    doc_w = " The decay w in fc"
+    doc_rc = "The cut-off for smooth function fc for NRL, fc(rc) = 0."
 
     powerlaw = [
         Argument("rs", float, optional=True, default=6.0, doc=doc_rs),
@@ -648,14 +667,15 @@ def hopping():
     
 
 def loss_options():
-    doc_method = "The loss function type, defined by a string like `<fitting target>_<loss type>`, Default: `eigs_l2dsf`. supported loss functions includes:\n\n\
-    - `eig_l2`: The l2 norm of predicted and labeled eigenvalues.\n\n\
-    - `eigs_l2d`: The l2 norm and the random differences of the predicted and labeled eigenvalues.\n\n\
-    - `block_l2`: \n\n\
-        Notice: The loss option define here only affect the training loss function, the loss for evaluation will always be `eig_l2`, as it compute the standard MSE of fitted eigenvalues."
-    doc_train = ""
-    doc_validation = ""
-    doc_reference = ""
+    doc_method = """The loss function type, defined by a string like `<fitting target>_<loss type>`, Default: `eigs_l2dsf`. supported loss functions includes:\n\n\
+                    - `eigvals`: The mse loss predicted and labeled eigenvalues and Delta eigenvalues between different k.
+                    - `hamil`: 
+                    - `hamil_abs`:
+                    - `hamil_blas`:
+                """
+    doc_train = "Loss options for training."
+    doc_validation = "Loss options for validation."
+    doc_reference = "Loss options for reference data in training."
 
     hamil = [
         Argument("onsite_shift", bool, optional=True, default=False, doc="Whether to use onsite shift in loss function. Default: False"),
@@ -972,7 +992,7 @@ def fmm():
         Argument("err", [int, float], optional=True, default=1e-5, doc=doc_err)
     ]
 
-def normalize_run(data):
+def run_options():
     doc_task = "the task to run, includes: band, dos, pdos, FS2D, FS3D, ifermi"
     doc_structure = "the structure to run the task"
     doc_gui = "To use the GUI or not"
@@ -984,9 +1004,13 @@ def normalize_run(data):
         AtomicData_options_sub()
     ]
 
-    base = Argument("base", dict, args)
-    data = base.normalize_value(data)
-    base.check_value(data, strict=True)
+    return Argument("run_op", dict, args)
+
+def normalize_run(data):
+
+    run_op = run_options()
+    data = run_op.normalize_value(data)
+    run_op.check_value(data, strict=True)
     
     return data
 
@@ -1016,14 +1040,18 @@ def task_options():
         ],optional=False, doc=doc_task)
 
 def band():
-    doc_kline_type = ""
-    doc_kpath = ""
-    doc_klabels = ""
-    doc_emin=""
-    doc_emax=""
-    doc_E_fermi = ""
-    doc_ref_band = ""
-    doc_nel_atom = "Dict: The valence electron number of each type of atom."
+    doc_kline_type ="""The different type to build kpath line mode.
+                    - "abacus" : the abacus format 
+                    - "vasp" : the vasp format
+                    - "ase" : the ase format
+                    """
+    doc_kpath = "for abacus, this is list, for vasp it is a string to specifc the kpath."
+    doc_klabels = "the labels for high symmetry kpoint"
+    doc_emin="the min energy to show the band plot"
+    doc_emax="the max energy to show the band plot"
+    doc_E_fermi = "the fermi level used to plot band"
+    doc_ref_band = "the reference band structure to be ploted together with dptb bands."
+    doc_nel_atom = "the valence electron number of each type of atom."
     
     return [
         Argument("kline_type", str, optional=False, doc=doc_kline_type),
@@ -1272,16 +1300,15 @@ def normalize_bandinfo(data):
     return data
 
 def bandinfo_sub():
-    doc_nkpoints = ""
-    doc_nbands = ""
-    doc_band_min = ""
-    doc_band_max = ""
-    doc_emin = ""
-    doc_emax = ""
+    doc_band_min = """the minum band index for the training band window with respected to the correctly selected DFT bands.
+                   `important`: before setting this tag you should make sure you have already  exclude all the irrelevant in your training data.
+                                This logic for band_min and max is based on the simple fact the total number TB bands > the bands you care.   
+                   """
+    doc_band_max = "The maxmum band index for training band window"
+    doc_emin = "the minmum energy window, 0 meand the min value of the band at index band_min"
+    doc_emax = "the max energy window, emax value is respect to the min value of the band at index band_min"
     
     args = [
-        Argument("nkpoints", int, optional=True, doc=doc_nkpoints, default=0),
-        Argument("nbands", int, optional=True, doc=doc_nbands, default=0),
         Argument("band_min", int, optional=True, doc=doc_band_min, default=0),
         Argument("band_max", [int, None], optional=True, doc=doc_band_max, default=None),
         Argument("emin", [float, None], optional=True, doc=doc_emin,default=None),
@@ -1291,10 +1318,10 @@ def bandinfo_sub():
     return Argument("bandinfo", dict, optional=True, sub_fields=args, sub_variants=[], doc="")
 
 def AtomicData_options_sub():
-    doc_r_max = ""
-    doc_er_max = ""
-    doc_oer_max = ""
-    doc_pbc = ""
+    doc_r_max = "the cutoff value for bond considering in TB model."
+    doc_er_max = "The cutoff value for environment for each site for env correction model. should set for nnsk+env correction model."
+    doc_oer_max = "The cutoff value for onsite environment for nnsk model, for now only need to set in strain and NRL mode."
+    doc_pbc = "The periodic condition for the structure, can bool or list of bool to specific x,y,z direction."
     
     args = [
         Argument("r_max", [float, int, dict], optional=False, doc=doc_r_max, default=4.0),
@@ -1305,7 +1332,7 @@ def AtomicData_options_sub():
 
     return Argument("AtomicData_options", dict, optional=False, sub_fields=args, sub_variants=[], doc="")
 
-def normalize_setinfo(data):
+def set_info_options():
     doc_nframes = "Number of frames in this trajectory."
     doc_natoms = "Number of atoms in each frame."
     doc_pos_type = "Type of atomic position input. Can be frac / cart / ase."
@@ -1317,7 +1344,12 @@ def normalize_setinfo(data):
         bandinfo_sub(),
         AtomicData_options_sub()
     ]
-    setinfo = Argument("setinfo", dict, sub_fields=args)
+
+    return Argument("setinfo", dict, sub_fields=args)
+
+def normalize_setinfo(data):
+
+    setinfo = set_info_options()
     data = setinfo.normalize_value(data)
     setinfo.check_value(data, strict=True)