%=======================================================================
% Chapter 10, code 5: Schrodinger equation as a matrix eigenvalue problem.
% Dimensionless units: hbar = mass = 1.
%   -0.5*psi'' + V(x)*psi = E*psi, with psi(a) = psi(b) = 0.
% Method: centered finite differences and symmetric Jacobi rotations.
% Jacobi follows Chapter 8, with a stable small-angle rotation formula.
%
% Example: harmonic oscillator, V(x) = 0.5*x^2.
% Exact energies on the infinite line: E_j = j+0.5, j = 0,1,2,... .
% A finite interval approximates the infinite line for the lowest states.
% Check both grid spacing and interval size when testing accuracy.
%
% MATLAB alternatives: [vectors,D] = eig(H) gives all eigenpairs.
% For a large sparse Hamiltonian, eigs can compute selected eigenpairs.
% Sort the energies and their matching eigenvector columns afterward.
%
% To solve another problem:
%   1. Edit the interval, grid size, and number of states below.
%   2. Edit potential immediately below the driver.
%   3. Disable check_exact or replace the exact-energy comparison.
% Jacobi computes all eigenpairs of a dense matrix. Keep grids modest.
%=======================================================================
clear;

%-------------------- Problem and grid: edit here -----------------------
a = -6.0;
b = 6.0;
n = 120;                     % n intervals, n-1 interior unknowns
states = 3;
abserr = 1.0e-18;            % Bound on SUM of squared off-diagonal entries
max_sweeps = 100;
check_exact = true;           % Harmonic-oscillator energy comparison
make_plot = true;
%-----------------------------------------------------------------------
interior = n-1;
if n < 2 || n ~= floor(n) || states < 1 || states > interior || ...
        states ~= floor(states) || b <= a
    error('Use integer n >= 2, integer 1 <= states <= n-1, and b > a.');
end
h = (b-a)/n;
x = linspace(a,b,n+1)';
v = zeros(n+1,1);
for i = 1:n+1
    v(i) = potential(x(i));
end

% Only interior values are unknown. Zero endpoints contribute nothing
% to the first and last interior equations. MATLAB interior index i
% corresponds to the full-grid point x(i+1).
matrix = zeros(interior,interior);
for i = 1:interior
    matrix(i,i) = 1/(h*h)+v(i+1);
end
for i = 1:interior-1
    matrix(i,i+1) = -0.5/(h*h);
    matrix(i+1,i) = matrix(i,i+1);
end

fprintf('Schrodinger equation: finite difference matrix and Jacobi eigenvalues\n');
fprintf('Grid intervals: %8d\n',n);
fprintf('Matrix size:    %8d\n',interior);
fprintf('Grid spacing:   %12.6f\n',h);
[matrix,vectors,sweeps,rotations,offdiag_norm,info] = ...
    jacobi(matrix,abserr,max_sweeps);
if info ~= 0
    error('Jacobi failed: check controls or increase the sweep limit.');
end
energies = diag(matrix);
[energies,vectors] = sort_eigenpairs(energies,vectors);

% Eigenvectors have sum(vectors.^2) = 1. Wavefunctions need
% h*sum(psi.^2) = 1, also the trapezoidal integral with zero endpoints.
psi = zeros(n+1,states);
for j = 1:states
    norm_value = sqrt(h*sum(vectors(:,j).^2));
    psi(2:n,j) = vectors(:,j)/norm_value;
    % Fix the arbitrary overall sign: largest-magnitude component positive.
    [~,pivot] = max(abs(psi(:,j)));
    if psi(pivot,j) < 0
        psi(:,j) = -psi(:,j);
    end
end

fprintf('\n');
if check_exact
    fprintf('%7s %16s %16s %12s\n','state j','numerical E','exact E','abs error');
else
    fprintf('%7s %16s\n','state j','numerical E');
end
for j = 1:states
    if check_exact
        exact = j-1+0.5;
        fprintf('%7d %16.8f %16.8f %12.4e\n', ...
            j-1,energies(j),exact,abs(energies(j)-exact));
    else
        fprintf('%7d %16.8f\n',j-1,energies(j));
    end
end

% Check with the original tridiagonal operator, rather than the matrix
% overwritten by Jacobi. These checks concern matrix eigenpairs;
% they do not measure grid or finite-interval errors.
max_residual = 0;
for j = 1:states
    for i = 2:n
        residual = (1/(h*h)+v(i)-energies(j))*psi(i,j) ...
                   -0.5*(psi(i-1,j)+psi(i+1,j))/(h*h);
        max_residual = max(max_residual,abs(residual));
    end
end
max_orthogonality = 0;
for j = 1:states
    for k = 1:states
        overlap = h*sum(psi(:,j).*psi(:,k));
        target = double(j == k);
        max_orthogonality = max(max_orthogonality,abs(overlap-target));
    end
end
fprintf('Jacobi sweeps:            %8d\n',sweeps);
fprintf('Largest matrix residual:  %12.4e\n',max_residual);
fprintf('Largest overlap defect:   %12.4e\n',max_orthogonality);

