doubleSphericalLensPSF.m

function [...
    stats_real, stats_ideal, I, I_color...
] = doubleSphericalLensPSF(...
    lens_params, ray_params, image_params, X_lights, z_film,...
    lights_filter, request_spline_smoothing, varargin...
)
% DOUBLESPHERICALLENSPSF  Generate images of point light sources by raytracing
%
% ## Syntax
% [...
%     stats_real, stats_ideal, I, I_color...
% ] = doubleSphericalLensPSF(...
%     lens_params, ray_params, image_params, X_lights, z_film,...
%     lights_filter, request_spline_smoothing [, depth_factors, verbose]...
% )
%
% ## Description
% [...
%     stats_real, stats_ideal, I, I_color...
% ] = doubleSphericalLensPSF(...
%     lens_params, ray_params, image_params, X_lights, z_film,...
%     lights_filter, request_spline_smoothing [, depth_factors, verbose]...
% )
%   Simulate the image of grids of point light sources by raytracing. (One
%   to four output arguments can be requested.)
%
% ## Input Arguments
%
% lens_params -- Lens parameters structure
%   A description of a lens formed from two spherical surfaces.
%   Passed as a structure with the following fields:
%   - lens_radius: The radius of the lens (i.e. half the height of the lens
%     when viewed edge-on)
%   - axial_thickness: The thickness of the lens along its optical axis.
%   - radius_front: The radius of curvature of the front surface of the
%     lens
%   - radius_back: The radius of curvature of the back surface of the
%     lens
%   - ior_lens: The refractive indices of the lens, one for each wavelength
%     of the light to be simulated. A row vector of length 'k'.
%   - wavelengths: The wavelengths of light corresponding to the elements
%     of `ior_lens`. A row vector of length 'k'. This parameter is used for
%     figure legends only, not for calculations.
%   - wavelengths_to_rgb: RGB colours to be used when displaying the
%     images generated by raytracing with the indices of refraction in
%     `ior_lens`. The i-th row of this k x 3 matrix represents the RGB
%     colour corresponding to the i-th wavelength. `wavelengths_to_rgb`
%     allows colour images to be produced by adding together the
%     contributions of each wavelength to the red, green, and blue colour
%     channels.
%
% ray_params -- Raytracing parameters structure
%   A structure with the following fields:
%   - n_incident_rays: The number of rays to sample, over the front
%     aperture of the lens, for each light source in the scene. Each sample
%     produces one ray from the point light source through the front
%     surface of the lens. The front aperture is uniformly sampled, but
%     samples are culled if they are occluded by the front lens surface,
%     from the perspective of the point light source.
%   - sample_random: Whether to sample rays at random over the front
%     aperture of the lens (true) or in a polar grid pattern (false). In
%     either case, the rays will be sampled uniformly per unit area on the
%     front aperture.
%   - ior_environment: The refractive index of the medium surrounding the
%     lens on both sides
%
% image_params -- Image formation parameters
%   A structure with the following fields, describing how to generate
%   images from raytracing results:
%   - image_sampling: A two-element vector containing the image height,
%     and width, respectively, of the final image (in units of pixels).
%   - image_bounds: A 4-element vector, in the form of the output argument
%     of 'imageBoundaries()'. (Refer to the documentation of
%     'imageBoundaries.m' for details.) Alternatively, this information can
%     be automatically estimated if there is more than one light source in
%     the scene, in which case the field can be empty (`[]`).
%   - normalize_psfs_before_combining: Some point spread functions are very
%     sharp, compared to others. To assist with visualization, the combined
%     image can be generated by normalizing each PSF by its maximum
%     intensity before adding the individual images together. However, the
%     final image will be a little misleading. Setting this parameter to
%     `false` will result in an image formed from raw PSF intensities.
%   - normalize_color_images_globally: Colour images are normalized so that
%     colour channel values are no greater than one. Normalization is
%     either performed by dividing by the maximum channel value over all
%     scene depths (true), or the maximum channel value of each scene depth
%     (false). Therefore, if this parameter is true, colour images produced
%     for different depths are directly comparable, but images for
%     out-of-focus depths may be quite dim.
%
% X_lights -- Light positions
%   The positions, expressed in cartesian coordinates (x, y, z), of the
%   lights in the scene. `X_lights(i, :, j)` is the 3D position of the i-th
%   light in the grid of lights placed at the j-th depth.
%
% z_film -- Image plane location
%   The z-coordinate of the image plane.
%
% lights_filter -- Light gaps filter
%   For efficiency, a full grid of lights producing point spread functions
%   over the image plane may not be desirable. `lights_filter` is a logical
%   vector of length `size(X_lights, 1)` indicating which light positions
%   are to be ignored. For instance, lights close to the optical axis may
%   be ignored, as their images show little chromatic aberration.
%
% request_spline_smoothing -- Smooth image intensities
%   A flag determining whether image intensities are subject to thin-plate
%   spline smoothing. For low numbers of rays, `request_spline_smoothing`
%   can be set to reduce noise, but for high numbers of rays, spline
%   smoothing will be computationally expensive.
%
%   Presently, if spline smoothing is not requested, but the image output
%   arguments are generated (refer to the notes on Efficiency, below),
%   spline smoothing will be triggered just for the purposes of producing
%   images.
%
% depth_factors -- Depths expressed in focal lengths
%   The depths of the light sources, measured in multiples of the first
%   focal length of the lens from the first principal plane, for instance.
%   `depth_factors(j)` corresponds to the z-values in `X_lights(:, :, j)`.
%
%   `depth_factors` is a vector of values to substitute for light source
%   depths, for display purposes, and is not needed when
%   debugging/visualization output is disabled.
%
% verbose -- Debugging flags
%   If recognized fields of `verbose` are true, corresponding graphical
%   output will be generated for debugging purposes.
%
%   All debugging flags default to false if `verbose` is not passed.
%
% ## Output Arguments
%
% stats_real -- Simulated point spread function statistics
%   Point spread function statistics computed for the images produced by
%   each light source, for each wavelength, and for each depth, output as a
%   structure array. `stats_real(i, k, j)` is the `stats` output argument
%   of 'analyzePSF.m' for the i-th light, emitting the k-th wavelength, and
%   positioned at the j-th depth.
%
% stats_ideal -- Theoretical image statistics
%   Point spread function statistics for the images produced by each light
%   source, as predicted by the thick lens equation. `stats_ideal` has
%   the same format as `stats_real`.
%
% I -- Simulated greyscale images
%   A 4D array containing the images formed by combining all point spread
%   functions for all lights at each depth. `I(:, :, k, j)` is the
%   greyscale image for the k-th wavelength produced by the grid of lights
%   placed at the j-th depth.
%
%   This output argument is only available if there are multiple lights at
%   each depth.
%
% I_color -- Simulated colour images
%   A 4D array containing the images formed by combining all point spread
%   functions for all lights at each depth. `I_color(:, :, :, j)` is the
%   colour image for the produced by the grid of lights placed at the j-th
%   depth. The third dimension of `I_color` indexes colour channels.
%   `I_color` is produced by combining the images for individual
%   wavelengths in `I` according to the RGB colour values for the different
%   wavelengths in `lens_params.wavelengths_to_rgb`.
%
%   This output argument is only available if there are multiple lights at
%   each depth.
%
% ## Notes
% - Automatic estimation of the image boundaries, for `I_color` and `I`,
%   will yield poor results if all of the lights are on a single line.
%
% ### Coordinate System
% - The radii of both faces of the lens are positive when the lens is
%   biconvex.
% - The positive z-axis points towards the front of the lens, along
%   the optical axis, assuming the front of the lens is convex.
%
% ### Polychromatic light sources
%
% I output colour images by mixing together point spread functions for
% individual wavelengths, but this is purely for visualization. All
% mathematical analysis of chromatic aberration is done under narrowband
% sensor response assumptions. An advantage of this segregation is it is
% possible to estimate the theoretical chromatic aberration between
% individual wavelengths, ignoring the properties of the camera sensor.
% Furthermore, it is possible to run this function on an arbitrary number
% of wavelengths, and obtain as many estimates of chromatic aberration (not
% only three).
%
% See 'doubleSphericalLensPSF2()' for analysis of chromatic aberration in
% colour images.
%
% ### Efficiency
% - The `I` or `I_color` output arguments are expensive to produce. They
%   are generated if explicitly requested as output, or if any of the
%   following flags are set:
%   - `verbose.display_each_psf`
%   - `verbose.display_all_psf_each_ior`
%   - `verbose.display_all_psf_each_depth`
%
%   Of course, for a single light source, these output arguments are
%   unavailable, as mentioned above. The image boundaries are calculated
%   from the spacing between the image positions of multiple sources,
%   predicted by the thick lens equation. For a single light source, point
%   spread functions can be visualized using `verbose.display_each_psf`,
%   but there is presently no convenient way to define image boundaries.
% - Thin-plate spline smoothing is prohibitively expensive for large
%   numbers of rays (as discussed further in the documentation of
%   `request_spline_smoothing`, above).
% - For a single light source, `verbose.display_each_psf` will presently
%   trigger a scattered data interpolation operation, if
%   `request_spline_smoothing` is `false`.
%
% See also opticsFromLens, doubleSphericalLens, densifyRays, analyzePSF,
% imagingScenario, doubleSphericalLensPSF2, imageBoundaries

