tailwater

Client + post-processing toolkit for the Tailwater Wannier-Hamiltonian inference API.

Full documentation: tailwater.readthedocs.io — installation, quick start, exporting models (sparse .npz / HDF5 / _hr.dat / pybinding / PythTB / Kwant), performance guides, and the complete API reference.

tailwater lets you upload a crystal structure to the Tailwater API, receive a tight-binding Hamiltonian, optionally fine-tune the output heads on customer-side targets, and run band-structure / DOS / surface-state analyses locally — all from one pip-installable package.

Installation

pip install tailwater

Optional extras:

pip install "tailwater[scatter]"        # if torch_scatter import fails
pip install "tailwater[seekpath]"       # enables auto k-path mode of bulk_band_structure
pip install "tailwater[dev]"            # pytest, ruff, build, twine

The tight-binding-library converters (to_pb / to_pythtb / to_kwant, and tb_model.load(...).to_pb()) need their target package, imported lazily so the base install stays light:

pip install pybinding-dev                 # pybinding  (a separate package, not a tailwater extra)
pip install pythtb                        # PythTB
conda install -c conda-forge kwant        # Kwant

as_tbmodels / to_hr_dat / to_hdf5 need only tbmodels, which is a core dependency (installed automatically).

Tested on Python 3.9–3.12.

API access

The Tailwater inference API is hosted at https://api.tailwater.io — this is the default endpoint tw_api_call(...) talks to, so the basic usage below needs no extra configuration beyond your credentials.

• Credentials. Authentication is HTTP Basic (username + password). Email the Tailwater team to request an account; you'll be issued a username and a one-time-displayed password.

• Billing. Each successful inference call decrements your server-side credit balance by one. Health checks (/healthz) and balance lookups (/credits/) are free.

• Checking your balance:

  from tailwater import remaining_credits
  print(remaining_credits("user", "pw"))   # -> int
  

Three workflow layers

1. HTTP client — talk to the API

from pymatgen.core import Structure
from tailwater import tw_api_call

structure = Structure.from_file("MyMaterial.cif")

# tw_api_call ALWAYS returns a dict of extracted paths. Every response
# includes a "win" key — the canonical wannier90.win file the server
# actually ran inference on — alongside the mode-specific artifact(s).

# (a) default: tbmodels HDF5 hr-model + .win
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat")
# paths = {"hdf5": "...", "win": "..."}

# (b) backbone embeddings + .win
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat",
                    return_embeddings=True)
# paths = {"embeddings": "...", "win": "..."}

# (c) project bundle: all artifacts + .win in a single call
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat",
                    project=True)
# paths = {"hdf5": "...", "embeddings": "...", "npz": "...", "win": "..."}

# (d) sparse output — always keep the O(N) SparseHR .npz (large systems)
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat",
                    output_format="sparse")
# paths = {"npz": "...", "meta": "...", "win": "..."}

Five output modes are available — return_embeddings, return_input, return_graph_output, project, or default HDF5. Set them as keyword arguments to tw_api_call. See tailwater.tw_api_call.__doc__ for the priority order.

output_format ("auto" | "sparse" | "hdf5", default "auto") controls how the Hamiltonian is transported and delivered:

• "auto" — request the sparse wannier90_hr.npz (O(N) egress). Systems below 30 atoms are transparently converted back to dense HDF5 (r["hdf5"] still works, and r["npz"] is kept too); larger systems stay sparse under r["npz"], with a printed note on how to convert / analyse them.
• "sparse" — always keep the raw .npz (a SparseHR — see the Sparse Hamiltonians section below, or the exporting-models guide), whatever the size.
• "hdf5" — always deliver dense tbmodels HDF5 (the pre-0.9 behaviour).

A server that predates the sparse backend ignores the flag and returns HDF5, so every mode degrades cleanly. See Sparse Hamiltonians below for what to do with the .npz.

Each successful call decrements your server-side credit balance by one. Failures surface as PermissionError (401, bad password) or RuntimeError (402, out of credits / other 5xx).

2. Subspace projection — fine-tune the heads near the Fermi level

From the project=True bundle (embeddings.pt + wannier90_hr.npz) you can fine-tune the output heads to reproduce the Hamiltonian's eigenvalues inside a narrow energy window — a compact, downfolded model for that window. It runs entirely on the artifacts the API returned:

from tailwater import subspace_projection

subspace_projection(
    start_lr     = 1e-4,
    end_lr       = 1e-5,
    num_epochs   = 20,
    energy_range = (-2.0, 2.0),      # eV, relative to E_F
    decay_sigma  = 0.5,
    device       = "cpu",
    save_path    = "./projection_out",
    embed_path   = paths["embeddings"],
    hr_npz_path  = paths["npz"],     # the sparse Hamiltonian is the fit target
)

