Coverage for C:\src\imod-python\imod\data\sample_data.py: 65%

75 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +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 

25REGISTRY = pooch.create( 

26 path=pooch.os_cache("imod"), 

27 base_url="https://github.com/Deltares/imod-artifacts/raw/main/", 

28 version=None, 

29 version_dev="main", 

30 env="IMOD_DATA_DIR", 

31) 

32with importlib.resources.files("imod.data") as pkg_dir: 

33 REGISTRY.load_registry(pkg_dir / "registry.txt") 

34 

35 

36def twri_output(path: Union[str, Path]) -> None: 

37 fname_twri = REGISTRY.fetch("ex01-twri-output.zip") 

38 with ZipFile(fname_twri) as archive: 

39 archive.extractall(path) 

40 

41 

42def hondsrug_initial() -> xr.Dataset: 

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

44 return xr.open_dataset(fname) 

45 

46 

47def hondsrug_layermodel() -> xr.Dataset: 

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

49 return xr.open_dataset(fname) 

50 

51 

52def hondsrug_meteorology() -> xr.Dataset: 

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

54 return xr.open_dataset(fname) 

55 

56 

57def hondsrug_river() -> xr.Dataset: 

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

59 return xr.open_dataset(fname) 

60 

61 

62def hondsrug_drainage() -> xr.Dataset: 

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

64 return xr.open_dataset(fname) 

65 

66 

67def head_observations() -> pd.DataFrame: 

68 fname = REGISTRY.fetch("head-observations.csv") 

69 df = pd.read_csv(fname) 

70 # Manually convert time column to datetime type because pandas >2.0 doesn't 

71 # do this automatically anymore upon reading. 

72 df["time"] = pd.to_datetime(df["time"]) 

73 return df 

74 

75 

76def fluxes() -> xr.Dataset: 

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

78 return xr.open_dataset(fname) 

79 

80 

81def ahn() -> xr.Dataset: 

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

83 return xr.open_dataset(fname) 

84 

85 

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

87 fname_lakes_shp = REGISTRY.fetch("lakes_shp.zip") 

88 with ZipFile(fname_lakes_shp) as archive: 

89 archive.extractall(path) 

90 return gpd.read_file(Path(path) / "lakes.shp") 

91 

92 

93def circle() -> xu.Ugrid2d: 

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

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

96 

97 nodes = np.loadtxt(fname_nodes) 

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

99 

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

101 

102 

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

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

105 

106 with ZipFile(fname_model) as archive: 

107 archive.extractall(path) 

108 

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

110 

111 

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

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

114 

115 with ZipFile(fname_simulation) as archive: 

116 archive.extractall(path) 

117 

118 return Modflow6Simulation.from_file(Path(path) / "mf6-hondsrug-example.toml") 

119 

120 

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

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

123 

124 with ZipFile(fname_simulation) as archive: 

125 archive.extractall(path) 

126 

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