Source code for pyamg.gallery.stencil

"""Construct sparse matrix from a local stencil"""

__docformat__ = "restructuredtext en"

import numpy as np
import scipy.sparse as sparse

__all__ = ['stencil_grid']


[docs]def stencil_grid(S, grid, dtype=None, format=None): """Construct a sparse matrix form a local matrix stencil Parameters ---------- S : ndarray matrix stencil stored in N-d array grid : tuple tuple containing the N grid dimensions dtype : data type of the result format : string sparse matrix format to return, e.g. "csr", "coo", etc. Returns ------- A : sparse matrix Sparse matrix which represents the operator given by applying stencil S at each vertex of a regular grid with given dimensions. Notes ----- The grid vertices are enumerated as arange(prod(grid)).reshape(grid). This implies that the last grid dimension cycles fastest, while the first dimension cycles slowest. For example, if grid=(2,3) then the grid vertices are ordered as (0,0), (0,1), (0,2), (1,0), (1,1), (1,2). This coincides with the ordering used by the NumPy functions ndenumerate() and mgrid(). Examples -------- >>> from pyamg.gallery import stencil_grid >>> stencil = [-1,2,-1] # 1D Poisson stencil >>> grid = (5,) # 1D grid with 5 vertices >>> A = stencil_grid(stencil, grid, dtype=float, format='csr') >>> A.todense() matrix([[ 2., -1., 0., 0., 0.], [-1., 2., -1., 0., 0.], [ 0., -1., 2., -1., 0.], [ 0., 0., -1., 2., -1.], [ 0., 0., 0., -1., 2.]]) >>> stencil = [[0,-1,0],[-1,4,-1],[0,-1,0]] # 2D Poisson stencil >>> grid = (3,3) # 2D grid with shape 3x3 >>> A = stencil_grid(stencil, grid, dtype=float, format='csr') >>> A.todense() matrix([[ 4., -1., 0., -1., 0., 0., 0., 0., 0.], [-1., 4., -1., 0., -1., 0., 0., 0., 0.], [ 0., -1., 4., 0., 0., -1., 0., 0., 0.], [-1., 0., 0., 4., -1., 0., -1., 0., 0.], [ 0., -1., 0., -1., 4., -1., 0., -1., 0.], [ 0., 0., -1., 0., -1., 4., 0., 0., -1.], [ 0., 0., 0., -1., 0., 0., 4., -1., 0.], [ 0., 0., 0., 0., -1., 0., -1., 4., -1.], [ 0., 0., 0., 0., 0., -1., 0., -1., 4.]]) """ S = np.asarray(S, dtype=dtype) grid = tuple(grid) if not (np.asarray(S.shape) % 2 == 1).all(): raise ValueError('all stencil dimensions must be odd') if len(grid) != np.ndim(S): raise ValueError('stencil dimension must equal number of grid\ dimensions') if min(grid) < 1: raise ValueError('grid dimensions must be positive') N_v = np.prod(grid) # number of vertices in the mesh N_s = (S != 0).sum() # number of nonzero stencil entries # diagonal offsets diags = np.zeros(N_s, dtype=int) # compute index offset of each dof within the stencil strides = np.cumprod([1] + list(reversed(grid)))[:-1] indices = tuple(i.copy() for i in S.nonzero()) for i, s in zip(indices, S.shape): i -= s // 2 # i = (i - s) // 2 # i = i // 2 # i = i - (s // 2) for stride, coords in zip(strides, reversed(indices)): diags += stride * coords data = S[S != 0].repeat(N_v).reshape(N_s, N_v) indices = np.vstack(indices).T # zero boundary connections for index, diag in zip(indices, data): diag = diag.reshape(grid) for n, i in enumerate(index): if i > 0: s = [slice(None)] * len(grid) s[n] = slice(0, i) diag[s] = 0 elif i < 0: s = [slice(None)]*len(grid) s[n] = slice(i, None) diag[s] = 0 # remove diagonals that lie outside matrix mask = abs(diags) < N_v if not mask.all(): diags = diags[mask] data = data[mask] # sum duplicate diagonals if len(np.unique(diags)) != len(diags): new_diags = np.unique(diags) new_data = np.zeros((len(new_diags), data.shape[1]), dtype=data.dtype) for dia, dat in zip(diags, data): n = np.searchsorted(new_diags, dia) new_data[n, :] += dat diags = new_diags data = new_data return sparse.dia_matrix((data, diags), shape=(N_v, N_v)).asformat(format)
if __name__ == '__main__': D = 2 if D == 1: # 1D Laplacian S = np.array([-1, 2, -1]) grid = (4,) if D == 2: # 2D Laplacian S = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]]) # S = array([[-1, -1, -1], # [-1, 8, -1], # [-1, -1, -1]]) grid = (2, 1) if D == 3: S = np.array([[[0, 0, 0], [0, -1, 0], [0, 0, 0]], [[0, -1, 0], [-1, 6, -1], [0, -1, 0]], [[0, 0, 0], [0, -1, 0], [0, 0, 0]]]) grid = (3, 4, 5) A = stencil_grid(S, grid) print A.todense()