Added model export functionality to convert trained models to Wannier90 and PythTB formats. (#300)

* feat: Add export to Wannier90  and PythTB functionality

* Add exmple for topythtb

* feat: update PythTB export functionality

* update docs

* fix towannier

* refactor: remove debug prints from load_data_for_model function

Removed multiple debug print statements that were used for troubleshooting data object attributes and transformations. The debug output was checking for cell attribute existence, data object keys, and transformation results. These temporary debugging statements are no longer needed and have been cleaned up to improve code readability.

* refactor: simplify ToWannier90 data handling and structure reconstruction

Remove complex ASE atoms reconstruction logic in _get_data_and_blocks method. Store essential data (positions, cell, atom types, symbols) as instance attributes for reuse. Simplify return values by removing structase parameter. Streamline orbital mapping and hopping collection logic in write_hr method.

* refactor: reorganize examples directory structure and notebooks

- Rename ToPythTB directory to ToW90_PythTB for clarity
- Move and update ToPythTB.ipynb with corrected paths and execution counts
- Add new ToW90_pythtb.ipynb for Wannier90 export verification
- Update .gitignore to exclude example output files (*.xyz, *.win)

* docs: add Wannier90 API example notebook

Add a new Jupyter notebook demonstrating how to export DeePTB models to Wannier90 format and validate the results by comparing band structures. The notebook includes step-by-step instructions for loading a model, exporting to _hr.dat format, and verifying correctness through PythTB interface comparison. Update the API documentation index to include the new example.

* test: Install optional pythtb dependencies for testing.

* fix: remove redundant multiplication by zero in Fermi energy subtraction

The multiplication by 0 in the Fermi energy subtraction was redundant and has been removed. This simplifies the code without changing its behavior.

* feat: refactor orbital handling and improve file resource management

- Add get_orbitals_for_type helper function to centralize orbital expansion logic
- Update load_data_for_model to use context manager for HDF5 file handling
- Refactor ToWannier90 class to use new orbital helper function
- Improve type hints with Optional for better parameter clarity
- Fix trailing whitespace in documentation file

* feat: simplify Wannier90 export interface

Remove redundant struct_path parameter from write_win and write_centres methods in ToWannier90 class. Update export entrypoint and documentation notebooks to reflect the simplified API. Also improve device handling in ToPythTB class with better default value management.

* refactor: simplify Wannier90 export method calls

Remove redundant structure file parameter from write_win and write_centres methods as it's already handled internally by the exporter. This streamlines the API and reduces parameter duplication.

Author: QG-phy

.gitignore

@@ -1,4 +1,6 @@
 test*.ipynb
+examples/**/*centres.xyz
+examples/**/*.win
 **/processed*/*
 dptb/tests/**/*.yaml
 dptb/tests/**/pre_*.pt

docs/index.rst

@@ -40,7 +40,7 @@ For more details, see our papers:
    quick_start/easy_install
    quick_start/input
    quick_start/hands_on/index
-   quick_start/basic_api
+   quick_start/api/index
 
 
 .. toctree::

docs/quick_start/basic_api.md docs/quick_start/api/basic_api.mddocs/quick_start/basic_api.md renamed to docs/quick_start/api/basic_api.md

@@ -0,0 +1,9 @@
+=================================================
+A quick Example
+=================================================
+
+.. toctree::
+   basic_api
+   pythtb_api
+   wan90_api
+   

docs/quick_start/api/index.rst

@@ -0,0 +1,114 @@
+import json
+import torch
+import os
+import argparse
+import logging
+from dptb.nn.build import build_model
+from dptb.utils.tools import j_loader, j_must_have
+from dptb.postprocess.interfaces import ToWannier90, ToPythTB
+
+log = logging.getLogger(__name__)
+
+def export(INPUT: str, init_model: str = None, structure: str = None, output: str = "./", format: str = "wannier90", **kwargs):
+    """
+    Export the DeePTB model to external formats.
+    """
+    
+    # 1. Load config
+    if INPUT and INPUT.endswith(".json"):
+        jdata = j_loader(INPUT)
+    else:
+        # Fallback if no json provided (though usually required)
+        # If user provides structure + model but no config? 
+        # DeePTB typically needs config for model architecture parameters stored in it (or in checkpoint?)
+        # build_model usually needs common_options.
+        # If init_model is a .pth, it might contain config?
+        # dptb.nn.build.build_model(checkpoint=...) handles loading from checkpoint which might have config.
+        jdata = {}
+
+    # Determine device
+    device_str = jdata.get("device", "cpu")
+    device = torch.device(device_str)
+    
+    # Load model
+    ckpt_path = init_model if init_model else jdata.get("model_ckpt")
+    if not ckpt_path:
+        log.error("Model checkpoint not found. Provide via -i or in config json via 'model_ckpt'.")
+        return
+    
+    # Common options for build_model
+    # We might need to extract them from jdata or checkpoint
+    # build_model(checkpoint=...) attempts to load config from checkpoint if available?
+    # run.py passes in_common_options loaded from json.
+    
+    in_common_options = {}
+    if jdata.get("device"): in_common_options["device"] = jdata["device"]
+    if jdata.get("dtype"): in_common_options["dtype"] = jdata["dtype"]
+    
+    try:
+        model = build_model(checkpoint=ckpt_path, common_options=in_common_options)
+        model.to(device)
+        model.eval()
+    except Exception as e:
+        log.error(f"Failed to load model: {e}")
+        return
+    
+    # Determine structure input
+    struct_path = structure if structure else jdata.get("structure")
+    if not struct_path:
+        log.error("Structure file not specified. Provide via -stu or in config json.")
+        return
+        
+    # Ensure output directory exists
+    if not os.path.exists(output):
+        os.makedirs(output)
+        
+    log.info(f"Loaded model from {ckpt_path}")
+    log.info(f"Processing structure {struct_path}")
+    log.info(f"Exporting to format: {format}")
+    
+    # 2. Dispatch to Interface
+    
+    if format.lower() in ["wannier90", "w90", "tb2j", "wannierberri"]:
+        exporter = ToWannier90(model, device=device)
+        
+        prefix = os.path.splitext(os.path.basename(struct_path))[0]
+        hr_file = os.path.join(output, f"{prefix}_hr.dat")
+        win_file = os.path.join(output, f"{prefix}.win")
+        cen_file = os.path.join(output, f"{prefix}_centres.xyz")
+        
+        ad_options = jdata.get("AtomicData_options", {})
+        e_fermi = jdata.get("e_fermi", 0.0)
+        
+        try:
+            exporter.write_hr(struct_path, filename=hr_file, AtomicData_options=ad_options, e_fermi=e_fermi)
+            exporter.write_win(filename=win_file)
+            exporter.write_centres(filename=cen_file)
+        except Exception as e:
+            log.error(f"Error during export: {e}")
+        
+    elif format.lower() in ["pythtb"]:
+        try:
+            exporter = ToPythTB(model, device=device)
+        except ImportError:
+            return # Logged in __init__
+            
+        ad_options = jdata.get("AtomicData_options", {})
+        e_fermi = jdata.get("e_fermi", 0.0)
+        
+        try:
+            tb_model = exporter.get_model(struct_path, AtomicData_options=ad_options, e_fermi=e_fermi)
+            
+            prefix = os.path.splitext(os.path.basename(struct_path))[0]
+            pkl_file = os.path.join(output, f"{prefix}_pythtb.pkl")
+            
+            import pickle
+            with open(pkl_file, 'wb') as f:
+                pickle.dump(tb_model, f)
+                
+            log.info(f"Saved PythTB model object to {pkl_file}")
+        except Exception as e:
+            log.error(f"Error during PythTB export: {e}")
+        
+    else:
+        log.error(f"Unknown format: {format}")

docs/quick_start/api/pythtb_api.ipynb

@@ -13,7 +13,9 @@
 from dptb.utils.loggers import set_log_handles
 from dptb.utils.config_check import check_config_train
 from dptb.entrypoints.collectskf import skf2pth, skf2nnsk
+
 from dptb.entrypoints.emp_sk import to_empsk
+from dptb.entrypoints.export import export
 
 from dptb import __version__
 
@@ -443,6 +445,54 @@ def main_parser() -> argparse.ArgumentParser:
         type=float,
         help="Enable SOC, default 0.2 if no value is given. Example: --soc or --soc=0.5"
     )
+
+
+    # export parser
+    parser_export = subparsers.add_parser(
+        "export",
+        parents=[parser_log],
+        help="Export DeePTB model to external formats (Wannier90, PythTB, TB2J)",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    
+    parser_export.add_argument(
+        "INPUT", help="the input parameter file in json format",
+        type=str,
+        default=None
+    )
+    
+    parser_export.add_argument(
+        "-i",
+        "--init-model",
+        type=str,
+        default=None,
+        help="Initialize the model by the provided checkpoint.",
+    )
+
+    parser_export.add_argument(
+        "-stu",
+        "--structure",
+        type=str,
+        default=None,
+        help="the structure file name."
+    )
+    
+    parser_export.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        default="./",
+        help="The output directory."
+    )
+    
+    parser_export.add_argument(
+        "-f",
+        "--format",
+        type=str,
+        default="wannier90",
+        help="Target format: wannier90 (default) or pythtb."
+    )
+
     return parser
 
 def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
@@ -509,3 +559,6 @@ def main():
 
     elif args.command == 'esk':
         to_empsk(**dict_args)
+
+    elif args.command == 'export':
+        export(**dict_args)