Skip to content

Commit a193de5

Browse files
authored
PV contributions to energy (pyscf#201)
* PV-energy from density matrix * Fix issues in pv * ruff check on test_energy-pv file * differentiate according to the existence of dm
1 parent f0ccb26 commit a193de5

2 files changed

Lines changed: 216 additions & 31 deletions

File tree

‎pyscf/pv/energy.py‎

100755100644
Lines changed: 202 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,9 @@
1919
'''
2020

2121
import numpy
22+
23+
from pyscf import gto, lib
2224
from pyscf.scf import dhf
23-
from pyscf import lib, gto
2425

2526

2627
def fc_integrals(mol, mf, atom, **kwargs):
@@ -72,6 +73,50 @@ def fc_integrals(mol, mf, atom, **kwargs):
7273

7374
return fc_integrals
7475

76+
77+
78+
def gamma5_fc_integrals(mol, mf, atom, **kwargs):
79+
"""
80+
Compute the integrals for PV contributions to energy for a given atom.
81+
The integrals are calculated in the four-component spinor basis.
82+
Point nuclear charge distribution is assumed.
83+
84+
:param mol: Molecule object containing information about the molecular system.
85+
:type mol: pyscf.gto.Mole
86+
:param mf: Mean-field object, obtained from a Dirac-Hartree-Fock (DHF) calculation.
87+
:type mf: pyscf.scf.hf.SCF
88+
:param atom: Index of the atom for which the integrals are computed.
89+
:type atom: int
90+
:return: Matrix of integrals (4-component spinor form) for the specified atom.
91+
:rtype: numpy.ndarray (complex128), shape (n4c, n4c)
92+
93+
Example:
94+
mol = gto.Mole()
95+
# Setup mol with geometry and basis set...
96+
mf = scf.DHF(mol)
97+
mf.kernel()
98+
atom_index = 0
99+
fc_matrix = PV_integrals(mol, mf, atom_index)
100+
"""
101+
n4c = mf.mo_coeff.shape[0]
102+
n2c = n4c // 2
103+
104+
105+
# Get the Fermi Contact integrals for the selected atom
106+
fc_ao = fc_integrals(mol, mf, atom)
107+
108+
# Prepare the empty matrix for the FC*Gamma5 integrals (LL, LS, SL, SS blocks)
109+
gamma5_fc_integrals = numpy.zeros((n4c, n4c), dtype=numpy.complex128)
110+
111+
# LS and SL blocks
112+
gamma5_fc_integrals[:n2c, n2c:] = fc_ao[:n2c, n2c:].copy()
113+
gamma5_fc_integrals[n2c:, :n2c] = fc_ao[n2c:, :n2c:].copy()
114+
115+
116+
117+
118+
return gamma5_fc_integrals
119+
75120
def fc_expval(mol, mf, atom):
76121
"""
77122
Calculate the Fermi Contact (FC) expectation values for each occupied orbital in a molecule,
@@ -117,18 +162,32 @@ def fc_expval(mol, mf, atom):
117162
fc_ao_SS = fc_ao[n2c:, n2c:]
118163

119164
for k in range(nocc):
120-
expval_LL_k = numpy.einsum('pi,ij,qj->pq', Lo[:, k].conjugate(), fc_ao_LL, Lo[:, k])
121-
expval_SS_k = numpy.einsum('pi,ij,qj->pq', So[:, k].conjugate(), fc_ao_SS, Lo[:, k])
165+
expval_LL_k = numpy.einsum('i,ij,j->', Lo[:,k].conj(), fc_ao_LL, Lo[:,k])
166+
expval_SS_k = numpy.einsum('i,ij,j->', So[:,k].conj(), fc_ao_SS, So[:,k])
167+
122168
expval_perorb[k] = expval_LL_k + expval_SS_k
123169

170+
124171
return fac * expval_perorb.real
125172

126173

127174

128-
def Epv_atom(mol, mf, atom_index):
175+
def Epv_atom(mol, mf, atom_index, dm=None):
129176
"""
130177
Calculate the parity-violating (PV) contribution to energy for a given atom in a molecule.
131178
179+
The function supports:
180+
- orbital expectation values from DHF orbitals
181+
- density matrix contraction in AO or MO basis
182+
183+
For a one-electron operator:
184+
E_PV = Tr(D * h_PV)
185+
186+
If dm is provided:
187+
- AO density matrices are transformed to MO basis
188+
- MO density matrices (e.g. MP2/CCSD PySCF densities) are used directly
189+
190+
132191
It then returns the contributions from each occupied orbital.
133192
(First term of Eq. 4 - https://doi.org/10.1002/wcms.1396)
134193
@@ -151,13 +210,16 @@ def Epv_atom(mol, mf, atom_index):
151210
"""
152211
n4c, nmo = mf.mo_coeff.shape
153212
n2c = n4c // 2
154-
nNeg = nmo // 2 # Molecular orbitals of negative energy
155-
nocc = mf.mol.nelectron
156213

214+
# Molecular orbitals of negative energy
215+
nNeg = nmo // 2
157216

158-
# Get the positive components of the MO spinors
159-
Lo = mf.mo_coeff[:n2c, nNeg:nNeg + nocc]
160-
So = mf.mo_coeff[n2c:, nNeg:nNeg + nocc]
217+
# Number of occupied electrons
218+
nocc = mf.mol.nelectron
219+
220+
# ------------------------------------------------------------
221+
# Nuclear weak charge prefactor
222+
# ------------------------------------------------------------
161223

162224
# Atom masses and atomic numbers
163225
masses = mf.mol.atom_mass_list(isotope_avg=False)
@@ -166,36 +228,126 @@ def Epv_atom(mol, mf, atom_index):
166228
# Neutrons per atom
167229
neutrons = masses - atomic_numbers
168230

169-
S2THETAW = 0.23122 # AS DIRAC24 (CODATA 2018)
231+
S2THETAW = 0.23122 # (CODATA 2018)
170232

171233
# Weak charge of the nucleus of the selected atom
172234
QW = (1 - 4 * S2THETAW) * atomic_numbers[atom_index] - neutrons[atom_index]
173235

174-
# Get the Fermi Contact integrals for the selected atom
175-
fc_ao = fc_integrals(mol, mf, atom_index)
176-
177-
# Expectation values
178-
expval_LS = numpy.einsum('ij,ji->i', Lo.conjugate().T @ fc_ao[:n2c, n2c:], So)
179-
expval_SL = numpy.einsum('ij,ji->i', So.conjugate().T @ fc_ao[n2c:, :n2c], Lo)
180-
181-
# Sum LS and SL terms
182-
expval_perorb = expval_LS + expval_SL
183-
184236
# Prefactor
185237
fac = (2.2225 * 10 ** (-14) * QW) / (2 * numpy.sqrt(2))
186238

239+
if dm is None:
240+
# ------------------------------------------------------------
241+
# PV operator in AO basis
242+
# ------------------------------------------------------------
243+
244+
# Get the Fermi Contact integrals for the selected atom
245+
fc_ao = fc_integrals(mol, mf, atom_index)
246+
247+
# ------------------------------------------------------------
248+
# Orbital expectation values (4c-DHF or 4c-DFT reference)
249+
# ------------------------------------------------------------
250+
251+
# Get the positive components of the MO spinors
252+
Lo = mf.mo_coeff[:n2c, nNeg:nNeg + nocc]
253+
So = mf.mo_coeff[n2c:, nNeg:nNeg + nocc]
254+
255+
# Expectation values
256+
expval_LS = numpy.einsum('ij,ji->i', Lo.conjugate().T @ fc_ao[:n2c, n2c:], So)
257+
expval_SL = numpy.einsum('ij,ji->i', So.conjugate().T @ fc_ao[n2c:, :n2c], Lo)
258+
expval_perorb = expval_LS + expval_SL
259+
260+
# ------------------------------------------------------------
261+
# Density matrix contraction
262+
# ------------------------------------------------------------
263+
264+
if dm is not None:
265+
266+
# Positive-energy MO coefficient matrix
267+
C = mf.mo_coeff[:, nNeg:]
268+
269+
nmo_pos = C.shape[1]
270+
271+
# Get the Gamma5 Fermi-Contact integrals for the selected atom (in AO basis)
272+
gamma5_fc_ao = gamma5_fc_integrals(mol, mf, atom_index)
273+
274+
# Transform operator AO -> MO
275+
# Operator transformation:
276+
#
277+
# hᴹᴼ = C† hᴬᴼ C
278+
#
279+
# where:
280+
# C : MO coefficient matrix
281+
# hᴬᴼ : operator matrix in AO basis
282+
# hᴹᴼ : operator matrix in MO basis
283+
gamma5_fc_mo = C.conj().T @ gamma5_fc_ao @ C
284+
285+
# --------------------------------------------------------
286+
# Determine density matrix representation
287+
# --------------------------------------------------------
288+
289+
if dm.shape[0] == n4c:
290+
# Transform AO density matrix to MO
291+
# 𝑫ᴹᴼ = C† S 𝑫ᴬᴼ S C
292+
#
293+
# where:
294+
# C : MO coefficient matrix
295+
# S : AO overlap matrix
296+
# 𝑫 : one-particle density matrix
297+
S = mf.get_ovlp()
298+
dm_mo = C.conj().T @ S @ dm @ S @ C
299+
300+
elif dm.shape[0] == nmo_pos:
301+
# Density already in MO basis
302+
# (MP2/CCSD densities)
303+
dm_mo = dm
304+
else:
305+
raise ValueError(
306+
"Density matrix dimension does not match AO or positive-energy MO space"
307+
)
308+
309+
# Total expectation value
310+
# E = Tr(𝑫 h)
311+
#expval_dm = numpy.trace(dm_mo @ gamma5_fc_mo)
312+
313+
# Orbital decomposition in MO basis
314+
# For correlated densities (MP2/CCSD), the orbital decomposition
315+
# is representation-dependent. Only the total trace is invariant.
316+
expval_perorb = numpy.einsum(
317+
'ij,ji->i',
318+
gamma5_fc_mo,
319+
dm_mo
320+
)
321+
322+
323+
# ------------------------------------------------------------
324+
# Return
325+
# ------------------------------------------------------------
326+
187327
return fac * expval_perorb.real
188328

189329

190330

191-
def Epv_molecule(mol, mf):
331+
332+
333+
334+
335+
def Epv_molecule(mol, mf, dm=None):
192336
"""
193337
Calculate the weak charge parity-violating (PV) contributions for all atoms in the molecule
194338
within a punctual nuclear charge distribution model.
195339
196340
This function iterates over all atoms in the molecule and computes the PV contributions
197341
for each atom.
198342
343+
If dm is None:
344+
returns orbital contributions from the reference DHF/DFT orbitals.
345+
346+
If dm is provided:
347+
returns contributions obtained by contracting the PV operator with
348+
the supplied density matrix (AO or MO representation).
349+
350+
199351
:param mol: Molecule object containing information about the molecular system.
200352
:type mol: pyscf.gto.Mole
201353
:param mf: Mean-field object, typically obtained from a Dirac-Hartree-Fock (DHF) calculation.
@@ -205,9 +357,34 @@ def Epv_molecule(mol, mf):
205357
:rtype: numpy.ndarray (real), shape (n_atoms, n_occ)
206358
"""
207359

208-
nocc = mf.mol.nelectron
209-
result = numpy.zeros((mf.mol.natm, nocc))
210-
for i in range(mf.mol.natm):
211-
result[i, :] = Epv_atom(mol, mf, i)
360+
natm = mf.mol.natm
361+
362+
# Determine output size from the calculation type
363+
if dm is None:
364+
n4c, nmo = mf.mo_coeff.shape
365+
nocc = mf.mol.nelectron
366+
ncol = nocc
367+
368+
else:
369+
n4c, nmo = mf.mo_coeff.shape
370+
nNeg = nmo // 2
371+
nmo_pos = nmo - nNeg
372+
373+
# AO or MO density
374+
if dm.shape[0] == n4c or dm.shape[0] == nmo_pos:
375+
ncol = nmo_pos
376+
377+
else:
378+
raise ValueError(
379+
"Density matrix dimension does not match AO "
380+
"or positive-energy MO space"
381+
)
382+
383+
384+
result = numpy.zeros((natm, ncol))
385+
386+
for i in range(natm):
387+
result[i, :] = Epv_atom(mol, mf, i, dm)
388+
212389
return result
213390

‎pyscf/pv/test/test_energy-pv.py‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
#!/usr/bin/env python
21
# Copyright 2014-2022 The PySCF Developers. All Rights Reserved.
32
#
43
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,13 +13,14 @@
1413
# limitations under the License.
1514
#
1615

17-
import numpy
1816
import unittest
19-
from pyscf import gto
20-
from pyscf import scf
21-
from pyscf import lib
17+
18+
import numpy
19+
20+
from pyscf import gto, lib, scf
2221
from pyscf.pv.energy import Epv_molecule
2322

23+
2424
def setUpModule():
2525
global mol, mf
2626

@@ -56,7 +56,15 @@ def tearDownModule():
5656
class KnownValues(unittest.TestCase):
5757
def test_pv(self):
5858
# Reference results (for PV) compared with DIRAC24
59-
Epv = Epv_molecule(mol, mf)
59+
Epv = Epv_molecule(mol, mf, dm=None)
60+
Epv_Oxigen = numpy.sum(Epv, axis=1)[0]
61+
self.assertAlmostEqual(Epv_Oxigen, -2.745821e-21, 5)
62+
self.assertAlmostEqual(mf.e_tot,-317.831176378,8)
63+
64+
def test_pv_dm(self):
65+
# Reference results (for PV) compared with DIRAC24, using the density matrix
66+
density_matrix = mf.make_rdm1()
67+
Epv = Epv_molecule(mol, mf, dm=density_matrix)
6068
Epv_Oxigen = numpy.sum(Epv, axis=1)[0]
6169
self.assertAlmostEqual(Epv_Oxigen, -2.745821e-21, 5)
6270
self.assertAlmostEqual(mf.e_tot,-317.831176378,8)

0 commit comments

Comments
 (0)