mf6_netcdf01_tutorial.py

# ---
# jupyter:
#   jupytext:
#     cell_metadata_filter: -all
#     formats: ipynb,py:light
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.17.2
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# # MODFLOW 6: MODFLOW 6 NetCDF inputs from FloPy simulation
#
# ## Write MODFLOW 6 NetCDF simulation
#
# This tutorial demonstrates how to generate a MODFLOW 6 NetCDF file
# from an existing FloPy simulation. Three variations will be shown.
# In the first, FloPy will generate the file with a modified
# `write_simulation()` call.  The second method is more interactive,
# providing an opputinity to modify the dataset before it is written
# to NetCDF.  The third example will show how to update the simulation
# and dataset for a few specific parameters.
#
# Support for generating a MODFLOW 6 NetCDF input without a defined
# FloPy mf6 model or package instances is briefly discussed at the
# end of the tutorial.
#
# For more information on supported MODFLOW 6 NetCDF formats see:
# [MODFLOW NetCDF Format](https://github.com/MODFLOW-ORG/modflow6/wiki/MODFLOW-NetCDF-Format).
#
# Note that NetCDF is only supported by the Extended version of MODFLOW 6.
# A nightly windows build of Extended MODFLOW 6 is available from
# [nightly build](https://github.com/MODFLOW-ORG/modflow6-nightly-build).

# package import
import sys
from pathlib import Path
from pprint import pformat, pprint
from tempfile import TemporaryDirectory

import git
import numpy as np
import pooch
import xarray as xr

import flopy

print(sys.version)
print(f"flopy version: {flopy.__version__}")

sim_name = "uzf01"

# ### Check if we are in the repository and define the data path.

try:
    root = Path(git.Repo(".", search_parent_directories=True).working_dir)
except:
    root = None

data_path = root / "examples" / "data" / "mf6" / "netcdf" if root else Path.cwd()

file_names = {
    "mfsim.nam": None,
    "uzf01.dis": None,
    "uzf01.ghb.obs": None,
    "uzf01.ghbg": None,
    "uzf01.ic": None,
    "uzf01.ims": None,
    "uzf01.nam": None,
    "uzf01.npf": None,
    "uzf01.obs": None,
    "uzf01.oc": None,
    "uzf01.sto": None,
    "uzf01.tdis": None,
    "uzf01.uzf": None,
    "uzf01.uzf.obs": None,
}

for fname, fhash in file_names.items():
    pooch.retrieve(
        url=f"https://github.com/modflowpy/flopy/raw/develop/examples/data/{sim_name}/{fname}",
        fname=fname,
        path=data_path / sim_name,
        known_hash=fhash,
    )

# ### Create simulation workspace

# create temporary directories
temp_dir = TemporaryDirectory()
workspace = Path(temp_dir.name)

# ### Load and run baseline simulation
#
# For the purposes of this tutorial, the specifics of the simulation
# other than it is a candidate for NetCDF input are not a focus. It
# is a NetCDF input candidate because it defines a supported model type
# (`GWF6`) with a structured discretization and packages that support
# NetCDF input parameters.
#
# More information about package NetCDF support in MODFLOW 6 can be
# found in the `MODFLOW 6 - Description of Input and Output` (`mf6io.pdf`),
# also available at the nightly build repository linked above.

# load and run the ASCII based simulation
sim = flopy.mf6.MFSimulation.load(sim_ws=data_path / sim_name)
sim.set_sim_path(workspace)
sim.write_simulation()
success, buff = sim.run_simulation(silent=True, report=True)
assert success, pformat(buff)

# ## Create NetCDF based simulation method 1 (non-interactive)
#
# The most straightforward way to create a NetCDF simulation
# from the loaded simulation is to provide a `netcdf` argument
# to `write_simulation()` and define it to be either `structured`
# or `layered`, depending on the desired format of the generated
# NetCDF file.
#
# The name of the created file can be specified by first setting the
# model `name_file.nc_filerecord` attribute to the desired name. If
# this step is not taken, the default name of `{model_name}.input.nc`
# is used.

# create directory for the NetCDF sim
sim.set_sim_path(workspace / "netcdf1")
# set model name file nc_filerecord attribute to export name
gwf = sim.get_model("uzf01")
gwf.name_file.nc_filerecord = "uzf01.structured.nc"
# write simulation with structured NetCDF file
sim.write_simulation(netcdf="structured")

# ### Run MODFLOW 6 simulation with NetCDF input
#
# The simulation generated by this tutorial should run with the
# Extended MODFLOW 6 executable, available from the nightly-build
# repository (linked above).

