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, and irreps_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, and atoms. 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), plus data / 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.hdf5 under the same key, so callers can ignore the symmetry detail and just load r["hdf5"]. With symmetrize=True the bundle is:

    {“hdf5”: “…”, # the (possibly Kramers-fixed) model

    ”win”: “…”, # canonical .win “symmetrize_note”: “…”} # symmetry findings + diagnostics

    Set symmetrize=False to get the raw prediction (no fix, no symmetry check). Loses to project and the return_* 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 / project Hamiltonian 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.hdf5 client-side and returned under the "hdf5" key exactly as before (the .npz is 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).

    • ”sparse” — always keep the raw .npz (no client-side HDF5

      conversion), whatever the system size. Returned under "npz".

    • ”hdf5” — always deliver dense tbmodels HDF5 (server-side for

      small systems; client-side conversion of any .npz as a fallback). May fail / OOM for very large systems — that is what sparsity is for.

    Only affects the default full-inference and project modes; the return_* / symmetrize modes 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" (a tailwater.SparseHR); "hdf5" is additionally present whenever a dense conversion was made (small systems, or output_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:

dict

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).

Parameters:
Return type:

Optional[int]

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) expects k in rad/length — i.e. in the basis of the Cartesian reciprocal-lattice vectors b_i, not the fractional k_i Wannier90 and tbmodels use by default. The conversion is:

k_cart = 2π · inv(LM) @ k_frac

where LM has 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: object

Loader 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(), and to_kwant() methods for conversion to pybinding, PythTB, and Kwant respectively. All standard tbmodels.Model functionality is preserved.

Return type:

tbmodels.Model