plot_results_colloc_vs_ivp.py

from inner_optimisation_scipy import *
import inner_optimisation_scipy
# from generate_data import *
from pints_classes import *
import pandas as pd
import pickle as pkl
# get the colours out
import matplotlib.colors as mcolors
import seaborn as sns
import os
from tqdm import tqdm
from matplotlib.colors import LogNorm, Normalize
from matplotlib.ticker import MaxNLocator
from load_protocols import *


# definitions
def V(t):
    return volts_interpolated_for_plots(t/ 1000)


SolveInner = False
Validation = True
PlotParams = False
PlotCosts = False
# set up variables for the simulation
tlim = [300, 14899]
times = np.linspace(*tlim, tlim[-1] - tlim[0], endpoint=False)
volts_interpolated_for_plots = volts_interpolated
voltage = V(times)  # must read voltage at the correct times to match the output
del tlim
model_name = 'Wang' # this is the generative model name, can be HH or Kemp
snr_db = 20 # signal to noise ratio in dB
## set up the parameters for the fitted model
fitted_model = hh_model
Thetas_ODE = thetas_hh_baseline
state_names = ['a', 'r'] # how many states we have in the model that we are fitting
# outer optimisation settings
inLogScale = True  # is the search of thetas in log scale
# set the exponents
lambda_exps = [0,1,2,3,4,5,6,7,8]  # gradient matching weight - test
# lambda_exps = [8,7,6,5,4]  # gradient matching weight - test
nIters = 500
# get the colours for each lambda to plot all results in one figure
# colours = list(mcolors.TABLEAU_COLORS.values())
# colours = colours + colours
# colours = colours[:len(lambda_exps)+1]
colours = plt.cm.PuOr(np.linspace(0, 1, len(lambda_exps)))
####################################################################################################################
### from this point no user changes are required
####################################################################################################################
# load the protocols
load_protocols
voltage = V(times)
# generate the segments with B-spline knots and intialise the betas for  - we won't actully need knots for the optimisation
jump_indeces, times_roi, voltage_roi, knots_roi, collocation_roi, spline_order = generate_knots(times)
jumps_odd = jump_indeces[0::2]
jumps_even = jump_indeces[1::2]
nSegments = len(jump_indeces[:-1])
print('The time axis is split into ' + str(nSegments) + ' segments based on protocol steps.')
nBsplineCoeffs = (len(knots_roi[0]) - spline_order - 1) * len(state_names)
init_betas_roi = nSegments * [0.5 * np.ones(nBsplineCoeffs)]
print('Number of B-spline coeffs per segment: ' + str(nBsplineCoeffs) +'.')

# generate a solution here - we dont need it, just to resolve dependency in data splitting for now
if model_name.lower() not in available_models:
    raise ValueError(f'Unknown model name: {model_name}. Available models are: {available_models}.')
elif model_name.lower() == 'hh':
    thetas_true = thetas_hh_baseline
elif model_name.lower() == 'kemp':
    thetas_true = thetas_kemp
    nIters = 500
elif model_name.lower() == 'wang':
    thetas_true = thetas_wang
    nIters = 350

