diff --git a/README.md b/README.md index bfe17acbe2..e06237cfe3 100644 --- a/README.md +++ b/README.md @@ -209,6 +209,8 @@ Additional dependencies to use optional FloPy helper methods are listed below. | `Raster()` in `flopy.utils.Raster` | **rasterio**, **affine**, and **scipy** | | `Raster().sample_polygon()` in `flopy.utils.Raster` | **shapely** | | `Raster().crop()` in `flopy.utils.Raster` | **shapely** | +| `._zverts_smooth()` in `flopy.discretization.structuredgrid` `StructuredGrid` class | **scipy.interpolate** | +| `.array_at_verts()` in `flopy.discretization.structuredgrid` `StructuredGrid` class | **scipy.interpolate** | How to Cite ----------------------------------------------- diff --git a/autotest/t050_test.py b/autotest/t050_test.py index a8de5558ab..64761c221d 100644 --- a/autotest/t050_test.py +++ b/autotest/t050_test.py @@ -38,16 +38,16 @@ def test_vtk_export_array2d(): # export and check m.dis.top.export(output_dir, name='top', fmt='vtk') filetocheck = os.path.join(output_dir, 'top.vtu') - totalbytes = os.path.getsize(filetocheck) - # assert(totalbytes==352026) + # totalbytes = os.path.getsize(filetocheck) + # assert(totalbytes==351846) nlines = count_lines_in_file(filetocheck) assert(nlines==2846) # with smoothing m.dis.top.export(output_dir, fmt='vtk', name='top_smooth', smooth=True) filetocheck = os.path.join(output_dir, 'top_smooth.vtu') - totalbytes1 = os.path.getsize(filetocheck) - # assert(totalbytes1==351829) + # totalbytes1 = os.path.getsize(filetocheck) + # assert(totalbytes1==357587) nlines1 = count_lines_in_file(filetocheck) assert(nlines1==2846) @@ -65,8 +65,8 @@ def test_vtk_export_array3d(): # export and check m.upw.hk.export(output_dir, fmt='vtk', name='hk') filetocheck = os.path.join(output_dir, 'hk.vtu') - totalbytes = os.path.getsize(filetocheck) - # assert(totalbytes==992576) + # totalbytes = os.path.getsize(filetocheck) + # assert(totalbytes==992036) nlines = count_lines_in_file(filetocheck) assert(nlines==8486) @@ -74,8 +74,8 @@ def test_vtk_export_array3d(): m.upw.hk.export(output_dir, fmt='vtk', name='hk_points', point_scalars=True) filetocheck = os.path.join(output_dir, 'hk_points.vtu') - totalbytes1 = os.path.getsize(filetocheck) - # assert(totalbytes1==1321502) + # totalbytes1 = os.path.getsize(filetocheck) + # assert(totalbytes1==1354204) nlines1 = count_lines_in_file(filetocheck) assert(nlines1==10605) @@ -83,10 +83,11 @@ def test_vtk_export_array3d(): m.upw.hk.export(output_dir, fmt='vtk', name='hk_points_bin', point_scalars=True, binary=True) filetocheck = os.path.join(output_dir, 'hk_points_bin.vtu') - totalbytes2 = os.path.getsize(filetocheck) + # totalbytes2 = os.path.getsize(filetocheck) # assert(totalbytes2==629401) - nlines2 = count_lines_in_file(filetocheck, binary=True) - assert(nlines2==1869) + # nlines2 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines2==2257) + assert(os.path.exists(filetocheck)) return @@ -103,28 +104,30 @@ def test_vtk_transient_array_2d(): # export and check m.rch.rech.export(output_dir, fmt='vtk') filetocheck = os.path.join(output_dir, 'rech_01.vtu') - totalbytes = os.path.getsize(filetocheck) - # assert(totalbytes==355324) + # totalbytes = os.path.getsize(filetocheck) + # assert(totalbytes==355144) nlines = count_lines_in_file(filetocheck) assert(nlines==2851) filetocheck = os.path.join(output_dir, 'rech_01097.vtu') - totalbytes1 = os.path.getsize(filetocheck) - # assert(totalbytes1==354622) + # totalbytes1 = os.path.getsize(filetocheck) + # assert(totalbytes1==354442) nlines1 = count_lines_in_file(filetocheck) assert(nlines1==2851) # with binary m.rch.rech.export(output_dir_bin, fmt='vtk', binary=True) filetocheck = os.path.join(output_dir_bin, 'rech_01.vtu') - totalbytes2 = os.path.getsize(filetocheck) + # totalbytes2 = os.path.getsize(filetocheck) # assert(totalbytes2==168339) - nlines2 = count_lines_in_file(filetocheck, binary=True) - assert(nlines2==762) + # nlines2 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines2==846) + assert(os.path.exists(filetocheck)) filetocheck = os.path.join(output_dir_bin, 'rech_01097.vtu') - totalbytes3 = os.path.getsize(filetocheck) + # totalbytes3 = os.path.getsize(filetocheck) # assert(totalbytes3==168339) - nlines3 = count_lines_in_file(filetocheck, binary=True) - assert(nlines3==762) + # nlines3 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines3==846) + assert(os.path.exists(filetocheck)) return @@ -140,8 +143,8 @@ def test_vtk_export_packages(): output_dir = os.path.join(cpth, 'DIS') m.dis.export(output_dir, fmt='vtk') filetocheck = os.path.join(output_dir, 'DIS.vtu') - totalbytes = os.path.getsize(filetocheck) - # assert(totalbytes==1020397) + # totalbytes = os.path.getsize(filetocheck) + # assert(totalbytes==1019857) nlines = count_lines_in_file(filetocheck) assert(nlines==8496) @@ -149,8 +152,8 @@ def test_vtk_export_packages(): output_dir = os.path.join(cpth, 'UPW') m.upw.export(output_dir, fmt='vtk', point_scalars=True) filetocheck = os.path.join(output_dir, 'UPW.vtu') - totalbytes1 = os.path.getsize(filetocheck) - # assert(totalbytes1==2485991) + # totalbytes1 = os.path.getsize(filetocheck) + # assert(totalbytes1==2531309) nlines1 = count_lines_in_file(filetocheck) assert(nlines1==21215) @@ -158,8 +161,8 @@ def test_vtk_export_packages(): output_dir = os.path.join(cpth, 'BAS') m.bas6.export(output_dir, fmt='vtk', smooth=True) filetocheck = os.path.join(output_dir, 'BAS6.vtu') - totalbytes2 = os.path.getsize(filetocheck) - # assert(totalbytes2==1002054) + # totalbytes2 = os.path.getsize(filetocheck) + # assert(totalbytes2==1035056) nlines2 = count_lines_in_file(filetocheck) assert(nlines2==8491) @@ -167,13 +170,13 @@ def test_vtk_export_packages(): output_dir = os.path.join(cpth, 'DRN') m.drn.export(output_dir, fmt='vtk') filetocheck = os.path.join(output_dir, 'DRN_01.vtu') - totalbytes3 = os.path.getsize(filetocheck) - # assert(totalbytes3==20702) + # totalbytes3 = os.path.getsize(filetocheck) + # assert(totalbytes3==20670) nlines3 = count_lines_in_file(filetocheck) assert(nlines3==191) filetocheck = os.path.join(output_dir, 'DRN_01097.vtu') - totalbytes4 = os.path.getsize(filetocheck) - # assert(totalbytes4==20702) + # totalbytes4 = os.path.getsize(filetocheck) + # assert(totalbytes4==20670) nlines4 = count_lines_in_file(filetocheck) assert(nlines4==191) @@ -181,19 +184,21 @@ def test_vtk_export_packages(): output_dir = os.path.join(cpth, 'DIS_bin') m.dis.export(output_dir, fmt='vtk', binary=True) filetocheck = os.path.join(output_dir, 'DIS.vtu') - totalbytes5 = os.path.getsize(filetocheck) + # totalbytes5 = os.path.getsize(filetocheck) # assert(totalbytes5==519516) - nlines5 = count_lines_in_file(filetocheck, binary=True) - assert(nlines5==1545) + # nlines5 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines5==1797) + assert(os.path.exists(filetocheck)) # upw with point scalars and binary output_dir = os.path.join(cpth, 'UPW_bin') m.upw.export(output_dir, fmt='vtk', point_scalars=True, binary=True) filetocheck = os.path.join(output_dir, 'UPW.vtu') - totalbytes6 = os.path.getsize(filetocheck) + # totalbytes6 = os.path.getsize(filetocheck) # assert(totalbytes6==1349801) - nlines6 = count_lines_in_file(filetocheck, binary=True) - assert(nlines6==4004) + # nlines6 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines6==4392) + assert(os.path.exists(filetocheck)) return @@ -217,7 +222,7 @@ def test_vtk_mf6(): # check one filetocheck = os.path.join(cpth, 'twrihfb2015', 'npf.vtr') - totalbytes = os.path.getsize(filetocheck) + # totalbytes = os.path.getsize(filetocheck) # assert(totalbytes==21609) nlines = count_lines_in_file(filetocheck) assert(nlines==76) @@ -244,8 +249,8 @@ def test_vtk_binary_head_export(): (0, 1089)]) filetocheck = os.path.join(otfolder, filenametocheck) - totalbytes = os.path.getsize(filetocheck) - # assert(totalbytes==993755) + # totalbytes = os.path.getsize(filetocheck) + # assert(totalbytes==993215) nlines = count_lines_in_file(filetocheck) assert(nlines==8486) @@ -256,8 +261,8 @@ def test_vtk_binary_head_export(): 1089)], point_scalars=True, nanval=-999.99) filetocheck = os.path.join(otfolder, filenametocheck) - totalbytes1 = os.path.getsize(filetocheck) - # assert(totalbytes1==1332153) + # totalbytes1 = os.path.getsize(filetocheck) + # assert(totalbytes1==1365320) nlines1 = count_lines_in_file(filetocheck) assert(nlines1==10605) @@ -268,8 +273,8 @@ def test_vtk_binary_head_export(): 1089)], smooth=True, nanval=-999.99) filetocheck = os.path.join(otfolder, filenametocheck) - totalbytes2 = os.path.getsize(filetocheck) - # assert(totalbytes2==993551) + # totalbytes2 = os.path.getsize(filetocheck) + # assert(totalbytes2==1026553) nlines2 = count_lines_in_file(filetocheck) assert(nlines2==8486) @@ -280,10 +285,11 @@ def test_vtk_binary_head_export(): 1089)], smooth=True, binary=True, nanval=-999.99) filetocheck = os.path.join(otfolder, filenametocheck) - totalbytes3 = os.path.getsize(filetocheck) + # totalbytes3 = os.path.getsize(filetocheck) # assert(totalbytes3==493853) - nlines3 = count_lines_in_file(filetocheck, binary=True) - assert(nlines3==1529) + # nlines3 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines3==1825) + assert(os.path.exists(filetocheck)) # with smoothing and binary, single time otfolder = os.path.join(cpth, 'heads_test_4') @@ -291,10 +297,11 @@ def test_vtk_binary_head_export(): point_scalars=False, smooth=True, binary=True, nanval=-999.99) filetocheck = os.path.join(otfolder, 'freyberg_Heads_KPER1_KSTP1.vtu') - totalbytes4 = os.path.getsize(filetocheck) + # totalbytes4 = os.path.getsize(filetocheck) # assert(totalbytes4==493853) - nlines4 = count_lines_in_file(filetocheck, binary=True) - assert(nlines4==1535) + # nlines4 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines4==1831) + assert(os.path.exists(filetocheck)) return @@ -313,8 +320,8 @@ def test_vtk_cbc(): vtk.export_cbc(m, cbcfile, otfolder, kstpkper=[(0, 0), (0, 1), (0, 2)], point_scalars=True) filetocheck = os.path.join(otfolder, filenametocheck) - totalbytes = os.path.getsize(filetocheck) - # assert(totalbytes==2626880) + # totalbytes = os.path.getsize(filetocheck) + # assert(totalbytes==2664009) nlines = count_lines_in_file(filetocheck) assert(nlines==19093) @@ -324,10 +331,11 @@ def test_vtk_cbc(): kstpkper=[(0, 0), (0, 1), (0, 2)], point_scalars=True, binary=True) filetocheck = os.path.join(otfolder, filenametocheck) - totalbytes1 = os.path.getsize(filetocheck) + # totalbytes1 = os.path.getsize(filetocheck) # assert(totalbytes1==1205818) - nlines1 = count_lines_in_file(filetocheck, binary=True) - assert(nlines1==2514) + # nlines1 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines1==3310) + assert(os.path.exists(filetocheck)) # with point scalars and binary, only one budget component otfolder = os.path.join(cpth, 'freyberg_CBCTEST_bin2') @@ -335,42 +343,72 @@ def test_vtk_cbc(): kstpkper=(0, 0), text='CONSTANT HEAD', point_scalars=True, binary=True) filetocheck = os.path.join(otfolder, filenametocheck) - totalbytes2 = os.path.getsize(filetocheck) + # totalbytes2 = os.path.getsize(filetocheck) # assert(totalbytes2==10142) - nlines2 = count_lines_in_file(filetocheck, binary=True) - assert(nlines2==62) + # nlines2 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines2==66) + assert(os.path.exists(filetocheck)) return def test_vtk_vector(): from flopy.utils import postprocessing as pp + from flopy.utils import geometry # test mf 2005 freyberg mpth = os.path.join('..', 'examples', 'data', 'freyberg_multilayer_transient') namfile = 'freyberg.nam' cbcfile = os.path.join(mpth, 'freyberg.cbc') + hdsfile = os.path.join(mpth, 'freyberg.hds') m = flopy.modflow.Modflow.load(namfile, model_ws=mpth, verbose=False, load_only=['dis', 'bas6']) - q = pp.get_specific_discharge(m, cbcfile=cbcfile) + q = list(pp.get_specific_discharge(m, cbcfile=cbcfile)) + q[0], q[1] = geometry.rotate(q[0], q[1], 0., 0., + m.modelgrid.angrot_radians) output_dir = os.path.join(cpth, 'freyberg_vector') filenametocheck = 'discharge.vtu' # export and check with point scalar vtk.export_vector(m, q, output_dir, 'discharge', point_scalars=True) filetocheck = os.path.join(output_dir, filenametocheck) - totalbytes = os.path.getsize(filetocheck) - # assert(totalbytes==2249214) + # totalbytes = os.path.getsize(filetocheck) + # assert(totalbytes==2281411) nlines = count_lines_in_file(filetocheck) assert(nlines==10605) # with point scalars and binary - vtk.export_vector(m, q, output_dir + '_bin', 'discharge', point_scalars=True, + vtk.export_vector(m, q, output_dir + '_bin', 'discharge', + point_scalars=True, binary=True) + filetocheck = os.path.join(output_dir + '_bin', filenametocheck) + # totalbytes1 = os.path.getsize(filetocheck) + # assert(totalbytes1==942413) + # nlines1 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines1==4014) + assert(os.path.exists(filetocheck)) + + # with values directly given at vertices + q = list(pp.get_specific_discharge(m, cbcfile=cbcfile, hdsfile=hdsfile, + position='vertices')) + q[0], q[1] = geometry.rotate(q[0], q[1], 0., 0., + m.modelgrid.angrot_radians) + output_dir = os.path.join(cpth, 'freyberg_vector') + filenametocheck = 'discharge_verts.vtu' + vtk.export_vector(m, q, output_dir, 'discharge_verts') + filetocheck = os.path.join(output_dir, filenametocheck) + # totalbytes2 = os.path.getsize(filetocheck) + # assert(totalbytes2==2039528) + nlines2 = count_lines_in_file(filetocheck) + assert(nlines2==10598) + + # with values directly given at vertices and binary + vtk.export_vector(m, q, output_dir + '_bin', 'discharge_verts', binary=True) filetocheck = os.path.join(output_dir + '_bin', filenametocheck) - totalbytes1 = os.path.getsize(filetocheck) - # assert(totalbytes1==917033) - nlines1 = count_lines_in_file(filetocheck, binary=True) - assert(nlines1==2725) + # totalbytes3 = os.path.getsize(filetocheck) + # assert(totalbytes3==891486) + # nlines3 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines3==3113) + assert(os.path.exists(filetocheck)) return @@ -385,7 +423,7 @@ def test_vtk_vti(): # export and check m.export(output_dir, fmt='vtk') filetocheck = os.path.join(output_dir, filenametocheck) - totalbytes = os.path.getsize(filetocheck) + # totalbytes = os.path.getsize(filetocheck) # assert(totalbytes==6322) nlines = count_lines_in_file(filetocheck) assert(nlines==21) @@ -393,7 +431,7 @@ def test_vtk_vti(): # with point scalar m.export(output_dir + '_points', fmt='vtk', point_scalars=True) filetocheck = os.path.join(output_dir + '_points', filenametocheck) - totalbytes1 = os.path.getsize(filetocheck) + # totalbytes1 = os.path.getsize(filetocheck) # assert(totalbytes1==16382) nlines1 = count_lines_in_file(filetocheck) assert(nlines1==38) @@ -401,16 +439,17 @@ def test_vtk_vti(): # with binary m.export(output_dir + '_bin', fmt='vtk', binary=True) filetocheck = os.path.join(output_dir + '_bin', filenametocheck) - totalbytes2 = os.path.getsize(filetocheck) + # totalbytes2 = os.path.getsize(filetocheck) # assert(totalbytes2==4617) - nlines2 = count_lines_in_file(filetocheck, binary=True) - assert(nlines2==18) + # nlines2 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines2==18) + assert(os.path.exists(filetocheck)) # force .vtr filenametocheck = 'DIS.vtr' m.export(output_dir, fmt='vtk', vtk_grid_type='RectilinearGrid') filetocheck = os.path.join(output_dir, filenametocheck) - totalbytes3 = os.path.getsize(filetocheck) + # totalbytes3 = os.path.getsize(filetocheck) # assert(totalbytes3==7146) nlines3 = count_lines_in_file(filetocheck) assert(nlines3==56) @@ -419,7 +458,7 @@ def test_vtk_vti(): filenametocheck = 'DIS.vtu' m.export(output_dir, fmt='vtk', vtk_grid_type='UnstructuredGrid') filetocheck = os.path.join(output_dir, filenametocheck) - totalbytes4 = os.path.getsize(filetocheck) + # totalbytes4 = os.path.getsize(filetocheck) # assert(totalbytes4==67905) nlines4 = count_lines_in_file(filetocheck) assert(nlines4==993) @@ -429,19 +468,20 @@ def test_vtk_vti(): T = (m.bcf6.tran.array, m.bcf6.tran.array, np.zeros(m.modelgrid.shape)) vtk.export_vector(m, T, output_dir, 'T', point_scalars=True) filetocheck = os.path.join(output_dir, filenametocheck) - totalbytes5 = os.path.getsize(filetocheck) + # totalbytes5 = os.path.getsize(filetocheck) # assert(totalbytes5==12621) nlines5 = count_lines_in_file(filetocheck) assert(nlines5==20) - # vector binary + # vector with point scalars and binary vtk.export_vector(m, T, output_dir + '_bin', 'T', point_scalars=True, binary=True) filetocheck = os.path.join(output_dir + '_bin', filenametocheck) - totalbytes6 = os.path.getsize(filetocheck) + # totalbytes6 = os.path.getsize(filetocheck) # assert(totalbytes6==16716) - nlines6 = count_lines_in_file(filetocheck, binary=True) - assert(nlines6==18) + # nlines6 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines6==18) + assert(os.path.exists(filetocheck)) return @@ -456,7 +496,7 @@ def test_vtk_vtr(): # export and check m.export(output_dir, fmt='vtk') filetocheck = os.path.join(output_dir, filenametocheck) - totalbytes = os.path.getsize(filetocheck) + # totalbytes = os.path.getsize(filetocheck) # assert(totalbytes==79953) nlines = count_lines_in_file(filetocheck) assert(nlines==87) @@ -464,7 +504,7 @@ def test_vtk_vtr(): # with point scalar m.export(output_dir + '_points', fmt='vtk', point_scalars=True) filetocheck = os.path.join(output_dir + '_points', filenametocheck) - totalbytes1 = os.path.getsize(filetocheck) + # totalbytes1 = os.path.getsize(filetocheck) # assert(totalbytes1==182168) nlines1 = count_lines_in_file(filetocheck) assert(nlines1==121) @@ -472,16 +512,17 @@ def test_vtk_vtr(): # with binary m.export(output_dir + '_bin', fmt='vtk', binary=True) filetocheck = os.path.join(output_dir + '_bin', filenametocheck) - totalbytes2 = os.path.getsize(filetocheck) - # assert(totalbytes2==47778) - nlines2 = count_lines_in_file(filetocheck, binary=True) - assert(nlines2==28) + # totalbytes2 = os.path.getsize(filetocheck) + # assert(totalbytes2==47874) + # nlines2 = count_lines_in_file(filetocheck, binary=True) + # assert(nlines2==28) + assert(os.path.exists(filetocheck)) # force .vtu filenametocheck = 'EVT_01.vtu' m.export(output_dir, fmt='vtk', vtk_grid_type='UnstructuredGrid') filetocheck = os.path.join(output_dir, filenametocheck) - totalbytes3 = os.path.getsize(filetocheck) + # totalbytes3 = os.path.getsize(filetocheck) # assert(totalbytes3==78762) nlines3 = count_lines_in_file(filetocheck) assert(nlines3==1105) diff --git a/flopy/discretization/grid.py b/flopy/discretization/grid.py index 9853a95738..4f0ca183fa 100644 --- a/flopy/discretization/grid.py +++ b/flopy/discretization/grid.py @@ -136,7 +136,11 @@ def __init__(self, grid_type=None, top=None, botm=None, idomain=None, LENUNI = {"u": 0, "f": 1, "m": 2, "c": 3} self.use_ref_coords = True self._grid_type = grid_type + if top is not None: + top = top.astype(float) self._top = top + if botm is not None: + botm = botm.astype(float) self._botm = botm self._idomain = idomain diff --git a/flopy/discretization/structuredgrid.py b/flopy/discretization/structuredgrid.py index f8a77b45e1..a6ad0492f7 100644 --- a/flopy/discretization/structuredgrid.py +++ b/flopy/discretization/structuredgrid.py @@ -2,6 +2,34 @@ import numpy as np from .grid import Grid, CachedData +def array_at_verts_basic2d(a): + """ + Computes values at cell vertices on 2d array using neighbor averaging. + + Parameters + ---------- + a : ndarray + Array values at cell centers, could be a slice in any orientation. + + Returns + ------- + averts : ndarray + Array values at cell vertices, shape (a.shape[0]+1, a.shape[1]+1). + """ + assert a.ndim == 2 + shape_verts2d = (a.shape[0]+1, a.shape[1]+1) + + # create a 3D array of size (nrow+1, ncol+1, 4) + averts3d = np.full(shape_verts2d + (4,), np.nan) + averts3d[:-1, :-1, 0] = a + averts3d[:-1, 1:, 1] = a + averts3d[1:, :-1, 2] = a + averts3d[1:, 1:, 3] = a + + # calculate the mean over the last axis, ignoring NaNs + averts = np.nanmean(averts3d, axis=2) + + return averts class StructuredGrid(Grid): """ @@ -42,16 +70,18 @@ def __init__(self, delc=None, delr=None, top=None, botm=None, idomain=None, super(StructuredGrid, self).__init__('structured', top, botm, idomain, lenuni, epsg, proj4, prj, xoff, yoff, angrot) - self.__delc = delc - self.__delr = delr if delc is not None: self.__nrow = len(delc) + self.__delc = delc.astype(float) else: self.__nrow = nrow + self.__delc = delc if delr is not None: self.__ncol = len(delr) + self.__delr = delr.astype(float) else: self.__ncol = ncol + self.__delr = delr if top is not None: assert self.__nrow * self.__ncol == len(np.ravel(top)) if botm is not None: @@ -122,6 +152,39 @@ def delc(self): def delr(self): return copy.deepcopy(self.__delr) + @property + def delz(self): + cache_index = 'delz' + if cache_index not in self._cache_dict or \ + self._cache_dict[cache_index].out_of_date: + delz = self.top_botm[:-1, :, :] - self.top_botm[1:, :, :] + self._cache_dict[cache_index] = CachedData(delz) + if self._copy_cache: + return self._cache_dict[cache_index].data + else: + return self._cache_dict[cache_index].data_nocopy + + @property + def top_botm_withnan(self): + """ + Same as top_botm array but with NaN where idomain==0 both above and + below a cell. + """ + cache_index = 'top_botm_withnan' + if cache_index not in self._cache_dict or \ + self._cache_dict[cache_index].out_of_date: + is_inactive_above = np.full(self.top_botm.shape, True) + is_inactive_above[:-1, :, :] = self._idomain==0 + is_inactive_below = np.full(self.top_botm.shape, True) + is_inactive_below[1:, :, :] = self._idomain==0 + where_to_nan = np.logical_and(is_inactive_above, is_inactive_below) + top_botm_withnan = np.where(where_to_nan, np.nan, self.top_botm) + self._cache_dict[cache_index] = CachedData(top_botm_withnan) + if self._copy_cache: + return self._cache_dict[cache_index].data + else: + return self._cache_dict[cache_index].data_nocopy + @property def xyzvertices(self): """ @@ -194,6 +257,86 @@ def zedges(self): else: return self._cache_dict[cache_index].data_nocopy + @property + def zverts_smooth(self): + """ + Get a unique z of cell vertices for smooth (instead of stepwise) layer + elevations using bilinear interpolation. + + Returns + ------- + zverts : ndarray, shape (nlay+1, nrow+1, ncol+1) + z of cell vertices. NaN values are assigned in accordance with + inactive cells defined by idomain. + """ + cache_index = 'zverts_smooth' + if cache_index not in self._cache_dict or \ + self._cache_dict[cache_index].out_of_date: + zverts_smooth = self._zverts_smooth() + self._cache_dict[cache_index] = CachedData(zverts_smooth) + if self._copy_cache: + return self._cache_dict[cache_index].data + else: + return self._cache_dict[cache_index].data_nocopy + + def _zverts_smooth(self): + """ + For internal use only. The user should call zverts_smooth. + """ + # initialize the result array + shape_verts = (self.nlay+1, self.nrow+1, self.ncol+1) + zverts_basic = np.empty(shape_verts, dtype='float64') + + # assign NaN to top_botm where idomain==0 both above and below + if self._idomain is not None: + _top_botm = self.top_botm_withnan + + # perform basic interpolation (this will be useful in all cases) + # loop through layers + for k in range(self.nlay+1): + zvertsk = array_at_verts_basic2d(_top_botm[k, : , :]) + zverts_basic[k, : , :] = zvertsk + + if self.is_regular(): + # if the grid is regular, basic interpolation is the correct one + zverts = zverts_basic + else: + # cell centers + xcenters, ycenters = self.get_local_coords(self.xcellcenters, + self.ycellcenters) + # flip y direction because RegularGridInterpolator requires + # increasing input coordinates + ycenters = np.flip(ycenters, axis=0) + _top_botm = np.flip(_top_botm, axis=1) + xycenters = (ycenters[:, 0], xcenters[0, :]) + + # vertices + xverts, yverts = self.get_local_coords(self.xvertices, + self.yvertices) + xyverts = np.ndarray((xverts.size, 2)) + xyverts[:, 0] = yverts.ravel() + xyverts[:, 1] = xverts.ravel() + + # interpolate + import scipy.interpolate as interp + shape_verts2d = (self.nrow+1, self.ncol+1) + zverts = np.empty(shape_verts, dtype='float64') + # loop through layers + for k in range(self.nlay+1): + # interpolate layer elevation + zcenters_k = _top_botm[k, : , :] + interp_func = interp.RegularGridInterpolator(xycenters, + zcenters_k, bounds_error=False, fill_value=np.nan) + zverts_k = interp_func(xyverts) + zverts_k = zverts_k.reshape(shape_verts2d) + zverts[k, : , :] = zverts_k + + # use basic interpolation for remaining NaNs at boundaries + where_nan = np.isnan(zverts) + zverts[where_nan] = zverts_basic[where_nan] + + return zverts + @property def xyzcellcenters(self): """ @@ -431,7 +574,9 @@ def from_gridspec(cls, gridspec_file, lenuni=0): # Exporting def write_shapefile(self, filename='grid.shp', epsg=None, prj=None): - """Write a shapefile of the grid with just the row and column attributes""" + """ + Write a shapefile of the grid with just the row and column attributes. + """ from ..export.shapefile_utils import write_grid_shapefile if epsg is None and prj is None: epsg = self.epsg @@ -455,13 +600,13 @@ def is_regular(self): is_regular_y = np.count_nonzero(rel_diff_y > rel_tol) == 0 # Regularity test in z direction - thickness = (self.top[0, 0] - self.botm[0, 0, 0]) - rel_diff_z1 = (self.top - self.botm[0, :, :] - thickness) / thickness - failed = np.abs(rel_diff_z1) > rel_tol + rel_diff_thick0 = (self.delz[0, :, :] - self.delz[0, 0, 0]) \ + / self.delz[0, 0, 0] + failed = np.abs(rel_diff_thick0) > rel_tol is_regular_z = np.count_nonzero(failed) == 0 - for k in range(self.nlay - 1): - rel_diff_zk = (self.botm[k, :, :] - self.botm[k + 1, :, :] - - thickness) / thickness + for k in range(1, self.nlay): + rel_diff_zk = (self.delz[k, :, :] - self.delz[0, :, :]) \ + / self.delz[0, :, :] failed = np.abs(rel_diff_zk) > rel_tol is_regular_z = is_regular_z and np.count_nonzero(failed) == 0 @@ -476,19 +621,294 @@ def is_rectilinear(self): rel_tol = 1.e-5 # Rectilinearity test in z direction - thickness = (self.top[0, 0] - self.botm[0, 0, 0]) - rel_diff_z1 = (self.top - self.botm[0, :, :] - thickness) / thickness - failed = np.abs(rel_diff_z1) > rel_tol - is_rectilinear_z = np.count_nonzero(failed) == 0 - for k in range(self.nlay - 1): - thickness_k = (self.botm[k, 0, 0] - self.botm[k + 1, 0, 0]) - rel_diff_zk = (self.botm[k, :, :] - self.botm[k + 1, :, :] - - thickness_k) / thickness_k + is_rect_z = True + for k in range(self.nlay): + rel_diff_zk = (self.delz[k, :, :] - self.delz[k, 0, 0]) \ + / self.delz[k, 0, 0] failed = np.abs(rel_diff_zk) > rel_tol - is_rectilinear_z = is_rectilinear_z and \ - np.count_nonzero(failed) == 0 + is_rect_z = is_rect_z and np.count_nonzero(failed) == 0 + + return is_rect_z + + def array_at_verts_basic(self, a): + """ + Computes values at cell vertices using neighbor averaging. + + Parameters + ---------- + a : ndarray + Array values at cell centers. + + Returns + ------- + averts : ndarray + Array values at cell vertices, shape + (a.shape[0]+1, a.shape[1]+1, a.shape[2]+1). NaN values are assigned + in accordance with inactive cells defined by idomain. + """ + assert a.ndim == 3 + shape_verts = (a.shape[0]+1, a.shape[1]+1, a.shape[2]+1) + + # set to NaN where idomain==0 + a[self._idomain==0] = np.nan + + # create a 4D array of size (nlay+1, nrow+1, ncol+1, 8) + averts4d = np.full(shape_verts + (8,), np.nan) + averts4d[:-1, :-1, :-1, 0] = a + averts4d[:-1, :-1, 1:, 1] = a + averts4d[:-1, 1:, :-1, 2] = a + averts4d[:-1, 1:, 1:, 3] = a + averts4d[1:, :-1, :-1, 4] = a + averts4d[1:, :-1, 1:, 5] = a + averts4d[1:, 1:, :-1, 6] = a + averts4d[1:, 1:, 1:, 7] = a + + # calculate the mean over the last axis, ignoring NaNs + averts = np.nanmean(averts4d, axis=3) + + return averts - return is_rectilinear_z + def array_at_verts(self, a): + """ + Computes values at cell vertices using trilinear interpolation. + + Parameters + ---------- + a : ndarray + Array values. Allowed shapes are: (nlay, nrow, ncol), + (nlay, nrow, ncol+1), (nlay, nrow+1, ncol) and + (nlay+1, nrow, ncol). + * When the shape is (nlay, nrow, ncol), the input values are + considered at cell centers. + * When the shape is extended in one direction, the input values are + considered at the center of cell faces in this direction. + + Returns + ------- + averts : ndarray + Array values interpolated at cell vertices, shape + (nlay+1, nrow+1, ncol+1). NaN values are assigned in accordance + with inactive cells defined by idomain. + """ + # define shapes + shape_ext_x = (self.nlay, self.nrow, self.ncol+1) + shape_ext_y = (self.nlay, self.nrow+1, self.ncol) + shape_ext_z = (self.nlay+1, self.nrow, self.ncol) + shape_verts = (self.nlay+1, self.nrow+1, self.ncol+1) + + # perform basic interpolation (this will be useful in all cases) + if a.shape == self.shape: + averts_basic = self.array_at_verts_basic(a) + elif a.shape == shape_ext_x: + averts_basic = np.empty(shape_verts, dtype=a.dtype) + for j in range(self.ncol+1): + averts_basic[:, :, j] = array_at_verts_basic2d(a[:, :, j]) + elif a.shape == shape_ext_y: + averts_basic = np.empty(shape_verts, dtype=a.dtype) + for i in range(self.nrow+1): + averts_basic[:, i, :] = array_at_verts_basic2d(a[:, i, :]) + elif a.shape == shape_ext_z: + averts_basic = np.empty(shape_verts, dtype=a.dtype) + for k in range(self.nlay+1): + averts_basic[k, :, :] = array_at_verts_basic2d(a[k, :, :]) + + if self.is_regular(): + # if the grid is regular, basic interpolation is the correct one + averts = averts_basic + else: + # get input coordinates + xcenters, ycenters = self.get_local_coords(self.xcellcenters, + self.ycellcenters) + zcenters = self.zcellcenters + if a.shape == self.shape: + xinput = xcenters * np.ones(self.shape) + yinput = ycenters * np.ones(self.shape) + zinput = zcenters + # set array to NaN where inactive + if self._idomain is not None: + a = np.where(self._idomain == 0, np.nan, a) + elif a.shape == shape_ext_x: + xinput = np.reshape(self.xyedges[0], (1, 1, self.ncol+1)) + xinput = xinput * np.ones(shape_ext_x) + yinput = np.reshape(self.ycellcenters[:, 0], (1, self.nrow, 1)) + yinput = yinput * np.ones(shape_ext_x) + zinput = self.array_at_faces(zcenters, 'x', withnan=False) + elif a.shape == shape_ext_y: + xinput = np.reshape(self.xcellcenters[0, :], (1, 1, self.ncol)) + xinput = xinput * np.ones(shape_ext_y) + yinput = np.reshape(self.xyedges[1], (1, self.nrow+1, 1)) + yinput = yinput * np.ones(shape_ext_y) + zinput = self.array_at_faces(zcenters, 'y', withnan=False) + elif a.shape == shape_ext_z: + xinput = xcenters * np.ones(shape_ext_z) + yinput = ycenters * np.ones(shape_ext_z) + zinput = self.top_botm + else: + raise ValueError('Incompatible array shape') + + # flip y and z directions because RegularGridInterpolator requires + # increasing input coordinates + xinput = np.flip(xinput, axis=[0, 1]) + yinput = np.flip(yinput, axis=[0, 1]) + zinput = np.flip(zinput, axis=[0, 1]) + _a = np.flip(a, axis=[0, 1]) + + # get output coordinates (i.e. vertices) + xoutput, youtput = self.get_local_coords(self.xvertices, + self.yvertices) + xoutput = xoutput * np.ones(shape_verts) + youtput = youtput * np.ones(shape_verts) + zoutput = self.zverts_smooth + xyzoutput = np.ndarray((zoutput.size, 3)) + xyzoutput[:, 0] = zoutput.ravel() + xyzoutput[:, 1] = youtput.ravel() + xyzoutput[:, 2] = xoutput.ravel() + + # interpolate + import scipy.interpolate as interp + if self.is_rectilinear(): + xyzinput = (zinput[:, 0, 0], yinput[0, :, 0], xinput[0, 0, :]) + interp_func = interp.RegularGridInterpolator(xyzinput, _a, + bounds_error=False, fill_value=np.nan) + else: + # format inputs, excluding NaN + valid_input = np.logical_not(np.isnan(_a)) + xyzinput = np.ndarray((np.count_nonzero(valid_input), 3)) + xyzinput[:, 0] = zinput[valid_input] + xyzinput[:, 1] = yinput[valid_input] + xyzinput[:, 2] = xinput[valid_input] + _a = _a[valid_input] + interp_func = interp.LinearNDInterpolator(xyzinput, _a, + fill_value=np.nan) + averts = interp_func(xyzoutput) + averts = averts.reshape(shape_verts) + + # use basic interpolation for remaining NaNs at boundaries + where_nan = np.isnan(averts) + averts[where_nan] = averts_basic[where_nan] + + # assign NaN where idomain==0 at all 8 neighbors (these should be + # the same locations as in averts_basic) + averts[np.isnan(averts_basic)] = np.nan + + return averts + + def array_at_faces(self, a, direction, withnan=True): + """ + Computes values at the center of cell faces using linear interpolation. + + Parameters + ---------- + a : ndarray + Values at cell centers, shape (nlay, row, ncol). + direction : str, possible values are 'x', 'y' and 'z' + Direction in which values will be interpolated at cell faces. + withnan : bool + If True (default), the result value will be set to NaN where the + cell face sits between inactive cells. If False, not. + + Returns + ------- + afaces : ndarray + Array values interpolated at cell vertices, shape as input extended + by 1 along the specified direction. + + """ + assert a.shape == self.shape + + # get the dimension that corresponds to the direction + dir_to_dim = {'x': 2, 'y': 1, 'z': 0} + dim = dir_to_dim[direction] + + # extended array with ghost cells on both sides having zero values + ghost_shape = list(a.shape) + ghost_shape[dim] += 2 + a_ghost = np.zeros(ghost_shape, dtype=a.dtype) + + # extended delta with ghost cells on both sides having zero values + delta_ghost = np.zeros(ghost_shape, dtype=a.dtype) + + # inactive bool array + if withnan and self._idomain is not None: + inactive = self._idomain == 0 + + if dim == 0: + # fill array with ghost cells + a_ghost[1:-1, :, :] = a + a_ghost[0, :, :] = a[0, :, :] + a_ghost[-1, :, :] = a[-1, :, :] + + # calculate weights + delta_ghost[1:-1, :, :] = self.delz + weight2 = delta_ghost[:-1, :, :] / (delta_ghost[:-1, :, :] + \ + delta_ghost[1:, :, :]) + weight1 = 1. - weight2 + + # interpolate + afaces = a_ghost[:-1, :, :]*weight1 + a_ghost[1:, :, :]*weight2 + + # assign NaN where idomain==0 on both sides + if withnan and self._idomain is not None: + inactive_faces = np.full(afaces.shape, True) + inactive_faces[:-1, :, :] = np.logical_and( + inactive_faces[:-1, :, :], inactive) + inactive_faces[1:, :, :] = np.logical_and( + inactive_faces[1:, :, :], inactive) + afaces[inactive_faces] = np.nan + + elif dim == 1: + # fill array with ghost cells + a_ghost[:, 1:-1, :] = a + a_ghost[:, 0, :] = a[:, 0, :] + a_ghost[:, -1, :] = a[:, -1, :] + + # calculate weights + delc = np.reshape(self.delc, (1, self.nrow, 1)) + delc_3D = delc * np.ones(self.shape) + delta_ghost[:, 1:-1, :] = delc_3D + weight2 = delta_ghost[:, :-1, :] / (delta_ghost[:, :-1, :] + \ + delta_ghost[:, 1:, :]) + weight1 = 1. - weight2 + + # interpolate + afaces = a_ghost[:, :-1, :]*weight1 + a_ghost[:, 1:, :]*weight2 + + # assign NaN where idomain==0 on both sides + if withnan and self._idomain is not None: + inactive_faces = np.full(afaces.shape, True) + inactive_faces[:, :-1, :] = np.logical_and( + inactive_faces[:, :-1, :], inactive) + inactive_faces[:, 1:, :] = np.logical_and( + inactive_faces[:, 1:, :], inactive) + afaces[inactive_faces] = np.nan + + elif dim == 2: + # fill array with ghost cells + a_ghost[:, :, 1:-1] = a + a_ghost[:, :, 0] = a[:, :, 0] + a_ghost[:, :, -1] = a[:, :, -1] + + # calculate weights + delr = np.reshape(self.delr, (1, 1, self.ncol)) + delr_3D = delr * np.ones(self.shape) + delta_ghost[:, :, 1:-1] = delr_3D + weight2 = delta_ghost[:, :, :-1] / (delta_ghost[:, :, :-1] + \ + delta_ghost[:, :, 1:]) + weight1 = 1. - weight2 + + # interpolate + afaces = a_ghost[:, :, :-1]*weight1 + a_ghost[:, :, 1:]*weight2 + + # assign NaN where idomain==0 on both sides + if withnan and self._idomain is not None: + inactive_faces = np.full(afaces.shape, True) + inactive_faces[:, :, :-1] = np.logical_and( + inactive_faces[:, :, :-1], inactive) + inactive_faces[:, :, 1:] = np.logical_and( + inactive_faces[:, :, 1:], inactive) + afaces[inactive_faces] = np.nan + + return afaces if __name__ == "__main__": import matplotlib.pyplot as plt diff --git a/flopy/export/vtk.py b/flopy/export/vtk.py index fa383c67a5..9fc21b5a10 100644 --- a/flopy/export/vtk.py +++ b/flopy/export/vtk.py @@ -395,6 +395,8 @@ def __init__(self, model, verbose=None, nanval=-1e+20, smooth=False, self.ncol = self.modelgrid.ncol self.shape = (self.nlay, self.nrow, self.ncol) self.shape2d = (self.shape[1], self.shape[2]) + self.shape_verts = (self.shape[0]+1, self.shape[1]+1, self.shape[2]+1) + self.shape_verts2d = (self.shape_verts[1], self.shape_verts[2]) self.nanval = nanval self.cell_type = 11 @@ -403,6 +405,8 @@ def __init__(self, model, verbose=None, nanval=-1e+20, smooth=False, self.smooth = smooth self.point_scalars = point_scalars + self.has_cell_data = False + self.has_point_data = False # check if structured grid, vtk only supports structured grid assert (isinstance(self.modelgrid, StructuredGrid)) @@ -437,7 +441,7 @@ def __init__(self, model, verbose=None, nanval=-1e+20, smooth=False, def _vtk_grid_type(self, vtk_grid_type='auto'): """ - Determine the vtk grid type and corresponding file extension. + Determines the vtk grid type and corresponding file extension. Parameters ---------- @@ -502,7 +506,7 @@ def _vtk_grid_type(self, vtk_grid_type='auto'): def _format_array(self, a, array2d=False): """ - Format array for vtk output + Formats array for vtk output. Parameters ---------- @@ -516,30 +520,55 @@ def _format_array(self, a, array2d=False): Return ------ - Formatted array (a copy is made) + Formatted array (note a copy is made) """ # if array is 2d reformat to 3d array if array2d: - assert a.shape == self.shape2d - array = np.full(self.shape, self.nanval) + if a.shape == self.shape2d: + array = np.full(self.shape, self.nanval) + elif a.shape == self.shape_verts2d: + array = np.full(self.shape_verts, self.nanval) + else: + raise ValueError('Incompatible array size') array[0, :, :] = a a = array - try: - assert a.shape == self.shape - except AssertionError: - return - - # set to nan where nanval or where ibound==0 - # note np.where makes a copy of the array - where_to_nan = np.logical_or(a==self.nanval, self.ibound==0) + # deal with inactive cells + inactive3d = self.ibound==0 + if a.shape == self.shape: + # set to nan where nanval or where ibound==0 + where_to_nan = np.logical_or(a==self.nanval, inactive3d) + self.has_cell_data = True + elif a.shape == self.shape_verts: + # set to nan where ibound==0 at all 8 neighbors + where_to_nan = np.full(self.shape_verts, True) + where_to_nan[:-1, :-1, :-1] = inactive3d + where_to_nan[:-1, :-1, 1:] = np.logical_and( + where_to_nan[:-1, :-1, 1:], inactive3d) + where_to_nan[:-1, 1:, :-1] = np.logical_and( + where_to_nan[:-1, 1:, :-1], inactive3d) + where_to_nan[:-1, 1:, 1:] = np.logical_and( + where_to_nan[:-1, 1:, 1:], inactive3d) + where_to_nan[1:, :-1, :-1] = np.logical_and( + where_to_nan[1:, :-1, :-1], inactive3d) + where_to_nan[1:, :-1, 1:] = np.logical_and( + where_to_nan[1:, :-1, 1:], inactive3d) + where_to_nan[1:, 1:, :-1] = np.logical_and( + where_to_nan[1:, 1:, :-1], inactive3d) + where_to_nan[1:, 1:, 1:] = np.logical_and( + where_to_nan[1:, 1:, 1:], inactive3d) + self.has_point_data = True + self.smooth = True + else: + # incompatible size, skip this array + return None a = np.where(where_to_nan, np.nan, a) return a def add_array(self, name, a, array2d=False): """ - Adds an array to the vtk object + Adds an array to the vtk object. Parameters ---------- @@ -548,6 +577,7 @@ def add_array(self, name, a, array2d=False): name of the array a : flopy array the array to be added to the vtk object + shape should match either grid cells or grid vertices array2d : bool True if the array is 2d """ @@ -562,7 +592,7 @@ def add_array(self, name, a, array2d=False): def add_vector(self, name, v, array2d=False): """ - Adds a vector (i.e., a tuple of arrays) to the vtk object + Adds a vector (i.e., a tuple of arrays) to the vtk object. Parameters ---------- @@ -571,6 +601,7 @@ def add_vector(self, name, v, array2d=False): name of the vector v : tuple of arrays the vector to be added to the vtk object + shape of arrays should match either grid cells or grid vertices array2d : bool True if the vector components are 2d arrays """ @@ -629,7 +660,7 @@ def write(self, output_file, timeval=None): # get the verts and iverts to be output verts, iverts, _ = \ - self.get_3d_vertex_connectivity(actwcells=actwcells3d) + self._get_3d_vertex_connectivity(actwcells=actwcells3d) # check if there is data to be written out if len(verts) == 0: @@ -725,81 +756,111 @@ def write(self, output_file, timeval=None): # end coordinates xml.close_element('Coordinates') - # cell data - xml.open_element('CellData') + if self.has_cell_data: + # cell data + xml.open_element('CellData') - # loop through stored arrays - for name, a in self.arrays.items(): - if self.vtk_grid_type == 'UnstructuredGrid': - xml.write_array(a, actwcells=actwcells3d, Name=name, - NumberOfComponents='1') - else: - # flip "a" so coordinates increase along with indices as in vtk - a = np.flip(a, axis=0) - a = np.flip(a, axis=1) - xml.write_array(a, Name=name, NumberOfComponents='1') + # loop through stored arrays + for name, a in self.arrays.items(): + if a.shape == self.shape_verts: + # these are dealt with later + continue + if self.vtk_grid_type == 'UnstructuredGrid': + xml.write_array(a, actwcells=actwcells3d, Name=name, + NumberOfComponents='1') + else: + # flip "a" so coordinates increase along with indices as in + # vtk + a = np.flip(a, axis=[0, 1]) + xml.write_array(a, Name=name, NumberOfComponents='1') - # loop through stored vectors - for name, v in self.vectors.items(): - ncomp = len(v) - v_as_array = np.moveaxis(np.array(v), 0, -1) - if self.vtk_grid_type == 'UnstructuredGrid': - shape4d = actwcells3d.shape + (ncomp,) - actwcells4d = actwcells3d.reshape(actwcells3d.shape + (1,)) - actwcells4d = np.broadcast_to(actwcells4d, shape4d) - xml.write_array(v_as_array, actwcells=actwcells4d, Name=name, - NumberOfComponents=ncomp) - else: - # flip "v" so coordinates increase along with indices as in vtk - v_as_array = np.flip(v_as_array, axis=0) - v_as_array = np.flip(v_as_array, axis=1) - xml.write_array(v_as_array, Name=name, - NumberOfComponents=ncomp) + # loop through stored vectors + for name, v in self.vectors.items(): + if v[0].shape == self.shape_verts: + # these are dealt with later + continue + ncomp = len(v) + v_as_array = np.moveaxis(np.array(v), 0, -1) + if self.vtk_grid_type == 'UnstructuredGrid': + shape4d = actwcells3d.shape + (ncomp,) + actwcells4d = actwcells3d.reshape(actwcells3d.shape + (1,)) + actwcells4d = np.broadcast_to(actwcells4d, shape4d) + xml.write_array(v_as_array, actwcells=actwcells4d, + Name=name, NumberOfComponents=ncomp) + else: + # flip "v" so coordinates increase along with indices as in + # vtk + v_as_array = np.flip(v_as_array, axis=[0, 1]) + xml.write_array(v_as_array, Name=name, + NumberOfComponents=ncomp) - # end cell data - xml.close_element('CellData') + # end cell data + xml.close_element('CellData') - if self.point_scalars: + if self.point_scalars or self.has_point_data: # point data (i.e., values at vertices) xml.open_element('PointData') # loop through stored arrays for name, a in self.arrays.items(): - # get the array values onto vertices - if self.vtk_grid_type == 'UnstructuredGrid': - _, _, zverts = self.get_3d_vertex_connectivity( - actwcells=actwcells3d, zvalues=a) - a = np.array(list(zverts.values())) + if a.shape == self.shape: + if not self.point_scalars: + continue + # get the array values onto vertices + if self.vtk_grid_type == 'UnstructuredGrid': + _, _, averts = self._get_3d_vertex_connectivity( + actwcells=actwcells3d, zvalues=a) + a = np.array(list(averts.values())) + else: + a = self.modelgrid.array_at_verts(a) + a = np.flip(a, axis=[0, 1]) else: - a = self.extendedDataArray(a) - # flip "a" so coordinates increase along with indices as in - # vtk - a = np.flip(a, axis=0) - a = np.flip(a, axis=1) - # write to file + if self.vtk_grid_type == 'UnstructuredGrid': + # still need to do this to be consistent with + # connectivity (i.e. 8 points for every cell) + _, _, averts = self._get_3d_vertex_connectivity( + actwcells=actwcells3d, zvalues=a) + a = np.array(list(averts.values())) + else: + # flip "a" so coordinates increase along with indices + # as in vtk + a = np.flip(a, axis=[0, 1]) xml.write_array(a, Name=name, NumberOfComponents='1') # loop through stored vectors for name, v in self.vectors.items(): - # get the vector values onto vertices - v_ext = () - for vcomp in v: - if self.vtk_grid_type == 'UnstructuredGrid': - _, _, zverts = self.get_3d_vertex_connectivity( - actwcells=actwcells3d, zvalues=vcomp) - vcomp = np.array(list(zverts.values())) - else: - vcomp = self.extendedDataArray(vcomp) - v_ext = v_ext + (vcomp,) + if v[0].shape == self.shape: + if not self.point_scalars: + continue + # get the vector values onto vertices + v_verts = () + for vcomp in v: + if self.vtk_grid_type == 'UnstructuredGrid': + _, _, averts = self._get_3d_vertex_connectivity( + actwcells=actwcells3d, zvalues=vcomp) + vcomp = np.array(list(averts.values())) + else: + vcomp = self.modelgrid.array_at_verts(vcomp) + vcomp = np.flip(vcomp, axis=[0, 1]) + v_verts = v_verts + (vcomp,) + v = v_verts + else: + v_verts = () + for vcomp in v: + if self.vtk_grid_type == 'UnstructuredGrid': + # still need to do this to be consistent with + # connectivity (i.e. 8 points for every cell) + _, _, averts = self._get_3d_vertex_connectivity( + actwcells=actwcells3d, zvalues=vcomp) + vcomp = np.array(list(averts.values())) + else: + vcomp = np.flip(vcomp, axis=[0, 1]) + v_verts = v_verts + (vcomp,) + v = v_verts # write to file - ncomp = len(v_ext) - v_ext_as_array = np.moveaxis(np.array(v_ext), 0, -1) - if self.vtk_grid_type != 'UnstructuredGrid': - # flip "v" so coordinates increase along with indices as in - # vtk - v_ext_as_array = np.flip(v_ext_as_array, axis=0) - v_ext_as_array = np.flip(v_ext_as_array, axis=1) - xml.write_array(v_ext_as_array, Name=name, + ncomp = len(v) + v_as_array = np.moveaxis(np.array(v), 0, -1) + xml.write_array(v_as_array, Name=name, NumberOfComponents=ncomp) # end point data @@ -823,40 +884,61 @@ def _configure_data_arrays(self): Compares arrays and active cells to find where active data exists, and what cells to output. """ - # get 1d shape + # build 1d index array shape1d = self.shape[0] * self.shape[1] * self.shape[2] - - # build index array - ot_idx_array = np.zeros(shape1d, dtype=np.int) + actwcells1d = np.zeros(shape1d, dtype=np.int) + if self.has_point_data: + shape1d_verts = self.shape_verts[0] * self.shape_verts[1] * \ + self.shape_verts[2] + actwcells1d_verts = np.zeros(shape1d_verts, dtype=np.int) # loop through arrays - for name in self.arrays: - array = self.arrays[name] + for a in self.arrays.values(): # make array 1d - a = array.ravel() + a1d = a.ravel() # get the indexes where there is data - idxs = np.argwhere(np.logical_not(np.isnan(a))) + idxs = np.argwhere(np.logical_not(np.isnan(a1d))) # set the active array to 1 - ot_idx_array[idxs] = 1 + if a.shape == self.shape: + actwcells1d[idxs] = 1 + elif self.has_point_data: + actwcells1d_verts[idxs] = 1 # loop through vectors - for name in self.vectors: - for vcomp in self.vectors[name]: + for v in self.vectors.values(): + for vcomp in v: # make array 1d - a = vcomp.ravel() + vcomp1d = vcomp.ravel() # get the indexes where there is data - idxs = np.argwhere(np.logical_not(np.isnan(a))) + idxs = np.argwhere(np.logical_not(np.isnan(vcomp1d))) # set the active array to 1 - ot_idx_array[idxs] = 1 - - # reset the shape of the active data array - ot_idx_array = ot_idx_array.reshape(self.shape) - - return ot_idx_array - - def get_3d_vertex_connectivity(self, actwcells=None, zvalues=None): + if vcomp.shape == self.shape: + actwcells1d[idxs] = 1 + elif self.has_point_data: + actwcells1d_verts[idxs] = 1 + + # reshape to 3D array + actwcells3d = actwcells1d.reshape(self.shape) + if self.has_point_data: + actwcells3d_verts = actwcells1d_verts.reshape(self.shape_verts) + # activate cells that are neighbor of 8 active vertices + activate = np.full(self.shape, True) + activate[actwcells3d_verts[:-1, :-1, :-1] == 0] = False + activate[actwcells3d_verts[:-1, :-1, 1:] == 0] = False + activate[actwcells3d_verts[:-1, 1:, :-1] == 0] = False + activate[actwcells3d_verts[:-1, 1:, 1:] == 0] = False + activate[actwcells3d_verts[1:, :-1, :-1] == 0] = False + activate[actwcells3d_verts[1:, :-1, 1:] == 0] = False + activate[actwcells3d_verts[1:, 1:, :-1] == 0] = False + activate[actwcells3d_verts[1:, 1:, 1:] == 0] = False + activate[self.ibound==0] = False + actwcells3d[activate] = 1 + + return actwcells3d + + def _get_3d_vertex_connectivity(self, actwcells=None, zvalues=None): """ - Builds x,y,z vertices + Builds x,y,z vertices. Parameters ---------- @@ -888,10 +970,14 @@ def get_3d_vertex_connectivity(self, actwcells=None, zvalues=None): # if smoothing interpolate the z values if self.smooth: if zvalues is not None: - # use the given data array values - zVertices = self.extendedDataArray(zvalues) + if zvalues.shape==self.shape: + # interpolate using the given values + zVertices = self.modelgrid.array_at_verts(zvalues) + else: + # in this case the given values are already at vertices + zVertices = zvalues else: - zVertices = self.extendedDataArray(self.modelgrid.top_botm) + zVertices = self.modelgrid.zverts_smooth else: zVertices = None @@ -955,32 +1041,6 @@ def get_3d_vertex_connectivity(self, actwcells=None, zvalues=None): zvertsdict[cellid] = zverts return vertsdict, ivertsdict, zvertsdict - def extendedDataArray(self, dataArray): - - if dataArray.shape[0] == self.nlay+1: - dataArray = dataArray - else: - listArray = [dataArray[0]] - for lay in range(dataArray.shape[0]): - listArray.append(dataArray[lay]) - dataArray = np.stack(listArray) - - matrix = np.zeros([self.nlay+1, self.nrow+1, self.ncol+1]) - for lay in range(self.nlay+1): - for row in range(self.nrow+1): - for col in range(self.ncol+1): - indexList = [[row-1, col-1], [row-1, col], [row, col-1], - [row, col]] - neighList = [] - for index in indexList: - if index[0] in range(self.nrow) and index[1] in \ - range(self.ncol): - neighList.append(dataArray[lay, index[0], - index[1]]) - neighList = np.array(neighList) - matrix[lay, row, col] = np.nanmean(neighList) - return matrix - def _get_names(in_list): ot_list = [] diff --git a/flopy/utils/postprocessing.py b/flopy/utils/postprocessing.py index 01cc20f07e..a8dd3620d4 100644 --- a/flopy/utils/postprocessing.py +++ b/flopy/utils/postprocessing.py @@ -534,7 +534,7 @@ def get_extended_budget(cbcfile, precision='single', idx=None, def get_specific_discharge(model, cbcfile, precision='single', idx=None, kstpkper=None, totim=None, boundary_ifaces=None, - hdsfile=None): + hdsfile=None, position='centers'): """ Get the discharge vector at cell centers at a given time. For "classical" MODFLOW versions, we calculate it from the flow rate across cell faces. @@ -584,6 +584,9 @@ def get_specific_discharge(model, cbcfile, precision='single', idx=None, hdsfile is also required if the budget term 'HEAD DEP BOUNDS', 'RIVER LEAKAGE' or 'DRAINS' is present in boundary_ifaces and that the corresponding value is a list. + position : str + Position at which the specific discharge will be calculated. Possible + values are "centers" (default), "faces" and "vertices". Returns ------- @@ -603,7 +606,7 @@ def get_specific_discharge(model, cbcfile, precision='single', idx=None, cbf = bf.CellBudgetFile(cbcfile, precision=precision) rec_names = cbf.get_unique_record_names(decode=True) classical_budget_terms = ['FLOW RIGHT FACE', 'FLOW FRONT FACE', - 'FLOW RIGHT FACE'] + 'FLOW RIGHT FACE'] classical_budget = False for budget_term in classical_budget_terms: matched_name = [s for s in rec_names if budget_term in s] @@ -618,9 +621,8 @@ def get_specific_discharge(model, cbcfile, precision='single', idx=None, if classical_budget: # get extended budget Qx_ext, Qy_ext, Qz_ext = get_extended_budget(cbcfile, - precision=precision, idx=idx, kstpkper=kstpkper, - totim=totim, boundary_ifaces=boundary_ifaces, - hdsfile=hdsfile, model=model) + precision=precision, idx=idx, kstpkper=kstpkper, totim=totim, + boundary_ifaces=boundary_ifaces, hdsfile=hdsfile, model=model) # get saturated thickness (head - bottom elev for unconfined layer) if hdsfile is None: @@ -628,24 +630,49 @@ def get_specific_discharge(model, cbcfile, precision='single', idx=None, else: sat_thk = get_saturated_thickness(head, model, model.hdry) - # calculate qx at cell center - delc = np.reshape(model.modelgrid.delc, (1, model.modelgrid.nrow, 1)) - cross_area_x = np.empty(model.modelgrid.shape, dtype=float) + # inform modelgrid of no-flow and dry cells + modelgrid = model.modelgrid + if modelgrid._idomain is None: + modelgrid._idomain = model.dis.ibound + if hdsfile is not None: + noflo_or_dry = np.logical_or(head==model.hnoflo, head==model.hdry) + modelgrid._idomain[noflo_or_dry] = 0 + + # get cross section areas along x + delc = np.reshape(modelgrid.delc, (1, modelgrid.nrow, 1)) + cross_area_x = np.empty(modelgrid.shape, dtype=float) cross_area_x = delc * sat_thk - qx = 0.5 * (Qx_ext[:, :, 1:] + Qx_ext[:, :, :-1]) / cross_area_x - # calculate qy at cell center - delr = np.reshape(model.modelgrid.delr, (1, 1, model.modelgrid.ncol)) - cross_area_y = np.empty(model.modelgrid.shape, dtype=float) + # get cross section areas along y + delr = np.reshape(modelgrid.delr, (1, 1, modelgrid.ncol)) + cross_area_y = np.empty(modelgrid.shape, dtype=float) cross_area_y = delr * sat_thk - qy = 0.5 * (Qy_ext[:, 1:, :] + Qy_ext[:, :-1, :]) / cross_area_y - # calculate qz at cell center - cross_area_z = np.empty(model.modelgrid.shape, dtype=float) - cross_area_z = delc * delr - qz = 0.5 * (Qz_ext[1:, :, :] + Qz_ext[:-1, :, :]) / cross_area_z + # get cross section areas along z + cross_area_z = np.ones(modelgrid.shape) * delc * delr + + # calculate qx, qy, qz + if position == 'centers': + qx = 0.5 * (Qx_ext[:, :, 1:] + Qx_ext[:, :, :-1]) / cross_area_x + qy = 0.5 * (Qy_ext[:, 1:, :] + Qy_ext[:, :-1, :]) / cross_area_y + qz = 0.5 * (Qz_ext[1:, :, :] + Qz_ext[:-1, :, :]) / cross_area_z + elif position == 'faces' or position == 'vertices': + cross_area_x = modelgrid.array_at_faces(cross_area_x, 'x') + cross_area_y = modelgrid.array_at_faces(cross_area_y, 'y') + cross_area_z = modelgrid.array_at_faces(cross_area_z, 'z') + qx = Qx_ext / cross_area_x + qy = Qy_ext / cross_area_y + qz = Qz_ext / cross_area_z + else: + raise ValueError('"' + position + '" is not a valid value for ' + 'position') + if position == 'vertices': + qx = modelgrid.array_at_verts(qx) + qy = modelgrid.array_at_verts(qy) + qz = modelgrid.array_at_verts(qz) else: + # check valid options if boundary_ifaces is not None: import warnings warnings.warn('the boundary_ifaces option is not implemented ' @@ -653,6 +680,12 @@ def get_specific_discharge(model, cbcfile, precision='single', idx=None, 'budget is not recorded as FLOW RIGHT FACE, ' 'FLOW FRONT FACE and FLOW LOWER FACE; it will be ' 'ignored', UserWarning) + if position != 'centers': + raise NotImplementedError('position can only be "centers" for ' + '"non-classical" MODFLOW versions where ' + 'the budget is not recorded as FLOW ' + 'RIGHT FACE, FLOW FRONT FACE and FLOW ' + 'LOWER FACE') is_spdis = [s for s in rec_names if 'DATA-SPDIS' in s] if not is_spdis: @@ -675,10 +708,13 @@ def get_specific_discharge(model, cbcfile, precision='single', idx=None, qz.shape = shape # set no-flow and dry cells to NaN - if hdsfile is not None: + if hdsfile is not None and position == 'centers': noflo_or_dry = np.logical_or(head==model.hnoflo, head==model.hdry) qx[noflo_or_dry] = np.nan qy[noflo_or_dry] = np.nan qz[noflo_or_dry] = np.nan return qx, qy, qz + + +