% Bernard Llanos
% Supervised by Dr. Y.H. Yang
% University of Alberta, Department of Computing Science
% File created June 22, 2017

%% Parse parameters

nargoutchk(1,4);
narginchk(7,9);

ray_params = lensParamsToRayParams(ray_params, lens_params, z_film);

ior_lens = lens_params.ior_lens;
n_ior_lens = length(ior_lens);

wavelengths = lens_params.wavelengths;
wavelengths_to_rgb = lens_params.wavelengths_to_rgb;

image_sampling = image_params.image_sampling;
normalize_psfs_before_combining = image_params.normalize_psfs_before_combining;
normalize_color_images_globally = image_params.normalize_color_images_globally;

n_lights = size(X_lights, 1);
single_source = (n_lights == 1);

image_output_requested = nargout > 2;

if ~isempty(varargin)
    if length(varargin) == 2
        depth_factors = varargin{1};
        verbose = varargin{2};
    else
        error('Both `depth_factors`, and `verbose` must be passed together.')
    end
    verbose_ray_tracing = verbose.verbose_ray_tracing;
    verbose_ray_interpolation = verbose.verbose_ray_interpolation;
    verbose_psf_analysis = verbose.verbose_psf_analysis;
    display_each_psf = verbose.display_each_psf;
    display_all_psf_each_ior = verbose.display_all_psf_each_ior;
    display_all_psf_each_depth = verbose.display_all_psf_each_depth;
    display_summary = verbose.display_summary;
