% Fast poisson solver
% 4-16-2008
clear all; close all; clc

% Problem parameters
Lx = 20;  Ly = 20;
nx = 128; ny = 128;
x2 = linspace(-Lx/2, Lx/2, nx+1);  x = x2(1:nx);
y2 = linspace(-Ly/2, Ly/2, ny+1);  y = y2(1:ny);

% Construct the grid and the rhs omega
[X,Y] = meshgrid(x,y);
omega = exp(-X.^2 - Y.^2);
figure(1);
surf(x,y,omega); % Plot omega
shading flat;

% Construct wave number vectors
kx = (2*pi/Lx) * [0:nx/2-1 -nx/2:-1];
ky = (2*pi/Ly) * [0:ny/2-1 -ny/2:-1];
kx(1) = 10^(-4);  ky(1) = 10^(-4);

% Solve the PDE
[KX,KY] = meshgrid(kx,ky);
psi = real( ifft2( -fft2(omega)./(KX.^2 + KY.^2) ) );
psi = psi-max(max(psi));  % Correct for the nonuniqueness for plotting

figure(2);
surf(x,y,psi);
shading flat;

% Take derivatives of the stream function
psi_x = imag(ifft(-i*KX.*fft(psi)));
psi_y = real(ifft(-i*KY.*fft(psi)));
figure(3);
subplot(1,2,1)
surf(x,y,psi_x);
shading flat;
subplot(1,2,2)
surf(x,y,psi_y);
shading flat;