Post-processing: bands, DOS, surface states, Fermi arcs

Calculator classes

Each calculator accepts either an HDF5 path (str) or an in-memory tbmodels.Model; the .run() method returns a typed Result dataclass with raw NumPy arrays AND matplotlib Figures.

class tailwater.wannier_wizard.BulkDOS(model_or_path: str | Model, k_mesh: Tuple[int, int, int] = (4, 4, 4), energies: Tuple[float, float] = (-5.0, 5.0), NC: int = 2048, NV: int = 4, N_tilde: int | None = None, device: str = 'cpu', verbose: bool = True)[source]

Bases: object

k-mesh-averaged density of states via KPM.

For each k in a uniform Monkhorst-Pack-like mesh, builds the periodic-cell Hamiltonian H(k), runs KPM with NV random phase vectors, accumulates Chebyshev moments, then reconstructs the DOS on the requested energy grid.

Construction

BulkDOS(model_or_path, k_mesh=(4,4,4), energies=(-5, 5),

NC=2048, NV=4, device=’cpu’)

Run

result = calc.run() result.energies, result.dos, result.figure

param model_or_path:

type model_or_path:

Union[str, Model]

param k_mesh:

type k_mesh:

Tuple[int, int, int]

param energies:

type energies:

Tuple[float, float]

param NC:

type NC:

int

param NV:

type NV:

int

param N_tilde:

type N_tilde:

Optional[int]

param device:

type device:

str

param verbose:

type verbose:

bool

run()[source]
Return type:

BulkDOSResult

class tailwater.wannier_wizard.SurfaceSpectralDensity(model_or_path: str | Model, surface: ndarray, LZ: int, energies: Tuple[float, float], k_path: Sequence[Sequence[float]], k_labels: List[str] | None = None, N_path: int = 101, NC: int = 4096, NV: int = 4, device: str = 'cpu', verbose: bool = True, rng_seed: int | None = None)[source]

Bases: object

Surface DOS along a k-path via KPM, for top and bottom surfaces.

Construction

SurfaceSpectralDensity(

model_or_path, surface=np.eye(3), LZ=5, energies=(E_F-1, E_F+1), k_path=[[0,0,0], [0,0,0.5], …], k_labels=None, N_path=101, NC=2**12, NV=4, device=’cpu’,

)

surface is a 3x3 integer matrix passed to _reorient_model. Use np.eye(3) to leave the unit cell unchanged. LZ is the slab thickness in unit-cell layers along the (post-reorient) z-axis.

param model_or_path:

type model_or_path:

Union[str, Model]

param surface:

type surface:

ndarray

param LZ:

type LZ:

int

param energies:

type energies:

Tuple[float, float]

param k_path:

type k_path:

Sequence[Sequence[float]]

param k_labels:

type k_labels:

Optional[List[str]]

param N_path:

type N_path:

int

param NC:

type NC:

int

param NV:

type NV:

int

param device:

type device:

str

param verbose:

type verbose:

bool

param rng_seed:

type rng_seed:

Optional[int]

run()[source]
Return type:

SurfaceSpectralDensityResult

class tailwater.wannier_wizard.SurfaceGreensFunction(model_or_path: str | Model, surface: ndarray, energies: Sequence[float], k_path: Sequence[Sequence[float]], N_path: int = 101, k_labels: List[str] | None = None, thickness: int = 6, NN: int = 5, eps: float = 0.005, delta: float = 0.0, device: str = 'cpu', verbose: bool = True, chunk_size: int = 256, n_jobs: int = 1)[source]

Bases: object

Surface Green’s function along a k-path via Lopez-Sancho recursion.

Construction

SurfaceGreensFunction(

model_or_path, surface=np.eye(3), energies=np.linspace(-1, 1, 101), k_path=[[0,0,0], [0,0,0.5], …], N_path=101, k_labels=None, thickness=6, NN=5, eps=0.005, delta=0.0, device=’cpu’, chunk_size=256, n_jobs=1,

)

Performance

The Lopez-Sancho recursion is the bottleneck — roughly 95% of wall-clock time. Two knobs control how it’s parallelized:

  • chunk_size (default 256): for each k-point, batches the energy axis through _recursion_torch_batched() so all energies share a single Python dispatch per Lopez-Sancho step. Larger chunks ⇒ less overhead, more memory. The default keeps peak memory bounded even for dense energy grids on thick slabs.

  • n_jobs (default 1): when set to -1 or any integer > 1, fans the k-points out across worker processes via joblib. Each worker runs one BLAS thread to avoid oversubscription. On multi-core CPUs this is the biggest single win — k-points are fully independent, so speedup scales close to linearly with the number of physical cores until pickling overhead bites (typically beyond ~16 workers on this problem class).

