Coverage for C:\src\imod-python\imod\data\sample_data.py: 52%
110 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-13 11:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-13 11:15 +0200
1"""
2Functions to load sample data.
3"""
5import importlib
6from pathlib import Path
7from typing import Union
8from zipfile import ZipFile
10import numpy as np
11import pandas as pd
12import pooch
13import xarray as xr
14import xugrid as xu
16from imod.formats.prj import open_projectfile_data
17from imod.mf6 import Modflow6Simulation
18from imod.util.imports import MissingOptionalModule
20try:
21 import geopandas as gpd
22except ImportError:
23 gpd = MissingOptionalModule("geopandas")
26def create_pooch_registry() -> pooch.core.Pooch:
27 registry = pooch.create(
28 path=pooch.os_cache("imod"),
29 base_url="https://github.com/Deltares/imod-artifacts/raw/main/",
30 version=None,
31 version_dev="main",
32 env="IMOD_DATA_DIR",
33 )
34 return registry
37def load_pooch_registry(registry: pooch.core.Pooch) -> pooch.core.Pooch:
38 with importlib.resources.files("imod.data") as pkg_dir:
39 registry.load_registry(pkg_dir / "registry.txt")
40 return registry
43REGISTRY = create_pooch_registry()
44REGISTRY = load_pooch_registry(REGISTRY)
47def twri_output(path: Union[str, Path]) -> None:
48 fname_twri = REGISTRY.fetch("ex01-twri-output.zip")
49 with ZipFile(fname_twri) as archive:
50 archive.extractall(path)
53def hondsrug_initial() -> xr.Dataset:
54 fname = REGISTRY.fetch("hondsrug-initial.nc")
55 return xr.open_dataset(fname)
58def hondsrug_layermodel() -> xr.Dataset:
59 fname = REGISTRY.fetch("hondsrug-layermodel.nc")
60 return xr.open_dataset(fname)
63def hondsrug_meteorology() -> xr.Dataset:
64 fname = REGISTRY.fetch("hondsrug-meteorology.nc")
65 return xr.open_dataset(fname)
68def hondsrug_river() -> xr.Dataset:
69 fname = REGISTRY.fetch("hondsrug-river.nc")
70 return xr.open_dataset(fname)
73def hondsrug_drainage() -> xr.Dataset:
74 fname = REGISTRY.fetch("hondsrug-drainage.nc")
75 return xr.open_dataset(fname)
78def head_observations() -> pd.DataFrame:
79 fname = REGISTRY.fetch("head-observations.csv")
80 df = pd.read_csv(fname)
81 # Manually convert time column to datetime type because pandas >2.0 doesn't
82 # do this automatically anymore upon reading.
83 df["time"] = pd.to_datetime(df["time"])
84 return df
87def fluxes() -> xr.Dataset:
88 fname = REGISTRY.fetch("fluxes.nc")
89 return xr.open_dataset(fname)
92def ahn() -> xr.Dataset:
93 fname = REGISTRY.fetch("ahn.nc")
94 return xr.open_dataset(fname)
97def lakes_shp(path: Union[str, Path]) -> "geopandas.GeoDataFrame": # type: ignore # noqa
98 fname_lakes_shp = REGISTRY.fetch("lakes_shp.zip")
99 with ZipFile(fname_lakes_shp) as archive:
100 archive.extractall(path)
101 return gpd.read_file(Path(path) / "lakes.shp")
104def circle() -> xu.Ugrid2d:
105 fname_nodes = REGISTRY.fetch("circle-nodes.txt")
106 fname_triangles = REGISTRY.fetch("circle-triangles.txt")
108 nodes = np.loadtxt(fname_nodes)
109 triangles = np.loadtxt(fname_triangles).astype(np.int32)
111 return xu.Ugrid2d(*nodes.T, -1, triangles)
114def imod5_projectfile_data(path: Union[str, Path]) -> dict:
115 fname_model = REGISTRY.fetch("iMOD5_model.zip")
117 with ZipFile(fname_model) as archive:
118 archive.extractall(path)
120 return open_projectfile_data(Path(path) / "iMOD5_model_pooch" / "iMOD5_model.prj")
123def hondsrug_simulation(path: Union[str, Path]) -> Modflow6Simulation:
124 fname_simulation = REGISTRY.fetch("hondsrug-simulation.zip")
126 with ZipFile(fname_simulation) as archive:
127 archive.extractall(path)
129 simulation = Modflow6Simulation.from_file(Path(path) / "mf6-hondsrug-example.toml")
130 # The model was written before the xt3d_option and rhs_option arguments were
131 # added to iMOD Python. Set missing options to False.
132 simulation["GWF"]["npf"].set_xt3d_option(is_xt3d_used=False, is_rhs=False)
134 return simulation
137def hondsrug_crosssection(path: Union[str, Path]) -> "geopandas.GeoDataFrame": # type: ignore # noqa
138 fname_simulation = REGISTRY.fetch("hondsrug-crosssection.zip")
140 with ZipFile(fname_simulation) as archive:
141 archive.extractall(path)
143 return gpd.read_file(Path(path) / "crosssection.shp")
146def hondsrug_layermodel_topsystem() -> xr.Dataset:
147 """
148 This is a modified version of the hondsrug_layermodel, used for the
149 topsystem example in the user guide. n_max_old original layers are
150 taken and subdivided into n_new layers. This makes for more layers around
151 the topsystem.
152 """
153 fname = REGISTRY.fetch("hondsrug-layermodel.nc")
154 layer_model = xr.open_dataset(fname)
155 # Make layer model more interesting for this example by subdividing layers
156 # into n_new layers.
157 n_new = 4
158 n_max_old = 5
160 # Loop over original layers until n_max_old and subdivide each into n_new
161 # layers.
162 new_ds_ls = []
163 for i in range(n_max_old):
164 sub_iter = np.arange(n_new) + 1
165 layer_coord = sub_iter + i * (n_max_old - 1)
166 distribution_factors = 1 / n_new * sub_iter
167 da_distribution = xr.DataArray(
168 distribution_factors, coords={"layer": layer_coord}, dims=("layer",)
169 )
170 layer_model_sel = layer_model.sel(layer=i + 1, drop=True)
171 # Compute thickness
172 D = layer_model_sel["top"] - layer_model_sel["bottom"]
174 new_ds = xr.Dataset()
175 new_ds["k"] = xr.ones_like(da_distribution) * layer_model_sel["k"]
176 new_ds["idomain"] = xr.ones_like(da_distribution) * layer_model_sel["idomain"]
177 # Put da_distribution in front of equation to enforce dims as (layer, y, x)
178 new_ds["top"] = (da_distribution - 1 / n_new) * -D + layer_model_sel["top"]
179 new_ds["bottom"] = da_distribution * -D + layer_model_sel["top"]
181 new_ds_ls.append(new_ds)
183 return xr.concat(new_ds_ls, dim="layer")
186def hondsrug_planar_river() -> xr.Dataset:
187 """
188 This is the hondsrug river dataset with the following modifications:
190 1) Aggregated over layer dimension to create planar grid.
191 2) Stages raised towards the top of the model, as the original stages are
192 most of the time laying at bottom elevation, making for boring examples.
194 """
195 fname = REGISTRY.fetch("hondsrug-river.nc")
196 planar_river = xr.open_dataset(fname).max(dim="layer")
198 fname = REGISTRY.fetch("hondsrug-layermodel.nc")
199 top = xr.open_dataset(fname)["top"].sel(layer=1)
201 planar_river["stage"] = (top - planar_river["stage"]) / 2 + planar_river["stage"]
203 return planar_river