function hw9 ( )

%*****************************************************************************80
%
%% hw9 models heat advection in a stream function/vorticity flow.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    26 June 2020
%
%  Author:
%
%    John Burkardt
%
  fprintf ( 1, '\n' );
  fprintf ( 1, 'hw9:\n' );
  fprintf ( 1, '  Solve the driven cavity problem using the stream\n' );
  fprintf ( 1, '  function/vorticity formulation, with heat advection.\n' );
%
%  Space discretization.
%
  nr = 21;
  nc = 21;
  [ X, Y ] = xy_square ( nr, nc );
  dy = 1.0 / ( nr - 1 );
  dx = 1.0 / ( nc - 1 );
%
%  Time discretization.
%
  nt = 101;
  tmin = 0.0;
  tmax = 0.10;
  t = linspace ( tmin, tmax, nt );
  dt = ( tmax - tmin ) / ( nt - 1 );
%
%  The fluid viscosity.
%
  nu = 0.2;
%
%  The horizontal velocity along the lid.
%
  utop = 1.0;
%
%  Get matrix for the equation
%    Del^2 PSI = - OMEGA
%  On boundary nodes, we specify PSI=0 (walls) or PSI=slid (top).
%
  A = psi_laplacian ( nr, nc, dx, dy );
%
%  Get symbolic indices for nodes in the 5-point stencil for OMEGA.
%
  [ C, N, S, E, W ] = omega_nodes ( nr, nc );
%
%  Time loop.
%
  for k = 1 : nt
%
%  Stream function.
%
    if ( k == 1 )

      SN = zeros ( nr, nc );
      SN(1,1:nc) = 2.0 * dx * utop;

    else

      SO = SN;
      on = - ON;
      for i = 1 : nr
        for j = 1 : nc
          if ( i == 1 )
            on(i,j) = 2.0 * dx * utop;
          elseif ( i == nr | j == 1 | j == nc )
            on(i,j) = 0.0;
          end
        end
      end

      on = reshape ( on, nr * nc, 1 );
      sn = A \ on;
      SN = reshape ( sn, nr, nc );

    end
%
%  Vorticity.
%
    if ( k == 1 )

      ON = zeros ( nr, nc );

    else

      OO = ON;
      ON(C) = OO(C) + dt * ( ...
        - ( 1.0 / dx / dy / 4.0 ) * ( SN(E) - SN(W) ) .* ( OO(N) - OO(S) ) ...
        + ( 1.0 / dy / dx / 4.0 ) * ( SN(N) - SN(S) ) .* ( OO(E) - OO(W) ) ...
        + ( nu / dx^2 ) * ( OO(E) - 2.0 * OO(C) + OO(W) ) ...
        + ( nu / dy^2 ) * ( OO(N) - 2.0 * OO(C) + OO(S) ) );

    end
%
%  Set a nonzero vorticity along the top, corresponding to the
%  forced horizontal velocity UTOP.
%
    for j = 1 : nc
      ON(1,j) = - 2.0 * ( SN(1,j) - SN(2,j) ) / dy^2 - 2.0 * utop / dy;
    end
%
%  Temperature initial condition is 0, plus a hot blob.
%
    if ( k == 1 )
      TN = zeros ( nr, nc );
      i = floor(2*nr/10);
      j = floor(9*nr/10);
      TN(i-2:i+2,j-2:j+2) = 200.0;
%
%  Update temperature using the advection equation.
%
    else
      TO = TN;
      TN(C) = 0.25 * ( TO(N) + TO(S) + TO(E) + TO(W)) ...
        - ( dt / dx / 2.0 ) * U(C) .* ( TO(E) - TO(W) ) ...
        - ( dt / dy / 2.0 ) * V(C) .* ( TO(N) - TO(S) );
    end
%
%  Plot.
%
    [ U, V ] = uv_from_psi ( X, Y, SN );

    figure ( 1 )

    label = sprintf ( 'Driven cavity, k = %d, t = %g\n', k, t(k) );
    sgtitle ( label )

    subplot ( 2, 2, 1 );
    quiver ( X, Y, U, V );
    title ( 'Velocity field' );
    grid ( 'on' );
    axis ( 'equal' );
    axis ( 'tight' );

    UN = U ./ sqrt ( U.^2 + V.^2 );
    VN = V ./ sqrt ( U.^2 + V.^2 );

    I = find ( ( U.^2 + V.^2 ) == 0 );
    UN(I) = 0.0;
    VN(I) = 0.0;

    subplot ( 2, 2, 2 );
    quiver ( X, Y, UN, VN );
    title ( 'Direction field' );
    grid ( 'on' );
    axis ( 'equal' );
    axis ( 'tight' );

    subplot ( 2, 2, 3 )
    contourf ( X, Y, SN );
    axis ( 'square' );
    title ( 'Stream function \psi(x,y)' );

    subplot ( 2, 2, 4 )
    contourf ( X, Y, TN, 'edgecolor', 'none' );
    axis ( 'square' );
    title ( 'Temperature' );

  end
