Grids in two space dimensions

AMath 585, Winter Quarter 2020 at the University of Washington. Developed by R.J. LeVeque and distributed under the BSD license. You are free to modify and use as you please, with attribution.

These notebooks are all available on Github.

This notebook illustrates how to define a function $u(x,y)$ on a two-dimensional grid, with some demos of different ways to plot such functions.

In [1]:
%matplotlib inline
In [2]:
from pylab import *
from scipy import sparse

Defining 2D grid functions

In two space dimensions we will be approximating functions of $(x,y)$ by arrays with two indices so that $U_{ij} \approx u(x_i, y_j)$. So it is natural to define a 2D numpy array U so that U[i,j] = $U_{ij}$.

First suppose we have some function like $u(x,y) = x + 10y$ that we want to approximate on a grid with mx=9 points in x on the interval $0\leq x\leq 1$ and my=3 interior points in y on the interval $5\leq y \leq 6$. Then we can do the following:

In [3]:
mx = 9
x = linspace(0,1,mx+2)
print('x = ',x)

my = 3
y = linspace(5,6,my+2)
print('y = ',y)

# turn into 2D arrays:
X,Y = meshgrid(x,y,indexing='ij')

print('X = \n', X)
print('Y = \n', Y)
print()

x =  [0.  0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1. ]
y =  [5.   5.25 5.5  5.75 6.  ]
X = 
 [[0.  0.  0.  0.  0. ]
 [0.1 0.1 0.1 0.1 0.1]
 [0.2 0.2 0.2 0.2 0.2]
 [0.3 0.3 0.3 0.3 0.3]
 [0.4 0.4 0.4 0.4 0.4]
 [0.5 0.5 0.5 0.5 0.5]
 [0.6 0.6 0.6 0.6 0.6]
 [0.7 0.7 0.7 0.7 0.7]
 [0.8 0.8 0.8 0.8 0.8]
 [0.9 0.9 0.9 0.9 0.9]
 [1.  1.  1.  1.  1. ]]
Y = 
 [[5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]
 [5.   5.25 5.5  5.75 6.  ]]

Note that the arrays X,Y are defined so that (X[i,j], Y[i,j]) = (x[i], y[j]) is the (i,j) grid point:

In [4]:
i = 2
j = 1
print('x[i] = ',x[i], '  y[j] = ', y[j])
print('X[i,j] = ',X[i,j], '  Y[i,j] = ',Y[i,j])

x[i] =  0.2   y[j] =  5.25
X[i,j] =  0.2   Y[i,j] =  5.25

Having 2D arrays X,Y makes it easier to evaluate a function at all the grid points, e.g.

In [5]:
U = X + 10*Y
print('U = \n', U)

U = 
 [[50.  52.5 55.  57.5 60. ]
 [50.1 52.6 55.1 57.6 60.1]
 [50.2 52.7 55.2 57.7 60.2]
 [50.3 52.8 55.3 57.8 60.3]
 [50.4 52.9 55.4 57.9 60.4]
 [50.5 53.  55.5 58.  60.5]
 [50.6 53.1 55.6 58.1 60.6]
 [50.7 53.2 55.7 58.2 60.7]
 [50.8 53.3 55.8 58.3 60.8]
 [50.9 53.4 55.9 58.4 60.9]
 [51.  53.5 56.  58.5 61. ]]
In [6]:
i=5
j=3
print('x[i] = ',x[i], '  y[j] = ', y[j])
print('U[i,j] = ', U[i,j])

x[i] =  0.5   y[j] =  5.75
U[i,j] =  58.0

Plotting the grid lines / points

We can use the fact that plot(X,Y) plots each column of X against the corresponding column of Y to produce one set of grid lines, and then transpose the matrices to plot the others, e.g.

In [7]:
figure(figsize=(12,4))

subplot(1,3,1)
plot(X,Y,'k')

subplot(1,3,2)
plot(X.T,Y.T,'k')

subplot(1,3,3)
plot(X,Y,'k')
plot(X.T,Y.T,'k')

# put red dots at the interior points:
plot(X[1:-1,1:-1],Y[1:-1,1:-1],'ro');

Make a little function to plot grids this way:

We also use axis('scaled') so that squares turn out square (one unit in x is the same length as in y).

In [8]:
def plot_grid(X,Y):
    plot(X,Y,'k')
    plot(X.T,Y.T,'k')
    #plot(X[1:-1,1:-1],Y[1:-1,1:-1],'ro') # for the red dots
    axis('scaled')

Contour plots

The matplotlib contour command can be used to make contour plots:

In [9]:
plot_grid(X,Y)
contour(X,Y,U,range(50,62,1),colors='g');

Filled contours

The contourf command can be used to plot filled contours. Here we use the default colormap:

In [10]:
plot_grid(X,Y)
contourf(X,Y,U,range(50,63,2))
colorbar()
Out[10]:
<matplotlib.colorbar.Colorbar at 0x1190d0c88>

Many other colormaps are available or you can generate your own. See in particular:

In [11]:
cmap_cool=get_cmap('cool')  # uses function from matplotlib.pyplot
In [12]:
plot_grid(X,Y)
contourf(X,Y,U,range(50,63,2), cmap=cmap_cool)
colorbar();

Choice of colormaps

For many years "jet" was the default colormap in Matlab and also Matplotlib, but this is known to be a poor choice for many reasons. Depending on what you are trying to show, the choice of a good colormap can be critical. See for example,

In [13]:
plot_grid(X,Y)
contourf(X,Y,U,range(50,63,2), cmap=get_cmap('jet'))
colorbar();
title('With jet colormap');

pcolor plots

The pcolor and pcolormesh functions are similar but the latter is more efficient if the grid is uniform. By default it gives uniform colors or patches defined by X,Y, but perhaps not what you expect to see. Although X, Y, U are all $11\times 5$ arrays, pcolor gives a $10\times 4$ array of patches, colored according to the value of U at the bottom left corner of a patch (see the documentation).

In [14]:
pcolormesh(X,Y,U,edgecolors='r')
colorbar();

With more work you can get pcolormesh to plot patches that are centered at the X,Y points. In particular, X,Y have to be one row and column larger than U in order to use all the data in U.

In [15]:
dx = x[1] - x[0]
dy = y[1] - y[0]
mx = len(x) - 2
my = len(y) - 2
xcells = linspace(x[0]-dx/2, x[-1]+dx/2, mx+3)
ycells = linspace(y[0]-dy/2, y[-1]+dy/2, my+3)
Xcells,Ycells = meshgrid(xcells,ycells,indexing='ij')
pcolor(Xcells,Ycells,U,edgecolors='r')
plot(X,Y,'ro')
#axis('scaled')

# plot a white box around the actual domain,
# assuming `X,Y` points lie on boundary:
plot([x[0],x[-1],x[-1],x[0],x[0]],  [y[0],y[0],y[-1],y[-1],y[0]], 'w');

# restrict plot to interior points if desired:
#xlim(x[0],x[-1])
#ylim(y[0],y[-1])

print('Xcells.shape = ', Xcells.shape)
print('U.shape = ', U.shape)

Xcells.shape =  (12, 6)
U.shape =  (11, 5)

For our plotting purposes it's easiest to use contourf.