else
    verbose_ray_tracing = false;
    verbose_ray_interpolation = false;
    verbose_psf_analysis = false;
    display_each_psf = false;
    display_all_psf_each_ior = false;
    display_all_psf_each_depth = false;
    display_summary = false;
end


%% Imaging setting configuration, and thick lens equation results

n_depths = size(X_lights, 3);
if n_depths == 1
    X_lights_matrix = X_lights;
else
    X_lights_matrix = reshape(permute(X_lights, [1, 3, 2]), [], 3);
end

n_lights_and_depths = n_lights * n_depths;

% "Theoretical" image positions, from the thick lens equation
stats_ideal_matrix = preallocateStats([n_lights_and_depths, n_ior_lens]);
for k = 1:n_ior_lens
    [ imageFn, ~, ~, U, U_prime ] = opticsFromLens(...
        ray_params.ior_environment,...
        ior_lens(k),...
        ray_params.ior_environment,...
        ray_params.radius_front, ray_params.radius_back,...
        ray_params.d_lens...
    );
    psfFn = opticsToPSF( imageFn, U, U_prime, lens_params.lens_radius, z_film );
    stats_ideal_matrix(:, k) = psfFn(X_lights_matrix);
end

stats_ideal = permute(reshape(stats_ideal_matrix, n_lights, n_depths, n_ior_lens), [1, 3, 2]);

% Find image boundaries
if isempty(image_params.image_bounds) && single_source
    error('Image boundaries cannot be automatically estimated for a single light source.')
end
image_bounds = imageBoundaries( image_params.image_bounds, stats_ideal );

%% Remove filtered-out light positions
lights_filter_rep = repmat(lights_filter, n_depths, 1);
stats_ideal = stats_ideal(lights_filter, :, :);
stats_ideal_matrix = stats_ideal_matrix(lights_filter_rep, :);
X_lights = X_lights(lights_filter, :, :);
X_lights_matrix = X_lights_matrix(lights_filter_rep, :);
n_lights = size(X_lights, 1);

%% Trace rays through the lens and form rays into an image
request_images = ~single_source && (...
    display_each_psf || display_all_psf_each_ior ||...
    display_all_psf_each_depth || image_output_requested ...
);
if request_images
    n_channels = size(wavelengths_to_rgb, 2);
    I = zeros([image_sampling, n_ior_lens, n_depths]);
    I_color = zeros([image_sampling, n_channels, n_depths]);
else
    I = [];
    I_color = [];
