% http://www.math.utah.edu/~wright/courses/5620/homework/fd2poisson.m

% Numerical approximation to Poisson's equation over the square [0,a]x[0,a] with
% Dirichlet boundary conditions.  Uses a uniform mesh with (n+2)x(n+2) total
% points (i.e, n x n interior grid points).
% Input:
%        pfunc : the RHS of poisson equation
%        bfunc : the boundary function representing the Dirichlet B.C.
%            a : the size of the square
%            n : n+2 is the number of points in either direction of the mesh.
% Ouput:
%            u : the numerical solution of Poisson equation at the mesh points.
%          x,y : the uniform mesh.
%
function [u, x, y] = myfd2poisson(pfunc, bfunc, a, n)

	h = a/(n+1);   % Mesh spacing
	[x, y] = meshgrid(0:h:a, 0:h:a);   % Uniform mesh, including boundary points.

	% Compute u on the boundary from the Dirichlet boundary condition
	ub = zeros(n, n);
	idx = 2:n+1;
	idy = 2:n+1;

	% West and East boundaries
	ub(:,1) = bfunc(x(idx,1),y(idy,1));
	ub(:,n) = bfunc(x(idx,n+2),y(idy,n+2));

	% North and South boundaries
	ub(1,1:n) = ub(1,1:n) + bfunc(x(1,idx),y(1,idy)); 
	ub(n,1:n) = ub(n,1:n) + bfunc(x(n+2,idx),y(n+2,idy));

	% Convert ub to a vector using column reordering
	ub = (1/h^2)*reshape(ub, n*n, 1);

	% Evaluate the RHS of Poisson's equation at the interior points.
	f = pfunc(x(idx,idy),y(idx,idy));

	% Convert f to a vector using column reordering
	f = reshape(f, n*n, 1);

	% Create the D2x and D2y matrices

	% sparse matrix version
	d1 = [ones(n-1, 1); 0];
	d3 = [0; ones(n-1, 1)];
	d2 = -2*ones(n, 1);
	A = spdiags([d1 d2 d3], [-1 0 1], n, n);
	D2x = 1/h^2*kron(A, eye(n));
	D2y = 1/h^2*kron(eye(n), A);

	% Solve the system
	u = (D2x + D2y)\(f-ub);

	% Convert u from a column vector to a matrix to make it easier to work with
	% for plotting.
	u = reshape(u, n, n);

	% Append the boundary values from the Dirichlet condition.
	u = [bfunc(x(1,1:n+2),y(1,1:n+2));...
		 bfunc(x(2:n+1,1),y(2:n+1,1)) u bfunc(x(2:n+1,n+2),y(2:n+1,n+2));...
		 bfunc(x(n+2,1:n+2),y(n+2,1:n+2))];

end