Per epoch the script prints the mean eigenvalue loss. When done, three files are written to save_path:

Bring your own targets. To fit your own band structure instead of the API's prediction, pass your own DFT eigenvalues via make_eigenvalue_only_data, or supply your own Hamiltonian as a SparseHR .npz through hr_npz_path.

3. Post-processing — bulk DOS, surface states, Fermi arcs

import numpy as np
from tailwater import (
    tb_model,
    BulkDOS,
    SurfaceSpectralDensity,
    SurfaceGreensFunction,
    FermiArcMap,
)

# Bulk DOS — reads the sparse .npz the API returns, in states/eV per cell
# (it integrates to num_wann; run() prints that sum rule). Exact diagonalisation
# below ~4000 orbitals, KPM above — method="auto" decides.
result = BulkDOS("outputs/wannier90_hr.npz", k_mesh=(8, 8, 8), energies=(-4, 4)).run()
result.figure.savefig("bulk_dos.png")

# The surface calculators below take a dense tbmodels model
model = tb_model.load("outputs/wannier90_hr.hdf5")
np.savez("bulk_dos.npz", **result.as_dict())

# Surface spectral density along a k-path (KPM)
result = SurfaceSpectralDensity(
    model, surface=np.eye(3), LZ=5,
    energies=(-1, 1),
    k_path=[[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0]],
    k_labels=["M", r"$\Gamma$", "K"],
    N_path=101, NC=2**12, NV=4,
).run()
result.figure_top.savefig("surface_top.png")
result.figure_bottom.savefig("surface_bottom.png")

# Surface Green's function (Lopez-Sancho)
result = SurfaceGreensFunction(
    model, surface=np.eye(3),
    energies=np.linspace(-1, 1, 201),
    k_path=[[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0]],
    k_labels=["M", r"$\Gamma$", "K"],
    N_path=101, thickness=6, NN=30, eps=0.005,
).run()
np.savez("surface_gf.npz", **result.as_dict())

# 2D Fermi-arc map at one energy
result = FermiArcMap(
    model, surface=np.eye(3), energy=0.0,
    Nx=50, Ny=50, thickness=6,
).run()
result.figure_top_interpolated.savefig("fermi_arc_top.png")

# Bulk band structure along a manual k-path
from tailwater import bulk_band_structure
fig = bulk_band_structure(
    model,
    k_points = [[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0], [0, 0, 0]],
    k_labels = ["M", r"$\Gamma$", "K", r"$\Gamma$"],
    spacing  = 0.01,
    fermi_level = 0.0,
    e_range  = (-3, 3),
)
fig.savefig("bands.png")

# Or use seekpath to auto-determine the high-symmetry path
from pymatgen.core import Structure
structure = Structure.from_file("MyMaterial.cif")
fig = bulk_band_structure(model, auto=True, structure=structure,
                          spacing=0.02, e_range=(-3, 3))
fig.savefig("bands_auto.png")

BulkDOS, BandStructure, compute_band_edges and align_to_vbm take a sparse .npz path or a SparseHR natively — no conversion, and none of the dense-model cost. They also accept an HDF5 path or an in-memory tbmodels.Model. So do SurfaceGreensFunction and FermiArcMap, which build their Lopez-Sancho layer blocks straight from the bulk H(R) and never construct a supercell at all. Only SurfaceSpectralDensity still needs a dense model, because its KPM recursion wants the whole slab Hamiltonian rather than two blocks of it. Every .run() returns a typed Result dataclass with raw NumPy arrays and matplotlib Figure objects.

4. Local DOS — which atom, which shell, where in the cell

tailwater.ldos resolves the density of states by atom, species, orbital shell, or any predicate you pass, and renders it in real space. Sparse-native; method="auto" uses exact eigenvector projections below 4000 orbitals and group-restricted KPM above.

from tailwater import local_dos, ldos_map, ldos_volume

# DOS per species — the curves sum to the total DOS, and each integrates
# to that group's own orbital count. `by` also takes "atom" or "shell",
# or pass select=<predicate> for an arbitrary group.
res = local_dos("wannier90_hr.npz", by="species", energies=(-6, 6), n_energies=400)
res.figure.savefig("ldos_species.png")

# Per-atom weight at one energy, for a structure viewer
m = ldos_map("wannier90_hr.npz", energy=-1.0)
m.to_xyz("ldos.xyz")            # extended XYZ, LDOS as a per-atom property (OVITO)

# Real-space field rho(r) on one unit cell, with periodic image summation
vol = ldos_volume("wannier90_hr.npz", energy=-1.0)
vol.to_cube("ldos.cube")        # OVITO, VESTA, VMD and py3Dmol read this directly