end
stats_real = preallocateStats([n_lights, n_ior_lens, n_depths]);
for j = 1:n_depths
    for k = 1:n_ior_lens
        ray_params.ior_lens = ior_lens(k);
        for i = 1:n_lights
            ray_params.source_position = X_lights(i, :, j);
            [ ...
                image_position, ray_irradiance, ~, incident_position_cartesian ...
                ] = doubleSphericalLens( ray_params, verbose_ray_tracing );

            if request_images
                if request_spline_smoothing
                    [ ~, v_adj, image_spline, I_ikj ] = densifyRays(...
                        incident_position_cartesian,...
                        ray_params.radius_front,...
                        image_position,...
                        ray_irradiance,...
                        image_bounds, image_sampling,...
                        verbose_ray_interpolation ...
                    );
                else
                    [ image_values, v_adj, ~, I_ikj ] = densifyRays(...
                        incident_position_cartesian,...
                        ray_params.radius_front,...
                        image_position,...
                        ray_irradiance,...
                        image_bounds, image_sampling,...
                        verbose_ray_interpolation ...
                    );
                end

                if normalize_psfs_before_combining
                    I_ikj_scaled = I_ikj ./ max(max(I_ikj));
                else
                    I_ikj_scaled = I_ikj;
                end
                I(:, :, k, j) = I(:, :, k, j) + I_ikj_scaled;
                for c = 1:n_channels
                    I_color(:, :, c, j) = I_color(:, :, c, j) +...
                        (I_ikj_scaled .* wavelengths_to_rgb(k, c));
                end

                if display_each_psf
                    figure
                    ax = gca;
                    imagesc(ax,...
                        [image_bounds(1), image_bounds(1) + image_bounds(3)],...
                        [image_bounds(2) + image_bounds(4), image_bounds(2)],...
                        I_ikj...
                        );
                    colormap gray
                    ax.YDir = 'normal';
                    xlabel('X');
                    ylabel('Y');
                    c = colorbar;
                    c.Label.String = 'Irradiance';
                    title(...
                        sprintf('Estimated PSF for a point source at position\n[%g, %g, %g] (%g focal lengths, IOR %g)',...
                        X_lights(i, 1, j), X_lights(i, 2, j), X_lights(i, 3, j),...
                        depth_factors(j), ior_lens(k)...
                        ));
                    axis equal
                end
            else
                if request_spline_smoothing
                    [ ~, v_adj, image_spline ] = densifyRays(...
                        incident_position_cartesian,...
                        ray_params.radius_front,...
                        image_position,...
                        ray_irradiance,...
                        verbose_ray_interpolation ...
                    );
                else
                    [ image_values, v_adj ] = densifyRays(...
                        incident_position_cartesian,...
                        ray_params.radius_front,...
                        image_position,...
                        ray_irradiance,...
                        verbose_ray_interpolation ...
                    );
                end
            end

            if single_source && display_each_psf
                figure
                if request_spline_smoothing
                    pts = fnplt(image_spline);
                else
                    plot_resolution = [200 200];
                    x = linspace(...
                        min(image_position(:, 1)),...
                        max(image_position(:, 1)), plot_resolution(2)...
                    );
                    y = linspace(...
                        min(image_position(:, 2)),...
                        max(image_position(:, 2)), plot_resolution(1)...
                    );
                    [pts{1},pts{2}] = meshgrid(x,y);
                    % In the future, this operation could be deduplicated
                    % between this function and `analyzePSF()` (and perhaps
                    % also `densifyRays()`, if applicable).
                    psf_interpolant = scatteredInterpolant(...
                        image_position, image_values...
                    );
                    pts{3} = psf_interpolant(pts{1},pts{2});
                end
                surf(pts{1}, pts{2}, pts{3}, 'EdgeColor', 'none');
                colorbar
                xlabel('X');
                ylabel('Y');
                zlabel('Irradiance');
                colormap summer
                c = colorbar;
                c.Label.String = 'Irradiance';
                title(...
                    sprintf('Estimated PSF for a point source at position\n[%g, %g, %g] (%g focal lengths, IOR %g)',...
                    X_lights(i, 1, j), X_lights(i, 2, j), X_lights(i, 3, j),...
                    depth_factors(j), ior_lens(k)...
                ));
            end

            if request_spline_smoothing
                stats_real(i, k, j) = analyzePSF(...
                    image_spline, image_position, v_adj,...
                    verbose_psf_analysis...
                );
            else
                stats_real(i, k, j) = analyzePSF(...
                    image_values, image_position, v_adj,...
                    verbose_psf_analysis...
                );
            end
        end

        % Visualize the results, for this wavelength
        if request_images && display_all_psf_each_ior
            figure
            ax = gca;
            imagesc(...
                [image_bounds(1), image_bounds(1) + image_bounds(3)],...
                [image_bounds(2) + image_bounds(4), image_bounds(2)],...
                I(:, :, k, j)...
                );
            colormap gray
            ax.YDir = 'normal';
            xlabel('X');
            ylabel('Y');
            c = colorbar;
            c.Label.String = 'Irradiance';
            hold on
            mean_position_ideal = vertcat(stats_ideal(:, k, j).mean_position);
            mean_position_real = vertcat(stats_real(:, k, j).mean_position);
            scatter(mean_position_ideal(:, 1), mean_position_ideal(:, 2), [], wavelengths_to_rgb(k, :), 'o');
            scatter(mean_position_real(:, 1), mean_position_real(:, 2), [], wavelengths_to_rgb(k, :), '.');
            legend('Thick lens formula', 'Raytracing centroids');
            title(sprintf(...
                'Images of point sources at %g focal lengths, for \\lambda = %g nm',...
                depth_factors(j), wavelengths(k)));
            axis equal
            hold off
        end
    end
