%% Lab 3: Trichromacy/SVD, Regression
% Updated 2022
% -- solutions and figures from slides-- % 
clear
close all
rng(2)

%% Part 1 Trichromacy
% In this part of the lab, we will go over a simple trichromacy example. 
% Let's remember how a trichromacy test works. 
% Subjects are given a set of primaries and are 
% asked to match the colors by adjusting the weights 
% on a set of pre-determined primaries. Let's call the 
% primaries $P$ and the test light $\vec{l}$. The humans 
% will select the weights $\vec{w}$ such that $P\vec{w}$ 
% produces a light that looks similar to $\vec{l}$. We'll 
% denote this $P\vec{w} \sim \vec{l}$ 

%%
% We assume humans are some linear system $H$ - this matrix 
% $H$ will take in a test light $\vec{l}$ and tell us how to 
% weight the primaries. So $\vec{w} = H\vec{l}$. Putting this 
% all together, we have $P\vec{w} = PH\vec{l} \sim \vec{l}$

l = rand(31, 1);

%%
% We are given primaries and our human linear system $H$ below: 
primaries = [[3.72665317e-06, 4.57833362e-01, 4.39369336e-02];
       [1.86644691e-05, 6.06530660e-01, 1.11089965e-02];
       [8.36483472e-05, 7.54839602e-01, 2.18749112e-03];
       [3.35462628e-04, 8.82496903e-01, 3.35462628e-04];
       [1.20385999e-03, 9.69233234e-01, 4.00652974e-05];
       [3.86592014e-03, 1.00000000e+00, 3.72665317e-06];
       [1.11089965e-02, 9.69233234e-01, 2.69957850e-07];
       [2.85655008e-02, 8.82496903e-01, 1.52299797e-08];
       [6.57285286e-02, 7.54839602e-01, 6.69158609e-10];
       [1.35335283e-01, 6.06530660e-01, 2.28973485e-11];
       [2.49352209e-01, 4.57833362e-01, 6.10193668e-13];
       [4.11112291e-01, 3.24652467e-01, 6.10193668e-13];
       [6.06530660e-01, 2.16265167e-01, 2.28973485e-11];
       [8.00737403e-01, 1.35335283e-01, 6.69158609e-10];
       [9.45959469e-01, 7.95595087e-02, 1.52299797e-08];
       [1.00000000e+00, 4.39369336e-02, 2.69957850e-07];
       [9.45959469e-01, 2.27941809e-02, 3.72665317e-06];
       [8.00737403e-01, 1.11089965e-02, 4.00652974e-05];
       [6.06530660e-01, 5.08606923e-03, 3.35462628e-04];
       [4.11112291e-01, 2.18749112e-03, 2.18749112e-03];
       [2.49352209e-01, 8.83826307e-04, 1.11089965e-02];
       [1.35335283e-01, 8.83826307e-04, 4.39369336e-02];
       [6.57285286e-02, 2.18749112e-03, 1.35335283e-01];
       [2.85655008e-02, 5.08606923e-03, 3.24652467e-01];
       [1.11089965e-02, 1.11089965e-02, 6.06530660e-01];
       [3.86592014e-03, 2.27941809e-02, 8.82496903e-01];
       [1.20385999e-03, 4.39369336e-02, 1.00000000e+00];
       [3.35462628e-04, 7.95595087e-02, 8.82496903e-01];
       [8.36483472e-05, 1.35335283e-01, 6.06530660e-01];
       [1.86644691e-05, 2.16265167e-01, 3.24652467e-01];
       [3.72665317e-06, 3.24652467e-01, 1.35335283e-01]];

