objective_calc.m

% This function aims to calculate the size of the penalty terms in the
% objective function for a model
%
% It also has the option to generate a dataset from an input model and mesh
% flags are used to specify which functions are desired
%
% Inputs : mesh - generated using forward_modelling_2 using an input data
% file containing the measurements.
%
% Model - this is the resistivity model to be used. Currently generated by
% reading mesh into res2dmod_generator and saving the interpolated output

function objective_calc()
clearvars
close all
ref_grid_list = [{'mesh_define\big_block_l1_ls0lag'},{'mesh_define\big_block_l1_100lag'},{'mesh_define\gauss150_l1_100lag'},{'periodic\block2_1pc_l1_ls0lag'},{'periodic\golden_test\gb1_1pc_l1_ls0lag'}];

% Read in model
filename = 'mesh_define\gb1_refgrid.mat';
reference_grid = ref_grid_list{5};% 1-sqgrid; 2-recgrid; 3-vargrid; 4-refgrid; 5 - gb1 refgrid
user = getenv('username');
obj_flag = 1; % = 1 if want to calculate objective functions
fwd_flag = 0; % = 1 if want to calculate forward problem and sFave data
p_increment_plot_flag = 1;
p_plot = 0; % Value of mu to plot
p_conv = 1e-15;
tgv_itr =  100;
save_flag = 1; % saves data to subfolder in mesh_define with original filename

periodic_flag = 1;
fd_weight_flag = 0;
tgv_lagrn = [0.6];
grad_adj_flag = 1; % =1 if want no distances to be taken into account in cx, cy
fsz = 20; %fontsize
mkrsz = 10;
lw = 1.2;

