fix(mcscf): include spin penalty in GASSCF Newton CI derivatives

Author: yideng2025

pyscf/mcscf/fci_gas.py

@@ -236,6 +236,32 @@ def _spin_penalty_parameters(solver, norb, nelec):
     return shift, target, minimum, tuple(float(x) for x in eigenvalues)
 
 
+def _spin_penalty_action(contract_ss, vector, parameters):
+    """Apply the GAS spin penalty, without adding the physical Hamiltonian.
+
+    Shared by the fixed-orbital solver and the Newton CI derivatives. The
+    supplied contraction owns/reuses its spin plan; the input is not modified.
+    """
+    shift, target, minimum, _ = parameters
+    delta = numpy.asarray(contract_ss(vector)) - target * vector
+    if target >= minimum + 0.1:
+        delta = numpy.asarray(contract_ss(delta)) - target * delta
+    return shift * delta
+
+
+def _spin_penalty_diagonal(spin_diagonal, parameters):
+    """Return the shared Davidson/Newton penalty preconditioner diagonal.
+
+    For the quadratic penalty this is (diag(S^2)-target)^2, an inexpensive
+    approximation, not diag((S^2-target)^2). Operator actions remain exact.
+    """
+    shift, target, minimum, _ = parameters
+    delta = spin_diagonal - target
+    if target >= minimum + 0.1:
+        delta = delta * delta
+    return shift * delta
+
+
 class _GasSpinPlan:
     """Reusable block-sparse ``S^2`` contraction plan.
 
@@ -1173,11 +1199,8 @@ def option(name, default):
                     if full_spin_pspace:
                         s2 = spin_plan.matrix()
                         s2 = s2[numpy.ix_(addresses, addresses)]
-                        delta_s2 = s2 - spin_target * numpy.eye(gas.ndet)
-                        if linear_spin_penalty:
-                            h0 += spin_shift * delta_s2
-                        else:
-                            h0 += spin_shift * numpy.dot(delta_s2, delta_s2)
+                        h0 += _spin_penalty_action(
+                            s2.dot, numpy.eye(gas.ndet), spin_penalty)
                         self.spin_penalty_method = "exact-small-space"
                     eigenvalues, eigenvectors = numpy.linalg.eigh(h0)
                     e = eigenvalues[:nroots]
@@ -1263,33 +1286,14 @@ def hop(vec):
                             h2e, vec, norb, nelec, plan=plan).reshape(-1)
                         if spin_penalty is None:
                             return result
-                        ss_vector = numpy.asarray(
-                            spin_plan.contract(vec)).reshape(-1)
-                        if linear_spin_penalty:
-                            ss_vector -= spin_target * vec
-                            result += spin_shift * ss_vector
-                            return result
-                        tmp = ss_vector - spin_target * vec
-                        correction = numpy.asarray(
-                            spin_plan.contract(tmp)).reshape(-1)
-                        correction -= spin_target * tmp
-                        result += spin_shift * correction
+                        result += _spin_penalty_action(
+                            spin_plan.contract, vec, spin_penalty)
                         return result
 
                     preconditioner_diagonal = hdiag
                     if spin_penalty is not None:
-                        spin_diagonal = spin_plan.diagonal_vector()
-                        delta_diagonal = spin_diagonal - spin_target
-                        if linear_spin_penalty:
-                            penalty_diagonal = spin_shift * delta_diagonal
-                        else:
-                            # This is the inexpensive diagonal approximation to
-                            # (S^2-target)^2.  The projected trial vectors
-                            # improve access to the target sector; this term
-                            # only improves Davidson conditioning.
-                            penalty_diagonal = (
-                                spin_shift * delta_diagonal * delta_diagonal)
-                        preconditioner_diagonal = hdiag + penalty_diagonal
+                        preconditioner_diagonal = hdiag + _spin_penalty_diagonal(
+                            spin_plan.diagonal_vector(), spin_penalty)
 
                     def precond(dx, e, *args):
                         denom = preconditioner_diagonal - e

pyscf/mcscf/gasscf.py

@@ -39,6 +39,7 @@
 import hashlib
 import json
 import sys
+from types import FunctionType
 
 import numpy
 
@@ -154,12 +155,106 @@ def _gasscf_newton_kernel(casscf, *args, **kwargs):
     kwargs["verbose"] = verbose
     # Native Newton consumes objectives throughout its complete trajectory.
     # Convert only the returned tuple before mc1step assigns/logs public slots.
-    result = newton_casscf.kernel(casscf, *args, **kwargs)
+    native_kernel = newton_casscf.kernel
+    if hasattr(casscf.fcisolver, 'ss_penalty'):
+        # PySCF kernel -> update_orb_ci -> gen_g_hop uses module globals,
+        # not casscf.gen_g_hop. Reuse the original function code/defaults in
+        # a private namespace, changing only the response entry point. Never
+        # patch the imported module: other CASSCF calculations keep their
+        # native operators, including during callbacks and nested calls.
+        namespace = dict(vars(newton_casscf))
+        namespace['gen_g_hop'] = gen_g_hop
+        for name in ('update_orb_ci', 'kernel'):
+            original = getattr(newton_casscf, name)
+            rebound = FunctionType(original.__code__, namespace, original.__name__,
+                                   original.__defaults__, original.__closure__)
+            rebound.__kwdefaults__ = original.__kwdefaults__
+            namespace[name] = rebound
+        native_kernel = namespace['kernel']
+    result = native_kernel(casscf, *args, **kwargs)
     physical, gas_physical = gasci._publish_energy_results(
         casscf, result[1], result[2])
     return (result[0], physical, gas_physical) + result[3:]
 
 
+def gen_g_hop(mc, mo, ci0, eris, verbose=None):
+    """Add spin-penalty CI derivatives to the native Newton operators.
+
+    The spin penalty P has no orbital dependence. Keep contract_2e physical:
+    native Newton uses that hook for both H and orbital variations of H.
+    Only the CI gradient, CI-CI Hessian and CI preconditioner need corrections.
+    For a normalized root c, p = <c|P|c> and r = Pc - pc, these are
+    2 w r and 2 w [(P-p) v - r(c.v) - c(r.v)], respectively, in the native
+    Newton normalization convention. The keyframe gradient uses the same P.
+    """
+    gradient, update, hop, hdiag = newton_casscf.gen_g_hop(
+        mc, mo, ci0, eris, verbose)
+    solver = mc.fcisolver
+    parameters = fci_gas._spin_penalty_parameters(
+        solver, mc.ncas, mc._effective_nelecas())
+    if parameters is None or parameters[0] == 0:
+        return gradient, update, hop, hdiag
+
+    # The spin plan owns only Python/NumPy data and remains valid after the
+    # solver drops its cache. No borrowed C pointers escape in these closures.
+    if solver.cache_plans:
+        plan = solver._get_spin_plan(mc.ncas, mc.nelecas)
+    else:
+        plan = solver.make_spin_plan(mc.ncas, mc.nelecas)
+
+    def penalty(vector):
+        return fci_gas._spin_penalty_action(plan.contract, vector, parameters)
+
+    roots = [ci0] if solver.nroots == 1 else ci0
+    roots = [numpy.asarray(c).ravel() for c in roots]
+    weights = tuple(getattr(mc, 'weights', (1.,)))
+    ngorb = gradient.size - sum(c.size for c in roots)
+    diagonal = fci_gas._spin_penalty_diagonal(
+        plan.diagonal_vector(), parameters)
+    gradient = gradient.copy()
+    hdiag = hdiag.copy()
+    terms = []
+    start = ngorb
+    for c, weight in zip(roots, weights):
+        block = slice(start, start + c.size)
+        start += c.size
+        if weight == 0:
+            continue
+        pc = penalty(c)
+        energy = c.dot(pc)
+        residual = pc - energy * c
+        gradient[block] += 2 * weight * residual
+        hdiag[block] += 2 * weight * (diagonal - energy - 2 * residual * c)
+        terms.append((block, c, weight, energy, residual))
+
+    def penalized_hop(vector):
+        result = hop(vector)
+        for block, c, weight, energy, residual in terms:
+            v = vector[block]
+            result[block] += 2 * weight * (
+                penalty(v) - energy * v - residual * c.dot(v)
+                - c * residual.dot(v))
+        return result
+
+    def penalized_update(u, ci):
+        result = update(u, ci)
+        current = [ci] if solver.nroots == 1 else ci
+        start = ngorb
+        for c, weight in zip(current, weights):
+            c = numpy.asarray(c).ravel()
+            block = slice(start, start + c.size)
+            start += c.size
+            if weight == 0:
+                continue
+            # Native g_update normalizes each current CI vector as well.
+            c = c / numpy.linalg.norm(c)
+            pc = penalty(c)
+            result[block] += 2 * weight * (pc - c.dot(pc) * c)
+        return result
+
+    return gradient, penalized_update, penalized_hop, hdiag
+
+
 class _GASFCISolver(fci_gas.FCISolver):
     """GASCI solver shell with GASSCF-owned helper plans.
 