H = [[-7.86629629e-03, -1.64055828e-02, -2.95981171e-02, ...
        -3.93796959e-02, -4.09500966e-02, -3.38263122e-02,...
        -2.37004317e-02, -9.91485201e-03,  1.20239124e-02,...
         3.57069852e-02,  6.48046756e-02,  1.04178973e-01,...
         1.44668517e-01,  1.68752352e-01,  1.77676195e-01,...
         1.74234972e-01,  1.60424243e-01,  1.37083289e-01,...
         1.06634557e-01,  7.34436273e-02,  4.39069236e-02,...
         2.20195756e-02,  8.79040897e-03,  2.16407666e-03,...
        -5.16366094e-04, -1.13994950e-03, -8.66775757e-04,...
        -5.47922896e-04, -3.41163670e-04, -1.58712800e-04,...
        -1.06894912e-04];...
       [ 3.18589820e-02,  6.65498781e-02,  1.21057897e-01,...
         1.65625669e-01,  1.82539231e-01,  1.67568336e-01,...
         1.46179420e-01,  1.26060988e-01,  8.45603935e-02,...
         4.92830803e-02,  2.78158110e-02,  1.36406802e-02,...
         2.84097397e-03, -2.50989671e-03, -5.22701906e-03,...
        -6.80644635e-03, -6.62236256e-03, -6.79197033e-03,...
        -5.95825114e-03, -4.92531575e-03, -3.87453776e-03,...
        -2.91506416e-03, -2.10206731e-03, -1.40963160e-03,...
        -9.04883888e-04, -5.42254762e-04, -3.06237143e-04,...
        -1.59404323e-04, -8.40136175e-05, -4.06022050e-05,...
        -1.99329051e-05];...
       [ 2.71777556e-02,  5.62179675e-02,  1.00315960e-01,...
         1.29839162e-01,  1.27729775e-01,  9.36612642e-02,...
         4.66875314e-02, -1.36220201e-02, -8.72000273e-02,...
        -1.55793381e-01, -2.33802345e-01, -3.27595055e-01,...
        -3.87724450e-01, -3.65221662e-01, -2.79396978e-01,...
        -1.43816136e-01,  3.56879942e-02,  2.48737645e-01,...
         4.69506595e-01,  6.59391091e-01,  7.74562774e-01,...
         7.85955117e-01,  6.97520716e-01,  5.36579177e-01,...
         3.79600809e-01,  2.42798994e-01,  1.41234085e-01,...
         7.53007753e-02,  4.06534560e-02,  1.95333127e-02,...
         1.01241222e-02]];...

%%
% Let's generate a new light that looks like our test_light, 
% and plot it to compare it with the original light. 
% Do they look alike? 

% TOD0: use P, H, and our test light l to generate a new light combo
% that is "perceptually indistinguishable" from l, called matched_light
matched_light = %YOUR CODE HERE

% Do these lights have similar spectra? Plot l and matched_light to see.
figure;
plot(l)
hold on;
plot(matched_light)
xlabel("Wavelength (bins)")
ylabel("Intensity")
legend("Test Light", "Matched Light")

%%
% Print the weights for the matched light and the original 
% test light. Are they the same?

%TODO: find the weights associated with the test light, and the weights
% associated with the matched light
weights_for_test_light = %YOUR CODE HERE
weights_for_matched_light = %YOUR CODE HERE
disp(weights_for_test_light)
disp(weights_for_matched_light)

%%
% Let's take the SVD of $H$ - this will tell us what $H$ is 
% doing "behind the scenes". How many singular values does 
% $H$ have? Does that line up with what you expected?

%TODO: take the SVD and display the number of singular values
[U, S, V] =  %YOUR CODE HERE
singular_values = %YOUR CODE HERE
number_of_singular_values = length(singular_values)

%%
% Now let's think about what information gets thrown away 
% by our matrix $H$. Remember that the vectors that form the 
% basis of the nullspace of $H$ are the vectors in $V$ that 
% correspond with singular values of 0. Which vectors of V 
% correspond with singular values of 0? How do we know they 
% are in the nullspace?

% TODO: find the nullspace basis vectors using V and S
nullspace_basis_vectors = %YOUR CODE HERE
round(H * nullspace_basis_vectors, 5);

%%
% Using this information, construct a new light $\vec{n}$ 
% which is some weighted combination of basis vectors, 
% and add it to your test light, $\vec{l}$. Show that 
% these are metamers by demonstrating that $H\vec{n} = H\vec{l}$

% TODO: generate some weighted combo of nullspace basis vectors and
% and add it to your test light
n = %YOUR CODE HERE
weights_for_new_light = %YOUR CODE HERE
disp(weights_for_test_light)
disp(weights_for_new_light)

%%
% Plot your new light and compare it with your old light. 
% Is the new light a color that could be physically realizable? 
% If not, try to find some other new light that is physically 
% realizable and also is a metamer with the original test light.
figure;
%YOUR CODE HERE
xlabel("Wavelength (bins)")
ylabel("Intensity")
legend("Test Light", "New Light", "Physically Realizable New Light")

%% Clear your workspace and close all figures
clear all
close all

%% Part 2 Regression
%% Visual neuron 
% In this section, we will model the response of a simulated neuron
% responding to the luminance of a visual stimulus, and characterize its
% "spiking" with linear regression
%%
% Create test data simulating a neuron firing in response to a light stimulus:

luminance = (1:10)'; % Arbitrary units to represent input intensities

%%
% 
% what is the spike response to different luminance? Model it as a linear
% amplification plus baseline and noise. Include a term to introduce
% experimental error to random trials
baseline = 10; noise = 2; amp = 2; experimentalError = (rand(10,1)>.5).*randn(10,1)*5;
spikes = (amp * luminance + baseline + randn(length(luminance),1).*noise) + experimentalError ; 

%%