Typical recipe: SurfaceGreensFunction(model, ..., n_jobs=-1).

param model_or_path:

type model_or_path:

Union[str, Model]

param surface:

type surface:

ndarray

param energies:

type energies:

Sequence[float]

param k_path:

type k_path:

Sequence[Sequence[float]]

param N_path:

type N_path:

int

param k_labels:

type k_labels:

Optional[List[str]]

param thickness:

type thickness:

int

param NN:

type NN:

int

param eps:

type eps:

float

param delta:

type delta:

float

param device:

type device:

str

param verbose:

type verbose:

bool

param chunk_size:

type chunk_size:

int

param n_jobs:

type n_jobs:

int

run()[source]
Return type:

SurfaceGreensFunctionResult

class tailwater.wannier_wizard.FermiArcMap(model_or_path: str | Model, surface: ndarray, energy: float, Nx: int, Ny: int, thickness: int = 6, NN: int = 5, eps: float = 0.005, delta: float = 0.0, device: str = 'cuda', verbose: bool = True, chunk_size: int = 128, n_jobs: int = 1)[source]

Bases: object

Surface spectral function at a SINGLE energy on a 2D k-grid.

Same Lopez-Sancho machinery as SurfaceGreensFunction, but the k-grid is the 2D BZ slice at k_z = 0 (post-reorient), spanning [-0.5, 0.5] in both k_x and k_y. Produces four matplotlib figures: raw and griddata-interpolated maps for both surfaces.

Parameters:

See Fermi arcs and 2D surface spectral maps for a worked Bi2Se3 example, surface-basis guidance, and tips for choosing the energy and grid resolution.

run()[source]
Return type:

FermiArcMapResult

class tailwater.wannier_wizard.BandStructure(model_or_path: str | Model, k_points: Sequence[Sequence[float]], k_labels: List[str] | None = None, spacing: float = 0.01, fermi_level: float = 0.0, e_range: Tuple[float, float] | None = None, verbose: bool = True, linewidth: float = 1.6)[source]

Bases: object

Bulk band structure along a user-supplied or auto-generated k-path.

Construction

BandStructure(model_or_path,

k_points=[[0,0,0], [0,0,0.5], [0.5,0.5,0]], k_labels=[r”$Gamma$”, “Z”, “M”], spacing=0.01)

Or with seekpath-derived high-symmetry path:

BandStructure.auto(model_or_path, structure, spacing=0.01)

The spacing argument is the maximum allowed step in fractional reciprocal coordinates between adjacent k-points along the path (pass a smaller value for a denser path). The total number of samples is ceil(path_length / spacing).

Run

result = bs.run() result.eigenvalues, result.k_dist, result.figure

param model_or_path:

type model_or_path:

Union[str, Model]

param k_points:

type k_points:

Sequence[Sequence[float]]

param k_labels:

type k_labels:

Optional[List[str]]

param spacing:

type spacing:

float

param fermi_level:

type fermi_level:

float

param e_range:

type e_range:

Optional[Tuple[float, float]]

param verbose:

type verbose:

bool

param linewidth:

type linewidth:

float

classmethod auto(model_or_path: str | Model, structure, spacing: float = 0.01, fermi_level: float = 0.0, e_range: Tuple[float, float] | None = None, with_time_reversal: bool = True, verbose: bool = True, linewidth: float = 1.6)[source]

Build a BandStructure whose k-path is determined by seekpath.

structure is forwarded to _seekpath_auto_path — either a pymatgen.Structure or the (lattice, frac_coords, atomic_numbers) tuple seekpath accepts directly.

Parameters:
Return type:

BandStructure

run()[source]
Return type:

BandStructureResult

Convenience function

tailwater.wannier_wizard.bulk_band_structure(model_or_path: str | Model, k_points: Sequence[Sequence[float]] | None = None, k_labels: List[str] | None = None, spacing: float = 0.01, *, auto: bool = False, structure=None, fermi_level: float = 0.0, e_range: Tuple[float, float] | None = None, return_raw: bool = False, verbose: bool = True, with_time_reversal: bool = True, linewidth: float = 1.6)[source]

Compute the bulk band structure of a tight-binding model along a k-path.

Two modes:

  • MANUAL (default): pass k_points (and optionally k_labels).
    bulk_band_structure(model, k_points=[[0,0,0], [0,0,0.5]],

    k_labels=[r”$Gamma$”, “Z”])

  • AUTO (auto=True): the high-symmetry path is determined by seekpath from a structure argument (pymatgen.Structure or a seekpath-format tuple).

    bulk_band_structure(model, auto=True, structure=mp_structure)

