cw04_DL_Fit.m

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% This script is part of the UCL MedICSS 2022 "Estimation of brain tissue 
% microstructure with dMRI" project. The repository of the project is
% available at: https://github.com/CIG-UCL/MedICSS_2022_microImag
% 
% Author: Michele Guerreri (m.guerreri@ucl.ac.uk)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% course work # 4: Deep learning fit

% The aim of this course work is to learn how to estimate NODDI model 
% parameters using a deep learning based approach.
% We will explore the training set, see how to define and train a
% multi-layer perceptron (MLP) network and test it on a test set.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% How to use this code?
% 1. You should read the comments in the script.
% 2. There are some questions through the script, write your answer into the
%    dedicated space. Don't worry if you don't know the answer, you are
%    here to learn!!
% 3. The code is divided in sections. Every time you start a new section you
%    should uncomment the code and complete the missing parts (which are
%    highlighted). Then run the section code and go to the next section.
% 4. Repeat step 3. until the end of the script.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%% 4.0 Cleaning and loading

%{

% First thing let's clean our workspace, clear the command window and close
% all open figures:

clc
clear
close all

% Start from the right directory
toCourseWorkDIR()

% load useful variables previously created
load(fullfile('coursework_outputs', 'cw1.mat'), 'gccVox')
load(fullfile('coursework_outputs', 'cw2.mat'), 'noddi', 'acqProtocol')
load(fullfile('coursework_outputs', 'cw3.mat'), 'gccVox_norm', 'noddiFit_gcc', 'noddiLoglik_gcc');

%}

%% 4.1 Training set

%{

% Remember, the problem of deep learning based parameter estimation is:
% to learn a function (1) that can predict the parameters of interest (2) 
% from the data observed (3), given a set of observed data and parameters 
% of interest pairs (4) exampling the relation between the two.

% So, unlike the previous conventional approach with DL we need a set of
% parameters and connected signal exampling the relation between the two.

% Let's creat this "training" dataset.

% We will create N=10e4 voxels randomly, extracting the value of each 
% parameter separately. We will use specific density dist functions. 

% Once we have the parameter combintations we will generate the signal
% using NODDI model and the acquisition protocol we used so far.

% use the function "generateModelTrainingSet" to do so. Let's also use a
% random seed so that we will  all have the same "random" parameters
% extracted.

N = 10e4;
seed1 = 123;
[modelParam, modelSignal] = generateModelTrainingSet(); % complete this

% Let's explore the parameter distribution. Use the plotDist() function to see
% the generated parameter distributions

% first let's convert the kappa values into odi values
modelParam_odi = kappa2odi(); % complete this

% chech the function kappa2odi to know the relation between the odi and
% kappa. If kappa is defined [0 +inf] what is the odi range?
% ANSWER: 

plotDist(); % complete this


% Describe the distributions characterizing each parameter. Do they make
% sense to you?
% ANSWER:


% you can give a try with modelParam variable and check what distribution
% of kappa corresponds to a uniform distribution of odi

plotDist(); % complete this

% plot the signal of the first voxel to check that everything is all right. Use b0=1.
fig_train_compact = figure('Position', [400 100 500 800], 'color', [ 1 1 1]);
fibredir = GetFibreOrientation(); % complete this
VoxelDataViewer(); % complete this

% let's now add some noise to the data to make them look as if they were
% data acquired experimentally.
% We use the "addNoise2Signal" function.

% We use a signal to noise ratio (SNR) of 30 which is similar to the SNR
% usually found in invivo data.
snr = 30;

% again, we set a random seed
seed2 = 234;
 
modelSignal_noisy = addNoise2Signal(); % complete this

% OPTIONAL: have a look at the addNoise2Signal function. How is noise
% generated? What type of distribution is that? Any clue why?
% ANSWER: 

% Let's plot the first voxel noisy signal
fig_noisy_compact = figure('Position', [400 100 500 800], 'color', [ 1 1 1]);
VoxelDataViewer(); % complete this

% Compare it with the GCC voxel signal plot. Do they look similar?
load(fullfile('coursework_outputs', 'cw1.mat'), 'fig_gcc_compact');

%}

