% stommel.m: Solve Stommel's (1948) model for ocean circulation to
% see the intensified Western Boundary currents

lambda = 10^9;
b = 2*pi*10^8;
D = 2*10^4;
F = 1.0;
R = 0.02;
beta = (10)^-13;

alpha = D*beta/R;
gamma = -pi*F/b/R;

% equation to solve is Laplcian psi + alpha psi_x = gamma sin
% pi*y/b with 0 boundary conditions

% Create grids (NxN)
N = 40;
dx = lambda/N;
dy = b/N;
[x_gr,y_gr] = meshgrid(linspace(0,lambda,N+1),linspace(0,b,N+1));

% Create the right-hand side vector
rhs = gamma.*sin(pi*reshape(y_gr./b,(N+1)^2,1));
size(rhs);

% the laplacian operator has five diagonals.  This does NOT account
% for boundary conditions.
operator = diag((-2/dx^2 -2/dy^2)*ones((N+1)^2,1)) + ...
                diag((1/dy^2)*ones((N+1)^2-1,1),-1) + ...
                diag((1/dy^2)*ones((N+1)^2-1,1),1) + ...
                diag((1/dx^2)*ones((N+1)^2-(N+1),1),N+1) + ...
                diag((1/dx^2)*ones((N+1)^2-(N+1),1),-(N+1));

% the x-derivative has only two diagonals (centered differences)
% Again, no BCs
operator = operator + alpha*(diag((1/2/dx)*ones((N+1)^2-(N+1),1),N+1) + ...
                   diag((-1/2/dx)*ones((N+1)^2-(N+1),1),-(N+1)));

% operator = laplace_operator + alpha*x_deriv_operator;


% Correct the operator and rhs to enforce the boundary conditions
for i=2:N
  rhs(i)   = 0;
  operator(i,i)   = 1;
  operator(i,i-1) = 0;
  operator(i,i+1) = 0;
  operator(i,i+(N+1)) = 0;

  rhs(i+N*(N+1))   = 0;
  operator(i+N*(N+1),i+N*(N+1))   = 1;
  operator(i+N*(N+1),i+N*(N+1)-1) = 0;
  operator(i+N*(N+1),i+N*(N+1)+1) = 0;  
  operator(i+N*(N+1),i+N*(N+1)-(N+1)) = 0;    
end

for i=2:N
  rhs(1+(i-1)*(N+1))   = 0;
  operator(1+(i-1)*(N+1),1+(i-1)*(N+1))   = 1;
  operator(1+(i-1)*(N+1),1+(i-1)*(N+1)-(N+1)) = 0;
  operator(1+(i-1)*(N+1),1+(i-1)*(N+1)+(N+1)) = 0;  
  operator(1+(i-1)*(N+1),2+(i-1)*(N+1)) = 0;    

  rhs((N+1)+(i-1)*(N+1))   = 0;
  operator((N+1)+(i-1)*(N+1),(N+1)+(i-1)*(N+1))   = 1;
  operator((N+1)+(i-1)*(N+1),(N+1)+(i-1)*(N+1)-(N+1)) = 0;
  operator((N+1)+(i-1)*(N+1),(N+1)+(i-1)*(N+1)+(N+1)) = 0;  
  operator((N+1)+(i-1)*(N+1),(N)+(i-1)*(N+1)) = 0;    
end

rhs(1)   = 0;
operator(1,1)      = 1;
operator(1,2)      = 0;
operator(1,N+2)   = 0;

rhs(N+1)   = 0;
operator(N+1,N+1)     = 1;
operator(N+1,N)       = 0;
operator(N+1,2*(N+1)) = 0;

rhs(1+(N)*(N+1))                = 0;
operator(1+(N)*(N+1),1+(N)*(N+1))   = 1;
operator(1+(N)*(N+1),2+(N)*(N+1))   = 0;
operator(1+(N)*(N+1),1+(N-1)*(N+1)) = 0;

rhs((N+1)*(N+1))                = 0;
operator((N+1)*(N+1),(N+1)*(N+1))   = 1;
operator((N+1)*(N+1),(N)*(N+1))     = 0;
operator((N+1)*(N+1),(N+1)*(N+1)-1) = 0;

% Now just solve operator x = rhs

% Using Matlab's built-in solver
%psi_matlab = operator\rhs;

%Using SOR
tol = 1e-15
psi = rhs;
psi_new = zeros(size(psi));
sigma = 0;
omega = 0.9;

err = norm(operator*psi-rhs)
err_vec = [err];
while err >tol
  for i = 1:(N+1)^2
    sigma = operator(i,max(1,i-(N+1)):i-1)*psi_new(max(1,i-(N+1)):i-1) ...
	    + operator(i,i+1:min((N+1)^2,i+(N+1)))*psi(i+1:min((N+1)^2,i+(N+1)));
    sigma = (rhs(i) - sigma)/operator(i,i);
    psi_new(i) = psi(i) + omega*(sigma-psi(i));
  end
  figure(1);mesh(reshape(psi-psi_new,N+1,N+1))
  figure(2);mesh(reshape(psi,N+1,N+1))
  err = norm(operator*psi-rhs)
  err_vec = [err_vec err];
  psi = psi_new;
end


% Plot the solution
psi_gr = reshape(psi,N+1,N+1);
figure(2);contourf(x_gr,y_gr,psi_gr);title('solution?')
[v,u]=gradient(psi_gr,dx,dy);
u = -u;
figure(3);quiver(x_gr,y_gr,u,v,5)