diff --git a/docs/examples/image/qs.png b/docs/examples/image/qs.png new file mode 100644 index 00000000..0d64e08c Binary files /dev/null and b/docs/examples/image/qs.png differ diff --git a/docs/examples/image/qsprj.png b/docs/examples/image/qsprj.png new file mode 100644 index 00000000..f71b239a Binary files /dev/null and b/docs/examples/image/qsprj.png differ diff --git a/docs/examples/image/quickstart.png b/docs/examples/image/quickstart.png new file mode 100644 index 00000000..77737ad8 Binary files /dev/null and b/docs/examples/image/quickstart.png differ diff --git a/docs/examples/qsplot.py b/docs/examples/qsplot.py new file mode 100644 index 00000000..02d82045 --- /dev/null +++ b/docs/examples/qsplot.py @@ -0,0 +1,111 @@ +import os + +import cartopy.crs as ccrs +import flopy +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr +from matplotlib.collections import PolyCollection + +QS_ROOT = os.path.abspath(os.path.join(__file__, "..")) +modelname = "mymodel" +ws = os.path.join(QS_ROOT, "quickstart_data") + + +def _pcolormesh(grid, da, ax): + vertices = [] + for face in grid.iverts: + f = [] + for i in face: + f.append(grid.verts[i]) + vertices.append(f) + vertices = np.array(vertices) + + collection = PolyCollection(vertices) + collection.set_array(da.to_numpy().ravel()) + p = ax.add_collection(collection, autolim=False) + + xmin, xmax, ymin, ymax = grid.extent + ax.set_xlim(xmin, xmax) + ax.set_ylim(ymin, ymax) + cverts = (xmin, ymin), (xmax, ymax) + ax.update_datalim(cverts) + + return p + + +# create budget reader +bpth = os.path.join(ws, f"{modelname}.bud") +bobj = flopy.utils.CellBudgetFile(bpth, precision="double") + +# set specific discharge +spdis = bobj.get_data(text="DATA-SPDIS")[0] + +# create head reader +hpth = os.path.join(ws, f"{modelname}.hds") +hobj = flopy.utils.HeadFile(hpth, precision="double") + +# set heads +heads = hobj.get_alldata() + +# create grb reader +grbpth = os.path.join(ws, f"{modelname}.dis.grb") +grbobj = flopy.mf6.utils.MfGrdFile(grbpth) +grid = None + +# flow vector component arrays +uflow = None +vflow = None + +# create grid object +if "NROW" in grbobj._datadict and "NCOL" in grbobj._datadict: + grid = flopy.discretization.StructuredGrid.from_binary_grid_file(grbpth) + u = [] + v = [] + for r in spdis: + u.append(r[3]) + v.append(r[4]) + uflow = np.array(u).reshape(grid.nrow, grid.ncol) + vflow = np.array(v).reshape(grid.nrow, grid.ncol) +else: + raise Exception("unsupported discretization") + +# set data coordinate arrays +xcrs = grid.xycenters[0] +ycrs = grid.xycenters[1] + +# create qx, qy dataarrys +qx = xr.DataArray(uflow, dims=("y", "x"), coords={"x": xcrs, "y": ycrs}) +qy = xr.DataArray(vflow, dims=("y", "x"), coords={"x": xcrs, "y": ycrs}) + +# create head data array +da = xr.DataArray( + heads[0][0], + dims=("y", "x"), + coords={"x": xcrs, "y": ycrs}, + name="head", +) + +# quickstart mesh with contours and flow arrows +fig, ax = plt.subplots() +_pcolormesh(grid, da, ax) +qz = np.sin(np.sqrt(uflow**2 + vflow**2)) +plt.quiver(xcrs, ycrs, qx, qy, color="w") +cbar = plt.colorbar(label=da.name) +plt.contour(xcrs, ycrs, qz, 4, cmap="viridis") +# plt.contourf(xcrs, ycrs, qz, 5, cmap='viridis', alpha=0.4) +qs_pth = os.path.join(QS_ROOT, "image", "qs.png") +fig.savefig(qs_pth) + +# projection example with cartopy +fig, ax = plt.subplots(subplot_kw={"projection": ccrs.PlateCarree()}) +da.plot(ax=ax, transform=ccrs.PlateCarree()) +ax.coastlines() +ax.stock_img() +ax.set_extent([-5, 15, -4, 14], crs=ccrs.PlateCarree()) +gln = ax.gridlines(draw_labels=True) +gln.top_labels = False +gln.right_labels = False +ax.set_title("Head") +qsprj_pth = os.path.join(QS_ROOT, "image", "qsprj.png") +fig.savefig(qsprj_pth) diff --git a/docs/examples/quickstart.py b/docs/examples/quickstart.py index aca9e221..89d41496 100644 --- a/docs/examples/quickstart.py +++ b/docs/examples/quickstart.py @@ -1,3 +1,7 @@ +from pathlib import Path + +import flopy +import matplotlib.pyplot as plt import numpy as np from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc @@ -37,3 +41,43 @@ # check OC assert oc.data["save_head"][0] == "all" assert oc.data.save_head.sel(per=0) == "all" + +# PLOT +# create budget reader +bpth = Path("./quickstart_data/mymodel.bud") +bobj = flopy.utils.CellBudgetFile(bpth, precision="double") + +# set specific discharge +spdis = bobj.get_data(text="DATA-SPDIS")[0] + +# create head reader +hpth = Path("./quickstart_data/mymodel.hds") +hobj = flopy.utils.HeadFile(hpth, precision="double") + +# set heads +heads = hobj.get_alldata() + +# create grid +grbpth = Path("./quickstart_data/mymodel.dis.grb") +grid = flopy.discretization.StructuredGrid.from_binary_grid_file(grbpth) + +# TODO: get_specific_discharge is dependent on flopy3 model +# qx, qy, qz = flopy.utils.postprocessing.get_specific_discharge(spdis, gwf) + +# set discharge component arrays +u = [] +v = [] +for r in spdis: + u.append(r[3]) + v.append(r[4]) +qx = np.array(u).reshape(grid.nrow, grid.ncol) +qy = np.array(v).reshape(grid.nrow, grid.ncol) + +fig, ax = plt.subplots() +pmv = flopy.plot.PlotMapView(modelgrid=grid, ax=ax) +pmv.plot_array(heads[0][0]) +pmv.plot_grid(colors="white") +pmv.contour_array(heads[0][0], levels=[0.2, 0.4, 0.6, 0.8], linewidths=3.0) +pmv.plot_vector(qx, qy, normalize=True, color="white") +qs_pth = Path("./image/quickstart.png") +fig.savefig(qs_pth)