Model assembly and hr-file I/O

Helpers for building a tbmodels.Model from the API’s dense head predictions and writing it out as HDF5 or a Wannier90-style _hr.dat file.

See Exporting models: sparse .npz, HDF5, _hr.dat, pybinding, PythTB, Kwant for the customer-facing walkthrough, including the model.to_pb() pybinding-conversion path.

tailwater.hr_export.write_hr_output(hr_model: Model, out_path: str, fmt: str = 'hdf5')[source]

Persist a tbmodels.Model. fmt is “hdf5” or “hr_dat”.

Parameters:

  • hr_model (Model)

  • out_path (str)

  • fmt (str)

Return type:

str

tailwater.hr_export.build_hr_model(edge_pred, onsite_pred, gdata, LM, atoms: List[Tuple[str, List[float]]], hop_threshold: float = 0.01)[source]

Build a tbmodels.Model from the model’s dense predictions.

Parameters:

  • edge_pred (torch.Tensor or ndarray, shape [num_edges, 18, 18, 2]) – (or anything reshape-compatible). Real/imag in the last dim. Self-loop entries are overwritten by onsite_pred internally — caller may pass the raw head output unchanged.

  • onsite_pred (torch.Tensor or ndarray, shape [num_atoms, 18, 18, 2].)

  • gdata (PyG Data with edge_index, edge_vectors, inv_data,) – node_features. Same object the model consumed.

  • LM (3x3 lattice matrix (rows = lattice vectors, Å). Passed) – to tbmodels.Model as uc.

  • atoms ([(symbol, [x, y, z]), ...]. Per-atom Cartesian) – positions; used as the sublattice positions of the tbmodels.Model so the per-atom orbitals carry the same geometric labels they have in the structure.

  • hop_threshold (drop hops with |val| <= this (eV). Default 0.01.)

Returns:

hr_model

Return type:

tbmodels.Model populated with on-site energies and hops.

tailwater.hr_export.build_hr_model_fast(edge_pred, onsite_pred, gdata, LM, atoms: List[Tuple[str, List[float]]], hop_threshold: float = 0.01)[source]

Vectorized equivalent of build_hr_model — same output, much faster.

The reference build_hr_model loops l × atm1 × s1o × atm2 × s2o in pure Python, which is ~num_R * N^2 * 18^2 inner iterations and crosses into multi-minute territory for ~50-atom inputs. This function keeps the EXACT same iteration order (so tbmodels’ first-add-wins semantics and the model’s non-Hermitian inter-atom predictions produce the same final recorded values) but lifts the threshold check, the active-mask filter, the magnitude computation, the value gather, and the orbital-index lookup out of Python and into NumPy. The surviving per-hop Python loop only visits hops that actually got recorded.

Order preservation is what makes this drop-in safe: np.nonzero on a [num_R, N, 18, N, 18] mask returns indices in C (row-major) order, which is identical to the reference’s outer-to-inner nested-loop ordering of those five axes.

Approximate speedup on a 50-atom material: ~100-300x for the hop insertion phase. Output (the tbmodels.Model) is byte-identical to build_hr_model’s output as long as pb.Lattice / tbmodels.Model.add_hop are deterministic — both are.

Parameters:
Return type:

Model