%% Homework 2 - Optimization Methods
% Author: Claudio Maggioni
%
% Note: exercises are not in the right order due to matlab constraints of
% functions inside of scripts.

clear
clc
close all
plots = 1;

%% 1.4 - Solution for 1D Poisson for N=1000 using CG

n = 1000;
[A, b] = build_poisson(n);

A(2, 1) = 0;
A(n-1, n) = 0;

[x, ys, gnorms] = CGSolve(A, b, zeros(n,1), n, 1e-8);
display(x);

%% 1.5 - Plots for the 1D Poisson solution

if plots
    plot(0:(size(ys,1)-1), ys);
    sgtitle("Objective function values per iteration");
    axis([-1 500 -inf +inf]);

    figure;
    semilogy(0:(size(gnorms,1)-1), gnorms);
    sgtitle("Log of gradient norm per iteration");
    axis([-1 500 -inf +inf]);
end
%% 2.1 - Matrix definitions

A1 = diag([1:10]);
A2 = diag(ones(1,10));
A3 = diag([1, 1, 1, 3, 4, 5, 5, 5, 10, 10]);
A4 = diag([1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0]);

%% 2.2 - Application of CG for each AX/b couple

rng(0); % random seed fixed due to reproducibility purposes
b = rand(10,1);
display(b);

n = 10;
[x1, ys1, gnorms1, xs1] = CGSolve(A1, b, zeros(n,1), n, 1e-8);
[x2, ys2, gnorms2, xs2] = CGSolve(A2, b, zeros(n,1), n, 1e-8);
[x3, ys3, gnorms3, xs3] = CGSolve(A3, b, zeros(n,1), n, 1e-8);
[x4, ys4, gnorms4, xs4] = CGSolve(A4, b, zeros(n,1), n, 1e-8);

%% 2.3 - Logarithm energy norm of the error computation

if plots 
    enl_plot(x1, xs1, A1);
    sgtitle("Log energy norm of the error per iter. (matrix A1)");
    enl_plot(x2, xs2, A2);
    sgtitle("Log energy norm of the error per iter. (matrix A2)");
    enl_plot(x3, xs3, A3)
    sgtitle("Log energy norm of the error per iter. (matrix A3)");
    enl_plot(x4, xs4, A4);
    sgtitle("Log energy norm of the error per iter. (matrix A4)");
end

function enl_plot(xsol, xs, A)
    enls = zeros(size(xs, 2), 1);
    for i = 1:size(xs, 2)
        x = xs(:, i);
        enls(i) = log((xsol - x)' * A * (xsol - x));
    end
    
    figure;
    plot(0:(size(xs, 2)-1), enls, '-k.');
    axis([-1 11 -35 +2]);
end

%% 1.1 - Build the Poisson matrix A and vector b

function [A,b] = build_poisson(n)
    A = diag(2 * ones(1,n));
    
    A(1,1) = 1;
    A(n,n) = 1;
    
    for i = 2:n-1
        A(i, i+1) = -1;
        A(i, i-1) = -1;
    end
    
    h = 1 / (n - 1);
    
    b = h^2 * ones(n, 1);
    b(1) = 0;
    b(n) = 0;
end

%% 1.3 - Implementation of Conjugate Gradient

function [x, ys, gnorms, xs] = CGSolve(A, b, x, max_itr, tol)
    ys = zeros(max_itr + 1, 1);
    gnorms = zeros(max_itr + 1, 1);
    xs = zeros(size(x, 1), max_itr + 1);
    
    r = A * x - b;
    p = -r;
    
    gnorms(1) = norm(p, 2);
    xs(:, 1) = x;
    ys(1) = 1/2 * dot(A*x, x) - dot(b, x);
    
    for i = 1:(max_itr+1)
        
        alpha = -r' * p / (p' * A * p);
        x = x + alpha * p;
        r = A * x - b;
        beta = (r' * A * p) / (p' * A * p);  
        p = -r + beta * p;
        
        gnorms(i+1) = norm(p, 2);
        xs(:, i+1) = x;
        ys(i+1) = 1/2 * dot(A*x, x) - dot(b, x);
        
        if gnorms(i+1) < tol
            ys = ys(1:(i+1));
            gnorms = gnorms(1:(i+1));
            xs = xs(:, 1:(i+1));
            break
        end
    end
end