Skip to content

Update tests #968

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sunqm merged 3 commits into from
Oct 16, 2026
+85 −56
Merged

Update tests #968

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
eadc33a
Adjust ucdft tests
sunqm Oct 16, 2026
9407777
typo when printing UHF mo_energy
sunqm Oct 16, 2026
2af407c
Fix
sunqm Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 81 additions & 52 deletions gpu4pyscf/dft/tests/test_ucdft.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,21 +12,21 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import cupy as cp
import unittest
import pyscf
from pyscf import lib
from gpu4pyscf.dft import uks, ucdft
from gpu4pyscf.dft.cdft_soscf import newton_cdft as newton_penalty
from gpu4pyscf.dft.cdft_soscf_full import newton_cdft
import pytest

atom = '''
O 0 0 0
H 0. -0.757 0.587
H 0. 0.757 0.587
'''
bas='def2-tzvp'
auxbas='def2-tzvp-jkfit'
charge_constraints_2 = [ [0, 1], [8.1, 0.95] ]
charge_constraints_3 = [ [0, 1, 2], [8.1, 0.95, 0.95] ]
charge_constraints_2_ao = [
Expand All @@ -43,6 +43,7 @@ def run_dft(mol, xc, method, charge_constraints,
penalty_weight=penalty,
projection_method=projection_method
)
mf = mf.density_fit(auxbasis=auxbas)
mf.xc = xc
mf.grids.atom_grid = (99, 590)
mf.conv_tol = 1.0E-8
Expand Down Expand Up @@ -99,24 +100,22 @@ def setUpClass(cls):
charge_constraints_2, projection_method='becke')
cls.output_lagrange_nested_cons2_minao = run_dft(mol, 'b3lyp', 'lagrange',
charge_constraints_2, projection_method='minao')
cls.output_lagrange_nested_cons2_minao_ao = run_dft(mol, 'b3lyp', 'lagrange',
charge_constraints_2_ao, projection_method='minao')

@classmethod
def tearDownClass(cls):
cls.mol.stdout.close()

def test_energy(self):
ref = -76.37687528310500 # ref from becke partition, 0.1e-10 penalty weight
self.assertAlmostEqual(self.output_penalty_cons2['e_tot'], ref, 1)
ref = -76.376907212893 # DF Becke Lagrange reference
self.assertAlmostEqual(self.output_penalty_cons2['e_tot'], ref, delta=0.05)
self.assertAlmostEqual(self.output_lagrange_soscf_cons2['e_tot'], ref, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2['e_tot'], ref, 6)

def test_homo_lumo(self):
homo_ref = -0.59463327378979
lumo_ref = -0.08711895811466
self.assertAlmostEqual(self.output_penalty_cons2['homo'], homo_ref, 1)
self.assertAlmostEqual(self.output_penalty_cons2['lumo'], lumo_ref, 1)
homo_ref = -0.594634337972
lumo_ref = -0.087117515816
self.assertAlmostEqual(self.output_penalty_cons2['homo'], homo_ref, delta=0.05)
self.assertAlmostEqual(self.output_penalty_cons2['lumo'], lumo_ref, delta=0.05)
self.assertAlmostEqual(self.output_lagrange_soscf_cons2['homo'], homo_ref, 6)
self.assertAlmostEqual(self.output_lagrange_soscf_cons2['lumo'], lumo_ref, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2['homo'], homo_ref, 6)
Expand All @@ -126,26 +125,30 @@ def test_charge(self):
charge_ref_O = 8.1
charge_ref_H1 = 0.95

self.assertAlmostEqual(self.output_penalty_cons2['O_charge'], charge_ref_O, 1)
# A finite quadratic penalty leaves a residual: near the constrained
# solution, delta_N ~ V / (2 * weight). For O, V ~ -0.45 and weight=10
# predict about -0.0225 electrons; this is not a rounding tolerance.
self.assertAlmostEqual(self.output_penalty_cons2['O_charge'], charge_ref_O, delta=0.03)
self.assertAlmostEqual(self.output_lagrange_soscf_cons2['O_charge'], charge_ref_O, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2['O_charge'], charge_ref_O, 6)
self.assertAlmostEqual(self.output_penalty_cons2['H1_charge'], charge_ref_H1, 2)
self.assertAlmostEqual(self.output_penalty_cons2['H1_charge'], charge_ref_H1, delta=0.005)
self.assertAlmostEqual(self.output_lagrange_soscf_cons2['H1_charge'], charge_ref_H1, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2['H1_charge'], charge_ref_H1, 6)

