% BestInversionComparisonFigure_v1.m 

% Lea-Sophie Höyns, 01/21/2024

% This script makes a figure comparing the inversion results for best-value
% lambda, max lambda, and min lambda.  For each lambda value, it shows drag
% coefficient, drag, and velocity misfit.  

% Parts of this script are taken from the script InversionComparisonFigure_v1.m 
% by M. Wolovick from 12/9/2022 (https://doi.org/10.5281/zenodo.7798650).

% Colormaps citation: 
% [1] Zhaoxu Liu / slandarer (2024). 200 colormap (https://www.mathworks.com/matlabcentral/fileexchange/120088-200-colormap), 
% MATLAB Central File Exchange. Retrieved February 13, 2024.
% [2] Crameri, Fabio. (2021). Scientific colour maps (7.0.1). Zenodo. https://doi.org/10.5281/zenodo.5501399

% Figure for the experiment: Ncuas, m=3
% a.-c. Drag coefficient map for lambda_min, lambda_best, lambda_max
% d.-f. Drag map for lambda_min, lambda_best, lambda_max
% g.-i. Velocity misfit for lambda_min, lambda_best, lambda_max on symmetric log-scale

clearvars
tic
close all

%% PATHS
issmscriptpath='/albedo/work/projects/landice/ISSM/ISSM_installations/ISSM_v4p22/trunk/';
functionpath='~/1_Models/WAIS_FrictionOptimization_v2/Scripts_MW/Functions';
slancm='~/1_Models/WAIS_FrictionOptimization_v2/Scripts_MW/Functions/slanCM';
addpath(genpath(issmscriptpath),functionpath,slancm)
% File names:
lcurvefile='/albedo/scratch/user/lhoeyns/1_ISSM_execution/Zwischenspeicher/OldLcurveRange/Run_1/WAIS.dragInv_HO_m3_Ncuas_1DTemp_Lcurve_17-Jan-2024.mat';
figname='/albedo/home/lhoeyns/1_Models/WAIS_FrictionOptimization_v2/Scripts_MW/Plotscripts/Figurescripts-Manuscript/Zenodo/BestInversionComparisonFigure_v1_final.png';

%% PARAMETERS

% Ticks and limits:
xtick=500;                         % km
ytick=500;                         % km
% Caxis limits: 
dragcoeffclims=[-4,-0.5];          % [1x2]
draglims=[0,200];                  % [1x2] kPa
dragtick=50;                       % kPa 
veldifflims=[1,1e3];               % [1x2] m/yr (goes into symlog scale)

% Aesthetic settings
% Colors:
backplanecolor=[.75,.75,.75];      % valid color string or RGB
% Colormaps:
dragcoeffcmap=crameri('-romaO');   % valid colormap string
flipdragcoeffcmap=0;               % logical
dragcmap=slanCM('twilight_s');     % valid colormap string
flipdragcmap=0;                    % logical
velcmap=crameri('-bam');           % valid colormap string

% Fontsizes:
axesfontsize=14;                   % points
titlefontsize=14;                  % points
colorbarfontsize=15;               % points
% Linewidths:
axeslinewidth=2;                   % points

% Output size and resolution:
TileSpacing='tight';
Padding='compact';
resolution=300;                    % dpi

%% PREPARATION

% Load basic model:
load(lcurvefile,'BasicModel','MyInversionParameters','*best*','Lambda')

% Rename model structure:
md=BasicModel;

% Unpack inversion parameters:
unpack(MyInversionParameters)

% Identify max/min acceptable models:
bestmodelind_max=find(Lambda<=bestlambda_maxallowable,1,'last');
bestmodelind_min=find(Lambda>=bestlambda_minallowable,1,'first');

% Compute mesh properties:
% Extract 2D mesh:
try
    x2d=md.mesh.x2d;
    y2d=md.mesh.y2d;
    elements2d=md.mesh.elements2d;
    oldnumlayers=md.mesh.numberoflayers;
catch
    x2d=md.mesh.x;
    y2d=md.mesh.y;
    elements2d=md.mesh.elements;
    oldnumlayers=1;
end
numelements2d=size(elements2d,1);
numvertices2d=length(x2d);
% Compute edge lengths:
a=sqrt((x2d(elements2d(:,1))-x2d(elements2d(:,2))).^2+(y2d(elements2d(:,1))-y2d(elements2d(:,2))).^2);
b=sqrt((x2d(elements2d(:,2))-x2d(elements2d(:,3))).^2+(y2d(elements2d(:,2))-y2d(elements2d(:,3))).^2);
c=sqrt((x2d(elements2d(:,3))-x2d(elements2d(:,1))).^2+(y2d(elements2d(:,3))-y2d(elements2d(:,1))).^2);
% Compute area:
s=(a+b+c)/2;
Area=sqrt(s.*(s-a).*(s-b).*(s-c));
% Compute hydraulic diameter:
HydraulicDiameter_elements=2*Area./s;
% Compute element centers:
CentX=mean(x2d(elements2d),2);
CentY=mean(y2d(elements2d),2);
% Find mesh limits:
xlims=[min(x2d),max(x2d)];
ylims=[min(y2d),max(y2d)];
% Extract masks on 2D mesh and truncate:
OceanLevelSet_vertices=max(-1,min(1,md.mask.ocean_levelset(1:numvertices2d)));
IceLevelSet_vertices=max(-1,min(1,md.mask.ice_levelset(1:numvertices2d)));
% Compute ocean and ice area fractions:
OceanAreaFraction_elements=(1-mean(OceanLevelSet_vertices(elements2d),2))/2;
IceAreaFraction_elements=(1-mean(IceLevelSet_vertices(elements2d),2))/2;

% Generate the active mask:
ActiveMask_drag=IceLevelSet_vertices<=0&OceanLevelSet_vertices>=0;
ActiveMask_elements=mean(ActiveMask_drag(elements2d),2);

% Unpack the good data mask:
GoodData_obs_vertices=md.miscellaneous.dummy(1:numvertices2d);

%% MAKE Comparison FIGURE
FigH = figure('Position', get(0, 'Screensize'));
tiledlayout(3,3,'TileSpacing',TileSpacing,'Padding',Padding) 

%% MAKE Drag Coeff SUBFIGURE

% Generate drag coeff units string:
if ~strcmp(nsource,'none')
    if r==1
        if m==1
            cunitsstring='(m/yr)^{-1}';
        else
            cunitsstring=['(m/yr)^{-1/',num2str(m),'}'];
        end
    else
        if m==1
            cunitsstring=['Pa^{',num2str(1-r),'}(m/yr)^{-1}'];
        else
            cunitsstring=['Pa^{',num2str(1-r),'}(m/yr)^{-1/',num2str(m),'}'];
        end
    end
else
    if m==1
        cunitsstring='Pa(m/yr)^{-1}';
    else
        cunitsstring=['Pa(m/yr)^{-1/',num2str(m),'}'];
    end
end

% Loop through models:
for thismodel=1:3
    
    % Get this model ind:
    if thismodel==1
        thismodelind=bestmodelind_min;
        lambdadescrip='\lambda_{min}';
    elseif thismodel==2
        thismodelind=bestmodelind;
        lambdadescrip='\lambda_{best}';
    else
        thismodelind=bestmodelind_max;
        lambdadescrip='\lambda_{max}';
    end
    
    % Construct this file name:
    slashnum=strfind(lcurvefile,'/');
    thisbase=10^floor(log10(Lambda(thismodelind)));
    thisname=[lcurvefile(slashnum(end)+1:end-4),'_Lambda_',num2str(thisbase*.01*round(100*Lambda(thismodelind)/thisbase)),'.mat'];

    % Load this model:
    load([lcurvefile(1:slashnum(end)),thisname],'md')

    % Extract variables we need:
    if numlayers>1 && ~strcmp(flowequation,'SSA')
        ModelDragCoeff=md.results.StressbalanceSolution.FrictionCoefficient(md.mesh.vertexonbase==1);
    else
        ModelDragCoeff=md.results.StressbalanceSolution.FrictionCoefficient;
    end
     
    % Compute drag coeff ticks and limits:
    if thismodel==1
        % Compute limits from percentiles:
        cclims=2*log10(prctile(ModelDragCoeff(ActiveMask_drag)/(secondsperyear^(1/(2*m))),[.1,99.9]));
        % Round up/down to powers of ten:
        cclims=[floor(cclims(1)),ceil(cclims(2))];
        % Compute ticks and labels for color bar:
        % Compute first base and tick for color bar:
        thiscbase=10^cclims(1);
        thistick=thiscbase;
        % Create all ticks and labels for color bar:
        ccticks=zeros(100,1);
        cclabels=cell(100,1);
        for ii=1:100
            % Record this tick:
            ccticks(ii)=log10(thistick);
            % Check if we hit the next power of ten:
            if abs(thistick-10*thiscbase)/thiscbase<1e-12 || (ii==1 && abs(thistick-thiscbase)/thiscbase<1e-12)
                % Generate label:
                cclabels{ii}=['10^{',num2str(round(log10(thistick))),'}'];
                % Advance base:
                thiscbase=thistick;
            end
            % Advance tick:
            thistick=thistick+thiscbase;
            % Break from loop and truncate:
            if thistick>(1+1e-12)*(10^cclims(2))
                ccticks=ccticks(1:ii);
                cclabels=cclabels(1:ii);
                break
            end
        end
     end
     
     % Make subplot:
     nexttile
     % Draw figure:
     patch('Vertices',[x2d/1000,y2d/1000],'Faces',elements2d(IceAreaFraction_elements>.5&OceanAreaFraction_elements<.5,:),'EdgeColor','none','FaceColor','interp','FaceVertexCData',2*log10(ModelDragCoeff/(secondsperyear^(1/(2*m)))))
     
     % Set up color scale:
     eval(['thiscmap=colormap(dragcoeffcmap);'])
     if flipdragcoeffcmap
         colormap(gca,flipud(thiscmap))
     else
         colormap(gca,thiscmap)
     end
     colormap(gca,crameri('-romaO'));
     caxis(dragcoeffclims)
     %set(gca,'ColorScale','log')
     %caxis(gca,cclims)
     if thismodel==3
         hc=colorbar;
         %set(hc,'YTick',ccticks)
         %set(hc,'YTickLabel',cclabels)
         set(get(hc,'Ylabel'),'String',cunitsstring,'FontSize',colorbarfontsize)
         set(hc,'TickDir','out')
         set(hc,'FontSize',colorbarfontsize,'LineWidth',axeslinewidth)
         set(hc,'Ticks',[-4 -3.5 -3 -2.5 -2 -1.5 -1 -0.5])

     end
     % Set up axes:
     xlim(xlims/1000)
     ylim(ylims/1000)
     set(gca,'XTick',xtick*[ceil(xlims(1)/(1e3*xtick)):1:floor(xlims(2)/(1e3*xtick))])
     set(gca,'YTick',ytick*[ceil(ylims(1)/(1e3*ytick)):1:floor(ylims(2)/(1e3*ytick))])
     if thismodel==1
        ylabel('y (km)','FontSize',axesfontsize)
     else
         set(gca,'YTickLabel',[])
     end
     %xlabel('x (km)','FontSize',axesfontsize)
     set(gca,'XTickLabel',[])
     set(gca,'PlotBoxAspectRatio',[(xlims(2)-xlims(1))/(ylims(2)-ylims(1)),1,1])
     %set(gca,'TickDir','out')
     set(gca,'FontSize',axesfontsize,'LineWidth',axeslinewidth)
     box on 
     % Make title:
     title(['\rm (',char(96+thismodel),') Drag coefficient for ',lambdadescrip,'=',num2str(num2str(thisbase*.01*round(100*Lambda(thismodelind)/thisbase)))],'FontSize',titlefontsize)
end  

%% MAKE Drag SUBFIGURE

% Loop through models:
for thismodel=1:3
    
    % Get this model ind:
    if thismodel==1
        thismodelind=bestmodelind_min;
        lambdadescrip='\lambda_{min}';
    elseif thismodel==2
        thismodelind=bestmodelind;
        lambdadescrip='\lambda_{best}';
    else
        thismodelind=bestmodelind_max;
        lambdadescrip='\lambda_{max}';
    end
    
    % Construct this file name:
    slashnum=strfind(lcurvefile,'/');
    thisbase=10^floor(log10(Lambda(thismodelind)));
    thisname=[lcurvefile(slashnum(end)+1:end-4),'_Lambda_',num2str(thisbase*.01*round(100*Lambda(thismodelind)/thisbase)),'.mat'];

    % Load this model:
    load([lcurvefile(1:slashnum(end)),thisname],'md')

    % Extract variables we need:
     if numlayers>1 && ~strcmp(flowequation,'SSA')
        ModelVel=md.results.StressbalanceSolution.Vel(md.mesh.vertexonsurface==1);
        ModelBasalVel=md.results.StressbalanceSolution.Vel(md.mesh.vertexonbase==1);
        ObsVel=md.inversion.vel_obs;%(md.mesh.vertexonsurface==1);
        ModelDragCoeff=md.results.StressbalanceSolution.FrictionCoefficient(md.mesh.vertexonbase==1);
        if ~strcmp(nsource,'none')
            EffectivePressure_vertices=md.friction.effective_pressure(md.mesh.vertexonbase==1);
        end
     else
        ModelVel=md.results.StressbalanceSolution.Vel;
        ModelBasalVel=ModelVel;
        ObsVel=md.inversion.vel_obs;
        ModelDragCoeff=md.results.StressbalanceSolution.FrictionCoefficient;
        if ~strcmp(nsource,'none')
            EffectivePressure_vertices=md.friction.effective_pressure;
        end
     end
     
     % Compute drag:
     Drag=(ModelDragCoeff.^2).*((ModelBasalVel/secondsperyear).^(1/m));
     if ~strcmp(nsource,'none')
         Drag=Drag.*(EffectivePressure_vertices.^r);
     end
     % Set drag to zero under floating shelf and in ice-free areas:
     Drag(OceanLevelSet_vertices<0|IceLevelSet_vertices>0)=0;
          
     % Make subplot:
     nexttile 
     % Draw figure:
     patch('Vertices',[x2d/1000,y2d/1000],'Faces',elements2d(IceAreaFraction_elements>.5&OceanAreaFraction_elements<.5,:),'EdgeColor','none','FaceColor','interp','FaceVertexCData',Drag/1e3)
     
     % Set up color scale:
     %eval(['thiscmap=colormap(dragcmap);'])
     colormap(gca,slanCM('twilight_s'))
     caxis(gca,draglims)
     if thismodel==3
         hc=colorbar;
         set(get(hc,'Ylabel'),'String','kPa','FontSize',colorbarfontsize)
         set(hc,'YTick',dragtick*[ceil(draglims(1)/dragtick):1:floor(draglims(2)/dragtick)])
         set(hc,'TickDir','out')
         set(hc,'FontSize',colorbarfontsize,'LineWidth',axeslinewidth)
         set(hc,'Ticks',[0 50 100 150 200],'TickLabels',{'0','50','100','150','\geq 200'})
     end
     % Set up axes:
     xlim(xlims/1000)
     ylim(ylims/1000)
     set(gca,'XTick',xtick*[ceil(xlims(1)/(1e3*xtick)):1:floor(xlims(2)/(1e3*xtick))])
     set(gca,'YTick',ytick*[ceil(ylims(1)/(1e3*ytick)):1:floor(ylims(2)/(1e3*ytick))])
     if thismodel==1
         ylabel('y (km)','FontSize',axesfontsize)
     else
        set(gca,'YTickLabel',[])
     end
     %xlabel('x (km)','FontSize',axesfontsize)
     set(gca,'XTickLabel',[])
     set(gca,'PlotBoxAspectRatio',[(xlims(2)-xlims(1))/(ylims(2)-ylims(1)),1,1])
     %set(gca,'TickDir','out')
     set(gca,'FontSize',axesfontsize,'LineWidth',axeslinewidth)
     box on 
     % Make title:
     title(['\rm (',char(96+3+thismodel),') Drag for ',lambdadescrip,'=',num2str(num2str(thisbase*.01*round(100*Lambda(thismodelind)/thisbase)))],'FontSize',titlefontsize)
     
end 

%% MAKE Velocity Error SUBFIGURE

% Generate caxis ticks and labels for velocity difference:
% Round axes limits to powers of 10:
veldifflims=10.^[floor(log10(veldifflims(1))),ceil(log10(veldifflims(2)))];
% Determine number of decades on each side of zero:
numposdecades=round(log10(veldifflims(2)/veldifflims(1)));
% Pre-allocate colorbar tick vectors:
veldiffticks=zeros(2*9*numposdecades+1,1);
veldiffticklabels=cell(size(veldiffticks));
% Loop through positive values:
for ii=1:9*numposdecades
    % when ii=1, we are doing 2*thisbase
    % Compute this base:
    thisbase=10^floor(log10(veldifflims(1))+ii/9);
    % Compute tick value on linear scale:
    thistickvalue=thisbase*(rem(ii,9)+1);
    % Compute tick value on positive log scale:
    veldiffticks(9*numposdecades+1+ii)=log(thistickvalue/veldifflims(1))/log(veldifflims(2)/veldifflims(1));
    % Compute tick value on negative log scale:
    veldiffticks(9*numposdecades+1-ii)=-log(thistickvalue/veldifflims(1))/log(veldifflims(2)/veldifflims(1));
    % Generate label text:
    if rem(ii,9)==0
        veldiffticklabels{9*numposdecades+1+ii}=['+10^{',num2str(round(log10(thisbase))),'}'];
        veldiffticklabels{9*numposdecades+1-ii}=['-10^{',num2str(round(log10(thisbase))),'}'];
    end
end
% Generate "zero" tick and label:
veldiffticks(9*numposdecades+1)=0;
veldiffticklabels{9*numposdecades+1}=['\pm 10^{',num2str(log10(veldifflims(1))),'}'];

% Loop through models:
for thismodel=1:3
    
    % Get this model ind:
    if thismodel==1
        thismodelind=bestmodelind_min;
        lambdadescrip='\lambda_{min}';
    elseif thismodel==2
        thismodelind=bestmodelind;
        lambdadescrip='\lambda_{best}';
    else
        thismodelind=bestmodelind_max;
        lambdadescrip='\lambda_{max}';
    end
    
    % Construct this file name:
    slashnum=strfind(lcurvefile,'/');
    thisbase=10^floor(log10(Lambda(thismodelind)));
    thisname=[lcurvefile(slashnum(end)+1:end-4),'_Lambda_',num2str(thisbase*.01*round(100*Lambda(thismodelind)/thisbase)),'.mat'];

    % Load this model:
    load([lcurvefile(1:slashnum(end)),thisname],'md')


    % Extract variables we need:
    if numlayers>1 && ~strcmp(flowequation,'SSA')
        ModelVel=md.results.StressbalanceSolution.Vel(md.mesh.vertexonsurface==1);
        ModelBasalVel=md.results.StressbalanceSolution.Vel(md.mesh.vertexonbase==1);
        ObsVel=md.inversion.vel_obs;%(md.mesh.vertexonsurface==1);
        ModelDragCoeff=md.results.StressbalanceSolution.FrictionCoefficient(md.mesh.vertexonbase==1);
        if ~strcmp(nsource,'none')
            EffectivePressure_vertices=md.friction.effective_pressure(md.mesh.vertexonbase==1);
        end
    else
        ModelVel=md.results.StressbalanceSolution.Vel;
        ModelBasalVel=ModelVel;
        ObsVel=md.inversion.vel_obs;
        ModelDragCoeff=md.results.StressbalanceSolution.FrictionCoefficient;
        if ~strcmp(nsource,'none')
            EffectivePressure_vertices=md.friction.effective_pressure;
        end
    end
     
    % Compute velocity difference:
    VelDiff=ModelVel-ObsVel;
     
    % Make subplot:
    nexttile

    % Convert real values to scaled values:
    VelDiff_scaled=zeros(numvertices2d,1);
    VelDiff_scaled(VelDiff<0)=-log(max(1,-VelDiff(VelDiff<0)/veldifflims(1)))/log(veldifflims(2)/veldifflims(1));
    VelDiff_scaled(VelDiff>0)=log(max(1,VelDiff(VelDiff>0)/veldifflims(1)))/log(veldifflims(2)/veldifflims(1));
    VelDiff_scaled=VelDiff_scaled.*(GoodData_obs_vertices.^2);
    % Draw main figure:
    patch('Vertices',[x2d/1000,y2d/1000],'Faces',elements2d(IceAreaFraction_elements>.5&OceanAreaFraction_elements<.5,:),'EdgeColor','none','FaceColor','interp','FaceVertexCData',VelDiff_scaled)

    % Set up axes:
    xlim(xlims/1000)
    ylim(ylims/1000)
    set(gca,'XTick',xtick*[ceil(xlims(1)/(1e3*xtick)):1:floor(xlims(2)/(1e3*xtick))])
    set(gca,'YTick',ytick*[ceil(ylims(1)/(1e3*ytick)):1:floor(ylims(2)/(1e3*ytick))])
    if thismodel==1
        ylabel('y (km)','FontSize',axesfontsize)
    else
        set(gca,'YTickLabel',[])
    end
    xlabel('x (km)','FontSize',axesfontsize)
    set(gca,'PlotBoxAspectRatio',[(xlims(2)-xlims(1))/(ylims(2)-ylims(1)),1,1])
    set(gca,'TickDir','out')
    set(gca,'FontSize',axesfontsize,'LineWidth',axeslinewidth)
    %set(gca,'PositionConstraint','innerposition')
    % Set up color scale:
    colormap(gca,velcmap) %FadedJet
    caxis(gca,[-1,1])
    if thismodel==3
        hc=colorbar;
        set(get(hc,'Ylabel'),'String','m/yr (v_{s}-v^{obs})','FontSize',colorbarfontsize)
        set(hc,'YTick',veldiffticks)
        set(hc,'YTickLabel',veldiffticklabels)
        set(hc,'TickDir','out')
        set(hc,'FontSize',colorbarfontsize,'LineWidth',axeslinewidth)
    end
    box on 
    % Make title:
    title(['\rm (',char(96+6+thismodel),') Velocity error for ',lambdadescrip,'=',num2str(num2str(thisbase*.01*round(100*Lambda(thismodelind)/thisbase)))],'FontSize',titlefontsize)
end   

%% SAVE FIGURE
saveas(FigH,figname,'png');
exportgraphics(FigH,figname,'resolution',resolution)

% Final display:
disp('Done!')
toc
