import numpy as np
import numpy.random as random
import matplotlib.pyplot as plt
%matplotlib inline

dim = 3

"""
Explanation: Diffusion Monte Carlo Method
The goal of the diffusion Monte Carlo method (DMC) is to get to the ground state wave function of a quantum system. Using the eigenstates $\psi_n$ of a system, every state can be represented as a linear combination of these eigenstates.
$$\Psi(x,0) = \sum_n c_n \psi_n(x)$$
Using the time propagation operator $U(t, t_0) = e^{-i H (t-t_0)}$, the solution for the real time Schrödinger equation reads
$$\Psi(x,t) = U(t,0)\Psi(x,0) = \sum_n c_n e^{-i E_n t / \hbar } \psi_n (x).$$
When this equation is continued into imaginary time via $t \rightarrow -i \tau$, the resulting equation
$$\Psi(x,\tau) = \sum_n c_n e^{- E_n \tau / \hbar } \psi_n (x)$$
is an exponentially damped equation with the eigenenergies $E_n$ as specific damping coefficients for the separate eigenstates. Taking the limit $\tau \to \infty$ in the damping equation leads to every state's amplitude being damped, but the ground state being damped the least
$$\lim_{\tau \to \infty} U(\tau, 0) \Psi(x,0) = \lim_{\tau\to\infty} \sum_n c_n e^{-E_n \tau / \hbar} \psi_n(x, \tau) = c_0\psi_0$$ 
One method to numerically reach the ground state, the time propagation operator in imaginary time step width $\Delta \tau$ is repeatedly applied to a random initial state and the state normalized thereafter until we are left with the ground state.
$$c_0\psi_0 = \prod_n U(n\Delta \tau, (n-1)\Delta\tau) \Psi(x,0)$$
A different approach, presented in the following, is to model the time evolution as a diffusion process using an ensemble of diffusive particles, reaching the ground state through increased damping of higher-energy states as well. 

Problem with imaginary time propagation: solving the Schrödinger equation is not stable
Use diffusion instead, therefore...

Modeling Imaginary Time Propagation as Diffusion
The diffusion Monte Carlo Method uses a restated form of the Schrödinger equation as a diffusion equation to calculate the ground state wave function via a Monte Carlo process. The Schrödinger equation for a free particle in one dimension is
$$i\hbar \frac{\partial \psi (x,t)}{\partial t} = - \frac{\hbar ^2}{2m} \frac{\partial ^2 \psi (x,t)}{\partial x^2}$$
and can be rewritten as a diffusion equation 
$$\frac{\partial \psi (x,t)}{\partial t} = \frac{i \hbar ^2}{2m} \frac{\partial ^2 \psi (x,t)}{\partial x^2} = \gamma_{im} \frac{\partial ^2 \psi (x,t)}{\partial x^2},$$
where $\gamma_{im}$ is the imaginary diffusion constant 
$$\gamma_{im} = \frac{i \hbar ^2}{2m}.$$
In order to model the diffusion process and exploit the simulation possibilities, the imaginary diffusion constant has to be turned into a rel constant. Therefore, the operator in real time is analytically continued in imaginary time via
$$t \rightarrow i \tau,$$
leading to a diffusion equation in real space and imaginary time
$$\frac{\partial \psi (x,\tau)}{\partial \tau} = \frac{\hbar}{2m} \frac{\partial ^2 \psi (x,\tau)}{\partial x^2}.$$
We can therefore model the motion of a quantum particle in real space by simulation the diffusion of a cloud of particles in imaginary time. //TODO: explain
Diffusion In a Potential
Implementing Diffusion Monte Carlo
Using an ensemble of random walkers, the diffusion Monte Carlo algorithm is implemented. Here, the motion of the particles based on the diffusion equation models the kinetic energy part of the diffusion equation. The potential part is modeled by adding or removing walkers from the ensemble.
In the following, a diffusion Monte Carlo simulation for the 3-dimensional harmonic oscillator is developed, which leads to the ground state energy $E_0$ and ground state wave function $\psi_0$. Since this problem, can be solved analytically, the results can be checked against the analytical results
$$E_0 = \frac{2}{3}, \hspace{2em} \psi_0 = \frac{e^{-r^2/2}}{(2\pi)^{3/2}}$$
using $m=\omega=\hbar=1.$
End of explanation
"""


def pot(x):
    return 0.5 * x.dot(x)

