MATLAB与Python实战——二维向量与等高线的绘制

代码前部分是求解偏微分方程，后面一部分是绘图。

在绘图的过程中，主要难点有两个：

• 第一， 如何使用python或MATLAB绘制等高线以及热力图，并把它们叠加在一张图上。
• 第二，如何使用python或MATLAB绘制二维的向量图。

MATLAB实现

代码：

% Solution of 2D Stokes and continuity equations with finite differences
% on a regular grid using stream function - vorticity formulation
% for a medium with constant viscosity

% Clean all variables
clear all;
% Clear all figures
% clf;

% Numerical model parameters
% Model size, m
xsize=1000000; % Horizontal
ysize=1500000; % Vertical

% Numbers of nodes
xnum=41; % Horizontal
ynum=31; % Vertical
% Grid step
xstp=xsize/(xnum-1); % Horizontal
ystp=ysize/(ynum-1); % Vertical

% Model viscosity
eta=1e+21;

% Gravity acceleration directed downward
gy=10; % m/s^2

% Making vectors for nodal points positions
x=0:xstp:xsize; % Horizontal
y=0:ystp:ysize; % Vertical

% Creating array for density structure (two vertical layers)
rho=zeros(ynum,xnum);
for i=1:1:ynum
  for j=1:1:xnum
    % Horizontal position of the nodal point
    if(x(j)<xsize/2)
        rho(i,j)=3200;  % left layer
    else
        rho(i,j)=3300;  % right layer
    end
  end
end

% Matrix of coefficients initialization
L=sparse(xnum*ynum,xnum*ynum);
% L = zeros(xnum*ynum,xnum*ynum);
% Vector of right part initialization
R=zeros(xnum*ynum,1);

% Solving Poisson equation for worticity
% d2OMEGA/dx2+d2OMEGA/dy2=gy*dRHO/dx
% Composing matrix of coefficients L()
% and vector (column) of right parts R()
% Boundary conditions: OMEGA=0
% Process all Grid points
for i=1:1:ynum
  for j=1:1:xnum
    % Global index for current node
    k=(j-1)*ynum+i;
%     if i == 1
%         k
%     end
    %Boundary nodes OMEGA=0
    if(i==1 || i==ynum || j==1 || j==xnum)
        L(k,k)=1;
        R(k,1)=0;
    %Internal nodes
    else
        %Left part: d2OMEGA/dx2+d2OMEGA/dy2
        L(k,k-ynum)=1/xstp^2;
        L(k,k-1)=1/ystp^2;
        L(k,k)=-2/xstp^2-2/ystp^2;
        L(k,k+1)=1/ystp^2;
        L(k,k+ynum)=1/xstp^2;
        % Right part: gy*dRHO/dx
        R(k,1)=gy/eta*(rho(i,j+1)-rho(i,j-1))/2/xstp;
    end
  end
end
   
%Obtaining vector of solutions S()
S=L\R;

% Reload solutions S() to 2D vorticity array OMEGA()
OMEGA=zeros(ynum,xnum);
% Process all Grid points
for i=1:1:ynum
  for j=1:1:xnum
    % Global index for current node
    k=(j-1)*ynum+i;
    OMEGA(i,j)=S(k);
  end
end

% Solving Poisson equation for stream function
% d2PSI/dx2+d2PSI/dy2=OMEGA
% Simplified procedure as below is only possible 
% when boundary conditions for OMEGA and PSI are the same
% othervise i,j cycle is needed as in the previous case
%L=L; % Left parts remain the same
R=S; % Creates right parts from previous solution

%Obtaining vector of solutions S()
S=L\R;

% Reload solutions S() to 2D stream function array PSI()
PSI=zeros(ynum,xnum);
% Process all Grid points
for i=1:1:ynum
  for j=1:1:xnum
    % Global index for current node
    k=(j-1)*ynum+i;
    PSI(i,j)=S(k);
  end
end


