style: optmize the import of each submodule

dptb/data/__init__.py

@@ -14,14 +14,13 @@
     AtomicInMemoryDataset,
     NpzDataset,
     ASEDataset,
-    HDF5Dataset,
     ABACUSDataset,
     ABACUSInMemoryDataset,
     DefaultDataset
 )
 from .dataloader import DataLoader, Collater, PartialSampler
-from .build import dataset_from_config
-from .test_data import EMTTestDataset
+from .build import build_dataset
+from .transforms import OrbitalMapper
 
 __all__ = [
     AtomicData,
@@ -33,17 +32,16 @@
     AtomicInMemoryDataset,
     NpzDataset,
     ASEDataset,
-    HDF5Dataset,
     ABACUSDataset,
     ABACUSInMemoryDataset,
     DefaultDataset,
     DataLoader,
     Collater,
     PartialSampler,
-    dataset_from_config,
+    OrbitalMapper,
+    build_dataset,
     _NODE_FIELDS,
     _EDGE_FIELDS,
     _GRAPH_FIELDS,
     _LONG_FIELDS,
-    EMTTestDataset,
 ]

dptb/nn/__init__.py

@@ -1,11 +1,21 @@
 from .build import build_model
-from .deeptb import NNENV
+from .deeptb import NNENV, MIX
 from .nnsk import NNSK
+from .dftbsk import DFTBSK
+from .hamiltonian import E3Hamiltonian, SKHamiltonian
+from .hr2hk import HR2HK
+from .energy import Eigenvalues
 
 __all__ = [
     build_model,
+    E3Hamiltonian,
+    SKHamiltonian,
+    HR2HK,
+    Eigenvalues,
     NNENV,
     NNSK,
+    MIX,
+    DFTBSK,
 ]
 """
 

dptb/postprocess/__init__.py

@@ -0,0 +1,11 @@
+from .bandstructure import Band
+from .totbplas import TBPLaS
+from .write_ham import write_ham
+
+
+__all__ = [
+    Band,
+    TBPLaS,
+    write_ham,
+
+]

dptb/postprocess/bandstructure/__init__.py

@@ -0,0 +1 @@
+from .band import Band