"""
Explanation: Potential of the harmonic oscillator in three dimensions
$$V(\mathbf x) = \frac{1}{2}\mathbf x ^2$$
End of explanation
"""


dt = 0.05
N_T = 300
T_steps = 4000
alpha = 0.2

"""
Explanation: Set the time step size, the target number of walkers as well as the target number of time steps
End of explanation
"""


def diffusion(r, dt):
    return r + random.normal(size=(r.shape[0],3)) * np.sqrt(dt)

"""
Explanation: The first 20% of the time steps are used as thermalization steps, after that, the measured values are reset and the real simulation continues
In a timestep
1. Compute one DMC step on each walker
2. Remove dead walkers
3. Adjust $E_T$ to drive $N$ towards $N_T$
4. Accumulate data to measure $\left<E\right>$, its variance and the ground state wave function
In the diffusion step, each walker is moved randomly in the search space, in the physical picture due to the kinetic energy term. The random step is choosen from a Gaussian distribution with variance $\Delta t$.
End of explanation
"""


def branching(r, E_T, dt):
    
    r_new = np.zeros((0,3))
    q = np.zeros((r.shape[0],))
    
    for j in range(r.shape[0]):
        
        # branching factor
        q[j] = np.exp(-dt * (pot(r[j,:]) - E_T))
        
        # branching
        if q[j] - int(q[j]) > random.uniform():
            count_new = int(q[j]) + 1
        else:
            count_new = int(q[j])
            
        # generate new walker array
        for c in range(count_new):
            r_new = np.append(r_new,r[j:j+1,:],axis=0)
            
    return r_new, r_new.shape[0]

"""
Explanation: In the branching step, the number of walkers is modified by the parameter $E_T$, in the physical picture due to the potential energy term. The branching is implemented by computing the branching factor $q$ via
$$q=e^{-\Delta\tau (V(x) - E_T)}$$
which then determines if a walker is cloned, survives or dies. If $q<1$, the walker dies with a probability of $1-q$. If $q > 1$, the walker is copied.
End of explanation
"""


def adjust(N_T, N):
    return alpha * np.log(N_T / float(N))

"""
Explanation: In the adjustment step, the value of $E_T$ is adjusted to drive the number of walkers towards the desired number $N_T$. If $N > N_T$, the number of walkers is too high. We therefore increase $E_T$, reducing $q$ and the number of walkers. If, on the other hand, $N < N_T$, walkers need to be created. Then, $E_T$ is decreased, increasing $q$ as well as the number of walkers. The factor $\alpha$ controls the rate of change and has to be adjusted according to the simulation settings.
$$E_T \rightarrow E_T + \alpha \ln \left( \frac{N_T}{N} \right)$$
End of explanation
"""


class Accumulator(object):
    def __init__(self):
        self.E_sum = 0
        self.E_squared_sum = 0
        self.r_max = 4.0
        self.N_psi = 100
        self.psi = np.zeros((self.N_psi,))
        
    def reset(self):
        self.E_sum = 0
        self.E_squared_sum = 0
        self.psi = np.zeros((self.N_psi,))
        
    def handle_data(self, E_T, r):
        self.E_sum += E_T
        self.E_squared_sum += E_T**2
    
        for j in range(r.shape[0]):
            r_squared = r[j,:].dot(r[j,:])
            i_bin = int(np.sqrt(r_squared) / self.r_max * self.N_psi)
            if i_bin < self.N_psi:
                self.psi[i_bin] += 1

# Initialize walkers
r = np.zeros((N_T, 3))
N = N_T
E_T = 0 # initial guess for the ground state energy

thermal_steps = int(0.2*T_steps)

acc = Accumulator()

# Time step
for i in range(thermal_steps+T_steps):
    # Diffusion step
    r = diffusion(r, dt)
    
    # Branching step
    r, N = branching(r, E_T, dt)
    
    # Adjustment
    E_T += adjust(N_T, N)
    
    # Accumulation
    if i == thermal_steps:
        acc.reset()   
    acc.handle_data(E_T, r)

E_avg = acc.E_sum/T_steps
E_var = acc.E_squared_sum/T_steps - E_avg**2

print E_avg
print E_var
print acc.psi

plt.plot(acc.psi)
plt.show()

"""
Explanation: During each step of the production phase, the data is accumulated to evaluate the energy, energy variance and wave function $\psi$.
End of explanation
"""


