tailwater 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- tailwater-0.2.0/CHANGELOG.md +30 -0
- tailwater-0.2.0/LICENSE +15 -0
- tailwater-0.2.0/MANIFEST.in +7 -0
- tailwater-0.2.0/PKG-INFO +343 -0
- tailwater-0.2.0/README.md +296 -0
- tailwater-0.2.0/examples/01_basic_api_call.py +29 -0
- tailwater-0.2.0/examples/02_subspace_projection.py +43 -0
- tailwater-0.2.0/examples/03_surface_analysis.py +63 -0
- tailwater-0.2.0/examples/04_band_structure.py +56 -0
- tailwater-0.2.0/pyproject.toml +99 -0
- tailwater-0.2.0/setup.cfg +4 -0
- tailwater-0.2.0/src/tailwater/__init__.py +114 -0
- tailwater-0.2.0/src/tailwater/client.py +610 -0
- tailwater-0.2.0/src/tailwater/constants.py +30 -0
- tailwater-0.2.0/src/tailwater/finetune_heads.py +552 -0
- tailwater-0.2.0/src/tailwater/heads_only_model.py +317 -0
- tailwater-0.2.0/src/tailwater/hr_export.py +273 -0
- tailwater-0.2.0/src/tailwater/make_heads_only.py +34 -0
- tailwater-0.2.0/src/tailwater/subspace_utils.py +832 -0
- tailwater-0.2.0/src/tailwater/wannier_wizard.py +1466 -0
- tailwater-0.2.0/src/tailwater.egg-info/PKG-INFO +343 -0
- tailwater-0.2.0/src/tailwater.egg-info/SOURCES.txt +24 -0
- tailwater-0.2.0/src/tailwater.egg-info/dependency_links.txt +1 -0
- tailwater-0.2.0/src/tailwater.egg-info/requires.txt +26 -0
- tailwater-0.2.0/src/tailwater.egg-info/top_level.txt +1 -0
- tailwater-0.2.0/tests/test_imports.py +82 -0
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
# Changelog
|
|
2
|
+
|
|
3
|
+
All notable changes to the `tailwater` package. This project follows
|
|
4
|
+
[Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
|
5
|
+
|
|
6
|
+
## [0.2.0]
|
|
7
|
+
|
|
8
|
+
### Changed
|
|
9
|
+
- **Default API endpoint is now `https://api.tailwater.io`**
|
|
10
|
+
(previously `http://127.0.0.1:8000`). Existing notebooks that already set
|
|
11
|
+
`TW_API_URL` or pass `api_url=` are unaffected. To keep pointing at a local
|
|
12
|
+
dev server, set `TW_API_URL=http://127.0.0.1:8000` or pass
|
|
13
|
+
`api_url="http://127.0.0.1:8000"` per call.
|
|
14
|
+
|
|
15
|
+
### Removed
|
|
16
|
+
- Stale `TW_API` legacy-callable reference dropped from `README.md` and
|
|
17
|
+
`docs/api/client.rst` (the symbol itself was removed in an earlier change,
|
|
18
|
+
but the doc references lingered and produced Sphinx warnings).
|
|
19
|
+
|
|
20
|
+
### Docs
|
|
21
|
+
- Added an "API access" section to `README.md` and a "Getting API access"
|
|
22
|
+
section to `docs/installation.rst` covering the default endpoint,
|
|
23
|
+
credentials flow (HTTP Basic, request from the Tailwater team), credit
|
|
24
|
+
metering, and how to override the endpoint for local/private deployments.
|
|
25
|
+
- `docs/quickstart.rst` opens with a brief "you need credentials" note and
|
|
26
|
+
links into the installation page for endpoint configuration.
|
|
27
|
+
|
|
28
|
+
## [0.1.0]
|
|
29
|
+
|
|
30
|
+
Initial release.
|
tailwater-0.2.0/LICENSE
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
Apache License
|
|
2
|
+
Version 2.0, January 2004
|
|
3
|
+
http://www.apache.org/licenses/
|
|
4
|
+
|
|
5
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
6
|
+
you may not use this file except in compliance with the License.
|
|
7
|
+
You may obtain a copy of the License at
|
|
8
|
+
|
|
9
|
+
http://www.apache.org/licenses/LICENSE-2.0
|
|
10
|
+
|
|
11
|
+
Unless required by applicable law or agreed to in writing, software
|
|
12
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
13
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
14
|
+
See the License for the specific language governing permissions and
|
|
15
|
+
limitations under the License.
|
tailwater-0.2.0/PKG-INFO
ADDED
|
@@ -0,0 +1,343 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: tailwater
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Client + post-processing toolkit for the Tailwater Wannier-Hamiltonian inference API
|
|
5
|
+
Author: Tailwater
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Project-URL: Homepage, https://github.com/tailwater/tailwater
|
|
8
|
+
Project-URL: Documentation, https://github.com/tailwater/tailwater#readme
|
|
9
|
+
Project-URL: Issues, https://github.com/tailwater/tailwater/issues
|
|
10
|
+
Keywords: wannier,tight-binding,DFT,GNN,tbmodels,KPM,surface states,fermi arc,topological
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
14
|
+
Classifier: Operating System :: OS Independent
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
20
|
+
Classifier: Topic :: Scientific/Engineering :: Physics
|
|
21
|
+
Classifier: Topic :: Scientific/Engineering :: Chemistry
|
|
22
|
+
Requires-Python: >=3.9
|
|
23
|
+
Description-Content-Type: text/markdown
|
|
24
|
+
License-File: LICENSE
|
|
25
|
+
Requires-Dist: numpy>=1.20
|
|
26
|
+
Requires-Dist: scipy>=1.7
|
|
27
|
+
Requires-Dist: requests>=2.28
|
|
28
|
+
Requires-Dist: tqdm>=4.60
|
|
29
|
+
Requires-Dist: matplotlib>=3.5
|
|
30
|
+
Requires-Dist: h5py>=3.7
|
|
31
|
+
Requires-Dist: pymatgen>=2023.0
|
|
32
|
+
Requires-Dist: tbmodels>=1.4
|
|
33
|
+
Requires-Dist: torch>=2.0
|
|
34
|
+
Requires-Dist: torch-geometric>=2.0
|
|
35
|
+
Requires-Dist: e3nn>=0.5
|
|
36
|
+
Provides-Extra: pybinding
|
|
37
|
+
Requires-Dist: pybinding>=0.9; extra == "pybinding"
|
|
38
|
+
Provides-Extra: scatter
|
|
39
|
+
Requires-Dist: torch-scatter>=2.0; extra == "scatter"
|
|
40
|
+
Provides-Extra: seekpath
|
|
41
|
+
Requires-Dist: seekpath>=2.0; extra == "seekpath"
|
|
42
|
+
Provides-Extra: dev
|
|
43
|
+
Requires-Dist: pytest>=7.0; extra == "dev"
|
|
44
|
+
Requires-Dist: ruff>=0.1; extra == "dev"
|
|
45
|
+
Requires-Dist: build>=1.0; extra == "dev"
|
|
46
|
+
Requires-Dist: twine>=4.0; extra == "dev"
|
|
47
|
+
|
|
48
|
+
# tailwater
|
|
49
|
+
|
|
50
|
+
Client + post-processing toolkit for the **Tailwater** Wannier-Hamiltonian inference API.
|
|
51
|
+
|
|
52
|
+
`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.
|
|
53
|
+
|
|
54
|
+
---
|
|
55
|
+
|
|
56
|
+
## Installation
|
|
57
|
+
|
|
58
|
+
```bash
|
|
59
|
+
pip install tailwater
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
Optional extras:
|
|
63
|
+
|
|
64
|
+
```bash
|
|
65
|
+
pip install "tailwater[pybinding]" # enables tb_model.to_pb()
|
|
66
|
+
pip install "tailwater[scatter]" # if torch_scatter import fails
|
|
67
|
+
pip install "tailwater[seekpath]" # enables auto k-path mode of bulk_band_structure
|
|
68
|
+
pip install "tailwater[dev]" # pytest, ruff, build, twine
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
Tested on Python 3.9–3.12.
|
|
72
|
+
|
|
73
|
+
---
|
|
74
|
+
|
|
75
|
+
## API access
|
|
76
|
+
|
|
77
|
+
The Tailwater inference API is hosted at **`https://api.tailwater.io`** — this
|
|
78
|
+
is the default endpoint `tw_api_call(...)` talks to, so the basic usage below
|
|
79
|
+
needs no extra configuration beyond your credentials.
|
|
80
|
+
|
|
81
|
+
- **Credentials.** Authentication is **HTTP Basic** (username + password).
|
|
82
|
+
Email the Tailwater team to request an account; you'll be issued a username
|
|
83
|
+
and a one-time-displayed password.
|
|
84
|
+
- **Billing.** Each successful inference call decrements your server-side
|
|
85
|
+
credit balance by one. Health checks (`/healthz`) and balance lookups
|
|
86
|
+
(`/credits/`) are free.
|
|
87
|
+
- **Checking your balance:**
|
|
88
|
+
|
|
89
|
+
```python
|
|
90
|
+
from tailwater import remaining_credits
|
|
91
|
+
print(remaining_credits("user", "pw")) # -> int
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
- **Overriding the endpoint** (local dev server, a private deployment, or a
|
|
95
|
+
staging URL) — either set the env var once per session:
|
|
96
|
+
|
|
97
|
+
```python
|
|
98
|
+
import os
|
|
99
|
+
os.environ["TW_API_URL"] = "http://127.0.0.1:8000"
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
or pass it per call:
|
|
103
|
+
|
|
104
|
+
```python
|
|
105
|
+
paths = tw_api_call(structure, "user", "pw", "./out", "my_mat",
|
|
106
|
+
api_url="http://127.0.0.1:8000")
|
|
107
|
+
```
|
|
108
|
+
|
|
109
|
+
---
|
|
110
|
+
|
|
111
|
+
## Three workflow layers
|
|
112
|
+
|
|
113
|
+
### 1. HTTP client — talk to the API
|
|
114
|
+
|
|
115
|
+
```python
|
|
116
|
+
from pymatgen.core import Structure
|
|
117
|
+
from tailwater import tw_api_call
|
|
118
|
+
|
|
119
|
+
structure = Structure.from_file("MyMaterial.cif")
|
|
120
|
+
|
|
121
|
+
# tw_api_call ALWAYS returns a dict of extracted paths. Every response
|
|
122
|
+
# includes a "win" key — the canonical wannier90.win file the server
|
|
123
|
+
# actually ran inference on — alongside the mode-specific artifact(s).
|
|
124
|
+
|
|
125
|
+
# (a) default: tbmodels HDF5 hr-model + .win
|
|
126
|
+
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat")
|
|
127
|
+
# paths = {"hdf5": "...", "win": "..."}
|
|
128
|
+
|
|
129
|
+
# (b) backbone embeddings + .win
|
|
130
|
+
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat",
|
|
131
|
+
return_embeddings=True)
|
|
132
|
+
# paths = {"embeddings": "...", "win": "..."}
|
|
133
|
+
|
|
134
|
+
# (c) project bundle: all three artifacts + .win in a single call
|
|
135
|
+
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat",
|
|
136
|
+
project=True)
|
|
137
|
+
# paths = {"hdf5": "...", "embeddings": "...",
|
|
138
|
+
# "graph_output": "...", "win": "..."}
|
|
139
|
+
```
|
|
140
|
+
|
|
141
|
+
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.
|
|
142
|
+
|
|
143
|
+
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).
|
|
144
|
+
|
|
145
|
+
### 2. Subspace projection — fine-tune heads on supplier-side embeddings
|
|
146
|
+
|
|
147
|
+
Once you have the project bundle, you can fine-tune the output heads to fit a narrow energy window near the Fermi level:
|
|
148
|
+
|
|
149
|
+
```python
|
|
150
|
+
from tailwater import subspace_projection
|
|
151
|
+
|
|
152
|
+
subspace_projection(
|
|
153
|
+
start_lr = 5e-5,
|
|
154
|
+
end_lr = 5e-7,
|
|
155
|
+
num_epochs = 20,
|
|
156
|
+
energy_range = (-2.0, 2.0), # eV, relative to E_F
|
|
157
|
+
decay_sigma = 1.0,
|
|
158
|
+
device = "cpu",
|
|
159
|
+
save_path = "./projection_out",
|
|
160
|
+
embed_path = paths["embeddings"],
|
|
161
|
+
graph_output_path = paths["graph_output"],
|
|
162
|
+
loss_mode = "subspace", # default
|
|
163
|
+
)
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
Per epoch the script prints the mean eigenvalue loss. When done, three files are written to `save_path`:
|
|
167
|
+
|
|
168
|
+
| File | Contents |
|
|
169
|
+
|---|---|
|
|
170
|
+
| `HeadsFT_final.pth` | fine-tuned heads weights + metadata |
|
|
171
|
+
| `{stem}_pred.hdf5` | projected, subspace-restricted `tbmodels.Model` |
|
|
172
|
+
| `{stem}.basis.json` | mapping from subspace indices to `(atom, spatial, spin)` labels |
|
|
173
|
+
|
|
174
|
+
Three loss modes are exposed:
|
|
175
|
+
- `"subspace"` (default) — H-MSE + weighted eigenvalue loss within the energy window
|
|
176
|
+
- `"eig_only"` — eigenvalue-only fine-tune; no Hamiltonian targets needed
|
|
177
|
+
- `"full"` — plain H-MSE across all orbitals
|
|
178
|
+
|
|
179
|
+
### 3. Post-processing — bulk DOS, surface states, Fermi arcs
|
|
180
|
+
|
|
181
|
+
```python
|
|
182
|
+
import numpy as np
|
|
183
|
+
from tailwater import (
|
|
184
|
+
tb_model,
|
|
185
|
+
BulkDOS,
|
|
186
|
+
SurfaceSpectralDensity,
|
|
187
|
+
SurfaceGreensFunction,
|
|
188
|
+
FermiArcMap,
|
|
189
|
+
)
|
|
190
|
+
|
|
191
|
+
# Load the HDF5 the API produced — returns a tbmodels.Model with .to_pb()
|
|
192
|
+
model = tb_model.load("outputs/wannier90_hr.hdf5")
|
|
193
|
+
|
|
194
|
+
# Bulk DOS (KPM, k-mesh averaged)
|
|
195
|
+
result = BulkDOS(model, k_mesh=(8, 8, 8), energies=(-4, 4),
|
|
196
|
+
NC=2048, NV=4).run()
|
|
197
|
+
result.figure.savefig("bulk_dos.png")
|
|
198
|
+
np.savez("bulk_dos.npz", **result.as_dict())
|
|
199
|
+
|
|
200
|
+
# Surface spectral density along a k-path (KPM)
|
|
201
|
+
result = SurfaceSpectralDensity(
|
|
202
|
+
model, surface=np.eye(3), LZ=5,
|
|
203
|
+
energies=(-1, 1),
|
|
204
|
+
k_path=[[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0]],
|
|
205
|
+
k_labels=["M", r"$\Gamma$", "K"],
|
|
206
|
+
N_path=101, NC=2**12, NV=4,
|
|
207
|
+
).run()
|
|
208
|
+
result.figure_top.savefig("surface_top.png")
|
|
209
|
+
result.figure_bottom.savefig("surface_bottom.png")
|
|
210
|
+
|
|
211
|
+
# Surface Green's function (Lopez-Sancho)
|
|
212
|
+
result = SurfaceGreensFunction(
|
|
213
|
+
model, surface=np.eye(3),
|
|
214
|
+
energies=np.linspace(-1, 1, 201),
|
|
215
|
+
k_path=[[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0]],
|
|
216
|
+
k_labels=["M", r"$\Gamma$", "K"],
|
|
217
|
+
N_path=101, thickness=6, NN=5, eps=0.005,
|
|
218
|
+
).run()
|
|
219
|
+
np.savez("surface_gf.npz", **result.as_dict())
|
|
220
|
+
|
|
221
|
+
# 2D Fermi-arc map at one energy
|
|
222
|
+
result = FermiArcMap(
|
|
223
|
+
model, surface=np.eye(3), energy=0.0,
|
|
224
|
+
Nx=50, Ny=50, thickness=6,
|
|
225
|
+
).run()
|
|
226
|
+
result.figure_top_interpolated.savefig("fermi_arc_top.png")
|
|
227
|
+
|
|
228
|
+
# Bulk band structure along a manual k-path
|
|
229
|
+
from tailwater import bulk_band_structure
|
|
230
|
+
fig = bulk_band_structure(
|
|
231
|
+
model,
|
|
232
|
+
k_points = [[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0], [0, 0, 0]],
|
|
233
|
+
k_labels = ["M", r"$\Gamma$", "K", r"$\Gamma$"],
|
|
234
|
+
spacing = 0.01,
|
|
235
|
+
fermi_level = 0.0,
|
|
236
|
+
e_range = (-3, 3),
|
|
237
|
+
)
|
|
238
|
+
fig.savefig("bands.png")
|
|
239
|
+
|
|
240
|
+
# Or use seekpath to auto-determine the high-symmetry path
|
|
241
|
+
from pymatgen.core import Structure
|
|
242
|
+
structure = Structure.from_file("MyMaterial.cif")
|
|
243
|
+
fig = bulk_band_structure(model, auto=True, structure=structure,
|
|
244
|
+
spacing=0.02, e_range=(-3, 3))
|
|
245
|
+
fig.savefig("bands_auto.png")
|
|
246
|
+
```
|
|
247
|
+
|
|
248
|
+
Each post-processing class accepts either an HDF5 path (`str`) or an in-memory `tbmodels.Model`. The `.run()` method returns a typed `Result` dataclass with raw NumPy arrays and matplotlib `Figure` objects.
|
|
249
|
+
|
|
250
|
+
---
|
|
251
|
+
|
|
252
|
+
## API reference (top-level imports)
|
|
253
|
+
|
|
254
|
+
```python
|
|
255
|
+
# HTTP client + HDF5 loader
|
|
256
|
+
tw_api_call(structure, user, password, output_path, filename, ...)
|
|
257
|
+
tb_model.load(path_to_hdf5)
|
|
258
|
+
remaining_credits(user, password)
|
|
259
|
+
|
|
260
|
+
# Heads-only inference model
|
|
261
|
+
HeadsOnly(irreps_in)
|
|
262
|
+
CovariantOnsiteHead(irreps_in)
|
|
263
|
+
CovariantEdgeHead(irreps_in)
|
|
264
|
+
load_heads_only_checkpoint(path)
|
|
265
|
+
save_heads_only_checkpoint(full_state_dict, irreps_in_str, save_path)
|
|
266
|
+
|
|
267
|
+
# Subspace fine-tuning
|
|
268
|
+
subspace_projection(start_lr, end_lr, num_epochs, energy_range,
|
|
269
|
+
decay_sigma, device, save_path,
|
|
270
|
+
embed_path, graph_output_path, loss_mode="subspace")
|
|
271
|
+
|
|
272
|
+
# Subspace losses (advanced)
|
|
273
|
+
Subspace_H_MSE_Loss(gdata, edge_pred, onsite_pred, e_lo, e_hi)
|
|
274
|
+
Subspace_EigLoss(gdata, edge_pred, onsite_pred, kvec, neighbrs, e_lo, e_hi)
|
|
275
|
+
Eigenvalue_Only_Loss(gdata, edge_pred, onsite_pred, e_lo, e_hi)
|
|
276
|
+
make_eigenvalue_only_data(gdata, kvecs, eigs_per_k, e_lo, e_hi)
|
|
277
|
+
build_subspace_active_mask(node_features, onsite_target, e_lo, e_hi)
|
|
278
|
+
write_subspace_basis_file(out_path, active_mask, atoms, LM, ...)
|
|
279
|
+
|
|
280
|
+
# tbmodels.Model assembly from raw head output
|
|
281
|
+
build_hr_model (edge_pred, onsite_pred, gdata, LM, atoms)
|
|
282
|
+
build_hr_model_fast(edge_pred, onsite_pred, gdata, LM, atoms) # vectorized
|
|
283
|
+
write_hr_output(hr_model, out_path, fmt="hdf5"|"hr_dat")
|
|
284
|
+
|
|
285
|
+
# Post-processing calculators (each has a .run() method returning a Result)
|
|
286
|
+
BulkDOS(model_or_path, k_mesh, energies, NC, NV, device)
|
|
287
|
+
SurfaceSpectralDensity(model_or_path, surface, LZ, energies, k_path, ...)
|
|
288
|
+
SurfaceGreensFunction(model_or_path, surface, energies, k_path, thickness, NN, eps, ...)
|
|
289
|
+
FermiArcMap(model_or_path, surface, energy, Nx, Ny, thickness, NN, eps, ...)
|
|
290
|
+
generate_k_path(k_points, N_path, labels=None, rec_vecs=None)
|
|
291
|
+
|
|
292
|
+
# Constants
|
|
293
|
+
NUM_ELEMENTS # 109
|
|
294
|
+
NeighBrs # [17, 3] integer R-vector table
|
|
295
|
+
```
|
|
296
|
+
|
|
297
|
+
---
|
|
298
|
+
|
|
299
|
+
## End-to-end example
|
|
300
|
+
|
|
301
|
+
```python
|
|
302
|
+
import numpy as np
|
|
303
|
+
from pymatgen.core import Structure
|
|
304
|
+
from tailwater import (
|
|
305
|
+
tw_api_call, subspace_projection, tb_model, SurfaceGreensFunction,
|
|
306
|
+
)
|
|
307
|
+
|
|
308
|
+
# 1. Send the structure to the API (one credit, three artifacts)
|
|
309
|
+
structure = Structure.from_file("MyMaterial.cif")
|
|
310
|
+
paths = tw_api_call(
|
|
311
|
+
structure, user="user", password="pw",
|
|
312
|
+
output_path="./outputs", filename="my_mat",
|
|
313
|
+
project=True,
|
|
314
|
+
)
|
|
315
|
+
|
|
316
|
+
# 2. Fine-tune the heads to fit a near-Fermi window
|
|
317
|
+
subspace_projection(
|
|
318
|
+
start_lr=5e-5, end_lr=5e-7, num_epochs=20,
|
|
319
|
+
energy_range=(-2.0, 2.0), decay_sigma=1.0,
|
|
320
|
+
device="cpu",
|
|
321
|
+
save_path="./out_subspace",
|
|
322
|
+
embed_path=paths["embeddings"],
|
|
323
|
+
graph_output_path=paths["graph_output"],
|
|
324
|
+
)
|
|
325
|
+
|
|
326
|
+
# 3. Run surface-GF analysis on the projected hr-model
|
|
327
|
+
model = tb_model.load("./out_subspace/embeddings_pred.hdf5")
|
|
328
|
+
result = SurfaceGreensFunction(
|
|
329
|
+
model, surface=np.eye(3),
|
|
330
|
+
energies=np.linspace(-1, 1, 201),
|
|
331
|
+
k_path=[[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0]],
|
|
332
|
+
k_labels=["M", r"$\Gamma$", "K"],
|
|
333
|
+
).run()
|
|
334
|
+
result.figure_top.savefig("surface_top.png")
|
|
335
|
+
```
|
|
336
|
+
|
|
337
|
+
See `examples/` for runnable scripts covering each layer in isolation.
|
|
338
|
+
|
|
339
|
+
---
|
|
340
|
+
|
|
341
|
+
## License
|
|
342
|
+
|
|
343
|
+
Apache 2.0.
|
|
@@ -0,0 +1,296 @@
|
|
|
1
|
+
# tailwater
|
|
2
|
+
|
|
3
|
+
Client + post-processing toolkit for the **Tailwater** Wannier-Hamiltonian inference API.
|
|
4
|
+
|
|
5
|
+
`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.
|
|
6
|
+
|
|
7
|
+
---
|
|
8
|
+
|
|
9
|
+
## Installation
|
|
10
|
+
|
|
11
|
+
```bash
|
|
12
|
+
pip install tailwater
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
Optional extras:
|
|
16
|
+
|
|
17
|
+
```bash
|
|
18
|
+
pip install "tailwater[pybinding]" # enables tb_model.to_pb()
|
|
19
|
+
pip install "tailwater[scatter]" # if torch_scatter import fails
|
|
20
|
+
pip install "tailwater[seekpath]" # enables auto k-path mode of bulk_band_structure
|
|
21
|
+
pip install "tailwater[dev]" # pytest, ruff, build, twine
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
Tested on Python 3.9–3.12.
|
|
25
|
+
|
|
26
|
+
---
|
|
27
|
+
|
|
28
|
+
## API access
|
|
29
|
+
|
|
30
|
+
The Tailwater inference API is hosted at **`https://api.tailwater.io`** — this
|
|
31
|
+
is the default endpoint `tw_api_call(...)` talks to, so the basic usage below
|
|
32
|
+
needs no extra configuration beyond your credentials.
|
|
33
|
+
|
|
34
|
+
- **Credentials.** Authentication is **HTTP Basic** (username + password).
|
|
35
|
+
Email the Tailwater team to request an account; you'll be issued a username
|
|
36
|
+
and a one-time-displayed password.
|
|
37
|
+
- **Billing.** Each successful inference call decrements your server-side
|
|
38
|
+
credit balance by one. Health checks (`/healthz`) and balance lookups
|
|
39
|
+
(`/credits/`) are free.
|
|
40
|
+
- **Checking your balance:**
|
|
41
|
+
|
|
42
|
+
```python
|
|
43
|
+
from tailwater import remaining_credits
|
|
44
|
+
print(remaining_credits("user", "pw")) # -> int
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
- **Overriding the endpoint** (local dev server, a private deployment, or a
|
|
48
|
+
staging URL) — either set the env var once per session:
|
|
49
|
+
|
|
50
|
+
```python
|
|
51
|
+
import os
|
|
52
|
+
os.environ["TW_API_URL"] = "http://127.0.0.1:8000"
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
or pass it per call:
|
|
56
|
+
|
|
57
|
+
```python
|
|
58
|
+
paths = tw_api_call(structure, "user", "pw", "./out", "my_mat",
|
|
59
|
+
api_url="http://127.0.0.1:8000")
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
---
|
|
63
|
+
|
|
64
|
+
## Three workflow layers
|
|
65
|
+
|
|
66
|
+
### 1. HTTP client — talk to the API
|
|
67
|
+
|
|
68
|
+
```python
|
|
69
|
+
from pymatgen.core import Structure
|
|
70
|
+
from tailwater import tw_api_call
|
|
71
|
+
|
|
72
|
+
structure = Structure.from_file("MyMaterial.cif")
|
|
73
|
+
|
|
74
|
+
# tw_api_call ALWAYS returns a dict of extracted paths. Every response
|
|
75
|
+
# includes a "win" key — the canonical wannier90.win file the server
|
|
76
|
+
# actually ran inference on — alongside the mode-specific artifact(s).
|
|
77
|
+
|
|
78
|
+
# (a) default: tbmodels HDF5 hr-model + .win
|
|
79
|
+
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat")
|
|
80
|
+
# paths = {"hdf5": "...", "win": "..."}
|
|
81
|
+
|
|
82
|
+
# (b) backbone embeddings + .win
|
|
83
|
+
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat",
|
|
84
|
+
return_embeddings=True)
|
|
85
|
+
# paths = {"embeddings": "...", "win": "..."}
|
|
86
|
+
|
|
87
|
+
# (c) project bundle: all three artifacts + .win in a single call
|
|
88
|
+
paths = tw_api_call(structure, "user", "pw", "./outputs", "my_mat",
|
|
89
|
+
project=True)
|
|
90
|
+
# paths = {"hdf5": "...", "embeddings": "...",
|
|
91
|
+
# "graph_output": "...", "win": "..."}
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
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.
|
|
95
|
+
|
|
96
|
+
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).
|
|
97
|
+
|
|
98
|
+
### 2. Subspace projection — fine-tune heads on supplier-side embeddings
|
|
99
|
+
|
|
100
|
+
Once you have the project bundle, you can fine-tune the output heads to fit a narrow energy window near the Fermi level:
|
|
101
|
+
|
|
102
|
+
```python
|
|
103
|
+
from tailwater import subspace_projection
|
|
104
|
+
|
|
105
|
+
subspace_projection(
|
|
106
|
+
start_lr = 5e-5,
|
|
107
|
+
end_lr = 5e-7,
|
|
108
|
+
num_epochs = 20,
|
|
109
|
+
energy_range = (-2.0, 2.0), # eV, relative to E_F
|
|
110
|
+
decay_sigma = 1.0,
|
|
111
|
+
device = "cpu",
|
|
112
|
+
save_path = "./projection_out",
|
|
113
|
+
embed_path = paths["embeddings"],
|
|
114
|
+
graph_output_path = paths["graph_output"],
|
|
115
|
+
loss_mode = "subspace", # default
|
|
116
|
+
)
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
Per epoch the script prints the mean eigenvalue loss. When done, three files are written to `save_path`:
|
|
120
|
+
|
|
121
|
+
| File | Contents |
|
|
122
|
+
|---|---|
|
|
123
|
+
| `HeadsFT_final.pth` | fine-tuned heads weights + metadata |
|
|
124
|
+
| `{stem}_pred.hdf5` | projected, subspace-restricted `tbmodels.Model` |
|
|
125
|
+
| `{stem}.basis.json` | mapping from subspace indices to `(atom, spatial, spin)` labels |
|
|
126
|
+
|
|
127
|
+
Three loss modes are exposed:
|
|
128
|
+
- `"subspace"` (default) — H-MSE + weighted eigenvalue loss within the energy window
|
|
129
|
+
- `"eig_only"` — eigenvalue-only fine-tune; no Hamiltonian targets needed
|
|
130
|
+
- `"full"` — plain H-MSE across all orbitals
|
|
131
|
+
|
|
132
|
+
### 3. Post-processing — bulk DOS, surface states, Fermi arcs
|
|
133
|
+
|
|
134
|
+
```python
|
|
135
|
+
import numpy as np
|
|
136
|
+
from tailwater import (
|
|
137
|
+
tb_model,
|
|
138
|
+
BulkDOS,
|
|
139
|
+
SurfaceSpectralDensity,
|
|
140
|
+
SurfaceGreensFunction,
|
|
141
|
+
FermiArcMap,
|
|
142
|
+
)
|
|
143
|
+
|
|
144
|
+
# Load the HDF5 the API produced — returns a tbmodels.Model with .to_pb()
|
|
145
|
+
model = tb_model.load("outputs/wannier90_hr.hdf5")
|
|
146
|
+
|
|
147
|
+
# Bulk DOS (KPM, k-mesh averaged)
|
|
148
|
+
result = BulkDOS(model, k_mesh=(8, 8, 8), energies=(-4, 4),
|
|
149
|
+
NC=2048, NV=4).run()
|
|
150
|
+
result.figure.savefig("bulk_dos.png")
|
|
151
|
+
np.savez("bulk_dos.npz", **result.as_dict())
|
|
152
|
+
|
|
153
|
+
# Surface spectral density along a k-path (KPM)
|
|
154
|
+
result = SurfaceSpectralDensity(
|
|
155
|
+
model, surface=np.eye(3), LZ=5,
|
|
156
|
+
energies=(-1, 1),
|
|
157
|
+
k_path=[[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0]],
|
|
158
|
+
k_labels=["M", r"$\Gamma$", "K"],
|
|
159
|
+
N_path=101, NC=2**12, NV=4,
|
|
160
|
+
).run()
|
|
161
|
+
result.figure_top.savefig("surface_top.png")
|
|
162
|
+
result.figure_bottom.savefig("surface_bottom.png")
|
|
163
|
+
|
|
164
|
+
# Surface Green's function (Lopez-Sancho)
|
|
165
|
+
result = SurfaceGreensFunction(
|
|
166
|
+
model, surface=np.eye(3),
|
|
167
|
+
energies=np.linspace(-1, 1, 201),
|
|
168
|
+
k_path=[[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0]],
|
|
169
|
+
k_labels=["M", r"$\Gamma$", "K"],
|
|
170
|
+
N_path=101, thickness=6, NN=5, eps=0.005,
|
|
171
|
+
).run()
|
|
172
|
+
np.savez("surface_gf.npz", **result.as_dict())
|
|
173
|
+
|
|
174
|
+
# 2D Fermi-arc map at one energy
|
|
175
|
+
result = FermiArcMap(
|
|
176
|
+
model, surface=np.eye(3), energy=0.0,
|
|
177
|
+
Nx=50, Ny=50, thickness=6,
|
|
178
|
+
).run()
|
|
179
|
+
result.figure_top_interpolated.savefig("fermi_arc_top.png")
|
|
180
|
+
|
|
181
|
+
# Bulk band structure along a manual k-path
|
|
182
|
+
from tailwater import bulk_band_structure
|
|
183
|
+
fig = bulk_band_structure(
|
|
184
|
+
model,
|
|
185
|
+
k_points = [[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0], [0, 0, 0]],
|
|
186
|
+
k_labels = ["M", r"$\Gamma$", "K", r"$\Gamma$"],
|
|
187
|
+
spacing = 0.01,
|
|
188
|
+
fermi_level = 0.0,
|
|
189
|
+
e_range = (-3, 3),
|
|
190
|
+
)
|
|
191
|
+
fig.savefig("bands.png")
|
|
192
|
+
|
|
193
|
+
# Or use seekpath to auto-determine the high-symmetry path
|
|
194
|
+
from pymatgen.core import Structure
|
|
195
|
+
structure = Structure.from_file("MyMaterial.cif")
|
|
196
|
+
fig = bulk_band_structure(model, auto=True, structure=structure,
|
|
197
|
+
spacing=0.02, e_range=(-3, 3))
|
|
198
|
+
fig.savefig("bands_auto.png")
|
|
199
|
+
```
|
|
200
|
+
|
|
201
|
+
Each post-processing class accepts either an HDF5 path (`str`) or an in-memory `tbmodels.Model`. The `.run()` method returns a typed `Result` dataclass with raw NumPy arrays and matplotlib `Figure` objects.
|
|
202
|
+
|
|
203
|
+
---
|
|
204
|
+
|
|
205
|
+
## API reference (top-level imports)
|
|
206
|
+
|
|
207
|
+
```python
|
|
208
|
+
# HTTP client + HDF5 loader
|
|
209
|
+
tw_api_call(structure, user, password, output_path, filename, ...)
|
|
210
|
+
tb_model.load(path_to_hdf5)
|
|
211
|
+
remaining_credits(user, password)
|
|
212
|
+
|
|
213
|
+
# Heads-only inference model
|
|
214
|
+
HeadsOnly(irreps_in)
|
|
215
|
+
CovariantOnsiteHead(irreps_in)
|
|
216
|
+
CovariantEdgeHead(irreps_in)
|
|
217
|
+
load_heads_only_checkpoint(path)
|
|
218
|
+
save_heads_only_checkpoint(full_state_dict, irreps_in_str, save_path)
|
|
219
|
+
|
|
220
|
+
# Subspace fine-tuning
|
|
221
|
+
subspace_projection(start_lr, end_lr, num_epochs, energy_range,
|
|
222
|
+
decay_sigma, device, save_path,
|
|
223
|
+
embed_path, graph_output_path, loss_mode="subspace")
|
|
224
|
+
|
|
225
|
+
# Subspace losses (advanced)
|
|
226
|
+
Subspace_H_MSE_Loss(gdata, edge_pred, onsite_pred, e_lo, e_hi)
|
|
227
|
+
Subspace_EigLoss(gdata, edge_pred, onsite_pred, kvec, neighbrs, e_lo, e_hi)
|
|
228
|
+
Eigenvalue_Only_Loss(gdata, edge_pred, onsite_pred, e_lo, e_hi)
|
|
229
|
+
make_eigenvalue_only_data(gdata, kvecs, eigs_per_k, e_lo, e_hi)
|
|
230
|
+
build_subspace_active_mask(node_features, onsite_target, e_lo, e_hi)
|
|
231
|
+
write_subspace_basis_file(out_path, active_mask, atoms, LM, ...)
|
|
232
|
+
|
|
233
|
+
# tbmodels.Model assembly from raw head output
|
|
234
|
+
build_hr_model (edge_pred, onsite_pred, gdata, LM, atoms)
|
|
235
|
+
build_hr_model_fast(edge_pred, onsite_pred, gdata, LM, atoms) # vectorized
|
|
236
|
+
write_hr_output(hr_model, out_path, fmt="hdf5"|"hr_dat")
|
|
237
|
+
|
|
238
|
+
# Post-processing calculators (each has a .run() method returning a Result)
|
|
239
|
+
BulkDOS(model_or_path, k_mesh, energies, NC, NV, device)
|
|
240
|
+
SurfaceSpectralDensity(model_or_path, surface, LZ, energies, k_path, ...)
|
|
241
|
+
SurfaceGreensFunction(model_or_path, surface, energies, k_path, thickness, NN, eps, ...)
|
|
242
|
+
FermiArcMap(model_or_path, surface, energy, Nx, Ny, thickness, NN, eps, ...)
|
|
243
|
+
generate_k_path(k_points, N_path, labels=None, rec_vecs=None)
|
|
244
|
+
|
|
245
|
+
# Constants
|
|
246
|
+
NUM_ELEMENTS # 109
|
|
247
|
+
NeighBrs # [17, 3] integer R-vector table
|
|
248
|
+
```
|
|
249
|
+
|
|
250
|
+
---
|
|
251
|
+
|
|
252
|
+
## End-to-end example
|
|
253
|
+
|
|
254
|
+
```python
|
|
255
|
+
import numpy as np
|
|
256
|
+
from pymatgen.core import Structure
|
|
257
|
+
from tailwater import (
|
|
258
|
+
tw_api_call, subspace_projection, tb_model, SurfaceGreensFunction,
|
|
259
|
+
)
|
|
260
|
+
|
|
261
|
+
# 1. Send the structure to the API (one credit, three artifacts)
|
|
262
|
+
structure = Structure.from_file("MyMaterial.cif")
|
|
263
|
+
paths = tw_api_call(
|
|
264
|
+
structure, user="user", password="pw",
|
|
265
|
+
output_path="./outputs", filename="my_mat",
|
|
266
|
+
project=True,
|
|
267
|
+
)
|
|
268
|
+
|
|
269
|
+
# 2. Fine-tune the heads to fit a near-Fermi window
|
|
270
|
+
subspace_projection(
|
|
271
|
+
start_lr=5e-5, end_lr=5e-7, num_epochs=20,
|
|
272
|
+
energy_range=(-2.0, 2.0), decay_sigma=1.0,
|
|
273
|
+
device="cpu",
|
|
274
|
+
save_path="./out_subspace",
|
|
275
|
+
embed_path=paths["embeddings"],
|
|
276
|
+
graph_output_path=paths["graph_output"],
|
|
277
|
+
)
|
|
278
|
+
|
|
279
|
+
# 3. Run surface-GF analysis on the projected hr-model
|
|
280
|
+
model = tb_model.load("./out_subspace/embeddings_pred.hdf5")
|
|
281
|
+
result = SurfaceGreensFunction(
|
|
282
|
+
model, surface=np.eye(3),
|
|
283
|
+
energies=np.linspace(-1, 1, 201),
|
|
284
|
+
k_path=[[0, 0.5, 0], [0, 0, 0], [0.333, 0.333, 0]],
|
|
285
|
+
k_labels=["M", r"$\Gamma$", "K"],
|
|
286
|
+
).run()
|
|
287
|
+
result.figure_top.savefig("surface_top.png")
|
|
288
|
+
```
|
|
289
|
+
|
|
290
|
+
See `examples/` for runnable scripts covering each layer in isolation.
|
|
291
|
+
|
|
292
|
+
---
|
|
293
|
+
|
|
294
|
+
## License
|
|
295
|
+
|
|
296
|
+
Apache 2.0.
|