Source code for qrisp.vqe.problems.electronic_structure

"""
\********************************************************************************
* Copyright (c) 2025 the Qrisp authors
*
* This program and the accompanying materials are made available under the
* terms of the Eclipse Public License 2.0 which is available at
* http://www.eclipse.org/legal/epl-2.0.
*
* This Source Code may also be made available under the following Secondary
* Licenses when the conditions for such availability set forth in the Eclipse
* Public License, v. 2.0 are satisfied: GNU General Public License, version 2
* with the GNU Classpath Exception which is
* available at https://www.gnu.org/software/classpath/license.html.
*
* SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
********************************************************************************/
"""

import numpy as np

from qrisp import h, x, cx, ry, control, conjugate
from qrisp.operators.qubit import QubitOperator, QubitTerm
from qrisp.operators.fermionic import *
from functools import cache
import itertools
import math

#
# helper functions
#

def verify_symmetries(two_int):
    """
    Checks the symmetries of the two_electron integrals tensor in physicist's notation.

    Parameters
    ----------
    int_two : numpy.ndarray
        The two-electron integrals w.r.t. spin orbitals in physicist's notation.

    Returns
    -------

    """

    M = two_int.shape[0]
    for i in range(M):
        for j in range(M):
            for k in range(M):
                for l in range(M):
                    test1 = abs(two_int[i][j][k][l]-two_int[j][i][l][k])
                    test2 = abs(two_int[i][j][k][l]-two_int[k][l][i][j])
                    test3 = abs(two_int[i][j][k][l]-two_int[l][k][j][i])
                    if test1>1e-9 or test2>1e-9 or test3>1e-9:
                        return False
    return True

def delta(i,j):
    if i==j:
        return 1
    else:
        return 0

#
# spacial orbitals to spin orbitals
# 

def omega(x):
    return x%2

def spacial_to_spin(one_int,two_int):
    r"""
    Transforms one- and two-electron integrals w.r.t. $M$ spacial orbitals $\psi_0,\dotsc,\psi_{M-1}$ to 
    one- and two-electron integrals w.r.t. $2M$ spin orbitals $\chi_{2i}=\psi_{i}\alpha$, 
    $\chi_{2i+1}=\psi_{i}\beta$ for $i=0,\dotsc,M-1$.

    That is, given one- and two-electron integrals for spacial orbitals:

    .. math::

        h_{ij} &= \braket{\psi_i|h|\psi_j} \\
        h_{ijkl} &= \braket{\psi_i\psi_j|\psi_k\psi_l}

    this methods computes one- and two-electron integrals for spin orbitals:

    .. math::

        \tilde h_{ijkl} &= \braket{\chi_i|h|\chi_j} \\
        \tilde h_{ijkl} &= \braket{\chi_i\chi_j|\chi_k\chi_l}

    Parameters
    ----------
    one_int : numpy.ndarray
        The one-electron integrals w.r.t. spacial orbitals.
    two_int : numpy.ndarray
        The two-electron integrals w.r.t. spacial orbitals.

    Returns
    -------
    one_int_spin : numpy.ndarray
        The one-electron integrals w.r.t. spin orbitals.
    two_int_spin : numpy.ndarray
        The two-electron integrals w.r.t. spin orbitals.

    """

    num_spacial_orbs = one_int.shape[0]
    num_spin_orbs = 2*num_spacial_orbs

    # Initialize the spin-orbital one-electron integral tensor
    one_int_spin = np.zeros((num_spin_orbs, num_spin_orbs))

    for i in range(num_spacial_orbs):
        for j in range(num_spacial_orbs):

            one_int_spin[2*i][2*j] = one_int[i][j]
            
            one_int_spin[2*i+1][2*j+1] = one_int[i][j]

    # Initialize the spin-orbital two-electron integral tensor
    two_int_spin = np.zeros((num_spin_orbs, num_spin_orbs, num_spin_orbs, num_spin_orbs))

    for i in range(num_spacial_orbs):
        for j in range(num_spacial_orbs):  
            for k in range(num_spacial_orbs):
                for l in range (num_spacial_orbs):

                    two_int_spin[2*i][2*j+1][2*k+1][2*l] = two_int[i][j][k][l]

                    two_int_spin[2*i+1][2*j][2*k][2*l+1] = two_int[i][j][k][l]

                    two_int_spin[2*i][2*j][2*k][2*l] = two_int[i][j][k][l]

                    two_int_spin[2*i+1][2*j+1][2*k+1][2*l+1] = two_int[i][j][k][l]
    
    return one_int_spin, two_int_spin


