global_ocean.128x64x15/diags_matlab/mit_barostream.m

function [psi, psimask] = mit_barostream(u,varargin)
%function [psi, psimask] = mit_barostream(u,umask,dy,dz);
% or
%function [psi, psimask] = mit_barostream(u,gridinformation);
% global barotropic stream function for mitgcm model, time slabs are
% handled as cell objects, integration from north to the south (by
% convention). 

  if nargin == 2
    g = varargin{1};
    umask = g.umask;
    dy = g.dyg;
    dz = g.dz;
  elseif nargin == 4
    umask = varargin{1};
    dy = varargin{2};
    dz = varargin{3};
  else
    error('need 2 (one of which is the grid structure) or 4 arguments')
  end
  
  [nx ny nz] = size(umask);
  for kz=1:nz
    dydzs(:,:,kz) = (umask(:,:,kz).*dy)*dz(kz);
  end

  % mask for stream function
  pmask = change(squeeze(umask(:,:,1)),'==',NaN,0);
  % add psi-point to the north of all wet points
  pmask(1:nx,2:ny) = pmask(1:nx,2:ny)+pmask(1:nx,1:ny-1);
  pmask = change(pmask,'==',0,NaN);
  pmask = change(pmask,'~=',NaN,1);
  % integrate from the north to the south (by convention), change
  % integration direction by flipping the array ubar (transposed because
  % of MITgcm conventions)
  if iscell(u)
    nt = length(u);
    psi = cell(size(u));
    for k=1:nt
      udxdz = change(u{k}.*dydzs,'==',NaN,0);
      ubar = squeeze(sum(udxdz,3));
      psi{k} = fliplr(cumsum(fliplr(ubar),2)).*pmask;
    end
  else
    nt = size(u,4);
    psi = repmat(NaN,[nx ny nt]);
    udxdz = change(u.*repmat(dydzs,[1 1 1 nt]),'==',NaN,0);
    ubar = squeeze(sum(udxdz,3));
    for kt = 1:nt
      psi(:,:,kt) = fliplr(cumsum(fliplr(squeeze(ubar(:,:,kt))),2)).*pmask;
    end
  end
  if nargout == 2
    psimask = pmask;
  end
    
  return
1	function [psi, psimask] = mit_barostream(u,varargin)
2	%function [psi, psimask] = mit_barostream(u,umask,dy,dz);
3	% or
4	%function [psi, psimask] = mit_barostream(u,gridinformation);
5	% global barotropic stream function for mitgcm model, time slabs are
6	% handled as cell objects, integration from north to the south (by
7	% convention).
8
9	if nargin == 2
10	g = varargin{1};
11	umask = g.umask;
12	dy = g.dyg;
13	dz = g.dz;
14	elseif nargin == 4
15	umask = varargin{1};
16	dy = varargin{2};
17	dz = varargin{3};
18	else
19	error('need 2 (one of which is the grid structure) or 4 arguments')
20	end
21
22	[nx ny nz] = size(umask);
23	for kz=1:nz
24	dydzs(:,:,kz) = (umask(:,:,kz).dy)dz(kz);
25	end
26
27	% mask for stream function
28	pmask = change(squeeze(umask(:,:,1)),'==',NaN,0);
29	% add psi-point to the north of all wet points
30	pmask(1:nx,2:ny) = pmask(1:nx,2:ny)+pmask(1:nx,1:ny-1);
31	pmask = change(pmask,'==',0,NaN);
32	pmask = change(pmask,'~=',NaN,1);
33	% integrate from the north to the south (by convention), change
34	% integration direction by flipping the array ubar (transposed because
35	% of MITgcm conventions)
36	if iscell(u)
37	nt = length(u);
38	psi = cell(size(u));
39	for k=1:nt
40	udxdz = change(u{k}.*dydzs,'==',NaN,0);
41	ubar = squeeze(sum(udxdz,3));
42	psi{k} = fliplr(cumsum(fliplr(ubar),2)).*pmask;
43	end
44	else
45	nt = size(u,4);
46	psi = repmat(NaN,[nx ny nt]);
47	udxdz = change(u.*repmat(dydzs,[1 1 1 nt]),'==',NaN,0);
48	ubar = squeeze(sum(udxdz,3));
49	for kt = 1:nt
50	psi(:,:,kt) = fliplr(cumsum(fliplr(squeeze(ubar(:,:,kt))),2)).*pmask;
51	end
52	end
53	if nargout == 2
54	psimask = pmask;
55	end
56
57	return