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

1""" 

2Functions to load sample data. 

3""" 

4 

5import importlib 

6from pathlib import Path 

7from typing import Union 

8from zipfile import ZipFile 

9 

10import numpy as np 

11import pandas as pd 

12import pooch 

13import xarray as xr 

14import xugrid as xu 

15 

16from imod.formats.prj import open_projectfile_data 

17from imod.mf6 import Modflow6Simulation 

18from imod.util.imports import MissingOptionalModule 

19 

20try: 

21 import geopandas as gpd 

22except ImportError: 

23 gpd = MissingOptionalModule("geopandas") 

24 

25 

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 

35 

36 

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 

41 

42 

43REGISTRY = create_pooch_registry() 

44REGISTRY = load_pooch_registry(REGISTRY) 

45 

46 

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) 

51 

52 

53def hondsrug_initial() -> xr.Dataset: 

54 fname = REGISTRY.fetch("hondsrug-initial.nc") 

55 return xr.open_dataset(fname) 

56 

57 

58def hondsrug_layermodel() -> xr.Dataset: 

59 fname = REGISTRY.fetch("hondsrug-layermodel.nc") 

60 return xr.open_dataset(fname) 

61 

62 

63def hondsrug_meteorology() -> xr.Dataset: 

64 fname = REGISTRY.fetch("hondsrug-meteorology.nc") 

65 return xr.open_dataset(fname) 

66 

67 

68def hondsrug_river() -> xr.Dataset: 

69 fname = REGISTRY.fetch("hondsrug-river.nc") 

70 return xr.open_dataset(fname) 

71 

72 

73def hondsrug_drainage() -> xr.Dataset: 

74 fname = REGISTRY.fetch("hondsrug-drainage.nc") 

75 return xr.open_dataset(fname) 

76 

77 

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 

85 

86 

87def fluxes() -> xr.Dataset: 

88 fname = REGISTRY.fetch("fluxes.nc") 

89 return xr.open_dataset(fname) 

90 

91 

92def ahn() -> xr.Dataset: 

93 fname = REGISTRY.fetch("ahn.nc") 

94 return xr.open_dataset(fname) 

95 

96 

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") 

102 

103 

104def circle() -> xu.Ugrid2d: 

105 fname_nodes = REGISTRY.fetch("circle-nodes.txt") 

106 fname_triangles = REGISTRY.fetch("circle-triangles.txt") 

107 

108 nodes = np.loadtxt(fname_nodes) 

109 triangles = np.loadtxt(fname_triangles).astype(np.int32) 

110 

111 return xu.Ugrid2d(*nodes.T, -1, triangles) 

112 

113 

114def imod5_projectfile_data(path: Union[str, Path]) -> dict: 

115 fname_model = REGISTRY.fetch("iMOD5_model.zip") 

116 

117 with ZipFile(fname_model) as archive: 

118 archive.extractall(path) 

119 

120 return open_projectfile_data(Path(path) / "iMOD5_model_pooch" / "iMOD5_model.prj") 

121 

122 

123def hondsrug_simulation(path: Union[str, Path]) -> Modflow6Simulation: 

124 fname_simulation = REGISTRY.fetch("hondsrug-simulation.zip") 

125 

126 with ZipFile(fname_simulation) as archive: 

127 archive.extractall(path) 

128 

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) 

133 

134 return simulation 

135 

136 

137def hondsrug_crosssection(path: Union[str, Path]) -> "geopandas.GeoDataFrame": # type: ignore # noqa 

138 fname_simulation = REGISTRY.fetch("hondsrug-crosssection.zip") 

139 

140 with ZipFile(fname_simulation) as archive: 

141 archive.extractall(path) 

142 

143 return gpd.read_file(Path(path) / "crosssection.shp") 

144 

145 

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 

159 

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"] 

173 

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"] 

180 

181 new_ds_ls.append(new_ds) 

182 

183 return xr.concat(new_ds_ls, dim="layer") 

184 

185 

186def hondsrug_planar_river() -> xr.Dataset: 

187 """ 

188 This is the hondsrug river dataset with the following modifications: 

189 

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. 

193 

194 """ 

195 fname = REGISTRY.fetch("hondsrug-river.nc") 

196 planar_river = xr.open_dataset(fname).max(dim="layer") 

197 

198 fname = REGISTRY.fetch("hondsrug-layermodel.nc") 

199 top = xr.open_dataset(fname)["top"].sel(layer=1) 

200 

201 planar_river["stage"] = (top - planar_river["stage"]) / 2 + planar_river["stage"] 

202 

203 return planar_river