[docs]
def electronic_data(mol):
    """
    A function that utilizes `restricted Hartree-Fock (RHF) <https://pyscf.org/user/scf.html>`_
    calculation in the `PySCF <https://pyscf.org>`_ quantum chemistry package to obtain the electronic data for
    defining an electronic structure problem.

    Parameters
    ----------
    mol : pyscf.gto.Mole
        The `molecule <https://pyscf.org/user/gto.html#>`_.

    Returns
    -------
    data : dict
        A dictionary specifying the electronic data for a molecule. The following data is provided:
        
        * ``one_int`` : numpy.ndarray
            The one-electron integrals w.r.t. spin orbitals (in physicists' notation).
        * ``two_int`` : numpy.ndarray
            The two-electron integrals w.r.t. spin orbitals (in physicists' notation).
        * ``num_orb`` : int
            The number of spin orbitals.
        * ``num_elec`` : int
            The number of electrons.
        * ``energy_nuc``
            The nuclear repulsion energy.
        * ``energy_hf``    
            The Hartree-Fock ground state energy.

    """

    from pyscf import scf, ao2mo

    data = {}

    threshold = 1e-9
    def apply_threshold(matrix, threshold):
        matrix[np.abs(matrix) < threshold] = 0
        return matrix
    
    # Set verbosity level to 0 to suppress output
    mol.verbose = 0
    
    # Perform a Hartree-Fock calculation
    mf = scf.RHF(mol)
    energy_hf = mf.kernel()

    # Extract one-electron integrals
    one_int = apply_threshold(mf.mo_coeff.T @ mf.get_hcore() @ mf.mo_coeff,threshold)
    # Extract two-electron integrals (electron repulsion integrals)
    two_int = ao2mo.kernel(mol,mf.mo_coeff)
    # Full tensor with chemist's notation
    two_int = apply_threshold(ao2mo.restore(1,two_int,mol.nao_nr()),threshold)
    # Full tensor with physicist's notation
    two_int = np.transpose(two_int,(0,2,3,1))

    # Convert spacial orbital to spin orbitals
    one_int, two_int  = spacial_to_spin(one_int,two_int)

    data['mol'] = mol
    data['one_int'] = one_int
    data['two_int'] = two_int
    data['num_orb'] = 2*mf.mo_coeff.shape[0]  # Number of spin orbitals
    data['num_elec'] = mol.nelectron
    data['energy_nuc'] = mol.energy_nuc()
    data['energy_hf'] = energy_hf

    return data