Parameters:

  • model_or_path (str or tbmodels.Model) – Path to an HDF5 file produced by the API, or an in-memory model.

  • k_points (sequence of [k1, k2, k3] in fractional reciprocal coords) – Required when auto=False.

  • k_labels (list of str, optional) – Label per node, e.g. [r”$Gamma$”, “M”, “K”, r”$Gamma$”].

  • spacing (float) – Maximum step (in fractional reciprocal coords) between adjacent samples on the path. Smaller spacing = denser sampling.

  • auto (bool) – If True, use seekpath to determine path + labels (overrides k_points / k_labels). Requires the seekpath package and a structure argument.

  • structure (pymatgen.Structure or seekpath tuple, optional) – Required when auto=True.

  • fermi_level (float) – Subtracted from every eigenvalue before plotting.

  • e_range ((float, float), optional) – Y-axis limits for the figure.

  • return_raw (bool) – If True, return the full BandStructureResult (with raw arrays AND the figure). If False (default), return only the matplotlib Figure.

  • verbose (bool) – Toggle the tqdm progress bar.

  • with_time_reversal (bool) – Passed to seekpath in auto mode.

  • linewidth (float, default 1.6) – Band line width (and E_F reference line width). The default matches the WannierTools gnuplot bulkek convention (w lp lw 2) for a bolder, publication-style plot; lower it (e.g. 1.0) for dense band manifolds.

Returns:

  • matplotlib.figure.Figure if return_raw is False

  • BandStructureResult if return_raw is True

Result dataclasses

class tailwater.wannier_wizard.BulkDOSResult(energies: ndarray, dos: ndarray, figure: Figure)[source]

Bases: object

Output of BulkDOS.run().

Parameters:
energies: ndarray
dos: ndarray
figure: Figure
as_dict()[source]
Return type:

dict

class tailwater.wannier_wizard.SurfaceSpectralDensityResult(k_vec: ndarray, k_dist: ndarray, k_node: ndarray, energies: ndarray, dos_top: ndarray, dos_bottom: ndarray, figure_top: Figure, figure_bottom: Figure)[source]

Bases: object

Output of SurfaceSpectralDensity.run().

Parameters:
k_vec: ndarray
k_dist: ndarray
k_node: ndarray
energies: ndarray
dos_top: ndarray
dos_bottom: ndarray
figure_top: Figure
figure_bottom: Figure
as_dict()[source]
Return type:

dict

class tailwater.wannier_wizard.SurfaceGreensFunctionResult(k_vec: ndarray, k_dist: ndarray, k_node: ndarray, energies: ndarray, spectral_top: ndarray, spectral_bottom: ndarray, figure_top: Figure, figure_bottom: Figure)[source]

Bases: object

Output of SurfaceGreensFunction.run().

Parameters:
k_vec: ndarray
k_dist: ndarray
k_node: ndarray
energies: ndarray
spectral_top: ndarray
spectral_bottom: ndarray
figure_top: Figure
figure_bottom: Figure
as_dict()[source]
Return type:

dict

class tailwater.wannier_wizard.FermiArcMapResult(kx_grid: ndarray, ky_grid: ndarray, pos_x: ndarray, pos_y: ndarray, spectral_top: ndarray, spectral_bottom: ndarray, figure_top: Figure, figure_bottom: Figure, figure_top_interpolated: Figure, figure_bottom_interpolated: Figure)[source]

Bases: object

Output of FermiArcMap.run().

Parameters:

  • kx_grid (ndarray)

  • ky_grid (ndarray)

  • pos_x (ndarray)

  • pos_y (ndarray)

  • spectral_top (ndarray)

  • spectral_bottom (ndarray)

  • figure_top (Figure)

  • figure_bottom (Figure)

  • figure_top_interpolated (Figure)

  • figure_bottom_interpolated (Figure)

kx_grid: ndarray
ky_grid: ndarray
pos_x: ndarray
pos_y: ndarray
spectral_top: ndarray
spectral_bottom: ndarray
figure_top: Figure
figure_bottom: Figure
figure_top_interpolated: Figure
figure_bottom_interpolated: Figure
as_dict()[source]
Return type:

dict

class tailwater.wannier_wizard.BandStructureResult(k_vec: ndarray, k_dist: ndarray, k_node: ndarray, k_labels: List[str], eigenvalues: ndarray, figure: Figure)[source]

Bases: object

Output of BandStructure.run() / bulk_band_structure().