# success, buff = sim.run_simulation(silent=True, report=True)
# assert success, pformat(buff)

# ### Repeat method 1 with layered mesh NetCDF format

# create directory for the NetCDF sim
sim.set_sim_path(workspace / "netcdf2")
# set model name file nc_filerecord attribute to export name
gwf = sim.get_model("uzf01")
gwf.name_file.nc_filerecord = "uzf01.layered.nc"
# write simulation with with layered mesh NetCDF file
sim.write_simulation(netcdf="layered")

# ### Run MODFLOW 6 simulation with NetCDF input
#
# The simulation generated by this tutorial should run with the
# Extended MODFLOW 6 executable, available from the nightly-build
# repository (linked above).

# success, buff = sim.run_simulation(silent=True, report=True)
# assert success, pformat(buff)

# ## Create NetCDF based simulation method 2 (interactive)
#
# In this method we will set the FloPy `netcdf` argument to `` (empty
# string) when `write_simulation()` is called. As such, FloPy will
# not generate the NetCDF file automatically. We will manage the NetCDF
# file generation ourselves in method 2.
#
# Reset the simulation path and set the `GWF` name file `nc_filerecord`
# attribute to the name of the intended input NetCDF file. Display
# the resultant name file changes.
#
# When we write the updated simulation, all packages that support NetCDF
# input parameters will be written such that NetCDF parameters expect the
# data source to be a NetCDF file. We will therefore need to create a
# NetCDF input file containing arrays for the `DIS`, `NPF`, `IC`, `STO`,
# and `GHBG` packages. We will still use FloPy package objects to set the
# parameter data in the dataset; however an `update_data=False` argument
# could be passed to the `update_dataset()` call if this was not desired.

# create directory for the NetCDF sim
sim.set_sim_path(workspace / "netcdf3")
# set model name file nc_filerecord attribute to export name
gwf = sim.get_model("uzf01")
gwf.name_file.nc_filerecord = "uzf01.structured.nc"
# write simulation with ASCII inputs tagged for NetCDF
# but do not create NetCDF file
sim.write_simulation(netcdf="")

# ### Show name file with NetCDF input configured

# show name file with NetCDF input configured
with open(workspace / "netcdf3" / "uzf01.nam", "r") as fh:
    print(fh.read())

# ### Show example package file with NetCDF keywords

# show example package file with NetCDF input configured
with open(workspace / "netcdf3" / "uzf01.ic", "r") as fh:
    print(fh.read())

# ### Create dataset
#
# Create the base xarray dataset from the modelgrid object. This
# will add required dimensions and coordinate variables to the
# dataset according to the grid specification. Modeltime is needed
# for timeseries support.

# create the dataset
ds = gwf.modelgrid.dataset(modeltime=gwf.modeltime)

# ### Access model NetCDF attributes
#
# Internally, FloPy generates and uses NetCDF metadata dictionaries
# to update datasets. Both model and package objects can generate
# the dictionaries and they contain related but distinct NetCDF
# details. Model object dictionaries contain file scoped attribute
# information while package dictionaries are organized by NetCDF
# variable and contain variable scoped attribute information and
# details related to creating the variable, including dimensions,
# shape and data type.
#
# The dictionaries are available via the `netcdf_meta()` call. Their
# content also varies depending on the desired type of dataset (i.e.
# `structured` or `layered`). In this step we will access the model
# NetCDF metadata and use it to update dataset scoped attributes.
#
# First, retrieve and store the NetCDF metadata dictionary and display
# its contents. Then, in the following step, update the dataset with
# the model scoped attributes defined in the dictionary.
#
# These operations can also both be accomplished by calling
# `update_dataset()` on the model object, albeit without the
# opportunity to modify the intermediate metadata dictionary.
# Examples of this approach (with package objects) are shown below.

# retrieve the model NetCDF meta dictionary
nc_meta = gwf.netcdf_meta()
pprint(nc_meta)

# update dataset directly with required attributes
for a in nc_meta["attrs"]:
    ds.attrs[a] = nc_meta["attrs"][a]

# ### Update the dataset with supported `DIS` arrays
#
# First, we will show how package NetCDF parameters can be added to the
# dataset without using the NetCDF metadata dictionary. We will use the
# metadata dictionary when updating the dataset with NPF parameter data.
#
# Add NetCDF supported data arrays in package to dataset.

# update dataset with `DIS` arrays
dis = gwf.get_package("dis")
ds = dis.update_dataset(ds)