[docs]
def create_electronic_hamiltonian(arg, active_orb=None, active_elec=None):
    """
    Creates the qubit Hamiltonian for an electronic structure problem. 
    If an Active Space (AS) is specified, the Hamiltonian is calculated following this `paper <https://arxiv.org/abs/2009.01872>`_.
    
    Parameters
    ----------
    arg : pyscf.gto.Mole or dict
        A PySCF `molecule <https://pyscf.org/user/gto.html#>`_ or
        a dictionary specifying the electronic data for a molecule. The following data is required:
        
        * ``one_int`` : numpy.ndarray
            The one-electron integrals w.r.t. spin orbitals (in physicists' notation).
        * ``two_int`` : numpy.ndarray
            The two-electron integrals w.r.t. spin orbitals (in physicists' notation).
        * ``num_orb`` : int
            The number of spin orbitals.
        * ``num_elec`` : int
            The number of electrons.

    active_orb : int, optional
        The number of active spin orbitals.
    active_elec : int, optional
        The number of active electrons.

    Returns
    -------
    H : :ref:`FermionicOperator`
        The fermionic Hamiltonian.
    
    Examples
    --------

    We calucalte the fermionic Hamiltonian for the Hydrogen molecule, and transform it to a Pauli Hamiltonian via Jordan-Wigner transform.

    ::

        from pyscf import gto
        from qrisp.vqe.problems.electronic_structure import *

        mol = gto.M(
            atom = '''H 0 0 0; H 0 0 0.74''',
            basis = 'sto-3g')

        H = create_electronic_hamiltonian(mol)
        H.to_qubit_operator()   

    Yields:

    .. math::

        -&0.812170607248714 - 0.0453026155037992 X_0X_1Y_2Y_3 + 0.0453026155037992 X_0Y_1Y_2X_3 - 0.0453026155037992 Y_0Y_1X_2X_3

        &+0.171412826447769 Z_0 + 0.168688981703612 Z_0Z_1 + 0.120625234833904 Z_0Z_2 + 0.165927850337703 Z_0Z_3 + 0.171412826447769 Z_1

        &+0.165927850337703 Z_1Z_2 + 0.120625234833904 Z_1Z_3 - 0.223431536908133 Z_2 + 0.174412876122615Z Z_2Z_3 - 0.223431536908133 Z_3

    """

    import pyscf

    if isinstance(arg,pyscf.gto.Mole):
        data = electronic_data(arg)
    elif isinstance(arg,dict):
        data = arg
        if not verify_symmetries(data['two_int']):
            raise Warning("Failed to verify symmetries for two-electron integrals")
    else:
        raise TypeError("Cannot create electronic Hamiltonian from type "+str(type(arg)))

    one_int = data['one_int']
    two_int = data['two_int']
    M = data['num_orb']
    N = data['num_elec']
    K = active_orb
    L = active_elec

    if K is None or L is None:
        K = M
        L = N

    if L>N or K>M or K<L or K+N-L>M:
        raise Exception("Invalid number of active electrons or orbitals")

    # number of inactive electrons 
    I = N-L

    # inactive Fock operator
    F = one_int.copy()
    for p in range(M):
        for q in range(M):
            for i in range(I):
                #F[p][q] += (two_int[i][p][i][q]-two_int[i][q][p][i])
                F[p][q] += (two_int[i][p][q][i]-two_int[i][q][i][p])

    # inactive energy
    E = 0
    for j in range(I):
        E += (one_int[j][j]+F[j][j])/2

    # Hamiltonian
    H=E
    res_dict = {}
    for i in range(K):
        for j in range(K):
            if F[I+i][I+j]!=0:
                term = FermionicTerm([(j, False), (i, True)])
                res_dict[term] = F[I+i][I+j]
                #H += F[I+i][I+j]*c(i)*a(j)
    
    
    for i in range(K):
        for j in range(K): 
            for k in range(K):
                for l in range(K):
                    if two_int[I+i][I+j][I+k][I+l]!=0 and i!=j and k!=l:
                        term = FermionicTerm([(l, False), (k, False), (j, True), (i, True)])
                        res_dict[term] = (0.5*two_int[I+i][I+j][I+k][I+l])
                        # H += (0.5*two_int[I+i][I+j][I+k][I+l])*c(i)*c(j)*a(k)*a(l)
    temp_H = FermionicOperator(res_dict)
    H = E+ temp_H
    return H.reduce()


#
# ansatz
#

def conjugator(i,j):
    h(i)
    cx(i,j)

def pswap(phi,i,j):
    with conjugate(conjugator)(i,j):
        ry(-phi/2, [i,j])  

def conjugator2(i,j,k,l):
    cx(i,j)
    cx(k,l) 

def pswap2(phi,i,j,k,l):
    with conjugate(conjugator2)(i,j,k,l):
        with control([j,l],ctrl_state='00'):
            pswap(phi,i,k)


