% Fourier solution of scalar wave equation initial value problem.
function psi = waves2d(psi0, psidot0, times, c)
% Input parameters:
%  psi0 = initial wave function (in 2 dimensions)
%  psidot0 = initial time derivative of wave function
%     (must be the same size as the psi0 array).
%  c = wave speed in pixels per unit time. Default = 1.
%  times = array of times for which the solution is desired.

if not(exist('c', 'var'))
   c = 1;  % Default wave speed
end

psi0_f = fftn(psi0);
psidot0_f = fftn(psidot0);

% Work out the wavenumbers for each point in the FFT plane
[kx, ky, k] = k2d(size(psi0_f));

psi=[];
Nt = numel(times);
for i = 1:Nt
   t = times(i);
   psi_f = psi0_f .* cos(c*k*t) + ...
           psidot0_f .* sin(c*k*t) ./ (c * k);
   psi_f(1,1) = psi0_f(1,1) + psidot0_f(1,1) * t; % handle k=0 singularity.
   psi(:,:,i) = real(ifftn(psi_f));
end
