HTTP client + HDF5 loader
Client-side helpers for hitting the Tailwater inference API.
Two entry points:
- tw_api_call(…) — accepts an in-memory pymatgen Structure
(no disk round-trip), uses requests for the POST (proper status-code + error handling), and supports five inference modes:
(default): receive a tbmodels HDF5 hr-model.
return_embeddings = True: receive a .pt file containing the pre-head backbone embeddings, suitable for energetic-subspace fine-tuning via finetune_heads.py / finetune_subspace.py.
return_input = True: receive a .pt file containing the raw GNN input graph — no model inference runs. Useful for debugging the structure-to-graph pipeline or for fully-offline inference.
return_graph_output = True: receive a .pt file containing the model’s dense head outputs (edge_pred, onsite_pred) — full forward runs but tbmodels assembly is skipped. Use this to debug the tbmodels-construction step locally without re-running the model.
project = True: receive a SINGLE zip containing all of {wannier90_hr.hdf5, embeddings.pt, graph_output.pt} — the exact bundle finetune_heads.subspace_projection needs. Costs one credit; saves two follow-up API round trips. The zip is auto-extracted into output_path and the function returns a dict mapping artifact -> filesystem path.
symmetrize = True: receive a zip containing a Kramers-degeneracy- enforced tbmodels HDF5, the raw HDF5, and a per-k helper script. The server detects spatial inversion (P) / C2 around z (C₂ᶻ); if present it applies a minimum-perturbation spectral fix on an adaptive k-mesh (Δk ≈ 0.1 Å⁻¹ by default) — bands stay as close to the raw prediction as possible while doublets become Kramers-paired. If neither symmetry is present the raw model is returned unchanged (a note explains why generic-k splittings must not be averaged out for non-PT crystals). For exact Kramers at arbitrary k, call the bundled kramers_helper.per_k_kramers_fix(raw, k) on the raw HDF5. Costs one credit.
If multiple flags are True the most expensive request wins:
- project > symmetrize > return_input
> return_embeddings > return_graph_output > default HDF5.
- tb_model.load(…) — local HDF5 loader, not an API call. Reads a
tight-binding model produced by the API and returns the standard tbmodels.Model with an instance-bound .to_pb() method that converts it to a pybinding.Lattice for visualization / transport workflows.
Both API entry points use HTTP Basic auth — credentials are checked against the server-side users.db, and each accepted call decrements the caller’s credit balance by 1 (enforced by the require_credit dependency on the server). On 401 (bad credentials) or 402 (out of credits) we surface a clean Python exception so calling code can react.
- tailwater.client.tw_api_call(structure: Structure, user: str, password: str, output_path: str, filename: str, return_embeddings: bool = False, return_input: bool = False, return_graph_output: bool = False, project: bool = False, symmetrize: bool = False, api_url: str = 'https://api.tailwater.io', timeout: float = 600.0, save_cif: bool = True, keep_zip: bool = False, dev: bool = False, model: str | None = None, output_format: str = 'auto')[source]
Submit a pymatgen Structure to the API and save the response.
Five output modes are available, each tapping into a different stage of the inference pipeline. They are mutually exclusive; if multiple flags are True, project wins, then return_input > return_embeddings > return_graph_output > default full HDF5.
(default) full inference -> tbmodels HDF5 hr-model (.hdf5)
return_embeddings = True -> pre-head backbone embeddings (.pt)
- return_input = True -> raw GNN input graph (.pt),
no model inference runs at all
- return_graph_output = True -> dense model output
(edge_pred, onsite_pred) plus graph metadata, before tbmodels assembly. Use this to debug the tbmodels build step locally without re-running the model.
- project = True -> BUNDLE mode: a single zip
containing all of {full HDF5, embeddings.pt, graph_output.pt} — every artifact finetune_heads.subspace_projection needs, in one API call (one credit). The zip is extracted into output_path and a dict mapping artifact name to path is returned instead of a single string.
- symmetrize = True -> SYMMETRIZATION mode: a single
zip containing the symmetrized HDF5 (post-WannSymm), the raw (pre-symm) HDF5, the wannsymm.in actually used, and the wannsymm stdout/stderr log. One credit per call. Use this when you want the predicted Hamiltonian to obey the crystal’s point / space group symmetries exactly, as a post-processing cleanup on top of inference.
- Parameters:
structure (pymatgen.core.structure.Structure) – The structure to model. Serialized in memory via
Structure.as_dict()-> JSON; no on-disk intermediate file is created.user (str) – HTTP Basic auth credentials. Must match a row in the server’s users.db. Each successful call decrements the user’s credit balance by 1 on the server side.
password (str) – HTTP Basic auth credentials. Must match a row in the server’s users.db. Each successful call decrements the user’s credit balance by 1 on the server side.
output_path (str) – Local directory where the response will be saved (created if missing).
filename (str) – Filename stem (without extension). The extension is chosen by the mode: “.hdf5” for the full hr-model, “.pt” for either the embeddings or the input-graph .pt files.
return_embeddings (bool, default False) – Hit the embeddings endpoint instead of full inference. The .pt file is a dict with keys
data(PyG Data object with .f_out and .edge_feat populated),LM(lattice),atoms, andirreps_in.return_input (bool, default False) – Hit the raw-input endpoint — no model inference runs. The .pt file is a dict with keys
data(PyG Data object with the structural fields only: node_features, edge_index, edge_vectors, inv_data, atom_number),LM, andatoms. Use this to inspect the parsed graph (e.g. before feeding it through your own model + heads), or for offline experimentation that doesn’t need a full server-side inference call.return_graph_output (bool, default False) – Run the full model but DON’T assemble tbmodels. The .pt file is a dict with keys
sparse_edge_list(list of [18,18] complex CSR matrices, one per edge),sparse_onsite_list(one [18,18] CSR per atom), plusdata/LM/atoms. Use this to debug the tbmodels assembly step (phase convention, sublattice positions, duplicate-hop handling) by feeding the sparse matrices into your own / a modified hr_export.build_hr_model locally — no model re-runs needed per attempt.project (bool, default False) –
Bundle mode for the subspace-projection workflow. Server runs the full pipeline ONCE and returns a single zip containing wannier90_hr.hdf5 + embeddings.pt + graph_output.pt. The zip is extracted into output_path and the function returns a dict instead of a single path:
{“hdf5”: “…”, “embeddings”: “…”, “graph_output”: “…”}
Costs one credit per call regardless of how many artifacts. Wins over the other return_* flags if multiple are True.
symmetrize (bool, default False) –
Kramers-degeneracy enforcement. When True the server applies the minimum-perturbation spectral fix to the prediction if the crystal has spatial inversion (P) or C2 around z (C₂ᶻ); if not, the raw model is returned unchanged with a note explaining why generic-k splittings (Rashba / Weyl-style) must not be averaged out. Either way you get a single
wannier90_hr.hdf5under the same key, so callers can ignore the symmetry detail and just loadr["hdf5"]. Withsymmetrize=Truethe bundle is:- {“hdf5”: “…”, # the (possibly Kramers-fixed) model
”win”: “…”, # canonical .win “symmetrize_note”: “…”} # symmetry findings + diagnostics
Set
symmetrize=Falseto get the raw prediction (no fix, no symmetry check). Loses toprojectand thereturn_*flags if any of those is also True. For exact Kramers at arbitrary k (band paths, BZ integration on non-mesh k), hit the PT endpoint directly — it bundles the raw HDF5 + a per-k helper script.keep_zip (bool, default False) – When project=True, controls whether the downloaded .zip is retained after extraction. Default False (delete the zip; keep only the three unpacked artifacts).
api_url (str) – Base URL of the API. Defaults to
https://api.tailwater.io(the production deployment). Almost no one should need to set this — only pass it if the Tailwater team specifically pointed you at a different endpoint.timeout (float) – Request timeout in seconds. Backbone inference on a 50-atom material is typically <60 s on CPU; the default 600 s is generous for batched / cold-start cases.
save_cif (bool, default True) – If True, also write the structure to
{output_path}/Structure.cif. Set False to skip.dev (bool, default False) – Opt into the server’s canonical-cell position-wrap fix (sent as
?dev=true). Corrects band structures for inputs whose atoms sit on/over the unit-cell boundary (e.g. fractional coords numerically ~1.0). Default False reproduces the current production behavior, and the flag is harmlessly ignored by servers that predate the patch.model (str, optional (default None)) –
Model checkpoint version. When None (the default) the SDK does NOT forward ?model= to the server, so the server’s own default applies — i.e. whichever checkpoint the operator most recently promoted to default with DEFAULT_MODEL in RunAPI.py. Pass a specific version string to force a particular checkpoint:
”V0.0” → evMace_Epoch_51.pth (the original GWANN release).
- ”V0.1” → Mace_FT2_Gaps_Epoch_7.pth (FT2-Gaps fine-tune;
the current production default since 2026-06-15).
Unknown versions return 400 with the list of valid choices. Older deployments without the registry silently ignore the flag (FastAPI tolerates unknown query params), so forwarding is backward-safe.
output_format (str, default "auto") –
How the full-inference /
projectHamiltonian is transported and delivered. Sent to the server as?format=for the H(R)-returning endpoints; a server that predates the sparse backend ignores it and returns dense HDF5, so every mode degrades cleanly.- ”auto” (default) — request the sparse
wannier90_hr.npz(O(N) egress). SMALL systems (< 30 atoms) are then converted to
wannier90_hr.hdf5client-side and returned under the"hdf5"key exactly as before (the.npzis KEPT too, under"npz"). LARGE systems are left sparse — returned under"npz"— with a printed note on how to convert or analyse them (pybinding / Kwant /_hr.dat/ HDF5).
- ”auto” (default) — request the sparse
- ”sparse” — always keep the raw
.npz(no client-side HDF5 conversion), whatever the system size. Returned under
"npz".
- ”sparse” — always keep the raw
- ”hdf5” — always deliver dense tbmodels HDF5 (server-side for
small systems; client-side conversion of any
.npzas a fallback). May fail / OOM for very large systems — that is what sparsity is for.
Only affects the default full-inference and
projectmodes; thereturn_*/symmetrizemodes ignore it.
- Returns:
- Always a dict. Keys depend on the mode:
default (small, auto) -> {“hdf5”: “…”, “npz”: “…”, “win”: “…”} default (large/sparse)-> {“npz”: “…”, “win”: “…”} return_input -> {“input”: “…”, “win”: “…”} return_embeddings -> {“embeddings”: “…”, “win”: “…”} return_graph_output -> {“graph_output”: “…”, “win”: “…”} project -> {“npz”: “…”, “embeddings”: “…”, “win”: “…”}
(small systems also get a converted “hdf5”)
- symmetrize=True
(default) -> {“hdf5”: “…”, “win”: “…”, “symmetrize_note”: “…”}
With the sparse backend the raw Hamiltonian is always available under
"npz"(atailwater.SparseHR);"hdf5"is additionally present whenever a dense conversion was made (small systems, oroutput_format="hdf5"). Against a pre-sparse server only"hdf5"is returned, as before. The"win"key always points at the canonical wannier90.win file the server actually ran inference on — useful for reproducing the exact graph the server built from your structure (positions, lattice, projections) in any downstream tool.- Return type:
- Raises:
PermissionError – On HTTP 401 — bad username/password.
RuntimeError – On HTTP 402 — out of credits. On any other non-2xx response — surfaces the server’s detail message for debugging.
- tailwater.client.remaining_credits(user: str, password: str, api_url: str = 'https://api.tailwater.io')[source]
Return the caller’s current credit balance via GET /credits/.
This is a free, read-only check — the server’s /credits/ route authenticates but does NOT consume a credit. Returns None if the server doesn’t expose the endpoint (e.g. an older deployment that predates it, which answers 404) or on any non-auth error; raises PermissionError on 401 (bad credentials).
- tailwater.client.k_cart_from_frac(k_frac, lattice_vectors)[source]
Convert a fractional k-point to Cartesian (rad/length) for pybinding.
Pybinding’s
set_wave_vector(k)expectskin rad/length — i.e. in the basis of the Cartesian reciprocal-lattice vectorsb_i, not the fractionalk_iWannier90 and tbmodels use by default. The conversion is:k_cart = 2π · inv(LM) @ k_frac
where
LMhas the real-space lattice vectors as rows.- Parameters:
k_frac (array-like, shape (3,) or (N, 3)) – Fractional k (or batch of k-points), in the same units
tbmodels.Model.hamilton(k)expects.lattice_vectors (array-like, shape (3, 3)) – Real-space lattice vectors as rows (e.g.
model.uc).
- Returns:
Cartesian k in rad/length, ready for
pb.Model.set_wave_vector.- Return type:
np.ndarray of shape
(3,)or(N, 3)
Example
import numpy as np, pybinding as pb from tailwater import tb_model, k_cart_from_frac model = tb_model.load("wannier90_hr.hdf5") lat = model.to_pb() pmod = pb.Model(lat, pb.translational_symmetry()) # Sample H(k) at Gamma → M (Bi2Se3) on a fractional path: k_path_frac = np.array([[0,0,0], [0.5, 0, 0]]) bands = [] for kf in k_path_frac: pmod.set_wave_vector(k_cart_from_frac(kf, model.uc)) bands.append(np.sort(np.linalg.eigvalsh(pmod.hamiltonian.todense())))
- class tailwater.client.tb_model[source]
Bases:
objectLoader namespace for the API’s tight-binding HDF5 outputs.
Usage
from Tailwater import tb_model
model = tb_model.load(“wannier90_hr.hdf5”)
# All standard tbmodels.Model methods/attributes still work: bands = np.stack([model.eigenval(k) for k in k_path]) hops = model.hop size = model.size
# Plus three converters to other tight-binding libraries: pb_lat = model.to_pb() # pybinding.Lattice py_model = model.to_pythtb() # pythtb.tb_model kwant_b = model.to_kwant() # kwant.Builder (3D periodic)
# All three accept an optional lattice-vector override: pb_lat = model.to_pb (lattice_vectors=np.diag([3.5, 3.5, 12.0])) py_model = model.to_pythtb(lattice_vectors=np.diag([3.5, 3.5, 12.0])) kwant_b = model.to_kwant (lattice_vectors=np.diag([3.5, 3.5, 12.0]))
The returned object still passes
isinstance(model, tbmodels.Model)— we attach the converters as bound instance methods rather than swapping the class. Loading multiple HDF5 files in the same Python session is safe: each instance carries its own bindings.- static load(path_to_hdf5: str)[source]
Load a tight-binding model from an HDF5 file with the
to_pb()/to_pythtb()/to_kwant()converters attached.- Parameters:
path_to_hdf5 (str) – Path to an HDF5 file produced by the API (
/upload_json_process_and_download_dat/or extracted from the project bundle).- Returns:
The loaded model, with instance-bound
to_pb(),to_pythtb(), andto_kwant()methods for conversion to pybinding, PythTB, and Kwant respectively. All standardtbmodels.Modelfunctionality is preserved.- Return type:
tbmodels.Model