Feat: Add new unifined postprocess class to handle DeePTB model. (#301)

* new_post_v0

* new post v0.2

* feat: add support for built-in polynomial checkpoints

Enhance build_model to handle 'poly2' and 'poly4' as special checkpoint names, automatically resolving them to their corresponding base model files in the dftb directory. This simplifies model initialization for common polynomial baselines.

* feat: add GUI detection and refactor calculator interface

Add is_gui_available() function to detect matplotlib GUI display availability in various environments including Jupyter notebooks and different operating systems. Refactor HamiltonianCalculator from Protocol to ABC with proper abstract methods for better type safety and inheritance. Update DeePTBAdapter to inherit from the new abstract base class.

* feat: Implement unified post-processing for band structure and DOS, including new examples and tests.

* feat: add get_hk method to HamiltonianCalculator interface

Add new abstract method get_hk() to calculate H(k) and S(k) matrices at specified k-points. Rename get_hamiltonian_blocks() to get_hr() for clarity. Implement get_hk() in DeePTBAdapter with support for custom k-point injection and proper tensor handling. Remove redundant get_hamiltonian_at_k() method from BandAccessor. Add data property to TBSystem for atomic data access.

* feat: add eigenstates calculation and DOS analysis capabilities

- Add abstract get_eigenstates method to HamiltonianCalculator interface
- Implement eigenstates calculation in DeePTBAdapter using Eigh solver
- Add DosData class for structured DOS results storage
- Add DOS plotting functionality with matplotlib integration
- Refactor solver initialization to support both Eigenvalues and Eigh solvers
- Remove unused logging imports and update error handling

* refactor: improve DOS calculation interface and configuration

Rename set_dos to set_dos_config for clarity in DOS configuration setup. Add reuse parameter to get_dos and get_bands methods to control recalculation. Improve error messages and add dos_data property for easier access to computed DOS results. Refactor DOS initialization to separate kpoints setup from configuration.

* refactor(test): replace mock tests with integration tests using real data

Replace mock-based unit tests with integration tests that use actual silicon example data. Add pytest fixtures for efficient system initialization, test band structure and DOS calculations with real model, and include plotting functionality tests. This provides more realistic validation of the unified postprocessing module.

* feat: add unified post-processing tutorial notebook

Add a comprehensive Jupyter notebook demonstrating the new TBSystem class for unified post-processing in DeePTB. The tutorial covers initialization, band structure calculation, DOS computation, and visualization using the centralized interface that manages both atomic structure and Hamiltonian model.

* feat: ignore band structure and DOS plot files in .gitignore

* fix: ensure dtype consistency in tensor product operations

Fix dtype mismatches in tensor product calculations by explicitly setting dtype to match input tensors. Updates J_full_small initialization and _Jd tensor conversion in batch_wigner_D function, and enforces double precision in test cases to prevent numerical precision issues.

* feat: add band_data initialization and simplify tensor type check

- Initialize _band_data attribute in BandAccessor.__init__ for future use
- Remove NestedTensor type check in test_postprocess_unified.py to simplify validation
- Ensure consistent tensor type handling in Hamiltonian calculations

* feat: add PythTB export tutorial notebook

Add comprehensive tutorial demonstrating how to export DeePTB models to PythTB format. The notebook covers loading pre-trained models, calculating band structures with both DeePTB and PythTB, and comparing results to ensure compatibility. This enables users to leverage external tools for post-processing trained tight-binding models.

* feat: implement Fermi level calculation and integration

Add proper Fermi level calculation functionality to replace hardcoded values. The band structure and DOS properties now use the system's Fermi level when available, falling back to 0.0 when not set. Includes a new utility function for calculating Fermi energy from eigenvalues using various smearing methods and temperature parameters. Also improves DOS plotting with inward tick direction.

* feat: add export functionality for TBSystem to third-party formats

Add ExportAccessor class to

* feat: support dict input for export interfaces

Add support for dictionary input in ToWannier90 and ToPythTB export classes. Skip data loading when input is already a dictionary to avoid redundant processing. Update ExportAccessor to use explicit parameter names for better clarity.

* feat: add PythTB-Wannier postprocessing example and rename existing notebook

- Add new_postprocess_pythtb_wannier.ipynb demonstrating PythTB and Wannier90 integration
- Rename new_postprocess.ipynb to new_postprocess_band_dos.ipynb for clarity
- New example includes band structure calculation with k-path configuration and Fermi level determination

* feat: add Wannier90 integration and export functionality

Add new cells to demonstrate Wannier90 integration with PythTB, including export functionality and model loading. Update execution counts and fix reference from tbsys.band_data to tbsys.band for consistency.

* test: update export tests to match implementation changes

Updated test assertions in test_export_unified.py to align with the actual implementation of export methods. Fixed parameter names and added missing arguments like overlap, e_fermi, and filename parameters to ensure tests accurately validate the export functionality.

* refactor: simplify HR2HK onsite block construction and SOC handling

Simplify the HR2HK class by removing redundant overlap-specific code paths and consolidating onsite block construction. The SOC handling is now restricted to Hamiltonian-only cases with clearer documentation about current limitations. This refactoring improves code readability and maintainability while preserving the same functionality.

* feat(nn): add gauge convention support to HR2HK module

Add support for two gauge conventions in HR2HK:
- Wannier90 Gauge (gauge=False) using cell shift vectors
- Physical/Periodic Gauge (gauge=True) using edge vectors

The gauge parameter controls phase factor calculation in k-space
transformation. When derivative=True, automatically enables gauge=True
mode. Added comprehensive test suite for both conventions.

* add new key for AtomicDataDict: HAMILTONIAN_DERIV_KEY = "hamiltonian_derivative" # dH(k)/dk
OVERLAP_DERIV_KEY = "overlap_derivative"  # dS(k)/dk

* feat(nn): add Hamiltonian k-derivative computation to HR2HK

Add support for computing dH/dk derivatives in the HR2HK module. The implementation calculates the k-space derivatives of the Hamiltonian matrix using the analytical formula dH/dk = -i * R * H_R * exp(-i k·R), where R is the real-space hopping vector. The derivatives are computed for all three Cartesian directions and stored in the output data dictionary. This feature is controlled by the derivative parameter and uses the new out_derivative_field parameter to specify the output key. The gauge parameter default is changed to False to maintain consistency with the Wannier90 convention.

* feat: add derivative support to get_hk method in DeePTBAdapter

Enhance the get_hk method to support Hamiltonian and overlap matrix derivatives by adding a new with_derivative parameter. This allows computing derivatives of H(k) and S(k) with respect to k-points, which is useful for band structure analysis and property calculations. The implementation uses the HR2HK module for both regular and derivative computations, maintaining consistency with the existing codebase structure.

* refactor: restructure derivative handling in DeePTBAdapter

Restructured the derivative handling logic in the DeePTBAdapter class to improve code clarity and maintainability. The changes include moving derivative-related operations to a dedicated section and using .get() for safer access to optional derivative keys. This refactoring maintains the same functionality while making the code flow more logical and reducing potential errors when overlap derivatives are not present.

* correct v0

* correct v0

* jit

* v1

* feat: remove commented derivative code in HR2HK

Clean up HR2HK module by removing commented-out derivative computation code that was no longer needed. The active derivative computation logic remains intact and functional.

* refactor: rename optical to optical_conductivity and update class name

Renamed optical.py to optical_conductivity.py for clarity and updated OpticalAccessor to ACAccessor. Modified TBSystem to use accond property for accessing optical conductivity functionality, improving code organization and naming consistency.

* feat: improve optical conductivity calculation with complex Lorentzian

Enhance optical conductivity computation by implementing complex Lorentzian broadening to capture both real and imaginary parts. Update accumulation methods (loop and JIT) to use complex form 1/(E - ω + iη) for physical accuracy. Add comprehensive documentation for parameters and return values. Remove redundant vectorized method for cleaner code structure.

* feat: add spin degeneracy parameter to optical conductivity

Add g_s parameter to control spin degeneracy factor in optical conductivity calculations. Replace hardcoded spin factor of 2.0 with configurable g_s parameter (default 2.0). Update calculation to use 2πg_s/volume factor. Change default method from 'vectorized' to 'loop' for consistency.

* test: add optical conductivity test for silicon model

Add comprehensive unit test for optical conductivity calculations using a silicon tight-binding model. The test validates Fermi level calculation, optical conductivity computation with both JIT and loop methods, and consistency between different broadening functions (Lorentzian and Gaussian). Includes reference value comparisons and physical property validation.

* fix: improve broadening validation and example notebook

- Fix inconsistent indentation in optical_conductivity.py ACAccessor
- Add explicit validation for broadening type with clear error message
- Update example notebook with non-orthogonal model demonstration
- Reset execution counts and add missing import warnings

* update(example): refresh plot output in optical_cond.ipynb

The notebook's execution counts and matplotlib plot outputs have been updated. This reflects re-running the code cells to generate the latest visualization, ensuring the saved notebook state is consistent with its code.

* test: relax tolerance for optical conductivity tests

Adjusted assertion tolerances from 1e-5 to 1e-4 in optical conductivity unit tests to accommodate numerical precision differences while maintaining test validity.

* feat: add overlap override support for DeePTB calculator

Enable override of overlap matrices in DeePTB models by adding override_overlap parameter to DeePTBAdapter and TBSystem. This allows users to provide custom overlap files even when the original model includes overlap support, improving flexibility for tight-binding calculations.</think>
Based on the provided git diff, generate a concise and descriptive commit message.

The commit message should:
1. Has a short title (50-72 characters)
2. The commit message should adhere to the conventional commit format
3. Describe what was changed and why
4. De clear and informative

'git --no-pager diff --staged --diff-filter=d' Output:

diff --git a/dptb/sktb/neighbors.py b/dptb/sktb/neighbors.py
index 5128b57..689ec19 1006 Slater-Koster tight binding model module."""
+
+import numpy as np
import torch
from torch import nn
from torch.autograd import grad
from prefetch_generator import prefetch
from dptb.nn.ops.e3gnn import (
@@ -263,6 +264,11 @@ class NearestNeighborsFinder(nn.Module):
         self.max_num_neighbors = max_num_neighbors
         self.reference_distance = reference_distance
         self.neighbor_indexing = neighbor_indexing
+        if use_efficient_format:
+            if self.max_num_neighbors is None:
+                self.max_num_neighbors = 200
+                log.info(f"use_efficient_format=True but max_num_neighbors is not specified, setting to default value 200.")
         self.use_efficient_format = use_efficient_format

     def forward(
@@ -525,6 +55, neighbor_list, atom_mask, positions, cell, offset_vecs = build_neighbor_list(
             positions=positions,
             cell=missing_masked_cell,
-            neighbor_cutoff=neighbor_cutoff,
+            neighbor_cutoff=neighbor_cutoff,</think>
<think></think>
feat: add efficient neighbor format with default max neighbors

Introduce use_efficient_format option to NearestNeighborsFinder for optimized neighbor list handling. Automatically sets max_num_neighbors to 200 when not specified to improve performance with efficient format processing.

* refactor: rename TBSystems example notebooks for clarity

- Rename new_postprocess_band_dos.ipynb to tbsys_band_dos.ipynb
- Rename new_postprocess_pythtb_wannier.ipynb to tbsys_to_pythtb_wannier.ipynb
- Update kpath configuration and output filename in band DOS example

Author: QG-phy

.gitignore

@@ -1,5 +1,7 @@
 debug_**
 test*.ipynb
+*_bands.png
+*_dos.png
 examples/**/*centres.xyz
 examples/**/*.win
 **/processed*/*

dptb/data/_keys.py

@@ -32,8 +32,11 @@
 KPOINT_KEY = "kpoint"
 
 HAMILTONIAN_KEY = "hamiltonian"
+HAMILTONIAN_DERIV_KEY = "hamiltonian_derivative" # dH(k)/dk
 
 OVERLAP_KEY = "overlap"
+OVERLAP_DERIV_KEY = "overlap_derivative"  # dS(k)/dk
+
 # [n_batch, 3] bool tensor
 PBC_KEY: Final[str] = "pbc"
 # [n_atom, 1] long tensor

dptb/nn/build.py

@@ -5,7 +5,7 @@
 import torch
 from dptb.utils.tools import j_must_have, j_loader
 import copy
-
+import os
 log = logging.getLogger(__name__)
 
 def build_model(
@@ -43,6 +43,10 @@ def build_model(
 
     # load the model_options and common_options from checkpoint if not provided
     if not from_scratch:
+        if checkpoint in ['poly2', 'poly4']:
+            modelname = f'base_{checkpoint}.pth'
+            checkpoint = os.path.join(os.path.dirname(__file__), 'dftb', modelname)
+            
         if checkpoint.split(".")[-1] == "json":
             ckptconfig = j_loader(checkpoint)
         else:

dptb/nn/hr2hk.py

@@ -8,6 +8,8 @@
 
 
 class HR2HK(torch.nn.Module):
+    # this is actually a general FFT from real space hamiltonian/overlap to kspace hamiltonian/overlap
+    # the more correct name should be HSR2HSK. But to keep consistent with previous naming convention, we still use HR2HK here.
     def __init__(
             self, 
             basis: Dict[str, Union[str, list]]=None,
@@ -18,9 +20,18 @@ def __init__(
             overlap: bool = False,
             dtype: Union[str, torch.dtype] = torch.float32, 
             device: Union[str, torch.device] = torch.device("cpu"),
+            derivative:bool = False,
+            out_derivative_field: str = AtomicDataDict.HAMILTONIAN_DERIV_KEY,
+            gauge: bool = False 
             ):
+        # gauge: False -> Tight-binding Convention I:  Wannier90 Gauge 
+        # gauge: True  -> Tight-binding Convention II: "Physical Gauge"/"Periodic Gauge"
         super(HR2HK, self).__init__()
-
+    
+        if derivative:
+            gauge = True
+        self.gauge = gauge
+        self.derivative = derivative
         if isinstance(dtype, str):
             dtype = getattr(torch, dtype)
         self.dtype = dtype
@@ -44,15 +55,17 @@ def __init__(
         self.edge_field = edge_field
         self.node_field = node_field
         self.out_field = out_field
-
-
-
+        self.out_derivative_field = out_derivative_field
 
     def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
 
         # construct bond wise hamiltonian block from obital pair wise node/edge features
         # we assume the edge feature have the similar format as the node feature, which is reduced from orbitals index oj-oi with j>i
 
+        # Ensure edge_vectors are computed if using gauge mode
+        if self.gauge:
+            data = AtomicDataDict.with_edge_vectors(data, with_lengths=True)
+        
         orbpair_hopping = data[self.edge_field]
         orbpair_onsite = data.get(self.node_field)
         bondwise_hopping = torch.zeros((len(orbpair_hopping), self.idp.full_basis_norb, self.idp.full_basis_norb), dtype=self.dtype, device=self.device)
@@ -67,15 +80,12 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         soc = data.get(AtomicDataDict.NODE_SOC_SWITCH_KEY, False)
         if isinstance(soc, torch.Tensor):
             soc = soc.all()
-        if soc:
-            # if self.overlap:
-                # print("Overlap for SOC is realized by kronecker product.")
-
+        if soc: 
+            # this soc only support sktb.
             orbpair_soc = data[AtomicDataDict.NODE_SOC_KEY]
             soc_upup_block = torch.zeros((len(data[AtomicDataDict.ATOM_TYPE_KEY]), self.idp.full_basis_norb, self.idp.full_basis_norb), dtype=self.ctype, device=self.device)
             soc_updn_block = torch.zeros((len(data[AtomicDataDict.ATOM_TYPE_KEY]), self.idp.full_basis_norb, self.idp.full_basis_norb), dtype=self.ctype, device=self.device)
 
-
         ist = 0
         for i,iorb in enumerate(self.idp.full_basis):
             jst = 0
@@ -92,45 +102,53 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
 
                 if i <= j:
                     bondwise_hopping[:,ist:ist+2*li+1,jst:jst+2*lj+1] = factor * orbpair_hopping[:,self.idp.orbpair_maps[orbpair]].reshape(-1, 2*li+1, 2*lj+1)
+                    onsite_block[:,ist:ist+2*li+1,jst:jst+2*lj+1] = factor * orbpair_onsite[:,self.idp.orbpair_maps[orbpair]].reshape(-1, 2*li+1, 2*lj+1)
 
-
-                # constructing onsite blocks
-                if self.overlap:
-                    # if iorb == jorb:
-                    #     onsite_block[:, ist:ist+2*li+1, jst:jst+2*lj+1] = factor * torch.eye(2*li+1, dtype=self.dtype, device=self.device).reshape(1, 2*li+1, 2*lj+1).repeat(onsite_block.shape[0], 1, 1)
-                    if i <= j:
-                        onsite_block[:,ist:ist+2*li+1,jst:jst+2*lj+1] = factor * orbpair_onsite[:,self.idp.orbpair_maps[orbpair]].reshape(-1, 2*li+1, 2*lj+1)
-
-                    if soc and i == j:
-                        soc_updn_tmp = orbpair_soc[:, self.idp.orbpair_soc_maps[orbpair]].reshape(-1, 2*li+1, 2*(2*lj+1))
-                        # j==i -> 2*lj+1 == 2*li+1
-                        soc_upup_block[:, ist:ist+2*li+1, jst:jst+2*lj+1] = soc_updn_tmp[:, :2*li+1, :2*lj+1]
-                        soc_updn_block[:, ist:ist+2*li+1, jst:jst+2*lj+1] = soc_updn_tmp[:, :2*li+1, 2*lj+1:]
-                else:
-                    if i <= j:
-                        onsite_block[:,ist:ist+2*li+1,jst:jst+2*lj+1] = factor * orbpair_onsite[:,self.idp.orbpair_maps[orbpair]].reshape(-1, 2*li+1, 2*lj+1)
-                    
-                    if soc and i==j:
+                if soc and i==j and not self.overlap:
+                        # For now, The SOC part is only added to Hamiltonian, not overlap matrix.
+                        # For now, The SOC only has onsite contribution.
                         soc_updn_tmp = orbpair_soc[:,self.idp.orbpair_soc_maps[orbpair]].reshape(-1, 2*li+1, 2*(2*lj+1))
                         soc_upup_block[:,ist:ist+2*li+1,jst:jst+2*lj+1] = soc_updn_tmp[:, :2*li+1,:2*lj+1]
                         soc_updn_block[:,ist:ist+2*li+1,jst:jst+2*lj+1] = soc_updn_tmp[:, :2*li+1,2*lj+1:]
+
+                # constructing onsite blocks
+                #if self.overlap:
+                #    # if iorb == jorb:
+                #    #     onsite_block[:, ist:ist+2*li+1, jst:jst+2*lj+1] = factor * torch.eye(2*li+1, dtype=self.dtype, device=self.device).reshape(1, 2*li+1, 2*lj+1).repeat(onsite_block.shape[0], 1, 1)
+                #    if i <= j:
+                #        onsite_block[:,ist:ist+2*li+1,jst:jst+2*lj+1] = factor * orbpair_onsite[:,self.idp.orbpair_maps[orbpair]].reshape(-1, 2*li+1, 2*lj+1)
+                #    if soc and i == j:
+                #        soc_updn_tmp = orbpair_soc[:, self.idp.orbpair_soc_maps[orbpair]].reshape(-1, 2*li+1, 2*(2*lj+1))
+                #        # j==i -> 2*lj+1 == 2*li+1
+                #        soc_upup_block[:, ist:ist+2*li+1, jst:jst+2*lj+1] = soc_updn_tmp[:, :2*li+1, :2*lj+1]
+                #        soc_updn_block[:, ist:ist+2*li+1, jst:jst+2*lj+1] = soc_updn_tmp[:, :2*li+1, 2*lj+1:]
+                #else:
+                #    if i <= j:
+                #        onsite_block[:,ist:ist+2*li+1,jst:jst+2*lj+1] = factor * orbpair_onsite[:,self.idp.orbpair_maps[orbpair]].reshape(-1, 2*li+1, 2*lj+1)
+                #    
+                #    if soc and i==j:
+                #        soc_updn_tmp = orbpair_soc[:,self.idp.orbpair_soc_maps[orbpair]].reshape(-1, 2*li+1, 2*(2*lj+1))
+                #        soc_upup_block[:,ist:ist+2*li+1,jst:jst+2*lj+1] = soc_updn_tmp[:, :2*li+1,:2*lj+1]
+                #        soc_updn_block[:,ist:ist+2*li+1,jst:jst+2*lj+1] = soc_updn_tmp[:, :2*li+1,2*lj+1:]
 
                 jst += 2*lj+1
             ist += 2*li+1
         self.onsite_block = onsite_block
         self.bondwise_hopping = bondwise_hopping
-        if soc:
-            # 先保存已有的
+        if soc and not self.overlap:
+            # store for later use
+            # for now, soc only contribute to Hamiltonain, thus for overlap not store soc parts.
             self.soc_upup_block = soc_upup_block
             self.soc_updn_block = soc_updn_block
 
         # R2K procedure can be done for all kpoint at once.
         all_norb = self.idp.atom_norb[data[AtomicDataDict.ATOM_TYPE_KEY]].sum()
         block = torch.zeros(kpoints.shape[0], all_norb, all_norb, dtype=self.ctype, device=self.device)
-        # block = torch.complex(block, torch.zeros_like(block))
-        # if data[AtomicDataDict.NODE_SOC_SWITCH_KEY].all():
-        #     block_uu = torch.zeros(data[AtomicDataDict.KPOINT_KEY].shape[0], all_norb, all_norb, dtype=self.ctype, device=self.device)
-        #     block_ud = torch.zeros(data[AtomicDataDict.KPOINT_KEY].shape[0], all_norb, all_norb, dtype=self.ctype, device=self.device)
+        
+        # Initialize derivative blocks if needed: dH/dk = [dH/dkx, dH/dky, dH/dkz]
+        if self.derivative:
+            dblock = torch.zeros(kpoints.shape[0], all_norb, all_norb, 3, dtype=self.ctype, device=self.device)
+        
         atom_id_to_indices = {}
         ist = 0
         for i, oblock in enumerate(onsite_block):
@@ -139,21 +157,7 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
             block[:,ist:ist+masked_oblock.shape[0],ist:ist+masked_oblock.shape[1]] = masked_oblock.squeeze(0)
             atom_id_to_indices[i] = slice(ist, ist+masked_oblock.shape[0])
             ist += masked_oblock.shape[0]
-        
-        # if data[AtomicDataDict.NODE_SOC_SWITCH_KEY].all():
-        #     ist = 0
-        #     for i, soc_block in enumerate(soc_upup_block):
-        #         mask = self.idp.mask_to_basis[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()[i]]
-        #         masked_soc_block = soc_block[mask][:,mask]
-        #         block_uu[:,ist:ist+masked_soc_block.shape[0],ist:ist+masked_soc_block.shape[1]] = masked_soc_block.squeeze(0)
-        #         ist += masked_soc_block.shape[0]
-        #     ist = 0
-        #     for i, soc_block in enumerate(soc_updn_block):
-        #         mask = self.idp.mask_to_basis[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()[i]]
-        #         masked_soc_block = soc_block[mask][:,mask]
-        #         block_ud[:,ist:ist+masked_soc_block.shape[0],ist:ist+masked_soc_block.shape[1]] = masked_soc_block.squeeze(0)
-        #         ist += masked_soc_block.shape[0]
-
+    
         for i, hblock in enumerate(bondwise_hopping):
             iatom = data[AtomicDataDict.EDGE_INDEX_KEY][0][i]
             jatom = data[AtomicDataDict.EDGE_INDEX_KEY][1][i]
@@ -163,12 +167,37 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
             jmask = self.idp.mask_to_basis[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()[jatom]]
             masked_hblock = hblock[imask][:,jmask]
 
-            block[:,iatom_indices,jatom_indices] += masked_hblock.squeeze(0).type_as(block) * \
-                torch.exp(-1j * 2 * torch.pi * (kpoints @ data[AtomicDataDict.EDGE_CELL_SHIFT_KEY][i])).reshape(-1,1,1)
+            if self.gauge:
+                # phase factor according to convention II
+                # k and R are in fractional coordinates, need to convert to cartesian
+                edge_vec = data[AtomicDataDict.EDGE_VECTORS_KEY][i]  # Cartesian coordinates
+                phase_factor = torch.exp(-1j * 2 * torch.pi * (
+                    kpoints @ data[AtomicDataDict.CELL_KEY].inverse().T @ edge_vec)).reshape(-1,1,1)
+                # Compute derivative: dH/dk_alpha = -i * R_alpha * H_R * exp(-i k·R)
+                # where R is edge_vec in Cartesian coordinates
+                if self.derivative:
+                    # derivative_factor shape: [n_kpoints, 1, 1, 3]
+                    # - i * R * exp(-i k·R) = -i * R * phase_factor
+                    derivative_factor = (-1.0j * edge_vec).reshape(1, 1, 1, 3) * phase_factor.unsqueeze(-1)
+            else:
+                phase_factor = torch.exp(-1j * 2 * torch.pi * (
+                    kpoints @ data[AtomicDataDict.EDGE_CELL_SHIFT_KEY][i])).reshape(-1,1,1)
+                
+            block[:,iatom_indices,jatom_indices] += masked_hblock.squeeze(0).type_as(block) * phase_factor
+            
+            if self.derivative and self.gauge:
+                # Add derivative contribution
+                dblock[:,iatom_indices,jatom_indices,:] += masked_hblock.squeeze(0).type_as(dblock).unsqueeze(-1) * derivative_factor
 
         block = block + block.transpose(1,2).conj()
         block = block.contiguous()
 
+        # Hermitianize derivative blocks: dH/dk should also be Hermitian
+        if self.derivative:
+            for alpha in range(3):
+                dblock[:,:,:,alpha] = dblock[:,:,:,alpha] + dblock[:,:,:,alpha].transpose(1,2).conj()
+            dblock = dblock.contiguous()
+        
         if soc:
             if self.overlap:
                 # ========== S_soc = S ⊗ I₂ : N×N S(k) to 2N×2N kronecker product ==========
@@ -182,10 +211,6 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
                 data[self.out_field] = S_soc
             else:
                 HK_SOC = torch.zeros(kpoints.shape[0], 2*all_norb, 2*all_norb, dtype=self.ctype, device=self.device)
-                #HK_SOC[:,:all_norb,:all_norb] = block + block_uu
-                #HK_SOC[:,:all_norb,all_norb:] = block_ud
-                #HK_SOC[:,all_norb:,:all_norb] = block_ud.conj()
-                #HK_SOC[:,all_norb:,all_norb:] = block + block_uu.conj()
                 ist = 0
                 assert len(soc_upup_block) == len(soc_updn_block)
                 for i in range(len(soc_upup_block)):
@@ -207,6 +232,10 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
                 data[self.out_field] = HK_SOC
         else:
             data[self.out_field] = block
+        
+        # Store derivative if computed
+        if self.derivative:
+            data[self.out_derivative_field] = dblock
 
         return data
 

dptb/nn/tensor_product.py

@@ -74,10 +74,10 @@ def batch_wigner_D(l_max, alpha, beta, gamma, _Jd):
     D_total = sum(dims)
 
     # Construct block-diagonal J matrix
-    J_full_small = torch.zeros(D_total, D_total, device=device)
+    J_full_small = torch.zeros(D_total, D_total, device=device, dtype=alpha.dtype)
     for l in range(l_max + 1):
         start = offsets[l]
-        J_full_small[start:start+2*l+1, start:start+2*l+1] = _Jd[l]
+        J_full_small[start:start+2*l+1, start:start+2*l+1] = _Jd[l].to(dtype=alpha.dtype)
 
     J_full = J_full_small.unsqueeze(0).expand(N, -1, -1)
     angle_stack = torch.cat([alpha, beta, gamma], dim=0) 

dptb/postprocess/common.py

@@ -6,9 +6,10 @@
 from typing import Union, Optional
 from copy import deepcopy
 from ase.io import read
-
+import sys
 from dptb.data import AtomicData, AtomicDataDict, block_to_feature
 from dptb.utils.argcheck import get_cutoffs_from_model_options
+import matplotlib.pyplot as plt
 
 log = logging.getLogger(__name__)
 
@@ -110,3 +111,61 @@ def load_data_for_model(
     # Actually, ElecStruCal.get_data does NOT run self.model(data). It runs self.model.idp(data).
     # self.get_eigs runs self.model(data).
     return data_obj
+
+def is_gui_available():
+    """
+    Detect if GUI display is available for matplotlib.
+    
+    Returns:
+        bool: True if GUI is available, False otherwise
+    """
+    try:
+        # Check if we're in a Jupyter notebook environment
+        if 'ipykernel' in sys.modules or 'IPython' in sys.modules:
+            # In Jupyter, we can typically show plots
+            return True
+        
+        # Check DISPLAY environment variable (Unix-like systems)
+        if sys.platform.startswith('linux') or sys.platform.startswith('darwin'):
+            display = os.environ.get('DISPLAY')
+            if display is None:
+                return False
+        
+        # Try to get the current matplotlib backend
+        backend = plt.get_backend().lower()
+        
+        # Non-interactive backends
+        non_gui_backends = ['agg', 'pdf', 'ps', 'svg', 'cairo', 'gdk', 'template']
+        if any(non_gui in backend for non_gui in non_gui_backends):
+            return False
+        
+        # Try to create a test figure to see if it works
+        # This is a more robust check
+        try:
+            import matplotlib
+            # Save current backend
+            current_backend = matplotlib.get_backend()
+            
+            # Try to use a GUI backend if not already
+            if 'agg' in backend.lower():
+                # Try common GUI backends
+                for test_backend in ['TkAgg', 'Qt5Agg', 'Qt4Agg', 'WXAgg']:
+                    try:
+                        matplotlib.use(test_backend, force=True)
+                        test_fig = plt.figure()
+                        plt.close(test_fig)
+                        matplotlib.use(current_backend, force=True)
+                        return True
+                    except:
+                        continue
+                return False
+            else:
+                # Current backend seems to be GUI-based
+                return True
+                
+        except Exception:
+            return False
+            
+    except Exception:
+        # If any error occurs, assume no GUI is available
+        return False