On MgSe this reads out the chemistry directly: at −1 eV the weight is Se-p and the field peaks on the Se site; at +3 eV it is Mg s′/p′ — the conduction band is carried by shells the 18-slot basis did not contain. See examples/11_local_dos.py.

The V4 enlarged basis (26 slots)

The served V4 model gives 16 elements (Na Mg K Ca Co Ni Cu Zn Rh Pd Ag Cd Ir Pt Au Hg) an enlarged 26-slot basis — the legacy 18 plus an s′/p′ shell — while everything else keeps 18. Two consequences for your code:

• Do not assume 18 orbitals per atom. Read active_projections from the returned meta.json; it is the server's authoritative per-atom slot listing, and restrict_orbitals(..., active_projections=...) takes it directly.
• Check the energy reference. meta.json states it: VBM = 0 for gapped systems, E_F = 0 for metals. It is not inferable from the Hamiltonian, and tailwater.fermi needs it (or an explicit electron count) to place the Fermi level.

Fine-tuning heads on a V4 model works through the same multi_finetune entry points, which take their basis width from the graph.

Sparse Hamiltonians (SparseHR) & format conversion

For large systems the API returns the Hamiltonian in sparse form — a wannier90_hr.npz (COO hops + on-site diagonal + geometry) that is O(N) in RAM/egress instead of O(N²). You get it from output_format="sparse" (always), or from "auto" for systems ≥ 30 atoms (see the client section above).

Load a .npz with SparseHR, then compute spectra directly — including for num_wann far larger than a dense H(k) could hold — or convert to any supported tight-binding format:

from tailwater import SparseHR

shr = SparseHR.load("outputs/wannier90_hr.npz")
shr.num_wann, shr.nnz                       # size / number of stored hops
Hk  = shr.Hk([0.0, 0.0, 0.0])               # scipy sparse H(k) at fractional k (Γ)
w   = shr.eigsh_near_fermi([0, 0, 0], e_fermi=0.0, num=20)  # 20 states nearest E_F (shift-invert)
Rd  = shr.hr_dict()                         # {R_tuple: scipy.sparse.csr_matrix}
ev  = shr.eigvals_grid([[0,0,0], [0.5,0,0]])# dense eigenvalues on a list of k-points

Convert to any format — one call, auto-detecting the input

The top-level converters accept either a sparse input (a SparseHR or a .npz path) or a dense one (a tbmodels.Model, an .hdf5, or a Wannier90 _hr.dat) and dispatch automatically — the same call works regardless of what you're holding:

from tailwater import as_tbmodels, to_hr_dat, to_hdf5, to_pb, to_pythtb, to_kwant

npz = "outputs/wannier90_hr.npz"
model     = as_tbmodels(npz)                 # tbmodels.Model (dense)
to_hr_dat(npz, "wannier90_hr.dat")           # Wannier90 _hr.dat
to_hdf5(npz,   "wannier90_hr.hdf5")          # tbmodels HDF5
pb_lattice = to_pb(npz)                       # pybinding.Lattice
py_model   = to_pythtb(npz)                    # pythtb model
syst, lat  = to_kwant(npz)                     # kwant (Builder, lattice)

# the identical calls work on a dense input, too:
to_hr_dat("wannier90_hr.hdf5", "wannier90_hr.dat")
pb_lattice = to_pb(model)

• pybinding & Kwant are built straight from the sparse COO (no dense matrix), so they scale to large num_wann. _hr.dat and HDF5 are dense on-disk formats (size ≈ num_R · num_wann²) and are guarded for very large systems — pass max_wann= to override the guard if you really mean it.
• pybinding / Kwant / PythTB are optional and imported lazily: pip install pybinding-dev, conda install -c conda-forge kwant, pip install pythtb.
• SparseHR also has to_tbmodels() / to_hdf5() / to_hr_dat() / to_pb() / to_kwant() methods if you already hold the object; the top-level functions are the format-agnostic entry points.

BulkDOS, BandStructure, compute_band_edges and align_to_vbm need no conversion at all — hand them the .npz or the SparseHR. Convert only for the surface calculators (SurfaceSpectralDensity, SurfaceGreensFunction, FermiArcMap), which build their slabs with tbmodels.

API reference (top-level imports)

# HTTP client + HDF5 loader
tw_api_call(structure, user, password, output_path, filename,
            output_format="auto", ...)      # "auto" | "sparse" | "hdf5"
tb_model.load(path_to_hdf5)
remaining_credits(user, password)

# Sparse Hamiltonian (from output_format="sparse") + format-detecting converters
SparseHR.load(path_to_npz)          # -> SparseHR: .Hk, .eigsh_near_fermi, .hr_dict, .eigvals_grid,
                                    #    .to_tbmodels/.to_hdf5/.to_hr_dat/.to_pb/.to_kwant
