# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Sampling the potential energy surface 

# ## Introduction 
#
# This interactive notebook demonstrates how to utilize the Potential Energy Surface (PES) samplers algorithm of Qiskit Nature to generate the dissociation profile of a molecule. We use the Born-Oppenhemier Potential Energy Surface (BOPES)and demonstrate how to exploit bootstrapping and extrapolation to reduce the total number of function evaluations in computing the PES using the Variational Quantum Eigensolver (VQE). 
#

# +
# import common packages
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from functools import partial

# qiskit
from qiskit.utils import QuantumInstance
from qiskit import Aer
from qiskit.algorithms.minimum_eigen_solvers import NumPyMinimumEigensolver, VQE
from qiskit.algorithms.optimizers import SLSQP
from qiskit.circuit.library import ExcitationPreserving
from qiskit import BasicAer
from qiskit.algorithms import NumPyMinimumEigensolver, VQE
from qiskit.algorithms.optimizers import SLSQP

# qiskit nature imports
from qiskit_nature.problems.second_quantization import ElectronicStructureProblem
from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit_nature.mappers.second_quantization import JordanWignerMapper
from qiskit_nature.algorithms import GroundStateEigensolver
from qiskit_nature.drivers import PySCFDriver, UnitsType, Molecule
from qiskit_nature.algorithms.pes_samplers import BOPESSampler, Extrapolator

import warnings
warnings.simplefilter('ignore', np.RankWarning)
# -

# Here, we use the H2 molecule as a model system for testing.

# +
stretch1 = partial(Molecule.absolute_stretching, atom_pair=(1, 0))
mol = Molecule(geometry=[('H', [0., 0., 0.]),
                        ('H', [0., 0., 0.3])],
                       degrees_of_freedom=[stretch1],
                       )

# pass molecule to PSYCF driver
driver = PySCFDriver(molecule=mol)
# -

print(mol.geometry)

# ### Make a perturbation to the molecule along the absolute_stretching dof

# +
mol.perturbations = [0.2]
print(mol.geometry)

mol.perturbations = [0.6]
print(mol.geometry)
# -

# # Calculate bond dissociation profile using BOPES Sampler
#
# Here, we pass the molecular information and the VQE to a built-in type called the BOPES Sampler. The BOPES Sampler allows the computation of the potential energy surface for a specified set of degrees of freedom/points of interest.

# ## First we compare no bootstrapping vs bootstrapping 
#
# Bootstrapping the BOPES sampler involves utilizing the optimal variational parameters for a given degree of freedom, say r (ex. interatomic distance) as the initial point for VQE at a later degree of freedom, say r + $\epsilon$. By default, if boostrapping is set to True, all previous optimal parameters are used as initial points for the next runs.

# +
es_problem = ElectronicStructureProblem(driver)

qubit_converter = QubitConverter(JordanWignerMapper())
quantum_instance = QuantumInstance(backend=Aer.get_backend('statevector_simulator'))
solver = VQE(quantum_instance=quantum_instance)

me_gsc = GroundStateEigensolver(qubit_converter, solver)

# +
distance1 = partial(Molecule.absolute_distance, atom_pair=(1, 0))
mol = Molecule(geometry=[('H', [0., 0., 0.]),
                        ('H', [0., 0., 0.3])],
                       degrees_of_freedom=[distance1],
                       )

# pass molecule to PSYCF driver
driver = PySCFDriver(molecule=mol)

es_problem = ElectronicStructureProblem(driver)


# Specify degree of freedom (points of interest)
points = np.linspace(0.25, 2, 30)
results_full = {} # full dictionary of results for each condition
results = {} # dictionary of (point,energy) results for each condition
conditions = {False: 'no bootstrapping', True: 'bootstrapping'}


for value, bootstrap in conditions.items():
    # define instance to sampler
    bs = BOPESSampler(
        gss=me_gsc
        ,bootstrap=value
        ,num_bootstrap=None
        ,extrapolator=None)
    # execute
    res = bs.sample(es_problem,points)
    results_full[f'{bootstrap}'] =  res.raw_results
    results[f'points_{bootstrap}'] = res.points
    results[f'energies_{bootstrap}'] = res.energies
# -

# # Compare to classical eigensolver
#

# define numpy solver
solver_numpy = NumPyMinimumEigensolver()
me_gsc_numpy = GroundStateEigensolver(qubit_converter, solver_numpy)
bs_classical = BOPESSampler(
               gss=me_gsc_numpy
               ,bootstrap=False
               ,num_bootstrap=None
               ,extrapolator=None)