def test_multiplier(self):
# Multipliers are nonvariational and more sensitive than total energies.
ref_O = -4.50487015E-01
ref_H = -2.37985800E-08

self.assertAlmostEqual(self.output_lagrange_soscf_cons2['v_lagrange'][0], ref_O, 6)
self.assertAlmostEqual(self.output_lagrange_soscf_cons2['v_lagrange'][1], ref_H, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2['v_lagrange'][0], ref_O, delta=1e-6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2['v_lagrange'][1], ref_H, delta=1e-6)
self.assertAlmostEqual(self.output_lagrange_soscf_cons2['v_lagrange'][0], ref_O, delta=1e-4)
self.assertAlmostEqual(self.output_lagrange_soscf_cons2['v_lagrange'][1], ref_H, delta=1e-4)
self.assertAlmostEqual(self.output_lagrange_nested_cons2['v_lagrange'][0], ref_O, delta=1e-4)
self.assertAlmostEqual(self.output_lagrange_nested_cons2['v_lagrange'][1], ref_H, delta=1e-4)

def test_minao_projection(self):
ref_energy = -76.3639776639758
ref_homo = -0.123418866057557
ref_lumo = 0.058258581369145
ref_energy = -76.364007406852
ref_homo = -0.123407786265
ref_lumo = 0.058274316958
ref_O_charge = 8.1
ref_H1_charge = 0.95
ref_O_multiplier = 0.31278893
Expand All @@ -155,23 +158,16 @@ def test_minao_projection(self):
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao['lumo'], ref_lumo, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao['O_charge'], ref_O_charge, 4)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao['H1_charge'], ref_H1_charge, 4)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao['v_lagrange'][0], ref_O_multiplier, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao['v_lagrange'][1], ref_H_multiplier, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao['v_lagrange'][0], ref_O_multiplier, delta=1e-4)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao['v_lagrange'][1], ref_H_multiplier, delta=1e-4)

self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao_ao['e_tot'], ref_energy, 7)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao_ao['homo'], ref_homo, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao_ao['lumo'], ref_lumo, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao_ao['O_charge'], ref_O_charge, 4)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao_ao['H1_charge'], ref_H1_charge, 4)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao_ao['v_lagrange'][0], ref_O_multiplier, 6)
self.assertAlmostEqual(self.output_lagrange_nested_cons2_minao_ao['v_lagrange'][1], ref_H_multiplier, 6)

def test_minao_projection_1ao(self):
mf = ucdft.CDFT_UKS(self.mol,
charge_constraints=[ ['0 O 2pz'], [1.2] ],
method='lagrange',
projection_method='minao'
)
mf = mf.density_fit(auxbasis=auxbas)
mf.xc = 'b3lyp'
mf.grids.atom_grid = (99, 590)
mf.conv_tol = 1.0E-8
Expand All @@ -187,23 +183,22 @@ def test_minao_projection_1ao(self):
O_charge = cp.trace(dm[0] @ projs[0]) + cp.trace(dm[1] @ projs[0])

self.assertAlmostEqual(float(O_charge), 1.2, 6)
self.assertAlmostEqual(float(e_tot), -76.3840250035048, 7)
self.assertAlmostEqual(float(homo), -0.236743012466074, 6)
self.assertAlmostEqual(float(lumo), 0.0494326184909096, 6)
self.assertAlmostEqual(float(v_lagrange[0]), 0.288396, 6)
self.assertAlmostEqual(float(e_tot), -76.384052803986, 7)
self.assertAlmostEqual(float(homo), -0.236737660768, 6)
self.assertAlmostEqual(float(lumo), 0.049445292477, 6)
self.assertAlmostEqual(float(v_lagrange[0]), 0.288396, delta=1e-4)

