% Attenuation Compensation Using Time Variant Filtering Example
%
% This example demonstrates how the acoustic attenuation present in the
% photoacoustic forward problem can be compensated for using time-variant
% filtering. It builds on the 2D Time Reversal Reconstruction For A
% Circular Sensor and Attenuation Compensation Using Time Reversal
% examples.
%
% For a more detailed discussion of this example and the underlying
% techniques, see B. E. Treeby "Acoustic attenuation compensation in
% photoacoustic tomography using time-variant filtering," J. Biomed. Opt.,
% vol. 18, no. 3, p.036008, 2013.
%
% author: Bradley Treeby
% date: 14th August 2014
% last update: 20th August 2014
%  
% This function is part of the k-Wave Toolbox (http://www.k-wave.org)
% Copyright (C) 2009-2014 Bradley Treeby and Ben Cox

% This file is part of k-Wave. k-Wave is free software: you can
% redistribute it and/or modify it under the terms of the GNU Lesser
% General Public License as published by the Free Software Foundation,
% either version 3 of the License, or (at your option) any later version.
% 
% k-Wave is distributed in the hope that it will be useful, but WITHOUT ANY
% WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
% FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for
% more details. 
% 
% You should have received a copy of the GNU Lesser General Public License
% along with k-Wave. If not, see <http://www.gnu.org/licenses/>. 

clear all;

% set to true to enable use of the parallel computing toolbox
USE_PARALLEL_COMPUTING_TOOLBOX = false;

% =========================================================================
% FORWARD SIMULATION
% =========================================================================

% define the size of the simulation grid and the PML
simulation_size = 512;              % [grid points]
PML_size = 20;                      % [grid points]
x = 52e-3;                          % [m]
y = x;                              % [m]

% reduce the number of grid points in Nx and Ny by the size of the PML so
% that using 'PMLInside' set to false will still give the correct
% simulation size 
Nx = simulation_size - 2*PML_size;  % [grid points]    
Ny = Nx;                            % [grid points]
dx = x/Nx;                          % [m]
dy = dx;                            % [m]

% create the computational grid
kgrid = makeGrid(Nx, dx, Ny, dy);

% define the properties of the propagation medium
medium_non_absorbing.sound_speed = 1510;	% [m/s]

% create a duplicate of the propagation medium structure and append the
% absorption properties
medium = medium_non_absorbing;
medium.alpha_power = 1.5;      
medium.alpha_coeff = 3;              % [dB/(MHz^y cm)]

% store maximum supported frequency
f_max = kgrid.k_max*medium.sound_speed/(2*pi);

% load the shepp logan phantom (note if the simulation or PML sizes are
% changed, the loaded data will need to be resized)
load EXAMPLE_shepp_logan

% smooth the phantom and assign it to the initial pressure
shepp_logan = smooth(kgrid, shepp_logan, true);
source.p0 = shepp_logan;

% define a circular Cartesian sensor mask
sensor_radius = 25e-3;              % [m]
sensor_points = 200;
cart_sensor_mask = makeCartCircle(sensor_radius, sensor_points);
sensor.mask = cart_sensor_mask;

% create the time array used for the simulation, with t_max defined using
% Huygens' principle to avoid artifact trapping in the reconstruction
t_max = 2*sensor_radius/medium.sound_speed;
[kgrid.t_array, dt] = makeTime(kgrid, medium.sound_speed, [], t_max);

% set the input options, switching off the smoothing (the input has already
% been smoothed), setting the PML to be outside the defined grid, casting
% to 'single' to speed up the example, and switching off visualisation
input_args = {'Smooth', false, 'PMLInside', false, 'PMLSize', PML_size, ...
    'DataCast', 'single', 'PlotSim', true};

% run the forward simulation
sensor_data_lossy = kspaceFirstOrder2D(kgrid, medium, source, sensor, input_args{:});

% add noise to the recorded sensor data
signal_to_noise_ratio = 40;         % [dB]
sensor_data_lossy = addNoise(sensor_data_lossy, signal_to_noise_ratio, 'peak');

% =========================================================================
% IMAGE RECONSTRUCTION WITHOUT ATTENUATION COMPENSATION
% =========================================================================

% create a second computation grid for the reconstruction to avoid the
% inverse crime
PML_size = 25;                      % [grid points]
Nx = simulation_size - 2*PML_size;  % [grid points]
Ny = Nx;                            % [grid points]
dx = x/Nx;                          % [m]
dy = dx;                            % [m]
kgrid_recon = makeGrid(Nx, dx, Ny, dy);

% attach the original time array
kgrid_recon.t_array = kgrid.t_array;

% remove the initial pressure field from the source structure
source = rmfield(source, 'p0');

% create a continuous binary sensor mask with the same radius as the
% Cartesian sensor mask used in the forward simulation
pixel_radius = round(sensor_radius/kgrid_recon.dx);
binary_sensor_mask = makeCircle(kgrid_recon.Nx, kgrid_recon.Ny, floor(kgrid_recon.Nx/2) + 1, floor(kgrid_recon.Ny/2) + 1, pixel_radius);

% assign the sensor mask to the sensor structure
sensor.mask = binary_sensor_mask;

% interpolate the simulated sensor data onto a continuous binary sensor
% mask to remove any gaps and assign to the time reversal field
sensor.time_reversal_boundary_data = interpCartData(kgrid_recon, sensor_data_lossy, cart_sensor_mask, binary_sensor_mask, 'linear');

% re-assign the input options (the PML_size has changed)
input_args = {'Smooth', false, 'PMLInside', false, 'PMLSize', PML_size,...
    'DataCast', 'single', 'PlotSim', true};

% run the time-reversal reconstruction using the non absorbing medium
p0_recon = kspaceFirstOrder2D(kgrid_recon, medium_non_absorbing, source, sensor, input_args{:});

% =========================================================================
% IMAGE RECONSTRUCTION WITH ATTENUATION COMPENSATION - FIXED
% =========================================================================

% correct for acoustic attenuation using time variant filtering regularised
% by a Tukey window with a fixed cutoff frequency
sensor_data_comp = attenComp(sensor_data_lossy, kgrid.dt, medium.sound_speed,...
    medium.alpha_coeff, medium.alpha_power, 'FilterCutoff', [3e6 3e6]);

% interpolate the simulated sensor data onto a continuous binary sensor
% mask to remove any gaps and assign to the time reversal field
sensor.time_reversal_boundary_data = interpCartData(kgrid_recon, sensor_data_comp, cart_sensor_mask, binary_sensor_mask, 'linear');
            
% run the time-reversal reconstruction
p0_recon_comp_fixed = kspaceFirstOrder2D(kgrid_recon, medium_non_absorbing, source, sensor, input_args{:});

% =========================================================================
% IMAGE RECONSTRUCTION WITH ATTENUATION COMPENSATION - AVERAGE
% =========================================================================

% correct for acoustic attenuation using time variant filtering regularised
% by a Tukey window with a time-variant cutoff frequency based on the
% average time-frequency distribution of the signals
sensor_data_comp = attenComp(sensor_data_lossy, kgrid.dt, ...
                medium.sound_speed, medium.alpha_coeff, medium.alpha_power, ...
                'FrequencyMultiplier', 3, 'Plot', true);

% interpolate the simulated sensor data onto a continuous binary sensor
% mask to remove any gaps and assign to the time reversal field
sensor.time_reversal_boundary_data = interpCartData(kgrid_recon, sensor_data_comp, cart_sensor_mask, binary_sensor_mask, 'linear');

% run the time-reversal reconstruction
p0_recon_comp_avg = kspaceFirstOrder2D(kgrid_recon, medium_non_absorbing, source, sensor, input_args{:});

% =========================================================================
% IMAGE RECONSTRUCTION WITH ATTENUATION COMPENSATION - INDIVIDUAL
% =========================================================================

% correct for acoustic attenuation using time variant filtering regularised
% by a Tukey window with a time-variant cutoff frequency based on the
% time-frequency distribution for each signal
sensor_data_comp = zeros(size(sensor_data_lossy));

if ~USE_PARALLEL_COMPUTING_TOOLBOX
    
    % update command line
    disp('Appying Time Variant Filter... ');
    
    % loop through signals
    for index = 1:size(sensor_data_lossy, 1)

        % correct signal
        sensor_data_comp(index, :) = attenComp(sensor_data_lossy(index, :), kgrid.dt,...
            medium.sound_speed, medium.alpha_coeff, medium.alpha_power, ...
            'FrequencyMultiplier', 3);

        % display signals
        if index == 1
            figure;
        end
        plot(kgrid.t_array*1e6, sensor_data_lossy(index, :), 'k-', kgrid.t_array*1e6, sensor_data_comp(index, :), 'r-');
        title(['Signal ' num2str(index) ' of ' num2str(sensor_points)]);
        drawnow;

    end
    
    % update command line
    disp(['  completed in ' scaleTime(toc)]);

else

    % update command line
    disp('Appying Time Variant Filter... ');    
    
    % start matlab pool
    parpool;
    
    % loop through signals using parfor
    parfor index = 1:size(sensor_data_lossy, 1)
        
        % correct signal
        sensor_data_comp(index, :) = attenComp(sensor_data_lossy(index, :), kgrid.dt,...
            medium.sound_speed, medium.alpha_coeff, medium.alpha_power, ...
            'FrequencyMultiplier', 3, 'DisplayUpdates', false); %#ok<PFBNS>
        
    end
    
    % delete matlab pool
    delete(gcp('nocreate'));
    
    % update command line
    disp(['  completed in ' scaleTime(toc)]);
    
end

% interpolate the simulated sensor data onto a continuous binary sensor
% mask to remove any gaps and assign to the time reversal field
sensor.time_reversal_boundary_data = interpCartData(kgrid_recon, sensor_data_comp, cart_sensor_mask, binary_sensor_mask, 'linear');

% run the time-reversal reconstruction
p0_recon_comp_tvf = kspaceFirstOrder2D(kgrid_recon, medium_non_absorbing, source, sensor, input_args{:});

% =========================================================================
% VISUALISATION
% =========================================================================

% set the pressure outside the sensor mask to be zero
binary_sensor_map = makeDisc(kgrid_recon.Nx, kgrid_recon.Ny, kgrid_recon.Nx/2, kgrid_recon.Ny/2, pixel_radius - 50);
p0_recon(binary_sensor_map ~= 1) = 0;
p0_recon_comp_tvf(binary_sensor_map ~= 1) = 0;
p0_recon_comp_avg(binary_sensor_map ~= 1) = 0;
p0_recon_comp_fixed(binary_sensor_map ~= 1) = 0;

% shrink the data sets
phantom_padding = 50;
shepp_logan = shepp_logan(1+ phantom_padding:end-phantom_padding, 1+ phantom_padding:end-phantom_padding);
p0_recon = p0_recon(1+ phantom_padding:end-phantom_padding, 1+ phantom_padding:end-phantom_padding);
p0_recon_comp_tvf = p0_recon_comp_tvf(1+ phantom_padding:end-phantom_padding, 1+ phantom_padding:end-phantom_padding);
p0_recon_comp_avg = p0_recon_comp_avg(1+ phantom_padding:end-phantom_padding, 1+ phantom_padding:end-phantom_padding);
p0_recon_comp_fixed = p0_recon_comp_fixed(1+ phantom_padding:end-phantom_padding, 1+ phantom_padding:end-phantom_padding);

% create plot axes
x_vec = kgrid_recon.x_vec(1+ phantom_padding:end-phantom_padding)*1e3;
y_vec = kgrid_recon.y_vec(1+ phantom_padding:end-phantom_padding)*1e3;
y_vec_sl = kgrid.y_vec(1+ phantom_padding:end-phantom_padding)*1e3;

% ----

% plot the original shepp logan phantom
figure;
subplot(1, 2, 1)
imagesc(y_vec, x_vec, shepp_logan, [0 1]);
axis image;
colormap(flipud(gray));
ylabel('x-position [mm]');
xlabel('y-position [mm]');
title('Initial Pressure Distribution');
scaleFig(1.2, 0.8);

% ----

% plot the reconstruction with no filter
figure;
subplot(1, 2, 1)
imagesc(y_vec, x_vec, p0_recon, [0 1]);
axis image;
colormap(flipud(gray));
ylabel('x-position [mm]');
xlabel('y-position [mm]');
title('Reconstruction - No Filter');

% plot a profile
subplot(1, 2, 2)
plot(y_vec_sl, shepp_logan(end/2,:), 'k');
hold on;
plot(y_vec, p0_recon(end/2,:), 'r-');
axis square
set(gca, 'XLim', [-20, 20], 'YLim', [-0.1 1.1]);
ylabel('Pressure Magnitude [au]');
xlabel('y-position [mm]');
title('Profile');
scaleFig(1.2, 0.8);
 
% ----

% plot the reconstruction with fixed tvf filter
figure;
subplot(1, 2, 1)
imagesc(y_vec, x_vec, p0_recon_comp_fixed, [0 1]);
axis image;
colormap(flipud(gray));
ylabel('x-position [mm]');
xlabel('y-position [mm]');
title('Reconstruction - Fixed Time Variant Filter');

% plot a profile
subplot(1, 2, 2)
plot(y_vec_sl, shepp_logan(end/2,:), 'k');
hold on;
plot(y_vec, p0_recon_comp_fixed(end/2,:), 'r-');
axis square
set(gca, 'XLim', [-20, 20], 'YLim', [-0.1 1.1]);
ylabel('Pressure Magnitude [au]');
xlabel('y-position [mm]');
title('Profile');
scaleFig(1.2, 0.8);

% ----

% plot the reconstruction with average tvf filter
figure;
subplot(1, 2, 1)
imagesc(y_vec, x_vec, p0_recon_comp_avg, [0 1]);
axis image;
colormap(flipud(gray));
ylabel('x-position [mm]');
xlabel('y-position [mm]');
title('Reconstruction - Average Time Variant Filter');

% plot a profile
subplot(1, 2, 2)
plot(y_vec_sl, shepp_logan(end/2,:), 'k');
hold on;
plot(y_vec, p0_recon_comp_avg(end/2,:), 'r-');
axis square
set(gca, 'XLim', [-20, 20], 'YLim', [-0.1 1.1]);
ylabel('Pressure Magnitude [au]');
xlabel('y-position [mm]');
title('Profile');
scaleFig(1.2, 0.8);

% ----

% plot the reconstruction with tvf filter
figure;
subplot(1, 2, 1)
imagesc(y_vec, x_vec, p0_recon_comp_tvf, [0 1]);
axis image;
colormap(flipud(gray));
ylabel('x-position [mm]');
xlabel('y-position [mm]');
title('Reconstruction - Time Variant Filter');

% plot a profile
subplot(1, 2, 2)
plot(y_vec_sl, shepp_logan(end/2,:), 'k');
hold on;
plot(y_vec, p0_recon_comp_tvf(end/2,:), 'r-');
axis square
set(gca, 'XLim', [-20, 20], 'YLim', [-0.1 1.1]);
ylabel('Pressure Magnitude [au]');
xlabel('y-position [mm]');
title('Profile');
scaleFig(1.2, 0.8);