VQE-LD

def pauli_string_fermion_operator(p, op_type):
    """
    Map a fermionic creation (0) or annihilation (1) operator at orbital p
    to its equivalent Pauli string via Jordan-Wigner transform.

    Parameters:
        p (int): orbital index
        op_type (int): 0 for creation, 1 for annihilation

    Returns:
        qml.operation.Tensor: Pauli string operator
    """
    f = (2 * op_type - 1) * 1j  # +i for creation, -i for annihilation
    Q_p = 0.5 * (qml.PauliX(p) + f * qml.PauliY(p))

    for i in range(p):
        Q_p = qml.PauliZ(i) @ Q_p

    return Q_p


def pauli_string_excitation_operator(p, q):
    """
    Constructs the Pauli string representation of the fermionic excitation operator a_p† a_q
    using the Jordan-Wigner transformation.

    This operator corresponds to exciting an electron from orbital q to orbital p
    (i.e., annihilation at q followed by creation at p).

    Args:
        p (int): Index of the target orbital (creation operator a_p†).
        q (int): Index of the source orbital (annihilation operator a_q).

    Returns:
        qml.operation.Tensor: The corresponding Pauli string operator representing a_p† a_q.

    Notes:
        - This is a non-Hermitian operator representing a single excitation.
        - To obtain a Hermitian excitation generator (for variational ansätze), you would
          need to combine this operator with its Hermitian conjugate (a_q† a_p).
        - The mapping relies on the Jordan-Wigner transformation, implemented via the
          `pauli_string_fermion_operator` function for single fermionic operators.
    """
    return pauli_string_fermion_operator(p, CREATION) @ pauli_string_fermion_operator(
        q, ANNIHILATION
    )


def rdm1_real_operator(p, q):
    """
    Constructs the Hermitian operator O_real = 1/2 (a_p† a_q + a_q† a_p),
    whose expectation value yields the real part of the 1-RDM element.

    Args:
        p (int): Index of the first orbital.
        q (int): Index of the second orbital.

    Returns: The Hermitian qubit operator corresponding to 1/2 (a_p† a_q + a_q† a_p).
    """
    op = pauli_string_excitation_operator(p, q)
    op_dag = pauli_string_excitation_operator(q, p)
    O_real = 0.5 * (op + op_dag)
    return O_real


def rdm1_imag_operator(p, q):
    """
    Constructs the Hermitian operator O_imag = (1/2i)(a_p† a_q - a_q† a_p),
    whose expectation value yields the imaginary part of the 1-RDM element.

    Args:
        p (int): Index of the first orbital.
        q (int): Index of the second orbital.

    Returns: The Hermitian qubit operator corresponding to (1/2i)(a_p† a_q - a_q† a_p).
    """
    op = pauli_string_excitation_operator(p, q)
    op_dag = pauli_string_excitation_operator(q, p)
    O_imag = 0.5j * (op - op_dag)  # same as (1/(2i))(op - op_dag)
    return O_imag


def rdm1_observables(qubits_active):
    """
    Construct all the observables needed to compute the 1-RDM.
    Diagonal elements correspond to occupation numbers,
    while off-diagonal elements are separated into real and imaginary components.

    Args: Number of active qubits in the system.

    Returns:
        observables (dict): Dictionary mapping keys to PennyLane observables.
        ordered_keys (list): Sorted list of keys for consistent ordering.
    """
    observables = {}
    for p in range(qubits_active):
        for q in range(p, qubits_active):
            observables[(p, q, "re")] = rdm1_real_operator(p, q)
            observables[(p, q, "im")] = rdm1_imag_operator(p, q)

    return observables, observables.keys()

geom = """
     C  0.0000000  0.1525520  0.0000000
     H  1.1165590  0.3217700  0.0000000
     H -0.5550270 -1.0611280  0.0000000
     H  0.3813020 -1.1309260  0.0000000
     H -0.4714170  0.4774870  0.9592110
     H -0.4714170  0.4774870 -0.9592110
     """