@pytest.mark.slow
def test_canonical_mo_energy(self):
mf = self.output_lagrange_nested_cons2_minao_ao['mf']
mf.xc = 'b3lyp'
mf.grids.atom_grid = (99, 590)
mf.conv_tol = 1.0E-8
mf.max_cycle = 100
mf = newton_cdft(mf)
mf.kernel()
mo_energy = mf.get_canonical_mo()[0]
mf = self.output_lagrange_nested_cons2_minao['mf']
mo_energy, mo_coeff = mf.get_canonical_mo()
dm = mf.make_rdm1()
fock = (mf.get_hcore() + mf.get_veff(self.mol, dm)).get()
overlap = mf.get_ovlp().get()
np.testing.assert_allclose(
fock[0] @ mo_coeff[0], (overlap @ mo_coeff[0]) * mo_energy[0], atol=1e-9, rtol=0)
np.testing.assert_allclose(
fock[1] @ mo_coeff[1], (overlap @ mo_coeff[1]) * mo_energy[1], atol=1e-9, rtol=0)
mf.analyze()
assert abs(lib.fp(mo_energy) - (-92.71420354230172)) < 1e-6

def test_spin_constraint(self):
mol = pyscf.M(
Expand All @@ -227,24 +222,27 @@ def test_spin_constraint(self):
charge = 0,
output = '/dev/null'
)
mf_atom = uks.UKS(mol_atom)
mf_atom = uks.UKS(mol_atom).density_fit(auxbasis=auxbas)
mf_atom.xc = 'b3lyp'
mf_atom.grids.atom_grid = (99, 590)
mf_atom.kernel()

mf = ucdft.CDFT_UKS(mol,
spin_constraints= [[0], [2.001657]],
)
mf = mf.density_fit(auxbasis=auxbas)
mf.xc = 'b3lyp'
mf.grids.atom_grid = (99, 590)
mf.conv_tol = 1e-12
mf.conv_tol = 1e-9
mf.kernel()
mo_energy = mf.get_canonical_mo()[0]
assert abs(float(mf_atom.e_tot)*2 - float(mf.e_tot)) < 2e-5
assert abs(float(mo_energy[0][7]) - float(mf_atom.mo_energy[1][2])) < 1e-3
assert abs(float(mo_energy[0][8]) - float(mf_atom.mo_energy[1][3])) < 1e-3
assert abs(float(mo_energy[1][7]) - float(mf_atom.mo_energy[1][2])) < 1e-3
assert abs(float(mo_energy[1][8]) - float(mf_atom.mo_energy[1][3])) < 1e-3
# Dissociation and near-degenerate orbital splittings are sensitive
# to the finite grid; use separate energy and orbital tolerances.
assert abs(float(mf_atom.e_tot)*2 - float(mf.e_tot)) < 1e-4
assert abs(float(mo_energy[0][7]) - float(mf_atom.mo_energy[1][2])) < 5e-3
assert abs(float(mo_energy[0][8]) - float(mf_atom.mo_energy[1][3])) < 5e-3
assert abs(float(mo_energy[1][7]) - float(mf_atom.mo_energy[1][2])) < 5e-3
assert abs(float(mo_energy[1][8]) - float(mf_atom.mo_energy[1][3])) < 5e-3

def test_mixed_charge_and_spin_constraints(self):
charge_cons = [[0], [8.1]]
Expand All @@ -255,6 +253,7 @@ def test_mixed_charge_and_spin_constraints(self):
spin_constraints=spin_cons,
method='lagrange',
projection_method='minao')
mf = mf.density_fit(auxbasis=auxbas)
mf.xc = 'b3lyp'
mf.grids.atom_grid = (99, 590)
mf.kernel()
Expand All @@ -268,7 +267,7 @@ def test_mixed_charge_and_spin_constraints(self):
self.assertAlmostEqual(float(n_O), 8.1, 4)
self.assertAlmostEqual(float(m_O), 0.2, 4)
self.assertEqual(len(mf.v_lagrange), 2)
self.assertAlmostEqual(float(mf.e_tot), -76.36119086716313, 6)
self.assertAlmostEqual(float(mf.e_tot), -76.361220462326, 6)

def test_group_constraints(self):
group_cons = [[[1, 2]], [1.5]]
Expand All @@ -277,6 +276,7 @@ def test_group_constraints(self):
charge_constraints=group_cons,
method='lagrange',
projection_method='minao')
mf = mf.density_fit(auxbasis=auxbas)
mf.xc = 'b3lyp'
mf.grids.atom_grid = (75, 302)
mf.kernel()
Expand All @@ -285,7 +285,36 @@ def test_group_constraints(self):
projs = mf.build_projectors()
n_H_sum = cp.trace(dm[0] @ projs[0]) + cp.trace(dm[1] @ projs[0])
self.assertAlmostEqual(float(n_H_sum), 1.5, 5)
self.assertAlmostEqual(float(mf.e_tot), -76.44530418143451, 6)
self.assertAlmostEqual(float(mf.e_tot), -76.445332942880, 6)

