CoCalc -- classical

GitHub Repository: NVIDIA/cuda-q-academic
Path: blob/main/chemistry-simulations/aux_files/krylov/qchem/classical_pyscf.py
⁵⁸³ views
1
# ============================================================================ #
2
# Copyright (c) 2024 - 2025 NVIDIA Corporation & Affiliates.                   #
3
# All rights reserved.                                                         #
4
#                                                                              #
5
# This source code and the accompanying materials are made available under     #
6
# the terms of the Apache License 2.0 which accompanies this distribution.     #
7
# ============================================================================ #
8

9
import json
10
from functools import reduce
11

12
try:
13
    import numpy as np
14
except ValueError:
15
    print('numpy should be installed.')
16

17
try:
18
    from pyscf import gto, scf, cc, ao2mo, mp, mcscf, fci
19
except ValueError:
20
    print('PySCF should be installed. Use pip install pyscf')
21

22
from pyscf.tools import molden
23

24
#######################################################
25
# Generate the spin molecular orbital hamiltonian
26

27

28
def generate_molecular_spin_ham_restricted(h1e, h2e, ecore):
29

30
    # This function generates the molecular spin Hamiltonian
31
    # H = E_core+sum_{`pq`}  h_{`pq`} a_p^dagger a_q +
32
    #                          0.5 * h_{`pqrs`} a_p^dagger a_q^dagger a_r a_s
33
    # h1e: one body integrals h_{`pq`}
34
    # h2e: two body integrals h_{`pqrs`}
35
    # `ecore`: constant (nuclear repulsion or core energy in the active space Hamiltonian)
36

37
    # Total number of qubits equals the number of spin molecular orbitals
38
    nqubits = 2 * h1e.shape[0]
39

40
    # Initialization
41
    one_body_coeff = np.zeros((nqubits, nqubits))
42
    two_body_coeff = np.zeros((nqubits, nqubits, nqubits, nqubits))
43

44
    ferm_ham = []
45

