human
I am working on the file "Diffusion Monte Carlo.ipynb".
The first task is:
During each step of the production phase, the data is accumulated to evaluate the energy, energy variance and wave function $\psi$.
Can you write Python code for it?
assistant
#Here is the Python code you requested.

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