%% 4.2 DL model

%{

% Once we obtained the training set let's define the DL network, i.e. the
% function that we will use to predict the parameters given the dMRI
% signal.

% We will use the function "buildDL_function()"
% To define the network we need to specify the input size and the output
% size.

% The particular model we are going to build uses sigle voxel diffusion MRI
% data to predict the non-orientation dependent NODDI parameters of that same voxel.
% (we are not going to estimate phi and theta)

% What is the input of the network? What is the input size?
% ANSWER: 

% What is the output to the network? What is the output size?
% ANSWER:

numOfFeatures = ; % complete this
numOfOutputs = ; % 2 are theta and phi 

% build the network
layers = buildDL_function(numOfFeatures, numOfOutputs); % complete this

% use the following command to visualize the network architecture
analyzeNetwork(layers);

% How many layers are there? What type of layers? How many learnable
% parameters? Are you able to deduce the number of parameters given the
% architecture?
% 

%}

%% 4.3 Network training

%{

% the next step is to update the function (network) we just defined so that 
% we can learn the mapping between input data and output parameters.

% we use the training set we dfefined to do so. 
% First we need to divide the training set into training and validation sets.
% The validation set is useful to understand if we are overfitting the data.

% We use trainValSplit() function in a hold out modality. We use 20% of the
% data for validation. Again, we fix the random seed.
hPercentage = 0.2;
seed3 = 345;
% the train params are only the three we are willing to estimate
trainParam = ; % complete this
[xTrain, yTrain, xValidation, yValidation] =   trainValSplit(); % complete this

% Dose the output make sense to you?

% Next we need to define the training options:
mlpOptions = trainingOptions('adam', ...
    'MaxEpochs',200, ...
    'MiniBatchSize',2500, ...
    'InitialLearnRate',1e-2, ...
    'ValidationPatience', 10, ...
    'ValidationData', {xValidation, yValidation}, ...
    'Plots','training-progress');

% Do you know what each field refers to?

% Finally, we run the network training. For time reason let it run for a
% couple of minutes and stop it with the button at the top right corner of
% the training GUI.
rng(456);
net = trainNetwork(); % complete this

% I completed the training for you.
load('..\..\Data\DL_paramEstimation\trained_network.mat', 'net')

% Have a look at the training progress window. Try to understand what each
% line means. What the iterations are? What the epochs are?

%}

%% 4.4 Training assessment

%{

% It seems the training is doing something interesting, but how can we
% judge the quality of the training?
% ANSWER: We can have a look at the performance on the validation set to start.

% We can input the signal of the validation set into the network and
% compare the predicted parameters with the ground truth ones:
valPrediction = predict(net, xValidation); % complete this

% Let's compute the overall RMSE...
val_rmse = ; % complete this

% What value do you get?
% 

% how does the validation RMSE looks compared to the one obtained for the
% synthetic voxel in course work 3? Does it make sense two compare the two?
% ANSWER: 

% Also we can have a look at the plots of ground truth vs predicted value:
fig_netPredictionQC = figure('Position', [400 100 1300 400], 'color', [ 1 1 1]);
subplot(1,3,1), plot_gt_vs_prediction(, 'NDI (ficvf)'); % complete this
subplot(1,3,2), plot_gt_vs_prediction(, 'ODI'); % complete this
subplot(1,3,3), plot_gt_vs_prediction(, 'FWF (fiso)'); % complete this

% How do the plots look?
% ANSWER: 

figure,plot_gt_vs_prediction(, 'NDI VS FWF'); % complete this
figure,plot_gt_vs_prediction(, 'ODI VS FWF'); % complete this
figure,plot_gt_vs_prediction(, 'ODI VS NDI'); % complete this

%}

%% 4.5 Extend to an unseen voxel