if __name__ == '__main__':
    # folderForOutput = 'Test_plot_output'
    folderForOutput = 'Lambda_exploration_results_' + model_name
    if not os.path.exists(folderForOutput):
        os.makedirs(folderForOutput)

    # create objects for all figures
    ## for the states and outputs
    fig, axes = plt.subplot_mosaic([['a)', 'b)'], ['c)', 'd)'], ['e)', 'f)']], layout='constrained', sharex=True, figsize=(15, 6))
    y_labels = ['$I_{collocation}$', '$I_{fitted}$', '$a_{collocation}$', '$a_{fitted}$', '$r_{collocation}$', '$r_{fitted}$']
    for _, ax in axes.items():
        for iSegment, SegmentStart in enumerate(jumps_odd):
            ax.axvspan(times[SegmentStart], times[jumps_even[iSegment]], facecolor='0.2', alpha=0.075)
    # state errors
    fig1, axes1 = plt.subplot_mosaic([['a)', 'b)'], ['c)', 'd)'], ['e)', 'f)']], layout='constrained', sharex=True, figsize=(15, 6))
    y_labels1 = ['$I_{true} - I_{collocation}$', '$I_{true} - I_{fitted}$', 'Gradient error: $da(B) - RHS(B)$', 'Gradient error: $dr(B) - RHS(B)$',
                 'Collocation  error: $a - Phi B_a$', 'Collocation error: $r - Phi B_r$']
    for _, ax in axes1.items():
        for iSegment, SegmentStart in enumerate(jumps_odd):
            ax.axvspan(times[SegmentStart], times[jumps_even[iSegment]], facecolor='0.2', alpha=0.075)
    # create figure with two subplots for validation protocol and modelling error
    fig7, axes7 = plt.subplots(3, 1, figsize=(10, 6))
    y_labels7 = ['Current','Residual', 'RMSE']
    # plot the current generated by the true model on validation protocol
    volts_interpolated_for_plots = volts_interpolated_ap
    voltage_ap = V(times_ap)
    solution, current_model = generate_synthetic_data(model_name, thetas_true, times_ap)
    current_ap_noiseless = current_model(times_ap, solution, thetas_true)
    axes7[0].plot(times_ap, current_ap_noiseless, '-k', label=r'True current', linewidth=2, alpha=0.27)
    ####################################################################################################################
    # plot parameter values in log scale
    fig2, axes2 = plt.subplots(len(Thetas_ODE), 1, figsize=(10,2 * len(Thetas_ODE)), sharex=True)
    # plot parameter values in decimal scale
    fig3, axes3 = plt.subplots(len(Thetas_ODE), 1, figsize=(10,2 * len(Thetas_ODE)), sharex=True)
    # create plot for inner costs:
    fig4, ax4 = plt.subplots(figsize=(10,3))
    ax4.plot(range(nIters), np.zeros(nIters), '--k', label='Truth', linewidth=1.5, alpha=0.27) # for making legend
    # create plot for outer costs:
    fig5, ax5 = plt.subplots(figsize=(10,3))
    # create plot for gradient costs:
    fig6, ax6 = plt.subplots(figsize=(10,3))
    ####################################################################################################################
    # iterate over lambdas
    counter = 0 # for plotting current only once and colours of plots
    # create dataframe with column names: lambda, best inner, best data, best  gradient
    df_best = pd.DataFrame(columns=['lambda', '$L(B \mid \hat{\Theta}^{*}, Y)$', '$L_{data}(\hat{\Theta}^{*}, \hat{B}(\hat{\Theta}^{*}), Y)$', '$L_{ODE}(\hat{\Theta}^{*}, \hat{B}(\hat{\Theta}^{*}), Y) $'])
    for weight in tqdm(lambda_exps):
        lambd = 10 ** weight
        # load the folders
        folderName = 'Results_gen_model_'+model_name+'_lambda_' + str(int(lambd))
        if counter == 0:
            with open(folderName + '/synthetic_data.pkl', 'rb') as f:
                # load the model output
                model_output = pkl.load(f)
            times, voltage, current_true, states_true, thetas_true, knots_roi, snr_db = model_output
            # send stuff into inner optimisation
            inner_optimisation_scipy.voltage = voltage  # send the voltage to the inner optimisation module
            inner_optimisation_scipy.current_true = current_true  # send the current to the inner optimisation module
            # plot true current - only once is fine!!
            axes['a)'].plot(times, current_true, '--k', label=r'True current', linewidth=1.5, alpha=0.27)
            axes['b)'].plot(times, current_true, '--k', label=r'True current', linewidth=1.5, alpha=0.27)
        ####################################################################################################################
        # load the optimisation metrix from the csv file in the folderName directory
        df = pd.read_csv(folderName + '/iterations_both_states.csv')
        # get the number of thetas for the fitted model
        columnNames = df.columns
        nThetas  = sum('Theta_' in s for s in columnNames)
        ####################################################################################################################
        ### remake all the lists for plotting from the dataframe
        # create a list of InnerCosts_all from the dataframe by grouping entries in Inner Cost column by iteration
        InnerCosts_all = [group["Inner Cost"].tolist() for i, group in df.groupby("Iteration")]
        # create a list of OuterCosts_all from the dataframe by grouping entries in Outer Cost column by iteration
        OuterCosts_all = [group["Outer Cost"].tolist() for i, group in df.groupby("Iteration")]
        # create a list of GradCosts_all from the dataframe by grouping entries in Gradient cost column by iteration
        GradCosts_all = [group["Gradient Cost"].tolist() for i, group in df.groupby("Iteration")]
        # create a list of thetas from the dataframe by grouping entries in columns Theta_1 to Theta_8 by iteration
        theta_visited = [group[["Theta_" + str(i) for i in range(nThetas)]].values.tolist() for i, group in
                         df.groupby("Iteration")]
        # get a number of iterations
        nIter = len(df["Iteration"].unique())
        theta_best = []
        f_outer_best = []
        f_inner_best = []
        f_gradient_best = []
        for iIter in range(nIter):
            OuterCosts = OuterCosts_all[iIter]
            InnerCosts = InnerCosts_all[iIter]
            GradCosts = GradCosts_all[iIter]
            index_best = OuterCosts.index(np.nanmin(OuterCosts))
            theta_best.append(theta_visited[iIter][index_best][:])
            f_outer_best.append(OuterCosts[index_best])
            f_inner_best.append(InnerCosts[index_best])
            f_gradient_best.append(GradCosts[index_best])
        theta_best = np.array(theta_best)
        f_outer_best = np.array(f_outer_best)
        f_inner_best = np.array(f_inner_best)
        f_gradient_best = np.array(f_gradient_best)
        ################################################################################################################
        # change the coltage protocol to staircase
        volts_interpolated_for_plots = volts_interpolated
        Thetas_ODE = theta_best[-1]
        param_names = [f'p_{i}' for i in range(1, len(Thetas_ODE) + 1)]
        if SolveInner:
            tic = tm.time()
            # generate solution - we dont need it beyond sending it into data splitting
            solution, current_model = generate_synthetic_data(model_name, thetas_true, times)
            # set variables definitions for inner optimisation module
            inner_optimisation_scipy.nSegments = nSegments
            inner_optimisation_scipy.state_names = state_names
            inner_optimisation_scipy.spline_order = spline_order
            inner_optimisation_scipy.nBsplineCoeffs = nBsplineCoeffs
            inner_optimisation_scipy.fitted_model = fitted_model
            states_roi, states_known_roi, current_roi = split_generated_data_into_segments(solution, current_true,
                                                                                           jump_indeces, times)
            ## simulate the optimised model using B-splines
            # dont forget to convert to decimal scale if the search was done in log scale
            if inLogScale:
                # convert thetas to decimal scale for inner optimisation
                Thetas_ODE = np.exp(Thetas_ODE)
            else:
                Thetas_ODE = Thetas_ODE
            test_output = inner_optimisation(Thetas_ODE, lambd, times_roi, voltage_roi, current_roi, knots_roi,
                                             collocation_roi)
            betas_sample, inner_cost_sample, data_cost_sample, grad_cost_sample, state_fitted_at_sample = test_output
            state_all_segments = np.array(state_fitted_at_sample)
            current_all_segments = observation_direct_input(state_all_segments, voltage, Thetas_ODE)
            # get the derivative and the RHS
            rhs_of_roi = {key: [] for key in state_names}
            deriv_of_roi = {key: [] for key in state_names}
            for iSegment in range(nSegments):
                model_output_fit = simulate_segment(betas_sample[iSegment], times_roi[iSegment], knots_roi[iSegment],
                                                    first_spline_coeffs=None)
                state_at_sample, state_deriv_at_sample, rhs_at_sample = np.split(model_output_fit, 3, axis=1)
                if iSegment == 0:
                    index_start = 0  # from which timepoint to store the states
                else:
                    index_start = 1  # from which timepoint to store the states
                for iState, stateName in enumerate(state_names):
                    deriv_of_roi[stateName] += list(state_deriv_at_sample[index_start:, iState])
                    rhs_of_roi[stateName] += list(rhs_at_sample[index_start:, iState])
            ## end of loop over segments
            ## simulate the model using the best thetas and the ODE model used
            x0_optimised_ODE = state_all_segments[:, 0]
            solution_optimised = sp.integrate.solve_ivp(fitted_model, [0, times[-1]], x0_optimised_ODE, args=[Thetas_ODE],
                                                        dense_output=True, method='LSODA', rtol=1e-8, atol=1e-8)
            states_optimised_ODE = solution_optimised.sol(times)
            current_optimised_ODE = observation_direct_input(states_optimised_ODE, voltage, Thetas_ODE)
            toc = tm.time()
            print('Optimisation solved. Time elapsed: ' + str(toc-tic))
            ####################################################################################################################
            # add states to the first figure- compare the modelled current with the true current
            axes['a)'].plot(times, current_all_segments, '-', color = colours[counter],
                            label=r'$\lambda = $' + "{:.2e}".format(lambd), linewidth=1, alpha=0.5)
            axes['b)'].plot(times, current_optimised_ODE, '-', color = colours[counter],
                            label=r'$\lambda = $' + "{:.2e}".format(lambd), linewidth=1, alpha=0.5)
            # axes['a)'].set_xlim(times_of_segments[0], times_of_segments[-1])
            axes['a)'].set_xlim(1890, 1920)
            axes['c)'].plot(times, state_all_segments[0, :], '-', color = colours[counter],
                            label=r'$\lambda$ = ' +  "{:.2e}".format(lambd),
                            linewidth=1, alpha=0.5)
            axes['e)'].plot(times, state_all_segments[1, :], '-', color = colours[counter],
                            label=r'$\lambda$ = ' +  "{:.2e}".format(lambd),
                            linewidth=1, alpha=0.5)
            axes['d)'].plot(times, states_optimised_ODE[0, :], color = colours[counter],
                            label=r'$\lambda$ = ' +  "{:.2e}".format(lambd),
                            linewidth=1, alpha=0.5)
            axes['f)'].plot(times, states_optimised_ODE[1, :], color = colours[counter],
                            label=r'$\lambda$ = ' +  "{:.2e}".format(lambd),
                            linewidth=1, alpha=0.5)
            # plot errors
            axes1['a)'].plot(times, current_all_segments - current_true, '--', color = colours[counter],
                        label=r'$\lambda=$ ' +  "{:.2e}".format(lambd),linewidth=1, alpha=0.5)
            axes1['b)'].plot(times, current_optimised_ODE - current_true, '--',color = colours[counter],
                             label=r'$\lambda=$' +  "{:.2e}".format(lambd), linewidth=1, alpha=0.5)
            axes1['c)'].plot(times, np.array(rhs_of_roi[state_names[0]]) - np.array(deriv_of_roi[state_names[0]]),
                             '--', color = colours[counter],label=r'$\lambda$ = ' +  "{:.2e}".format(lambd),
                             linewidth=1,alpha=0.5)
            axes1['d)'].plot(times, np.array(rhs_of_roi[state_names[1]]) - np.array(deriv_of_roi[state_names[1]]),
                             '--', color = colours[counter],label=r'$\lambda$ = ' +  "{:.2e}".format(lambd),
                             linewidth=1, alpha=0.5)
            axes1['e)'].plot(times, state_all_segments[0, :] - states_optimised_ODE[0, :], '--',color = colours[counter],
                             label=r'$\lambda$ = ' +  "{:.2e}".format(lambd), linewidth=1, alpha=0.5)
            axes1['f)'].plot(times, state_all_segments[1, :] - states_optimised_ODE[1, :], '--',color = colours[counter],
                             label=r'$\lambda$ = ' +  "{:.2e}".format(lambd), linewidth=1, alpha=0.5)
            #############################################################################################################
        if Validation:
            # compute model with validation protocol
            volts_interpolated_for_plots = volts_interpolated_ap
            voltage_ap = V(times_ap)
            Thetas_ODE = theta_best[-1]
            if inLogScale:
                # convert thetas to decimal scale for inner optimisation
                Thetas_ODE = np.exp(Thetas_ODE)
            else:
                Thetas_ODE = Thetas_ODE
            fitted_model_name = 'hh'
            solution, current_model = generate_synthetic_data(fitted_model_name, Thetas_ODE, times_ap)
            states_optimised_ODE = solution.sol(times_ap)
            current_ODE_output = current_model(times_ap, solution, Thetas_ODE)
            current_residual = current_ODE_output - current_ap_noiseless
            current_rmse = np.sqrt(np.mean(current_residual** 2))
            # plot the current generated by the fitted model on validation protocol
            axes7[0].plot(times_ap, current_ODE_output, '--', color=colours[counter],
                          label=r'$\lambda = 10^{' + str(weight) + '}$',
                          linewidth=1, alpha=0.5)
            axes7[1].plot(times_ap, current_residual, '--', color=colours[counter],
                          linewidth=1, alpha=0.5)
            # add the RMSE to the plot of lambda vs RMSE
            # axes7[0].text(times_ap[-1], current_ODE_output[-1], "{:.2f}".format(current_rmse), fontsize=8)
            axes7[-1].scatter(weight, current_rmse, color=colours[counter], edgecolors='k', marker='o', alpha=0.7)
            ############################################################################################################
        # end of if condition
        ################################################################################################################
        axes3 = axes3.flatten()
        if PlotParams:
            # plot parameter values after search was done on decimal scale
            for iAx, ax in enumerate(axes2.flatten()):
                for iIter in range(len(theta_visited)):
                    x_visited_iter = [theta_visited[iIter][i][iAx] for i in range(len(theta_visited[iIter]))]
                    ax.scatter(iIter * np.ones(len(x_visited_iter)), x_visited_iter, color=colours[counter], marker='.', alpha=0.075,
                               linewidth=0)
                    axes3[iAx].scatter(iIter * np.ones(len(x_visited_iter)), np.exp(x_visited_iter), color=colours[counter],
                               marker='.', alpha=0.075,linewidth=0)
                ax.plot(theta_best[:, iAx], '-', color = colours[counter], linewidth=1,
                        label=r'$\lambda$ = ' +  "{:.2e}".format(lambd) +r", $\log(" + param_names[iAx] + ") = $ " + "{:.6f}".format(theta_best[-1, iAx]))
                axes3[iAx].plot(np.exp(theta_best[:, iAx]), '-', color=colours[counter], linewidth=1,
                        label=r'$\lambda$ = ' + "{:.2e}".format(lambd) + r", $" + param_names[iAx] + " = $" + "{:.6f}".format(np.exp(theta_best[-1, iAx])))
                ax.text(nIters, theta_best[-1, iAx], "{:.6f}".format(theta_best[-1, iAx]), fontsize=8)
                axes3[iAx].text(nIters, np.exp(theta_best[-1, iAx]), "{:.6f}".format(np.exp(theta_best[-1, iAx])), fontsize=8)
        # ####################################################################################################################
        if PlotCosts:
            # fill in a row of the dataframe
            df_best.loc[counter] = [r'$10^{' + str(weight)+ '}$', f_inner_best[-1], f_outer_best[-1], f_gradient_best[-1]]
            ## plot evolution of inner costs
            for iIter in range(len(f_outer_best)):
                ax4.scatter(iIter * np.ones(len(InnerCosts_all[iIter])), InnerCosts_all[iIter], color=colours[counter],
                            marker='.', alpha=0.075, linewidths=0)
            ax4.plot(f_inner_best, '-', color=colours[counter], linewidth=1.5,
                     label=r'$\lambda = 10^{' + str(weight)+ '}$')
            #
            # plot evolution of outer costs
            for iIter in range(len(f_outer_best)):
                ax5.scatter(iIter * np.ones(len(OuterCosts_all[iIter])), OuterCosts_all[iIter], color=colours[counter], marker='.',
                            alpha=0.075, linewidths=0)
            ax5.plot(f_outer_best, '-', color=colours[counter], linewidth=1,
                     label=r'$\lambda = 10^{' + str(weight)+ '}$')
            #
            # plot evolution of gradient costs
            for iIter in range(len(f_gradient_best)):
                ax6.scatter(iIter * np.ones(len(GradCosts_all[iIter])), GradCosts_all[iIter], color=colours[counter], marker='.',
                            alpha=0.075,linewidths=0)
            ax6.plot(f_gradient_best, '-',  color=colours[counter], linewidth=1,
                     label=r'$\lambda = 10^{' + str(weight)+ '}$')
        ################################################################################################################
        # done plotting
        counter += 1
    ## end of loop over lambda values
    ####################################################################################################################
    ## save the figures
    # states plot
    axes['a)'].set_ylim(-4, 4)
    axes['b)'].set_ylim(-4, 4)
    if SolveInner:
        iAx = 0
        for _, ax in axes.items():
            ax.set_ylabel(y_labels[iAx], fontsize=12)
            ax.set_facecolor('white')
            ax.grid(which='major', color='grey', linestyle='solid', alpha=0.2, linewidth=1)
            # if iAx == 0:
                # ax.legend(fontsize=12, loc='best',ncol=3)
            iAx += 1
        fig.tight_layout(pad=0.3)
        fig.savefig(folderForOutput + '/states_model_output.png', dpi=400)
        # save the error plot
        iAx = 0
        for _, ax in axes1.items():
            ax.set_ylabel(y_labels1[iAx], fontsize=12)
            ax.set_facecolor('white')
            ax.grid(which='major', color='grey', linestyle='solid', alpha=0.2, linewidth=1)
            if iAx == 0:
                ax.legend(fontsize=12, loc='best',ncol=3)
            iAx += 1
        fig1.tight_layout(pad=0.3)
        fig1.savefig(folderForOutput + '/errors_model_output.png', dpi=400)

    if Validation:
        # add the result from the model obtained by fitting model with IVP solver
        # load the model from the IVP folder
        folderName = 'Results_gen_model_' + model_name + '_IVP'
        # load the optimisation metrix from the csv file in the folderName directory
        df = pd.read_csv(folderName + '/iterations_both_states.csv')
        # get the number of thetas for the fitted model
        columnNames = df.columns
        nThetas  = sum('Theta_' in s for s in columnNames)
        ####################################################################################################################
        ### remake all the lists for plotting from the dataframe
        # create a list of OuterCosts_all from the dataframe by grouping entries in Outer Cost column by iteration
        OuterCosts_all = [group["Outer Cost"].tolist() for i, group in df.groupby("Iteration")]
        # create a list of thetas from the dataframe by grouping entries in columns Theta_1 to Theta_8 by iteration
        theta_visited = [group[["Theta_" + str(i) for i in range(nThetas)]].values.tolist() for i, group in
                         df.groupby("Iteration")]
        # get a number of iterations
        nIter = len(df["Iteration"].unique())
        theta_best = []
        f_outer_best = []
        for iIter in range(nIter):
            OuterCosts = OuterCosts_all[iIter]
            index_best = OuterCosts.index(np.nanmin(OuterCosts))
            theta_best.append(theta_visited[iIter][index_best][:])
            f_outer_best.append(OuterCosts[index_best])
        theta_best = np.array(theta_best)
        # compute model with validation protocol
        volts_interpolated_for_plots = volts_interpolated_ap
        voltage_ap = V(times_ap)
        Thetas_ODE = theta_best[-1]
        if inLogScale:
            # convert thetas to decimal scale for inner optimisation
            Thetas_ODE = np.exp(Thetas_ODE)
        else:
            Thetas_ODE = Thetas_ODE
        fitted_model_name = 'hh'
        solution, current_model = generate_synthetic_data(fitted_model_name, Thetas_ODE, times_ap)
        states_optimised_ODE = solution.sol(times_ap)
        current_ODE_output = current_model(times_ap, solution, Thetas_ODE)
        current_residual = current_ODE_output - current_ap_noiseless
        current_rmse = np.sqrt(np.mean(current_residual ** 2))
        # plot the current generated by the fitted model on validation protocol
        axes7[0].plot(times_ap, current_ODE_output, '--', color='m',
                      label=r'$\lambda = 10^{' + str(weight) + '}$',
                      linewidth=1, alpha=0.5)
        axes7[1].plot(times_ap, current_residual, '--', color='m',
                      linewidth=1, alpha=0.5)
        # add the RMSE as a dashed line to mark the reference value
        axes7[-1].plot([lambda_exps[0], lambda_exps[-1]], [current_rmse,current_rmse], '--', color='m', linewidth=1, alpha=0.5, label='IVP solver')
        ############################################################################################################
        # save the validation protocol and RMSE plot
        iAx = 0
        for ax in axes7:
            ax.set_ylabel(y_labels7[iAx], fontsize=12)
            ax.set_facecolor('white')
            ax.grid(which='major', color='grey', linestyle='solid', alpha=0.2, linewidth=1)
            iAx += 1
        axes7[0].set_ylim(-3, 3)
        axes7[0].set_xlim(times_ap[0], times_ap[-1])
        axes7[1].set_ylim(-1, 1)
        axes7[1].set_xlim(times_ap[0], times_ap[-1])
        axes7[1].set_xlabel('Time [ms]')
        axes7[-1].set_xlabel(r'$\log(\lambda)$')
        axes7[-1].set_ylabel('RMSE')
        axes7[-1].set_yscale('log')
        axes7[-1].legend(fontsize=12, loc='upper right')
        for ax in axes7.flat:
            ax.yaxis.set_label_coords(-0.05, 0.5)
        fig7.tight_layout()
        fig7.savefig(folderForOutput + '/validation_results.png', dpi=400)
    # end of if statement


    # save figure 2 - log scale
    if PlotParams:
        for iAx, ax in enumerate(axes2.flatten()):
            ax.set_ylabel(r'$\log(' + param_names[iAx] + ')$')
            ax.set_facecolor('white')
            ax.grid(which='major', color='grey', linestyle='solid', alpha=0.2, linewidth=1)
            # ax.legend(fontsize=12, loc='best', ncol=2)
            ax.set_xlim(0, nIters)
            iAx += 1
        ax.set_xlabel('Iteration')
        # fig2.tight_layout()
        fig2.savefig(folderForOutput + '/ODE_params_log_scale.png', dpi=400)

        # save figure 3 - decimal scale
        for iAx, ax in enumerate(axes3.flatten()):
            ax.set_ylabel(param_names[iAx])
            ax.set_yscale('log')
            ax.set_facecolor('white')
            ax.grid(which='major', color='grey', linestyle='solid', alpha=0.2, linewidth=1)
            ax.set_xlim(0, nIters)
            # ax.legend(fontsize=12, loc='best', ncol=3)
        ax.set_xlabel('Iteration')
        # fig3.tight_layout()
        fig3.savefig(folderForOutput + '/ODE_params.png', dpi=400)


    if PlotCosts:
        # convert lambda column in the dataframe to exponential scale
        # df_best['lambda'] = df_best['lambda'].apply(lambda x: "{:.2e}".format(x))
        # convert the data frame lambda column to index
        df_best = df_best.set_index('lambda')
        # plot the map of df_best with each column having its own colourbar
        fig, ax = plt.subplots(figsize=(8, 4))
        # display column names at the top for heatmap
        sns.heatmap(df_best, annot=True, fmt=".2e",norm=LogNorm(), cmap='magma', ax=ax,cbar=True,annot_kws={"fontsize":14})
        ax.set(xlabel="", ylabel="")
        ax.xaxis.tick_top()
        ax.set_facecolor('white')
        # ax.set_xlabel('Best cost')
        ax.set_ylabel(r'$\lambda$',fontsize=18)
        fig.tight_layout(pad=0.2)
        fig.savefig(folderForOutput + '/best_costs.png', dpi=400)

        # save results for figure 4
        ax4.set_yscale('log')
        ax4.set_xlabel('Iteration')
        ax4.set_ylabel('Inner cost ' +r'$(L_{data} + \lambda L_{ODE})$')
        ax4.set_facecolor('white')
        ax4.grid(which='major', color='grey', linestyle='solid', alpha=0.2, linewidth=1)
        ax4.set_xlim(0, nIters)
        ax4.legend(loc='upper center', bbox_to_anchor=(0.5, 1.3),ncol=4, fontsize=12,fancybox=True,facecolor='white')
        fig4.tight_layout(pad=0.3)
        fig4.savefig(folderForOutput + '/inner_cost_evolution.png', dpi=400)

        # save results for figure 5
        ax5.set_yscale('log')
        ax5.set_xlabel('Iteration')
        ax5.set_ylabel('Data cost'+r'$(L_{data})$')
        ax5.set_facecolor('white')
        ax5.grid(which='major', color='grey', linestyle='solid', alpha=0.2, linewidth=1)
        ax5.set_xlim(0, nIters)
        # ax5.legend(loc='best',fontsize=12)
        fig5.tight_layout()
        fig5.savefig(folderForOutput + '/outer_cost_evolution.png', dpi=400)

        # save results for figure 6
        ax6.set_yscale('log')
        ax6.set_xlabel('Iteration')
        ax6.set_ylabel('Gradient matching cost '+r'$(L_{ODE})$')
        ax6.set_facecolor('white')
        ax6.grid(which='major', color='grey', linestyle='solid', alpha=0.2, linewidth=1)
        ax6.set_xlim(0, nIters)
        # ax6.legend(loc='best',fontsize=12)
        fig6.tight_layout()
        fig6.savefig(folderForOutput + '/gradient_cost_evolution.png', dpi=400)