feat: pbc.lo - native Maximally Localized Wannier Functions

[ethanvo]

feat: pbc.lo — native Maximally Localized Wannier Functions

Adds a pure-Python MLWF implementation (NumPy/SciPy) as
pyscf.pbc.lo.kernel(kmf, proj_guess, ...). Takes a converged k-point
mean-field object and returns the Wannier gauge, centers, spreads, and
the Ω_I + Ω_OD + Ω_D decomposition. Algorithm is MV 1997 (PRB 56,
12847) plus SMV 2001 (PRB 65, 035109) for entangled bands. No
libwannier90 / Fortran dependency — everything flows from kmf.mo_coeff,
no file I/O at runtime.

Usage

from pyscf.pbc.lo import kernel

proj = [
    {'center': (x, y, z), 'l': 0, 'm': 0, 'zeta': 1.0},
    # hybrids: {'center': ..., 'zeta': ..., 'components': [(c, l, m), ...]}
]
U, centers, spreads, omega_i, omega_d, omega_od, converged = kernel(
    kmf, proj, dis_win=(emin, emax), conv_tol=1e-10)

Projection dicts use standard (l, m) (no Wannier90 m_r / axis
conventions to juggle). Supports s/p/d and arbitrary hybrids. Centers
Bohr, spreads Bohr², energy windows Hartree. Collinear only — call
twice with spin=0 / spin=1 for UHF.

Layout

pyscf/pbc/lo/
  mlwf.py         # kernel
  bvectors.py     # find_bvectors (B1 shell search)
  overlap.py      # compute_mmn via ft_aopair
  projection.py   # compute_amn + lowdin_orthonormalize
  spread.py       # Ω decomposition
  wannierise.py   # MV steepest descent
  disentangle.py  # SMV iteration

Validation

42 unit tests covering each stage of the pipeline, plus 2 optional
end-to-end tests that drive a compiled wannier90.x on inputs emitted
from our side (.mmn/.amn/.eig/.win) and compare final
centers/spreads. Enable with W90_EXE=/path/to/wannier90.x; skipped
otherwise.

2×2×2 H₂ crystal (8 Bohr cubic, STO-3G):

Case	Quantity	wannier90.x v4	pyscf.pbc.lo
Isolated 2→2	Total Ω	0.801808 Ų	0.801808 Ų
Isolated 2→2	Centers	(−0.1057, 0, 0), (0.8995, 0, 0) Å	match to 6 dp
Entangled 2→1	Total Ω	0.560957 Ų	0.560957 Ų
Entangled 2→1	Center	(0.3969, 0, 0) Å	match to 6 dp

Agreement to 6–7 decimals on both paths.

Running tests

$ python -m pytest pyscf/pbc/lo -q
42 passed, 2 skipped in ~30s

$ W90_EXE=/path/to/wannier90.x python -m pytest pyscf/pbc/lo -q
44 passed in ~30s

Not included / known limits

• dis_froz (frozen inner window) — raises NotImplementedError.
• dis_win assumes a uniform band-index set across k (fine for gapped
  windows; partial metals would need ragged arrays).
• Steepest descent only — no CG or line-search step-size adaptation.
• No postw90 features (band interpolation, DOS, Berry phase, etc.).

Test plan

• pip install -e . from repo root
• python -m pytest pyscf/pbc/lo -q → 42 passed, 2 skipped
• from pyscf.pbc.lo import kernel imports
• Optional: set W90_EXE → 44 passed

[tberkel]

Thanks for this PR! Native MLWFs in PySCF will be great to have. A few quick questions, though I admit I haven't read the code super closely.

1. Does this work with both FFTDF and GDF? In principle, it shouldn't matter, but since integration is done on a real-space grid, I'm not sure if FFTDF and GDF have comparable grids.
2. MLWFs are "just" Boys orbitals, and molecular PySCF has Boys localization. This doesn't seem to reuse any of those functions. Thoughts?

[ethanvo]

Thanks for the review!

FFTDF vs GDF. Works the same with both, the pipeline never touches kmf.with_df. compute_mmn calls ft_aopair (which builds its own machinery from the Cell) and compute_amn builds a fresh gen_grid.Grids. Whatever DF you ran the SCF with stays inside the SCF.

lo.boys. Boys minimizes ⟨r²⟩ − |⟨r⟩|² with real position-operator integrals; in PBC the position operator isn't well-defined, so MV replaces it with the finite-difference expression in M_mn(k, b) that couples k-points. Boys' CIAH-driven kernel(localizer, mo_coeff) is set up for one fixed-MO unitary, not per-k unitaries with cross-k connectivity. The CIAH solver itself is reusable — swapping it in for the steepest descent in wannierise.py is a clean follow-up. I kept this PR scoped to MV/SMV so it's easy to validate against wannier90.x.