%{

% We can now try to extend our analysis to unseen data. We will use the
% GCC voxel again.

% Try ro predict the parameters via the network using the DW GCCsignal as
% input. Be mindful to the signal you input...
gccParPrediction = predict(net, gccVox_norm'); % complete this

% This time we don't have the parameter ground truth. How can we assess the
% quality of the estimation?
% ANSWER:

% Let's use the parameters estimated by the network to estimate the signal.
% we use the theta and phi from the conventional fit approach.
noddiNet_gcc_odi = noddiFit_gcc;
noddiNet_gcc_odi([1 3 4]) = ; % complete this
noddiNet_gcc = odi2kappa(); % complete this

% compute the loglikelihood and get the fit quality plot. Compare both
% results with those obtained via the conventional fit approach.
loglik_gccnet = getFitRicianLogLik(); % complete this
fig_noddiQC_gccnet = QualityOfFit(); % complete this

load(fullfile('coursework_outputs', 'cw3.mat'), 'fig_noddiQC_gcc');

% ANSWER: 
%

% We showed that our DL network, trained on synthetic data can explain
% in-vivo data quite well.

%}

%% 4.6 Extend the analysis to in-vivo data

%{

% In the last step of our analysis we are going to deploy our DL model on a
% set of invivo data (which include our GCC voxel).

% We are going to use the data from one single slide.

% Let's load the diffusion data
dwi_path = fullfile('..','..','NODDI_example_dataset','output','NODDI_roi.mat');
load(dwi_path, 'roi', 'idx', 'mask');
% we are going to compare the results with those obtaind using the
% conventional fit approach via the NODDI toolbox
fittedParams_path = fullfile('..','..','NODDI_example_dataset','output','FittedParams.mat');
load(fittedParams_path, 'mlps', 'fobj_ml')

% you need first to normalize the data
roi_norm = ; % complete this

% Let's use the network to predict the parameters
tic
fittedParams_dl = predict(); % complete this
dl_fitTime = toc;

% how long did it take? We fitted ~5500 voxels. Try to work out the time
% we would need to do the same using the conventional fit approach based on
% the results of the previous course work.

% We can now have a look at the parameter maps
% we need to rearrange the voxel back in the right spatial order
noddiMaps_fit = zeros(size(mask,1), size(mask,2), numOfOutputs);
noddiMaps_dl = zeros(size(mask,1), size(mask,2), numOfOutputs);
% let's convert the fitted params from kappa values to odi
mlps_odi = kappa2odi(mlps, noddi.name);

parIdx = [1 3 4];
% Loop over all the voxels and the parameters
for ii = 1:size(idx,1)
    for jj = 1:numOfOutputs
        noddiMaps_fit(idx(ii,1),idx(ii,2),jj) = mlps_odi(ii,parIdx(jj));
        noddiMaps_dl(idx(ii,1),idx(ii,2),jj) = fittedParams_dl(ii,jj);
    end
end

% define the difference map
map_diff = noddiMaps_fit-noddiMaps_dl;

% put together all the maps in one figure
fig_outMaps = figure('Position', [400 100 1000   750], 'color', [ 0 0 0]);
for jj = 1:numOfOutputs
    subplot(3,3,jj);
    imshow( flipud(noddiMaps_fit(:,:,jj)'), [0 1], Colormap=gray),colorbar('color',[1 1 1]);
    subplot(3,3,numOfOutputs+jj);
    imshow( flipud(map_diff(:,:,jj)'), [-0.2 0.2], Colormap=turbo),colorbar('color',[1 1 1]);
    subplot(3,3,(2*numOfOutputs+jj));
    imshow( flipud(noddiMaps_dl(:,:,jj)'), [0 1], Colormap=gray),colorbar('color',[1 1 1]);
end

% To correctly compare the two approaches we should compute the
% loglikelihood for each voxel and used that fo comparison...
% Do it if you have time :D !

%}

%% 4.7 Save the output

%{

% Before running this section make sure you have finished your analysis as
% this will close all the open figures

% let's now save the variables we created into a mat fi
save(fullfile('coursework_outputs', 'cw4.mat'))

% and close the figures
close all

%}