From 8fb5eb1a5d4056a9ef6a8ad83acb4b9840abbd90 Mon Sep 17 00:00:00 2001 From: wtpeter Date: Sat, 12 Sep 2026 17:09:14 +0800 Subject: [PATCH 1/3] feat(soscf): add SGM optimizer for ROHF and ROKS Add analytic square-gradient minimization, SOMO subspace diagnostics, and directional curvature analysis with regression tests and a triplet example. Preserve the energy-gap preconditioner as the baseline before adding alternative preconditioners. --- examples/soscf/01_sgm.py | 76 +++++ src/nest/soscf/sgm.py | 529 +++++++++++++++++++++++++++++++ src/nest/soscf/tests/test_sgm.py | 270 ++++++++++++++++ 3 files changed, 875 insertions(+) create mode 100644 examples/soscf/01_sgm.py create mode 100644 src/nest/soscf/sgm.py create mode 100644 src/nest/soscf/tests/test_sgm.py diff --git a/examples/soscf/01_sgm.py b/examples/soscf/01_sgm.py new file mode 100644 index 0000000..35fa6ae --- /dev/null +++ b/examples/soscf/01_sgm.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +'''Run a restricted open-shell Kohn-Sham (ROKS) calculation using the SGM''' +# Q-Chem reference input for this CH4 core excitation benchmark: +# +# $rem +# METHOD BHHLYP +# BASIS aug-cc-pvtz +# SYMMETRY false +# SYM_IGNORE true +# NO_REORIENT True +# XC_GRID 000075000302 +# UNRESTRICTED false +# delta_scf true +# $end +# +# $delta_scf +# triplet restricted +# triplet_SCF_algorithm SGM +# somo_1 1 +# somo_2 6 +# $end +# +# Q-Chem output: +# Restricted open-shell triplet state = -29.947480 Ha +# SOMO(1): best initial orbital 5 overlap 0.999002 +# SOMO(2): best initial orbital 6 overlap 0.879378 +# Excitation energy = 10.554376 Ha / 287.199215 eV + +from pyscf import gto +from nest.soscf import sgm + + +atom = ''' +C 0.00000000 0.00000000 0.00000000 +H 0.62760000 0.62760000 0.62760000 +H 0.62760000 -0.62760000 -0.62760000 +H -0.62760000 0.62760000 -0.62760000 +H -0.62760000 -0.62760000 0.62760000 +''' + +mol = gto.M(atom=atom, charge=0, spin=0, basis='aug-cc-pvtz') +mf = mol.RKS(xc='BHandHLYP') +mf.grids.atom_grid = (75, 302) +mf.kernel() + +setocc = mf.to_uks().mo_occ +setocc[1][0] -= 1 +setocc[0][5] += 1 +ro_occ = setocc[0] + setocc[1] + +mol1 = gto.M(atom=atom, charge=0, spin=2, basis='aug-cc-pvtz') +mf1 = mol1.ROKS(xc='BHandHLYP') +mf1.grids.atom_grid = (75, 302) +mf1.mo_coeff = mf.mo_coeff +mf1.mo_occ = ro_occ + +sgm_mf = mf1.SGM() +sgm_mf.verbose = 4 +sgm_mf.kernel() + +print('SGM converged =', sgm_mf.converged) +print('SGM energy = %.12f Ha' % sgm_mf.e_tot) diff --git a/src/nest/soscf/sgm.py b/src/nest/soscf/sgm.py new file mode 100644 index 0000000..be9309d --- /dev/null +++ b/src/nest/soscf/sgm.py @@ -0,0 +1,529 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Author: Tai Wang & Codex +# Ref: JCTC 2020, 16, 1699 + +from functools import reduce + +import numpy +import scipy.linalg + +from pyscf import __config__, lib +from pyscf.lib import logger +from pyscf.scf import hf, hf_symm + + +_ARMIJO_C1 = 1e-4 +_MIN_ALPHA = 1e-12 +_CURVATURE_TOL = 1e-10 +_PRECOND_FLOOR = 1e-12 +_MAX_PRECOND = 1e4 + + +def gen_g_hop_rohf(mf, mo_coeff, mo_occ, fock_ao=None, h1e=None, + with_symmetry=True): + '''ROHF orbital gradient and full-K Hessian-vector product.''' + mol = mf.mol + mo_occ = numpy.asarray(mo_occ) + if h1e is None: + h1e = mf.get_hcore(mol) + if fock_ao is None or getattr(fock_ao, 'focka', None) is None: + dm0 = mf.make_rdm1(mo_coeff, mo_occ) + vhf = mf.get_veff(mol, dm0) + focka_ao = h1e + vhf[0] + fockb_ao = h1e + vhf[1] + else: + focka_ao, fockb_ao = fock_ao.focka, fock_ao.fockb + + focka = reduce(numpy.dot, (mo_coeff.conj().T, focka_ao, mo_coeff)) + fockb = reduce(numpy.dot, (mo_coeff.conj().T, fockb_ao, mo_coeff)) + occidxa = mo_occ > 0 + occidxb = mo_occ == 2 + viridxa = ~occidxa + viridxb = ~occidxb + uniq_var_a = viridxa[:, None] & occidxa + uniq_var_b = viridxb[:, None] & occidxb + uniq_var = uniq_var_a | uniq_var_b + + gmat = numpy.zeros_like(focka) + gmat[uniq_var_a] = focka[uniq_var_a] + gmat[uniq_var_b] += fockb[uniq_var_b] + g = gmat[uniq_var] + + focka_diag = focka.diagonal().real + fockb_diag = fockb.diagonal().real + h_diag_mat = numpy.zeros_like(focka_diag[:, None] - focka_diag) + h_diag_mat[uniq_var_a] = ( + focka_diag[:, None] - focka_diag)[uniq_var_a] + h_diag_mat[uniq_var_b] += ( + fockb_diag[:, None] - fockb_diag)[uniq_var_b] + h_diag = h_diag_mat[uniq_var] + + if with_symmetry and mol.symmetry: + orbsym = hf_symm.get_orbsym(mol, mo_coeff) + sym_forbid = (orbsym[:, None] != orbsym)[uniq_var] + g = g.copy() + h_diag = h_diag.copy() + g[sym_forbid] = 0 + h_diag[sym_forbid] = 0 + + mo_occ_a = occidxa.astype(numpy.double) + mo_occ_b = occidxb.astype(numpy.double) + vind = mf.gen_response((mo_coeff, mo_coeff), (mo_occ_a, mo_occ_b), + hermi=1, with_nlc=False) + + def h_op(x): + if with_symmetry and mol.symmetry: + x = x.copy() + x[sym_forbid] = 0 + kappa = hf.unpack_uniq_var(x, mo_occ) + + dm1a_mo = kappa * (mo_occ_a[None, :] - mo_occ_a[:, None]) + dm1b_mo = kappa * (mo_occ_b[None, :] - mo_occ_b[:, None]) + dm1 = numpy.asarray(( + reduce(numpy.dot, (mo_coeff, dm1a_mo, mo_coeff.conj().T)), + reduce(numpy.dot, (mo_coeff, dm1b_mo, mo_coeff.conj().T)))) + v1a, v1b = vind(dm1) + + hmat_a = focka.dot(kappa) - kappa.dot(focka) + hmat_b = fockb.dot(kappa) - kappa.dot(fockb) + hmat_a += reduce(numpy.dot, (mo_coeff.conj().T, v1a, mo_coeff)) + hmat_b += reduce(numpy.dot, (mo_coeff.conj().T, v1b, mo_coeff)) + + out = numpy.zeros_like(focka) + out[uniq_var_a] = hmat_a[uniq_var_a] + out[uniq_var_b] += hmat_b[uniq_var_b] + out = out[uniq_var] + if with_symmetry and mol.symmetry: + out = out.copy() + out[sym_forbid] = 0 + return out + + return g, h_op, h_diag + + +def gen_delta_curvature_rohf(mf, mo_coeff, mo_occ, with_symmetry=True): + '''Analytic directional curvature of Delta = g.T@g for real RO orbitals. + + Returns a callable curvature(x) -> (exact, gauss_newton), evaluated along + C(t) = C exp(t K(x)). The two values are 2*|g'|**2 + 2*g.T@g'' and + 2*|g'|**2, respectively. They coincide at an orbital stationary point. + + Supports ROHF and LDA/GGA ROKS (including hybrid exchange); NLC response + is omitted as in gen_g_hop_rohf. DFT requires third XC derivatives. + This is a diagnostic: one call computes two density responses and, for + DFT, a kxc contraction. It is not used by the SGM optimizer. + ''' + if numpy.iscomplexobj(mo_coeff): + raise NotImplementedError('Delta curvature requires real orbitals') + mol = mf.mol + c = numpy.asarray(mo_coeff) + mo_occ = numpy.asarray(mo_occ) + occ = numpy.asarray((mo_occ > 0, mo_occ == 2), dtype=float) + masks = (occ[:, :, None] == 0) & (occ[:, None, :] > 0) + unique = masks[0] | masks[1] + allowed = numpy.ones(numpy.count_nonzero(unique), dtype=bool) + if with_symmetry and mol.symmetry: + orbsym = hf_symm.get_orbsym(mol, mo_coeff) + allowed = (orbsym[:, None] == orbsym)[unique] + + def pack(f): + mat = numpy.where(masks[0], f[0], 0) + numpy.where(masks[1], f[1], 0) + return mat[unique] * allowed + + dm0 = mf.make_rdm1(mo_coeff, mo_occ) + vhf = mf.get_veff(mol, dm0) + fock = c.T @ (mf.get_hcore() + vhf) @ c + g = pack(fock) + vind = mf.gen_response((mo_coeff, mo_coeff), tuple(occ), hermi=1, with_nlc=False) + xctype = 'HF' + if isinstance(mf, hf.KohnShamDFT): + ni = mf._numint + xctype = ni._xc_type(mf.xc) + if xctype not in ('HF', 'LDA', 'GGA'): + raise NotImplementedError('Delta curvature supports HF, LDA and GGA') + ni.libxc.test_deriv_order(mf.xc, 3, raise_error=True) + + def curvature(x): + kappa = hf.unpack_uniq_var(numpy.asarray(x) * allowed, mo_occ) + # D(t) = C exp(tK) N exp(-tK) C.T. + dm1_mo = kappa * (occ[:, None, :] - occ[:, :, None]) + dm2_mo = kappa @ dm1_mo - dm1_mo @ kappa + dm1 = c @ dm1_mo @ c.T + dm2 = c @ dm2_mo @ c.T + v1 = vind(dm1) + v2 = vind(dm2) + + if xctype != 'HF': + # F'' = response(D'') + kxc[D', D']. eval_xc_eff differentiates + # with respect to rho and its Cartesian gradients, not sigma. + for ao, mask, weight, _ in ni.block_loop(mol, mf.grids, c.shape[0], 1): + ao_rho = ao[0] if xctype == 'LDA' else ao + rho0 = numpy.asarray([ni.eval_rho(mol, ao_rho, dm, mask, xctype, hermi=1) + for dm in dm0]) + rho1 = numpy.asarray([ni.eval_rho(mol, ao_rho, dm, mask, xctype, hermi=1) + for dm in dm1]) + kxc = ni.eval_xc_eff(mf.xc, rho0, deriv=3, xctype=xctype, spin=1)[3] + if xctype == 'LDA': + rho1 = rho1[:, None, :] + wv = numpy.einsum('axbyczg,byg,czg->axg', kxc, rho1, rho1, optimize=True) + wv *= weight + for spin in range(2): + if xctype == 'LDA': + v2[spin] += ao[0].T @ (wv[spin, 0, :, None] * ao[0]) + else: + wv[spin, 0] *= .5 + aow = numpy.einsum('xgi,xg->gi', ao, wv[spin]) + mat = ao[0].T @ aow + v2[spin] += mat + mat.T + + v1_mo = c.T @ v1 @ c + comm = fock @ kappa - kappa @ fock + g1 = pack(comm + v1_mo) + g2 = pack(comm @ kappa - kappa @ comm + + 2*(v1_mo @ kappa - kappa @ v1_mo) + c.T @ v2 @ c) + gauss_newton = 2 * numpy.dot(g1, g1) + return gauss_newton + 2*numpy.dot(g, g2), gauss_newton + + return curvature + + +class LBFGSHistory: + '''Bounded L-BFGS history for inverse-Hessian two-loop recursion.''' + + def __init__(self, max_size=8, min_curvature=1e-10): + self.max_size = max_size + self.min_curvature = min_curvature + self.s_list = [] + self.y_list = [] + self.rho_list = [] + + def __len__(self): + return len(self.s_list) + + def clear(self): + self.s_list.clear() + self.y_list.clear() + self.rho_list.clear() + + def push(self, step, grad_diff): + sy = numpy.dot(step, grad_diff) + # A relative test retains valid secant pairs as the gradients shrink. + scale = numpy.linalg.norm(step) * numpy.linalg.norm(grad_diff) + if sy <= self.min_curvature * scale: + return False, sy + + self.s_list.append(step.copy()) + self.y_list.append(grad_diff.copy()) + self.rho_list.append(1.0 / sy) + + if len(self.s_list) > self.max_size: + self.s_list.pop(0) + self.y_list.pop(0) + self.rho_list.pop(0) + + return True, sy + + def direction(self, grad, h0_inv): + '''Return p = -H_k^{-1} grad.''' + if not self.s_list: + return -(h0_inv * grad) + + nvec = len(self.s_list) + alpha = numpy.empty(nvec) + q = grad.copy() + + for i in range(nvec - 1, -1, -1): + alpha[i] = self.rho_list[i] * numpy.dot(self.s_list[i], q) + q -= alpha[i] * self.y_list[i] + + z = h0_inv * q + + for i in range(nvec): + beta = self.rho_list[i] * numpy.dot(self.y_list[i], z) + z += (alpha[i] - beta) * self.s_list[i] + + return -z + + +def _somo_overlaps(mo_ref, mo_coeff, mo_occ, s1e): + '''Cosines of principal angles between the two SOMO subspaces.''' + somo = numpy.asarray(mo_occ) == 1 + overlap = reduce(numpy.dot, (mo_ref[:, somo].conj().T, s1e, + mo_coeff[:, somo])) + return numpy.linalg.svd(overlap, compute_uv=False) + + +class SGM(lib.StreamObject): + '''Square Gradient Minimization mixin for ROHF-like SCF objects. + + Important parameters: + tol + Convergence threshold on sqrt(Delta), where Delta = g_orb.T@g_orb. + An explicitly set PySCF conv_tol_grad takes precedence over tol. + gradient_scale + Scalar c from the SGM paper. Q-Chem's DeltaSCF driver uses 0.75 + by default. It scales the gradient seen by L-BFGS and the L-BFGS + y_k history; the line search still uses the true Delta directional + derivative. + somo_overlap_tol + Warn if the smallest SOMO overlap singular value falls below + this threshold relative to the initial or previous orbitals. + Default 0.7 corresponds to a largest principal angle of about + 46 degrees. This is a diagnostic, not a state constraint. + + The returned object also inherits the input ROHF/ROKS class. Use + mf.SGM().set(...).run(mo_coeff, mo_occ), or SGM(mf).kernel(...). + NLC response is omitted. Only real orbital rotations are supported. + ''' + + __name_mixin__ = 'SGM' + + max_cycle = getattr(__config__, 'sgm_max_cycle', 200) + tol = getattr(__config__, 'sgm_tol', 1e-4) + lbfgs_memory = getattr(__config__, 'sgm_lbfgs_memory', 8) + gradient_scale = getattr(__config__, 'sgm_gradient_scale', 0.75) + somo_overlap_tol = getattr(__config__, 'sgm_somo_overlap_tol', 0.7) + canonicalization = getattr(__config__, + 'soscf_newton_ah_SOSCF_canonicalization', True) + + _keys = {'max_cycle', 'tol', 'lbfgs_memory', 'gradient_scale', + 'canonicalization', 'somo_overlap_tol'} + + gen_g_hop = staticmethod(gen_g_hop_rohf) + + def __new__(cls, mf): + if isinstance(mf, SGM): + return mf + assert isinstance(mf, hf.SCF) + if not mf.istype('ROHF'): + raise NotImplementedError('SGM currently supports ROHF/ROKS objects') + obj = object.__new__(cls) + return lib.set_class(obj, (cls, mf.__class__)) + + def __init__(self, mf): + if mf is self or isinstance(mf, SGM): + return + self.__dict__.update(mf.__dict__) + self._scf = mf + + def dump_flags(self, verbose=None): + super().dump_flags(verbose) + log = logger.new_logger(self, verbose) + log.info('SGM gradient tolerance = %g', + self.tol if self.conv_tol_grad is None else self.conv_tol_grad) + log.info('SGM gradient scale = %g', self.gradient_scale) + log.info('SGM L-BFGS memory = %d', self.lbfgs_memory) + log.info('SGM SOMO overlap warning threshold = %g', self.somo_overlap_tol) + log.info('SGM canonicalization = %s', self.canonicalization) + return self + + def _preconditioner(self, h_diag): + # Diagonal, frozen-Fock approximation to 2*J.T@J for Delta = g.T@g. + # J = d g / d kappa; this is not the full Delta Hessian away from g=0. + denom = numpy.maximum(2.0 * h_diag ** 2, _PRECOND_FLOOR) + return numpy.minimum(1.0 / denom, _MAX_PRECOND) + + def _trial_mo(self, mo_coeff, mo_occ, step): + kappa = hf.unpack_uniq_var(step, mo_occ) + mo_trial = numpy.dot(mo_coeff, scipy.linalg.expm(kappa)) + if self.mol.symmetry: + orbsym = hf_symm.get_orbsym(self.mol, mo_coeff) + mo_trial = lib.tag_array(mo_trial, orbsym=orbsym) + return mo_trial + + def _exact_sgm_state(self, mo_coeff, mo_occ, fock_ao=None, h1e=None): + g_orb, h_op, h_diag = self.gen_g_hop(self, mo_coeff, mo_occ, fock_ao, h1e) + delta = numpy.dot(g_orb, g_orb) + grad_delta = 2.0 * h_op(g_orb) + return h_diag, delta, grad_delta + + def _line_search(self, mo_coeff, mo_occ, direction, delta, grad_delta, + h1e, s1e): + alpha = 1.0 + dir_deriv = numpy.dot(direction, grad_delta) + + while alpha >= _MIN_ALPHA: + step = alpha * direction + mo_trial = self._trial_mo(mo_coeff, mo_occ, step) + + dm_trial = self.make_rdm1(mo_trial, mo_occ) + vhf_trial = self.get_veff(self.mol, dm_trial) + fock_trial = self.get_fock(h1e, s1e, vhf_trial, dm_trial) + g_trial = self.get_grad(mo_trial, mo_occ, fock_trial) + delta_trial = numpy.dot(g_trial, g_trial) + + if delta_trial <= delta + _ARMIJO_C1 * alpha * dir_deriv: + return alpha, step, mo_trial, fock_trial, vhf_trial + + alpha *= 0.5 + + return None + + def _descent_direction(self, history, grad_delta, opt_grad_delta, h0_inv, + log, cycle): + direction = history.direction(opt_grad_delta, h0_inv) + dir_deriv = numpy.dot(direction, grad_delta) + + if dir_deriv < 0: + return direction + + log.warn('SGM: non-descent L-BFGS direction at iter %d; reset history', + cycle) + history.clear() + + direction = -(h0_inv * opt_grad_delta) + dir_deriv = numpy.dot(direction, grad_delta) + if dir_deriv >= 0: + log.warn('SGM: fallback direction is not descent at iter %d', cycle) + return None + return direction + + def kernel(self, mo_coeff=None, mo_occ=None): + log = logger.new_logger(self, self.verbose) + if mo_coeff is None: + mo_coeff = self.mo_coeff + if mo_occ is None: + mo_occ = self.mo_occ + if mo_occ is None: + raise RuntimeError('mo_occ must be specified for SGM') + if mo_coeff is None: + raise RuntimeError('mo_coeff must be specified for SGM') + tol = self.tol if self.conv_tol_grad is None else self.conv_tol_grad + if not numpy.isfinite(self.gradient_scale) or self.gradient_scale <= 0: + raise ValueError('gradient_scale must be finite and positive') + self.dump_flags() + + mo_guess = mo_coeff.copy() + h1e = self.get_hcore() + s1e = self.get_ovlp() + dm = self.make_rdm1(mo_coeff, mo_occ) + vhf = self.get_veff(self.mol, dm) + fock = self.get_fock(h1e, s1e, vhf, dm) + e_tot = self.energy_tot(dm, h1e, vhf) + + t0 = (logger.process_clock(), logger.perf_counter()) + h_diag, delta, grad_delta = self._exact_sgm_state(mo_coeff, mo_occ, fock, h1e) + opt_grad_delta = self.gradient_scale * grad_delta + + history = LBFGSHistory(self.lbfgs_memory, _CURVATURE_TOL) + log.info('SGM: initial E = %17.12f Delta = %.3e |g| = %.3e', + e_tot, delta, numpy.sqrt(delta)) + + self.cycles = 0 + for cycle in range(self.max_cycle): + norm_g = numpy.sqrt(delta) + if norm_g < tol: + break + + h0_inv = self._preconditioner(h_diag) + direction = self._descent_direction(history, grad_delta, + opt_grad_delta, h0_inv, log, + cycle) + if direction is None: + break + + trial = self._line_search(mo_coeff, mo_occ, direction, + delta, grad_delta, h1e, s1e) + if trial is None and len(history): + log.warn('SGM: line search failed at iter %d; retry steepest ' + 'descent after clearing L-BFGS history', cycle) + history.clear() + direction = -(h0_inv * opt_grad_delta) + trial = self._line_search(mo_coeff, mo_occ, direction, + delta, grad_delta, h1e, s1e) + + if trial is None: + log.warn('SGM: line search failed at iter %d', cycle) + break + + alpha, step, mo_trial, fock, vhf_new = trial + dm = self.make_rdm1(mo_trial, mo_occ) + e_new = self.energy_tot(dm, h1e, vhf_new) + + h_diag_new, delta_new, grad_delta_new = self._exact_sgm_state( + mo_trial, mo_occ, fock, h1e) + opt_grad_delta_new = self.gradient_scale * grad_delta_new + + step_grad = opt_grad_delta_new - opt_grad_delta + accepted, sy = history.push(step, step_grad) + if not accepted: + if sy < 0: + history.clear() + log.debug1('SGM: skip L-BFGS pair at iter %d; sTy = %.4g', + cycle, sy) + + somo_initial = _somo_overlaps(mo_guess, mo_trial, mo_occ, s1e) + somo_previous = _somo_overlaps(mo_coeff, mo_trial, mo_occ, s1e) + if somo_initial.size: + log.info('SGM SOMO overlap singular values: initial %s; previous %s', + somo_initial, somo_previous) + if min(somo_initial[-1], somo_previous[-1]) < self.somo_overlap_tol: + log.warn('SGM: SOMO subspace deviation at iter %d: ' + 'min overlap initial = %.6f, previous = %.6f ' + '(threshold %.3f); inspect state character', + cycle, somo_initial[-1], somo_previous[-1], self.somo_overlap_tol) + + delta_e = e_new - e_tot + mo_coeff = mo_trial + h_diag = h_diag_new + delta = delta_new + grad_delta = grad_delta_new + opt_grad_delta = opt_grad_delta_new + e_tot = e_new + vhf = vhf_new + self.cycles = cycle + 1 + + log.info('SGM iter %3d: E = %17.12f dE = % .3e ' + 'Delta = %.3e |g| = %.3e alpha = %.3f ' + '|step| = %.3e', + cycle, e_tot, delta_e, delta, numpy.sqrt(delta), + alpha, numpy.linalg.norm(step)) + if callable(self.callback): + self.callback(locals()) + + log.timer('SGM optimization', *t0) + + conv = numpy.sqrt(delta) < tol + if conv: + log.info('SGM converged in %d iterations: |g| = %.6g < %.6g', + self.cycles, numpy.sqrt(delta), tol) + else: + log.info('SGM did not converge in %d iterations: |g| = %.6g', + self.cycles, numpy.sqrt(delta)) + + mo_energy, mo_coeff_canon = self.canonicalize(mo_coeff, mo_occ, fock) + if self.canonicalization: + log.info('Canonicalize SCF orbitals') + mo_coeff = mo_coeff_canon + + self.mo_coeff = mo_coeff + self.mo_occ = mo_occ + self.converged = conv + self.mo_energy = mo_energy + self.e_tot = e_tot + self._finalize() + return self.e_tot + + scf = kernel + + def undo_sgm(self): + obj = lib.view(self, lib.drop_class(self.__class__, SGM)) + del obj._scf + return obj + + +hf.SCF.SGM = lib.class_as_method(SGM) diff --git a/src/nest/soscf/tests/test_sgm.py b/src/nest/soscf/tests/test_sgm.py new file mode 100644 index 0000000..393f7af --- /dev/null +++ b/src/nest/soscf/tests/test_sgm.py @@ -0,0 +1,270 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import io +import unittest +from unittest import mock +import numpy + +from pyscf import gto, lib +from pyscf.lib import logger +from nest.soscf import sgm + + +# Q-Chem benchmark, HCHO Rydberg excitation: +# +# $rem +# METHOD BHHLYP +# BASIS cc-pvdz +# SYMMETRY false +# SYM_IGNORE true +# NO_REORIENT True +# XC_GRID 000075000302 +# UNRESTRICTED false +# delta_scf true +# $end +# +# $delta_scf +# triplet restricted +# triplet_SCF_algorithm SGM +# somo_1 8 +# somo_2 15 +# $end +# +# Q-Chem output: +# Restricted open-shell triplet state = -113.131417 Ha +# SOMO(1): best initial orbital 8 overlap 0.977366 +# SOMO(2): best initial orbital 9 overlap 0.991702 + + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + atom = ''' + C 0.00000000 0.00000000 -1.13947666 + O 0.00000000 0.00000000 1.14402883 + H 0.00000000 1.76627623 -2.23398653 + H 0.00000000 -1.76627623 -2.23398653 + ''' + cls.mol0 = gto.M(atom=atom, charge=0, spin=0, basis='cc-pvdz', + verbose=0, output='/dev/null') + cls.mf0 = cls.mol0.RKS(xc='BHandHLYP') + cls.mf0.grids.atom_grid = (75, 302) + cls.mf0.kernel() + + setocc = cls.mf0.to_uks().mo_occ + setocc[1][7] -= 1 + setocc[0][14] += 1 + cls.ro_occ = setocc[0] + setocc[1] + + cls.mol1 = gto.M(atom=atom, charge=0, spin=2, basis='cc-pvdz', + verbose=0, output='/dev/null') + + @classmethod + def tearDownClass(cls): + cls.mol0.stdout.close() + cls.mol1.stdout.close() + del cls.mol0, cls.mol1, cls.mf0, cls.ro_occ + + def test_delta_gradient(self): + mf = self.mol1.ROKS(xc='BHandHLYP') + mf.grids.atom_grid = (75, 302) + mf.mo_coeff = self.mf0.mo_coeff.copy() + mf.mo_occ = self.ro_occ.copy() + + opt = sgm.SGM(mf) + _h_diag, _delta, grad_delta = opt._exact_sgm_state(mf.mo_coeff, + mf.mo_occ) + rng = numpy.random.default_rng(12) + direction = rng.normal(size=grad_delta.size) + direction /= numpy.linalg.norm(direction) + + eps = 1e-4 + mo_plus = opt._trial_mo(mf.mo_coeff, mf.mo_occ, eps * direction) + mo_minus = opt._trial_mo(mf.mo_coeff, mf.mo_occ, -eps * direction) + g_plus = mf.get_grad(mo_plus, mf.mo_occ) + g_minus = mf.get_grad(mo_minus, mf.mo_occ) + fd = (numpy.dot(g_plus, g_plus) - numpy.dot(g_minus, g_minus)) / (2*eps) + analytic = numpy.dot(grad_delta, direction) + scale = max(abs(fd), abs(analytic), 1e-12) + + self.assertLess(abs(fd - analytic) / scale, 1e-6, + 'fd = %.12e, analytic = %.12e' % (fd, analytic)) + + def test_hcho_rydberg_triplet(self): + for tol, max_cycle in ((1e-4, 80), (1e-7, 120)): + with self.subTest(tol=tol): + mf = self.mol1.ROKS(xc='BHandHLYP') + mf.grids.atom_grid = (75, 302) + mf.mo_coeff = self.mf0.mo_coeff.copy() + mf.mo_occ = self.ro_occ.copy() + + opt = mf.SGM().set(tol=tol, max_cycle=max_cycle) + e_tot = opt.kernel() + + self.assertTrue(opt.converged) + self.assertLess(numpy.linalg.norm(opt.get_grad(opt.mo_coeff, opt.mo_occ)), tol) + self.assertLess(abs(e_tot - -113.131417), 1e-4) + self.assertLess(abs(e_tot - -113.131420102713), 1e-6) + + +class OptimizerChecks(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='sto-3g', + spin=2, verbose=0) + cls.mf = cls.mol.ROHF().run(conv_tol=1e-12) + cls.occ = cls.mf.mo_occ.copy() + opt = cls.mf.SGM() + size = cls.mf.get_grad(cls.mf.mo_coeff, cls.occ).size + step = numpy.random.default_rng(42).normal(size=size)*0.02 + cls.guess = opt._trial_mo(cls.mf.mo_coeff, cls.occ, step) + + def test_last_cycle_and_fock_reuse(self): + opt = self.mf.SGM().set(max_cycle=1, canonicalization=False) + h1e, s1e = opt.get_hcore(), opt.get_ovlp() + hdiag, delta, grad = opt._exact_sgm_state(self.guess, self.occ) + direction = -opt._preconditioner(hdiag)*opt.gradient_scale*grad + trial = opt._line_search(self.guess, self.occ, direction, delta, grad, h1e, s1e) + self.assertIsNotNone(trial) + alpha, _, cnew, fock, _ = trial + final_norm = numpy.linalg.norm(opt.get_grad(cnew, self.occ, fock)) + self.assertLess(final_norm, numpy.sqrt(delta)) + opt.tol = (numpy.sqrt(delta)+final_norm)/2 + trace = [] + opt.callback = lambda env: trace.append(env['delta']) + with mock.patch.object(opt, 'get_veff', wraps=opt.get_veff) as veff: + opt.kernel(self.guess, self.occ) + self.assertTrue(opt.converged) + self.assertEqual(opt.cycles, 1) + self.assertEqual(len(trace), 1) + # One initial build and one per line-search trial, none for the + # accepted-point energy or analytic gradient. + self.assertEqual(veff.call_count, 2+round(-numpy.log2(alpha))) + numpy.testing.assert_allclose(opt.mo_coeff, cnew, atol=1e-12) + + def test_pyscf_interfaces(self): + for mf in (self.mf, self.mol.ROKS(xc='BHandHLYP')): + opt = mf.SGM().set(max_cycle=0, canonicalization=False) + self.assertIsInstance(opt, lib.StreamObject) + self.assertIsInstance(opt, mf.__class__) + self.assertIs(sgm.SGM(opt), opt) + self.assertIs(opt.SGM(), opt) + self.assertIs(opt.scf.__func__, opt.kernel.__func__) + self.assertIs(opt.run(self.guess, self.occ), opt) + plain = opt.undo_sgm() + self.assertIsInstance(plain, mf.__class__) + self.assertNotIsInstance(plain, sgm.SGM) + numpy.testing.assert_allclose(plain.mo_coeff, opt.mo_coeff) + opt = self.mf.SGM().set(max_cycle=0, tol=1e-15, conv_tol_grad=100) + opt.scf(self.guess, self.occ) + self.assertTrue(opt.converged) + + def test_somo_subspace(self): + s = numpy.diag([1., 2., 3., 4.]) + c = numpy.diag(1/numpy.sqrt(s.diagonal())) + occ = numpy.array([2,1,1,0]) + theta = 0.83 + u = numpy.array([[numpy.cos(theta), -numpy.sin(theta)], + [numpy.sin(theta), numpy.cos(theta)]]) + rotated = c.copy() + rotated[:,1:3] = c[:,1:3]@u + numpy.testing.assert_allclose(sgm._somo_overlaps(c,rotated,occ,s), 1, atol=1e-14) + leaked = c.copy() + leaked[:,2] = numpy.cos(theta)*c[:,2]+numpy.sin(theta)*c[:,3] + numpy.testing.assert_allclose(sgm._somo_overlaps(c,leaked,occ,s), + [1,numpy.cos(theta)], atol=1e-14) + leaked[:,1:3] = leaked[:,1:3]@u + numpy.testing.assert_allclose(sgm._somo_overlaps(rotated,leaked,occ,s), + [1,numpy.cos(theta)], atol=1e-14) + + def test_lbfgs_curvature_scale(self): + for scale in (1., 1e-8): + history = sgm.LBFGSHistory() + step = scale*numpy.array([1., 2.]) + diff = scale*numpy.array([2., 3.]) + self.assertTrue(history.push(step, diff)[0]) + self.assertFalse(history.push(step, -diff)[0]) + self.assertFalse(history.push(step, numpy.zeros(2))[0]) + direction = history.direction(diff, numpy.ones(2)) + self.assertLess(numpy.dot(direction, diff), 0) + + def test_somo_warning_each_step(self): + opt = self.mf.SGM().set(max_cycle=2, tol=1e-12, somo_overlap_tol=0.999999999, + verbose=logger.WARN, stdout=io.StringIO()) + opt.kernel(self.guess,self.occ) + self.assertEqual(opt.cycles, 2) + output = opt.stdout.getvalue() + self.assertIn('SOMO subspace deviation at iter 0', output) + self.assertIn('SOMO subspace deviation at iter 1', output) + + def test_frozen_fock_preconditioner(self): + # Remove density response and off-diagonal Fock terms to isolate + # the scale of the mean-field preconditioner in all RO blocks. + opt = self.mf.SGM() + nmo = len(self.occ) + c = numpy.eye(nmo) + fa = numpy.diag(numpy.linspace(-2,3,nmo)) + fb = numpy.diag(numpy.linspace(-1,4,nmo)**3) + fock = lib.tag_array((fa+fb)/2, focka=fa, fockb=fb) + _, _, hdiag = opt.gen_g_hop(opt,c,self.occ,fock) + eps = 1e-4 + curvature = [] + for i in range(hdiag.size): + step = numpy.zeros(hdiag.size) + step[i] = eps + gp = opt.get_grad(opt._trial_mo(c,self.occ,step), self.occ, fock) + gm = opt.get_grad(opt._trial_mo(c,self.occ,-step), self.occ, fock) + curvature.append((gp@gp+gm@gm)/eps**2) + numpy.testing.assert_allclose(curvature,1/opt._preconditioner(hdiag),rtol=1e-6) + + def test_analytic_delta_curvature(self): + for xc in (None, 'LDA,VWN', 'BHandHLYP'): + with self.subTest(xc=xc): + mf = self.mol.ROHF() if xc is None else self.mol.ROKS(xc=xc) + if xc is not None: + mf.grids.level = 1 + opt = mf.SGM() + # A nonstationary reference exercises the g.T@g'' term, + # including third XC derivatives for LDA/GGA. + c = self.guess + g, hop, _ = opt.gen_g_hop(mf, c, self.occ) + curvature = sgm.gen_delta_curvature_rohf(mf, c, self.occ) + rng = numpy.random.default_rng(23) + for _ in range(2): + x = rng.normal(size=g.size) + x /= numpy.linalg.norm(x) + exact, gn = curvature(x) + self.assertAlmostEqual(gn, 2*numpy.linalg.norm(hop(x))**2, places=9) + eps = 5e-4 + plus = opt.get_grad(opt._trial_mo(c, self.occ, eps*x), self.occ) + minus = opt.get_grad(opt._trial_mo(c, self.occ, -eps*x), self.occ) + fd = (plus@plus + minus@minus - 2*g@g)/eps**2 + self.assertLess(abs(fd-exact)/max(1., abs(exact)), 2e-6) + + def test_delta_curvature_at_stationary_point(self): + c = self.mf.mo_coeff + g, hop, _ = self.mf.SGM().gen_g_hop(self.mf, c, self.occ) + curvature = sgm.gen_delta_curvature_rohf(self.mf, c, self.occ) + x = numpy.random.default_rng(12).normal(size=g.size) + x /= numpy.linalg.norm(x) + exact, gn = curvature(x) + self.assertLess(numpy.linalg.norm(g), 1e-6) + self.assertLess(abs(exact-gn)/max(1., gn), 1e-6) + + +if __name__ == '__main__': + print('Full tests for soscf.sgm') + unittest.main() From 5d1b405d1c6796b4dc843ee411e024d3d20a88ec Mon Sep 17 00:00:00 2001 From: wtpeter Date: Sat, 12 Sep 2026 17:09:27 +0800 Subject: [PATCH 2/3] feat(soscf): add GN and full-Fock SGM preconditioners Add an exact Gauss-Newton diagonal with configurable refresh intervals and a full-Fock approximation without additional density responses. Keep the gap preconditioner as the default. Validation: SGM tests passed (12 tests and 11 subtests), Ruff passed, and all three preconditioners converged the HCHO triplet benchmark to a gradient norm below 1e-7. --- src/nest/soscf/sgm.py | 68 ++++++++++++++++++++++++++++---- src/nest/soscf/tests/test_sgm.py | 43 ++++++++++++++++++-- 2 files changed, 100 insertions(+), 11 deletions(-) diff --git a/src/nest/soscf/sgm.py b/src/nest/soscf/sgm.py index be9309d..dbf292b 100644 --- a/src/nest/soscf/sgm.py +++ b/src/nest/soscf/sgm.py @@ -284,6 +284,17 @@ class SGM(lib.StreamObject): this threshold relative to the initial or previous orbitals. Default 0.7 corresponds to a largest principal angle of about 46 degrees. This is a diagnostic, not a state constraint. + preconditioner + 'gap' (default) uses 2*h_diag**2. 'fock' retains the full Fock + commutator in J but omits density response; it needs no extra + response evaluations. 'gn' uses the exact diagonal + of 2*J.T@J, where J is the orbital-gradient Jacobian. This needs + one Hessian-vector product per orbital variable, but no third + derivatives and no stored full Hessian. + preconditioner_update + Refresh the 'gn' diagonal every this many accepted steps. + Default 1; 0 computes it only at the initial point. L-BFGS + updates still run every step. 'gap' and 'fock' are always refreshed. The returned object also inherits the input ROHF/ROKS class. Use mf.SGM().set(...).run(mo_coeff, mo_occ), or SGM(mf).kernel(...). @@ -297,11 +308,14 @@ class SGM(lib.StreamObject): lbfgs_memory = getattr(__config__, 'sgm_lbfgs_memory', 8) gradient_scale = getattr(__config__, 'sgm_gradient_scale', 0.75) somo_overlap_tol = getattr(__config__, 'sgm_somo_overlap_tol', 0.7) + preconditioner = 'gap' + preconditioner_update = 1 canonicalization = getattr(__config__, 'soscf_newton_ah_SOSCF_canonicalization', True) _keys = {'max_cycle', 'tol', 'lbfgs_memory', 'gradient_scale', - 'canonicalization', 'somo_overlap_tol'} + 'canonicalization', 'somo_overlap_tol', 'preconditioner', + 'preconditioner_update'} gen_g_hop = staticmethod(gen_g_hop_rohf) @@ -327,14 +341,48 @@ def dump_flags(self, verbose=None): self.tol if self.conv_tol_grad is None else self.conv_tol_grad) log.info('SGM gradient scale = %g', self.gradient_scale) log.info('SGM L-BFGS memory = %d', self.lbfgs_memory) + log.info('SGM preconditioner = %s', self.preconditioner) + if self.preconditioner == 'gn': + log.info('SGM preconditioner update interval = %d (0: initial only)', + self.preconditioner_update) log.info('SGM SOMO overlap warning threshold = %g', self.somo_overlap_tol) log.info('SGM canonicalization = %s', self.canonicalization) return self - def _preconditioner(self, h_diag): + def _preconditioner(self, h_diag, h_op=None, mo_coeff=None, mo_occ=None, fock=None): # Diagonal, frozen-Fock approximation to 2*J.T@J for Delta = g.T@g. # J = d g / d kappa; this is not the full Delta Hessian away from g=0. - denom = numpy.maximum(2.0 * h_diag ** 2, _PRECOND_FLOOR) + denom = 2.0 * h_diag ** 2 + if self.preconditioner == 'gn': + direction = numpy.zeros_like(h_diag) + for i in range(h_diag.size): + direction[i] = 1 + column = h_op(direction) + denom[i] = 2.0 * numpy.dot(column, column) + direction[i] = 0 + elif self.preconditioner == 'fock': + f = mo_coeff.T @ numpy.asarray((fock.focka, fock.fockb)) @ mo_coeff + occ = numpy.asarray((mo_occ > 0, mo_occ == 2)) + masks = (~occ[:, :, None]) & occ[:, None, :] + unique = masks[0] | masks[1] + allowed = numpy.ones(h_diag.size, dtype=bool) + if self.mol.symmetry: + orbsym = hf_symm.get_orbsym(self.mol, mo_coeff) + allowed = (orbsym[:, None] == orbsym)[unique] + for index, (a, i) in enumerate(zip(*numpy.where(unique))): + if not allowed[index]: + denom[index] = 0 + continue + # Column of J_F: [F, K_ai], evaluated using the two nonzero + # entries of K_ai instead of a dense matrix multiplication. + comm = numpy.zeros_like(f) + comm[:, :, i] += f[:, :, a] + comm[:, :, a] -= f[:, :, i] + comm[:, a, :] -= f[:, i, :] + comm[:, i, :] += f[:, a, :] + column = numpy.sum(comm * masks, axis=0)[unique] * allowed + denom[index] = 2.0 * numpy.dot(column, column) + denom = numpy.maximum(denom, _PRECOND_FLOOR) return numpy.minimum(1.0 / denom, _MAX_PRECOND) def _trial_mo(self, mo_coeff, mo_occ, step): @@ -349,7 +397,7 @@ def _exact_sgm_state(self, mo_coeff, mo_occ, fock_ao=None, h1e=None): g_orb, h_op, h_diag = self.gen_g_hop(self, mo_coeff, mo_occ, fock_ao, h1e) delta = numpy.dot(g_orb, g_orb) grad_delta = 2.0 * h_op(g_orb) - return h_diag, delta, grad_delta + return h_diag, delta, grad_delta, h_op def _line_search(self, mo_coeff, mo_occ, direction, delta, grad_delta, h1e, s1e): @@ -405,6 +453,10 @@ def kernel(self, mo_coeff=None, mo_occ=None): tol = self.tol if self.conv_tol_grad is None else self.conv_tol_grad if not numpy.isfinite(self.gradient_scale) or self.gradient_scale <= 0: raise ValueError('gradient_scale must be finite and positive') + if self.preconditioner not in ('gap', 'fock', 'gn'): + raise ValueError("preconditioner must be 'gap', 'fock' or 'gn'") + if not isinstance(self.preconditioner_update, (int, numpy.integer)) or self.preconditioner_update < 0: + raise ValueError('preconditioner_update must be a nonnegative integer') self.dump_flags() mo_guess = mo_coeff.copy() @@ -416,7 +468,7 @@ def kernel(self, mo_coeff=None, mo_occ=None): e_tot = self.energy_tot(dm, h1e, vhf) t0 = (logger.process_clock(), logger.perf_counter()) - h_diag, delta, grad_delta = self._exact_sgm_state(mo_coeff, mo_occ, fock, h1e) + h_diag, delta, grad_delta, h_op = self._exact_sgm_state(mo_coeff, mo_occ, fock, h1e) opt_grad_delta = self.gradient_scale * grad_delta history = LBFGSHistory(self.lbfgs_memory, _CURVATURE_TOL) @@ -429,7 +481,9 @@ def kernel(self, mo_coeff=None, mo_occ=None): if norm_g < tol: break - h0_inv = self._preconditioner(h_diag) + if (cycle == 0 or self.preconditioner in ('gap', 'fock') + or (self.preconditioner_update and cycle % self.preconditioner_update == 0)): + h0_inv = self._preconditioner(h_diag, h_op, mo_coeff, mo_occ, fock) direction = self._descent_direction(history, grad_delta, opt_grad_delta, h0_inv, log, cycle) @@ -454,7 +508,7 @@ def kernel(self, mo_coeff=None, mo_occ=None): dm = self.make_rdm1(mo_trial, mo_occ) e_new = self.energy_tot(dm, h1e, vhf_new) - h_diag_new, delta_new, grad_delta_new = self._exact_sgm_state( + h_diag_new, delta_new, grad_delta_new, h_op = self._exact_sgm_state( mo_trial, mo_occ, fock, h1e) opt_grad_delta_new = self.gradient_scale * grad_delta_new diff --git a/src/nest/soscf/tests/test_sgm.py b/src/nest/soscf/tests/test_sgm.py index 393f7af..0cb563b 100644 --- a/src/nest/soscf/tests/test_sgm.py +++ b/src/nest/soscf/tests/test_sgm.py @@ -85,8 +85,8 @@ def test_delta_gradient(self): mf.mo_occ = self.ro_occ.copy() opt = sgm.SGM(mf) - _h_diag, _delta, grad_delta = opt._exact_sgm_state(mf.mo_coeff, - mf.mo_occ) + _h_diag, _delta, grad_delta, _hop = opt._exact_sgm_state(mf.mo_coeff, + mf.mo_occ) rng = numpy.random.default_rng(12) direction = rng.normal(size=grad_delta.size) direction /= numpy.linalg.norm(direction) @@ -104,7 +104,7 @@ def test_delta_gradient(self): 'fd = %.12e, analytic = %.12e' % (fd, analytic)) def test_hcho_rydberg_triplet(self): - for tol, max_cycle in ((1e-4, 80), (1e-7, 120)): + for tol, max_cycle in ((1e-4, 80), (1e-7, 200)): with self.subTest(tol=tol): mf = self.mol1.ROKS(xc='BHandHLYP') mf.grids.atom_grid = (75, 302) @@ -135,7 +135,7 @@ def setUpClass(cls): def test_last_cycle_and_fock_reuse(self): opt = self.mf.SGM().set(max_cycle=1, canonicalization=False) h1e, s1e = opt.get_hcore(), opt.get_ovlp() - hdiag, delta, grad = opt._exact_sgm_state(self.guess, self.occ) + hdiag, delta, grad, _hop = opt._exact_sgm_state(self.guess, self.occ) direction = -opt._preconditioner(hdiag)*opt.gradient_scale*grad trial = opt._line_search(self.guess, self.occ, direction, delta, grad, h1e, s1e) self.assertIsNotNone(trial) @@ -264,6 +264,41 @@ def test_delta_curvature_at_stationary_point(self): self.assertLess(numpy.linalg.norm(g), 1e-6) self.assertLess(abs(exact-gn)/max(1., gn), 1e-6) + def test_response_preconditioners(self): + opt = self.mf.SGM() + dm = opt.make_rdm1(self.guess, self.occ) + fock = opt.get_fock(dm=dm) + _, hop, hd = opt.gen_g_hop(opt, self.guess, self.occ, fock) + # Independent finite differences of the orbital gradient give J's + # columns, including off-diagonal response contributions. + eps = 1e-5 + for mode in ('gn', 'fock'): + with self.subTest(mode=mode): + opt.preconditioner = mode + inverse = opt._preconditioner(hd, hop, self.guess, self.occ, fock) + fd_diagonal = [] + frozen = fock if mode == 'fock' else None + for i in range(hd.size): + step = numpy.zeros(hd.size) + step[i] = eps + gp = opt.get_grad(opt._trial_mo(self.guess, self.occ, step), self.occ, frozen) + gm = opt.get_grad(opt._trial_mo(self.guess, self.occ, -step), self.occ, frozen) + column = (gp-gm)/(2*eps) + fd_diagonal.append(2*numpy.dot(column, column)) + numpy.testing.assert_allclose(1/inverse, fd_diagonal, rtol=1e-7) + + def test_gn_preconditioner_convergence_and_refresh(self): + for mode, interval in (('gn', 0), ('gn', 1), ('gn', 3), ('fock', 1)): + with self.subTest(mode=mode, interval=interval): + opt = self.mf.SGM().set(preconditioner=mode, preconditioner_update=interval, + tol=1e-7, max_cycle=100) + with mock.patch.object(opt, '_preconditioner', wraps=opt._preconditioner) as precond: + opt.kernel(self.guess, self.occ) + self.assertTrue(opt.converged) + self.assertLess(abs(opt.e_tot-self.mf.e_tot), 1e-9) + expected = 1 if interval == 0 else 1+(opt.cycles-1)//interval + self.assertEqual(precond.call_count, expected) + if __name__ == '__main__': print('Full tests for soscf.sgm') From d5b76aee6d7168d10f4eb73aeb34b157c91c1153 Mon Sep 17 00:00:00 2001 From: wtpeter Date: Wed, 16 Sep 2026 17:05:11 +0800 Subject: [PATCH 3/3] feat(deltascf): add FR and migrate SGM workflows --- examples/deltascf/01_sgm.py | 90 ++++++++ examples/deltascf/02_fr.py | 66 ++++++ examples/deltascf/03_tda_guess.py | 109 ++++++++++ examples/deltascf/04_nttda.py | 78 +++++++ examples/soscf/01_sgm.py | 76 ------- src/nest/deltascf/fr.py | 267 ++++++++++++++++++++++++ src/nest/{soscf => deltascf}/sgm.py | 162 ++------------- src/nest/deltascf/tests/test_fr.py | 187 +++++++++++++++++ src/nest/deltascf/tests/test_sgm.py | 121 +++++++++++ src/nest/soscf/tests/test_sgm.py | 305 ---------------------------- 10 files changed, 930 insertions(+), 531 deletions(-) create mode 100644 examples/deltascf/01_sgm.py create mode 100644 examples/deltascf/02_fr.py create mode 100644 examples/deltascf/03_tda_guess.py create mode 100644 examples/deltascf/04_nttda.py delete mode 100644 examples/soscf/01_sgm.py create mode 100644 src/nest/deltascf/fr.py rename src/nest/{soscf => deltascf}/sgm.py (67%) create mode 100644 src/nest/deltascf/tests/test_fr.py create mode 100644 src/nest/deltascf/tests/test_sgm.py delete mode 100644 src/nest/soscf/tests/test_sgm.py diff --git a/examples/deltascf/01_sgm.py b/examples/deltascf/01_sgm.py new file mode 100644 index 0000000..d7643ec --- /dev/null +++ b/examples/deltascf/01_sgm.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CH4 core excitation with ordinary SGM and a canonical-orbital guess. + +This is a restricted triplet calculation. +Only use the excitation energy when the excited-state calculation converges. +""" +# Q-Chem reference input for this CH4 core excitation benchmark: +# +# $rem +# METHOD BHHLYP +# BASIS aug-cc-pvtz +# SYMMETRY false +# SYM_IGNORE true +# NO_REORIENT True +# XC_GRID 000075000302 +# UNRESTRICTED false +# delta_scf true +# $end +# +# $delta_scf +# triplet restricted +# triplet_SCF_algorithm SGM +# somo_1 1 +# somo_2 6 +# $end +# +# Q-Chem output: +# Restricted open-shell triplet state = -29.947480 Ha +# SOMO(1): best initial orbital 5 overlap 0.999002 +# SOMO(2): best initial orbital 6 overlap 0.879378 +# Excitation energy = 10.554376 Ha / 287.199215 eV + +from pyscf import gto +from pyscf.data import nist +from nest.deltascf import sgm + +# 1. Calculate the closed-shell ground state. +mol = gto.M( + atom=""" + C 0.0000 0.0000 0.0000 + H 0.6276 0.6276 0.6276 + H 0.6276 -0.6276 -0.6276 + H -0.6276 0.6276 -0.6276 + H -0.6276 -0.6276 0.6276 + """, + basis='aug-cc-pvtz', + spin=0, + symmetry=False, +) +ground = mol.RKS(xc='BHandHLYP') +ground.grids.atom_grid = (75, 302) +ground.kernel() + +# 2. Make a triplet guess by changing two occupations from 2/0 to 1/1. +# Indices start at zero: orbital 0 is C 1s, orbital 5 is the ground-state LUMO. +hole_index = 0 +particle_index = 5 +initial_orbitals = ground.mo_coeff.copy() +initial_occupations = ground.mo_occ.copy() +initial_occupations[hole_index] = 1 +initial_occupations[particle_index] = 1 + +# 3. Optimize the excited-state orbitals using the original SGM preconditioner. +triplet_mol = mol.copy() +triplet_mol.spin = 2 # PySCF spin = N_alpha - N_beta = 2S. +triplet_scf = triplet_mol.ROKS(xc=ground.xc) +triplet_scf.grids.atom_grid = (75, 302) +excited = sgm.SGM(triplet_scf) +excited.verbose = 4 +excited.kernel(initial_orbitals, initial_occupations) + +print('SGM converged:', excited.converged) +print('Last energy (Ha):', excited.e_tot) +if excited.converged: + excitation_energy = (excited.e_tot - ground.e_tot) * nist.HARTREE2EV + print('Triplet excitation energy (eV):', excitation_energy) diff --git a/examples/deltascf/02_fr.py b/examples/deltascf/02_fr.py new file mode 100644 index 0000000..7b87212 --- /dev/null +++ b/examples/deltascf/02_fr.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CH4 core excitation with ordinary FR and a canonical-orbital guess. + +This is a restricted triplet calculation. +Only use the excitation energy when the excited-state calculation converges. +""" + +from pyscf import gto +from pyscf.data import nist +from nest.deltascf import fr + +# 1. Calculate the closed-shell ground state. +mol = gto.M( + atom=""" + C 0.0000 0.0000 0.0000 + H 0.6276 0.6276 0.6276 + H 0.6276 -0.6276 -0.6276 + H -0.6276 0.6276 -0.6276 + H -0.6276 -0.6276 0.6276 + """, + basis='aug-cc-pvtz', + spin=0, + symmetry=False, +) +ground = mol.RKS(xc='BHandHLYP') +ground.grids.atom_grid = (75, 302) +ground.kernel() + +# 2. Make a triplet guess by changing two occupations from 2/0 to 1/1. +# Indices start at zero: orbital 0 is C 1s, orbital 5 is the ground-state LUMO. +hole_index = 0 +particle_index = 5 +initial_orbitals = ground.mo_coeff.copy() +initial_occupations = ground.mo_occ.copy() +initial_occupations[hole_index] = 1 +initial_occupations[particle_index] = 1 + +# 3. Freeze both SOMOs, relax the other orbitals, then release with IMOM. +triplet_mol = mol.copy() +triplet_mol.spin = 2 # PySCF spin = N_alpha - N_beta = 2S. +triplet_scf = triplet_mol.ROKS(xc=ground.xc) +triplet_scf.grids.atom_grid = (75, 302) +excited = fr.FR(triplet_scf) +excited.verbose = 4 +excited.kernel(initial_orbitals, initial_occupations) + +print('Freeze converged:', excited.freeze_converged, 'cycles:', excited.freeze_cycles) +print('Release converged:', excited.converged, 'cycles:', excited.cycles) +print('Last energy (Ha):', excited.e_tot) +if excited.converged: + excitation_energy = (excited.e_tot - ground.e_tot) * nist.HARTREE2EV + print('Triplet excitation energy (eV):', excitation_energy) diff --git a/examples/deltascf/03_tda_guess.py b/examples/deltascf/03_tda_guess.py new file mode 100644 index 0000000..e671efe --- /dev/null +++ b/examples/deltascf/03_tda_guess.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""C2H3F C1 core excitation: a core-TDA guess followed by SGM or FR. + +The chosen TDA root defines an initial guess; it does not guarantee the +identity of the final SCF solution. Check convergence before using energies. +""" + +import numpy as np +from scipy.linalg import null_space +from pyscf import gto, tdscf +from pyscf.data import nist +from nest.deltascf import fr, sgm + +# 1. Calculate the closed-shell ground state. +mol = gto.M( + atom=""" + C1 0.000000 -0.246412 -1.271068 + C2 0.000000 0.457081 -0.154735 + F1 0.000000 -0.119195 1.052878 + H1 0.000000 0.272328 -2.210194 + H2 0.000000 -1.319906 -1.249847 + H3 0.000000 1.530323 -0.095954 + """, + basis={'C1': 'aug-cc-pvtz', 'C2': 'aug-cc-pvtz', + 'F1': 'cc-pvtz', 'H': 'cc-pvdz'}, + spin=0, + symmetry=False, +) +ground = mol.RKS(xc='BHandHLYP') +ground.kernel() + +# 2. Obtain a particle orbital from single-hole, triplet core-TDA. +# These zero-based indices apply to THIS molecule and basis. +hole_index = 2 # C1 1s orbital. +particle_index = 12 # First ground-state virtual; also the slot for our particle orbital. +tda_root = 1 # Second TDA root: the even-parity root with a large LUMO component. + +# Allow excitations only from the chosen core orbital, to all virtual orbitals. +# This freezing applies to TDA only, not to the subsequent SCF calculation. +frozen_occupied = [] +for orbital_index in range(particle_index): + if orbital_index != hole_index: + frozen_occupied.append(orbital_index) + +tda = tdscf.TDA(ground) +tda.singlet = False +tda.frozen = frozen_occupied +tda.kernel(nstates=3) + +for root_index, (energy, amplitudes) in enumerate(zip(tda.e, tda.xy)): + excitation_amplitudes, _ = amplitudes + particle_vector = excitation_amplitudes[0].copy() + particle_vector /= np.linalg.norm(particle_vector) + lumo_weight = particle_vector[0]**2 + print(f'TDA root {root_index}: {energy * nist.HARTREE2EV:.6f} eV, ' + f'ground-state LUMO weight = {lumo_weight:.4f}') + +# Only one occupied orbital participates, so X has a single row. Its entries +# are the coefficients of the particle orbital in the OLD virtual basis. +# Normalize explicitly because PySCF's restricted TDA uses 2 * ||X||^2 = 1. +excitation_amplitudes, _ = tda.xy[tda_root] +particle_vector = excitation_amplitudes[0].copy() +particle_vector /= np.linalg.norm(particle_vector) + +# Put the particle orbital in the original LUMO slot. Complete the remaining +# virtual columns with orthonormal combinations perpendicular to that vector. +# These other columns are a basis completion, not higher TDA excited states. +old_virtual_orbitals = ground.mo_coeff[:, particle_index:] +remaining_virtual_vectors = null_space(particle_vector.reshape(1, -1)) +initial_orbitals = ground.mo_coeff.copy() +initial_orbitals[:, particle_index] = old_virtual_orbitals @ particle_vector +initial_orbitals[:, particle_index + 1:] = old_virtual_orbitals @ remaining_virtual_vectors + +initial_occupations = ground.mo_occ.copy() +initial_occupations[hole_index] = 1 +initial_occupations[particle_index] = 1 + +# 3. Optimize the restricted triplet. Choose FR or SGM on the next line. +triplet_mol = mol.copy() +triplet_mol.spin = 2 +triplet_scf = triplet_mol.ROKS(xc=ground.xc) +excited = fr.FR(triplet_scf) # Or: sgm.SGM(triplet_scf) +excited.verbose = 4 +excited.conv_tol_grad = 1e-5 +# FR first freezes both singly occupied orbitals, then releases them with IMOM. +# verbose=4 prints both stage boundaries and the final SOMO-subspace overlaps. +excited.conv_tol = 1e-9 +excited.max_cycle = 200 +excited.kernel(initial_orbitals, initial_occupations) + +print('Optimizer converged:', excited.converged, 'cycles:', excited.cycles) +print('Last energy (Ha):', excited.e_tot) +if excited.converged: + excitation_energy = (excited.e_tot - ground.e_tot) * nist.HARTREE2EV + print('Triplet excitation energy (eV):', excitation_energy) diff --git a/examples/deltascf/04_nttda.py b/examples/deltascf/04_nttda.py new file mode 100644 index 0000000..8a344dc --- /dev/null +++ b/examples/deltascf/04_nttda.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run deltaS=-1 NTTDA with SGM/FR delta-SCF references for CH4. + +Optimize the core-excited triplet with SGM or FR, then use it as the +reference for singlet NTTDA. Use the CH4 geometry and C 1s/LUMO guess +from 01_sgm.py and 02_fr.py, with a 6-31G* basis and the default grid. +""" + +import numpy as np +from pyscf import gto +from pyscf.data import nist +from nest import nttda # Registers NTTDA on ROKS. +from nest.deltascf import fr, sgm + +# 1. Calculate the closed-shell ground state. +mol = gto.M( + atom=""" + C 0.0000 0.0000 0.0000 + H 0.6276 0.6276 0.6276 + H 0.6276 -0.6276 -0.6276 + H -0.6276 0.6276 -0.6276 + H -0.6276 -0.6276 0.6276 + """, + basis='631G*', + spin=0, + symmetry=False, +) +ground = mol.RKS(xc='BHandHLYP') +ground.kernel() + +# 2. Make a triplet guess by changing two occupations from 2/0 to 1/1. +# Indices start at zero: orbital 0 is C 1s, orbital 5 is the ground-state LUMO. +hole_index = 0 +particle_index = 5 +initial_orbitals = ground.mo_coeff.copy() +initial_occupations = ground.mo_occ.copy() +initial_occupations[hole_index] = 1 +initial_occupations[particle_index] = 1 + +# 3. Optimize the delta-SCF triplet reference with SGM or FR. +triplet_mol = mol.copy() +triplet_mol.spin = 2 + +triplet_scf = triplet_mol.ROKS(xc=ground.xc) +excited = fr.FR(triplet_scf) # or sgm.SGM(triplet_scf) +excited.verbose = 4 +excited.kernel(initial_orbitals, initial_occupations) +excited.analyze(verbose=4) + +# 4. Run deltaS=-1 NTTDA on the optimized delta-SCF reference. +reference = excited # for SGM, reference = excited.undo_sgm() +td = reference.NTTDA() +td.deltaS = -1 # S_reference=1 -> S_final=0 (singlets). +td.nstates = 10 # Include core-excited singlets above the lower valence roots. +td.kernel() +td.analyze(verbose=4) + +# Supplementary excitation-energy comparison. +# E_triplet - E_RKS defines the delta-SCF gap for the settings used here. +# omega is relative to the triplet, so add that gap for a common RKS zero. +delta_scf_ev = (reference.e_tot - ground.e_tot) * nist.HARTREE2EV +nttda_rks_ev = delta_scf_ev + td.e * nist.HARTREE2EV +nttda_excitation_ev = (td.e - td.e[0]) * nist.HARTREE2EV +print(f'delta-SCF triplet excitation = {delta_scf_ev:.6f} eV') diff --git a/examples/soscf/01_sgm.py b/examples/soscf/01_sgm.py deleted file mode 100644 index 35fa6ae..0000000 --- a/examples/soscf/01_sgm.py +++ /dev/null @@ -1,76 +0,0 @@ -#!/usr/bin/env python -# Copyright 2026 The NEST Developers. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -'''Run a restricted open-shell Kohn-Sham (ROKS) calculation using the SGM''' -# Q-Chem reference input for this CH4 core excitation benchmark: -# -# $rem -# METHOD BHHLYP -# BASIS aug-cc-pvtz -# SYMMETRY false -# SYM_IGNORE true -# NO_REORIENT True -# XC_GRID 000075000302 -# UNRESTRICTED false -# delta_scf true -# $end -# -# $delta_scf -# triplet restricted -# triplet_SCF_algorithm SGM -# somo_1 1 -# somo_2 6 -# $end -# -# Q-Chem output: -# Restricted open-shell triplet state = -29.947480 Ha -# SOMO(1): best initial orbital 5 overlap 0.999002 -# SOMO(2): best initial orbital 6 overlap 0.879378 -# Excitation energy = 10.554376 Ha / 287.199215 eV - -from pyscf import gto -from nest.soscf import sgm - - -atom = ''' -C 0.00000000 0.00000000 0.00000000 -H 0.62760000 0.62760000 0.62760000 -H 0.62760000 -0.62760000 -0.62760000 -H -0.62760000 0.62760000 -0.62760000 -H -0.62760000 -0.62760000 0.62760000 -''' - -mol = gto.M(atom=atom, charge=0, spin=0, basis='aug-cc-pvtz') -mf = mol.RKS(xc='BHandHLYP') -mf.grids.atom_grid = (75, 302) -mf.kernel() - -setocc = mf.to_uks().mo_occ -setocc[1][0] -= 1 -setocc[0][5] += 1 -ro_occ = setocc[0] + setocc[1] - -mol1 = gto.M(atom=atom, charge=0, spin=2, basis='aug-cc-pvtz') -mf1 = mol1.ROKS(xc='BHandHLYP') -mf1.grids.atom_grid = (75, 302) -mf1.mo_coeff = mf.mo_coeff -mf1.mo_occ = ro_occ - -sgm_mf = mf1.SGM() -sgm_mf.verbose = 4 -sgm_mf.kernel() - -print('SGM converged =', sgm_mf.converged) -print('SGM energy = %.12f Ha' % sgm_mf.e_tot) diff --git a/src/nest/deltascf/fr.py b/src/nest/deltascf/fr.py new file mode 100644 index 0000000..709e918 --- /dev/null +++ b/src/nest/deltascf/fr.py @@ -0,0 +1,267 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Ref: JCTC 2027, 22, 4609 + +"""First-order freeze-and-release SCF for restricted two-SOMO triplets. + +ROHF/ROKS adaptation of Qin and Suo's FR-TO delta-SCF workflow. The +frozen stage uses exact complementary-space diagonalization, rather than +finite projector shifts in the paper's unrestricted spin channel. +""" + +import numpy as np +from scipy.linalg import eigh +from pyscf import lib, scf, symm +from pyscf.lib import logger +from pyscf.scf import hf + + +def _set_imom(mf, mo_ref, setocc): + scf.addons.mom_occ(mf, mo_ref, setocc) + imom_occ = mf.get_occ + if mf.mol.symmetry: + reference_symmetry = mf.get_orbsym(mo_ref) + overlap = mf.get_ovlp() + + def get_occ(mo_energy=None, mo_coeff=None): + if not mf.mol.symmetry: + occupations = imom_occ(mo_energy, mo_coeff) + else: + # PySCF mom_occ ranks all orbitals together. Here retain the + # reference alpha/beta electron counts within each irrep, so + # overlap ranking cannot move electrons between symmetry blocks. + if mo_coeff is None: + mo_coeff = mf.mo_coeff + orbital_symmetry = mf.get_orbsym(mo_coeff) + spin_occupations = np.zeros((2, mo_coeff.shape[1])) + for irrep in np.unique(reference_symmetry): + candidates = np.flatnonzero(orbital_symmetry == irrep) + for spin in range(2): + occupied_reference = (reference_symmetry == irrep) & (setocc[spin] > 0) + noccupied = np.count_nonzero(occupied_reference) + if len(candidates) < noccupied: + raise ValueError('Not enough orbitals in an IMOM symmetry block') + overlaps = mo_ref[:, occupied_reference].T @ overlap @ mo_coeff[:, candidates] + weights = np.sum(overlaps**2, axis=0) + selected = candidates[np.argsort(-weights, kind='stable')[:noccupied]] + spin_occupations[spin, selected] = 1 + occupations = spin_occupations.sum(axis=0) + # Independent alpha/beta rankings can violate restricted nesting. + if np.count_nonzero(occupations == 1) != 2: + raise RuntimeError('IMOM selected nonnested alpha/beta spaces incompatible with a ROHF triplet') + return occupations + + mf.get_occ = get_occ + + +class FR(lib.StreamObject): + """Freeze both singly occupied orbitals, relax the others, then run IMOM. + + Use ``mf.FR().kernel(mo_coeff, mo_occ)`` on a plain ROHF/ROKS object. + Requires real S-orthonormal orbitals and spin=2. With symmetry enabled, + input orbitals must belong to individual irreps of an Abelian point + group (e.g. Cs, C2v, D2h). Both stages retain the initial alpha/beta + electron counts in each irrep. For linear molecules or atoms, select + an Abelian subgroup explicitly; Dooh, Coov and SO3 are not supported. + Both stages use ordinary SCF and fresh CDIIS histories, with no Hessian. + Energies always come from the physical Hamiltonian. + + freeze_tol (default 1e-5) bounds PySCF's projected orbital-gradient norm; + freeze_max_cycle (100) bounds the frozen stage. Failure raises before + release. freeze_cycles/freeze_converged describe that stage; cycles and + converged describe release, governed by the usual SCF tolerances. + callback receives both stages, identified by env['fr_stage']. + PySCF's extra relaxed convergence check is disabled in both stages. + At INFO verbosity (verbose=4), stage starts and the final SOMO-subspace + overlap singular values relative to the input orbitals are logged. + """ + + __name_mixin__ = 'FR' + freeze_tol = 1e-5 + freeze_max_cycle = 100 + freeze_cycles = 0 + freeze_converged = None + _keys = {'freeze_tol', 'freeze_max_cycle', 'freeze_cycles', 'freeze_converged'} + + def __new__(cls, mf): + if isinstance(mf, FR): + return mf + if not isinstance(mf, hf.SCF) or not mf.istype('ROHF') or hasattr(mf, '_scf'): + raise NotImplementedError('FR requires a plain ROHF/ROKS object') + return lib.set_class(object.__new__(cls), (cls, mf.__class__)) + + def __init__(self, mf): + if mf is not self: + self.__dict__.update(mf.__dict__) + + def kernel(self, mo_coeff=None, mo_occ=None): + c = np.array(self.mo_coeff if mo_coeff is None else mo_coeff, copy=True) + occ = np.array(self.mo_occ if mo_occ is None else mo_occ, copy=True) + if (self.mol.spin != 2 or occ.ndim != 1 or + not np.all(np.isin(occ, [0, 1, 2])) or + np.count_nonzero(occ == 1) != 2 or occ.sum() != self.mol.nelectron): + raise ValueError('FR requires a spin=2 occupation vector with exactly two SOMOs') + if c.ndim != 2 or c.shape != (self.mol.nao_nr(), occ.size) or np.iscomplexobj(c): + raise ValueError('FR requires real MO coefficients matching mo_occ') + s = self.get_ovlp() + if not np.allclose(c.T @ s @ c, np.eye(occ.size), atol=1e-8, rtol=0): + raise ValueError('FR requires S-orthonormal input orbitals') + orbital_symmetry = np.zeros(occ.size, dtype=int) + if self.mol.symmetry: + if self.mol.groupname not in ('C1', 'Ci', 'Cs', 'C2', 'C2v', 'C2h', 'D2', 'D2h'): + raise NotImplementedError('FR supports Abelian point groups; select an Abelian subgroup') + orbital_symmetry = symm.label_orb_symm( + self.mol, self.mol.irrep_id, self.mol.symm_orb, c, s=s, check=True) + c = lib.tag_array(c, orbsym=orbital_symmetry) + for name, irrep in zip(self.mol.irrep_name, self.mol.irrep_id): + alpha = np.count_nonzero((orbital_symmetry == irrep) & (occ > 0)) + beta = np.count_nonzero((orbital_symmetry == irrep) & (occ == 2)) + requested = self.irrep_nelec.get(name) + if requested is not None: + matches = (requested == alpha + beta if np.isscalar(requested) + else tuple(requested) == (alpha, beta)) + if not matches: + raise ValueError('Input occupations conflict with irrep_nelec for ' + name) + if self.freeze_tol <= 0 or self.freeze_max_cycle < 1: + raise ValueError('freeze_tol and freeze_max_cycle must be positive') + + self.converged = False + self.freeze_converged = False + self.freeze_cycles = 0 + # Drop only this mixin: the two stages retain all physical MF settings. + base = lib.view(self, lib.drop_class(self.__class__, FR)).copy() + for key in FR._keys: + base.__dict__.pop(key, None) + frozen, release = base.copy(), base.copy() + callback = self.callback + for stage, mf in [('freeze', frozen), ('release', release)]: + mf.conv_check = False + if self.diis: + if stage == 'freeze': + mf.diis = scf.diis.CDIIS(mf) + mf.diis.space = self.diis_space + mf.diis.rollback = self.diis_space_rollback + mf.diis.damp = self.diis_damp + else: + # Let the ordinary kernel build its orthogonalized DIIS + # residual, without carrying over the frozen history. + mf.diis = True + mf.DIIS = scf.diis.CDIIS + def stage_callback(env, stage=stage): + if callback is not None: + callback(dict(env, fr_stage=stage)) + mf.callback = stage_callback + + # The two singly occupied columns never change during the freeze + # iterations. All other orbitals are linear combinations of the + # INITIAL doubly occupied and virtual columns, an S-orthonormal basis + # for the allowed space. Put occupied indices first: eigh returns + # increasing energies, so the lowest eigenvectors fill these slots. + doubly_occupied_indices = np.flatnonzero(occ == 2) + virtual_indices = np.flatnonzero(occ == 0) + free_indices = np.concatenate((doubly_occupied_indices, virtual_indices)) + free_orbitals = c[:, free_indices] + free_blocks = [free_indices[orbital_symmetry[free_indices] == irrep] + for irrep in np.unique(orbital_symmetry[free_indices])] + original_get_grad = frozen.get_grad + + def frozen_eig(fock, overlap, **kwargs): + # In each free block B, B.T @ S @ B = I. Thus this is an ordinary + # eigenproblem, not an AO generalized eigenproblem. Separate + # irreps must not mix, even if their eigenvalues are degenerate. + coeff = c.copy() + energies = np.einsum('pi,pi->i', c, fock @ c) + for indices in free_blocks: + basis = c[:, indices] + eigenvalues, rotation = eigh(basis.T @ fock @ basis) + coeff[:, indices] = basis @ rotation + energies[indices] = eigenvalues + if self.mol.symmetry: + coeff = lib.tag_array(coeff, orbsym=orbital_symmetry) + return energies, coeff + + def frozen_get_occ(mo_energy=None, mo_coeff=None): + # frozen_eig already placed the lowest free eigenvectors into + # the doubly occupied slots; keep the initial occupation labels. + return occ.copy() + + def frozen_get_grad(coeff, occupations, fock=None): + # PySCF packs all nonredundant ROHF rotations into a 1-D vector: + # alpha: occupied -> virtual; beta: doubly -> singly/virtual. + # Select only doubly occupied -> virtual entries of THAT vector. + # Derive masks from the supplied occupations because symmetry + # SCF can reorder orbitals during its finalization step. + alpha_rotations = (occupations == 0)[:, None] & (occupations > 0)[None, :] + beta_rotations = (occupations != 2)[:, None] & (occupations == 2)[None, :] + all_rotations = alpha_rotations | beta_rotations + doubly_to_virtual = (occupations == 0)[:, None] & (occupations == 2)[None, :] + free_gradient_entries = doubly_to_virtual[all_rotations] + gradient = original_get_grad(coeff, occupations, fock) + return gradient[free_gradient_entries] + + # These three hooks implement the constrained eigenproblem, filling, + # and convergence test. Density/Fock/energy updates stay in PySCF. + frozen.eig = frozen_eig + frozen.get_occ = frozen_get_occ + frozen.get_grad = frozen_get_grad + if frozen.diis: + # CDIIS must ignore rotations involving fixed singly occupied + # orbitals too. For the physical ROHF Fock, its virtual/doubly + # occupied block is (F_alpha + F_beta)/2. Multiplying by the + # occupation difference 2 gives exactly PySCF's orbital gradient: + # ||B.T @ (F D S - S D F) @ B||_F = sqrt(2) * ||g_free||_2, + # where D is the total (alpha + beta) density. The same-irrep + # restriction is applied by CDIIS using the orbsym tag below. + if self.mol.symmetry: + free_orbitals = lib.tag_array(free_orbitals, orbsym=orbital_symmetry[free_indices]) + frozen.diis.Corth = free_orbitals + frozen.conv_tol_grad = self.freeze_tol + frozen.max_cycle = self.freeze_max_cycle + logger.info(self, 'FR Freeze stage begins: fix both singly occupied orbitals; ' + 'relax doubly occupied/virtual orbitals') + frozen.kernel(dm0=frozen.make_rdm1(c, occ)) + self.freeze_cycles = frozen.cycles + self.freeze_converged = frozen.converged + if not self.freeze_converged: + raise RuntimeError('FR frozen SCF did not converge; release was not started') + + # PySCF mom_occ captures these occupied reference spaces once (IMOM). + # Resetting the reference after freeze incorporates spectator relaxation. + logger.info(self, 'FR Release stage begins: relax all orbitals with the post-freeze IMOM reference') + # Symmetry SCF finalization sorts columns AND occupations. Use both + # returned arrays together, rather than reusing initial column labels. + setocc = np.asarray((frozen.mo_occ > 0, frozen.mo_occ == 2), dtype=float) + _set_imom(release, frozen.mo_coeff, setocc) + release.kernel(dm0=frozen.make_rdm1()) + # With conv_check=False the ordinary kernel already checks energy + # AND orbital-gradient convergence. Honor its result, including any + # user-supplied check_convergence callback, without a second criterion. + # Retain the ordinary MF results and caches, but not stage-local hooks. + for key, value in release.__dict__.items(): + if key not in ('callback', 'get_occ', 'diis', 'DIIS'): + self.__dict__[key] = value + _set_imom(self, frozen.mo_coeff, setocc) + # Use the final occupation labels: IMOM and symmetry finalization + # can move the singly occupied orbitals to different column indices. + somo_overlap = c[:, occ == 1].T @ s @ self.mo_coeff[:, self.mo_occ == 1] + logger.info(self, 'FR target SOMO overlap singular values: %s', + np.linalg.svd(somo_overlap, compute_uv=False)) + return self.e_tot + + scf = kernel + + +hf.SCF.FR = lib.class_as_method(FR) diff --git a/src/nest/soscf/sgm.py b/src/nest/deltascf/sgm.py similarity index 67% rename from src/nest/soscf/sgm.py rename to src/nest/deltascf/sgm.py index dbf292b..d88ee2b 100644 --- a/src/nest/soscf/sgm.py +++ b/src/nest/deltascf/sgm.py @@ -115,92 +115,6 @@ def h_op(x): return g, h_op, h_diag -def gen_delta_curvature_rohf(mf, mo_coeff, mo_occ, with_symmetry=True): - '''Analytic directional curvature of Delta = g.T@g for real RO orbitals. - - Returns a callable curvature(x) -> (exact, gauss_newton), evaluated along - C(t) = C exp(t K(x)). The two values are 2*|g'|**2 + 2*g.T@g'' and - 2*|g'|**2, respectively. They coincide at an orbital stationary point. - - Supports ROHF and LDA/GGA ROKS (including hybrid exchange); NLC response - is omitted as in gen_g_hop_rohf. DFT requires third XC derivatives. - This is a diagnostic: one call computes two density responses and, for - DFT, a kxc contraction. It is not used by the SGM optimizer. - ''' - if numpy.iscomplexobj(mo_coeff): - raise NotImplementedError('Delta curvature requires real orbitals') - mol = mf.mol - c = numpy.asarray(mo_coeff) - mo_occ = numpy.asarray(mo_occ) - occ = numpy.asarray((mo_occ > 0, mo_occ == 2), dtype=float) - masks = (occ[:, :, None] == 0) & (occ[:, None, :] > 0) - unique = masks[0] | masks[1] - allowed = numpy.ones(numpy.count_nonzero(unique), dtype=bool) - if with_symmetry and mol.symmetry: - orbsym = hf_symm.get_orbsym(mol, mo_coeff) - allowed = (orbsym[:, None] == orbsym)[unique] - - def pack(f): - mat = numpy.where(masks[0], f[0], 0) + numpy.where(masks[1], f[1], 0) - return mat[unique] * allowed - - dm0 = mf.make_rdm1(mo_coeff, mo_occ) - vhf = mf.get_veff(mol, dm0) - fock = c.T @ (mf.get_hcore() + vhf) @ c - g = pack(fock) - vind = mf.gen_response((mo_coeff, mo_coeff), tuple(occ), hermi=1, with_nlc=False) - xctype = 'HF' - if isinstance(mf, hf.KohnShamDFT): - ni = mf._numint - xctype = ni._xc_type(mf.xc) - if xctype not in ('HF', 'LDA', 'GGA'): - raise NotImplementedError('Delta curvature supports HF, LDA and GGA') - ni.libxc.test_deriv_order(mf.xc, 3, raise_error=True) - - def curvature(x): - kappa = hf.unpack_uniq_var(numpy.asarray(x) * allowed, mo_occ) - # D(t) = C exp(tK) N exp(-tK) C.T. - dm1_mo = kappa * (occ[:, None, :] - occ[:, :, None]) - dm2_mo = kappa @ dm1_mo - dm1_mo @ kappa - dm1 = c @ dm1_mo @ c.T - dm2 = c @ dm2_mo @ c.T - v1 = vind(dm1) - v2 = vind(dm2) - - if xctype != 'HF': - # F'' = response(D'') + kxc[D', D']. eval_xc_eff differentiates - # with respect to rho and its Cartesian gradients, not sigma. - for ao, mask, weight, _ in ni.block_loop(mol, mf.grids, c.shape[0], 1): - ao_rho = ao[0] if xctype == 'LDA' else ao - rho0 = numpy.asarray([ni.eval_rho(mol, ao_rho, dm, mask, xctype, hermi=1) - for dm in dm0]) - rho1 = numpy.asarray([ni.eval_rho(mol, ao_rho, dm, mask, xctype, hermi=1) - for dm in dm1]) - kxc = ni.eval_xc_eff(mf.xc, rho0, deriv=3, xctype=xctype, spin=1)[3] - if xctype == 'LDA': - rho1 = rho1[:, None, :] - wv = numpy.einsum('axbyczg,byg,czg->axg', kxc, rho1, rho1, optimize=True) - wv *= weight - for spin in range(2): - if xctype == 'LDA': - v2[spin] += ao[0].T @ (wv[spin, 0, :, None] * ao[0]) - else: - wv[spin, 0] *= .5 - aow = numpy.einsum('xgi,xg->gi', ao, wv[spin]) - mat = ao[0].T @ aow - v2[spin] += mat + mat.T - - v1_mo = c.T @ v1 @ c - comm = fock @ kappa - kappa @ fock - g1 = pack(comm + v1_mo) - g2 = pack(comm @ kappa - kappa @ comm - + 2*(v1_mo @ kappa - kappa @ v1_mo) + c.T @ v2 @ c) - gauss_newton = 2 * numpy.dot(g1, g1) - return gauss_newton + 2*numpy.dot(g, g2), gauss_newton - - return curvature - - class LBFGSHistory: '''Bounded L-BFGS history for inverse-Hessian two-loop recursion.''' @@ -275,8 +189,8 @@ class SGM(lib.StreamObject): Convergence threshold on sqrt(Delta), where Delta = g_orb.T@g_orb. An explicitly set PySCF conv_tol_grad takes precedence over tol. gradient_scale - Scalar c from the SGM paper. Q-Chem's DeltaSCF driver uses 0.75 - by default. It scales the gradient seen by L-BFGS and the L-BFGS + Scalar c from the SGM paper (default here: 0.75). + It scales the gradient seen by L-BFGS and the L-BFGS y_k history; the line search still uses the true Delta directional derivative. somo_overlap_tol @@ -284,17 +198,6 @@ class SGM(lib.StreamObject): this threshold relative to the initial or previous orbitals. Default 0.7 corresponds to a largest principal angle of about 46 degrees. This is a diagnostic, not a state constraint. - preconditioner - 'gap' (default) uses 2*h_diag**2. 'fock' retains the full Fock - commutator in J but omits density response; it needs no extra - response evaluations. 'gn' uses the exact diagonal - of 2*J.T@J, where J is the orbital-gradient Jacobian. This needs - one Hessian-vector product per orbital variable, but no third - derivatives and no stored full Hessian. - preconditioner_update - Refresh the 'gn' diagonal every this many accepted steps. - Default 1; 0 computes it only at the initial point. L-BFGS - updates still run every step. 'gap' and 'fock' are always refreshed. The returned object also inherits the input ROHF/ROKS class. Use mf.SGM().set(...).run(mo_coeff, mo_occ), or SGM(mf).kernel(...). @@ -308,14 +211,11 @@ class SGM(lib.StreamObject): lbfgs_memory = getattr(__config__, 'sgm_lbfgs_memory', 8) gradient_scale = getattr(__config__, 'sgm_gradient_scale', 0.75) somo_overlap_tol = getattr(__config__, 'sgm_somo_overlap_tol', 0.7) - preconditioner = 'gap' - preconditioner_update = 1 canonicalization = getattr(__config__, 'soscf_newton_ah_SOSCF_canonicalization', True) _keys = {'max_cycle', 'tol', 'lbfgs_memory', 'gradient_scale', - 'canonicalization', 'somo_overlap_tol', 'preconditioner', - 'preconditioner_update'} + 'canonicalization', 'somo_overlap_tol'} gen_g_hop = staticmethod(gen_g_hop_rohf) @@ -341,48 +241,14 @@ def dump_flags(self, verbose=None): self.tol if self.conv_tol_grad is None else self.conv_tol_grad) log.info('SGM gradient scale = %g', self.gradient_scale) log.info('SGM L-BFGS memory = %d', self.lbfgs_memory) - log.info('SGM preconditioner = %s', self.preconditioner) - if self.preconditioner == 'gn': - log.info('SGM preconditioner update interval = %d (0: initial only)', - self.preconditioner_update) log.info('SGM SOMO overlap warning threshold = %g', self.somo_overlap_tol) log.info('SGM canonicalization = %s', self.canonicalization) return self - def _preconditioner(self, h_diag, h_op=None, mo_coeff=None, mo_occ=None, fock=None): + def _preconditioner(self, h_diag): # Diagonal, frozen-Fock approximation to 2*J.T@J for Delta = g.T@g. # J = d g / d kappa; this is not the full Delta Hessian away from g=0. - denom = 2.0 * h_diag ** 2 - if self.preconditioner == 'gn': - direction = numpy.zeros_like(h_diag) - for i in range(h_diag.size): - direction[i] = 1 - column = h_op(direction) - denom[i] = 2.0 * numpy.dot(column, column) - direction[i] = 0 - elif self.preconditioner == 'fock': - f = mo_coeff.T @ numpy.asarray((fock.focka, fock.fockb)) @ mo_coeff - occ = numpy.asarray((mo_occ > 0, mo_occ == 2)) - masks = (~occ[:, :, None]) & occ[:, None, :] - unique = masks[0] | masks[1] - allowed = numpy.ones(h_diag.size, dtype=bool) - if self.mol.symmetry: - orbsym = hf_symm.get_orbsym(self.mol, mo_coeff) - allowed = (orbsym[:, None] == orbsym)[unique] - for index, (a, i) in enumerate(zip(*numpy.where(unique))): - if not allowed[index]: - denom[index] = 0 - continue - # Column of J_F: [F, K_ai], evaluated using the two nonzero - # entries of K_ai instead of a dense matrix multiplication. - comm = numpy.zeros_like(f) - comm[:, :, i] += f[:, :, a] - comm[:, :, a] -= f[:, :, i] - comm[:, a, :] -= f[:, i, :] - comm[:, i, :] += f[:, a, :] - column = numpy.sum(comm * masks, axis=0)[unique] * allowed - denom[index] = 2.0 * numpy.dot(column, column) - denom = numpy.maximum(denom, _PRECOND_FLOOR) + denom = numpy.maximum(2.0 * h_diag ** 2, _PRECOND_FLOOR) return numpy.minimum(1.0 / denom, _MAX_PRECOND) def _trial_mo(self, mo_coeff, mo_occ, step): @@ -397,7 +263,7 @@ def _exact_sgm_state(self, mo_coeff, mo_occ, fock_ao=None, h1e=None): g_orb, h_op, h_diag = self.gen_g_hop(self, mo_coeff, mo_occ, fock_ao, h1e) delta = numpy.dot(g_orb, g_orb) grad_delta = 2.0 * h_op(g_orb) - return h_diag, delta, grad_delta, h_op + return h_diag, delta, grad_delta def _line_search(self, mo_coeff, mo_occ, direction, delta, grad_delta, h1e, s1e): @@ -453,10 +319,6 @@ def kernel(self, mo_coeff=None, mo_occ=None): tol = self.tol if self.conv_tol_grad is None else self.conv_tol_grad if not numpy.isfinite(self.gradient_scale) or self.gradient_scale <= 0: raise ValueError('gradient_scale must be finite and positive') - if self.preconditioner not in ('gap', 'fock', 'gn'): - raise ValueError("preconditioner must be 'gap', 'fock' or 'gn'") - if not isinstance(self.preconditioner_update, (int, numpy.integer)) or self.preconditioner_update < 0: - raise ValueError('preconditioner_update must be a nonnegative integer') self.dump_flags() mo_guess = mo_coeff.copy() @@ -468,7 +330,7 @@ def kernel(self, mo_coeff=None, mo_occ=None): e_tot = self.energy_tot(dm, h1e, vhf) t0 = (logger.process_clock(), logger.perf_counter()) - h_diag, delta, grad_delta, h_op = self._exact_sgm_state(mo_coeff, mo_occ, fock, h1e) + h_diag, delta, grad_delta = self._exact_sgm_state(mo_coeff, mo_occ, fock, h1e) opt_grad_delta = self.gradient_scale * grad_delta history = LBFGSHistory(self.lbfgs_memory, _CURVATURE_TOL) @@ -481,9 +343,7 @@ def kernel(self, mo_coeff=None, mo_occ=None): if norm_g < tol: break - if (cycle == 0 or self.preconditioner in ('gap', 'fock') - or (self.preconditioner_update and cycle % self.preconditioner_update == 0)): - h0_inv = self._preconditioner(h_diag, h_op, mo_coeff, mo_occ, fock) + h0_inv = self._preconditioner(h_diag) direction = self._descent_direction(history, grad_delta, opt_grad_delta, h0_inv, log, cycle) @@ -508,7 +368,7 @@ def kernel(self, mo_coeff=None, mo_occ=None): dm = self.make_rdm1(mo_trial, mo_occ) e_new = self.energy_tot(dm, h1e, vhf_new) - h_diag_new, delta_new, grad_delta_new, h_op = self._exact_sgm_state( + h_diag_new, delta_new, grad_delta_new = self._exact_sgm_state( mo_trial, mo_occ, fock, h1e) opt_grad_delta_new = self.gradient_scale * grad_delta_new @@ -528,7 +388,7 @@ def kernel(self, mo_coeff=None, mo_occ=None): if min(somo_initial[-1], somo_previous[-1]) < self.somo_overlap_tol: log.warn('SGM: SOMO subspace deviation at iter %d: ' 'min overlap initial = %.6f, previous = %.6f ' - '(threshold %.3f); inspect state character', + '(threshold %.3f)', cycle, somo_initial[-1], somo_previous[-1], self.somo_overlap_tol) delta_e = e_new - e_tot @@ -570,6 +430,8 @@ def kernel(self, mo_coeff=None, mo_occ=None): self.mo_energy = mo_energy self.e_tot = e_tot self._finalize() + if self.chkfile: + self.dump_chk(self.chkfile) return self.e_tot scf = kernel diff --git a/src/nest/deltascf/tests/test_fr.py b/src/nest/deltascf/tests/test_fr.py new file mode 100644 index 0000000..29b98e3 --- /dev/null +++ b/src/nest/deltascf/tests/test_fr.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from unittest import mock + +import numpy as np +from scipy.linalg import expm +from pyscf import gto, scf, symm +from nest.deltascf import fr + + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.mol = gto.M(atom='O 0 0 0; H 0 -.757 .587; H 0 .757 .587', + basis='6-31g', spin=2, verbose=0) + + def test_freeze_release(self): + for xc, use_symmetry in ((None, False), ('BHandHLYP', False), + (None, True), ('BHandHLYP', True)): + with self.subTest(xc=xc, symmetry=use_symmetry): + mol = gto.M(atom=self.mol.atom, basis=self.mol.basis, spin=2, + symmetry=use_symmetry, verbose=0) + ground_mol = mol.copy() + ground_mol.spin = 0 + ground = ground_mol.RHF() if xc is None else ground_mol.RKS(xc=xc) + mf = mol.ROHF() if xc is None else mol.ROKS(xc=xc) + if xc is not None: + ground.grids.level = mf.grids.level = 0 + ground.kernel() + occ = ground.mo_occ.copy() + occ[[4, 5]] = 1 + c = ground.mo_coeff.copy() + s = mf.get_ovlp() + stages = [] + + def callback(env): + stages.append(env['fr_stage']) + coeff = env['mo_coeff'] + np.testing.assert_allclose(coeff.T @ s @ coeff, np.eye(occ.size), atol=1e-10) + if env['fr_stage'] == 'freeze': + np.testing.assert_array_equal(coeff[:, occ == 1], c[:, occ == 1]) + free_basis = env['mf_diis'].Corth + np.testing.assert_allclose(free_basis.T @ s @ c[:, occ == 1], 0, atol=1e-10) + # Independently verify that projected CDIIS residual + # and the physical frozen gradient have the same zero. + dm_total = env['dm'].sum(axis=0) + fock = env['fock'] # Physical Fock at CURRENT density. + residual = free_basis.T @ (fock @ dm_total @ s - s @ dm_total @ fock) @ free_basis + full_gradient = mf.get_grad(coeff, occ, fock) + all_rotations = ((occ == 0)[:, None] & (occ > 0)) | ((occ != 2)[:, None] & (occ == 2)) + free_entries = ((occ == 0)[:, None] & (occ == 2))[all_rotations] + frozen_gradient = full_gradient[free_entries] + self.assertAlmostEqual(np.linalg.norm(residual), + np.sqrt(2) * np.linalg.norm(frozen_gradient), places=10) + np.testing.assert_allclose(env['mf'].get_grad(coeff, occ, fock), frozen_gradient, atol=1e-12) + if stages.count('freeze') == 1: + # Energy finite differences are independent of + # the gradient packing and the DIIS construction. + direction = np.zeros_like(full_gradient) + direction[free_entries] = frozen_gradient + direction /= np.linalg.norm(direction) + rotation = scf.hf.unpack_uniq_var(direction, occ) + eps = 1e-4 + plus = mf.energy_tot(dm=mf.make_rdm1(coeff @ expm(eps * rotation), occ)) + minus = mf.energy_tot(dm=mf.make_rdm1(coeff @ expm(-eps * rotation), occ)) + self.assertAlmostEqual((plus-minus)/(2*eps), 2 * full_gradient @ direction, places=6) + else: + self.assertEqual(np.count_nonzero(env['mo_occ'] == 1), 2) + + opt = mf.FR().set(conv_tol=1e-10, conv_tol_grad=1e-6, max_cycle=100, + callback=callback) + self.assertIsInstance(opt, fr.FR) + self.assertTrue(opt.istype('ROHF')) + # First-order implementation must not require response kernels. + with mock.patch.object(type(mf), 'gen_response', side_effect=AssertionError('Hessian used')): + opt.kernel(c, occ) + self.assertTrue(opt.freeze_converged) + self.assertTrue(opt.converged) + self.assertIn('freeze', stages) + self.assertIn('release', stages) + self.assertLess(np.linalg.norm(opt.get_grad(opt.mo_coeff, opt.mo_occ)), 1e-6) + self.assertAlmostEqual(opt.e_tot, opt.energy_tot(dm=opt.make_rdm1()), places=10) + np.testing.assert_array_equal(opt.get_occ(opt.mo_energy, opt.mo_coeff), opt.mo_occ) + # Independent conventional triplet SCF reference for this valence state. + ref = mf.copy().set(conv_tol=1e-11, conv_tol_grad=1e-7) + ref.kernel() + self.assertTrue(ref.converged) + self.assertAlmostEqual(opt.e_tot, ref.e_tot, places=7) + self.assertIsNone(mf.mo_coeff) + self.assertIs(opt.callback, callback) + self.assertIs(opt.FR(), opt) + if use_symmetry: + initial_symmetry = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, c) + final_symmetry = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, opt.mo_coeff) + np.testing.assert_array_equal(opt.mo_coeff.orbsym, final_symmetry) + for irrep in mol.irrep_id: + for occupied in (lambda o: o > 0, lambda o: o == 2): + self.assertEqual(np.count_nonzero(occupied(occ) & (initial_symmetry == irrep)), + np.count_nonzero(occupied(opt.mo_occ) & (final_symmetry == irrep))) + + def test_failed_freeze_does_not_release(self): + ground_mol = self.mol.copy() + ground_mol.spin = 0 + ground = ground_mol.RHF().run() + occ = ground.mo_occ.copy() + occ[[0, 5]] = 1 + stages = [] + opt = fr.FR(self.mol.ROHF()).set(freeze_max_cycle=1, freeze_tol=1e-12, + callback=lambda env: stages.append(env['fr_stage'])) + with self.assertRaisesRegex(RuntimeError, 'release was not started'): + opt.kernel(ground.mo_coeff, occ) + self.assertFalse(opt.converged) + self.assertFalse(opt.freeze_converged) + self.assertEqual(stages, ['freeze']) + + def test_reject_invalid_reference(self): + mf = self.mol.ROHF().run() + with self.assertRaisesRegex(ValueError, 'S-orthonormal'): + mf.FR().kernel(mf.mo_coeff * 2, mf.mo_occ) + with self.assertRaisesRegex(ValueError, 'two SOMOs'): + mf.FR().kernel(mf.mo_coeff, np.zeros_like(mf.mo_occ)) + with self.assertRaises(NotImplementedError): + fr.FR(self.mol.UHF()) + + def test_symmetry_finalization_reorders_columns(self): + mol = gto.M(atom=self.mol.atom, basis=self.mol.basis, spin=2, symmetry=True, verbose=0) + mf = mol.ROHF().set(conv_tol=1e-11, conv_tol_grad=1e-8).run() + order = np.arange(mf.mo_occ.size)[::-1] + # In particular, singly occupied columns are no longer immediately + # after doubly occupied ones. PySCF finalization will reorder them. + coeff = np.asarray(mf.mo_coeff[:, order]) + occupations = mf.mo_occ[order] + opt = mf.FR().set(conv_tol=1e-10, conv_tol_grad=1e-7) + opt.kernel(coeff, occupations) + self.assertTrue(opt.converged) + self.assertAlmostEqual(opt.e_tot, mf.e_tot, places=9) + np.testing.assert_allclose(opt.make_rdm1(), mf.make_rdm1(), atol=1e-6) + np.testing.assert_array_equal(opt.get_occ(opt.mo_energy, opt.mo_coeff), opt.mo_occ) + + def test_symmetry_rejects_mixed_orbitals_and_conflicting_counts(self): + mol = gto.M(atom=self.mol.atom, basis=self.mol.basis, spin=2, symmetry=True, verbose=0) + mf = mol.ROHF().run() + labels = mf.mo_coeff.orbsym + first = 0 + second = np.flatnonzero(labels != labels[first])[0] + mixed = mf.mo_coeff.copy() + mixed[:, first] = (mf.mo_coeff[:, first] + mf.mo_coeff[:, second]) / np.sqrt(2) + mixed[:, second] = (mf.mo_coeff[:, first] - mf.mo_coeff[:, second]) / np.sqrt(2) + with self.assertRaisesRegex(ValueError, 'not symmetrized'): + mf.FR().kernel(mixed, mf.mo_occ) + name = mol.irrep_name[mol.irrep_id.index(int(labels[first]))] + mf.irrep_nelec = {name: (0, 0)} + with self.assertRaisesRegex(ValueError, 'conflict'): + mf.FR().kernel(mf.mo_coeff, mf.mo_occ) + + def test_imom_rejects_nonnested_spaces(self): + mf = self.mol.ROHF().run() + setocc = np.asarray((mf.mo_occ > 0, mf.mo_occ == 2), dtype=float) + + def invalid_imom(mf, coeff, occupations): + invalid = np.zeros_like(mf.mo_occ) + invalid[:3] = 2 + invalid[3:7] = 1 # Same electron count, but four SOMOs. + mf.get_occ = lambda *args: invalid + + with mock.patch.object(scf.addons, 'mom_occ', side_effect=invalid_imom): + fr._set_imom(mf, mf.mo_coeff, setocc) + with self.assertRaisesRegex(RuntimeError, 'nonnested'): + mf.get_occ(mf.mo_energy, mf.mo_coeff) + + +if __name__ == '__main__': + unittest.main() diff --git a/src/nest/deltascf/tests/test_sgm.py b/src/nest/deltascf/tests/test_sgm.py new file mode 100644 index 0000000..386f183 --- /dev/null +++ b/src/nest/deltascf/tests/test_sgm.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python +# Copyright 2026 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +import numpy + +from pyscf import gto +from nest.deltascf import sgm + + +# Q-Chem benchmark, HCHO Rydberg excitation: +# +# $rem +# METHOD BHHLYP +# BASIS cc-pvdz +# SYMMETRY false +# SYM_IGNORE true +# NO_REORIENT True +# XC_GRID 000075000302 +# UNRESTRICTED false +# delta_scf true +# $end +# +# $delta_scf +# triplet restricted +# triplet_SCF_algorithm SGM +# somo_1 8 +# somo_2 15 +# $end +# +# Q-Chem output: +# Restricted open-shell triplet state = -113.131417 Ha +# SOMO(1): best initial orbital 8 overlap 0.977366 +# SOMO(2): best initial orbital 9 overlap 0.991702 + + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + atom = ''' + C 0.00000000 0.00000000 -1.13947666 + O 0.00000000 0.00000000 1.14402883 + H 0.00000000 1.76627623 -2.23398653 + H 0.00000000 -1.76627623 -2.23398653 + ''' + cls.mol0 = gto.M(atom=atom, charge=0, spin=0, basis='cc-pvdz', + verbose=0, output='/dev/null') + cls.mf0 = cls.mol0.RKS(xc='BHandHLYP') + cls.mf0.grids.atom_grid = (75, 302) + cls.mf0.kernel() + + setocc = cls.mf0.to_uks().mo_occ + setocc[1][7] -= 1 + setocc[0][14] += 1 + cls.ro_occ = setocc[0] + setocc[1] + + cls.mol1 = gto.M(atom=atom, charge=0, spin=2, basis='cc-pvdz', + verbose=0, output='/dev/null') + + @classmethod + def tearDownClass(cls): + cls.mol0.stdout.close() + cls.mol1.stdout.close() + del cls.mol0, cls.mol1, cls.mf0, cls.ro_occ + + def test_delta_gradient(self): + mf = self.mol1.ROKS(xc='BHandHLYP') + mf.grids.atom_grid = (75, 302) + mf.mo_coeff = self.mf0.mo_coeff.copy() + mf.mo_occ = self.ro_occ.copy() + + opt = sgm.SGM(mf) + _h_diag, _delta, grad_delta = opt._exact_sgm_state(mf.mo_coeff, + mf.mo_occ) + rng = numpy.random.default_rng(12) + direction = rng.normal(size=grad_delta.size) + direction /= numpy.linalg.norm(direction) + + eps = 1e-4 + mo_plus = opt._trial_mo(mf.mo_coeff, mf.mo_occ, eps * direction) + mo_minus = opt._trial_mo(mf.mo_coeff, mf.mo_occ, -eps * direction) + g_plus = mf.get_grad(mo_plus, mf.mo_occ) + g_minus = mf.get_grad(mo_minus, mf.mo_occ) + fd = (numpy.dot(g_plus, g_plus) - numpy.dot(g_minus, g_minus)) / (2*eps) + analytic = numpy.dot(grad_delta, direction) + scale = max(abs(fd), abs(analytic), 1e-12) + + self.assertLess(abs(fd - analytic) / scale, 1e-6, + 'fd = %.12e, analytic = %.12e' % (fd, analytic)) + + def test_hcho_rydberg_triplet(self): + mf = self.mol1.ROKS(xc='BHandHLYP') + mf.grids.atom_grid = (75, 302) + mf.mo_coeff = self.mf0.mo_coeff.copy() + mf.mo_occ = self.ro_occ.copy() + + opt = sgm.SGM(mf) + opt.max_cycle = 80 + opt.verbose = 0 + e_tot = opt.kernel() + + self.assertTrue(opt.converged) + self.assertLess(abs(e_tot - -113.131417), 1e-4) + self.assertLess(abs(e_tot - -113.131420102713), 1e-6) + + +if __name__ == '__main__': + print('Full tests for deltascf.sgm') + unittest.main() diff --git a/src/nest/soscf/tests/test_sgm.py b/src/nest/soscf/tests/test_sgm.py deleted file mode 100644 index 0cb563b..0000000 --- a/src/nest/soscf/tests/test_sgm.py +++ /dev/null @@ -1,305 +0,0 @@ -#!/usr/bin/env python -# Copyright 2026 The NEST Developers. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import io -import unittest -from unittest import mock -import numpy - -from pyscf import gto, lib -from pyscf.lib import logger -from nest.soscf import sgm - - -# Q-Chem benchmark, HCHO Rydberg excitation: -# -# $rem -# METHOD BHHLYP -# BASIS cc-pvdz -# SYMMETRY false -# SYM_IGNORE true -# NO_REORIENT True -# XC_GRID 000075000302 -# UNRESTRICTED false -# delta_scf true -# $end -# -# $delta_scf -# triplet restricted -# triplet_SCF_algorithm SGM -# somo_1 8 -# somo_2 15 -# $end -# -# Q-Chem output: -# Restricted open-shell triplet state = -113.131417 Ha -# SOMO(1): best initial orbital 8 overlap 0.977366 -# SOMO(2): best initial orbital 9 overlap 0.991702 - - -class KnownValues(unittest.TestCase): - @classmethod - def setUpClass(cls): - atom = ''' - C 0.00000000 0.00000000 -1.13947666 - O 0.00000000 0.00000000 1.14402883 - H 0.00000000 1.76627623 -2.23398653 - H 0.00000000 -1.76627623 -2.23398653 - ''' - cls.mol0 = gto.M(atom=atom, charge=0, spin=0, basis='cc-pvdz', - verbose=0, output='/dev/null') - cls.mf0 = cls.mol0.RKS(xc='BHandHLYP') - cls.mf0.grids.atom_grid = (75, 302) - cls.mf0.kernel() - - setocc = cls.mf0.to_uks().mo_occ - setocc[1][7] -= 1 - setocc[0][14] += 1 - cls.ro_occ = setocc[0] + setocc[1] - - cls.mol1 = gto.M(atom=atom, charge=0, spin=2, basis='cc-pvdz', - verbose=0, output='/dev/null') - - @classmethod - def tearDownClass(cls): - cls.mol0.stdout.close() - cls.mol1.stdout.close() - del cls.mol0, cls.mol1, cls.mf0, cls.ro_occ - - def test_delta_gradient(self): - mf = self.mol1.ROKS(xc='BHandHLYP') - mf.grids.atom_grid = (75, 302) - mf.mo_coeff = self.mf0.mo_coeff.copy() - mf.mo_occ = self.ro_occ.copy() - - opt = sgm.SGM(mf) - _h_diag, _delta, grad_delta, _hop = opt._exact_sgm_state(mf.mo_coeff, - mf.mo_occ) - rng = numpy.random.default_rng(12) - direction = rng.normal(size=grad_delta.size) - direction /= numpy.linalg.norm(direction) - - eps = 1e-4 - mo_plus = opt._trial_mo(mf.mo_coeff, mf.mo_occ, eps * direction) - mo_minus = opt._trial_mo(mf.mo_coeff, mf.mo_occ, -eps * direction) - g_plus = mf.get_grad(mo_plus, mf.mo_occ) - g_minus = mf.get_grad(mo_minus, mf.mo_occ) - fd = (numpy.dot(g_plus, g_plus) - numpy.dot(g_minus, g_minus)) / (2*eps) - analytic = numpy.dot(grad_delta, direction) - scale = max(abs(fd), abs(analytic), 1e-12) - - self.assertLess(abs(fd - analytic) / scale, 1e-6, - 'fd = %.12e, analytic = %.12e' % (fd, analytic)) - - def test_hcho_rydberg_triplet(self): - for tol, max_cycle in ((1e-4, 80), (1e-7, 200)): - with self.subTest(tol=tol): - mf = self.mol1.ROKS(xc='BHandHLYP') - mf.grids.atom_grid = (75, 302) - mf.mo_coeff = self.mf0.mo_coeff.copy() - mf.mo_occ = self.ro_occ.copy() - - opt = mf.SGM().set(tol=tol, max_cycle=max_cycle) - e_tot = opt.kernel() - - self.assertTrue(opt.converged) - self.assertLess(numpy.linalg.norm(opt.get_grad(opt.mo_coeff, opt.mo_occ)), tol) - self.assertLess(abs(e_tot - -113.131417), 1e-4) - self.assertLess(abs(e_tot - -113.131420102713), 1e-6) - - -class OptimizerChecks(unittest.TestCase): - @classmethod - def setUpClass(cls): - cls.mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='sto-3g', - spin=2, verbose=0) - cls.mf = cls.mol.ROHF().run(conv_tol=1e-12) - cls.occ = cls.mf.mo_occ.copy() - opt = cls.mf.SGM() - size = cls.mf.get_grad(cls.mf.mo_coeff, cls.occ).size - step = numpy.random.default_rng(42).normal(size=size)*0.02 - cls.guess = opt._trial_mo(cls.mf.mo_coeff, cls.occ, step) - - def test_last_cycle_and_fock_reuse(self): - opt = self.mf.SGM().set(max_cycle=1, canonicalization=False) - h1e, s1e = opt.get_hcore(), opt.get_ovlp() - hdiag, delta, grad, _hop = opt._exact_sgm_state(self.guess, self.occ) - direction = -opt._preconditioner(hdiag)*opt.gradient_scale*grad - trial = opt._line_search(self.guess, self.occ, direction, delta, grad, h1e, s1e) - self.assertIsNotNone(trial) - alpha, _, cnew, fock, _ = trial - final_norm = numpy.linalg.norm(opt.get_grad(cnew, self.occ, fock)) - self.assertLess(final_norm, numpy.sqrt(delta)) - opt.tol = (numpy.sqrt(delta)+final_norm)/2 - trace = [] - opt.callback = lambda env: trace.append(env['delta']) - with mock.patch.object(opt, 'get_veff', wraps=opt.get_veff) as veff: - opt.kernel(self.guess, self.occ) - self.assertTrue(opt.converged) - self.assertEqual(opt.cycles, 1) - self.assertEqual(len(trace), 1) - # One initial build and one per line-search trial, none for the - # accepted-point energy or analytic gradient. - self.assertEqual(veff.call_count, 2+round(-numpy.log2(alpha))) - numpy.testing.assert_allclose(opt.mo_coeff, cnew, atol=1e-12) - - def test_pyscf_interfaces(self): - for mf in (self.mf, self.mol.ROKS(xc='BHandHLYP')): - opt = mf.SGM().set(max_cycle=0, canonicalization=False) - self.assertIsInstance(opt, lib.StreamObject) - self.assertIsInstance(opt, mf.__class__) - self.assertIs(sgm.SGM(opt), opt) - self.assertIs(opt.SGM(), opt) - self.assertIs(opt.scf.__func__, opt.kernel.__func__) - self.assertIs(opt.run(self.guess, self.occ), opt) - plain = opt.undo_sgm() - self.assertIsInstance(plain, mf.__class__) - self.assertNotIsInstance(plain, sgm.SGM) - numpy.testing.assert_allclose(plain.mo_coeff, opt.mo_coeff) - opt = self.mf.SGM().set(max_cycle=0, tol=1e-15, conv_tol_grad=100) - opt.scf(self.guess, self.occ) - self.assertTrue(opt.converged) - - def test_somo_subspace(self): - s = numpy.diag([1., 2., 3., 4.]) - c = numpy.diag(1/numpy.sqrt(s.diagonal())) - occ = numpy.array([2,1,1,0]) - theta = 0.83 - u = numpy.array([[numpy.cos(theta), -numpy.sin(theta)], - [numpy.sin(theta), numpy.cos(theta)]]) - rotated = c.copy() - rotated[:,1:3] = c[:,1:3]@u - numpy.testing.assert_allclose(sgm._somo_overlaps(c,rotated,occ,s), 1, atol=1e-14) - leaked = c.copy() - leaked[:,2] = numpy.cos(theta)*c[:,2]+numpy.sin(theta)*c[:,3] - numpy.testing.assert_allclose(sgm._somo_overlaps(c,leaked,occ,s), - [1,numpy.cos(theta)], atol=1e-14) - leaked[:,1:3] = leaked[:,1:3]@u - numpy.testing.assert_allclose(sgm._somo_overlaps(rotated,leaked,occ,s), - [1,numpy.cos(theta)], atol=1e-14) - - def test_lbfgs_curvature_scale(self): - for scale in (1., 1e-8): - history = sgm.LBFGSHistory() - step = scale*numpy.array([1., 2.]) - diff = scale*numpy.array([2., 3.]) - self.assertTrue(history.push(step, diff)[0]) - self.assertFalse(history.push(step, -diff)[0]) - self.assertFalse(history.push(step, numpy.zeros(2))[0]) - direction = history.direction(diff, numpy.ones(2)) - self.assertLess(numpy.dot(direction, diff), 0) - - def test_somo_warning_each_step(self): - opt = self.mf.SGM().set(max_cycle=2, tol=1e-12, somo_overlap_tol=0.999999999, - verbose=logger.WARN, stdout=io.StringIO()) - opt.kernel(self.guess,self.occ) - self.assertEqual(opt.cycles, 2) - output = opt.stdout.getvalue() - self.assertIn('SOMO subspace deviation at iter 0', output) - self.assertIn('SOMO subspace deviation at iter 1', output) - - def test_frozen_fock_preconditioner(self): - # Remove density response and off-diagonal Fock terms to isolate - # the scale of the mean-field preconditioner in all RO blocks. - opt = self.mf.SGM() - nmo = len(self.occ) - c = numpy.eye(nmo) - fa = numpy.diag(numpy.linspace(-2,3,nmo)) - fb = numpy.diag(numpy.linspace(-1,4,nmo)**3) - fock = lib.tag_array((fa+fb)/2, focka=fa, fockb=fb) - _, _, hdiag = opt.gen_g_hop(opt,c,self.occ,fock) - eps = 1e-4 - curvature = [] - for i in range(hdiag.size): - step = numpy.zeros(hdiag.size) - step[i] = eps - gp = opt.get_grad(opt._trial_mo(c,self.occ,step), self.occ, fock) - gm = opt.get_grad(opt._trial_mo(c,self.occ,-step), self.occ, fock) - curvature.append((gp@gp+gm@gm)/eps**2) - numpy.testing.assert_allclose(curvature,1/opt._preconditioner(hdiag),rtol=1e-6) - - def test_analytic_delta_curvature(self): - for xc in (None, 'LDA,VWN', 'BHandHLYP'): - with self.subTest(xc=xc): - mf = self.mol.ROHF() if xc is None else self.mol.ROKS(xc=xc) - if xc is not None: - mf.grids.level = 1 - opt = mf.SGM() - # A nonstationary reference exercises the g.T@g'' term, - # including third XC derivatives for LDA/GGA. - c = self.guess - g, hop, _ = opt.gen_g_hop(mf, c, self.occ) - curvature = sgm.gen_delta_curvature_rohf(mf, c, self.occ) - rng = numpy.random.default_rng(23) - for _ in range(2): - x = rng.normal(size=g.size) - x /= numpy.linalg.norm(x) - exact, gn = curvature(x) - self.assertAlmostEqual(gn, 2*numpy.linalg.norm(hop(x))**2, places=9) - eps = 5e-4 - plus = opt.get_grad(opt._trial_mo(c, self.occ, eps*x), self.occ) - minus = opt.get_grad(opt._trial_mo(c, self.occ, -eps*x), self.occ) - fd = (plus@plus + minus@minus - 2*g@g)/eps**2 - self.assertLess(abs(fd-exact)/max(1., abs(exact)), 2e-6) - - def test_delta_curvature_at_stationary_point(self): - c = self.mf.mo_coeff - g, hop, _ = self.mf.SGM().gen_g_hop(self.mf, c, self.occ) - curvature = sgm.gen_delta_curvature_rohf(self.mf, c, self.occ) - x = numpy.random.default_rng(12).normal(size=g.size) - x /= numpy.linalg.norm(x) - exact, gn = curvature(x) - self.assertLess(numpy.linalg.norm(g), 1e-6) - self.assertLess(abs(exact-gn)/max(1., gn), 1e-6) - - def test_response_preconditioners(self): - opt = self.mf.SGM() - dm = opt.make_rdm1(self.guess, self.occ) - fock = opt.get_fock(dm=dm) - _, hop, hd = opt.gen_g_hop(opt, self.guess, self.occ, fock) - # Independent finite differences of the orbital gradient give J's - # columns, including off-diagonal response contributions. - eps = 1e-5 - for mode in ('gn', 'fock'): - with self.subTest(mode=mode): - opt.preconditioner = mode - inverse = opt._preconditioner(hd, hop, self.guess, self.occ, fock) - fd_diagonal = [] - frozen = fock if mode == 'fock' else None - for i in range(hd.size): - step = numpy.zeros(hd.size) - step[i] = eps - gp = opt.get_grad(opt._trial_mo(self.guess, self.occ, step), self.occ, frozen) - gm = opt.get_grad(opt._trial_mo(self.guess, self.occ, -step), self.occ, frozen) - column = (gp-gm)/(2*eps) - fd_diagonal.append(2*numpy.dot(column, column)) - numpy.testing.assert_allclose(1/inverse, fd_diagonal, rtol=1e-7) - - def test_gn_preconditioner_convergence_and_refresh(self): - for mode, interval in (('gn', 0), ('gn', 1), ('gn', 3), ('fock', 1)): - with self.subTest(mode=mode, interval=interval): - opt = self.mf.SGM().set(preconditioner=mode, preconditioner_update=interval, - tol=1e-7, max_cycle=100) - with mock.patch.object(opt, '_preconditioner', wraps=opt._preconditioner) as precond: - opt.kernel(self.guess, self.occ) - self.assertTrue(opt.converged) - self.assertLess(abs(opt.e_tot-self.mf.e_tot), 1e-9) - expected = 1 if interval == 0 else 1+(opt.cycles-1)//interval - self.assertEqual(precond.call_count, expected) - - -if __name__ == '__main__': - print('Full tests for soscf.sgm') - unittest.main()