filename = 'ch10_schrodinger_matrix.dat';
fid = fopen(filename,'w');
if fid == -1
    error('Cannot open the data file.');
end
fprintf(fid,'# x V(x) psi_0 psi_1 ...; integral(psi_j^2 dx) = 1\n');
for j = 1:states
    fprintf(fid,'# state %6d E = %.16e\n',j-1,energies(j));
end
fprintf('\nWavefunctions: columns are x, V(x), then the normalized states.\n');
for i = 1:n+1
    fprintf('%10.6f %16.8e',x(i),v(i));
    fprintf(' %16.8e',psi(i,:));
    fprintf('\n');
    fprintf(fid,' %.16e',[x(i),v(i),psi(i,:)]);
    fprintf(fid,'\n');
end
fclose(fid);
fprintf('Data file: %s\n',filename);

if make_plot
    figure;
    for j = 1:states
        subplot(states,2,2*j-1);
        plot(x,psi(:,j),'k-','LineWidth',1.2);
        xlabel('x');
        ylabel('\psi_j(x)');
        title(sprintf('State %d: E = %.8f',j-1,energies(j)));
        grid on;

        subplot(states,2,2*j);
        plot(x,psi(:,j).^2,'k-','LineWidth',1.2);
        xlabel('x');
        ylabel('|\psi_j(x)|^2');
        title(sprintf('Probability density, state %d',j-1));
        grid on;
    end
end


function v = potential(x)
% Supply V in -0.5*psi'' + V*psi = E*psi.
% Default: dimensionless harmonic oscillator.
% Change this expression for another potential and disable check_exact.
v = 0.5*x*x;
end


function [values,vectors] = sort_eigenpairs(values,vectors)
% Sort energies in ascending order, moving their eigenvector columns too.
n = length(values);
for i = 1:n-1
    smallest = i;
    for j = i+1:n
        if values(j) < values(smallest)
            smallest = j;
        end
    end
    if smallest == i
        continue;
    end
    value = values(i);
    values(i) = values(smallest);
    values(smallest) = value;
    work = vectors(:,i);
    vectors(:,i) = vectors(:,smallest);
    vectors(:,smallest) = work;
end
end


function [a,vectors,sweeps,rotations,offdiag_norm,info] = jacobi(a,abserr,max_sweeps)
%=======================================================================
% Thresholded Jacobi rotations for a real symmetric matrix.
% a is overwritten by a nearly diagonal similar matrix; its diagonal
% contains the eigenvalues. vectors contains eigenvectors in columns.
% info: 0 converged, 1 invalid input, 2 sweep limit.
% abserr bounds the SUM of squared off-diagonal entries.
%
% Rotate rows and columns of a by the same angle and rotate columns of
% vectors. Recompute the actual off-diagonal sum after each sweep.
% The stable tangent formula avoids cancellation for unequal diagonals.
% Row and column operations use temporary arrays to preserve old values.
%=======================================================================
n = size(a,1);
vectors = eye(n);
sweeps = 0;
rotations = 0;
offdiag_norm = Inf;
info = 1;
if n < 1 || size(a,2) ~= n || abserr <= 0 || max_sweeps < 1 || ...
        max_sweeps ~= floor(max_sweeps)
    return;
end
if ~isequal(a,a') || ~all(isfinite(a(:)))
    return;
end

while true
    b2 = 0;
    for i = 1:n
        for j = 1:n
            if i ~= j
                b2 = b2+a(i,j)^2;
            end
        end
    end
    offdiag_norm = sqrt(b2);
    if b2 <= abserr
        info = 0;
        return;
    end
    if sweeps >= max_sweeps
        info = 2;
        return;
    end
    bar = 0.5*b2/(n*n);
    for i = 1:n-1
        for j = i+1:n
            if a(j,i)^2 <= bar
                continue;
            end
            b2 = max(0,b2-2*a(j,i)^2);
            bar = 0.5*b2/(n*n);

            beta = (a(j,j)-a(i,i))/(2*a(j,i));
            % Unlike sign(0), Fortran's sign(1,0) gives +1.
            direction = 1;
            if beta < 0
                direction = -1;
            end
            tangent = -direction/(abs(beta)+sqrt(1+beta^2));
            c = 1/sqrt(1+tangent^2);
            s = tangent*c;

            % Rotate rows i and j.
            row_i = c*a(i,:)+s*a(j,:);
            row_j = -s*a(i,:)+c*a(j,:);
            a(i,:) = row_i;
            a(j,:) = row_j;

            % Rotate columns i and j of the matrix and eigenvectors.
            column_i = c*a(:,i)+s*a(:,j);
            column_j = -s*a(:,i)+c*a(:,j);
            a(:,i) = column_i;
            a(:,j) = column_j;
            column_i = c*vectors(:,i)+s*vectors(:,j);
            column_j = -s*vectors(:,i)+c*vectors(:,j);
            vectors(:,i) = column_i;
            vectors(:,j) = column_j;
            rotations = rotations+1;
        end
    end
    sweeps = sweeps+1;
end
end