% Plot Luminance vs Spikes
figure(1); % we give these figures a designated handle so that we can repeatedly plot in them 
plot(luminance,spikes,'ko', 'MarkerSize',18); 
xlabel('Luminance')
ylabel('Spikes')
set(gca,'Fontsize',18)
set(gcf,'color','w')
xlim([0 11])
ylim([0 max(spikes)+10])

%%
% Find a value of beta (i.e., a slope of the regression line) such that
% when it's multiplied to each of the luminance values, we minimize the
% distance to the actual measurements.
% aka we look for beta_op which minimizes ||y - betaOpt * x||^2

% Find betaOpt (optimal beta) via SVD!
betaOpt = %YOUR CODE HERE

% alternative method: (solution found by calculus)
% beta_op = ySpikes'*xLuminance / (xLuminance'*xLuminance); %Normalized projection

%%
% Now that we have our optimal weight to scale each luminance value, how
% well does our beta do at predicting the spikes? 

prediction = %YOUR CODE HERE %our prediction is multiplying each luminance value by beta
predError = %YOUR CODE HERE % get distance between measured spikes and our model prediction
predErrorOpt = %YOUR CODE HERE % get the squared distance

%%

% Plot our prediction
hold on
plot(luminance, prediction, 'k', 'LineWidth', 3, ...
    'DisplayName', ['lin reg, sq error = ', num2str(predErrorOpt)])
legend('location', 'southoutside')

%%
%  Was this the best beta? 
% Test a range of other beta values to confirm that your beta truly
% minimizes ||y - betaOpt * x||^2
numBeta = 200; testBetas = linspace(0, betaOpt*2, numBeta); % Try 200 betas over a range of 0:2*betaOpt
sqrError = nan*ones(1,numBeta); % initialize 
for ii = 1:numBeta % Loop through all the betas
    prediction = testBetas(ii) .* luminance;
    % We need to now explicitly represent the distance metric
    sqrError(ii) = sum((prediction-spikes).^2); %Sum of squares
end

%%

% now check if predErrorOpt is indeed optimal:
if min(sqrError) >= predErrorOpt 
    % if the closest error you got by force was no smaller than your calculated beta opt....
   disp('seems like we found a great solution')
end

%%

% Now plot the error for different betas
figure(2)
plot(testBetas,sqrError,'b','linewidth',3) % plot the error of your set of test betas
line([betaOpt betaOpt], [min(ylim) max(ylim)], 'color','r','linewidth',2); % plot a line at your calculated optimal beta value
xlabel('Beta')
ylabel('Error metric')
set(gca,'fontsize',18)
set(gcf,'color','w')

%% Multiple Linear Regression: Visual neuron model with y-intercept
%%
% What if we would like to add a y-intercept? We can do this with
% multiple linear regression if we model our independent variable (x)
% as a matrix whose columns contain the x values of the first order
% polynomial yPredicted = Beta0*x^0 + Beta1*x^1 (think y = b + mx). Solving
% this multiple linear regression will produce two beta values, Beta0 (scaling
% the y intercept term) and Beta1 (scaling the slope of the regression line)

%%

% Create your regressor matrix and calculate the optimal betas with
% linear algebra

% create your "multivariate" data matrix X (containing x^0 and x^1 as
% columns)
X = %YOUR CODE HERE

% Calculate betaOpt with the svd
betaOptYInt = %YOUR CODE HERE

%%
% Now that we have our optimal weight to scale each luminance value, and our y-interecpt, how
% well does our beta (now a vector) do at predicting the spikes? 

prediction = %YOUR CODE HERE %our prediction is multiplying each luminance value by beta
predError = %YOUR CODE HERE % get distance between measured spikes and our model prediction
predErrorOptYInt = %YOUR CODE HERE % get the squared distance

% Plot our prediction with error val
figure(1)
hold on
plot(luminance, prediction, '-.', 'LineWidth', 3, ...
    'DisplayName', ['regression with y-int, sq error = ', num2str(predErrorOptYInt)])
legend('location', 'southoutside')

%%
% How did our beta do? Test it manually. First, create a meshgrid of beta values (equivalent to a 2D version of
% the beta vector from the last section)

nBetas = 100;
beta0 = linspace(-betaOptYInt(1),betaOptYInt(1)*2,nBetas);
beta1 = linspace(1,betaOptYInt(2)*2,nBetas);

% now compute the errors
allerr = nan*ones(nBetas);
for ii = 1:nBetas
    for jj = 1:nBetas
        bb = [beta0(ii);beta1(jj)];
        allerr(ii,jj) = sum((X*bb-spikes).^2);
    end
end

%%
% plot the contour and compare the old prediction error (without the
% y-int) with the new prediction error - did you do any better?
% what does it mean for our neuron that the model with the y-intercept has
% lower error?

