AOC 2023, Day 24, Part 2

%% Advent of Code 2023, Day 24, Part 2 - GLSDC (nonlinear iterated least squares)
tic;
Np = 3; % only need 3 points to find solution
% Np = N; % but works just fine with all 300 pts too
Nx = 6 + Np; % states are initial position, initial velocity, and Np collision times
Ny = 3*Np; % measurements are initial hailstone positions

% normalize inputs to avoid inverting singular matrix (scale back at end)
posscale = 10^-ceil(log10(max(abs(pos),[],'all')));
velscale = 10^-ceil(log10(max(abs(vel),[],'all')));
p = pos(1:Np,:)*posscale; % compute and cache scaled input positions
v = vel(1:Np,:)*velscale; % compute and cache scaled input velocities

% preallocate arrays and initialize state estimate
H = zeros(Ny,Nx); % Jacobian / linear observation matrix
y = zeros(Ny,1); % actual measurements
yhat = zeros(Ny,1); % predicted/estimated measurements
xhat = 0.1*(1:Nx)'; % ensure initial observation matrix is full rank
% xhat = randn(Nx,1);

Ni = 50; % should generally need less than 20
numiters = 0;
dx = Inf;
while norm(dx) > 1e-9 && numiters < Ni
    pidx = 1:3;
    vidx = 4:6;
    xs = xhat(pidx);
    vs = xhat(vidx);
    for i=1:Np
        tidx = 6+i;
        yidx = (3*(i-1))+(1:3);
        ti = xhat(tidx);
        xi = p(i,:)';
        vi = v(i,:)';
        y(yidx) = xi;
        yhat(yidx) = xs + vs*ti -vi*ti;
        H(yidx,pidx) = eye(3);
        H(yidx,vidx) = ti*eye(3);
        H(yidx,tidx) = vs - vi;
    end
    dy = y - yhat; % innovation
    dx = (H'*H)\H'*dy; % compute correction
    xhat = xhat + dx; % apply correction
    numiters = numiters + 1;
end
if numiters < Ni
    % scale back converged state and compute final answer
    phat = round(xhat(1:3)/posscale);
    % vhat = round(xhat(4:6)/velscale);
    % that = round(xhat(7:end)*velscale/posscale);
    part2 = int64(sum(phat))
else
    error('Could not converge on solution')
end
toc % approx 5 ms for 3 pts, 150 ms using all 300 pts