Variational Monte Carlo

Variational Monte Carlo#

import numpy as np 
import pylab as plt
import qutip
import scipy
import scipy.linalg
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 bidning.

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

\[ \Psi(r_1;\alpha) = \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.3)

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(-2.1*L/2,2.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 stanadard 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} \]

You can test your Laplacian by doing

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))

and should get -0.23846646879377426

###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.

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$.

###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 local 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 \]

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 approprite 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(p_pos,params,limit_r,integrand):
    integration_result, _ = tplquad(
        integrand,
        0, limit_r,      # Limits for r
        0, np.pi,         # Limits for theta
        0, 2 * np.pi,    # Limits for ph, 
        args=(p_pos, params,),epsabs=1.49e-01, epsrel=1.49e-01
    )
    return integration_result

alpha=0.5
L = 2.0
p_pos = np.array([[-L/2, 0, 0], [L/2, 0, 0]])
limit_r=50.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 $L=50$ compute the energy. I get -0.4316921028050629.

###ANSWER HERE

d. Optimizing the energy#

Now, let’s do this again looping over all alpha=np.linspace(0.1,0.5,20). To have it run at some reasonable rate, turn down $r_limit=10$. 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 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).

###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()) keep the new positions. Otherwise revert to the old positions.
- Add the current local energy to a 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 and acceptance ratios.

Let’s try this for our wave-function at $\alpha=0.5$ and run for 10,000 sweeps. You can compute an error bar by doing error_energy=np.std(energies[l:])/np.sqrt(len(energies[l:])) where l is a number selected to ignore the first few sweeps. 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 wide range of values.

###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,20). You should get something similar to above but you should have more accurate energies (assuming you were using r_limit=10 above). There will be error bars on your results thought.

###ANSWER HERE

###ANSWER HERE

c. A problematic wave-function#

Although we are getting reasonable results with our wave-function there is something that is qualitatively problematic about it. You might notice that our local energy has a lot of very large spikes (and the histogram has some pretty extreme values). This is a sign that the wave-function isn’t very good. Let’s try to understand what’s going on focusing on $\alpha=0.5$.

Consider the pieces of our integrands

\[ \tilde{T}=\int r^2 |\Psi(r)|^2 \frac{\nabla^2 \Psi(r)}{\Psi(r) } \]

\[ \tilde{V}=\int r^2 |\Psi(r)|^2 V(r) \]

\[ \int r^2 |\Psi(r)|^2 E_L(r) = \tilde{T} + \tilde{V} \]

and plot them function of xs=np.linspace(-2,2,100) with $y=z=0$ all on the same plot.

In addition, plot just $$ \int r^2 |\Psi(r)|^2 $$

What you should notice is that the local energy diverges at the proton locations because the potential energy gets very negative and the kinetic energy doesn’t compensate. Both the large potential energy and the lack of compensating kinetic energy are coming because the probability of the electron being near the protons is non-trivial due to the flat wave-function near the proton locations.

###ANSWER HERE

d. A Jastrow#

In order to fix this broken wave-function we should add a term to it which keeps the electron and proton away from each other. In fact, we can analytically calculate what the correct term should be in the limit where the electron-proton distance is small.

We can add a Jastrow factor

\[ \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 and again plot the integrand terms above (including a third separate plot which just shows the local energy). What do you notice that’s 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 he 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 phase point. We now have to search over the variational parameters $\alpha$ and $\beta$. Using alphas=np.linspace(0.01,0.5,20) and betas=np.linspace(0.1,10.0,20) search over the relevant space. For every beta plot a line which should be energy versus alpha. Also, store the energies in such a way that you can report the minimum over all the energies that you see (this is a little bit incorrect to do because of statistical fluctuations - don’t do it exactly this way for a research project). I am using 1000 sweeps per simulation. The exact answer (see wikipedia) is -0.597. You should find that you get something close.

I am labelling each of my lines with plt.plot(alphas,energies[alpha],label='beta='+str(beta)) and then pushing the legend outside of the box area with plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

Finally take the final wave-function and plot it as a function of $x$

###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 these lengths (and the energy at large separation is larger then this local minima), then we would expect that 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 this as a function of L and see if you get a local minima that would suggest bonding. You will probably find it useful to encapsulate the previous exercise into a function.

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?

###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 relations,

\[ 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))
print(VPsiOverPsi(e_pos, p_pos))
print(LocalEnergy(e_pos, p_pos, params, Psi))

and get

0.012657175654795516
-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.34209397913287004  +/- 0.019897206346676632

###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 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

###ANSWER HERE

Acknowledgements:

Bryan Clark (original)