figure(3);
[betaX,betaY] = meshgrid(beta0,beta1);
s = surf(betaX,betaY,allerr', 'DisplayName', 'sq err');
hold on
scatter3(betaOptYInt(1), betaOptYInt(2), predErrorOptYInt, 100,'r', 'filled', ...
        'DisplayName', 'Optimal beta')
hold on
scatter3(0, betaOpt, predErrorOpt, 100,'m', 'filled', ...
        'DisplayName', 'Old optimal beta')
xlabel('beta0')
ylabel('beta1')
zlabel('squared error')
rotate3d on
alpha 0.5
s.EdgeColor = 'none';
legend('location', 'northoutside')


%% Clear your workspace and close all figures
clear all
close all

%% Multiple Linear Regression: Auditory Neuron
% Now let's try an example of multiple linear regression. The structure
% mult_linreg contains three fields: freq1, freq2 (the dB intensities of two
% frequencies of sound stimuli) and response, the response of an auditory
% neuron. 

% First, load and plot the data ( in 3D!)
load('mult_linreg.mat');

figure(3)
plot3(data.freq1, data.freq2,data.response, '.', 'MarkerSize',20, 'DisplayName', 'Auditory neuron response')
xlabel('Freq 1')
ylabel('Freq 2')
zlabel('Firing rate')
set(gca,'fontsize',18)
rotate3d on
grid on

% Perform linear regression
betaOpt = %YOUR CODE HERE

% Calculate error between the prediction frorm our model, and the measured
% values
prediction = %YOUR CODE HERE %our prediction is multiplying each luminance value by beta
predError = %YOUR CODE HERE % get distance between measured spikes and our model prediction
predErrorOpt = %YOUR CODE HERE % get the squared distance

% plot our linear model (it's a plane!)
x = linspace(0, max(data.freq1*1.2),100)';
y = linspace(0, max(data.freq2*1.2),100)';

[xx, yy] = meshgrid(x,y);
zz = betaOpt(1)*xx + betaOpt(2)*yy;
figure(3)
hold on
s = surf(xx,yy,zz, 'DisplayName', ['lin reg, sq error = ', num2str(predErrorOpt)]);
alpha 0.5
s.EdgeColor = 'none';
s.FaceColor = [1,.8,.9];
legend('location','northoutside')


% Was this the best beta?
% first, create vectors of your possible betas
nBetas = 200;

beta0 = linspace(0,betaOpt(1)*2,nBetas);
beta1 = linspace(0,betaOpt(2)*2,nBetas);

% now compute the errors
allerr = nan*ones(nBetas);
for ii = 1:nBetas
    for jj = 1:nBetas
        bb = [beta0(ii);beta1(jj)];
        allerr(ii,jj) = sum(([data.freq1, data.freq2]*bb-data.response).^2);
    end
end

% plot the contour!
figure(5);

% in 3D, our possible errors are a surface instead of a line. We will use
% the meshgrid function to plot this surface.

[betaX,betaY] = meshgrid(beta0,beta1);
s = surf(betaX,betaY,allerr');
xlabel('beta0')
ylabel('beta1')
zlabel('squared error')
rotate3d on
alpha 0.5;
s.EdgeColor = 'none';
hold on
scatter3(betaOpt(1), betaOpt(2), predErrorOpt, 100,'k', 'filled', ...
        'DisplayName', 'optimal beta')


%% Multiple Linear Regression: Auditory neuron with baseline firing
% What if we suspect the auditory neuron has some baseline firing rate? We
% might want to include an intercept term in our regression. We can do this
% the same way we added an intercept term to the standard linear
% regression: by adding an x^0 term to our linear regression.

% Linear Regression 
newX = %YOUR CODE HERE
betaOpt = %YOUR CODE HERE

% Calculate error between the prediction frorm our model, and the measured
% values
newPrediction = %YOUR CODE HERE %our prediction is multiplying each luminance value by beta
newPredError = %YOUR CODE HERE % get distance between measured spikes and our model prediction
newPredErrorOpt = %YOUR CODE HERE % get the squared distance

% plot our linear model (it's a plane!)

newBeta0 = linspace(0,max(data.freq1)*1.2,nBetas)';
newBeta1 = linspace(0,max(data.freq2)*1.2,nBetas)';

[newBetaX,newBetaY] = meshgrid(newBeta0,newBeta1);
newS = betaOpt(1)*ones(size(newBetaX))+ betaOpt(2)*newBetaX + betaOpt(3)*newBetaY;
figure(3)
hold on
s = surf(newBetaX,newBetaY,newS, 'DisplayName', ['lin reg + zint, sq error = ', num2str(newPredErrorOpt)]);
alpha 0.5
s.EdgeColor = 'none';
s.FaceColor = [.9,.5,.1];
% How else might you improve the fit of our surface? 