% Compute vx,vy for internal nodes
vx=zeros(ynum,xnum);
vy=zeros(ynum,xnum);
% Process internal Grid points
for i=2:1:ynum-1
  for j=2:1:xnum-1
    % vx=dPSI/dy
    vx(i,j)=(PSI(i+1,j)-PSI(i-1,j))/2/xstp;
     % vy=-dPSI/dx
    vy(i,j)=-(PSI(i,j+1)-PSI(i,j-1))/2/xstp;
  end
end


%Plotting solution
% Making new figure
figure;

% Plotting vorticity as colormap
subplot(1,2,1);
pcolor(x/1000,y/1000,OMEGA);      % making colormap
shading interp;     % making smooth transitions between colors
colorbar;           % showing colorbar for the map% Plotting Stream function as contours

hold on;            % continuing plotting on the colormap 
[C, h] = contour (x/1000,y/1000,PSI,'k'); % drawing stream funtcion isolines
clabel(C,h,'Color','k');        % labelling the isolines
hold off;           % stop plotting on the colormap

caxis([min(min(OMEGA)) max(max(OMEGA))]); % use color limits for OMEGA
box on;             % making box around the plot
title('worticity, stream function'); % title of the plot
xlabel('x');        % title of horizontal axis
ylabel('y');        % title of vertical axis
axis([0 xsize/1000 0 ysize/1000]); % Making axes limits
axis ij image ;     % direct vertical axis downward, make proper dimensions

% Plotting density as colormap
subplot(1,2,2);
pcolor(x/1000,y/1000,rho);      % making colormap
shading interp;     % making smooth transitions between colors
colorbar;           % showing colorbar for the map

hold on;            % continuing plotting on the colormap 

% Plotting velocity vector as arrows using internal nodes only
quiver(x(2:1:xnum-1)/1000,y(2:1:ynum-1)/1000,vx(2:1:ynum-1,2:1:xnum-1),vy(2:1:ynum-1,2:1:xnum-1),'k'); % making field of arrows

hold off;           % stop plotting on the colormap

box on;             % making box around the plot
title('Density, velocity field');   % title of the plot
xlabel('x');        % title of horizontal axis
ylabel('y');        % title of vertical axis
axis([0 xsize/1000 0 ysize/1000]); % Making axes limits
axis ij image ;     % direct vertical axis downward, make proper dimensions

Python实现

代码：

# -*- coding: utf-8 -*-

# Numerical model parameters
# Model size, m
xsize=1000000; # Horizontal
ysize=1500000; # Vertical

# Numbers of nodes
xnum=41; # Horizontal
ynum=31; # Vertical
# Grid step
xstp=xsize/(xnum-1); # Horizontal
ystp=ysize/(ynum-1); # Vertical

# Model viscosity
eta=1e+21;

# Gravity acceleration directed downward
gy=10; # m/s^2

# Making vectors for nodal points positions
import numpy as np
x = np.arange(0,xsize+1,xstp); # Horizontal
y = np.arange(0,ysize+1,ystp); # Vertical

# Creating array for density structure (two vertical layers)
rho = np.zeros([ynum,xnum]);
for i in range(1,ynum+1):
    for j in range(1,xnum+1):
        # Horizontal position of the nodal point
        if x[j-1] < xsize/2:
            rho[i-1,j-1] = 3200   # left layer
        else:
            rho[i-1,j-1] = 3300   # right layer

# Matrix of coefficients initialization
L = np.zeros([xnum*ynum,xnum*ynum])
# Vector of right part initialization
R = np.zeros([xnum*ynum,1])

# Solving Poisson equation for worticity
# d2OMEGA/dx2+d2OMEGA/dy2=gy*dRHO/dx
# Composing matrix of coefficients L()
# and vector (column) of right parts R()
# Boundary conditions: OMEGA=0
# Process all Grid points
for i in range(1,ynum+1,1):
    for j in range(1,xnum+1,1):
        # Global index for current node
        k = (j-1)*ynum+i
