%------------------------------------------------------------
% Tony Hyun Kim
% 8.06 Term Paper on Hartree-Fock Approximation
%------------------------------------------------------------
% hartree_iter.m
% Illustration of the Hartree iteration of the helium atom
% (Figure 2 in the main paper)
%------------------------------------------------------------

subplot(121);

%------------------------------------------------------------
% Assignment of electron 1 wavefunction
%------------------------------------------------------------
r      = linspace(0,5);
V_nucl = -2./r;
plot(r,V_nucl);
ylim([-10 3]);
psi_a0 = sqrt(8/pi)*exp(-2*r);
hold on;
plot(r,psi_a0,'r');
grid on;
xlabel('Nondimensional radial distance, r');
ylabel('V_{nucl} [blue] and \psi_a^(0) [red]');
title('(a) Initial wavefunction assignment to electron 1');

%------------------------------------------------------------
% Computation of V_ee(r_2)
%------------------------------------------------------------
% Notes:
%   - Utilize spherical symmetry of psi_a0 by fixing r_2 to
%     remain on the x-axis
%   - |psi_a0|^2 at r = 3 is 1.5646E-5; hence we approximate 
%     the integral over all space by integration over a box
%     [-3,3]^3
%------------------------------------------------------------
N = 20; % Controls the degree of precision in the volume
        % integral (and also the computation time!)
x = linspace(-3,3,N);
y = linspace(-3,3,N);
z = linspace(-3,3,N);

dx = x(2)-x(1);
dy = y(2)-y(1);
dz = z(2)-z(1);
dV = dx*dy*dz;

V_ee = zeros(1,length(r));

% Loop over r2
for n = 1:100
    r2    = r(n);
    integ = 0;
    for i = 1:N           % Loop x
        for j = 1:N       % Loop y
            for k = 1:N   % Loop z
                r1       = sqrt(x(i)^2+y(j)^2+z(k)^2);
                psi_a_sq = 8/pi*exp(-4*r1);
                r12      = sqrt((x(i)-r2)^2+y(j)^2+z(k)^2);
                integ    = integ + dV*1/r12*psi_a_sq;
            end
        end
    end
    V_ee(n) = integ;
end

% Plot results
subplot(122);
plot(r,V_ee);
hold on;
plot(r,V_nucl);
ylim([-10 3]);
V_2 = V_ee + V_nucl;
plot(r,V_2,'LineWidth',1.5);
grid on;

%------------------------------------------------------------
% Assignment of electron 2 wavefunction
%------------------------------------------------------------
% Notes:
%   - We use an auxiliary Matlab script to perform the
%     integration of the Schroedinger equation, so that
%     this script does not become too complicated.
%     See "elec2_se.m"
%------------------------------------------------------------
psi_b = elec2_se(r,V_2);
plot(r,psi_b,'r');
xlabel('Nondimensional radial distance, r');
ylabel('V_{eff} = V_{nucl} + V_{ee} [thick blue] and \psi_b^(0) [red]');
title('(b) Corresponding wavefunction assignment to electron 2');