difftb 2025.12.16

Differentiable Density Functional Tight Binding in JAX

DiffTB: Differentiable Density Functional Tight Binding

A JAX-based implementation of DFTB with automatic differentiation support, enabling gradient-based optimization of Slater-Koster parameters against experimental data.

Features

• Fully Differentiable: Built on JAX for automatic differentiation through all calculations
• Spin Polarization: Support for magnetic systems with collinear spin
• Spin-Orbit Coupling: On-site L·S coupling for heavy elements (Bi, Te, etc.)
• Periodic Systems: k-space Hamiltonian construction with Brillouin zone sampling
• Inverse Problems: Tools for fitting SK parameters to experimental ARPES data

Installation

pip install difftb

For development:

git clone https://github.com/debangshu-ornl/difftb.git
cd difftb
pip install -e ".[dev]"

Quick Start

import jax.numpy as jnp
from difftb import (
    CrystalStructure,
    SlaterKosterParams,
    compute_bands,
    BrillouinZone,
)
from difftb.io.structure import bi2te3_structure
from difftb.io.slater_koster import create_test_sk_params

# Load structure
structure = bi2te3_structure()

# Create SK parameters (use create_test_sk_params for testing)
sk_params = create_test_sk_params(['Bi', 'Te'])

# Define k-path
from difftb.core.lattice import Lattice, hexagonal_path

lattice = Lattice.from_vectors(
    structure.lattice_vectors[0],
    structure.lattice_vectors[1],
    structure.lattice_vectors[2]
)
bz = BrillouinZone.from_path(lattice, hexagonal_path(), num_points=100)

# Compute band structure
# (Full example requires Hamiltonian construction - see examples/)

Core Components

Slater-Koster Tables

from difftb.io.slater_koster import SlaterKosterTable, interpolate_sk_integrals

# Interpolation is differentiable!
H_sk, S_sk = interpolate_sk_integrals(distances, sk_table)

Spin-Orbit Coupling

from difftb.physics.soc import build_spinor_hamiltonian, get_soc_params

soc_params = get_soc_params(['Bi', 'Te'])  # Uses tabulated ξ values
H_soc, S_soc = build_spinor_hamiltonian(H, S, elements, soc_params, orbital_info)

Parameter Fitting

from difftb.utils.optimization import fit_parameters, hierarchical_fit

# Define loss function (e.g., MSE to experimental ARPES)
def loss_fn(params):
    # Compute predicted spectrum
    # Compare to experiment
    return mse

# Fit parameters
result = fit_parameters(loss_fn, initial_params, max_iterations=1000)

# Hierarchical fitting for transferability
final_params, results = hierarchical_fit(
    loss_fn,
    param_groups=[bi_te_params, mn_params],
    group_names=['Bi-Te', 'Mn']
)

Acknowledgments

DiffTB builds on ideas from:

• Jrystal - JAX-based DFT
• dxtb - Differentiable xTB
• DFTB+ - Reference DFTB implementation