Subspace fine-tuning
Customer-side fine-tuning on pre-computed backbone embeddings.
This module exposes a single function — subspace_projection — that performs the end-to-end fine-tune + subspace export workflow on a SINGLE material, using two artifacts from one project=True API call:
- embed_path — the API’s embeddings.pt (PyG Data with f_out /
edge_feat + structural metadata).
- hr_npz_path — the API’s wannier90_hr.npz: the predicted
Hamiltonian as a sparse SparseHR. Its in-window eigenvalues are the target the heads are refined to reproduce (eigenvalue-only downfolding), so the heads learn a compact Hamiltonian for the chosen energy window.
Use it from another script as:
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 = “./finetune_out”, embed_path = “./package/embeddings.pt”, hr_npz_path = “./package/wannier90_hr.npz”,
)
Loss modes
With hr_npz_path, loss_mode is “eig_only”: the heads are fit to the predicted Hamiltonian’s eigenvalues that fall in the energy window. You can also bring your OWN eigenvalue targets (e.g. DFT band-structure data) for the same mode — see subspace_utils.make_eigenvalue_only_data.
Advanced: a dense target Hamiltonian may be supplied instead of the .npz, as graph_output_path (the API’s graph_output.pt), which additionally enables the H-MSE-based “subspace” and “full” loss modes. Provide exactly one of hr_npz_path or graph_output_path.
Per-epoch the function prints the mean per-material eigenvalue loss (for “subspace” / “eig_only” modes). For “full” mode it prints the total H MSE since there is no separate eigenvalue term.
Outputs (written to save_path)
HeadsFT_final.pth fine-tuned heads checkpoint + metadata {embedding_stem}_pred.hdf5 per-material projected tbmodels.Model
(subspace-restricted to orbitals whose PREDICTED onsite diagonal lies inside the energy window).
- {embedding_stem}.basis.json per-material basis-info JSON describing
the projected subspace (orbital labels, per-atom positions, energy window, etc.)
- tailwater.finetune_heads.subspace_projection(start_lr: float, end_lr: float, num_epochs: int, energy_range: Tuple[float, float], decay_sigma: float, device, save_path: str, embed_path: str, graph_output_path: str = None, loss_mode: str = 'subspace', *, hr_npz_path: str = None, heads_checkpoint: str = None, kgrid_n: int = 4, h_mse_weight: float = 0.001, eig_weight: float = 1.0, symmetrize_targets: bool = True)[source]
Fine-tune the heads to project a single material into an energy subspace.
Single-material workflow. The customer hits the API twice for the same structure:
- /upload_structure_and_download_embeddings/ -> embed_path (.pt)
backbone features f_out / edge_feat + structural metadata
- /upload_structure_and_download_graph_output/ -> graph_output_path (.pt)
full predicted Hamiltonian (edge_pred / onsite_pred) + structural metadata
The graph-output predictions are attached to the embedding’s PyG Data object as edge_targets — i.e. the model’s pre-fine-tune output is the “ground truth” the subspace fine-tune fits against. This is the self-distillation downfolding setup: the trained model predicts a full Hamiltonian, and the heads are refined so the SUBSPACE-restricted Hamiltonian reproduces the in-window eigenvalues of that full Hamiltonian as accurately as possible.
The LR is cosine-annealed from start_lr down to end_lr over num_epochs epochs. After every epoch the mean eigenvalue loss for the (single) material is printed.
Outputs written to save_path
HeadsFT_final.pth fine-tuned heads weights + metadata {stem}_pred.hdf5 projected tbmodels.Model
(SUBSPACE-restricted to orbitals whose predicted onsite-diagonal lies inside the energy window)
- {stem}.basis.json basis-info file mapping subspace
indices to (atom, spatial, spin) labels.
where {stem} is the basename of embed_path without extension.
- type start_lr:
- param start_lr:
- type start_lr:
initial LR for AdamW.
- type end_lr:
- param end_lr:
- type end_lr:
minimum LR at the end of cosine decay.
- type num_epochs:
- param num_epochs:
- type num_epochs:
number of training epochs.
- type energy_range:
- param energy_range:
- type energy_range:
(e_lo, e_hi) tuple, eV, relative to E_F = 0.
- type decay_sigma:
- param decay_sigma:
the window center.
- type decay_sigma:
Gaussian sigma for eigenvalue weights around
- type device:
torch.device or str.
- param device:
- type device:
torch.device or str.
- type save_path:
- param save_path:
- type save_path:
output directory (created if missing).
- type embed_path:
- param embed_path:
dict with data, LM, atoms, irreps_in).
- type embed_path:
path to ONE .pt embedding file (API-format
- type graph_output_path:
- param graph_output_path:
dict with edge_pred, onsite_pred, data, LM, atoms). This is the FULL model output that the projection refines toward. Provide EITHER this OR hr_npz_path (not both). Optional now that the sparse .npz target is supported.
- type graph_output_path:
path to ONE .pt graph-output file (API-format
- type loss_mode:
- param loss_mode:
Ignored / forced to “eig_only” when hr_npz_path is used (the .npz carries eigenvalues, not the dense target Hamiltonian the H-MSE term needs).
- type loss_mode:
“subspace” (default) | “full” | “eig_only”.
- type hr_npz_path:
- param hr_npz_path:
wannier90_hr.npz produced by the optimized API (the project=True bundle under output_format “auto”/”sparse”, or any tw_api_call “npz” result). Its in-window eigenvalues become the downfolding target — the whole project workflow then needs only embeddings.pt + the .npz, no dense graph_output.pt / HDF5. Mutually exclusive with graph_output_path.
- type hr_npz_path:
(keyword-only) path to the sparse
- type heads_checkpoint:
- param heads_checkpoint:
Defaults to the MACE-compatible HeadsOnly_MACE.pth that ships inside the installed package, so you don’t need to supply one for the standard workflow. Pass an explicit path only if you’re starting from a custom-fine-tuned heads checkpoint.
- type heads_checkpoint:
starting HeadsOnly checkpoint (path to a .pth).
- rtype:
str path to the saved HeadsFT_final.pth checkpoint.
- param kgrid_n:
- type kgrid_n:
- param h_mse_weight:
- type h_mse_weight:
- param eig_weight:
- type eig_weight:
- param symmetrize_targets:
- type symmetrize_targets: