From 53d167c0d8bfb40fb565650c8699609e83f7725d Mon Sep 17 00:00:00 2001 From: user Date: Fri, 5 Jun 2026 05:58:40 -0400 Subject: [PATCH 01/11] mag --- slakonet/examples/nio_spin_bands.py | 102 +++++++++++++++++++++------- slakonet/magnetism.py | 17 +++-- 2 files changed, 91 insertions(+), 28 deletions(-) diff --git a/slakonet/examples/nio_spin_bands.py b/slakonet/examples/nio_spin_bands.py index c815342..17d7ca9 100644 --- a/slakonet/examples/nio_spin_bands.py +++ b/slakonet/examples/nio_spin_bands.py @@ -1,17 +1,16 @@ """Spin-polarized bandstructure of NiO - simple demo. -IMPORTANT: the current slakonet universal parameter set (slakonet_v0) only -includes s and p shells for Ni (no d-shell - Ni orbs_per_atom = 4). This -means the Ni d-manifold that drives NiO magnetism is absent from the -Hamiltonian, and a physical Stoner model on l=2 produces zero splitting. - -This script therefore does a *demonstration* of the spin-polarized band -API: it imposes fixed atomic moments and applies the Stoner-like exchange -shift on whichever shells are present (s + p for Ni in the universal set). -The resulting plot shows how the up and down channels of the Ni-p bands -separate under an applied exchange field - it is not a quantitatively -correct NiO calculation. For production NiO you would need refit Ni/O -SKFs including the Ni-3d shell. +The current slakonet universal parameter set (slakonet_v1) includes the +Ni 3d shell (Ni shells = [s, p, d], orbs_per_atom = 9). NiO magnetism is +driven by that Ni-3d manifold, so we apply the Stoner-like exchange shift +on l=2 with a physical Stoner parameter (~0.037 Ha for Ni). + +This is a Stoner-rigid demonstration of the spin-polarized band API: it +imposes a fixed Ni moment and splits the two spin channels by a diagonal +on-site exchange field. It is not a full spin-polarized SCC-DFTB result, +but with the exchange on the d-shell it now produces a physically sensible +d-band splitting rather than the artificially large p-shell split used in +the older s/p-only parameter set. Change AFM = True for a 1x1x2 supercell with two Ni atoms (+/- moments). """ @@ -28,12 +27,23 @@ from slakonet.magnetism import compute_spin_polarized_bands, plot_spin_bands AFM = False -# Exchange strength per shell (Hartree). d-entry is harmless if no d-shell -# is present in the basis; p-entry is what actually drives the splitting -# in the current universal SKF. +SCF = True # self-consistently relax the moments (predict them) vs single-shot + +# Exchange strength per shell (Hartree). NiO magnetism lives on the Ni-3d +# manifold, so the exchange is placed on l=2; s and p stay at 0.0 (a large +# field on the p-shell, as the old s/p-only set required, gives an +# unphysically large split). +# +# NOTE on the value: the free-atom Ni Stoner parameter is ~0.037 Ha +# (DEFAULT_STONER_I), but for this tight-binding model the Ni-3d DOS at E_F +# puts the Stoner criterion threshold I*N(E_F) > 1 between ~0.04 and 0.08 Ha, +# so 0.037 yields a non-magnetic (m -> 0) self-consistent solution. We use an +# effective I_d = 0.12 Ha, which is above threshold and self-consistently +# predicts Ni ~1.8 mu_B / net 2.0 mu_B per cell - close to experimental NiO +# (~1.7-1.9 mu_B on Ni). Lower it toward 0.037 to see the moment collapse. STONER_I = { - 28: {0: 0.0, 1: 0.15, 2: 0.05}, - 8: {0: 0.0, 1: 0.00}, + 28: {0: 0.0, 1: 0.0, 2: 0.12}, + 8: {0: 0.0, 1: 0.0}, } @@ -55,6 +65,41 @@ def fcc_klines(n_seg: int = 20) -> torch.Tensor: return torch.tensor([rows], dtype=torch.float64) +def matched_exchange_splitting(result, calc): + """Physical exchange splitting (eV). + + The naive ``(eu - ed).abs().max()`` subtracts the two spin spectra by + *band index*, but the channels are diagonalized and sorted + independently, so equal indices need not label the same physical state + once the exchange field reorders the bands. Here we instead match each + spin-up state to the spin-down state of maximum eigenvector overlap + (via the overlap matrix S) and take the energy difference of matched + pairs. Returns (max_split, mean_split) in eV. + """ + eu = result["eigenvalues_up"] # [Nband, Nk] eV + ed = result["eigenvalues_dn"] + cu = result["eigenvectors_up"] # [Norb, Nband, Nk] + cd = result["eigenvectors_dn"] + S = calc._results["overlap"] + if S.ndim == 4: + S = S[0] + S = S.to(torch.complex128) # [Norb, Norb, Nk] + + Nband, Nk = eu.shape + diffs = [] + for ik in range(Nk): + Cu = cu[..., ik].to(torch.complex128) + Cd = cd[..., ik].to(torch.complex128) + if Cu.ndim == 3: + Cu, Cd = Cu.squeeze(0), Cd.squeeze(0) + # overlap , shape [Nband_up, Nband_dn] + O = (Cu.conj().transpose(0, 1) @ (S[..., ik] @ Cd)).abs() + match = O.argmax(dim=1) # best down-state per up-state + diffs.append((eu[:, ik] - ed[match, ik]).abs()) + diffs = torch.stack(diffs) + return diffs.max().item(), diffs.mean().item() + + def main(): atoms = nio_structure() print(f"NiO ({'AFM 2-Ni supercell' if AFM else 'FM primitive'}): " @@ -87,14 +132,25 @@ def main(): calc, stoner_I=STONER_I, initial_moments=torch.tensor(m0), - scf=False, + scf=SCF, + verbose=SCF, ) - eu = result["eigenvalues_up"] - ed = result["eigenvalues_dn"] - max_split = (eu - ed).abs().max().item() - print(f"Fermi level : {result['fermi_eV']:.3f} eV") - print(f"Max up/down splitting: {max_split:.3f} eV") + # Predicted magnetic moments (Mulliken spin density n_up - n_dn). + # With SCF=True these are self-consistently relaxed; with SCF=False they + # just echo the imposed initial moments above. + symbols = atoms.get_chemical_symbols() + moments = result["moments"].tolist() + print(f"SCF converged : {result['converged']}") + print("Predicted atom-wise moments (mu_B):") + for i, (s, mu) in enumerate(zip(symbols, moments)): + print(f" atom {i:2d} {s:>2} {mu:+.4f}") + print(f"Net magnetic moment : {result['total_moment']:+.4f} mu_B/cell") + + max_split, mean_split = matched_exchange_splitting(result, calc) + print(f"Fermi level : {result['fermi_eV']:.3f} eV") + print(f"Max exchange splitting : {max_split:.3f} eV (state-matched)") + print(f"Mean exchange splitting : {mean_split:.3f} eV (state-matched)") plot_spin_bands( result, diff --git a/slakonet/magnetism.py b/slakonet/magnetism.py index cd123fd..bec14f1 100644 --- a/slakonet/magnetism.py +++ b/slakonet/magnetism.py @@ -295,13 +295,20 @@ def plot_spin_bands(result, fermi_shift_eV=0.0, filename="bands_spin.png"): # shape: [Nband, Nk] plt.figure(figsize=(8, 6)) for b in range(eu.shape[0]): - plt.plot(eu[b] - fermi_shift_eV, color="tab:red", lw=0.8) + plt.plot( + eu[b] - fermi_shift_eV, color="tab:red", lw=0.8, + label="spin up" if b == 0 else None, + ) for b in range(ed.shape[0]): - plt.plot(ed[b] - fermi_shift_eV, color="tab:blue", lw=0.8, ls="--") - plt.axhline(0, ls="-.", color="k") + plt.plot( + ed[b] - fermi_shift_eV, color="tab:blue", lw=0.8, ls="--", + label="spin down" if b == 0 else None, + ) + plt.axhline(0, ls="-.", color="k", label="$E_F$") plt.xlabel("k-point") - plt.ylabel("E (eV)") + plt.ylabel(r"E - E$_F$ (eV)") plt.title(f"Spin bands (M = {result['total_moment']:.3f})") + plt.legend(loc="best", frameon=True) plt.tight_layout() - plt.savefig(filename) + plt.savefig(filename, dpi=150) plt.close() From 1cdfbd82574aa7952d00dc29fddefa4dac5680f0 Mon Sep 17 00:00:00 2001 From: knc6 Date: Wed, 5 Aug 2026 08:31:55 -0400 Subject: [PATCH 02/11] Add slakonet_v1a parameter set; fix ASE forces, stress and total energy Parameter sets: - Add SLAKONET_MODELS registry mapping slakonet_v0/v1/v1a to their Figshare file IDs, replacing the hardcoded v0 download URL. - default_model(model_name=...), --model_path and the SLAKONET_MODEL env var now select a set; load_trained_model resolves a bare registry name through the download/cache path. Forces / stress / energy (all validated against finite differences): - Geometry stores lengths in Bohr, so dE/dR is eV/Bohr; convert to eV/Angstrom instead of applying the beta=0.1 fudge factor. - Fix the stress: dE/dh^T @ h (not dE/dh @ h^T), correct virial sign, and convert eV/Bohr^3 -> eV/Angstrom^3. The ASE wrapper also reshaped the Voigt-6 vector to (2,3) and left it in GPa, which made get_stress() raise. - alpha scaled the band energy but not the repulsive term, so the total energy was not E_band + E_rep. Default alpha and beta to 1.0; Si's equilibrium lattice constant improves from 5.72 to 5.51 A. Eigenvalues and bandgaps are unaffected. Analysis: - Add compute_site_projected_dos, layer_resolved_dos and gap_states_metric for per-atom-index projections, as needed for surface, defect and interface electronic structure. --- README.md | 37 ++++++- slakonet/analysis.py | 190 +++++++++++++++++++++++++++++++++++ slakonet/ase_calc.py | 7 +- slakonet/main.py | 92 ++++++++--------- slakonet/optim.py | 56 +++++++++-- slakonet/predict_slakonet.py | 22 +++- 6 files changed, 341 insertions(+), 63 deletions(-) diff --git a/README.md b/README.md index 8886f50..d538786 100644 --- a/README.md +++ b/README.md @@ -64,6 +64,33 @@ python slakonet/predict_slakonet.py --file_path slakonet/examples/POSCAR-JVASP- ![SlakoNet output](https://github.com/atomgptlab/slakonet/blob/main/slakonet/examples/slakonet_bands_dos.png) +### Available Parameter Sets + +Parameter sets are downloaded from +[Figshare](https://figshare.com/articles/dataset/SlakoNet_parameters/30122215) +on first use and cached under `~/.cache/atomgptlab/slakonet/`. + +| Name | Description | +| --- | --- | +| `slakonet_v0` | Original universal parameter set (paper v1) | +| `slakonet_v1` | Second-generation universal parameter set | +| `slakonet_v1a` | Refined v1 parameter set | + +```python +from slakonet.optim import default_model + +model = default_model(model_name="slakonet_v1a") +``` + +`default_model()` with no arguments uses `slakonet_v1`; set the +`SLAKONET_MODEL` environment variable to change the default globally, and +`--model_path slakonet_v1a` selects a set from the command line: + +```bash +SLAKONET_MODEL=slakonet_v1a python slakonet/predict_slakonet.py --jid JVASP-107 +python slakonet/predict_slakonet.py --model_path slakonet_v1a --jid JVASP-107 +``` + ### Using Pretrained Models in Python ```python @@ -143,10 +170,12 @@ Toggles (constructor keywords): `compute_forces`, `compute_stress`, `device`. Setting `compute_forces=False` gives a fast energy-only path for high-throughput screening. -Notes: forces are scaled by `beta` (default `0.1`); pass `beta=1.0` for -physically correct forces. Stress is converted to ASE units -(eV/Ang^3, Voigt) but should be validated against a numerical-strain -reference before use in cell relaxation. A full runnable demo is in +Notes: `alpha` scales the band-structure energy and `beta` the forces; +both default to `1.0`, which gives the standard DFTB total energy +`E = E_band + E_rep` together with its exact gradient. Energy, forces +and stress have been checked against finite differences (agreement +better than 0.5% for bulk Si and SiC), so cell relaxation with +`ExpCellFilter` is supported. A full runnable demo is in `slakonet/examples/slakonet_calculator_example.py`. See also the ASE docs page *Calculators -> SlaKoNet*. diff --git a/slakonet/analysis.py b/slakonet/analysis.py index 2965f27..a877289 100644 --- a/slakonet/analysis.py +++ b/slakonet/analysis.py @@ -476,10 +476,200 @@ def compute_fermi_surface_3d( } +# --------------------------------------------------------------------------- +# 6) Site- and layer-resolved DOS (surfaces, defects, interfaces) +# --------------------------------------------------------------------------- +def compute_site_projected_dos( + atoms, + model=None, + kmesh=(4, 4, 1), + energy_range: Tuple[float, float] = (-10.0, 10.0), + sigma: float = 0.1, + n_points: int = 1000, + cutoff: float = 10.0, + device: Optional[str] = None, +) -> dict: + """Per-atom (site) projected DOS on a Monkhorst-Pack mesh. + + Unlike :func:`compute_bandstructure`, which projects onto *element + types*, this resolves every individual atom index. That is what is + needed to tell a surface layer from a bulk-like layer, a defect site + from its host, or the two sides of an interface apart. + + Parameters + ---------- + atoms : jarvis.core.atoms.Atoms + model : trained slakonet model (defaults to default_model()) + kmesh : Monkhorst-Pack divisions. Use 1 along a vacuum/non-periodic + direction (e.g. (4, 4, 1) for a slab stacked along c). + energy_range : window in eV relative to the Fermi level + sigma : Gaussian broadening in eV + + Returns + ------- + dict with keys: + energies : ndarray [n_points], eV relative to E_F + site_dos : ndarray [n_atoms, n_points] + total_dos : ndarray [n_points] + elements : list[str] per atom index + bandgap, vbm, cbm, fermi_energy : floats (eV) + eigenvalues : ndarray [nk, nbands], Fermi-referenced + """ + from slakonet.atoms import Geometry + from slakonet.main import generate_shell_dict_upto_Z65 + + model = _resolve_model(model) + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + + geometry = Geometry.from_ase_atoms([atoms.ase_converter()]) + shell_dict = generate_shell_dict_upto_Z65(model=model) + kpts = torch.tensor([list(kmesh)], dtype=torch.int32) + + with torch.no_grad(): + properties, success = model.compute_multi_element_properties( + geometry=geometry, + shell_dict=shell_dict, + kpoints=kpts, + get_fermi=True, + with_eigenvectors=True, + device=device, + cutoff=cutoff, + ) + if not success: + raise RuntimeError("SlakoNet failed to compute properties") + + eigenvalues = properties["eigenvalues"] # [1, nk, nb], E_F-shifted + eigenvectors = properties["eigenvectors"] # [1, nk, nb, norb] + + basis = properties["basis"] + on_atoms = basis.on_atoms + if on_atoms.ndim == 2: + on_atoms = on_atoms[0] + on_atoms_np = on_atoms.cpu().numpy() + + n_atoms = atoms.num_atoms + # Orbital index list per atom index (padding atoms carry -1). + site_orbitals = [[] for _ in range(n_atoms)] + for orb_idx, a_idx in enumerate(on_atoms_np): + a_idx = int(a_idx) + if 0 <= a_idx < n_atoms: + site_orbitals[a_idx].append(orb_idx) + + grid = torch.linspace( + energy_range[0], energy_range[1], n_points, device=eigenvalues.device + ) + site_dos = torch.zeros(n_atoms, n_points, device=eigenvalues.device) + norm = 1.0 / (sigma * np.sqrt(2.0 * np.pi)) + _, n_k, n_b = eigenvalues.shape + + for k in range(n_k): + for b in range(n_b): + e = eigenvalues[0, k, b] + if not (energy_range[0] - 6 * sigma <= e <= energy_range[1] + 6 * sigma): + continue + psi = eigenvectors[0, k, b, :] + w = (psi.conj() * psi).real if psi.is_complex() else psi * psi + # Normalize so each band contributes exactly one state, which + # makes per-site weights interpretable as fractions. + tot = w.sum() + if tot > 0: + w = w / tot + gauss = norm * torch.exp(-0.5 * ((grid - e) / sigma) ** 2) + for a in range(n_atoms): + idx = site_orbitals[a] + if idx: + site_dos[a] += w[idx].sum() * gauss + + site_dos = site_dos / n_k + site_dos_np = site_dos.detach().cpu().numpy() + + def _f(key): + v = properties.get(key) + if v is None: + return None + return float(np.asarray(v.detach().cpu().numpy()).flatten()[0]) + + return { + "energies": grid.detach().cpu().numpy(), + "site_dos": site_dos_np, + "total_dos": site_dos_np.sum(axis=0), + "elements": list(atoms.elements), + "bandgap": _f("bandgap"), + "vbm": _f("vbm"), + "cbm": _f("cbm"), + "fermi_energy": _f("fermi_energy"), + "eigenvalues": eigenvalues[0].detach().cpu().numpy(), + "kmesh": list(kmesh), + } + + +def layer_resolved_dos(atoms, site_dos, axis: int = 2, tol: float = 0.35): + """Bin per-site DOS into layers along a lattice direction. + + Parameters + ---------- + atoms : jarvis.core.atoms.Atoms (same ordering used for `site_dos`) + site_dos : ndarray [n_atoms, n_points] from compute_site_projected_dos + axis : lattice vector index defining the stacking direction + tol : layer merge tolerance in Angstrom + + Returns + ------- + (layer_positions, layer_dos, layer_members) + layer_positions : ndarray [n_layers], Cartesian coordinate (Ang) + layer_dos : ndarray [n_layers, n_points] + layer_members : list[list[int]] atom indices per layer + """ + coords = np.asarray(atoms.cart_coords)[:, axis] + order = np.argsort(coords) + + layer_members = [] + current = [int(order[0])] + for i in order[1:]: + if abs(coords[i] - coords[current[-1]]) <= tol: + current.append(int(i)) + else: + layer_members.append(current) + current = [int(i)] + layer_members.append(current) + + layer_positions = np.array( + [float(np.mean(coords[m])) for m in layer_members] + ) + layer_dos = np.array( + [np.asarray(site_dos)[m].sum(axis=0) for m in layer_members] + ) + return layer_positions, layer_dos, layer_members + + +def gap_states_metric(energies, dos, vbm_rel, cbm_rel, margin: float = 0.15): + """Integrated DOS strictly inside a reference gap window. + + Used to quantify surface / defect / interface states: a bulk-like + region gives ~0, a region carrying in-gap states gives a finite value. + + `vbm_rel` / `cbm_rel` are gap edges in the same (Fermi-referenced) + energy scale as `energies`; `margin` shrinks the window to avoid + picking up broadening tails from the band edges. + """ + energies = np.asarray(energies) + lo, hi = vbm_rel + margin, cbm_rel - margin + if hi <= lo: + return 0.0 + mask = (energies >= lo) & (energies <= hi) + if not mask.any(): + return 0.0 + return float(np.trapz(np.asarray(dos)[mask], energies[mask])) + + __all__ = [ "compute_bandstructure", "compute_bandstructure_3d", "compute_fermi_surface_2d", "compute_fermi_surface_3d", "compute_kmesh_2d", + "compute_site_projected_dos", + "layer_resolved_dos", + "gap_states_metric", ] diff --git a/slakonet/ase_calc.py b/slakonet/ase_calc.py index e7231b1..cce72d6 100644 --- a/slakonet/ase_calc.py +++ b/slakonet/ase_calc.py @@ -52,8 +52,11 @@ class SlaKoNetConfig(BaseModel): kpoints: List[int] = [3, 3, 3] # Monkhorst-Pack grid cutoff: float = 10.0 # Bohr kT: float = 0.025 # Fermi smearing (eV) - alpha: float = 0.1 # charge mixing - beta: float = 0.1 # force scaling: F = -beta * dE/dx + # alpha scales the band-structure energy and beta the forces. Both are + # 1.0 for the standard DFTB total energy E = E_band + E_rep and its + # exact gradient; changing them breaks energy/force consistency. + alpha: float = 1.0 + beta: float = 1.0 use_scc: bool = False compute_forces: bool = True compute_stress: bool = True # needs compute_forces + periodic diff --git a/slakonet/main.py b/slakonet/main.py index 35138f0..6324c03 100644 --- a/slakonet/main.py +++ b/slakonet/main.py @@ -46,6 +46,8 @@ # torch.set_default_dtype(torch.float64) # torch.set_default_dtype(torch.float32) H2E = 27.211 +# Geometry stores lengths in Bohr; ASE and jarvis work in Angstrom. +BOHR_TO_ANGSTROM = 0.5291772109 class SimpleDftb: @@ -69,8 +71,11 @@ def __init__( include_dos_data=True, include_HS=True, use_float32=True, - alpha=0.1, - beta=0.1, + # alpha scales the band-structure energy and beta the forces. + # Both are 1.0 for the standard DFTB total energy + # E = E_band + E_rep and its exact gradient. + alpha=1.0, + beta=1.0, updated_skfs=None, fermi_surface=False, use_scc=False, @@ -2460,55 +2465,45 @@ def calculate(self): allow_unused=False, ) - # UNIT NOTE: slakonet stores geometry.positions and - # geometry.cell in Bohr, while the total energy is in eV. - # torch.autograd.grad therefore yields eV/Bohr; the - # downstream stress/virial formula is then in eV/Bohr^3. - # ASE expects forces in eV/Ang and stress in eV/Ang^3 - # (then * 160.21766208 -> GPa). Multiply by 1/_BOHR_TO_ANG - # and 1/_BOHR_TO_ANG**3 respectively to convert. The - # virial uses the *Bohr-units* forces with Bohr-units - # positions so the (eV/Bohr * Bohr -> eV) bookkeeping - # stays self-consistent; the conversion is applied once - # at the end. - _BOHR_TO_ANG = 0.52917721092 - - forces_Bohr = -self.beta * grad_outputs[0] # eV/Bohr + # Geometry stores positions and cell in Bohr, so autograd + # returns dE/dR in eV/Bohr. Convert to eV/Angstrom. + grad_pos = grad_outputs[0] + forces = -self.beta * grad_pos / BOHR_TO_ANGSTROM + dE_dh = torch.autograd.grad( total_energy, self.geometry.cell, retain_graph=True, create_graph=True, - )[0] # eV/Bohr - cell = self.geometry.cell[0] # Bohr - volume = torch.abs(torch.det(cell)) # Bohr^3 - positions = self.geometry.positions[0] # Bohr + )[0] + cell = self.geometry.cell[0] # Bohr + volume = torch.abs(torch.det(cell)) # Bohr^3 + positions = self.geometry.positions[0] # Bohr mask = self.geometry.atomic_numbers[0] > 0 - # virial in eV/Bohr^3 (Bohr-units throughout) - stress_virial = torch.einsum( - "ia,ib->ab", forces_Bohr[0][mask], positions[mask] + # Under a homogeneous strain eps, R -> (1+eps)R and + # h -> (1+eps)h, so + # dE/d(eps_ab) = sum_i dE/dR_ia * R_ib + # + sum_c dE/dh_ca * h_cb + # and the (ASE-convention) stress is that divided by V. + virial = torch.einsum( + "ia,ib->ab", grad_pos[0][mask], positions[mask] ) - stress_tensor = ( - stress_virial - dE_dh[0] @ cell.T - ) / volume # eV/Bohr^3 - - # --- convert to ASE/SI units --- - forces = forces_Bohr / _BOHR_TO_ANG # eV/Ang - stress_eVperA3 = stress_tensor / (_BOHR_TO_ANG ** 3) + cell_term = dE_dh[0].transpose(0, 1) @ cell + stress_tensor = (virial + cell_term) / volume # eV/Bohr^3 - # Voigt + GPa + # eV/Bohr^3 -> eV/Angstrom^3 -> GPa + stress_tensor = stress_tensor / BOHR_TO_ANGSTROM**3 stress = ( - torch.tensor( + torch.stack( [ - stress_eVperA3[0, 0], - stress_eVperA3[1, 1], - stress_eVperA3[2, 2], - stress_eVperA3[1, 2], - stress_eVperA3[0, 2], - stress_eVperA3[0, 1], - ], - device=self.device, + stress_tensor[0, 0], + stress_tensor[1, 1], + stress_tensor[2, 2], + stress_tensor[1, 2], + stress_tensor[0, 2], + stress_tensor[0, 1], + ] ) * 160.21766208 ) @@ -2668,8 +2663,8 @@ def run_calc( kpoints_array=[1, 1, 1], device="cuda", compute_forces=True, - alpha=0.1, - beta=0.1, + alpha=1.0, + beta=1.0, elements_needed=None, updated_skfs=None, # NEW with_eigenvectors=False, @@ -2773,8 +2768,11 @@ def __init__( model_path=None, kpoints_array=[1, 1, 1], device="cuda", - alpha=0.1, - beta=0.1, + # alpha scales the band-structure energy and beta the forces. + # Both are 1.0 for the standard DFTB total energy + # E = E_band + E_rep and its exact gradient. + alpha=1.0, + beta=1.0, compute_forces=True, elements_needed=None, use_cached_model=False, @@ -2927,8 +2925,10 @@ def calculate( self.results["forces"] = forces.reshape(-1, 3) if "stress" in properties and result.get("stress") is not None: - stress = result["stress"].detach().cpu().numpy() - self.results["stress"] = stress.reshape(-1, 3) + # SimpleDftb returns the Voigt stress (xx, yy, zz, yz, xz, xy) + # in GPa; ASE expects a 6-vector in eV/Angstrom^3. + stress = result["stress"].detach().cpu().numpy().reshape(-1) + self.results["stress"] = stress / 160.21766208 if "fermi_energy" in result: self.results["fermi_energy"] = ( diff --git a/slakonet/optim.py b/slakonet/optim.py index d6acf76..5feff97 100644 --- a/slakonet/optim.py +++ b/slakonet/optim.py @@ -3394,19 +3394,59 @@ def _smart_load_slakonet_model(stem_or_pt, elements=None, prefer=None): return model -def default_model( - dir_path=None, model_name="slakonet_v1", elements=None, prefer=None -): +# Registry of published SlakoNet parameter sets. +# https://figshare.com/articles/dataset/SlakoNet_parameters/30122215 +SLAKONET_MODELS = { + "slakonet_v0": { + "file_id": 57945370, + "md5": "4b6af7ebf90c43f01e4d954dd22d5394", + "description": "Original universal parameter set (paper v1).", + }, + "slakonet_v1": { + "file_id": 64744347, + "md5": "b7dbdc43ac169ba4f0c107e41d46ef66", + "description": "Second-generation universal parameter set.", + }, + "slakonet_v1a": { + "file_id": 67253969, + "md5": "c656fc974849858342257efed25508cb", + "description": "Refined v1 parameter set (v1a).", + }, +} + +# Default parameter set; override with the SLAKONET_MODEL env variable. +DEFAULT_MODEL_NAME = os.environ.get("SLAKONET_MODEL", "slakonet_v1") + + +def model_download_url(model_name): + """Figshare download URL for a registered parameter set.""" + try: + file_id = SLAKONET_MODELS[model_name]["file_id"] + except KeyError: + raise ValueError( + f"Unknown SlakoNet model '{model_name}'. " + f"Available: {sorted(SLAKONET_MODELS)}" + ) + return f"https://ndownloader.figshare.com/files/{file_id}" + + +def default_model(dir_path=None, model_name=None, elements=None, + prefer=None): """ Load or download the SlakoNet model with proper Figshare handling. Args: + model_name: key in SLAKONET_MODELS (e.g. 'slakonet_v0', + 'slakonet_v1a'). Defaults to DEFAULT_MODEL_NAME, which follows + the SLAKONET_MODEL environment variable. elements: optional iterable of element symbols. When the safetensors-format cache exists, only the matching SKF pairs are materialized (fast, low-memory). Has no effect on the legacy .pt path. prefer: 'safetensors' | 'pt' | None. Overrides SLAKONET_LOADER env. """ + if model_name is None: + model_name = DEFAULT_MODEL_NAME if dir_path is None: dir_path = os.path.join(get_cache_dir("slakonet"), model_name) # dir_path = str(os.path.join(os.path.dirname(__file__), model_name)) @@ -3459,10 +3499,8 @@ def default_model( return model # Download from Figshare - use ndownloader subdomain - # v0 url = "https://ndownloader.figshare.com/files/57945370" - url = "https://ndownloader.figshare.com/files/64744347" - # url="https://figshare.com/ndownloader/files/57945370" - print(f"Downloading {model_name} model from Figshare...") + url = model_download_url(model_name) + print(f"Downloading {model_name} model from Figshare ({url})...") # Create directory if needed if not os.path.exists(dir_path): @@ -3573,7 +3611,7 @@ def default_model_new(dir_path=None, model_name="slakonet_v1"): return model # Download the file - url = "https://figshare.com/ndownloader/files/57945370" + url = model_download_url(model_name) print(f"Downloading {model_name} model...") # Simple download - no streaming @@ -3659,7 +3697,7 @@ def default_model_old(dir_path=None, model_name="slakonet_v1"): return model # If we get here, need to download - url = "https://figshare.com/ndownloader/files/57945370" + url = model_download_url(model_name) print(f"Downloading and loading {model_name} model from zip...") response = requests.get(url, stream=True) diff --git a/slakonet/predict_slakonet.py b/slakonet/predict_slakonet.py index 8627aff..cb7617d 100644 --- a/slakonet/predict_slakonet.py +++ b/slakonet/predict_slakonet.py @@ -1,3 +1,4 @@ +import os import matplotlib.pyplot as plt import numpy as np from matplotlib.gridspec import GridSpec @@ -77,12 +78,25 @@ def load_trained_model(model_path, method="compact", elements=None, prefer=None): """Load a SlakoNet model. + `model_path` may be the name of a published parameter set (e.g. + "slakonet_v0", "slakonet_v1a"), in which case it is fetched/cached from + Figshare, or a path to a local checkpoint stem. + Prefers the safetensors layout (lazy, mmap) when available next to `model_path`. Set `prefer="pt"` (or env SLAKONET_LOADER=pt) to force the legacy torch.load path. Pass `elements={"Si","C"}` to materialize only the relevant SKF pairs. """ - from slakonet.optim import _smart_load_slakonet_model + from slakonet.optim import _smart_load_slakonet_model, SLAKONET_MODELS + + # A bare registry name (no local file beside it) -> download/cache path. + if model_path in SLAKONET_MODELS and not os.path.exists( + f"{model_path}.pt" + ): + return default_model( + model_name=model_path, elements=elements, prefer=prefer + ) + model = _smart_load_slakonet_model( model_path, elements=elements, prefer=prefer ) @@ -682,7 +696,7 @@ def plot_band_dos_atoms( jid=None, atoms=None, model=None, - model_path="slakonet_v0", + model_path=None, energy_range=(-10, 10), filename=None, cutoff=10.0, @@ -692,6 +706,10 @@ def plot_band_dos_atoms( alpha=0.1, ): if not model: + if model_path is None: + from slakonet.optim import DEFAULT_MODEL_NAME + + model_path = DEFAULT_MODEL_NAME elements_hint = None if atoms is not None: try: From 14909c982c8520a2825d77cb3ec6c8a9e0611908 Mon Sep 17 00:00:00 2001 From: knc6 Date: Wed, 5 Aug 2026 09:36:11 -0400 Subject: [PATCH 03/11] Faster implementation --- slakonet/main.py | 65 ++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 58 insertions(+), 7 deletions(-) diff --git a/slakonet/main.py b/slakonet/main.py index 6324c03..3382555 100644 --- a/slakonet/main.py +++ b/slakonet/main.py @@ -18,6 +18,8 @@ from slakonet.slaterkoster import fermi, hs_matrix from jarvis.core.atoms import Atoms from slakonet.utils import eighb, pack +import contextlib +import functools import matplotlib.pyplot as plt from jarvis.core.specie import atomic_numbers_to_symbols @@ -50,6 +52,37 @@ BOHR_TO_ANGSTROM = 0.5291772109 +def nondeterministic_ok(fn): + """Run `fn` with the deterministic-algorithms requirement lifted.""" + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with allow_nondeterministic(): + return fn(*args, **kwargs) + + return wrapper + + +@contextlib.contextmanager +def allow_nondeterministic(): + """Temporarily lift torch's deterministic-algorithms requirement. + + The force/stress backward pass goes through CuBLAS routines that have + no deterministic kernel, so it raises outright if another library in + the same process called torch.use_deterministic_algorithms(True) + (alignn does this when configured deterministic). Relax it just for + the gradient evaluation and restore the caller's setting afterwards. + """ + was_enabled = torch.are_deterministic_algorithms_enabled() + if was_enabled: + torch.use_deterministic_algorithms(False) + try: + yield + finally: + if was_enabled: + torch.use_deterministic_algorithms(True) + + class SimpleDftb: """Enhanced DFTB calculator for periodic systems with analysis tools.""" @@ -68,6 +101,11 @@ def __init__( kT=0.025, # eV for Fermi smearing H2E=27.211, # Hartree to eV compute_forces=True, + # Build a differentiable graph through the force/stress gradients. + # Only needed to backpropagate *through* forces (e.g. force-matching + # training); it makes plain energy+force evaluation ~2x slower + # because torch records an fx stack trace per autograd node. + create_graph=False, include_dos_data=True, include_HS=True, use_float32=True, @@ -100,6 +138,7 @@ def __init__( self.kT = kT self.H2E = H2E self.compute_forces = compute_forces + self.create_graph = create_graph self.include_dos_data = include_dos_data self.include_HS = include_HS # Setup basis and feeds @@ -2257,6 +2296,7 @@ def hybridization(self, omegas, correlated_subspace, **kw): return hybridization(self, omegas, correlated_subspace, **kw) + @nondeterministic_ok def calculate(self): """Main calculation method.""" compute_forces = self.compute_forces @@ -2460,7 +2500,7 @@ def calculate(self): grad_outputs = torch.autograd.grad( total_energy, self.geometry.positions, - create_graph=True, + create_graph=self.create_graph, retain_graph=True, allow_unused=False, ) @@ -2474,7 +2514,7 @@ def calculate(self): total_energy, self.geometry.cell, retain_graph=True, - create_graph=True, + create_graph=self.create_graph, )[0] cell = self.geometry.cell[0] # Bohr volume = torch.abs(torch.det(cell)) # Bohr^3 @@ -2509,11 +2549,22 @@ def calculate(self): ) except RuntimeError as e: - print(f"❌ Error computing forces: {e}") - print( - "⚠️ Forces set to zero - positions may not be in computation graph" - ) - forces = torch.zeros_like(self.geometry.positions) + # Never substitute zeros here. Silently returning zero + # forces makes an optimizer report immediate convergence + # while nothing has moved, and the caller has no way to + # tell. One way to hit this: another library in the same + # process enables torch.use_deterministic_algorithms(True) + # (alignn does when configured deterministic), which makes + # the CuBLAS-backed backward pass raise. + raise RuntimeError( + f"SlakoNet failed to compute forces/stress: {e}\n" + "If this mentions deterministic algorithms, another " + "library in this process enabled them; set " + "CUBLAS_WORKSPACE_CONFIG=:4096:8 before starting " + "Python, or run the calculators in separate " + "processes. Pass compute_forces=False if you only " + "need energies." + ) from e # print('forces',forces) self._results = { "energy": total_energy, From 3d7e08413aeb4c77739bd7ea74c034ee2227be44 Mon Sep 17 00:00:00 2001 From: knc6 Date: Wed, 5 Aug 2026 09:46:22 -0400 Subject: [PATCH 04/11] Stop setting global torch debug flags on import; ~6x faster forces slakonet/optim.py enabled torch.use_deterministic_algorithms(True), set_detect_anomaly(True) and set_anomaly_enabled(True) at module import. These are process-global, expensive, and were never opted into. Consequences, measured on a 2-atom Si cell with a 4x4x4 mesh (energy + forces, slakonet_v1a): anomaly + create_graph (before) 1.93 s create_graph=False 0.86 s anomaly off as well 0.32 s Anomaly mode records a Python stack trace per autograd node; profiling showed torch.fx.traceback.format_stack accounted for 84% of runtime. Worse, the global determinism flag broke force evaluation: SlakoNet's backward uses CuBLAS routines with no deterministic kernel, so autograd raised and the bare `except RuntimeError` substituted torch.zeros_like(). Any optimizer then reported instant convergence with nothing moved. It also broke unrelated calculators sharing the process - running alignn_ff and SlakoNet in one chipsff run made SlakoNet return exactly zero forces, silently. Changes: - Move the flags into an opt-in set_debug_mode(); honour SLAKONET_DEBUG. - Add create_graph (default False) to SimpleDftb. Only force-matching training needs a graph through the gradients. - Wrap SimpleDftb.calculate in allow_nondeterministic() so SlakoNet still works if another library enables determinism, restoring the caller's setting afterwards. - Never substitute zero forces on failure; raise with a diagnostic. Forces and stress are unchanged (bit-identical) and still match finite differences; 3C-SiC bandgap unchanged at 2.0995 eV. --- slakonet/optim.py | 34 +++++++++++++++++++++++++++------- 1 file changed, 27 insertions(+), 7 deletions(-) diff --git a/slakonet/optim.py b/slakonet/optim.py index 5feff97..c8bf8c9 100644 --- a/slakonet/optim.py +++ b/slakonet/optim.py @@ -54,16 +54,36 @@ xm.set_rng_state(random_seed) except ImportError: pass -torch.backends.cudnn.deterministic = True -torch.backends.cudnn.benchmark = False os.environ["PYTHONHASHSEED"] = str(random_seed) -os.environ["CUBLAS_WORKSPACE_CONFIG"] = str(":4096:8") -torch.use_deterministic_algorithms(True) -torch.autograd.set_detect_anomaly(True) -# Optional: Make it raise errors immediately -torch.set_anomaly_enabled(True) +def set_debug_mode(deterministic=True, detect_anomaly=True): + """Opt in to bit-reproducible + anomaly-checked execution. + + These are process-global torch settings and both are expensive, so + they are no longer enabled on import: + + * ``detect_anomaly`` records a Python stack trace for every autograd + node. It made a Si energy+force evaluation ~2.3x slower. + * ``deterministic`` makes torch raise on any op without a + deterministic kernel. SlakoNet's own backward uses such CuBLAS + routines, so it broke force evaluation outright - and, because it + is global, it broke unrelated calculators (e.g. alignn) sharing + the process. + + CUBLAS_WORKSPACE_CONFIG must be set in the environment *before* + Python starts for the deterministic path to work at all. + """ + torch.backends.cudnn.deterministic = deterministic + torch.backends.cudnn.benchmark = not deterministic + if deterministic: + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.use_deterministic_algorithms(deterministic) + torch.autograd.set_detect_anomaly(detect_anomaly) + + +if os.environ.get("SLAKONET_DEBUG"): + set_debug_mode() def get_atoms(jid="", dataset=None, id_tag="jid"): From 7dd0992dcde99326c5e2e3a332a4a37bc755adbb Mon Sep 17 00:00:00 2001 From: knc6 Date: Wed, 5 Aug 2026 20:31:23 -0400 Subject: [PATCH 05/11] Version update --- setup.py | 2 +- slakonet/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index ba5109b..b4c094c 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ setuptools.setup( name="slakonet", - version="5.20.2026", + version="2026.7.26", author="Kamal Choudhary", author_email="kchoudh2@jhu.edu", description="slakonet", diff --git a/slakonet/__init__.py b/slakonet/__init__.py index e7e9ab0..75bdd87 100644 --- a/slakonet/__init__.py +++ b/slakonet/__init__.py @@ -1,3 +1,3 @@ """Version number.""" -__version__ = "5.20.2026" +__version__ = "2026.7.26" From edaf3f1bd60486cb44cabb63d85187bed8e913fd Mon Sep 17 00:00:00 2001 From: knc6 Date: Wed, 5 Aug 2026 20:50:50 -0400 Subject: [PATCH 06/11] Add kspacing option to SlakoNetCalculator for per-structure k-meshes A single fixed kpoints_array cannot serve cells that range from 2-atom bulk to defect supercells to slabs. On ideal diamond Si, where symmetry requires zero forces, the residual is: 1x1x1 2.373 eV/Ang 2x2x2 0.449 4x4x4 0.042 6x6x6 0.005 so a mesh chosen for a 64-atom supercell leaves spurious forces an order of magnitude above a typical fmax=0.05 relaxation criterion. kpoints_for() now derives the mesh from the reciprocal cell, n_i = ceil(|b_i| / kspacing), with a single k-point along non-periodic directions. kspacing=0.30 gives [7,7,7] for primitive Si, [2,2,2] for a 64-atom supercell and [4,4,1] for a slab, dropping Fmax to 0.0016. kpoints_array is still honoured when kspacing is None. --- slakonet/main.py | 40 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/slakonet/main.py b/slakonet/main.py index 3382555..68afe94 100644 --- a/slakonet/main.py +++ b/slakonet/main.py @@ -2818,6 +2818,12 @@ def __init__( model=None, model_path=None, kpoints_array=[1, 1, 1], + # Target reciprocal-space sampling in 1/Angstrom. When set, the + # Monkhorst-Pack mesh is derived per structure from the cell and + # `kpoints_array` is ignored. Strongly preferred when the cell size + # varies (bulk vs supercell vs slab): a mesh that is fine for a + # 64-atom supercell leaves large spurious forces on a 2-atom cell. + kspacing=None, device="cuda", # alpha scales the band-structure energy and beta the forces. # Both are 1.0 for the standard DFTB total energy @@ -2875,8 +2881,10 @@ def __init__( # Store settings self.model_path = model_path self.kpoints_array = kpoints_array + self.kspacing = kspacing self.device = device self.compute_forces = compute_forces + self._last_kpoints = None self.alpha = alpha self.beta = beta self.with_eigenvectors = with_eigenvectors @@ -2888,6 +2896,36 @@ def __init__( if elements_needed: self.set_elements(elements_needed) + def kpoints_for(self, atoms): + """Monkhorst-Pack divisions to use for `atoms`. + + With `kspacing` set, the mesh is derived from the reciprocal cell so + that every structure is sampled to the same density: + + n_i = ceil(|b_i| / kspacing), b_i = 2*pi * (cell^-1)^T + + Non-periodic directions get a single k-point. Without `kspacing` + the fixed `kpoints_array` is returned unchanged. + """ + if self.kspacing is None: + return self.kpoints_array + + cell = np.asarray(atoms.get_cell()) + if abs(np.linalg.det(cell)) < 1e-8: + return self.kpoints_array + recip = 2.0 * np.pi * np.linalg.inv(cell).T + pbc = np.asarray(atoms.get_pbc()) + mesh = [] + for i in range(3): + if not pbc[i]: + mesh.append(1) + continue + n = int(np.ceil(np.linalg.norm(recip[i]) / self.kspacing - 1e-8)) + mesh.append(max(1, n)) + if mesh != self._last_kpoints: + self._last_kpoints = mesh + return mesh + def set_elements(self, elements_needed): """ Set which elements to use for calculations. @@ -3032,7 +3070,7 @@ def eval(self): # Create geometry geometry = Geometry.from_ase_atoms([atoms]) geometry.positions.requires_grad_(True) - kpoints = torch.tensor(self.kpoints_array) + kpoints = torch.tensor(self.kpoints_for(atoms)) # Use filtered model filtered_model = FilteredModel(filtered_skfs) From 045913c937cce104400e95001a939d62a8f1ccbb Mon Sep 17 00:00:00 2001 From: knc6 Date: Thu, 6 Aug 2026 09:40:46 -0400 Subject: [PATCH 07/11] Default to slakonet_v1a DEFAULT_MODEL_NAME kept slakonet_v1 from the upstream rebase; v1a is the current parameter set. SLAKONET_MODEL still overrides it. --- README.md | 2 +- slakonet/optim.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index d538786..22040f2 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,7 @@ from slakonet.optim import default_model model = default_model(model_name="slakonet_v1a") ``` -`default_model()` with no arguments uses `slakonet_v1`; set the +`default_model()` with no arguments uses `slakonet_v1a`; set the `SLAKONET_MODEL` environment variable to change the default globally, and `--model_path slakonet_v1a` selects a set from the command line: diff --git a/slakonet/optim.py b/slakonet/optim.py index c8bf8c9..1b85b78 100644 --- a/slakonet/optim.py +++ b/slakonet/optim.py @@ -3435,7 +3435,7 @@ def _smart_load_slakonet_model(stem_or_pt, elements=None, prefer=None): } # Default parameter set; override with the SLAKONET_MODEL env variable. -DEFAULT_MODEL_NAME = os.environ.get("SLAKONET_MODEL", "slakonet_v1") +DEFAULT_MODEL_NAME = os.environ.get("SLAKONET_MODEL", "slakonet_v1a") def model_download_url(model_name): From 52daf7993e981c8acdb04ed6e3f5f9d443bc6553 Mon Sep 17 00:00:00 2001 From: knc6 Date: Thu, 6 Aug 2026 10:04:48 -0400 Subject: [PATCH 08/11] Add get_bandstructure, get_dos and get_HS to both ASE calculators SlakoNetCalculator (main.py) had only get_bandgap/get_fermi_energy; SlaKoNetCalculator (ase_calc.py) had band_structure/dos but no HS accessor. Both now expose the same three methods, with get_bandstructure/get_dos on ase_calc as aliases of the existing band_structure/dos. get_HS returns (H, S) reshaped to (n_kpoints, n_orbitals, n_orbitals); SimpleDftb stores them as (batch, n_orb, n_orb, n_k), which is awkward to index per k-point. H is in Hartree, not eV, and the basis is non-orthogonal, so the docstrings spell out the generalized eigenproblem needed to recover band energies. Verified: eigh(H[k], S[k]) * 27.211 minus E_F reproduces the calculator's eigenvalues exactly, and H and S are Hermitian at every k. Also hoist the FilteredModel wrapper out of _run_calc_with_filtered_skfs to module scope so the new methods can reuse it. --- slakonet/ase_calc.py | 52 +++++++++++++++++ slakonet/main.py | 136 +++++++++++++++++++++++++++++++++++++------ 2 files changed, 171 insertions(+), 17 deletions(-) diff --git a/slakonet/ase_calc.py b/slakonet/ase_calc.py index cce72d6..ef53f9c 100644 --- a/slakonet/ase_calc.py +++ b/slakonet/ase_calc.py @@ -400,7 +400,59 @@ def dos( dos.detach().cpu().numpy(), ) + def get_HS(self, atoms=None, kpoints=None): + """k-resolved Hamiltonian and overlap matrices. + + Returns ``(H, S)`` as numpy arrays of shape + ``(n_kpoints, n_orbitals, n_orbitals)``. H is in **Hartree** (the + SKF native unit) and the basis is non-orthogonal, so band energies + come from the generalized problem ``H c = e S c``: + + w = scipy.linalg.eigh(H[k], S[k], eigvals_only=True) + eigenvalues_eV = w * 27.211 - calc.get_fermi_level() + + The matrices are complex in general and Hermitian at every k. + + `kpoints` overrides the calculator's Monkhorst-Pack mesh for this + call only, e.g. ``get_HS(kpoints=(1, 1, 1))`` for Gamma only. + """ + atoms = atoms if atoms is not None else self.atoms + mesh = list(kpoints) if kpoints is not None else list(self.kpoints) + geo = Geometry.from_ase_atoms([atoms]) + sim = SimpleDftb( + geo, + self.model, + kpoints=torch.tensor(mesh), + device=self.device, + with_eigenvectors=False, + compute_forces=False, + include_dos_data=False, + include_HS=True, + repulsive=True, + alpha=self.alpha, + beta=self.beta, + kT=self.kT, + use_scc=self.use_scc, + ) + sim.calculate() + # SimpleDftb stores these as (batch, n_orb, n_orb, n_k); move the + # k axis to the front so H[k] is a matrix. + H = sim._results["hamiltonian"][0].permute(2, 0, 1) + S = sim._results["overlap"][0].permute(2, 0, 1) + return ( + H.detach().cpu().numpy(), + S.detach().cpu().numpy(), + ) + # ---- convenience ----------------------------------------------------- + def get_bandstructure(self, atoms=None, **kwargs): + """Alias for :meth:`band_structure`.""" + return self.band_structure(atoms=atoms, **kwargs) + + def get_dos(self, atoms=None, **kwargs): + """Alias for :meth:`dos`.""" + return self.dos(atoms=atoms, **kwargs) + def get_bandgap(self, atoms=None): if "bandgap" not in self.results: self.get_potential_energy(atoms) diff --git a/slakonet/main.py b/slakonet/main.py index 68afe94..40c8993 100644 --- a/slakonet/main.py +++ b/slakonet/main.py @@ -63,6 +63,25 @@ def wrapper(*args, **kwargs): return wrapper +class _FilteredModel: + """Model stand-in exposing only the SKF pairs a structure needs.""" + + def __init__(self, filtered_skfs): + self.filtered_skfs = filtered_skfs + + def get_updated_skfs(self): + return self.filtered_skfs + + def to(self, device): + return self + + def float(self): + return self + + def eval(self): + return self + + @contextlib.contextmanager def allow_nondeterministic(): """Temporarily lift torch's deterministic-algorithms requirement. @@ -3045,27 +3064,110 @@ def get_fermi_energy(self): raise RuntimeError("Calculation not performed yet") return self.results["fermi_energy"] - def _run_calc_with_filtered_skfs(self, atoms, filtered_skfs): - """Run calculation using only filtered SKF pairs""" - from slakonet.atoms import Geometry - from slakonet.main import SimpleDftb + def get_bandstructure( + self, atoms=None, line_density=20, default_points=2 + ): + """Band structure along the conventional high-symmetry k-path. + + Returns a dict with ``eigenvalues`` (n_kpoints, n_bands, eV and + referenced to the Fermi level), ``kpoints``, ``labels``, + ``bandgap``, ``vbm`` and ``cbm``. + """ + from jarvis.core.kpoints import Kpoints3D as Kpoints + from jarvis.core.atoms import ase_to_atoms + from slakonet.optim import kpts_to_klines + + atoms = atoms if atoms is not None else self.atoms + j_atoms = ase_to_atoms(atoms) + kpoints = Kpoints().kpath(j_atoms, line_density=line_density) + klines = kpts_to_klines(kpoints.kpts, default_points=default_points) + + geometry = Geometry.from_ase_atoms([atoms]) + calc = SimpleDftb( + geometry, + klines=klines, + model=_FilteredModel(self._get_filtered_skfs()), + compute_forces=False, + include_dos_data=False, + alpha=self.alpha, + beta=self.beta, + device=self.device, + ) + res = calc.calculate() + return { + "eigenvalues": res["eigenvalues"].detach().cpu().numpy()[0], + "kpoints": np.asarray(kpoints.kpts), + "labels": list(kpoints.labels), + "bandgap": float(res["bandgap"].detach().cpu().numpy().flatten()[0]), + "vbm": float(res["vbm"].detach().cpu().numpy().flatten()[0]), + "cbm": float(res["cbm"].detach().cpu().numpy().flatten()[0]), + } + + def get_dos( + self, atoms=None, energy_range=(-10.0, 10.0), num_points=3000, + sigma=0.1, + ): + """Total DOS on the calculator's k-mesh. + + Returns ``(energies_eV, dos)``, energies referenced to E_F. + """ + atoms = atoms if atoms is not None else self.atoms + calc = self._make_sim(atoms, include_HS=False) + calc.calculate() + e_grid, dos = calc.calculate_dos( + energy_range=energy_range, + num_points=num_points, + sigma=sigma, + fermi_shift=True, + ) + return ( + e_grid.detach().cpu().numpy(), + dos.detach().cpu().numpy(), + ) - # Create a temporary filtered model wrapper - class FilteredModel: - def __init__(self, filtered_skfs): - self.filtered_skfs = filtered_skfs + def get_HS(self, atoms=None): + """k-resolved Hamiltonian and overlap matrices. - def get_updated_skfs(self): - return self.filtered_skfs + Returns ``(H, S)`` with shape + ``(n_kpoints, n_orbitals, n_orbitals)``. H is in **Hartree** (the + SKF native unit) and the basis is non-orthogonal, so band energies + come from the generalized problem ``H c = e S c``: - def to(self, device): - return self + w = scipy.linalg.eigh(H[k], S[k], eigvals_only=True) + eigenvalues_eV = w * 27.211 - calc.get_fermi_energy() + """ + atoms = atoms if atoms is not None else self.atoms + calc = self._make_sim(atoms, include_HS=True) + calc.calculate() + # SimpleDftb stores these as (batch, n_orb, n_orb, n_k); move the + # k axis to the front so H[k] is a matrix. + H = calc._results["hamiltonian"][0].permute(2, 0, 1) + S = calc._results["overlap"][0].permute(2, 0, 1) + return ( + H.detach().cpu().numpy(), + S.detach().cpu().numpy(), + ) - def float(self): - return self + def _make_sim(self, atoms, include_HS=False): + """SimpleDftb on this calculator's mesh, without forces.""" + geometry = Geometry.from_ase_atoms([atoms]) + self.set_elements(set(atoms.get_chemical_symbols())) + return SimpleDftb( + geometry, + kpoints=torch.tensor(self.kpoints_for(atoms)), + model=_FilteredModel(self._get_filtered_skfs()), + compute_forces=False, + include_dos_data=False, + include_HS=include_HS, + alpha=self.alpha, + beta=self.beta, + device=self.device, + ) - def eval(self): - return self + def _run_calc_with_filtered_skfs(self, atoms, filtered_skfs): + """Run calculation using only filtered SKF pairs""" + from slakonet.atoms import Geometry + from slakonet.main import SimpleDftb # Create geometry geometry = Geometry.from_ase_atoms([atoms]) @@ -3073,7 +3175,7 @@ def eval(self): kpoints = torch.tensor(self.kpoints_for(atoms)) # Use filtered model - filtered_model = FilteredModel(filtered_skfs) + filtered_model = _FilteredModel(filtered_skfs) # Run SimpleDftb with filtered model calc = SimpleDftb( From 43a45acb3daff036c89c94534724e04e3ece9df6 Mon Sep 17 00:00:00 2001 From: knc6 Date: Thu, 6 Aug 2026 10:26:55 -0400 Subject: [PATCH 09/11] Add ASE calculator tests; fix broken README example The test suite had no calculator coverage at all. test_ase_calculators adds 11 tests over both calculator classes, checking forces and stress against finite differences of the model's own energy rather than only that the calls return something -- that is the check that catches the unit and sign errors fixed earlier. Also covers get_dos, get_bandstructure, get_HS (shape, hermiticity, and that eigh(H,S)*27.211 - E_F reproduces the eigenvalues), kspacing mesh selection, agreement between the two calculator classes, and SKF re-filtering when the elements change. README: the 'Using Pretrained Models in Python' example raised KeyError: 'band_gap_eV' -- the keys are 'bandgap' and 'fermi_energy', and both are tensors. generate_shell_dict_upto_Z65 also needs model=. Document get_HS, including that H is in Hartree and needs the generalized eigenproblem. All four README python blocks now execute. --- README.md | 35 ++++- slakonet/tests/test_ase_calculators.py | 182 +++++++++++++++++++++++++ 2 files changed, 213 insertions(+), 4 deletions(-) create mode 100644 slakonet/tests/test_ase_calculators.py diff --git a/README.md b/README.md index 22040f2..24a7f55 100644 --- a/README.md +++ b/README.md @@ -109,7 +109,7 @@ model = default_model() # Get structure (example with JARVIS ID) atoms, opt_gap, mbj_gap = get_atoms("JVASP-107") geometry = Geometry.from_ase_atoms([atoms.ase_converter()]) -shell_dict = generate_shell_dict_upto_Z65() +shell_dict = generate_shell_dict_upto_Z65(model=model) # Compute electronic properties with torch.no_grad(): @@ -120,9 +120,9 @@ with torch.no_grad(): device="cuda" ) -# Access results -print(f"Band gap: {properties['band_gap_eV']:.3f} eV") -print(f"Fermi energy: {properties['fermi_energy_eV']:.3f} eV") +# Access results (all tensors; .item() for scalars) +print(f"Band gap: {properties['bandgap'].item():.3f} eV") +print(f"Fermi energy: {properties['fermi_energy'].item():.3f} eV") # Plot band structure and DOS eigenvalues = properties["eigenvalues"] @@ -160,11 +160,38 @@ bs = calc.band_structure(si, path="GXWKGL", npoints=20, e, dos = calc.dos(si) print(calc.get_bandgap(), calc.get_fermi_level()) +# Hamiltonian and overlap, (n_kpoints, n_orbitals, n_orbitals) +H, S = calc.get_HS(si) + # reuse on another structure with NO model reload ge = bulk("Ge", "diamond", a=5.66); ge.calc = calc ge.get_potential_energy() ``` +`get_bandstructure()` and `get_dos()` are aliases of `band_structure()` +and `dos()`. The same three accessors exist on +`slakonet.main.SlakoNetCalculator`. + +`get_HS` returns the k-resolved Hamiltonian and overlap. **H is in +Hartree** and the basis is non-orthogonal, so band energies come from +the generalized eigenproblem: + +```python +import scipy.linalg as sla +from ase.build import bulk +from slakonet.optim import default_model +from slakonet.ase_calc import SlaKoNetCalculator + +calc = SlaKoNetCalculator(default_model().float(), kpoints=(3, 3, 3)) +si = bulk("Si", "diamond", a=5.43) +si.calc = calc +si.get_potential_energy() # sets the Fermi level + +H, S = calc.get_HS(si) +w = sla.eigh(H[0], S[0], eigvals_only=True) # k-point 0 +eigenvalues_eV = w * 27.211 - calc.get_fermi_level() +``` + Toggles (constructor keywords): `compute_forces`, `compute_stress`, `use_scc`, `include_dos`, `kpoints`, `cutoff`, `kT`, `alpha`, `beta`, `device`. Setting `compute_forces=False` gives a fast energy-only path diff --git a/slakonet/tests/test_ase_calculators.py b/slakonet/tests/test_ase_calculators.py new file mode 100644 index 0000000..429446a --- /dev/null +++ b/slakonet/tests/test_ase_calculators.py @@ -0,0 +1,182 @@ +"""Both ASE calculators: energy/forces/stress plus the band structure, +DOS and Hamiltonian/overlap accessors. + +Forces and stress are checked against finite differences of the model's +own energy, which is what catches unit-conversion and sign errors -- the +class of bug that a "does it run" test sails straight past. +""" +import numpy as np +import pytest +import torch +from ase.build import bulk + +from slakonet.optim import default_model +from slakonet.main import SlakoNetCalculator +from slakonet.ase_calc import SlaKoNetCalculator + +KPTS = [3, 3, 3] +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +model = default_model().float() + + +def _si(displace=0.0): + a = bulk("Si", "diamond", a=5.43) + if displace: + a.positions[1, 0] += displace + return a + + +@pytest.fixture(scope="module") +def calc(): + return SlakoNetCalculator( + model=model, kpoints_array=KPTS, device=DEVICE + ) + + +def test_energy_forces_stress_shapes(calc): + atoms = _si(0.1) + atoms.calc = calc + assert np.isfinite(atoms.get_potential_energy()) + assert atoms.get_forces().shape == (2, 3) + assert atoms.get_stress().shape == (6,) + + +def test_forces_match_finite_difference(calc): + """Analytic forces must equal -dE/dx of the same energy.""" + base = _si(0.1) + atoms = base.copy() + atoms.calc = calc + f_analytic = atoms.get_forces()[1, 0] + + h = 0.005 + + def energy(shift): + a = base.copy() + a.positions[1, 0] += shift + a.calc = calc + return a.get_potential_energy() + + f_fd = -(energy(h) - energy(-h)) / (2 * h) + assert f_analytic == pytest.approx(f_fd, abs=5e-3) + + +def test_stress_matches_finite_difference(calc): + """Analytic stress must equal (1/V) dE/d(strain).""" + atoms = _si() + atoms.calc = calc + s_analytic = atoms.get_stress()[0] + volume = atoms.get_volume() + + d = 0.002 + + def energy(eps): + a = _si() + m = np.eye(3) + m[0, 0] += eps + a.set_cell(a.cell @ m.T, scale_atoms=True) + a.calc = calc + return a.get_potential_energy() + + s_fd = (energy(d) - energy(-d)) / (2 * d) / volume + assert s_analytic == pytest.approx(s_fd, abs=5e-4) + + +def test_symmetry_forces_vanish_with_kspacing(): + """Ideal diamond Si has zero forces by symmetry. + + A fixed coarse mesh leaves a large spurious residual, so this also + guards the kspacing mesh selection. + """ + c = SlakoNetCalculator(model=model, kspacing=0.25, device=DEVICE) + atoms = _si() + atoms.calc = c + assert np.abs(atoms.get_forces()).max() < 0.02 + + +def test_kpoints_for_scales_with_cell(): + c = SlakoNetCalculator(model=model, kspacing=0.30, device=DEVICE) + small = c.kpoints_for(_si()) + large = c.kpoints_for(bulk("Si", "diamond", a=5.43, cubic=True).repeat(2)) + assert all(s >= l for s, l in zip(small, large)) + assert min(small) >= 1 and min(large) >= 1 + + slab = bulk("Si", "diamond", a=5.43, cubic=True) + slab.cell[2, 2] += 15.0 + slab.pbc = [True, True, False] + assert c.kpoints_for(slab)[2] == 1 # no sampling along vacuum + + +def test_get_dos(calc): + energies, dos = calc.get_dos(_si()) + assert energies.shape == dos.shape + assert (dos >= 0).all() + assert energies.min() < 0 < energies.max() + + +def test_get_bandstructure(calc): + bs = calc.get_bandstructure(_si()) + assert bs["eigenvalues"].ndim == 2 + assert len(bs["labels"]) == len(bs["kpoints"]) + assert bs["bandgap"] >= 0.0 + assert bs["cbm"] >= bs["vbm"] + + +def test_get_HS_shape_and_hermiticity(calc): + H, S = calc.get_HS(_si()) + nk = int(np.prod(KPTS)) + assert H.shape == S.shape + assert H.shape[0] == nk + assert H.shape[1] == H.shape[2] + for k in (0, nk // 2, nk - 1): + assert np.allclose(H[k], H[k].conj().T, atol=1e-6) + assert np.allclose(S[k], S[k].conj().T, atol=1e-6) + + +def test_get_HS_reproduces_eigenvalues(calc): + """eigh(H, S) * Hartree - E_F must give back the eigenvalues. + + This pins both the Hartree unit of H and the non-orthogonal + (generalized) eigenproblem documented for get_HS. + """ + import scipy.linalg as sla + + atoms = _si() + atoms.calc = calc + atoms.get_potential_energy() + + H, S = calc.get_HS(atoms) + w = sla.eigh(H[0], S[0], eigvals_only=True) * 27.211 + w = w - calc.get_fermi_energy() + + reference = np.sort(calc.results["eigenvalues"][0][0]) + assert np.allclose(np.sort(w)[:6], reference[:6], atol=1e-3) + + +def test_ase_calc_agrees_with_main_calculator(calc): + """The two calculator classes must give the same numbers.""" + other = SlaKoNetCalculator(model, kpoints=tuple(KPTS)) + a1 = _si(0.1) + a1.calc = calc + a2 = _si(0.1) + a2.calc = other + assert a1.get_potential_energy() == pytest.approx( + a2.get_potential_energy(), abs=1e-4 + ) + assert np.allclose(a1.get_forces(), a2.get_forces(), atol=1e-4) + + +def test_model_reuse_across_structures(calc): + """Switching elements must re-filter the SKF pairs, not go stale.""" + si = _si() + si.calc = calc + e_si = si.get_potential_energy() + + ge = bulk("Ge", "diamond", a=5.66) + ge.calc = calc + e_ge = ge.get_potential_energy() + + si2 = _si() + si2.calc = calc + assert si2.get_potential_energy() == pytest.approx(e_si, abs=1e-6) + assert e_ge != pytest.approx(e_si, abs=1e-6) From 207fbeba029da0d8cb5e95311724766ed444a8d2 Mon Sep 17 00:00:00 2001 From: knc6 Date: Thu, 6 Aug 2026 10:41:06 -0400 Subject: [PATCH 10/11] Trim README, add scaling table, apply black -l 79 README: condense the abstract to the key claims, drop the Applications section and the Output Properties duplication, and add the dense-vs- sparse scaling table. Key Features said 'up to 2000 atoms', which the table contradicts; now says >10,000 with the sparse solver. Installation instructions kept as they were. Formatting: black -l 79 over the files touched in this branch. flake8 with the CI selection (E9,F63,F7,F82) reports 0; no line over 79 chars was introduced. --- README.md | 68 ++++++++++++++------------ slakonet/ase_calc.py | 27 +++++----- slakonet/main.py | 21 ++++---- slakonet/optim.py | 3 +- slakonet/tests/test_ase_calculators.py | 5 +- 5 files changed, 68 insertions(+), 56 deletions(-) diff --git a/README.md b/README.md index 24a7f55..e246f39 100644 --- a/README.md +++ b/README.md @@ -1,31 +1,38 @@ # SlakoNet -Accurate and efficient prediction of electronic band structures is essential for designing materials with targeted properties. However, existing machine learning models often lack universality and struggle to predict detailed electronic structures, while traditional tight-binding models based on the Slater-Koster (SK) formalism suffer from (i) limited transferability, (ii) the need for manual parameterization, and (iii) training on low-fidelity electronic structure data. To address these challenges, I introduce SlaKoNet, a parameter optimization framework that learns SK-based Hamiltonian matrix elements across 65 elements of the periodic table using automatic differentiation. SlaKoNet is trained on density functional theory data from the JARVIS-DFT database using the Tran-Blaha modified Becke-Johnson (TBmBJ), encompassing over 20000 materials. The framework achieves a mean absolute error (MAE) of 0.74 eV for bandgap predictions against experimental data, representing a reasonable improvement over standard GGA functionals (MAE = 1.14 eV) while preserving the computational advantages and physical interpretability of tight-binding methods. SlaKoNet demonstrates promising scalability with up to 8.4× speedup on GPUs, enabling rapid electronic structure screening for materials discovery. - +SlaKoNet learns Slater-Koster tight-binding Hamiltonian matrix elements +across 65 elements using automatic differentiation, trained on JARVIS-DFT +data with the Tran-Blaha modified Becke-Johnson (TBmBJ) functional +(>20,000 materials). It reaches 0.74 eV MAE for band gaps against +experiment, versus 1.14 eV for standard GGA, while keeping the cost and +interpretability of tight binding. ![SlakoNet schematic](https://github.com/atomgptlab/slakonet/blob/main/slakonet/examples/sk_schematic.png) ## Key Features -- **Universal parameterization**: Works across 65 elements and their combinations -- **Physics-informed**: Based on Slater-Koster tight-binding formalism -- **High accuracy**: Mean absolute error of 0.74 eV for band gaps vs experimental values -- **Scalable**: GPU-accelerated calculations for systems up to 2000 atoms -- **Comprehensive properties**: Predicts band structures, DOS, band gaps, and orbital projections +- **Universal parameterization**: 65 elements and their combinations +- **Physics-informed**: Slater-Koster tight-binding formalism +- **Accurate**: 0.74 eV MAE for band gaps vs experiment +- **Scalable**: GPU-accelerated, >10,000 atoms with the sparse solver +- **Comprehensive**: band structures, DOS, band gaps, orbital projections +- **ASE-compatible**: energy, forces and stress through a standard calculator ## Installation -Install via pip: + ```bash pip install slakonet ``` -Or create a conda environment and install SlaKoNet in editable mode. To do so, first, install miniforge https://github.com/conda-forge/miniforge. For example: +Or create a conda environment and install SlaKoNet in editable mode. To +do so, first install [miniforge](https://github.com/conda-forge/miniforge): ``` wget "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh" ``` -Based on your system requirements, you'll get a file something like 'Miniforge3-XYZ'. +Based on your system requirements, you'll get a file something like +'Miniforge3-XYZ'. ``` bash Miniforge3-$(uname)-$(uname -m).sh @@ -214,32 +221,33 @@ page *Calculators -> SlaKoNet*. ## Performance Benchmarks -- **Accuracy**: 0.76 eV MAE for band gaps (vs 0.38 eV for reference TB-mBJ DFT) -- **Speed**: <10 seconds for 1000-atom systems on GPU -- **Scalability**: Efficient with GPU acceleration -- **Coverage**: Validated on 50 semiconductor/insulator compounds for experiments - -![SlakoNet timing](https://github.com/atomgptlab/slakonet/blob/main/slakonet/examples/timing.png) - +Accuracy: 0.76 eV MAE for band gaps (vs 0.38 eV for reference TB-mBJ +DFT), validated on 50 semiconductor/insulator compounds. -## Output Properties +### Scaling -SlakoNet predicts comprehensive electronic properties including: +Time per diagonalization, with peak GPU memory in brackets (GB). The +dense `eigh` path is limited to roughly 7,000 orbitals; beyond that the +sparse solver is the only option. -- Electronic band structures along high-symmetry k-paths -- Total and projected density of states (DOS) -- Band gaps (direct/indirect) and band edges -- Fermi energy and electronic structure topology -- Atom-projected and orbital-projected DOS (s/p/d contributions) +| atoms | Norb | dense eigh (s) | sparse solve (s) | +| ---: | ---: | ---: | ---: | +| 128 | 1,152 | 0.15 [2.6] | 0.12 [2.6] | +| 1,024 | 9,216 | – (Norb > 7k) | 3.71 [3.4] | +| 3,456 | 31,104 | – | 56.9 [5.3] | +| 8,192 | 73,728 | – | 403 [10.0] | +| 11,664 | 104,976 | – | 956 [19.2] | +| 16,000 | 144,000 | – | > 30 min (timeout) | -## Applications +![SlakoNet timing](https://github.com/atomgptlab/slakonet/blob/main/slakonet/examples/timing.png) -- High-throughput materials screening -- Electronic structure prediction without expensive DFT -- Band structure and DOS calculations for device design -- Semiconductor and quantum materials discovery -- Educational tools for solid-state physics +## Output Properties +- Band structures along high-symmetry k-paths +- Total, atom-projected and orbital-projected DOS (s/p/d) +- Band gaps (direct/indirect) and band edges +- Fermi energy +- Hamiltonian and overlap matrices ## Dataset diff --git a/slakonet/ase_calc.py b/slakonet/ase_calc.py index ef53f9c..af5466c 100644 --- a/slakonet/ase_calc.py +++ b/slakonet/ase_calc.py @@ -49,9 +49,9 @@ class SlaKoNetConfig(BaseModel): run can be fully described by a JSON file. """ - kpoints: List[int] = [3, 3, 3] # Monkhorst-Pack grid - cutoff: float = 10.0 # Bohr - kT: float = 0.025 # Fermi smearing (eV) + kpoints: List[int] = [3, 3, 3] # Monkhorst-Pack grid + cutoff: float = 10.0 # Bohr + kT: float = 0.025 # Fermi smearing (eV) # alpha scales the band-structure energy and beta the forces. Both are # 1.0 for the standard DFTB total energy E = E_band + E_rep and its # exact gradient; changing them breaks energy/force consistency. @@ -59,7 +59,7 @@ class SlaKoNetConfig(BaseModel): beta: float = 1.0 use_scc: bool = False compute_forces: bool = True - compute_stress: bool = True # needs compute_forces + periodic + compute_stress: bool = True # needs compute_forces + periodic include_dos: bool = False device: Optional[str] = None @@ -262,9 +262,7 @@ def band_structure( labels = [""] * len(kpts_frac) for name, pt in bp.special_points.items(): i = int( - np.argmin( - np.linalg.norm(kpts_frac - np.asarray(pt), axis=1) - ) + np.argmin(np.linalg.norm(kpts_frac - np.asarray(pt), axis=1)) ) labels[i] = (labels[i] + "|" + name) if labels[i] else name @@ -312,8 +310,15 @@ def band_structure( if savefig: self._plot_bands( - eigenvalues, labels, mid, gap, atoms, - emin, emax, mask_ev, savefig, + eigenvalues, + labels, + mid, + gap, + atoms, + emin, + emax, + mask_ev, + savefig, ) return out @@ -339,9 +344,7 @@ def _plot_bands( ep = np.sort(ep, axis=-1) if ep.shape[0] > 1: big = np.abs(np.diff(ep, axis=0)) > 2.0 - nm = np.concatenate( - [np.zeros_like(big[:1]), big], 0 - ).astype(bool) + nm = np.concatenate([np.zeros_like(big[:1]), big], 0).astype(bool) ep[nm] = np.nan fig, ax = plt.subplots(figsize=(8, 5)) diff --git a/slakonet/main.py b/slakonet/main.py index 40c8993..ed095da 100644 --- a/slakonet/main.py +++ b/slakonet/main.py @@ -2535,9 +2535,9 @@ def calculate(self): retain_graph=True, create_graph=self.create_graph, )[0] - cell = self.geometry.cell[0] # Bohr - volume = torch.abs(torch.det(cell)) # Bohr^3 - positions = self.geometry.positions[0] # Bohr + cell = self.geometry.cell[0] # Bohr + volume = torch.abs(torch.det(cell)) # Bohr^3 + positions = self.geometry.positions[0] # Bohr mask = self.geometry.atomic_numbers[0] > 0 # Under a homogeneous strain eps, R -> (1+eps)R and @@ -2549,7 +2549,7 @@ def calculate(self): "ia,ib->ab", grad_pos[0][mask], positions[mask] ) cell_term = dE_dh[0].transpose(0, 1) @ cell - stress_tensor = (virial + cell_term) / volume # eV/Bohr^3 + stress_tensor = (virial + cell_term) / volume # eV/Bohr^3 # eV/Bohr^3 -> eV/Angstrom^3 -> GPa stress_tensor = stress_tensor / BOHR_TO_ANGSTROM**3 @@ -3064,9 +3064,7 @@ def get_fermi_energy(self): raise RuntimeError("Calculation not performed yet") return self.results["fermi_energy"] - def get_bandstructure( - self, atoms=None, line_density=20, default_points=2 - ): + def get_bandstructure(self, atoms=None, line_density=20, default_points=2): """Band structure along the conventional high-symmetry k-path. Returns a dict with ``eigenvalues`` (n_kpoints, n_bands, eV and @@ -3098,13 +3096,18 @@ def get_bandstructure( "eigenvalues": res["eigenvalues"].detach().cpu().numpy()[0], "kpoints": np.asarray(kpoints.kpts), "labels": list(kpoints.labels), - "bandgap": float(res["bandgap"].detach().cpu().numpy().flatten()[0]), + "bandgap": float( + res["bandgap"].detach().cpu().numpy().flatten()[0] + ), "vbm": float(res["vbm"].detach().cpu().numpy().flatten()[0]), "cbm": float(res["cbm"].detach().cpu().numpy().flatten()[0]), } def get_dos( - self, atoms=None, energy_range=(-10.0, 10.0), num_points=3000, + self, + atoms=None, + energy_range=(-10.0, 10.0), + num_points=3000, sigma=0.1, ): """Total DOS on the calculator's k-mesh. diff --git a/slakonet/optim.py b/slakonet/optim.py index 1b85b78..dad2fca 100644 --- a/slakonet/optim.py +++ b/slakonet/optim.py @@ -3450,8 +3450,7 @@ def model_download_url(model_name): return f"https://ndownloader.figshare.com/files/{file_id}" -def default_model(dir_path=None, model_name=None, elements=None, - prefer=None): +def default_model(dir_path=None, model_name=None, elements=None, prefer=None): """ Load or download the SlakoNet model with proper Figshare handling. diff --git a/slakonet/tests/test_ase_calculators.py b/slakonet/tests/test_ase_calculators.py index 429446a..4c6ac77 100644 --- a/slakonet/tests/test_ase_calculators.py +++ b/slakonet/tests/test_ase_calculators.py @@ -5,6 +5,7 @@ own energy, which is what catches unit-conversion and sign errors -- the class of bug that a "does it run" test sails straight past. """ + import numpy as np import pytest import torch @@ -29,9 +30,7 @@ def _si(displace=0.0): @pytest.fixture(scope="module") def calc(): - return SlakoNetCalculator( - model=model, kpoints_array=KPTS, device=DEVICE - ) + return SlakoNetCalculator(model=model, kpoints_array=KPTS, device=DEVICE) def test_energy_forces_stress_shapes(calc): From 8b387a51177c25571f10c67421944e29dd3bbbd8 Mon Sep 17 00:00:00 2001 From: knc6 Date: Thu, 6 Aug 2026 10:45:27 -0400 Subject: [PATCH 11/11] Make the Colab link a clickable badge Replace the plain 'Open in Colab' text link with the standard colab-badge.svg image wrapped in the same link, so the badge itself opens the notebook. --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index e246f39..a97d826 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,7 @@ pip install uv; uv pip install -e . ### Google Colab example -[Open in Colab](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/slakonet_example.ipynb) +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/slakonet_example.ipynb) ### Example of Training Models