#        print(k)
        # Boundary nodes OMEGA=0
        if i==1 or i==ynum or j==1 or j==xnum:
            L[k-1,k-1] = 1;
            R[k-1,0] = 0;
        # Internal nodes
        else:
            # Left part: d2OMEGA/dx2+d2OMEGA/dy2
            L[k-1,k-ynum-1] = 1/xstp**2
            L[k-1,k-2] = 1/ystp**2
            L[k-1,k-1] = -2/xstp**2 - 2/ystp**2
            L[k-1,k] = 1/ystp**2
            L[k-1,k+ynum-1] = 1/xstp**2
            # Right part: gy*dRHO/dx
            R[k-1,0] = gy/eta*(rho[i-1,j] - rho[i-1,j-2])/2/xstp
        
# Obtaining vector of solutions S()
S = np.dot(np.linalg.inv(L),R)

# Reload solutions S() to 2D vorticity array OMEGA()
OMEGA = np.zeros([ynum,xnum])
# Process all Grid points
for i in range(1,ynum+1,1):
    for j in range(1,xnum+1,1):
        # Global index for current node
        k = (j-1)*ynum+i;
        OMEGA[i-1,j-1] = S[k-1];

# Solving Poisson equation for stream function
# d2PSI/dx2+d2PSI/dy2=OMEGA
# implified procedure as below is only possible 
# when boundary conditions for OMEGA and PSI are the same
# othervise i,j cycle is needed as in the previous case
# L=L; % Left parts remain the same
R = S  # Creates right parts from previous solution

# Obtaining vector of solutions S()
S = np.dot(np.linalg.inv(L),R)

# Reload solutions S() to 2D stream function array PSI()
PSI = np.zeros([ynum,xnum])
# Process all Grid points
for i in range(1,ynum+1,1):
    for j in range(1,xnum+1,1):
        # Global index for current node
        k = (j-1)*ynum+i
        PSI[i-1,j-1] = S[k-1]

# Compute vx,vy for internal nodes
vx = np.zeros([ynum,xnum])
vy = np.zeros([ynum,xnum])
# Process internal Grid points
for i in range(2,ynum,1):
    for j in range(2,xnum,1):
        # vx=dPSI/dy
        vx[i-1,j-1] = (PSI[i,j-1] - PSI[i-2,j-1])/2/xstp
        # vy=-dPSI/dx
        vy[i-1,j-1] = -(PSI[i-1,j] - PSI[i-1,j-2])/2/xstp

# Plotting solution
import matplotlib.pyplot as plt 

plt.figure()  # Making new figure
plt.pcolormesh(x/1000, y/1000, OMEGA,shading='gouraud')
plt.colorbar()  # showing colorbar for the map% Plotting Stream function as contours

# drawing stream funtcion isolines
C = plt.contour(x/1000, y/1000, PSI, 7, colors='black', linewidth=0.25, alpha=0.5)
plt.clabel(C, inline=True, fontsize=10)  # 添加文字标签 inlins表示等高线是穿过数字还是不穿过
plt.title('worticity, stream function')
plt.xlabel('x');        # title of horizontal axis
plt.ylabel('y');        # title of vertical axis
plt.axis([0,xsize/1000,0,ysize/1000])   # Making axes limits
plt.show()

# Plotting density as colormap
from pylab import quiver

plt.figure()
plt.pcolormesh(x/1000, y/1000, rho, shading='gouraud')
plt.colorbar()  # showing colorbar for the map% Plotting Stream function as contours

# Plotting velocity vector as arrows using internal nodes only
quiver(x[np.arange(2,xnum)]/1000,y[np.arange(2,ynum)]/1000,vx[2:ynum,2:xnum],vy[2:ynum,2:xnum]) # making field of arrows
plt.title('Density, velocity field')   # title of the plot
plt.xlabel('x')        # title of horizontal axis
plt.ylabel('y')        # title of vertical axis
plt.axis([0,xsize/1000,0,ysize/1000]) # Making axes limits
plt.show()