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.
@@ -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.
@@ -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.
@@ -0,0 +1,7 @@
1
+ include README.md
2
+ include CHANGELOG.md
3
+ include LICENSE
4
+ include pyproject.toml
5
+
6
+ recursive-include examples *.py *.md
7
+ recursive-include tests *.py
@@ -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.