numpy method to do ndarray to vector mapping as in Matlab's delsq demo?

Question

It Matlab you can easily set up a numbered grid and map operators to 'vectorized' form and back via various indexing tricks. It is not too hard to write this more explicitly but I am wondering if there is a numpy method for doing this? Perhaps via numpy.nditer?

This bit of code (see the mathworks example) gets you from vector u to Matrix U using the ordering stored in the matrix grid G:

U = G;
U(G>0) = full(u(G(G>0)));

Reference:

http://www.mathworks.com/products/matlab/examples.html?file=/products/demos/shipping/matlab/delsqdemo.html

Answer 1

The following reproduces the Demo. The numbering in G is different, but the numbers are just labels (the labeled grid that puzzled me).

import numpy as np
from scipy import sparse
from scipy.sparse import linalg
import matplotlib.pyplot as plt

def numgrid(n):
    """
    NUMGRID Number the grid points in a two dimensional region.
    G = NUMGRID('R',n) numbers the points on an n-by-n grid in
    an L-shaped domain made from 3/4 of the entire square.
    adapted from C. Moler, 7-16-91, 12-22-93.
    Copyright (c) 1984-94 by The MathWorks, Inc.
    """
    x = np.ones((n,1))*np.linspace(-1,1,n)
    y = np.flipud(x.T)
    G = (x > -1) & (x < 1) & (y > -1) & (y < 1) & ( (x > 0) | (y > 0))
    G = np.where(G,1,0) # boolean to integer
    k = np.where(G)
    G[k] = 1+np.arange(len(k[0]))
    return G

def delsq(G):
    """
    DELSQ  Construct five-point finite difference Laplacian.
    delsq(G) is the sparse form of the two-dimensional,
    5-point discrete negative Laplacian on the grid G.
    adapted from  C. Moler, 7-16-91.
    Copyright (c) 1984-94 by The MathWorks, Inc.
    """
    [m,n] = G.shape
    # Indices of interior points
    G1 = G.flatten()
    p = np.where(G1)[0]
    N = len(p)
    # Connect interior points to themselves with 4's.
    i = G1[p]-1
    j = G1[p]-1
    s = 4*np.ones(p.shape)

    # for k = north, east, south, west
    for k in [-1, m, 1, -m]:
       # Possible neighbors in k-th direction
       Q = G1[p+k]
       # Index of points with interior neighbors
       q = np.where(Q)[0]
       # Connect interior points to neighbors with -1's.
       i = np.concatenate([i, G1[p[q]]-1])
       j = np.concatenate([j,Q[q]-1])
       s = np.concatenate([s,-np.ones(q.shape)])
    # sparse matrix with 5 diagonals
    return sparse.csr_matrix((s, (i,j)),(N,N))

if __name__ == '__main__':
    print numgrid(6)
    print delsq(numgrid(6)).todense()
    G = numgrid(32)
    D = delsq(G)
    N = D.shape[0]
    rhs = np.ones((N,1))
    u = linalg.spsolve(D, rhs) # vector solution
    U = np.zeros(G.shape) # map u back onto G space
    U[G>0] = u[G[G>0]-1]
    plt.contour(U)
    plt.show()

Result: