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
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +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")
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")
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)
42def hondsrug_initial() -> xr.Dataset:
43 fname = REGISTRY.fetch("hondsrug-initial.nc")
44 return xr.open_dataset(fname)
47def hondsrug_layermodel() -> xr.Dataset:
48 fname = REGISTRY.fetch("hondsrug-layermodel.nc")
49 return xr.open_dataset(fname)
52def hondsrug_meteorology() -> xr.Dataset:
53 fname = REGISTRY.fetch("hondsrug-meteorology.nc")
54 return xr.open_dataset(fname)
57def hondsrug_river() -> xr.Dataset:
58 fname = REGISTRY.fetch("hondsrug-river.nc")
59 return xr.open_dataset(fname)
62def hondsrug_drainage() -> xr.Dataset:
63 fname = REGISTRY.fetch("hondsrug-drainage.nc")
64 return xr.open_dataset(fname)
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
76def fluxes() -> xr.Dataset:
77 fname = REGISTRY.fetch("fluxes.nc")
78 return xr.open_dataset(fname)
81def ahn() -> xr.Dataset:
82 fname = REGISTRY.fetch("ahn.nc")
83 return xr.open_dataset(fname)
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")
93def circle() -> xu.Ugrid2d:
94 fname_nodes = REGISTRY.fetch("circle-nodes.txt")
95 fname_triangles = REGISTRY.fetch("circle-triangles.txt")
97 nodes = np.loadtxt(fname_nodes)
98 triangles = np.loadtxt(fname_triangles).astype(np.int32)
100 return xu.Ugrid2d(*nodes.T, -1, triangles)
103def imod5_projectfile_data(path: Union[str, Path]) -> dict:
104 fname_model = REGISTRY.fetch("iMOD5_model.zip")
106 with ZipFile(fname_model) as archive:
107 archive.extractall(path)
109 return open_projectfile_data(Path(path) / "iMOD5_model_pooch" / "iMOD5_model.prj")
112def hondsrug_simulation(path: Union[str, Path]) -> Modflow6Simulation:
113 fname_simulation = REGISTRY.fetch("hondsrug-simulation.zip")
115 with ZipFile(fname_simulation) as archive:
116 archive.extractall(path)
118 return Modflow6Simulation.from_file(Path(path) / "mf6-hondsrug-example.toml")
121def hondsrug_crosssection(path: Union[str, Path]) -> "geopandas.GeoDataFrame": # type: ignore # noqa
122 fname_simulation = REGISTRY.fetch("hondsrug-crosssection.zip")
124 with ZipFile(fname_simulation) as archive:
125 archive.extractall(path)
127 return gpd.read_file(Path(path) / "crosssection.shp")