end

% Visualize the results, for each depth
if request_images
    if normalize_color_images_globally
        I_color = I_color ./ max(max(max(max(I_color))));
    else
        for j = 1:n_depths
            I_color(:, :, :, j) = I_color(:, :, :, j) ./ max(max(max(I_color(:, :, :, j))));
        end
    end

    if display_all_psf_each_depth
        for j = 1:n_depths
            figure
            ax = gca;
            image(...
                [image_bounds(1), image_bounds(1) + image_bounds(3)],...
                [image_bounds(2) + image_bounds(4), image_bounds(2)],...
                I_color(:, :, :, j)...
                );
            ax.YDir = 'normal';
            xlabel('X');
            ylabel('Y');
            hold on
            legend_strings = cell(n_ior_lens * 2, 1);
            for k = 1:n_ior_lens
                mean_position_real = vertcat(stats_real(:, k, j).mean_position);
                mean_position_ideal = vertcat(stats_ideal(:, k, j).mean_position);
                scatter(mean_position_ideal(:, 1), mean_position_ideal(:, 2), [], wavelengths_to_rgb(k, :), 'o');
                scatter(mean_position_real(:, 1), mean_position_real(:, 2), [], wavelengths_to_rgb(k, :), '.');
                legend_strings{2 * k - 1} = sprintf('Thick lens formula, \\lambda = %g nm', wavelengths(k));
                legend_strings{2 * k} = sprintf('Raytracing centroids, \\lambda = %g nm', wavelengths(k));
            end
            legend(legend_strings);
            title(sprintf(...
                'Images of point sources at %g focal lengths',...
                depth_factors(j)));
            axis equal
            hold off
        end
    end
end

%% Visualize the results, for all depths
stats_real_matrix = reshape(permute(stats_real, [1, 3, 2]), [], n_ior_lens);

if display_summary
    if single_source
        disp('Distances from first principal plane, in focal lengths:')
        disp(depth_factors)
        disp('Light source position [x, y, z]:')
        disp(X_lights_matrix)
        disp('Image positions:')
        for k = 1:n_ior_lens
            fprintf('Thick lens equation (lambda = %g nm):\n', wavelengths(k))
            disp(vertcat(stats_ideal_matrix(:, k).mean_position))
            fprintf('Raytracing centroids (lambda = %g nm):\n', wavelengths(k))
            disp(vertcat(stats_real_matrix(:, k).mean_position))
        end
    else
        figure
        hold on
        depth_factors_rep = repelem(depth_factors, n_lights);
        legend_strings = cell(n_ior_lens * 2, 1);
        for k = 1:n_ior_lens
            mean_position_ideal = vertcat(stats_ideal_matrix(:, k).mean_position);
            mean_position_real = vertcat(stats_real_matrix(:, k).mean_position);
            scatter3(mean_position_ideal(:, 1), mean_position_ideal(:, 2), depth_factors_rep, [], wavelengths_to_rgb(k, :), 'o');
            scatter3(mean_position_real(:, 1), mean_position_real(:, 2), depth_factors_rep, [], wavelengths_to_rgb(k, :), '.');
            legend_strings{2 * k - 1} = sprintf('Thick lens formula, \\lambda = %g nm', wavelengths(k));
            legend_strings{2 * k} = sprintf('Raytracing centroids, \\lambda = %g nm', wavelengths(k));
        end
        legend(legend_strings);
        title('Images of point sources seen through a thick lens')
        xlabel('X');
        ylabel('Y');
        zlabel('Light source distance (focal lengths)')
        hold off
    end
end

end