# ### Access `NPF` package NetCDF attributes
#
# Access package scoped NetCDF details by storing the dictionary returned
# from `netcdf_meta()`.
#
# The contents of the info dictionary are shown and then, in the following
# step, the dictionary and the dataset are passed to a helper routine that
# create the intended array variables.

# get npf package netcdf info
npf = gwf.get_package("npf")
nc_meta = npf.netcdf_meta()
pprint(nc_meta)

# ### Update package NetCDF metadata dictionary and dataset
#
# Here we update the metadata dictionary and then pass it directly to the
# `update_dataset()` function which uses it when adding variables to the
# dataset.
#
# We replace the default name for the `NPF K` input parameter and add
# the `standard_name` attribute to it's attribute dictionary. The
# dictionary is then passed to the `update_dataset()` function. Note the
# updated name is used in the subsequent block when updating the array
# values.

# update dataset with `NPF` arrays
nc_meta["k"]["varname"] = "npf_k_updated"
nc_meta["k"]["attrs"]["standard_name"] = "soil_hydraulic_conductivity_at_saturation"
ds = npf.update_dataset(ds, netcdf_meta=nc_meta)

# ### Show dataset `NPF K` parameter with updates

# print dataset npf k variable
print(ds["npf_k_updated"])

# ### Update the dataset with supported `IC` arrays

# ic
ic = gwf.get_package("ic")
ds = ic.update_dataset(ds)

# ### Update the dataset with supported `STO` arrays

# storage
sto = gwf.get_package("sto")
ds = sto.update_dataset(ds)

# ### Update the dataset with supported `GHBG` arrays

# update dataset with 'GHBG' arrays
ghbg = gwf.get_package("ghbg_0")
ds = ghbg.update_dataset(ds)

# ### Display generated dataset

# show the dataset
print(ds)

# ### Export generated dataset to NetCDF

# write dataset to NetCDF
ds.to_netcdf(
    workspace / "netcdf3" / "uzf01.structured.nc", format="NETCDF4", engine="netcdf4"
)

# ### Run MODFLOW 6 simulation with NetCDF input
#
# The simulation generated by this tutorial should run with the
# Extended MODFLOW 6 executable, available from the nightly-build
# repository (linked above).

# success, buff = sim.run_simulation(silent=True, report=True)
# assert success, pformat(buff)

# ## Support for independently creating and managing the dataset
#
# The NetCDF meta dictionaries are intended in part to support various
# workflows that might be used to generate a MODFLOW 6 input NetCDF file.
# A few more examples are shown here to demonstrate flexibility with
# respect to approach. First, we will use FloPy objects to retrieve
# instance NetCDF meta dictionaries but not to update the dataset. Finally,
# we will show how to retrieve NetCDF meta dictionaries when no FloPy
# objects are at hand.
#
# In particular, the next example shows how to add only select input
# parameters to the dataset. At this time, this sort of approach requires
# manual updating of the MODFLOW 6 ASCII input files, which is also
# shown here.

# ### Reload and write the simulation so that ASCII files can be edited

sim = flopy.mf6.MFSimulation.load(sim_ws=data_path / sim_name)
gwf = sim.get_model("uzf01")
gwf.name_file.nc_filerecord = "uzf01.structured.nc"
sim.set_sim_path(workspace / "netcdf4")
sim.write_simulation()

# ### Show name file with NetCDF input configured

# As the nc_filerecord attribute was configured, name file is already
# updated for NetCDF input
with open(workspace / "netcdf4" / "uzf01.nam", "r") as fh:
    print(fh.read())

# ### Create the dataset
#
# Although we are accessing the modelgrid through the FloPy model
# object, this does not need to be case.  Creating the modelgrid object
# independent of the model provides the same functionality.

modelgrid = gwf.modelgrid
modeltime = gwf.modeltime
ds = modelgrid.dataset(modeltime=gwf.modeltime)

# ### Retrieve model NetCDF meta dictionary

nc_meta = gwf.netcdf_meta()
pprint(nc_meta)

# ### Update dataset with required dataset scoped attributes

for a in nc_meta["attrs"]:
    ds.attrs[a] = nc_meta["attrs"][a]

# ### Retrieve NPF NetCDF meta dictionary

npf_nc_meta = gwf.get_package("npf").netcdf_meta()
pprint(npf_nc_meta)

# ### Set actual dimensions in order of npf_nc_meta["k"]["netcdf_shape"]

shape = (modelgrid.nlay, modelgrid.ncol, modelgrid.nrow)

# ### Set varname and default data using NPF NetCDF meta dictionary

