diff --git a/SimPEG/Examples/Utils_surface2ind_topo.py b/SimPEG/Examples/Utils_surface2ind_topo.py new file mode 100644 index 00000000..61a0f937 --- /dev/null +++ b/SimPEG/Examples/Utils_surface2ind_topo.py @@ -0,0 +1,42 @@ +from SimPEG import * +from SimPEG.Utils import surface2ind_topo + + +def run(plotIt=False, nx = 5, ny = 5): + """ + Here we show how to use :code:`Utils.surface2ind_topo` to identify cells below + a topographic surface. + + """ + + mesh = Mesh.TensorMesh([nx,ny], x0='CC') # 2D mesh + xtopo = np.linspace(mesh.gridN[:,0].min(), mesh.gridN[:,0].max()) + topo = 0.4*np.sin(xtopo*5) # define a topographic surface + + Topo = np.hstack([Utils.mkvc(xtopo,2),Utils.mkvc(topo,2)]) #make it an array + + indcc = surface2ind_topo(mesh, Topo,'CC') + + if plotIt: + from matplotlib.pylab import plt + from scipy.interpolate import interp1d + fig, ax = plt.subplots(1,1,figsize=(6,6)) + mesh.plotGrid(ax=ax, nodes=True, centers=True) + ax.plot(xtopo,topo,'k',linewidth=1) + # ax.plot(mesh.vectorNx, interp1d(xtopo,topo)(mesh.vectorNx),'--k',linewidth=3) + ax.plot(mesh.vectorCCx, interp1d(xtopo,topo)(mesh.vectorCCx),'--k',linewidth=3) + + + aveN2CC = Utils.sdiag(mesh.aveN2CC.T.sum(1))*mesh.aveN2CC.T + a = aveN2CC * indcc + a[a > 0] = 1. + a[a < 0.25] = np.nan + a = a.reshape(mesh.vnN, order='F') + masked_array = np.ma.array(a, mask=np.isnan(a)) + ax.pcolor(mesh.vectorNx,mesh.vectorNy,masked_array.T, cmap = plt.cm.gray,alpha=0.2) + plt.show() + + + +if __name__ == '__main__': + run(plotIt=True) diff --git a/SimPEG/Utils/__init__.py b/SimPEG/Utils/__init__.py index 18c1994f..c7597fee 100644 --- a/SimPEG/Utils/__init__.py +++ b/SimPEG/Utils/__init__.py @@ -7,3 +7,4 @@ from CounterUtils import * import ModelBuilder import SolverUtils from coordutils import * +from modelutils import * diff --git a/SimPEG/Utils/modelutils.py b/SimPEG/Utils/modelutils.py new file mode 100644 index 00000000..dad92fae --- /dev/null +++ b/SimPEG/Utils/modelutils.py @@ -0,0 +1,63 @@ +from matutils import mkvc, ndgrid +import numpy as np + +def surface2ind_topo(mesh, topo, gridLoc='CC'): +# def genActiveindfromTopo(mesh, topo): + """ + Get active indices from topography + """ + + + if mesh.dim == 3: + from scipy.interpolate import NearestNDInterpolator + Ftopo = NearestNDInterpolator(topo[:,:2], topo[:,2]) + + if gridLoc == 'CC': + XY = ndgrid(mesh.vectorCCx, mesh.vectorCCy) + Zcc = mesh.gridCC[:,2].reshape((np.prod(mesh.vnC[:2]), mesh.nCz), order='F') + + gridTopo = Ftopo(XY) + actind = [gridTopo[ixy] <= Zcc[ixy,:] for ixy in range(np.prod(mesh.vnC[0]))] + actind = np.hstack(actind) + + elif gridLoc == 'N': + + XY = ndgrid(mesh.vectorNx, mesh.vectorNy) + gridTopo = Ftopo(XY).reshape(mesh.vnN[:2], order='F') + + if mesh._meshType not in ['TENSOR', 'CYL', 'BASETENSOR']: + raise NotImplementedError('Nodal surface2ind_topo not implemented for %s mesh'%mesh._meshType) + + Nz = mesh.vectorNz[1:] # TODO: this will only work for tensor meshes + actind = np.array([False]*mesh.nC).reshape(mesh.vnC, order='F') + + for ii in range(mesh.nCx): + for jj in range(mesh.nCy): + actind[ii,jj,:] = [np.all(gridTopo[ii:ii+2, jj:jj+2] >= Nz[kk]) for kk in range(len(Nz)) ] + + elif mesh.dim == 2: + from scipy.interpolate import interp1d + Ftopo = interp1d(topo[:,0], topo[:,1]) + + if gridLoc == 'CC': + gridTopo = Ftopo(mesh.gridCC[:,0]) + actind = mesh.gridCC[:,1] <= gridTopo + + elif gridLoc == 'N': + + gridTopo = Ftopo(mesh.vectorNx) + if mesh._meshType not in ['TENSOR', 'CYL', 'BASETENSOR']: + raise NotImplementedError('Nodal surface2ind_topo not implemented for %s mesh'%mesh._meshType) + + Ny = mesh.vectorNy[1:] # TODO: this will only work for tensor meshes + actind = np.array([False]*mesh.nC).reshape(mesh.vnC, order='F') + + for ii in range(mesh.nCx): + actind[ii,:] = [np.all(gridTopo[ii:ii+2] > Ny[kk]) for kk in range(len(Ny)) ] + + else: + raise NotImplementedError('surface2ind_topo not implemented for 1D mesh') + + return mkvc(actind) + +