%
%  Save the last plot.
%
  filename = 'hw9.png';
  print ( '-dpng', filename );
  fprintf ( 1, '  Graphics saved as "%s"\n', filename );

  return
end
function [ U, V ] = uv_from_psi ( X, Y, PSI )

%*****************************************************************************80
%
%% uv_from_psi computes a velocity field (u,v) from a stream function.
%
%  Discussion:
%
%    We assume the x and y grid spacings are each uniform.
%
%    U =   d PSI/dy
%    V = - d PSI/dx
%
%  Modified:
%
%    12 June 2020
%
%  Input:
%
%    real X(nr,nc), Y(nc,nc): the node coordinates.
%
%    real PSI(nr,nc): the stream function.
%
%  Output:
%
%    real U(nr,nc), V(nr,nc): the velocity.
%
  [ nr, nc ] = size ( PSI );

  dx = X(1,2) - X(1,1);
  dy = Y(1,1) - Y(2,1);
  
  U = zeros ( nr, nc );
  V = zeros ( nr, nc );
%
%  For U, use centered differences except in first and last rows.
%  For Y derivatives, the larger Y value is at the SMALLER r value.
%
  for r = 1 : nr
    if ( r == 1 )
      U(r,1:nc) = ( PSI(r,1:nc) - PSI(r+1,1:nc) ) / dy;
    elseif ( r < nr )
      U(r,1:nc) = ( PSI(r-1,1:nc) - PSI(r+1,1:nc) ) / dy / 2.0;
    else
      U(r,1:nc) = ( PSI(r-1,1:nc) - PSI(r,1:nc) ) / dy;
    end
  end
%
%  For V, use centered differences except in first and last columns.
%  For X derivatives, the larger X value is at the larger c value.
%
  for c = 1 : nc
    if ( c == 1 )
      V(1:nr,c) = - ( PSI(1:nr,c+1) - PSI(1:nr,c) ) / dx;
    elseif ( c < nc )
      V(1:nr,c) = - ( PSI(1:nr,c+1) - PSI(1:nr,c-1) ) / dx / 2.0;
    else
      V(1:nr,c) = - ( PSI(1:nr,c) - PSI(1:nr,c-1) ) / dx;
    end
  end

  return
end
function A = psi_laplacian ( nr, nc, dx, dy )

%*****************************************************************************80
%
%% psi_laplacian returns the matrix A associated with Del^2 * PSI = OMEGA.
%
  A = sparse ( [], [], [], nr * nc, nr * nc, 5 * nr * nr );

  C = 0;
  for i = 1 : nr
    for j = 1 : nc
      C = C + 1;
      if ( i == 1 | i == nr | j == 1 | j == nc )
        A(C,C) = 1.0;
      else
        N = C - nc;
        S = C + nc;
        W = C - 1;
        E = C + 1;
        A(C,N) = 1.0 / dy^2;
        A(C,S) = 1.0 / dy^2;
        A(C,W) = 1.0 / dx^2;
        A(C,E) = 1.0 / dx^2;
        A(C,C) = - ( A(C,N) + A(C,S) + A(C,W) + A(C,E) );
      end
    end
  end

  return
end
function [ C, N, S, E, W ] = omega_nodes ( nr, nc )

%*****************************************************************************80
%
%% omega_nodes returns indices for the 5 point Laplacian stencil.
%
%  Discussion:
%
%    We assume the grid is numbered something like this example,
%    for which nr = 5 (five rows) and nc = 4 (four columns):
%
%   |  17 18 19 20
%   |  13 14 15 16
%   I   9 10 11 12
%   |   5  6  7  8
%   |   1  2  3  4
%   |
%   +-------J-----
%
%    We also assume that periodic boundary conditions have been imposed, so
%    that the next node to the right of node 20 is node 17, and the next
%    node above node 20 is 4, for instance.
%
%  Input:
%
%    integer i, j: the grid indices of a point.
%
%    integer nx: the number of nodes in the x direction.
%
%  Output:
%
%    integer C, N, S, E, W: the variable indices for the central,
%    north, south, east, and west nodes.
%
  n = ( nr - 2 ) * ( nc - 2 );

  C = zeros ( 1, n );
  N = zeros ( 1, n );
  S = zeros ( 1, n );
  E = zeros ( 1, n );
  W = zeros ( 1, n );

  k = 0;
  for i = 2 : nr - 1
    for j = 2 : nc - 1
      k = k + 1;
      C(k) = j     + ( i - 1 ) * nc;
      N(k) = j     + ( i     ) * nc;
      S(k) = j     + ( i - 2 ) * nc;
      E(k) = j + 1 + ( i - 1 ) * nc;
      W(k) = j - 1 + ( i - 1 ) * nc;

    end
  end

  return
end

