Variational Monte Carlo

Variational Monte Carlo#

import numpy as np 
import pylab as plt
import qutip
import scipy
import scipy.linalg
import stats
from tqdm.notebook import tqdm
import time
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import matplotlib.animation as animation
from typing import NamedTuple
from scipy.integrate import tplquad
from numba import njit
def resetMe(keepList=[]):
    ll=%who_ls
    keepList=keepList+['FuncAnimation','HTML','resetMe','scipy','np','plt','math','jax','jnp','jit','grad','HTML','animation','qutip','animation','FuncAnimation','NamedTuple']
    for iiii in keepList:
        if iiii in ll:
            ll.remove(iiii)
    for iiii in ll:
        jjjj="^"+iiii+"$"
        %reset_selective -f {jjjj}  
    ll=%who_ls
    plt.rcParams.update({"font.size": 14})
    return
#resetMe()
import datetime;datetime.datetime.now()
from functools import partial
#import jax as jnp

Our goal in this assignment is to understand molecules and binding.

Exercise 1. Dihydrogen cation#

We are going to start by looking at the Hydrogen H2 molecule but with only one electron - i.e. the dihydrogen cation. In exercise 3, we will work our way up to standard H2 with two electrons.

a. A variational wave-function: attempt 1#

In this assignment, we will allow our molecule be centered at the origin with protons at \(\pm L/2 \hat{x}\). We will start by working with \(L=2\) Bohr radii.

We now must decide on a variational wave-function. A reasonable first guess is to assume the electron is spending time in between the two protons - i.e. centered at zero with some spread around it.

We will begin by using the wave-function

\[ \Psi(r_1;\alpha) = \left(\frac{2\alpha}{\pi}\right)^{3/4} \exp[-\alpha(|r_1|^2 )] \]

where \(\alpha\) is a variational parameter. To store this parameter, we will define a named tuple

class WavefunctionParams(NamedTuple):
  """Wavefunction Parameters"""
  alpha: float = 0.5
  beta: float = 0.0
params=WavefunctionParams(alpha=0.5)

which can then be accessed by params.alpha. (beta will be used later - just ignore it for now)

Now write a function Psi(e_pos,p_pos,params) which takes the electron position, proton positions, and wave-function parameters and returns the wave-function amplitude. Note that our current wave-function just ignores the proton positions.

Because the hydrogen atoms are lined up on the x-axis, plot the wave-function as a function of \(x\) - i.e. xs=np.linspace(-3.1*L/2,3.1*L/2).

###ANSWER HERE

b. Evaluating the energy#

To use the variational principle, we should compute the energy of the wave-function:

\[ E = \frac{\int \Psi^*(R) [H\Psi](R) dx dy dz }{\int \Psi^*(R) \Psi(R) dR } \]

where \(R=(x,y,z)\) which we will rewrite as

\[ E = \frac{\int |\Psi(R)|^2 \frac{[H\Psi](R)}{\Psi(R)} dx }{\int |\Psi(R)|^2 dx} \equiv \frac{\int |\Psi(R)|^2 E_L(R)}{\int |\Psi(x)|^2 dx} \]

where the local energy

\[ E_L(R) = \frac{[H\Psi](R)}{\Psi(R)} = \frac{[-\nabla^2 \Psi /2+ V \Psi](R) }{\Psi(R)} \]

where we are using the Hamiltonian

\[ H = \sum_i -\frac{1}{2} \nabla_i^2 + V(R) \]

Let us start by writing a function Laplacian(e_pos,p_pos,params,Psi ,delta=1e-5) which takes the wave-function Psi and gives back the laplacian

\[ \nabla^2 \Psi = \frac{\partial^2}{\partial x^2}+ \frac{\partial^2}{\partial y^2}+ \frac{\partial^2}{\partial z^2} \]

You can do this using the standard finite difference stencil for the Laplacian

\[ \frac{\partial^2 F}{\partial x^2} \approx \frac{F(x+\delta) +F(x-\delta) -2F(x)}{\delta^2} \]

We wish to evaluate \(\nabla^2 \Psi / \Psi\) which we can do by evaluating

L=2.0
p_pos=np.array([[-L/2,0,0],[L/2,0,0]])
params=WavefunctionParams(alpha=0.5)
e_pos=np.array([[1.2,0.3,-0.4]])
print(Laplacian(e_pos, p_pos, params, Psi)/Psi(e_pos, p_pos, params))

and should get -1.3099993792655844

###ANSWER HERE

The next step is to get the potential piece of the energy

\[ \frac{V\Psi}{\Psi} = V(R) = -\sum_{Ie} \frac{1}{r_{Ie}} + \sum_{ee} \frac{1}{r_{ee}} + \sum_{II} \frac{1}{r_{II}} \]