% filepath = ['C:\Users\',user,'\OneDrive - The University of Nottingham\UoN Box Migration\data\resipy_models'];
filepath = ['D:\TGV_revision\thesis\'];

load_model = load(fullfile(filepath,filename));
if isfield(load_model, 'ip4di_direct')
    model = load_model.ip4di_direct;
    disp('loaded ip4di_direct')
else
    model = load_model.ip4di_im;
    disp('loaded ip4di_im')
end
% model = 10.^repmat(1.3 + (0.5./13)*(1:13), 45, 1)'; % for loglin_v
param = reshape(model',[],1); % same from as in ip4di. 
% param = reshape(model(:,3:end-2)',[],1); % same from as in ip4di. Cuts off extended edges
% xdim = size(model,2);
xdim = size(model,2);
ydim = size(model,1);

% Read in mesh, add model to mesh structure
mesh = []; input = [];
input.p_conv = p_conv;
input.p_init_flag = 1; % choice of how to initialise p in sp2. [0 is p = 0] [1 is p = grad(m)] [2 is result of an l2 iteration performed when p = 1] [3 is standard gaussian noise]
input.diff_type = 1; % 1 - forward-backward, 3 - half point
input.ghost_flag = 0;
input.sp2_itr = tgv_itr;

% load('D:\TGV_revision\thesis\neumann\big_block_l1_ls0lag')
load(fullfile(filepath, reference_grid))


% load(['C:\Users\',user,'\OneDrive - The University of Nottingham\UoN Box Migration\TGV_revision\resipy\periodic\ref\plume_layer_rpy1_1pc_l1_ls11lag']);
% load(['C:\Users\',user,'\OneDrive - The University of Nottingham\UoN Box Migration\TGV_revision\resipy\periodic\plume_layer_rpy1_1pc_l1_ls11lag']);
% load('test mesh');
mesh.res_param1 = param; mesh.init_res_param = param;  mesh.res_param2 = param; input.bgr_res_param = param; mesh.num_param = length(param);
mesh.num_param = final.num_param; mesh.param_x = final.param_x; mesh.param_y = final.param_y;

% % Calculate gradient matrix
% [mesh.cx, mesh.cy] = calc_grad(mesh);

tmp_x = unique(final.param_x);
tmp_y = unique(final.param_y);

% adjust image to fit
% tmp_x = tmp_x(4:end-4);
% tmp_y = tm;

%


if periodic_flag == 0
    [mesh.cx, mesh.cy] = c_calc(tmp_x, tmp_y);
else
    [mesh.cx, mesh.cy] = c_calc_periodic(tmp_x, tmp_y, fd_weight_flag);
end


% remove gradient distance normalisation 
if grad_adj_flag == 1
    mesh.cy(mesh.cy>0) = 1;
    mesh.cy(mesh.cy<0) = -1;
    mesh.cx(mesh.cx>0) = 1;
    mesh.cx(mesh.cx<0) = -1;
end

% figure
% imagesc(mesh.cx)
% title('cx')
% figure
% imagesc(mesh.cy)
% title('cy')



% Calculate the objective function terms
% TV term: |grad(m)|
grad_sumSq = (mesh.cx*log10(mesh.res_param1)).^2 + (mesh.cy*log10(mesh.res_param1)).^2;
grad2_sumSq = (mesh.cx'*mesh.cx*log10(mesh.res_param1)) + (mesh.cy'*mesh.cy*log10(mesh.res_param1));
dx_im = reshape(mesh.cx*log10(mesh.res_param1), length(tmp_x), [])';
dy_im = reshape(mesh.cy*log10(mesh.res_param1), length(tmp_x), [])';
TV_reg = sum(sqrt(grad_sumSq(:)));
L2_reg = sqrt(sum(grad_sumSq(:)));
dx = tmp_x(2) - tmp_x(1);
dy = tmp_y(2) - tmp_y(1);

% TGV term - need to calculate p from minimisation problem once.
px = zeros(length(mesh.res_param1),length(tgv_lagrn));
py = px;
gamma_p = 1e-15;               % cutoff for |x| -> 0
gamma_c = 1e-12;

if obj_flag == 1
    
    for i = 1:length(tgv_lagrn);
        input.tgv_lagrn = tgv_lagrn(i);
        [mesh, rms_p] = TGV_sp2(input, mesh, p_increment_plot_flag, fsz);
        TGV_t1_im = (sqrt((mesh.cx*log10(mesh.res_param1) - mesh.px).^2 + (mesh.cy*log10(mesh.res_param1) - mesh.py).^2));
        TGV_t2_im = (sqrt( (mesh.cx'*mesh.px).^2 + (mesh.cy'*mesh.py).^2 +  0.5*(mesh.cy'*mesh.py + mesh.cx'*mesh.px).^2));
        TGV_t1(i) = sum(TGV_t1_im(:));
        TGV_t2(i) = sum(TGV_t2_im(:));
        
        px(:,i) = mesh.px;
        py(:,i) = mesh.py;
        
        px_im = reshape(mesh.px,xdim,[])';
        py_im = reshape(mesh.py,xdim,[])';
        pz_im = zeros(size(px_im));
        z_cen = 100*ones(size(px_im));
        
        if p_increment_plot_flag == 1
            figure(5)
            surf(unique(mesh.param_x),unique(mesh.param_y),sqrt(px_im.^2 + py_im.^2),'edgecolor','none')
            view([0,0,1])
            title(['p, \mu = ',num2str(input.tgv_lagrn)],'interpreter','none')
            colorbar
            set(gca,'ydir','reverse')
            caxis([min(min(sqrt(px_im.^2 + py_im.^2))),max(max(sqrt(px_im.^2 + py_im.^2)))]);
            hold on
            quiver3(unique(mesh.param_x) + dx/2, unique(mesh.param_y) + dy/2, z_cen, px_im, py_im, pz_im,'color','m');
            axis image
            set(gca,'fontsize',fsz)
            hold off
        end
        
        if input.tgv_lagrn == 2
            TGV_t1_im_2 = reshape(TGV_t1_im,xdim,[])';
        end
        
    end
    
    
    TGV_reg = TGV_t1 + input.tgv_lagrn.*TGV_t2;
    p_im = sqrt(px_im.^2 + py_im.^2);
    grad_im = reshape(sqrt(grad_sumSq),xdim,[])';
    grad2_im = reshape((grad2_sumSq),xdim,[])';
    
    % Print results
    disp([filename, ' objective terms (\lambda_2 = 2)'])
    disp(['L2 objective = ',num2str(L2_reg)])
    disp(['TV objective = ',num2str(TV_reg)])
    disp(['TGV objective = ',num2str(sum(TGV_reg.*(tgv_lagrn == 2)))])
    
    figure(1)
    plot(tgv_lagrn,TV_reg.*ones(1,length(tgv_lagrn)),'-','markersize',mkrsz, 'linewidth', lw)
    hold on
    plot(tgv_lagrn,TGV_reg,'o-','markersize',mkrsz, 'linewidth', lw)
    plot(tgv_lagrn,TGV_t1,'^-','markersize',mkrsz, 'linewidth', lw)
    plot(tgv_lagrn,TGV_t2,'v-','markersize',mkrsz, 'linewidth', lw)
    %     plot(tgv_lagrn,L2_reg.*ones(1,length(tgv_lagrn)),'v-','markersize',mkrsz, 'linewidth', lw)
    legend('TV = |\nabla m|_{L1}','TGV_\mu = TGV^{(1)}_{\mu} + \muTGV^{(2)}_{\mu}','TGV^{(1)}_{\mu} = |\nabla m - p|_{l1}','TGV^{(2)}_{\mu} = 0.5|\nabla p + \nabla p^T|_{l1}')% 'total TGV objective',
%     title('Impact of \mu on the terms of the TGV functional for static models','fontsize',fsz)
    xlabel('\mu','fontsize',fsz)
    ylabel(['Sum across all model cells'],'fontsize',fsz)
    set(gca,'fontsize',fsz)
    hold off
    
    figure(2)
    surf(unique(mesh.param_x),unique(mesh.param_y),p_im,'edgecolor','none')
    view([0,0,1])
    title('p','fontsize',fsz)
    set(gca,'ydir','reverse')
    colorbar
    axis image
    set(gca,'fontsize',fsz)

    figure(3)
    surf(unique(mesh.param_x),unique(mesh.param_y),grad_im,'edgecolor','none')
    hold on
    dz_im = grad_im;
    dz_im(grad_im > 1) = 1;
    q_hand = quiver3(unique(mesh.param_x) + dx/2,unique(mesh.param_y) + dy/2, z_cen, dx_im, dy_im, dz_im,'color','m','LineWidth', 2);
    view([0,0,1])
    title('|\nabla m|','fontsize',fsz)
    set(gca,'ydir','reverse')
    colorbar
    axis image
    set(gca,'fontsize',fsz)
    
    figure(6)
    surf(unique(mesh.param_x),unique(mesh.param_y),grad2_im,'edgecolor','none')
    view([0,0,1])
    title('|\nabla^2 m|','fontsize',fsz)
    set(gca,'ydir','reverse')
    colorbar
    axis image
    set(gca,'fontsize',fsz)
    
    figure(7)
    surf(unique(mesh.param_x),unique(mesh.param_y),(grad_im - p_im),'edgecolor','none','LineWidth', 2)
    view([0,0,1])
    set(gca,'ydir','reverse')
    axis image
    set(gca,'fontsize',fsz)
    colorbar
    title('|\nabla m| - |p|')
    

    
    
%     figure(6)
%     surf(unique(mesh.param_x),unique(mesh.param_y),TGV_t1_im_2,'edgecolor','none')
%     view([0,0,1])
%     title('TGV t1 \lambda_2 = 2')
%     set(gca,'ydir','reverse')
%     colorbar
%     axis image
    
end

figure(4)
surf(unique(mesh.param_x),unique(mesh.param_y),reshape(log10(mesh.res_param1),xdim,[])','edgecolor','none')
view([0,0,1])
title('Model','fontsize',fsz)
set(gca,'ydir','reverse')
colorbar
axis image
set(gca,'fontsize',fsz)

if fwd_flag == 1
    % Calculate forward problem too - ideally should save data for inversion
    fem = forward_solver(input,mesh);
    
    fid = fopen([filepath,filename,'.d'],'w');
    fprintf(fid,'%g %g %g %g %g %g %g %g %g\n',[input.ax,input.az,input.bx,input.bz,input.mx,input.mz,input.nx,input.nz,fem.array_model_data]');
    fclose(fid);
end

if save_flag == 1  
    file_out = ['p_init',num2str(input.p_init_flag), '_mu', num2str(tgv_lagrn*10)];
    [save_path, save_name, ~] = fileparts(filename);
    cx = mesh.cx;
    cy = mesh.cy;
    
    save(fullfile(filepath, save_path, save_name, file_out), 'rms_p', 'px', 'py', 'xdim', 'cx', 'cy','model')    
end

clear -global
end


% Calculates cx, cy as per smooth matrix4 (with extra options removed)
function [cx, cy] = calc_grad(mesh)
[num_param, num_nodes, param_x, param_y, ~] = neumann_bc(mesh.num_param, mesh.param_x, mesh.param_y, mesh.res_param1, [0, 2]);

% factor by which surface layer damping is increased
surface_damp_strength = 15;

c=zeros(num_param,num_param);
cx=zeros(num_param,num_param);
cy=zeros(num_param,num_param);
dcx = 1;
dcy = 1;


% Add dummy rows/columns at the ends for length calculations
tmp_x = unique(param_x);
tmp_y = unique(param_y);


% In each row, = 1 where an adjacent parameter is present
for i=1:num_param
    
    current_x=param_x(i);
    current_y=param_y(i);
    indX = find(tmp_x == current_x);
    indY = find(tmp_y == current_y);
    % search all other parameters that have the same y and the x=ind+1
    if indX < (length(tmp_x))
        
        if input.fd_weight == 1
            %             dcx = (tmp_y(indY+1) - tmp_y(indY - 1))/2;
            dcx = (tmp_x(indX + 1) - tmp_x(indX))/2;
        end
        
        if input.damp_surface_flag == 1 && abs(current_y) == min(abs(param_y))
            cx(i,i) = -surface_damp_strength.*1./dcx;
        else
            cx(i,i) = -1./dcx;
        end
        
        for j=1:num_param
            
            
            if (param_y(j)==current_y) && (param_x(j)==tmp_x(indX+1))
                
                if input.damp_surface_flag == 1 && abs(current_y) == min(abs(param_y))
                    cx(i,j) = surface_damp_strength.*1./dcx;
                else
                    cx(i,j) = 1./dcx;
                    break
                end
            end
        end
        
    end
    
end



for i=1:num_param
    
    current_x=param_x(i);
    current_y=param_y(i);
    indX = find(tmp_x == current_x);
    indY = find(tmp_y == current_y);
    % search all other parameters that have the same x and the y=ind+1
    if indY < (length(tmp_y))
        
        if input.fd_weight == 1
            %             dcy = (tmp_x(indX + 1) - tmp_x(indX - 1))/2;
            dcy = tmp_y(indY+1) - tmp_y(indY);
        end
        
        if input.damp_surface_flag == 1 && abs(current_y) == min(abs(param_y))
            cy(i,i) = 1./dcy;
            %             cy(i,i) = -surface_damp_strength.*1./dcx;
        else
            cy(i,i) = -1/dcy;
        end
        
        for j=1:num_param
            if param_y(j)==tmp_y(indY+1) && param_x(j)==current_x
                if input.damp_surface_flag == 1 && abs(current_y) == min(abs(param_y))
                    cy(i,j) = 1./dcy;
                    %                     cy(i,i) = -surface_damp_strength.*1./dcx;
                else
                    cy(i,j) = 1/dcy;
                    break
                end
            end
        end
    end
end
end

function [mesh, rms_p] = TGV_sp2(input, mesh, plot_flag, fsz)

% assumes not using old neumann_bc ghost points
num_param = mesh.num_param;
domain_ind = 1:num_param;
param_x = mesh.param_x;
param_y = mesh.param_y;
res_param = mesh.res_param1;

grad_mx = mesh.cx*log10(res_param);
grad_my = mesh.cy*log10(res_param);



%initialise p - includes bc ghost point terms
if input.p_init_flag == 0 || input.p_init_flag == 2
    mesh.px = zeros(mesh.num_param,1);
    mesh.py = mesh.px;
elseif input.p_init_flag == 1 % p = grad(m) on domain of m.
    mesh.px = mesh.cx*res_param;
%     mesh.px = mesh.px(domain_ind);
    mesh.py = mesh.cy*res_param;
%     mesh.py = mesh.py(domain_ind);
elseif input.p_init_flag == 3 % p = homogeneous half space + gaussian noise
    noise_percent = 5;
    rand_seed = 38643;
    
    mesh.px = randn(mesh.num_param, 1);
    mesh.py = randn(mesh.num_param, 1);   
end

% Used in main calculation
gamma_p = 1e-15;               % cutoff for |x| -> 0
gamma_c = 1e-12;
% Used in performance metrics
p1 = zeros(2*mesh.num_param,1);
rms_p = zeros(1,input.sp2_itr+1);

% main sp2 loop
for i = 2:(input.sp2_itr+1)
    
    if input.diff_type ~= 3
        if input.ghost_flag == 1
            % add ghost points to p. Ax p,m have same domain, cx, cy are
            % the same, but different bc (first order homogeneous neumann).
            [~, ~, ~, ~, p_tmp] = neumann_bc(mesh.num_param, mesh.param_x, mesh.param_y, [mesh.px, mesh.py], [0, 1]);
            px_tmp = p_tmp(:,1);
            py_tmp = p_tmp(:,2);
        else
            px_tmp = mesh.px;
            py_tmp = mesh.py;
        end
        
        % calculate weights
        if input.p_init_flag == 2 && i == 2 % l2 1st iteration as initialisation
            Rm = speye(length(grad_mx));
            Rp = speye(length(px_tmp));
        else
            Rm = spdiags(1./sqrt((grad_mx - px_tmp).^2 + (grad_my - py_tmp).^2 + gamma_c^2), 0, length(grad_mx), length(grad_mx));
            Rp = spdiags(1./sqrt((mesh.cx'*px_tmp).^2 + (mesh.cy'*py_tmp).^2 + 0.5*(mesh.cx'*py_tmp + mesh.cy'*px_tmp).^2 + gamma_p.^2), 0, length(px_tmp), length(px_tmp) );
        end
        % Set up matrix equation for px, py and solve
        a11 = Rm + input.tgv_lagrn*(mesh.cx*Rp*mesh.cx' + 0.5*mesh.cy*Rp*mesh.cy');
        a12 = input.tgv_lagrn*0.5*mesh.cy*Rp*mesh.cx';
        a21 = input.tgv_lagrn*0.5*mesh.cx*Rp*mesh.cy';
        a22 = Rm + input.tgv_lagrn*(mesh.cy*Rp*mesh.cy' + 0.5*mesh.cx*Rp*mesh.cx');
        b1 = Rm*mesh.cx*log10(res_param);
        b2 = Rm*mesh.cy*log10(res_param);
        
        % Domain_ind first
        a11 = a11(domain_ind, domain_ind);
        a12 = a12(domain_ind, domain_ind);
        a21 = a21(domain_ind, domain_ind);
        a22 = a22(domain_ind, domain_ind);
        b1 = b1(domain_ind);
        b2 = b2(domain_ind);
        
        %                     A = [a11, a12; a21, a22];
        %                     b = [b1; b2];
        A = [a11, a12; a21, a22];
        b = [b1; b2];
        %                     condition = cond(A);
        %                     disp(['itr: ', num2str(i-1), ' | cond(A): ', num2str(condition,'%.2g')])
        p2 = A\b;
        %                     p2 = p2(domain_ind);
        clear A b b1 b2 a11 a12 a21 a22
    else
        [p2] = p_calc(mesh, 1, input.tgv_lagrn, gamma_c, gamma_p);
    end
    
    % spearate px, py
    mesh.px = p2(1:length(p2)/2);
    mesh.py = p2((length(p2)/2+1):length(p2));
    % eliminate ghost points
    %                         mesh.px = mesh.px(domain_ind);
    %                         mesh.py = mesh.py(domain_ind);
    p2 = [mesh.px; mesh.py];
    
    % Store results
    px(:,i) = mesh.px;
    py(:,i) = mesh.py;
    
    % root mean square difference between last and current
    % solution normalise by rms p value
    %                     rms_p(i) = sqrt(mean((p2 - p1).^2))/(0.5*sqrt(mean(p2.^2 + p1.^2));
    
    
    
    rms_p(i-1) = sqrt(mean((p2 - p1).^2))/(0.5*sqrt(mean( (mesh.cx*log10(res_param)).^2 + (mesh.cy*log10(res_param)).^2)));
    %         disp(['rms dp = ',num2str(p_rms(i))])
    %                 rms_change(i) = abs((p_rms(i) - p_rms(i-1)))./(p_rms(i) + p_rms(i-1)); % percentage rms change per itr
    
    disp(['tgv iteration ',num2str(i-1),' : rms dp = ',num2str(rms_p(i-1))])
    
    
    if ((i > 4) && (rms_p(i-1) < input.p_conv)) %  stop if change smaller than p_conv after first few itr
        sp2_i = i;  % store number of iterations needed
        break
%     elseif (i > 2 && (rms_p(i-1) > rms_p(i-2))) % stop if diverging
%         sp2_i = i-1;  % store number of iterations needed
%         mesh.px = p1(1:length(p1)/2);
%         mesh.py = p1((length(p1)/2+1):length(p1));
%         break
    end
    
    
    if i == (input.sp2_itr+1)
        sp2_i = i;
    end
    
    p1 = p2;
    
end


rms_p(i:end) = [];


if plot_flag == 1
    % Plot convergence
    figure(10)
    hold off
    semilogy(rms_p)
    title('rms change in p')
    xlabel('sp 2 iteration')
    ylabel('rms dp')
    set(gca,'fontsize',fsz)
end


end