[docs]
def create_QCCSD_ansatz(M,N):
    r"""
    This method creates a function for applying one layer of the `QCCSD ansatz <https://arxiv.org/abs/2005.08451>`_.

    The chemistry-inspired Qubit Coupled Cluster Single Double (QCCSD) ansatz evolves the initial state, 
    usually, the Hartree-Fock state

    .. math::

        \ket{\Psi_{\text{HF}}}=\ket{1_0,\dotsc,1_N,0_{N+1},\dotsc,0_M}
    
    under the action of parametrized (non-commuting) single and double excitation unitaries.
    
    The single (S) excitation unitaries $U_i^r$ implement a continuous swap for qubits $i$ and $r$:
    
    .. math::

        U_i^r(\theta) = \begin{pmatrix}
                        1&0&0&0\\
                        0&\cos(\theta)&-\sin(\theta)&0\\
                        0&\sin(\theta)&\cos(\theta)&0\\
                        0&0&0&1
                        \end{pmatrix}

    Similarly, the double (D) excitation unitaries $U_{ij}^{rs}(\theta)$ implement a continuous swap 
    for qubit pairs $i,j$ and $r,s$.
        
    Parameters
    ----------
    M : int
        The number of (active) spin orbitals.
    N : int
        The number of (active) electrons.

    Returns
    -------
    ansatz : function
        A function that can be applied to a :ref:`QuantumVariable` and a list of parameters.
    num_params : int
        The number of parameters.
    
    """

    spin_down_occupied = [i for i in range(N) if i%2==0]
    spin_down_virtual = [i for i in range(N,M) if i%2==0]
    spin_up_occupied = [i for i in range(N) if i%2==1]
    spin_up_virtual = [i for i in range(N,M) if i%2==1]

    num_singles = len(spin_down_occupied)*len(spin_down_virtual) + len(spin_up_occupied)*len(spin_up_virtual)

    num_doubles = len(spin_down_occupied)*len(spin_up_occupied)*len(spin_down_virtual)*len(spin_up_virtual) \
                    +math.comb(len(spin_down_occupied),2)*math.comb(len(spin_down_virtual),2) \
                    +math.comb(len(spin_up_occupied),2)*math.comb(len(spin_up_virtual),2)
    
    num_params = num_singles + num_doubles

    def ansatz(qv, theta):

        num_params = 0
        # Single excitations
        for i in spin_down_occupied:
            for j in spin_down_virtual:
                    pswap(theta[num_params],qv[i],qv[j])
                    num_params += 1
        
        for i in spin_up_occupied:
            for j in spin_up_virtual:
                    pswap(theta[num_params],qv[i],qv[j])
                    num_params += 1
        
        # Double excitation
        for i in spin_down_occupied:
            for j in spin_up_occupied:
                for k in spin_down_virtual:
                    for l in spin_up_virtual:
                        pswap2(theta[num_params],qv[i],qv[j],qv[k],qv[l])
                        num_params += 1

        for i,j in itertools.combinations(spin_down_occupied,2):
            for k,l in itertools.combinations(spin_down_virtual,2):
                pswap2(theta[num_params],qv[i],qv[j],qv[k],qv[l])
                num_params += 1

        for i,j in itertools.combinations(spin_up_occupied,2):
            for k,l in itertools.combinations(spin_up_virtual,2):
                pswap2(theta[num_params],qv[i],qv[j],qv[k],qv[l])
                num_params += 1

    return ansatz, num_params



[docs]
def create_hartree_fock_init_function(M, N):
    """
    Creates the function that, when applied to a :ref:`QuantumVariable`, initializes the Hartee-Fock state:
    Consistent with the Jordan-Wigner mapping, the first ``N`` qubits are initialized in the $\ket{1}$ state.

    Parameters
    ----------
    M : int
        The number of (active) spin orbitals.
    N : int
        The number of (active) electrons.

    Returns
    -------
    init_function : function
        A function that can be applied to a :ref:`QuantumVariable`.

    """

    def init_function(qv):
        for i in range(N):
            x(qv[i])

    """
    if mapping_type=='parity':
        def init_function(qv):
            for i in range(N//2):
                x(qv[2*i])
            # odd number of electrons
            if N%2==1:
                for i in range(N-1,M):
                    x(qv[i])
    """

    return init_function