as_tbmodels(src)                    # src = SparseHR/.npz  OR  tbmodels.Model/.hdf5/_hr.dat
to_hr_dat(src, path)   to_hdf5(src, path)
to_pb(src)   to_pythtb(src)   to_kwant(src)

# Heads-only inference model
HeadsOnly(irreps_in)
CovariantOnsiteHead(irreps_in)
CovariantEdgeHead(irreps_in)
load_heads_only_checkpoint(path)
save_heads_only_checkpoint(full_state_dict, irreps_in_str, save_path)

# Subspace fine-tuning — fit the heads to the project bundle's Hamiltonian
subspace_projection(start_lr, end_lr, num_epochs, energy_range, decay_sigma,
                    device, save_path, embed_path, *, hr_npz_path=...)
                    # advanced: a dense graph_output_path=... target instead of
                    #   hr_npz_path enables the "subspace" / "full" loss modes

# Subspace losses (advanced)
Subspace_H_MSE_Loss(gdata, edge_pred, onsite_pred, e_lo, e_hi)
Subspace_EigLoss(gdata, edge_pred, onsite_pred, kvec, neighbrs, e_lo, e_hi)
Eigenvalue_Only_Loss(gdata, edge_pred, onsite_pred, e_lo, e_hi)
make_eigenvalue_only_data(gdata, kvecs, eigs_per_k, e_lo, e_hi)
build_subspace_active_mask(node_features, onsite_target, e_lo, e_hi)
write_subspace_basis_file(out_path, active_mask, atoms, LM, ...)

# tbmodels.Model assembly from raw head output
build_hr_model     (edge_pred, onsite_pred, gdata, LM, atoms)
build_hr_model_fast(edge_pred, onsite_pred, gdata, LM, atoms)   # vectorized
write_hr_output(hr_model, out_path, fmt="hdf5"|"hr_dat")

# Post-processing calculators (each has a .run() method returning a Result)
BulkDOS(model_or_path, k_mesh, energies,                     # -> states/eV per cell
        method="auto"|"exact"|"kpm", resolution=None,        # resolution in eV
        sigma=None, NC=2048, NV=4, exact_max_orbitals=4000, device="cpu")
BandStructure(model_or_path, k_points, k_labels=None, spacing=0.01, ...)
SurfaceSpectralDensity(model_or_path, surface, LZ, energies, k_path, ...)
SurfaceGreensFunction(model_or_path, surface, energies, k_path, thickness, NN, eps, ...)
FermiArcMap(model_or_path, surface, energy, Nx, Ny, thickness, NN, eps, ...)
generate_k_path(k_points, N_path, labels=None, rec_vecs=None)

# Fermi / band-edge helpers (non-metals)
compute_band_edges(model_or_path, k_mesh=(4,4,4),            # -> {"vbm","cbm","gap","is_metal"}
                   n_occ=None, ref_tol=1e-5)                 # pass n_occ if the zero
                                                             # may not lie in the gap
align_to_vbm(model_or_path, k_mesh=(4,4,4),                  # -> same type in, shifted
             fermi_level=None, if_metal="warn")              # (output is already
                                                             #  VBM-referenced: rarely needed)

# Constants
NUM_ELEMENTS   # 109
NeighBrs       # [17, 3] integer R-vector table

End-to-end example

import numpy as np
from pymatgen.core import Structure
from tailwater import (
    tw_api_call, subspace_projection, tb_model, SurfaceGreensFunction,
)

# 1. Send the structure to the API (one credit, three artifacts)
structure = Structure.from_file("MyMaterial.cif")
paths = tw_api_call(
    structure, user="user", password="pw",
    output_path="./outputs", filename="my_mat",
    project=True,
)

# 2. Fine-tune the heads to fit a near-Fermi window
subspace_projection(
    start_lr=1e-4, end_lr=1e-5, num_epochs=20,
    energy_range=(-2.0, 2.0), decay_sigma=0.5,
    device="cpu",
    save_path="./out_subspace",
    embed_path=paths["embeddings"],
    hr_npz_path=paths["npz"],
)

# 3. Run surface-GF analysis on the projected hr-model
model = tb_model.load("./out_subspace/embeddings_pred.hdf5")
result = SurfaceGreensFunction(
    model, surface=np.eye(3),
    energies=np.linspace(-1, 1, 201),
    k_path=[[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0]],
    k_labels=["M", r"$\Gamma$", "K"],
).run()
result.figure_top.savefig("surface_top.png")

See examples/ for runnable scripts covering each layer in isolation, and Tutorials/ for worked Colab notebooks.

License

Apache 2.0.