geom_data = geom.split()
symbols = tuple(geom_data[4 * i] for i in range(len(geom_data) // 4))
coordinates = tuple(geom_data[4 * i + 1 : 4 * i + 4] for i in range(len(geom_data) // 4))
coordinates = jnp.array(coordinates, dtype=float).reshape((-1, 3)) * Ang2Bohr

charge = 1
mult = 1
nelec = 10
basis = "STO-3G"
active_electrons = 2
active_orbitals = 2

# Define the Molecule in PennyLane
qml_mol = qml.qchem.Molecule(symbols, coordinates, charge=charge, mult=mult, basis_name=basis)

# Calculate the Hamiltonian
H_active, qubits_active = qml.qchem.molecular_hamiltonian(
    qml_mol, active_electrons=active_electrons, active_orbitals=active_orbitals
)

# Get the Hartree-Fock state
ref_active = qml.qchem.hf_state(active_electrons, qubits_active)

# Define the device from PennyLane
dev = qml.device("default.qubit", wires=range(qubits_active))


@qml.qnode(dev, interface="jax")
def circuit(weights):
    """
    Executes a variational quantum circuit using an ansatz and computes
    the expectation value of the active space Hamiltonian.

    Args:
        weights (array): Variational parameters for the ansatz.

    Returns:
        float: Expectation value ⟨H_active⟩, i.e., the estimated energy of the quantum state.
    """
    qml.GateFabric(weights, wires=range(qubits_active), init_state=ref_active, include_pi=True)
    return qml.expval(H_active)


param_shape = qml.GateFabric.shape(n_layers=10, n_wires=qubits_active)

# Set observables to 1-RDM
observables, keys = rdm1_observables(qubits_active)
for k in keys:
    observables[k] = observables[k].simplify()


# Expectation value of the Hermitian one-body excitation operators
@qml.qnode(dev, interface="jax")
def rdm1_expectation(params):
    """
    Executes the variational quantum circuit and computes the expectation values of the
    one-body excitation operators needed to reconstruct the real part of the
    1-RDM for all orbital pairs (p, q).

    Args:
        params (array): Array of trainable parameters for the GateFabric ansatz.

    Returns:
        list[float]: A flat list of real-valued expectation values ⟨O_real⟩ for every
                     pair (p, q), corresponding to the real part of the 1-RDM.

    Key Steps:
        1. Prepares the quantum state using the GateFabric ansatz.
        2. For every pair of orbitals (p, q), construct the O_real and O_imag operators.
        3. Measures the expectation value.
        4. Returns all expectation values as a list.
    """
    qml.GateFabric(params, wires=range(qubits_active), init_state=ref_active, include_pi=True)
    return [qml.expval(observables[k]) for k in keys]


def compute_rdm1_active(params):
    """
    Computes the expectation values of the 1-RDM operators and organizes the results
    into a square matrix of shape (qubits_active × qubits_active).

    Args:
        params (array): Variational parameters for the quantum circuit.

    Returns:
        jnp.ndarray: The reconstructed real part of the one-electron 1-RDM
                     in the active space, as a (qubits_active, qubits_active) matrix.
    """
    values = rdm1_expectation(params)
    n = qubits_active
    rdm = jnp.zeros((n, n), dtype=complex)

    for val, key in zip(values, keys):
        p, q, part = key
        if part == "re":
            rdm = rdm.at[p, q].add(val)
            if p != q:
                rdm = rdm.at[q, p].add(val)
        else:  # Imaginary component
            rdm = rdm.at[p, q].add(1j * val)
            if p != q:
                rdm = rdm.at[q, p].add(-1j * val)

    return rdm

def term_energy(params):
    return jnp.real(circuit(params))


def term_rdm1(params, rdm1_prev):
    rdm1 = compute_rdm1_active(params)
    diff = rdm1 - rdm1_prev
    return jnp.sqrt(jnp.mean(jnp.abs(diff) ** 2))


grad_energy_fn = jax.jit(jax.grad(term_energy))
grad_rdm1_fn = jax.jit(lambda p, r_prev: jax.grad(term_rdm1, argnums=0)(p, r_prev))

# Iterations
total_iterations = 50000

# Threshold
Etol = 1e-6  # for energy
RDM1_tol = 1e-6  # for rmd1

# Initial value of variational parameters
params = jnp.zeros(param_shape)

# Classical optimizer
opt = optax.sgd(0.4)

# Keep track of the energies
energies = [circuit(params)]

# Set parameters
opt_state = opt.init(params)
current_opt = opt
rdm1_active_prev = compute_rdm1_active(jnp.zeros(param_shape)).copy()
rdm1_active = compute_rdm1_active(params)
rate = 0.6

f = open("log.txt", "a", encoding="utf-8")

for n in range(1, total_iterations + 1):
    # Gradient of the energy
    grad_E = grad_energy_fn(params)

    # Gradient of the RDM1 -> Only apply the RDM1 penalty from the second iteration
    # If it's the first iteration, grad_R is a zero vector
    if n > 1:
        grad_R = grad_rdm1_fn(params, rdm1_active_prev)
    else:
        grad_R = jnp.zeros_like(params)

    # Calculate the norm of each gradient
    norm_grad_E = jnp.linalg.norm(grad_E)
    norm_grad_R = jnp.linalg.norm(grad_R)

    # Define automatic weight for RDM1.
    # The goal is to balance the magnitude of grad_R with grad_E
    if norm_grad_R == 0:
        w_rdm1_auto = (norm_grad_E) * rate
    else:
        w_rdm1_auto = (norm_grad_E / (norm_grad_R)) * rate

    # Combine gradients
    grads = grad_E + w_rdm1_auto * grad_R

    # Update circuit parameters
    updates, opt_state = current_opt.update(grads, opt_state)
    params = optax.apply_updates(params, updates)

    # Compute updated energy and RDM1
    energy = circuit(params)
    rdm1_trial = compute_rdm1_active(params)  # new RDM1

    # Update references
    energies.append(energy)
    rdm1_active_prev = rdm1_active.copy()
    rdm1_active = rdm1_trial.copy()

    # Difference in RDM1
    diff = rdm1_trial - rdm1_active_prev
    rdm1_diff = jnp.sqrt(jnp.mean(diff**2))

    grad_norm = jnp.linalg.norm(grads)
    print(f"Iter {n} | Energy: {energy:.9f} Ha | ΔRDM1: {rdm1_diff:.6e} | ‖∇C‖: {grad_norm:.3e}")
    print(
        f"Iter {n} | Energy: {energy:.9f} Ha | ΔRDM1: {rdm1_diff:.6e} | ‖∇C‖: {grad_norm:.3e}",
        file=f,
    )

    # Convergence criterion
    if len(energies) > 1:
        delta_E = jnp.abs(energies[-1] - energies[-2])
        # If both energy and RDM1 converge, stop the loop
        if delta_E < Etol and rdm1_diff < RDM1_tol:
            print("Complete convergence in energy and RDM1.")
            break
f.close()

print(f"Iter {n} | Energy: {energy:.9f} Ha | ΔRDM1: {rdm1_diff:.6e} | ‖∇C‖: {grad_norm:.3e}")