Parameters:
k_vec: ndarray
k_dist: ndarray
k_node: ndarray
k_labels: List[str]
eigenvalues: ndarray
figure: Figure
as_dict()[source]
Return type:

dict

k-path helper

tailwater.wannier_wizard.generate_k_path(k_points: Sequence[Sequence[float]], N_path: int, labels: List[str] | None = None, rec_vecs: ndarray | None = None)[source]

Generate a k-path connecting k_points with about N_path total samples.

Parameters:

  • k_points (sequence of fractional k-vectors) – High-symmetry nodes the path visits, in order. Path segments between consecutive nodes are sampled in proportion to their Cartesian length when rec_vecs is supplied, else uniformly.

  • N_path (int) – Approximate total number of samples across all segments.

  • labels (list of str, optional) – Display labels for the nodes (e.g. [r"$\Gamma$", "K", "M"]).

  • rec_vecs (ndarray of shape (3, 3), optional) – Reciprocal lattice vectors. If provided, segment lengths are measured in Cartesian (1/Å) space — otherwise they default to fractional length, which can over/undersample anisotropic cells.

Return type:

Tuple[ndarray, ndarray, ndarray]

Returns:

  • k_vec (ndarray of shape (N, 3)) – Sampled k-points along the path.

  • k_dist (ndarray of shape (N,)) – Cumulative path length at each sample (for plotting the x-axis of a band-structure figure).

  • k_node (ndarray of shape (len(k_points),)) – Cumulative path length at each high-symmetry node (for x-tick positions on a band-structure figure).

Fermi / band-edge alignment

For semiconductors and insulators, anchor the energy zero to the valence band maximum so all calculators downstream share a physically natural reference — see Fermi alignment for semiconductors and insulators for the full guide.

tailwater.wannier_wizard.compute_band_edges(model_or_path: str | Model, k_mesh: Tuple[int, int, int] = (4, 4, 4))[source]

Locate VBM / CBM / gap on a uniform Monkhorst-Pack k-mesh.

Assumes the model’s current zero of energy is roughly E_F (the Tailwater training convention). Diagonalizes H(k) on every k in a k_mesh[0] x k_mesh[1] x k_mesh[2] uniform grid in fractional reciprocal coordinates, then takes:

  • VBM = the negative eigenvalue closest to zero (max of e < 0)

  • CBM = the positive eigenvalue closest to zero (min of e > 0)

  • gap = CBM - VBM

  • is_metal = (gap <= 0) — i.e. bands overlap E=0 across the mesh

Parameters:

  • model_or_path (str | tbmodels.Model) – Path to the HDF5 hr-model the API produced, or a tbmodels.Model already in memory.

  • k_mesh ((int, int, int)) – Grid density. Default (4, 4, 4) — denser meshes catch the VBM/CBM at off-symmetry k more accurately at small extra cost.

Return type:

dict

Returns:

  • dict with keys {“vbm”, “cbm”, “gap”, “is_metal”} (floats, except

  • is_metal which is bool; vbm/cbm/gap may be None in degenerate cases

  • where the spectrum has no eigenvalues on one side of zero).

tailwater.wannier_wizard.align_to_vbm(model_or_path: str | Model, k_mesh: Tuple[int, int, int] = (4, 4, 4), fermi_level: float | None = None, if_metal: str = 'warn')[source]

Return a NEW model with its on-site energies shifted so VBM = 0.

This re-anchors the energy scale so the band-edge sits exactly at zero — the natural reference for plotting / DOS / surface-state computations on semiconductors and insulators, instead of whatever DFT-chosen E_F the training data was referenced against.

Parameters:

  • model_or_path (str | tbmodels.Model) – Path to the HDF5 hr-model, or a tbmodels.Model in memory.

  • k_mesh ((int, int, int)) – k-mesh used to auto-detect the VBM (only consulted if fermi_level is None). Default (4, 4, 4).

  • fermi_level (float, optional) – If supplied, bypass auto-detection and shift on-site energies by -fermi_level (i.e. put your chosen Fermi value at the new zero). Useful if you already know E_F from a DFT calculation.

  • if_metal ({"warn", "raise", "skip"}) – What to do when compute_band_edges reports the spectrum has no clean gap around E=0 (signature of a metal, or of a non-metal whose current zero isn’t in the gap). Default "warn": emits a RuntimeWarning and returns the unshifted model so downstream code still runs. "raise" errors out; "skip" silently returns the unshifted model.

Returns:

A deep copy of the input model with its (0,0,0) hop block adjusted by shift * I so every eigenvalue at every k is offset by shift = -VBM. The input model is not mutated.

Return type:

tbmodels.Model