varname = npf_nc_meta["k"]["varname"]
data = np.full(
    shape,
    npf_nc_meta["k"]["attrs"]["_FillValue"],
    dtype=npf_nc_meta["k"]["xarray_type"],
)

# ### Add parameter with fill values to dataset

var_d = {varname: (npf_nc_meta["k"]["netcdf_shape"], data)}
ds = ds.assign(var_d)

# ### Update parameter attributes from NetCDF meta dictionary

for a in npf_nc_meta["k"]["attrs"]:
    ds[varname].attrs[a] = npf_nc_meta["k"]["attrs"][a]

# ### Update parameter data (set identical to ASCII sim)

k = np.ones(shape, dtype=float)
ds[varname].values = k

# ### Edit the NPF package ASCII file to specify K is in dataset

with open(workspace / "netcdf4" / "uzf01.npf", "w") as f:
    f.write("BEGIN options\n")
    f.write("END options\n\n")
    f.write("BEGIN griddata\n")
    f.write("  icelltype\n")
    f.write("    CONSTANT 1\n")
    f.write("  k netcdf\n")
    f.write("END griddata\n")

# ### Retrieve GHBG NetCDF meta dictionary

ghbg_nc_meta = gwf.get_package("ghbg").netcdf_meta()
pprint(ghbg_nc_meta)

# ### Set actual dimensions in order of ghbg_nc_meta["bhead"]["netcdf_shape"]

shape = (sum(modeltime.nstp), modelgrid.nlay, modelgrid.ncol, modelgrid.nrow)

# ### Set varnames and data with help of GHBG NetCDF meta dictionary

varnames = ["bhead", "cond"]
idata = [1.5, 1.0]
for i, v in enumerate(varnames):
    varname = ghbg_nc_meta[v]["varname"]
    data = np.full(
        shape,
        ghbg_nc_meta[v]["attrs"]["_FillValue"],
        dtype=ghbg_nc_meta[v]["xarray_type"],
    )

    # add parameters to dataset
    var_d = {varname: (ghbg_nc_meta[v]["netcdf_shape"], data)}
    ds = ds.assign(var_d)

    # apply parameter attributes
    for a in ghbg_nc_meta[v]["attrs"]:
        ds[varname].attrs[a] = ghbg_nc_meta[v]["attrs"][a]

    # update data to be consistent with ASCII simulation
    ds[varname].values[0, 99, 0, 0] = idata[i]

# ### Edit the GHBG package ASCII file to specify BHEAD and COND are in the dataset

with open(workspace / "netcdf4" / "uzf01.ghbg", "w") as f:
    f.write("BEGIN options\n")
    f.write("  READARRAYGRID\n")
    f.write("  PRINT_INPUT\n")
    f.write("  PRINT_FLOWS\n")
    f.write("  OBS6  FILEIN  uzf01.ghb.obs\n")
    f.write("END options\n\n")
    f.write("BEGIN period  1\n")
    f.write("  bhead  netcdf\n")
    f.write("  cond  netcdf\n")
    f.write("END period  1\n")

# ### Display the dataset

print(ds)

# ### Write the dataset

ds.to_netcdf(
    workspace / "netcdf4" / "uzf01.structured.nc", format="NETCDF4", engine="netcdf4"
)

# ### Run MODFLOW 6 simulation with NetCDF input
#
# The simulation generated by this tutorial should run with the
# Extended MODFLOW 6 executable, available from the nightly-build
# repository (linked above).

# success, buff = sim.run_simulation(silent=True, report=True)
# assert success, pformat(buff)

# ## NetCDF meta dictionary without FloPy objects
#
# The above method still uses FloPy objects to update the dataset arrays
# to values consistent with the state of the objects. The `netcdf_meta`
# dictionary is intended to supported creation of the dataset without
# an existing simulation defined. The base dataset can be defined with
# `modelgrid` and `modeltime` objects, while the full package NetCDF meta
# dictionary can retrieved with a static call to a model or package mf6
# type. The auxiliary input is optional but does show the variables that
# would be required if auxiliary variables were defined in the package.

# ### Demonstrate static call on MFPackage (structured dataset):

netcdf_meta = flopy.mf6.mfpackage.MFPackage.netcdf_package(
    mtype="GWF",
    ptype="NPF",
    auxiliary=["CONCENTRATION"],
)
pprint(netcdf_meta)

# ### Demonstrate static call on MFPackage (layered dataset):

netcdf_meta = flopy.mf6.mfpackage.MFPackage.netcdf_package(
    mtype="GWF", ptype="NPF", mesh="LAYERED", auxiliary=["CONCENTRATION"], nlay=2
)
pprint(netcdf_meta)