# execute
res_np = bs_classical.sample(es_problem, points)
results_full['np'] =  res_np.raw_results
results['points_np'] = res_np.points
results['energies_np'] = res_np.energies

# ## Plot results

fig = plt.figure()
for value, bootstrap in conditions.items():
    plt.plot(results[f'points_{bootstrap}'], results[f'energies_{bootstrap}'], label = f'{bootstrap}')
plt.plot(results['points_np'], results['energies_np'], label = 'numpy')
plt.legend()
plt.title('Dissociation profile')
plt.xlabel('Interatomic distance')
plt.ylabel('Energy')

# # Compare number of evaluations

for condition, result_full in results_full.items():
    print(condition)
    print('Total evaluations for ' + condition + ':')
    sum = 0
    for key in result_full:
        if condition is not 'np':
            sum += result_full[key].raw_result.cost_function_evals
        else:
            sum = 0
    print(sum)

# # Extrapolation 
#
# Here, an extrapolator added that will try to fit each (param,point) set to some specified function and suggest an initial parameter set for the next point (degree of freedom). 
#
# - Extrapolator is based on an external extrapolator that sets the 'window', that is, the number of previous points to use for extrapolation, while the internal extrapolator proceeds with the actual extrapolation.
# - In practice, the user sets the window by specifying an integer value to num_bootstrap - which is also the number of previous points to use for bootstrapping. Additionally, the external extrapolator defines the space within how to extrapolate - here PCA, clustering and the standard window approach. 
#
# In practice, if no extrapolator is defined and bootstrapping is True, then all previous points will be bootstrapped. If an extrapolator list is defined and no points are specified for bootstrapping, then the extrapolation will be done based on all previous points.

# 1. Window Extrapolator: An extrapolator which wraps another extrapolator, limiting the internal extrapolator's ground truth parameter set to a fixed window size
# 2. PCA Extrapolator: A wrapper extrapolator which reduces the points' dimensionality with PCA, performs extrapolation in the transformed pca space, and untransforms the results before returning.
# 3. Sieve Extrapolator: A wrapper extrapolator which performs an extrapolation, then clusters the extrapolated parameter values into two large and small clusters, and sets the small clusters' parameters to zero.
# 4. Polynomial Extrapolator: An extrapolator based on fitting each parameter to a polynomial function of a user-specified degree.
# 5. Differential Extrapolator: An extrapolator based on treating each param set as a point in space, and performing regression to predict the param set for the next point. A user-specified degree also adds derivatives to the values in the point vectors which serve as features in the training data for the linear regression.

# ## Here we test two different extrapolation techniques

# +
# define different extrapolators
degree = 1
extrap_poly = Extrapolator.factory("poly", degree = degree)
extrap_diff = Extrapolator.factory("diff_model", degree = degree)
extrapolators = {'poly': extrap_poly, 'diff_model': extrap_diff}

for key in extrapolators:
    extrap_internal = extrapolators[key]
    extrap = Extrapolator.factory("window", extrapolator = extrap_internal)
    # define extrapolator
    # BOPES sampler
    bs_extr = BOPESSampler(
        gss=me_gsc
        ,bootstrap=True
        ,num_bootstrap=None
        ,extrapolator=extrap)
    # execute
    res = bs_extr.sample(es_problem, points)
    
    results_full[f'{key}']= res.raw_results
    results[f'points_{key}'] = res.points
    results[f'energies_{key}'] = res.energies

# -

# ## Plot results

fig = plt.figure()
for value, bootstrap in conditions.items():
    plt.plot(results[f'points_{bootstrap}'], results[f'energies_{bootstrap}'], label = f'{bootstrap}')
plt.plot(results['points_np'], results['energies_np'], label = 'numpy')
for condition in extrapolators.keys():
    print(condition)
    plt.plot(results[f'points_{condition}'], results[f'energies_{condition}'], label = condition)
plt.legend()
plt.title('Dissociation profile')
plt.xlabel('Interatomic distance')
plt.ylabel('Energy')

# # Compare number of evaluations

for condition, result_full in results_full.items():
    print(condition)
    print('Total evaluations for ' + condition + ':')
    sum = 0
    for key in results_full[condition].keys():
        if condition is not 'np':
                sum += result_full[key].raw_result.cost_function_evals
        else:
                sum = 0
    print(sum)

import qiskit.tools.jupyter
# %qiskit_version_table
# %qiskit_copyright

StarcoderdataJupyterScriptsDedupFiltered