@@ -1511,7 +1606,7 @@ def state_specific_(self, *args, **kwargs):
 
     state_specific = state_specific_
 
-    gen_g_hop = newton_casscf.gen_g_hop
+    gen_g_hop = gen_g_hop
 
     def kernel(self, mo_coeff=None, ci0=None, callback=None):
         """Run full GASSCF orbital optimization with native CIAH.

pyscf/mcscf/test/test_gasscf.py

@@ -36,6 +36,7 @@
 from pyscf.tools import molden
 from pyscf.fci import addons as fci_addons
 from pyscf.fci import direct_spin1
+from pyscf.fci import spin_op
 from pyscf.mcscf import addons
 from pyscf.mcscf import addons_gas
 from pyscf.mcscf import df as mcdf
@@ -2920,6 +2921,180 @@ def test_fix_spin_rejects_spin_incomplete_gas_and_invalid_target(self):
         with self.assertRaisesRegex(ValueError, "target S"):
             complete.fix_spin_(shift=.2, ss=.5)
 
+    def _check_spin_penalty_newton_derivatives(self, use_df=False, weights=None):
+        mol = gto.M(atom='H 0 0 0; H 0 0 .8; H 0 0 1.8; H 0 0 2.6',
+                    basis='sto-3g', verbose=0)
+        mf = scf.RHF(mol).run()
+        mc = gasscf.GASSCF(
+            mf, 3, (1, 1), ncore=1, gas_orbs=(1, 2),
+            gas_restr=((0, 1), (2, 2)), gas_restr_type='cumulative-occ')
+        self.addCleanup(mc.close)
+        if use_df:
+            mc = mc.density_fit()
+            self.addCleanup(mc.close)
+        if weights is not None:
+            mc = mc.state_average(weights)
+            self.addCleanup(mc.close)
+        mc.validate_capabilities()
+        mo = mf.mo_coeff
+        eris = mc.ao2mo(mo)
+        h1, core = mc.get_h1gas(mo)
+        h2 = mc.get_h2gas(mo)
+        absorbed = direct_spin1.absorb_h1e(h1, h2, mc.ncas, mc.nelecas, .5)
+        # Independent full-FCI operators projected into a genuine restricted
+        # GAS space. Deliberately mix spin sectors: pure-spin eigenvectors can
+        # hide a missing penalty in the CI gradient and curvature.
+        with mc.fcisolver.make_space(mc.ncas, mc.nelecas) as space:
+            eye = numpy.eye(space.ndet)
+            def project_action(action, column):
+                full = fci_gas.gas2fci(column, space)
+                return fci_gas.fci2gas(action(full), space).ravel()
+            hmat = numpy.column_stack([
+                project_action(lambda c: direct_spin1.contract_2e(
+                    absorbed, c, mc.ncas, mc.nelecas), c) for c in eye])
+            s2 = numpy.column_stack([
+                project_action(lambda c: spin_op.contract_ss(
+                    c, mc.ncas, mc.nelecas), c) for c in eye])
+        nroots = mc.fcisolver.nroots
+        rng = numpy.random.default_rng(211)
+        columns = numpy.linalg.qr(rng.normal(size=(eye.shape[0], nroots)))[0]
+        roots = [columns[:, i].copy() for i in range(nroots)]
+        directions = []
+        for c in roots:
+            v = rng.normal(size=c.size)
+            v -= c.dot(v) * c
+            directions.append(v / numpy.linalg.norm(v))
+        ci = roots[0] if weights is None else roots
+        w = (1.,) if weights is None else weights
+        native = newton_casscf.gen_g_hop(mc, mo, ci, eris)
+        ngorb = native[0].size - sum(c.size for c in roots)
+        ci_direction = numpy.r_[numpy.zeros(ngorb), numpy.concatenate(directions)]
+        orbital_direction = numpy.r_[rng.normal(size=ngorb),
+                                    numpy.zeros(native[0].size-ngorb)]
+        u = scipy.linalg.expm(.02 * mc.unpack_uniq_var(orbital_direction[:ngorb]))
+
+        for target in (0., 2.):  # Linear minimum-spin and quadratic penalties.
+            mc.fix_spin_(shift=.2, ss=target)
+            delta = s2 - target * eye
+            penalty = .2 * (delta if target == 0 else delta @ delta)
+            operator = hmat + penalty
+            def energy(angle):
+                vectors = [numpy.cos(angle)*c + numpy.sin(angle)*v
+                           for c, v in zip(roots, directions)]
+                return core + sum(a * c.dot(operator @ c)
+                                  for a, c in zip(w, vectors))
+            previous = None
+            for cache in (True, False):
+                mc.cache_plans = cache
+                with self.subTest(df=use_df, weights=weights,
+                                  target=target, cache=cache):
+                    g, update, hop, diagonal = mc.gen_g_hop(mo, ci, eris)
+                    # Public contraction remains physical, even with fix_spin.
+                    numpy.testing.assert_allclose(
+                        mc.fcisolver.contract_2e(absorbed, roots[0],
+                                                mc.ncas, mc.nelecas),
+                        hmat @ roots[0], atol=1e-12, rtol=0)
+                    first = numpy.dot(g, ci_direction)
+                    second = numpy.dot(ci_direction, hop(ci_direction))
+                    self.assertGreater(abs(first-native[0].dot(ci_direction)), 1e-3)
+                    for step in (2e-4, 1e-4):
+                        plus, minus, center = energy(step), energy(-step), energy(0)
+                        self.assertAlmostEqual(first, (plus-minus)/(2*step), delta=2e-7)
+                        self.assertAlmostEqual(second, (plus-2*center+minus)/step**2,
+                                               delta=2e-6)
+                    # Compare the full tangent CI response with an independent
+                    # projected full-FCI matrix, not just one quadratic form.
+                    response = hop(ci_direction)[ngorb:]
+                    start = 0
+                    for a, c, v in zip(w, roots, directions):
+                        block = response[start:start+c.size]
+                        expected = 2*a*(operator @ v - c.dot(operator @ c)*v)
+                        numpy.testing.assert_allclose(
+                            block-c*c.dot(block), expected-c*c.dot(expected),
+                            atol=1e-11, rtol=0)
+                        start += c.size
+                    # P is orbital independent: no extra orbital or mixed block.
+                    numpy.testing.assert_allclose(g[:ngorb], native[0][:ngorb], atol=0, rtol=0)
+                    numpy.testing.assert_allclose(hop(orbital_direction),
+                                                  native[2](orbital_direction), atol=1e-12)
+                    numpy.testing.assert_allclose(hop(ci_direction)[:ngorb],
+                                                  native[2](ci_direction)[:ngorb], atol=1e-12)
+                    # Check the keyframe update at changed, unnormalized CI and
+                    # rotated orbitals. Only the normalized CI penalty changes.
+                    current = [(1.2+i)*(c+.13*v)
+                               for i, (c, v) in enumerate(zip(roots, directions))]
+                    current_arg = current[0] if weights is None else current
+                    actual = update(u, current_arg)
+                    expected = native[1](u, current_arg)
+                    start = ngorb
+                    for a, c in zip(w, current):
+                        c = c / numpy.linalg.norm(c)
+                        expected[start:start+c.size] += 2*a*(penalty @ c-c.dot(penalty @ c)*c)
+                        start += c.size
+                    numpy.testing.assert_allclose(actual, expected, atol=1e-11)
+                    # Quadratic diagonal is a documented preconditioner
+                    # approximation, shared with Davidson; the HVP is exact.
+                    pd = .2 * (s2.diagonal()-target)
+                    if target != 0:
+                        pd = .2 * (s2.diagonal()-target)**2
+                    expected = native[3].copy()
+                    start = ngorb
+                    for a, c in zip(w, roots):
+                        ep = c.dot(penalty @ c)
+                        residual = penalty @ c-ep*c
+                        expected[start:start+c.size] += 2*a*(pd-ep-2*residual*c)
+                        start += c.size
+                    numpy.testing.assert_allclose(diagonal, expected, atol=1e-12)
+                    if weights is not None and weights[-1] == 0:
+                        numpy.testing.assert_allclose(g[-roots[-1].size:], 0, atol=0)
+                        numpy.testing.assert_allclose(hop(ci_direction)[-roots[-1].size:], 0, atol=0)
+                    values = (g, hop(ci_direction+orbital_direction), diagonal)
+                    if previous is not None:
+                        for a, b in zip(values, previous):
+                            numpy.testing.assert_allclose(a, b, atol=1e-12)
+                    previous = values
+            mc.undo_fix_spin_()
+
+    def test_spin_penalty_newton_derivatives(self):
+        self._check_spin_penalty_newton_derivatives()
+
+    def test_spin_penalty_sa_newton_derivatives(self):
+        for weights in ((.3, .7), (1., 0.)):
+            self._check_spin_penalty_newton_derivatives(weights=weights)
+
+    def test_spin_penalty_df_newton_derivatives(self):
+        for weights in (None, (.3, .7), (1., 0.)):
+            self._check_spin_penalty_newton_derivatives(use_df=True, weights=weights)
+
+    def test_spin_penalty_kernel_uses_local_newton_response(self):
+        mol = gto.M(atom='H 0 0 0; H 0 0 .8; H 0 0 1.8; H 0 0 2.6',
+                    basis='sto-3g', verbose=0)
+        mf = scf.RHF(mol).run()
+        original = tuple(getattr(newton_casscf, name) for name in
+                         ('kernel', 'update_orb_ci', 'gen_g_hop'))
+        for weights in (None, (.3, .7), (1., 0.)):
+            mc = gasscf.GASSCF(mf, 3, (1, 1), ncore=1, gas_orbs=(1, 2),
+                              gas_restr=((0, 1), (2, 2)),
+                              gas_restr_type='cumulative-occ')
+            if weights is not None:
+                mc = mc.state_average(weights)
+            self.addCleanup(mc.close)
+            mc.fix_spin_(shift=.2, ss=0.)
+            mc.max_cycle_macro = mc.max_cycle_micro = 1
+            seen = []
+            def callback(env):
+                # Native module functions stay unchanged even while the
+                # calculation is running, not merely restored on return.
+                for name, function in zip(('kernel', 'update_orb_ci', 'gen_g_hop'), original):
+                    self.assertIs(getattr(newton_casscf, name), function)
+                self.assertAlmostEqual(env['e_tot'], mc.spin_energy_report()['objective'], 10)
+                seen.append(True)
+            with mock.patch.object(gasscf, 'gen_g_hop', wraps=gasscf.gen_g_hop) as response:
+                result = mc.kernel(callback=callback)
+                self.assertGreater(response.call_count, 0)
+            self.assertTrue(seen)
+            self.assertAlmostEqual(result[0], mc.spin_energy_report()['physical'], 10)
+
     def _check_physical_energy_result(self, mc, result):
         report = mc.spin_energy_report()
         roots = mc.ci if isinstance(mc.ci, (list, tuple)) else [mc.ci]