46
    for p in range(nqubits // 2):
47
        for q in range(nqubits // 2):
48

49
            # p & q have the same spin <a|a>= <b|b>=1
50
            # <a|b>=<b|a>=0 (orthogonal)
51
            one_body_coeff[2 * p, 2 * q] = h1e[p, q]
52
            temp = str(h1e[p, q]) + ' a_' + str(p) + '^dagger ' + 'a_' + str(q)
53
            ferm_ham.append(temp)
54
            one_body_coeff[2 * p + 1, 2 * q + 1] = h1e[p, q]
55
            temp = str(h1e[p, q]) + ' b_' + str(p) + '^dagger ' + 'b_' + str(q)
56
            ferm_ham.append(temp)
57

58
            for r in range(nqubits // 2):
59
                for s in range(nqubits // 2):
60

61
                    # Same spin (`aaaa`, `bbbbb`) <a|a><a|a>, <b|b><b|b>
62
                    two_body_coeff[2 * p, 2 * q, 2 * r,
63
                                   2 * s] = 0.5 * h2e[p, q, r, s]
64
                    temp = str(0.5 * h2e[p, q, r, s]) + ' a_' + str(
65
                        p) + '^dagger ' + 'a_' + str(
66
                            q) + '^dagger ' + 'a_' + str(r) + ' a_' + str(s)
67
                    ferm_ham.append(temp)
68
                    two_body_coeff[2 * p + 1, 2 * q + 1, 2 * r + 1,
69
                                   2 * s + 1] = 0.5 * h2e[p, q, r, s]
70
                    temp = str(0.5 * h2e[p, q, r, s]) + ' b_' + str(
71
                        p) + '^dagger ' + 'b_' + str(
72
                            q) + '^dagger ' + 'b_' + str(r) + ' b_' + str(s)
73
                    ferm_ham.append(temp)
74

75
                    # Mixed spin(`abab`, `baba`) <a|a><b|b>, <b|b><a|a>
76
                    #<a|b>= 0 (orthogonal)
77
                    two_body_coeff[2 * p, 2 * q + 1, 2 * r + 1,
78
                                   2 * s] = 0.5 * h2e[p, q, r, s]
79
                    temp = str(0.5 * h2e[p, q, r, s]) + ' a_' + str(
80
                        p) + '^dagger ' + 'a_' + str(
81
                            q) + '^dagger ' + 'b_' + str(r) + ' b_' + str(s)
82
                    ferm_ham.append(temp)
83
                    two_body_coeff[2 * p + 1, 2 * q, 2 * r,
84
                                   2 * s + 1] = 0.5 * h2e[p, q, r, s]
85
                    temp = str(0.5 * h2e[p, q, r, s]) + ' b_' + str(
86
                        p) + '^dagger ' + 'b_' + str(
87
                            q) + '^dagger ' + 'a_' + str(r) + ' a_' + str(s)
88
                    ferm_ham.append(temp)
89

90
    full_hamiltonian = " + ".join(ferm_ham)
91

92
    return one_body_coeff, two_body_coeff, ecore, full_hamiltonian
93

94

95
####################################################################
96

97
# A- Gas phase simulation
98
#############################
99
## Beginning of simulation
100
#############################
101

102
def get_mol_hamiltonian(xyz:str, spin:int, charge: int, basis:str, symmetry:bool = False, memory:float = 4000, cycles:int = 100, \
103
                              initguess:str = 'minao', nele_cas = None, norb_cas = None, MP2:bool = False, natorb:bool = False,\
104
                              casci:bool = False, ccsd:bool = False, casscf:bool = False, integrals_natorb:bool = False, \
105
                                integrals_casscf:bool = False, viz_orb:bool = False, verbose:bool = False):
106
    ################################
107
    # Initialize the molecule
108
    ################################
109
    filename = xyz.split('.')[0]
110

111
    if (nele_cas is None) and (norb_cas is not None):
112
        raise ValueError(
113
            "WARN: nele_cas is None and norb_cas is not None. "
114
            "nele_cas and norb_cas should be either both None or have values")
115

116
    if (nele_cas is not None) and (norb_cas is None):
117
        raise ValueError(
118
            "WARN: nele_cas is not None and norb_cas is None. "
119
            "nele_cas and norb_cas should be either both None or have values")
120

121
    ########################################################################
122
    # To add (coming soon)
123

124
    mol = gto.M(atom=xyz,
125
                spin=spin,
126
                charge=charge,
127
                basis=basis,
128
                max_memory=memory,
129
                symmetry=symmetry,
130
                output=filename + '-pyscf.log',
131
                verbose=4)
132

133
    ##################################
134
    # Mean field (HF)
135
    ##################################
136

137
    if spin == 0:
138
        myhf = scf.RHF(mol)
139
        myhf.max_cycle = cycles
140
        myhf.chkfile = filename + '-pyscf.chk'
141
        myhf.init_guess = initguess
142
        myhf.kernel()
143

144
        norb = myhf.mo_coeff.shape[1]
145
        if verbose:
146
            print('[pyscf] Total number of orbitals = ', norb)
147
    else:
148
        myhf = scf.ROHF(mol)
149
        myhf.max_cycle = cycles
150
        myhf.chkfile = filename + '-pyscf.chk'
151
        myhf.init_guess = initguess
152
        myhf.kernel()
153

154
        norb = myhf.mo_coeff[0].shape[1]
155
        if verbose:
156
            print('[pyscf] Total number of orbitals = ', norb)
157

158
    nelec = mol.nelectron
159
    if verbose:
160
        print('[pyscf] Total number of electrons = ', nelec)
161
    if verbose:
162
        print('[pyscf] HF energy = ', myhf.e_tot)
163
    if viz_orb:
164
        molden.from_mo(mol, filename + '_HF_molorb.molden', myhf.mo_coeff)
165

166
    ##########################
167
    # MP2
168
    ##########################
169
    if MP2:
170

171
        if spin != 0:
172
            raise ValueError("WARN: ROMP2 is unvailable in pyscf.")
173
        else:
174
            mymp = mp.MP2(myhf)
175
            mp_ecorr, mp_t2 = mymp.kernel()
176
            if verbose:
177
                print('[pyscf] R-MP2 energy= ', mymp.e_tot)
178

179
            if integrals_natorb or natorb:
180
                # Compute natural orbitals
181
                noons, natorbs = mcscf.addons.make_natural_orbitals(mymp)
182
                if verbose:
183
                    print(
184
                        '[pyscf] Natural orbital occupation number from R-MP2: '
185
                    )
186
                if verbose:
187
                    print(noons)
188
                if viz_orb:
189
                    molden.from_mo(mol, filename + '_MP2_natorb.molden',
190
                                   natorbs)
191

192
    #######################################
193
    # CASCI if active space is defined
194
    # FCI if the active space is None
195
    ######################################
196
    if casci:
197

198
        if nele_cas is None:
199
            myfci = fci.FCI(myhf)
200
            result = myfci.kernel()
201
            if verbose:
202
                print('[pyscf] FCI energy = ', result[0])
203

204
        else:
205
            if natorb and (spin == 0):
206
                mycasci = mcscf.CASCI(myhf, norb_cas, nele_cas)
207
                mycasci.kernel(natorbs)
208
                if verbose:
209
                    print('[pyscf] R-CASCI energy using natural orbitals= ',
210
                          mycasci.e_tot)
211

212
            elif natorb and (spin != 0):
213
                raise ValueError(
214
                    "WARN: Natural orbitals cannot be computed. ROMP2 is unvailable in pyscf."
215
                )
216

217
            else:
218
                mycasci_mo = mcscf.CASCI(myhf, norb_cas, nele_cas)
219
                mycasci_mo.kernel()
220
                if verbose:
221
                    print('[pyscf] R-CASCI energy using molecular orbitals= ',
222
                          mycasci_mo.e_tot)
223

224
    ########################
225
    # CCSD
226
    ########################
227
    if ccsd:
228

229
        if nele_cas is None:
230
            mycc = cc.CCSD(myhf)
231
            mycc.max_cycle = cycles
232
            mycc.kernel()
233
            if verbose:
234
                print('[pyscf] Total R-CCSD energy = ', mycc.e_tot)
235

236
        else:
237
            mc = mcscf.CASCI(myhf, norb_cas, nele_cas)
238
            frozen = []
239
            frozen += [y for y in range(0, mc.ncore)]
240
            frozen += [
241
                y for y in range(mc.ncore + norb_cas, len(myhf.mo_coeff))
242
            ]
243
            if natorb and (spin == 0):
244
                mycc = cc.CCSD(myhf, frozen=frozen, mo_coeff=natorbs)
245
                mycc.max_cycle = cycles
246
                mycc.kernel()
247
                if verbose:
248
                    print(
249
                        '[pyscf] R-CCSD energy of the active space using natural orbitals= ',
250
                        mycc.e_tot)
251

252
            elif natorb and (spin != 0):
253
                raise ValueError(
254
                    "WARN: Natural orbitals cannot be computed. ROMP2 is unvailable in pyscf."
255
                )
256

257
            else:
258
                mycc = cc.CCSD(myhf, frozen=frozen)
259
                mycc.max_cycle = cycles
260
                mycc.kernel()
261
                if verbose:
262
                    print(
263
                        '[pyscf] R-CCSD energy of the active space using molecular orbitals= ',
264
                        mycc.e_tot)
265

266
    #########################
267
    # CASSCF
268
    #########################
269
    if casscf:
270
        if nele_cas is None:
271
            raise ValueError("WARN: You should define the active space.")
272

273
        if natorb and (spin == 0):
274
            mycas = mcscf.CASSCF(myhf, norb_cas, nele_cas)
275
            mycas.max_cycle_macro = cycles
276
            mycas.kernel(natorbs)
277
            if verbose:
278
                print('[pyscf] R-CASSCF energy using natural orbitals= ',
279
                      mycas.e_tot)
280

281
        elif natorb and (spin != 0):
282
            raise ValueError(
283
                "WARN: Natural orbitals cannot be computed. ROMP2 is unvailable in pyscf."
284
            )
285

286
        else:
287
            mycas = mcscf.CASSCF(myhf, norb_cas, nele_cas)
288
            mycas.max_cycle_macro = cycles
289
            mycas.kernel()
290
            if verbose:
291
                print('[pyscf] R-CASSCF energy using molecular orbitals= ',
292
                      mycas.e_tot)
293

294
    ###################################
295
    # CASCI: `FCI` of the active space
296
    ##################################
297
    if casci and casscf:
298

299
        if natorb and (spin != 0):
300
            raise ValueError(
301
                "WARN: Natural orbitals cannot be computed. ROMP2 is unavailable in pyscf."
302
            )
303
        else:
304
            h1e_cas, ecore = mycas.get_h1eff()
305
            h2e_cas = mycas.get_h2eff()
306

307
            e_fci, fcivec = fci.direct_spin1.kernel(h1e_cas,
308
                                                    h2e_cas,
309
                                                    norb_cas,
310
                                                    nele_cas,
311
                                                    ecore=ecore)
312
            if verbose:
313
                print('[pyscf] R-CASCI energy using the casscf orbitals= ',
314
                      e_fci)
315

316
    ###################################################################################
317
    # Computation of one- and two- electron integrals for the active space Hamiltonian
318
    ###################################################################################
319

320
    if nele_cas is None:
321
        # Compute the 1e integral in atomic orbital then convert to HF basis
322
        h1e_ao = mol.intor("int1e_kin") + mol.intor("int1e_nuc")
323
        ## Ways to convert from `ao` to mo
324
        #h1e=`np.einsum('pi,pq,qj->ij', myhf.mo_coeff, h1e_ao, myhf.mo_coeff)`
325
        h1e = reduce(np.dot, (myhf.mo_coeff.T, h1e_ao, myhf.mo_coeff))
326
        #h1e=`reduce(np.dot, (myhf.mo_coeff.conj().T, h1e_ao, myhf.mo_coeff))`
327

328
        # Compute the 2e integrals then convert to HF basis
329
        h2e_ao = mol.intor("int2e_sph", aosym='1')
330
        h2e = ao2mo.incore.full(h2e_ao, myhf.mo_coeff)
331

332
        # `Reorder the chemist notation (pq|rs) ERI h_prqs to h_pqrs`
333
        # a_p^dagger a_r a_q^dagger a_s --> a_p^dagger a_q^dagger a_r a_s
334
        h2e = h2e.transpose(0, 2, 3, 1)
335

336
        nuclear_repulsion = myhf.energy_nuc()
337

338
        # Compute the molecular spin electronic Hamiltonian from the
339
        # molecular electron integrals
340
        obi, tbi, e_nn, ferm_ham = generate_molecular_spin_ham_restricted(
341
            h1e, h2e, nuclear_repulsion)
342

343
    else:
344

345
        if integrals_natorb:
346
            if spin != 0:
347
                raise ValueError("WARN: ROMP2 is unvailable in pyscf.")
348
            else:
349
                mc = mcscf.CASCI(myhf, norb_cas, nele_cas)
350
                h1e_cas, ecore = mc.get_h1eff(natorbs)
351
                h2e_cas = mc.get_h2eff(natorbs)
352
                h2e_cas = ao2mo.restore('1', h2e_cas, norb_cas)
353
                h2e_cas = np.asarray(h2e_cas.transpose(0, 2, 3, 1), order='C')
354

355
        elif integrals_casscf:
356
            if casscf:
357
                h1e_cas, ecore = mycas.get_h1eff()
358
                h2e_cas = mycas.get_h2eff()
359
                h2e_cas = ao2mo.restore('1', h2e_cas, norb_cas)
360
                h2e_cas = np.asarray(h2e_cas.transpose(0, 2, 3, 1), order='C')
361
            else:
362
                raise ValueError(
363
                    "WARN: You need to run casscf. Use casscf=True.")
364

365
        else:
366
            mc = mcscf.CASCI(myhf, norb_cas, nele_cas)
367
            h1e_cas, ecore = mc.get_h1eff(myhf.mo_coeff)
368
            h2e_cas = mc.get_h2eff(myhf.mo_coeff)
369
            h2e_cas = ao2mo.restore('1', h2e_cas, norb_cas)
370
            h2e_cas = np.asarray(h2e_cas.transpose(0, 2, 3, 1), order='C')
371

372
        # Compute the molecular spin electronic Hamiltonian from the
373
        # molecular electron integrals
374
        obi, tbi, core_energy, ferm_ham = generate_molecular_spin_ham_restricted(
375
            h1e_cas, h2e_cas, ecore)
376

377
    # Dump obi and `tbi` to binary file.
378
    # `obi.astype(complex).tofile(f'{filename}_one_body.dat')`
379
    # `tbi.astype(complex).tofile(f'{filename}_two_body.dat')`
380

381
    ######################################################
382
    # Dump energies / etc to a metadata file
383
    if nele_cas is None:
384
        # `metadata = {'num_electrons': nelec, 'num_orbitals': norb, 'nuclear_energy': e_nn, 'hf_energy': myhf.e_tot}`
385
        # with open(f'{filename}_metadata.`json`', 'w') as f:
386
        #        `json`.dump(metadata, f)
387
        return (obi, tbi, e_nn, nelec, norb, ferm_ham)
388

389
    else:
390
        # `metadata = {'num_electrons_cas': nele_cas, 'num_orbitals_cas': norb_cas, 'core_energy': ecore, 'hf_energy': myhf.e_tot}`
391
        # with open(f'{filename}_metadata.`json`', 'w') as f:
392
        #        `json`.dump(metadata, f)
393
        return (obi, tbi, ecore, nele_cas, norb_cas, ferm_ham)
394

395

396
#######################################################
397

398
Product

Resources

Company