where \(I\) is the ion location and \(e\) is the electron location. \(r_{ab}\) stands for the distance \(|r_a-r_b|\)

For the dihydrogen cation there is no electron-electron term and so we can ignore that.

Write a function def VPsiOverPsi(e_pos,p_pos) which returns this energy. Again, you can test this by evaluating

p_pos=np.array([[-L/2,0,0],[L/2,0,0]])
params=WavefunctionParams(alpha=0.5)
e_pos=np.array([[1.2,0.3,-0.4]])
print(VPsiOverPsi(e_pos, p_pos))

and you should get -1.8001955889484549

###ANSWER HERE

Finally, put it all together into a single function LocalEnergy(e_pos, p_pos, params,Psi) which computes the local energy. Make sure you remember the \(-1/2\) in front of the Laplacian term.

###ANSWER HERE

c. Computing the energy#

Now that we have the wave-function Psi and the local energy LocalEnergy we should be able to compute the energy by doing the two integrals

\[ N = \int |\Psi(x,y,z)|^2 E_L(x,y,z) dx dy dz \]

and

\[ D = \int |\Psi(x,y,z)|^2 dx dy dz \]

and evaluating \(N/D.\)

These are both three-dimensional integrals and so it’s possible to still compute them reasonably using a deterministic approach on a grid. Later on in this assignment, we will have to switch to a Monte Carlo approach. To compute an integral, we need to (for each integral) to use the integrate function below and implement an appropriate integrand function (using the spherical coordinates Jacobian that we’ve started you with)

def integrand(r, theta, phi, p_pos, params):
    """
    Calculate the integrand for the expectation value in spherical coordinates.
    """
    x = r * np.sin(theta) * np.cos(phi)
    y = r * np.sin(theta) * np.sin(phi)
    z = r * np.cos(theta)
    e_pos = np.array([[x, y, z]])
    jacobian = r**2 * np.sin(theta)
    ### compute the other pieces of the integrand 
    return jacobian * rest of integrand. 

def integrate_spherical(p_pos,params,limit_r,integrand):
    integration_result, _ = tplquad(
                integrand,
                0, 2*np.pi,                # phi
                lambda phi: 0, lambda phi: np.pi,  # theta
                lambda phi, theta: 0, lambda phi, theta: limit_r , # r
                args=(p_pos, params,),epsabs=1.49e-02, epsrel=1.49e-02
    )
    return integration_result

alpha=0.5
L = 2.0
p_pos = np.array([[-L/2, 0, 0], [L/2, 0, 0]])
limit_r=10.0
params=WavefunctionParams(alpha=alpha)
N=integrate_spherical(p_pos,params,limit_r,integrand_energy_spherical)
D=integrate_spherical(p_pos,params,limit_r,integrand_norm_spherical)
print(N/D)

Using \(\alpha=0.5\) and a sphere of limit_r=10 compute the energy. I get -0.43539807117579543 for \(N/D\). Notice that if you have the normalization correct you should be getting \(D=1\).

###ANSWER HERE

d. Optimizing the energy#

Now, let’s do this again looping over all alpha=np.linspace(0.1,0.5,20). We are looking for the variational parameter that makes this wave-function the smallest. Plot energy vs. \(\alpha\) . Fit a parobola to the function somewhere near the minima (poly=np.polyfit(x,y,2))and predict what the best \(\alpha\) is. Compute the energy again with this \(\alpha\) to get the lowest energy giving the variational space of this wave-function.

###ANSWER HERE

###ANSWER HERE

Finally plot out the wave-function on the x-axis (in the direction of the two protons) for the optimal \(\alpha\).

###ANSWER HERE

Exercise 2. Variational Monte Carlo#

a. Variational Monte Carlo#

So far we’ve evaluated our wave-function using quadrature. This works “ok” when you have one electron because you can afford to make a three-dimensional grid. Eventually, though, we will need to transition to Monte Carlo to evaluate our wave-functions; for example, to do the Hydrogen molecule with two electrons (as it would typically have), then we would need to do a six-dimensional integration.

In this exercise we will still work with the dihydrogen cation but switch over to variational Monte Carlo and make sure that we get the same energy.

The steps of variational Monte Carlo:

For many sweeps:
- For 10 steps per sweep:
  - Evaluate the old wavefunction amplitude wf_old=Psi(e_pos,p_pos,params)
  - Change the location of the electron by a random gaussian np.random.randn(3)*1.0
  - Evaluate the new wavefunction amplitude wf_new=Psi(e_pos,p_pos,params) at this new position
  - If (wf_new/wf_old)^2 is greater then a random number selected between 0 and 1 (np.random.rand()) accept and keep the new positions. Otherwise revert to the old positions.
- Add the current local energy to an energy list and the current location to a location list.

