diff --git a/README.md b/README.md index 8886f50..a97d826 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 @@ -48,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 @@ -64,6 +71,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_v1a`; 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 @@ -82,7 +116,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(): @@ -93,9 +127,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"] @@ -133,20 +167,49 @@ 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 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*. @@ -158,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 +Accuracy: 0.76 eV MAE for band gaps (vs 0.38 eV for reference TB-mBJ +DFT), validated on 50 semiconductor/insulator compounds. -![SlakoNet timing](https://github.com/atomgptlab/slakonet/blob/main/slakonet/examples/timing.png) - - -## 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/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" 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..af5466c 100644 --- a/slakonet/ase_calc.py +++ b/slakonet/ase_calc.py @@ -49,14 +49,17 @@ 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) - alpha: float = 0.1 # charge mixing - beta: float = 0.1 # force scaling: F = -beta * dE/dx + 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. + alpha: float = 1.0 + 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 @@ -259,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 @@ -309,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 @@ -336,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)) @@ -397,7 +403,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/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() diff --git a/slakonet/main.py b/slakonet/main.py index 35138f0..ed095da 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 @@ -46,6 +48,58 @@ # 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 + + +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 + + +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. + + 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: @@ -66,11 +120,19 @@ 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, - 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, @@ -95,6 +157,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 @@ -2252,6 +2315,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 @@ -2455,70 +2519,71 @@ 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, ) - # 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 + 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 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 + cell_term = dE_dh[0].transpose(0, 1) @ cell + stress_tensor = (virial + cell_term) / 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) - - # 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 ) 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, @@ -2668,8 +2733,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, @@ -2772,9 +2837,18 @@ 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=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, @@ -2826,8 +2900,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 @@ -2839,6 +2915,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. @@ -2927,8 +3033,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"] = ( @@ -2956,35 +3064,121 @@ 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. - # Create a temporary filtered model wrapper - class FilteredModel: - def __init__(self, filtered_skfs): - self.filtered_skfs = filtered_skfs + 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 - def get_updated_skfs(self): - return self.filtered_skfs + 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) - def to(self, device): - return self + 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 float(self): - return self + 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. - def eval(self): - return self + 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(), + ) + + def get_HS(self, atoms=None): + """k-resolved Hamiltonian and overlap matrices. + + 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``: + + 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 _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 _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]) 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) + filtered_model = _FilteredModel(filtered_skfs) # Run SimpleDftb with filtered model calc = SimpleDftb( diff --git a/slakonet/optim.py b/slakonet/optim.py index d6acf76..dad2fca 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"): @@ -3394,19 +3414,58 @@ 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_v1a") + + +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 +3518,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 +3630,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 +3716,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: diff --git a/slakonet/tests/test_ase_calculators.py b/slakonet/tests/test_ase_calculators.py new file mode 100644 index 0000000..4c6ac77 --- /dev/null +++ b/slakonet/tests/test_ase_calculators.py @@ -0,0 +1,181 @@ +"""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)