def test_charge_potential_signs(self):
mf = ucdft.CDFT_UKS(self.mol, charge_constraints=[[0], [8.1]])
w_charge = mf.build_projectors()[0]
potential = mf.get_constraint_potential([0.3])
expected = (0.3 * w_charge).get()
# Charge constraints shift both spin channels in the same direction.
np.testing.assert_allclose(potential[0].get(), expected, atol=1e-14, rtol=0)
np.testing.assert_allclose(potential[1].get(), expected, atol=1e-14, rtol=0)

def test_spin_potential_signs(self):
mf = ucdft.CDFT_UKS(self.mol, spin_constraints=[[1], [0.2]])
w_spin = mf.build_projectors()[0]
potential = mf.get_constraint_potential([-0.2])
expected_alpha = (-0.2 * w_spin).get()
# Spin constraints shift alpha and beta in opposite directions.
np.testing.assert_allclose(potential[0].get(), expected_alpha, atol=1e-14, rtol=0)
np.testing.assert_allclose(potential[1].get(), -expected_alpha, atol=1e-14, rtol=0)

def test_mixed_charge_and_spin_potential_signs(self):
mf = ucdft.CDFT_UKS(self.mol,
charge_constraints=[[0], [8.1]],
spin_constraints=[[1], [0.2]])
w_charge, w_spin = mf.build_projectors()
potential = mf.get_constraint_potential([0.3, -0.2])
expected_alpha = (0.3 * w_charge - 0.2 * w_spin).get()
expected_beta = (0.3 * w_charge + 0.2 * w_spin).get()
np.testing.assert_allclose(potential[0].get(), expected_alpha, atol=1e-14, rtol=0)
np.testing.assert_allclose(potential[1].get(), expected_beta, atol=1e-14, rtol=0)

def test_error_handling_mismatched_inputs(self):
groups = [[0], [1]]
Expand Down
6 changes: 3 additions & 3 deletions gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -578,14 +578,14 @@ def test_unique_image_pairs(self):
ret = multigrid._unique_image_pair(Ls)
assert ret[0].shape == (2431, 3)
assert abs(lib.fp(ret[0].get()) - 1.5047201402319172) < 1e-10
assert abs(lib.fp(ret[1].get()) - -483.0210637951298) < 1e-10
assert abs(lib.fp(ret[1].get()) - -483.0210637951298) < 1e-9

np.random.seed(2)
Ls = Ls[np.random.rand(len(Ls)) > .5]
ret = multigrid._unique_image_pair(Ls)
assert ret[0].shape == (2377, 3)
assert abs(lib.fp(ret[0].get()) - -34.99306750756055) < 1e-10
assert abs(lib.fp(ret[1].get()) - 1347.464804553046) < 1e-10
assert abs(lib.fp(ret[1].get()) - 1347.464804553046) < 1e-8

Lx = np.append(np.arange(0, 3), np.arange(-3, 0))
Ly = np.append(np.arange(0, 3), np.arange(-3, 0))
Expand All @@ -595,7 +595,7 @@ def test_unique_image_pairs(self):
ret = multigrid._unique_image_pair(Ls)
assert ret[0].shape == (1331, 3)
assert abs(lib.fp(ret[0].get()) - -1.355090495130784) < 1e-10
assert abs(lib.fp(ret[1].get()) - 2145.771285837819) < 1e-10
assert abs(lib.fp(ret[1].get()) - 2145.771285837819) < 1e-8

def test_image_pair_to_difference(self):
cell = gto.M(a=np.eye(3)*3, atom='He 0. 0. 0.', basis=[[0, [1, 1]]])
Expand Down
2 changes: 1 addition & 1 deletion gpu4pyscf/scf/uhf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -278,7 +278,7 @@ def get_occ(self, mo_energy=None, mo_coeff=None):
lumo = lumo_b = mo_energy[1, e_idx_b[n_b]].item()
lumo_a = None
if n_a < nmo:
lumo_a = mo_energy[1, e_idx_a[n_a]].item()
lumo_a = mo_energy[0, e_idx_a[n_a]].item()
lumo = min(lumo, lumo_a)
gap = (lumo - homo) * HARTREE2EV
self.scf_summary['gap'] = gap
Expand Down
Loading