Once this is done you can average the local energy list (ignoring the first few entries where the Monte Carlo was still equilibrating). Write a function MonteCarlo(p_pos,params,Psi,sweeps) which performs this Variational Monte Carlo and returns the energies, acceptance ratio of your moves, and locations.

Call your MonteCarlo function for our wave-function at \(\alpha=0.5\) using 10,000 sweeps. You can compute an error bar by doing error_energy=np.std(energies[l:])/np.sqrt(len(energies[l:])) where l=100 is a number of sweeps to ignore as the transient. Check that you are within 1-2 error bars of the value you computed above in exercise 1.

Also plot the energy as a function of sweeps and plot a histogram of the energies. You should notice that it is very spikey and the histogram has a long tail to the left.

###ANSWER HERE

b. Energy as a function of \(\alpha\)#

Now go ahead and produce the energy as a function of \(\alpha\) for alphas=np.linspace(0.01,0.5,40). Check that your data matches the data from above. Run for 20,000 sweeps for each point and only average your energy after ignoring the first 100 points (e.g. np.mean(energies[100:]) )

###ANSWER HERE

###ANSWER HERE

c. A problematic wave-function#

We would like to make our wave-function better. One property of the exact wave-function is that the local energy \(E_L(x,y,z) \equiv [H(\Psi_0)](x,y,z)/\Psi_0(x,y,z)\) is the same (a constant) no matter what coordinate we are evaluating it at. But we saw that our variational wave-function is very far from this limit. In fact, the spikiness and broad histogram of the local energy tells us that there are lots of configurations which are very different. We’d like to diagnose what is causing this extreme spikiness so we can work on making it better.

To do this, let us look at the local energy explicitly. Plot the local energy, the kinetic piece of the local energy and the potential piece of the local energy as a function of xs=np.linspace(-2,2,1000) all on the same plot . Pay special attention to where the electron and protons are at the same position.

###ANSWER HERE

d. A Jastrow#

What you should have seen in the last part is that the local energy gets very negative as the electron approaches the protons. This is happening because the potential piece is getting negative and the kinetic piece is not cancelling it out. To fix this, we can enforce a cusp condition on our wave-function by adding a Jastrow factor writing a new wave-function as

\[ \Psi_{new}(r;\alpha,\beta)=\Psi_{old}(r;\alpha)\exp[J_{ep}(r;\beta)] \]

where

\[ J_{ep}=-\frac{a_{ep}|r_i-I_j|} {1+b_{ep}|r_i-I_j|} \]

where \(a_{ep}=1\) are chosen to satisfy the electron-proton cusp conditions, respectively.

The parameter \(\beta\) gives additional freedom through the relations,

\[ b_{ep} = \left( \frac{a_{ep}}{\beta}\right)^{1/2} \]

Implement this Jastrow function (say using \(\beta =0.1\)) and again plot the local energy (with the kinetic and potential pieces). What do you notice that is different?

###ANSWER HERE

e. Using our Jastrow#

Now we are going to go ahead and use our Jastrow factor. Again, compute the value of our wavefunction with parameters \(\alpha=0.5\) and \(\beta=0.1\) Plot the local energy as a function of Monte Carlo step and the histogram. What do you notice now that is different from the earlier plots without the Jastrow?

###ANSWER HERE

f. Getting the energy again#

Now we want to find the best energy again for this bond-length. This time we have to search over the variational parameters \(\alpha\) and \(\beta\). Using alphas=np.linspace(0.01,0.5,20) and betas=np.linspace(5.0,15.0,20) search over the relevant space. Store the results in a 2D array all_energies=np.zeros((20,20)). I am using 1000 sweeps per simulation and ignoring the first 100 sweeps to evaluate the energy (treating this as a transient).

Plot the following things:

(1) For every beta plot a line which is the energy versus alpha for that beta

I did something like

for i,beta in enumerate(betas):
    plt.plot(alphas,allEnergies[i,:],label='beta='+str(beta))
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
plt.xlabel('alpha')
plt.ylabel('Energy')    
plt.show()

(2) Plot the 2D energies:

I did this using

plt.imshow(allEnergies,origin='lower',aspect='auto',cmap='hot',extent=[alphas[0],alphas[-1],betas[0],betas[-1]])
plt.colorbar(label='Energy')
plt.xlabel('Alpha')
plt.ylabel('Beta')
plt.title('Energy vs Alpha and Beta')
plt.show()

(3) Finally take the lowest energy wave-function:

To get the parameters for the lowest energy, I did

min_beta, min_alpha = np.unravel_index(np.argmin(allEnergies), allEnergies.shape)
best_beta=betas[min_beta]
best_alpha=alphas[min_alpha]

report its energy (to do this you want to find the best \(\alpha\) and \(\beta\) and rerun it otherwise your energy will be biased)
plot it as a function of \(x\).

The exact answer (see wikipedia) is -0.597. You should find that you get something close.