[docs]
def electronic_structure_problem(arg, active_orb=None, active_elec=None, ansatz_type='QCCSD', threshold=1e-4):
    r"""
    Creates a VQE problem instance for an electronic structure problem defined by the 
    one-electron and two-electron integrals for the spin orbitals (in physicists' notation).

    The problem Hamiltonian is given by:

    .. math::

        H = \sum\limits_{i,j=0}^{M-1}h_{i,j}a^{\dagger}_ia_j + \sum\limits_{i,j,k,l=0}^{M-1}h_{i,j,k,l}a^{\dagger}_ia^{\dagger}_ja_ka_l
    
    for one-electron integrals:

    .. math::

        h_{i,j}=\int\mathrm dx \chi_i^*(x)\chi_j(x)

    and two-electron integrals:

    .. math::

        h_{i,j,k,l} = \int\mathrm dx_1\mathrm dx_2\chi_i^*(x_1)\chi_j^*(x_2)\frac{1}{|r_1-r_2|}\chi_k(x_1)\chi_l(x_2)

    Parameters
    ----------
    arg : pyscf.gto.Mole or dict
        A PySCF `molecule <https://pyscf.org/user/gto.html#>`_ or
        a dictionary specifying the electronic data for a molecule. The following data is required:
        
        * ``one_int`` : numpy.ndarray
            The one-electron integrals w.r.t. spin orbitals (in physicists' notation).
        * ``two_int`` : numpy.ndarray
            The two-electron integrals w.r.t. spin orbitals (in physicists' notation).
        * ``num_orb`` : int
            The number of spin orbitals.
        * ``num_elec`` : int
            The number of electrons.

    active_orb : int, optional
        The number of active spin orbitals.
    active_elec : int, optional
        The number of active electrons.
    ansatz_type : string, optional
        The ansatz type. Availabe is ``QCCSD``. The default is ``QCCSD``.
    threshold : float, optional
        The threshold for the absolute value of the coefficients of Pauli products in the quantum Hamiltonian. The default is 1e-4.

    Returns
    -------
    VQEProblem
        The VQE problem instance.

    Examples
    --------

    We calculate the electronic energy for the Hydrogen molecule at bond distance 0.74 angstroms:

    ::

        from pyscf import gto
        from qrisp import QuantumVariable
        from qrisp.vqe.problems.electronic_structure import *

        mol = gto.M(
            atom = '''H 0 0 0; H 0 0 0.74''',
            basis = 'sto-3g')

        vqe = electronic_structure_problem(mol)
        vqe.set_callback()

        energy = vqe.run(QuantumVariable(4),depth=1,max_iter=50)
        print(energy)
        #Yields -1.8461290172512965
    
    """
    from qrisp.vqe import VQEProblem
    import pyscf
    if isinstance(arg,pyscf.gto.Mole):
        data = electronic_data(arg)
    elif isinstance(arg,dict):
        data = arg
        if not verify_symmetries(data['two_int']):
            raise Warning("Failed to verify symmetries for two-electron integrals")
    else:
        raise TypeError("Cannot instantiate VQEProblem from type "+str(type(arg)))
    
    M = data['num_orb']
    N = data['num_elec']
    K = active_orb
    L = active_elec

    if K is None or L is None:
        K = M
        L = N

    ansatz, num_params = create_QCCSD_ansatz(K,L)

    fermionic_hamiltonian = create_electronic_hamiltonian(data,K,L)
    hamiltonian = fermionic_hamiltonian.to_qubit_operator(mapping_type='jordan_wigner')
    hamiltonian.apply_threshold(threshold)

    return VQEProblem(hamiltonian, ansatz, num_params, init_function=create_hartree_fock_init_function(K,L))