% ch13_linear_fit.m
% Companion code for Chapter 13: Optimization and Data Fitting.
%
% Fit a straight line y = a0 + a1*x by least squares.
% Replace the example data below with measured x, y, and optional sigma_y.
% Assumptions: independent errors in y; errors in x are negligible.
% sigma_y contains known one-standard-deviation measurement uncertainties.
% With sigma_y = [], use ordinary least squares and estimate the common
% noise variance from the residuals. Parameter errors are standard errors,
% not confidence intervals. No fitting toolbox is required.

clear; close all; clc;

%% Example data and settings
seed = 1;
rng(seed, 'twister');
% Alternative: comment out the line above and uncomment the following.
% rng('shuffle');

x = linspace(0,10,15)';
y = 1.0 + 0.5*x + 0.5*randn(size(x));
sigma_y = 0.5*ones(size(x)); % Known uncertainties; [] if unknown
make_plot = true;

% For measured data, replace the example above, for instance:
% data = readmatrix('measurements.txt');
% x = data(:,1); y = data(:,2); sigma_y = data(:,3);
% Use sigma_y = [] for a two-column file without known uncertainties.

%% Fit
[a0,a1,info] = linear_fit(x,y,sigma_y);
fprintf('\nStraight-line least-squares fit: y = a0 + a1*x\n');
fprintf('Intercept a0: %.10g +/- %.6g\n',a0,info.standard_errors(1));
fprintf('Slope a1:     %.10g +/- %.6g\n',a1,info.standard_errors(2));
fprintf('Residual sum of squares: %.10g\n',info.sse);
fprintf('Residual standard deviation: %.6g\n',info.residual_std);
fprintf('Degrees of freedom: %d\n',info.dof);
if info.known_sigma
    fprintf('Chi-squared:         %.6g\n',info.chi2);
    fprintf('Reduced chi-squared: %.6g\n',info.chi2/info.dof);
    fprintf('Parameter errors use supplied uncertainties without rescaling.\n');
else
    fprintf('Parameter errors use the residual estimate of common noise variance.\n');
end
% Full parameter covariance matrix is available as info.covariance.

%% Optional data, fitted line, and residual plots
if make_plot
    xx = linspace(min(x),max(x),300);
    figure('Color','w');
    subplot(2,1,1);
    if info.known_sigma
        hd = errorbar(x,y,sigma_y,'ko','MarkerFaceColor','w','MarkerSize',6);
    else
        hd = plot(x,y,'ko','MarkerFaceColor','w','MarkerSize',6);
    end
    hold on;
    hf = plot(xx,a0+a1*xx,'k-','LineWidth',1.5);
    xlabel('x','FontSize',14); ylabel('y','FontSize',14);
    title('Straight-line least-squares fit','FontSize',14,'FontWeight','normal');
    legend([hd,hf],{'Data','Fitted line'},'Location','best');
    set(gca,'FontSize',14); grid on; box on;

    subplot(2,1,2);
    if info.known_sigma
        errorbar(x,info.residuals,sigma_y,'ko','MarkerFaceColor','w','MarkerSize',6);
    else
        plot(x,info.residuals,'ko','MarkerFaceColor','w','MarkerSize',6);
    end
    hold on; plot([min(x) max(x)],[0 0],'k--');
    xlabel('x','FontSize',14); ylabel('y - y_{fit}','FontSize',14);
    set(gca,'FontSize',14); grid on; box on;
end

%% Reusable numerical routine
function [a0,a1,info] = linear_fit(x,y,sigma_y)
% Weighted centered sums avoid subtracting large nearly equal raw sums.
% Known absolute sigma_y: covariance follows the supplied variances.
% Unknown common sigma_y: covariance is scaled by SSE/(N-2).
% Copy this routine to linear_fit.m for use from other programs.
    if nargin < 3
        sigma_y = [];
    end
    validateattributes(x,{'double'},{'real','vector','finite','nonempty'});
    validateattributes(y,{'double'},{'real','vector','finite','nonempty'});
    x=x(:); y=y(:); n=numel(x);
    if numel(y)~=n || n<3
        error('Supply matching x and y vectors with at least three points.');
    end
    known_sigma=~isempty(sigma_y);
    if known_sigma
        validateattributes(sigma_y,{'double'},{'real','vector','finite','positive'});
        sigma_y=sigma_y(:);
        if numel(sigma_y)~=n
            error('Supply one uncertainty for each measurement.');
        end
        w=1./sigma_y.^2;
    else
        w=ones(n,1);
    end
    S=sum(w);
    xbar=sum(w.*x)/S;
    ybar=sum(w.*y)/S;
    dx=x-xbar;
    Sxx=sum(w.*dx.^2);
    if ~isfinite(S) || S<=0 || ~isfinite(Sxx) || Sxx<=0
        error('Distinct x values and numerically usable weights are required.');
    end
    a1=sum(w.*dx.*(y-ybar))/Sxx;
    a0=ybar-a1*xbar;
    residuals=y-(ybar+a1*dx);
    sse=sum(residuals.^2);
    dof=n-2;
    covariance=[1/S+xbar^2/Sxx, -xbar/Sxx; -xbar/Sxx, 1/Sxx];
    if known_sigma
        chi2=sum((residuals./sigma_y).^2);
    else
        covariance=covariance*(sse/dof);
        chi2=NaN; % No known measurement variances for a chi-squared test.
    end
    info=struct('standard_errors',sqrt(diag(covariance)), ...
        'covariance',covariance,'residuals',residuals,'sse',sse, ...
        'residual_std',sqrt(sse/dof),'dof',dof, ...
        'known_sigma',known_sigma,'chi2',chi2);
end