# 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

```bash
pip install difftb
```

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

## Quick Start

```python
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
```python
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
```python
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
```python
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](https://github.com/sail-sg/jrystal) - JAX-based DFT
- [dxtb](https://github.com/grimme-lab/dxtb) - Differentiable xTB
- [DFTB+](https://dftbplus.org/) - Reference DFTB implementation