###ANSWER HERE

###ANSWER HERE

###ANSWER HERE

g. Bonding#

All of this work was to get the energy of the proton at a single bond-length. Now we would like to stretch (and compress) the bond and also get the energies. If there is a local minima at some value (and the energy at large separation is larger then this local minima), then the molecule should bind.

Go ahead and compute for L=[1.0,2.0,3.0] the minima energy (optimizing for each one over all \(\alpha\) and \(\beta\)) You already have the \(L=2\) point. Plot the best energies as a function of L and see if you get a local minima that would indicate bonding.

You can compare it against the exact energy with

cation=np.loadtxt("data",delimiter=",")
plt.plot(cation[:,0],cation[:,1])

You may notice that as you begin to stretch the molecule, things start to fall apart. Why is that? What could you do to fix this?

###ANSWER HERE

Exercise 3. Hydrogen Molecule#

Now that we’ve made some good progress with one electron, let’s go ahead and modify things to work with the Hydrogen molecule and two electrons - a single spin-up and spin-down.

Things we have to change:

Add an additional electron-electron Jastrow term
Make sure your functions work with two electrons - pay special attention that your Variatonal Monte Carlo is choosing 6 (not 3) random numbers to add.

Everything else should be essentially the same.

a. Modify the Jastrow#

We now need a Jastrow factor for electrons as well:

\[ \Psi_{new}(r_1,r_2;\alpha,\beta)=\Psi_{old}(r_1,r_2;\alpha,\beta )\exp[J_{ee}(r_1,r_2;\beta)] \]

where

\[ J_{ee}=\frac{a_{ee}|r_1-r_2|} {1+b_{ee}|r_1-r_2|} \]

where \(a_{ee}=1/2\) ,

The parameter \(\beta\) gives additional freedom through the relation,

\[ b_{ee} = \left( \frac{a_{ee}}{\beta} \right)^{1/2} \]

In addition, make sure your other Jastrow works for two electrons as does your wavefunction - i.e. there should now be a term that goes as \(\exp[-\alpha (r_1^2 + r_2^2)\)]

You can check it by making sure that

L=2.0
p_pos=np.array([[-L/2,0,0],[L/2,0,0]])
e_pos=np.array([[0.3,-0.5,2.1],[1.2,-0.2,1.1]])
params=WavefunctionParams(alpha=0.5,beta=0.2)
print(Psi(e_pos,p_pos,params))

gives you 0.007040289115058886

###ANSWER HERE

b. Modify your local energy#

Make the modifications to the Laplacian and the potential so it correctly computes the actual energy of your system. Make sure you use two particles in your laplacian and your potential energy including the electron-electron interaction now.

You can check it by doing

L=2.0
p_pos=np.array([[-L/2,0,0],[L/2,0,0]])
e_pos=np.array([[0.3,-0.5,2.1],[1.2,-0.2,1.1]])
params=WavefunctionParams(alpha=0.5,beta=0.2)
print(Laplacian(e_pos, p_pos, params, Psi)/Psi(e_pos, p_pos, params))
print(VPsiOverPsi(e_pos, p_pos))
print(LocalEnergy(e_pos, p_pos, params, Psi))

and get

1.7978204371923225
-0.8976856497576701
-1.7965958683538314

###ANSWER HERE

c. Do a Monte Carlo#

Modify your Monte Carlo to make sure it’s moving two electrons. Run your Monte Carlo on

L=2.0
p_pos=np.array([[-L/2,0,0],[L/2,0,0]])
e_pos=np.array([[0.3,-0.5,2.1],[1.2,-0.2,1.1]])
params=WavefunctionParams(alpha=0.5,beta=0.2)

and make sure that you get something within error bars of

-0.59998768751766 +/- 0.01068398300057258 

###ANSWER HERE

d. Get the best \(\alpha, \beta\) at this bond length#

Using betas=np.linspace(0.1,10.0,20) and alphas=np.linspace(0.01,0.5,20) get the optimal \(\alpha, \beta\) as well as the lowest energy at

L=2.0
p_pos=np.array([[-L/2,0,0],[L/2,0,0]])

###ANSWER HERE

###ANSWER HERE

###ANSWER HERE

e. Bonding#

Now compute the energy for the lengths Ls=[1.0,1.424,2.0,3.0,4.0,5.0,6.0]. You can compare them with the exact answer

data_H2=np.loadtxt("data_H2",delimiter=",")
plt.plot(data_H2[:,0],data_H2[:,1])

You may notice that you are getting reasonable answers near the equilibrium bond length (what equilibrium bond length did you find?) but getting something that is not reasonable at larger bond stretching. Why is this and how could you modify your wave-function to do a better job.

###ANSWER HERE

###ANSWER HERE

Acknowledgements:

Bryan Clark (original)