diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index 1b4992f32..56859502b 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -39,6 +39,10 @@ Added the final package has fully empty layers trimmed off afterwards, so this does not affect computed values. Set to False to keep the previous full-layer-coordinate behaviour. +- Added option ``chunk_layout`` to :func:`imod.idf.open` to control how the data is + loaded in chunks, which can improve performance and memory usage for large datasets. + The default option ``per_idf`` keeps the old behavior, by chunking per IDF file. + The alternative option ``per_timestep`` chunks per timestep. - Added :class:`imod.mf6.AdvectionUTVD` to the transport packages. This is the most accurate advection scheme in most cases, and is recommended for when sharp fronts need to be preserved or when computing on an unstructured grid. @@ -76,6 +80,9 @@ Changed (:class:`imod.mf6.River`, :class:`imod.mf6.GeneralHeadBoundary`, :class:`imod.mf6.Drainage`, :class:`imod.mf6.Recharge`) will also be masked where constant head cells are placed. +- :func:`imod.formats.prj.open_projectfile_data` now chunks the data by time + only instead of by both time and layer. This should enhance performance when + working with large datasets. [1.1.0] - 2026-08-03 -------------------- diff --git a/imod/formats/array_io/reading.py b/imod/formats/array_io/reading.py index 412513214..6ba82329f 100644 --- a/imod/formats/array_io/reading.py +++ b/imod/formats/array_io/reading.py @@ -280,8 +280,33 @@ def _dask(path, attrs=None, pattern=None, _read=None, header=None): return x, attrs -def _load(paths, use_cftime, _read, headers): - """Combine a list of paths to IDFs to a single xarray.DataArray""" +def load_as_components(paths, use_cftime, _read, headers): + """ + Combine a list of paths to IDFs to components that can be used to construct + a single xarray.DataArray. + + Parameters + ---------- + paths : list[str] + List of file paths to IDF files. + use_cftime : bool + Whether to use cftime for time coordinates. + _read : Callable + Function to read individual IDF files. + headers : list[dict] + List of headers corresponding to each IDF file. + + Returns + ------- + dask.array + Dask array containing the combined data from all IDF files. + dict + Dictionary of coordinates for the DataArray. + list[str] + List of dimension names for the DataArray. + str + Name of the DataArray. + """ # this function also works for single IDFs names = [h["name"] for h in headers] _all_equal(names, "names") @@ -320,8 +345,14 @@ def _load(paths, use_cftime, _read, headers): nested_dict.set_nested(groupby, groupbykeys, da) dask_arrays = nested_dict.sorted_nested_dict(groupby) dask_array = _ndconcat(dask_arrays, ndim) + return dask_array, coords, dims, names[0] + - out = xr.DataArray(dask_array, coords, dims, name=names[0]) +def _load(paths, use_cftime, _read, headers): + dask_array, coords, dims, name = load_as_components( + paths, use_cftime, _read, headers + ) + out = xr.DataArray(dask_array, coords, dims, name=name) first_attrs = headers[0] @@ -333,7 +364,27 @@ def _load(paths, use_cftime, _read, headers): return out -def _open(path, use_cftime, pattern, header, _read): +def handle_path( + path: pathlib.Path | str | list[pathlib.Path | str], +) -> list[pathlib.Path]: + """ + Handle a path input and return a list of pathlib.Path objects. + + Parameters + ---------- + path : pathlib.Path, str, or list of pathlib.Path + The input path(s) to handle. + + Returns + ------- + list[pathlib.Path] + List of pathlib.Path objects corresponding to the input path(s). + + Raises + ------ + FileNotFoundError + If no files matching the input path(s) are found. + """ if isinstance(path, pathlib.Path): path = str(path) @@ -342,8 +393,16 @@ def _open(path, use_cftime, pattern, header, _read): else: paths = [pathlib.Path(p) for p in glob.glob(path)] - headers = [header(p, pattern) for p in paths] n = len(paths) if n == 0: raise FileNotFoundError(f"Could not find any files matching {path}") + + return paths + + +def _open(path, use_cftime, pattern, header, _read): + paths = handle_path(path) + + headers = [header(p, pattern) for p in paths] + return _load(paths, use_cftime, _read, headers) diff --git a/imod/formats/idf.py b/imod/formats/idf.py index 68eb1e1c8..b797b0540 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -12,12 +12,13 @@ from collections.abc import Iterable from pathlib import Path from re import Pattern -from typing import Any, DefaultDict +from typing import Any, Callable, DefaultDict, Literal, Sequence, TypeAlias import dask import dask.array import numpy as np import xarray as xr +from numpy.typing import DTypeLike import imod from imod.formats import array_io @@ -26,8 +27,11 @@ # Make sure we can still use the built-in function... f_open = open +PatternType: TypeAlias = Pattern[str] | str +PathType: TypeAlias = str | Path -def header(path, pattern): + +def header(path: PathType, pattern: PatternType | None) -> dict[str, Any]: """Read the IDF header information into a dictionary""" attrs = imod.util.path.decompose(path, pattern) with f_open(path, "rb") as f: @@ -105,7 +109,14 @@ def header(path, pattern): return attrs -def _read(path, headersize, nrow, ncol, nodata, dtype): +def _read( + path: PathType, + headersize: int, + nrow: int, + ncol: int, + nodata: float, + dtype: DTypeLike, +) -> np.ndarray: """ Read a single IDF file to a numpy.ndarray @@ -132,8 +143,12 @@ def _read(path, headersize, nrow, ncol, nodata, dtype): return array_io.reading._to_nan(a, nodata) -# Open IDFs for multiple times and/or layers into one DataArray -def open(path, use_cftime=False, pattern=None): +def open( + path: PathType | list[PathType], + use_cftime=False, + pattern: PatternType | None = None, + chunk_layout: Literal["per_idf", "per_timestep"] = "per_idf", +) -> xr.DataArray: r""" Open one or more IDF files as an xarray.DataArray. @@ -163,6 +178,12 @@ def open(path, use_cftime=False, pattern=None): {name}_{time}_l{layer}, a custom pattern can be defined here either as a string, or as a compiled regular expression pattern. See the examples below. + chunk_layout : {"per_idf", "per_timestep"}, optional + Determines how the data should be chunked when loaded into memory. + ``"per_idf"`` chunks by individual IDF files, while + ``"per_timestep"`` chunks by timestep, grouping all layers within a + timestep into one chunk. This can help optimize memory usage and + computational efficiency. The default is ``"per_idf"``. Returns ------- @@ -213,6 +234,52 @@ def open(path, use_cftime=False, pattern=None): generally a fiddly process. Regex notation is also impossible to remember. The website https://regex101.com is a nice help. Alternatively, the most pragmatic solution may be to just rename your files. + + To increase performance when opening multiple IDF files, consider setting + the ``chunk_layout`` parameter to ``"per_timestep"`` to control how the + data is chunked in + memory. Data will be chunked along the time dimension, instead of per idf + file (so most often per time, per layer), which is the default. + + >>> head = imod.idf.open("head_2001*_l*.idf", chunk_layout="per_timestep") + + """ + + if chunk_layout == "per_idf": + return _open(path, use_cftime=use_cftime, pattern=pattern) + elif chunk_layout == "per_timestep": + return open_chunked_per_timestep( + path, use_cftime=use_cftime, pattern=pattern, headers=None + ) + else: + raise ValueError( + "Unsupported chunk_layout value: " + f"{chunk_layout}, supported are 'per_idf' or 'per_timestep'" + ) + + +# Open IDFs for multiple times and/or layers into one DataArray +def _open( + path: PathType | list[PathType], + use_cftime=False, + pattern: PatternType | None = None, +) -> xr.DataArray: + """ + Open an IDF file or list of IDF files and return as an xarray DataArray. + + Parameters + ---------- + path : PathType | list[PathType] + Path or list of paths to IDF files. + use_cftime : bool, optional + Whether to use cftime for datetime objects. Default is False. + pattern : PatternType | None, optional + Pattern to parse the filename into xarray coordinates. Default is None. + + Returns + ------- + xr.DataArray + The opened IDF data as an xarray DataArray. """ return array_io.reading._open(path, use_cftime, pattern, header, _read) @@ -223,9 +290,9 @@ def _more_than_one_unique_value(values: Iterable[Any]): def _merge_subdomains( - paths_per_subdomain: DefaultDict[Any, list[str]], + paths_per_subdomain: DefaultDict[Any, list[PathType]], use_cftime: bool, - pattern: str | Pattern, + pattern: PatternType | None, ): """ Open and spatially merge all subdomain IDF files for one timestep. @@ -236,7 +303,7 @@ def _merge_subdomains( """ das = [] for pathlist in paths_per_subdomain.values(): - da = open(pathlist, use_cftime=use_cftime, pattern=pattern) + da = _open(pathlist, use_cftime=use_cftime, pattern=pattern) if "subdomain" in da.dims: da = da.isel(subdomain=0, drop=True) das.append(da) @@ -244,19 +311,24 @@ def _merge_subdomains( def _merge_subdomains_values( - paths_per_subdomain: DefaultDict[Any, list[str]], + paths_per_subdomain: DefaultDict[Any, list[PathType]], use_cftime: bool, - pattern: str | Pattern, -): + pattern: PatternType, +) -> np.ndarray: """Wraps ``_merge_subdomains`` to return just a numpy array for ``dask.array.from_delayed``.""" data, _, _, _ = _merge_subdomains(paths_per_subdomain, use_cftime, pattern) + # ``from_delayed`` expects each task to resolve to a concrete ndarray chunk. + # Returning a dask array here creates nested lazy arrays, which can leak + # ``dask.array.Array`` values into downstream numba-compiled code paths. + if isinstance(data, dask.array.Array): + data = data.compute() return data def _merge_subdomains_to_dataarray( - paths_per_subdomain: DefaultDict[Any, list[str]], + paths_per_subdomain: DefaultDict[Any, list[PathType]], use_cftime: bool, - pattern: str | Pattern, + pattern: PatternType, ) -> xr.DataArray: """Wraps ``_merge_subdomains`` to return a DataArray for coordinate template.""" data, coords, dims, name = _merge_subdomains( @@ -271,7 +343,7 @@ def _merge_subdomains_to_dataarray( def check_subdomain_consistency( - parsed: list[dict[str, Any]], paths: list[str], pattern: str | Pattern + parsed: list[dict[str, Any]], paths: Sequence[PathType], pattern: PatternType ): """Check that each subdomain has the same number of IDF files.""" grouped = defaultdict(list) @@ -291,16 +363,201 @@ def check_subdomain_consistency( ) +def _open_idf_data_from_group( + group: tuple[list[str], list[dict[str, Any]]], + use_cftime: bool, +) -> np.ndarray: + """Open IDF files from a group and return as numpy array.""" + group_paths, group_headers = group + # load_as_components mutates header dictionaries while preparing lazy reads. + # Copy to keep repeated/template reads safe. + headers = [h.copy() for h in group_headers] + data, _, _, _ = array_io.reading.load_as_components( + group_paths, use_cftime=use_cftime, _read=_read, headers=headers + ) + # ``from_delayed`` expects each task to resolve to a concrete ndarray chunk. + # Returning a dask array here creates nested lazy arrays, which can leak + # ``dask.array.Array`` values into downstream numba-compiled code paths. + if isinstance(data, dask.array.Array): + data = data.compute() + return data + + +def _open_idf_from_group( + group: tuple[list[PathType], list[dict[str, Any]]], + use_cftime: bool, +) -> xr.DataArray: + """Open IDF files and return a DataArray template with header-derived dims/coords.""" + # load_as_components mutates header dictionaries while preparing lazy reads. + group_paths, group_headers = group + group_headers = [h.copy() for h in group_headers] + data, coords, dims, name = array_io.reading.load_as_components( + group_paths, use_cftime=use_cftime, _read=_read, headers=group_headers + ) + return xr.DataArray(data, coords, dims, name=name) + + +def _open_idf_chunked_per_timestep( + open_func: Callable[..., Any], + grouped_by_time: dict[Any, Any], + template: xr.DataArray, + has_time: bool, + use_cftime: bool, + *args, +): + """ + Open IDF files chunked per timestep using a template DataArray. + + Parameters + ---------- + open_func : Callable + Function to open individual IDF files. Needs to have group, use_cftime, + pattern as arguments. Needs to return dask array as data. + grouped_by_time : dict[Any, Any] + Dictionary grouping file paths by time. + template : xr.DataArray + Template DataArray to define the shape, dims, and coords. + has_time : bool + Whether the data has a time dimension. + use_cftime : bool + Whether to use cftime for time coordinates. + *args : tuple + Additional arguments to pass to the open_func. + + Returns + ------- + xr.DataArray + Combined DataArray with data from all subdomains, chunked per timestep. + + + """ + shape = template.shape # e.g. (1, nlayer, nrow, ncol) + dims = template.dims # e.g. ("time", "layer", "y", "x") + dtype = template.dtype + if has_time: + time_axis = list(dims).index("time") + else: + time_axis = -1 # steady-state, no time dimension + + # Sort and convert times before calling _merge_subdomains so that + # use_cftime is already correct when the template is built. + raw_times_sorted = sorted(grouped_by_time.keys()) + converted_times, use_cftime = imod.util.time._convert_datetimes( + raw_times_sorted, use_cftime + ) + + # Delayed tasks for each timestep, which will be concatenated into a single + # dask array. One delayed task per timestep → outer graph depth 3, O(n_time) + # tasks + dask_arrays = [] + for time_key in raw_times_sorted: + group = grouped_by_time[time_key] + timestep_data = dask.delayed(open_func)(group, use_cftime, *args) + dask_arrays.append( + dask.array.from_delayed(timestep_data, shape=shape, dtype=dtype) + ) + data = dask.array.concatenate(dask_arrays, axis=time_axis) + + # Build the full time coordinate + coords = dict(template.coords) + if has_time: + time_coord: xr.CFTimeIndex | np.ndarray + if use_cftime: + time_coord = xr.CFTimeIndex(converted_times) + else: + time_coord = np.array(converted_times, dtype="datetime64[ns]") + coords["time"] = time_coord + + return xr.DataArray(data, coords, dims, name=template.name, attrs=template.attrs) + + +def open_chunked_per_timestep( + path: PathType | list[PathType], + use_cftime: bool = False, + pattern: PatternType | None = None, + headers: list[dict[str, Any]] | None = None, +): + """ + Open IDF files chunked per timestep. + + Parameters + ---------- + path : str or Path or list of str or Path + Global path(s). + use_cftime : bool, optional + Whether to use cftime for time coordinates. + pattern : str, regex pattern, optional + If the filenames do match default naming conventions of + {name}_{time}_l{layer}, a custom pattern can be defined here either + as a string, or as a compiled regular expression pattern. See the + examples below. + headers: list, optional + List of headers. Can be provided to avoid re-reading headers from the + files. + + Returns + ------- + xarray.DataArray + + """ + paths = array_io.reading.handle_path(path) + if headers is None: + headers = [header(p, pattern) for p in paths] + + # If headers were supplied by the caller, they may already contain + # dimensions like time/layer that are absent in the filename itself. + # Prefer header metadata for grouping and template construction. + has_time = "time" in headers[0].get("dims", []) + + grouped_paths_by_time: DefaultDict[Any, list[PathType]] = defaultdict(list) + grouped_headers_by_time: DefaultDict[Any, list[dict[str, Any]]] = defaultdict(list) + for p, h in zip(paths, headers): + if has_time: + time_key = h["time"] + else: + # Work around for files without time dimension + # (imod.util.time._convert_datetimes special-cases this string) + time_key = "steady-state" + grouped_paths_by_time[time_key].append(p) + grouped_headers_by_time[time_key].append(h) + + raw_times_sorted = sorted(grouped_paths_by_time.keys()) + + grouped_by_time: dict[ + Any, + tuple[list[PathType], list[dict[str, Any]]], + ] = { + key: (grouped_paths_by_time[key], grouped_headers_by_time[key]) + for key in raw_times_sorted + } + + first_time_key = raw_times_sorted[0] + template = _open_idf_from_group( + grouped_by_time[first_time_key], + use_cftime=use_cftime, + ) + + return _open_idf_chunked_per_timestep( + _open_idf_data_from_group, + grouped_by_time, + template, + has_time, + use_cftime, + ) + + def open_subdomains( - path: str | Path, use_cftime: bool = False, pattern: str | Pattern = None + path: PathType | list[PathType], + use_cftime: bool = False, + pattern: PatternType | None = None, ) -> xr.DataArray: """ Combine IDF files of multiple subdomains. Parameters ---------- - path : str or Path - Global path. + path : PathType or list of PathType + Global path or list of paths to open. use_cftime : bool, optional pattern : str, regex pattern, optional If no pattern is provided, the function will first try: @@ -318,14 +575,13 @@ def open_subdomains( # merges the subdomains into one DataArray, in a delayed manner, chunked per # timestep. A lot of logic in this function is about grouping the files by # subdomain and time, and setting the right time coordinate again. - - paths = sorted(glob.glob(str(path))) + paths = array_io.reading.handle_path(path) if pattern is None: # If no pattern provided test if pattern = "{name}_c{species}_{time}_l{layer}_p{subdomain}" re_pattern_species = imod.util.path._custom_pattern_to_regex_pattern(pattern) - has_species = re_pattern_species.search(paths[0]) + has_species = re_pattern_species.search(str(paths[0])) if not has_species: pattern = "{name}_{time}_l{layer}_p{subdomain}" @@ -336,7 +592,7 @@ def open_subdomains( # Group by time (datetime.datetime from decompose), then by subdomain. # Each delayed task processes one timestep, keeping the outer graph at O(n_time). - grouped_by_time: DefaultDict[Any, DefaultDict[Any, list]] = defaultdict( + grouped_by_time: DefaultDict[Any, DefaultDict[Any, list[PathType]]] = defaultdict( lambda: defaultdict(list) ) @@ -352,9 +608,6 @@ def open_subdomains( # Sort and convert times before calling _merge_subdomains so that # use_cftime is already correct when the template is built. raw_times_sorted = sorted(grouped_by_time.keys()) - converted_times, use_cftime = imod.util.time._convert_datetimes( - raw_times_sorted, use_cftime - ) # Call _merge_subdomains eagerly for the first timestep to obtain a # coordinate template. No data is computed — only the coordinate arrays @@ -363,40 +616,24 @@ def open_subdomains( template = _merge_subdomains_to_dataarray( grouped_by_time[first_time_key], use_cftime, pattern ) - - shape = template.shape # e.g. (1, nlayer, nrow, ncol) - dims = template.dims # e.g. ("time", "layer", "y", "x") - dtype = template.dtype - if has_time: - time_axis = list(dims).index("time") - else: - time_axis = -1 # steady-state, no time dimension - - # Delayed tasks for each timestep, which will be concatenated into a single - # dask array. One delayed task per timestep → outer graph depth 3, O(n_time) - # tasks - merged = [] - for time_key in raw_times_sorted: - group = grouped_by_time[time_key] - timestep_data = dask.delayed(_merge_subdomains_values)( - group, use_cftime, pattern - ) - merged.append(dask.array.from_delayed(timestep_data, shape=shape, dtype=dtype)) - data = dask.array.concatenate(merged, axis=time_axis) - - # Build the full time coordinate - coords = dict(template.coords) - if has_time: - if use_cftime: - time_coord = xr.CFTimeIndex(converted_times) - else: - time_coord = np.array(converted_times, dtype="datetime64[ns]") - coords["time"] = time_coord - - return xr.DataArray(data, coords, dims, name=template.name, attrs=template.attrs) + # Collect additional arguments to pass to the open function. + args_for_func = (pattern,) + + return _open_idf_chunked_per_timestep( + _merge_subdomains_values, + grouped_by_time, + template, + has_time, + use_cftime, + *args_for_func, + ) -def open_dataset(globpath, use_cftime=False, pattern=None): +def open_dataset( + globpath: str | pathlib.Path, + use_cftime: bool = False, + pattern: PatternType | None = None, +): """ Open a set of IDFs to a dict of xarray.DataArrays. @@ -445,7 +682,7 @@ def open_dataset(globpath, use_cftime=False, pattern=None): # group the DataArrays together using their name # note that directory names are ignored, and in case of duplicates, the last one wins names = [imod.util.path.decompose(path, pattern)["name"] for path in paths] - paths_by_name = {name: [] for name in np.unique(names)} + paths_by_name: dict[str, list[PathType]] = {name: [] for name in np.unique(names)} for path, name in zip(paths, names, strict=True): paths_by_name[name].append(path) # load each group into a DataArray @@ -470,7 +707,12 @@ def open_dataset(globpath, use_cftime=False, pattern=None): return dataset_dict -def write(path, a, nodata=1.0e20, dtype=np.float32): +def write( + path: PathType, + a: xr.DataArray, + nodata: float = 1.0e20, + dtype: DTypeLike = np.float32, +): """ Write a 2D xarray.DataArray to a IDF file @@ -580,7 +822,7 @@ def write(path, a, nodata=1.0e20, dtype=np.float32): a.values.tofile(f) -def _as_voxeldata(a): +def _as_voxeldata(a: xr.DataArray): """ If "z" is present as a dimension, generate layer if necessary. Ensure that layer is the dimension (via swap_dims). Infer "dz" if necessary, and if @@ -627,7 +869,13 @@ def _as_voxeldata(a): return a -def save(path, a, nodata=1.0e20, pattern=None, dtype=np.float32): +def save( + path: PathType, + a: xr.DataArray, + nodata: float = 1.0e20, + pattern: PatternType | None = None, + dtype: DTypeLike = np.float32, +): """ Write a xarray.DataArray to one or more IDF files diff --git a/imod/formats/prj/prj.py b/imod/formats/prj/prj.py index dabaf2478..6eda34e95 100644 --- a/imod/formats/prj/prj.py +++ b/imod/formats/prj/prj.py @@ -568,10 +568,10 @@ def _create_dataarray_from_paths( factor = _get_array_transformation_parameters(headers, "factor", dim) addition = _get_array_transformation_parameters(headers, "addition", dim) da = _try_read_with_func( - imod.formats.array_io.reading._load, + imod.formats.idf.open_chunked_per_timestep, paths, use_cftime=False, - _read=imod.idf._read, + pattern=None, headers=headers, ) diff --git a/imod/tests/test_formats/test_idf.py b/imod/tests/test_formats/test_idf.py index af19ad894..be8bddeac 100644 --- a/imod/tests/test_formats/test_idf.py +++ b/imod/tests/test_formats/test_idf.py @@ -1,6 +1,9 @@ import datetime +import dask +import dask.array import numpy as np +import pandas as pd import pytest import xarray as xr from pytest import approx @@ -299,6 +302,80 @@ def test_open_subdomains_error(subdomains, expected, equidistant, tmp_path): idf.open_subdomains(tmp_path / "subdomains_*.idf") +class TemporalCases: + def create_da(self, ntime): + nlayer, nrow, ncol = 3, 6, 8 + dx = 1.0 + dy = -1.0 + xmin, xmax = 0.0, 8.0 + ymin, ymax = 0.0, 6.0 + layer = [1, 2, 3] + + time = pd.date_range("2000-01-01", periods=ntime, freq="D") + + kwargs = {"name": "temporal_data", "dims": ("time", "layer", "y", "x")} + kwargs["coords"] = util.spatial._xycoords((xmin, xmax, ymin, ymax), (dx, dy)) + kwargs["coords"]["layer"] = layer + kwargs["coords"]["time"] = time + kwargs["data"] = np.ones((ntime, nlayer, nrow, ncol), dtype=np.float64) + return xr.DataArray(**kwargs).cumsum(dim="time") + + def case_single_time(self): + ntime = 1 + return self.create_da(ntime=ntime), ntime + + def case_multiple_times(self): + ntime = 5 + return self.create_da(ntime=ntime), ntime + + +@parametrize_with_cases("temporal_data,ntime", cases=TemporalCases) +@parametrize("pattern", [None, r"{name}_{time}_l{layer}"]) +def test_open_chunked_per_timestep(temporal_data, ntime, pattern, tmp_path): + idf.save(tmp_path / "temporal_data", temporal_data) + + # Test with pattern + da = idf.open_chunked_per_timestep( + tmp_path / "temporal_data_*.idf", pattern=pattern + ) + + assert da.dims == ("time", "layer", "y", "x") + assert da.name == "temporal_data" + assert da.sizes["layer"] == 3 + assert da.sizes["y"] == 6 + assert da.sizes["x"] == 8 + assert da.sizes["time"] == ntime + time_chunk_shape = (1,) * ntime + assert da.chunks == ( + time_chunk_shape, + (da.sizes["layer"],), + (da.sizes["y"],), + (da.sizes["x"],), + ) + + # Compute and see if no error is thrown + da = da.compute() + + np.testing.assert_allclose(da, temporal_data) + + +def test_open_chunked_per_timestep__delayed_chunk_returns_numpy(tmp_path): + """ + Test that the dask chunks returned by open_chunked_per_timestep is a numpy array, + and not a dask array, as that would mean there are nested dask arrays, which + are problematic. + """ + temporal_data = TemporalCases().create_da(ntime=2) + idf.save(tmp_path / "temporal_data", temporal_data) + + da = idf.open_chunked_per_timestep(tmp_path / "temporal_data_*.idf") + first_chunk_delayed = da.data.to_delayed().ravel()[0] + first_chunk = dask.compute(first_chunk_delayed)[0] + + assert isinstance(first_chunk, np.ndarray) + assert not isinstance(first_chunk, dask.array.Array) + + def test_xycoords_equidistant(): dx, dy = 1.0, -1.0 xmin, xmax = 0.0, 4.0 @@ -357,6 +434,17 @@ def test_save_topbot__single_layer(test_da, tmp_path): assert da_read["dz"] == approx(1.0) +@parametrize("chunk_layout", ["per_idf", "per_timestep"]) +def test_saveopen_layer_chunking(test_layerda, chunk_layout, tmp_path): + da = test_layerda + idf.save(tmp_path / "layer", da) + actual = idf.open(tmp_path / "layer_l*.idf", chunk_layout=chunk_layout) + if chunk_layout == "per_idf": + assert actual.chunks[0] == (1,) * actual.sizes["layer"] + elif chunk_layout == "per_timestep": + assert actual.chunks[0] == (actual.sizes["layer"],) + + def test_save_topbot__layers(test_layerda, tmp_path): da = test_layerda da = da.assign_coords(z=("layer", np.arange(1.0, 6.0) - 0.5)) @@ -554,3 +642,58 @@ def test_open_subdomains_large_scale(tmp_path): assert result.sizes["x"] == grid_cols * ncol assert result.values.dtype == np.float32 assert np.all(result.values == 1.0) + + +@pytest.mark.timeout(600, method="thread") # 10 minutes +@pytest.mark.user_acceptance +def test_open_chunked_per_timestep_large_scale(tmp_path): + """ + Test the performance and correctness of opening large-scale IDF files + chunked by time. + """ + n_layers = 40 + n_times = 800 + nrow, ncol = 20, 20 # cells per subdomain + dx, dy = 1.0, -1.0 + + assert n_layers * n_times == 32_000 + + # Generate 50 daily timesteps starting 2000-01-01 + start = datetime.date(2000, 1, 1) + date_strs = [ + (start + datetime.timedelta(days=i)).strftime("%Y%m%d") for i in range(n_times) + ] + # Create the DataArray for each timestep + data = np.ones((nrow, ncol), dtype=np.float32) + xmin, xmax = 0, ncol * dx + ymin, ymax = 0, nrow * abs(dy) + coords = util.spatial._xycoords((xmin, xmax, ymin, ymax), (dx, dy)) + da = xr.DataArray(data, dims=("y", "x"), coords=coords, name="head") + + idf_dir = tmp_path / "idf_files" + idf_dir.mkdir(exist_ok=True) + + # Write all 32,000 IDF files + for layer in range(1, n_layers + 1): + for date_str in date_strs: + idf.write( + idf_dir / f"head_{date_str}_l{layer}.idf", + da, + ) + + # Open subdomains and load into memory + t0 = datetime.datetime.now() + result = idf.open_chunked_per_timestep(idf_dir / "head_*.idf").compute() + elapsed = datetime.datetime.now() - t0 + with open(tmp_path / "open_chunked_per_timestep_large_scale.log", "w") as f: + f.write( + f"Elapsed time for open_chunked_per_timestep + load: {elapsed.total_seconds():.2f}s\n" + ) + + assert isinstance(result, xr.DataArray) + assert result.dims == ("time", "layer", "y", "x") + assert result.sizes["time"] == n_times + assert result.sizes["layer"] == n_layers + assert result.sizes["y"] == nrow + assert result.sizes["x"] == ncol + assert result.values.dtype == np.float32 diff --git a/imod/tests/test_formats/test_prj.py b/imod/tests/test_formats/test_prj.py index e98fe913c..1bca91c17 100644 --- a/imod/tests/test_formats/test_prj.py +++ b/imod/tests/test_formats/test_prj.py @@ -593,6 +593,25 @@ def test_open_projectfile_data(self, projectfile, request): assert isinstance(content["pcg"], dict) assert set(repeats["rch"]) == {datetime(1899, 4, 1), datetime(1899, 10, 1)} + # Test if chunking of topsystem IDF data went correctly, should be + # chunked per timestep and have only 1 layer. + assert content["ghb"]["conductance"].shape == (3, 1, 2, 2) + assert content["ghb"]["conductance"].chunksizes["time"] == (1, 1, 1) + assert content["ghb"]["conductance"].chunksizes["layer"] == (1,) + assert content["ghb"]["conductance"].chunksizes["x"] == (2,) + assert content["ghb"]["conductance"].chunksizes["y"] == (2,) + # Load the data to ensure that dask arrays are correctly formed and + # can be evaluated. + content["ghb"]["conductance"].load() + + # Test if chunking BND data is expected, should be one chunk including + # the two layers. + assert content["bnd"]["ibound"].shape == (2, 2, 2) + assert content["bnd"]["ibound"].chunksizes["layer"] == (2,) + assert content["bnd"]["ibound"].chunksizes["x"] == (2,) + assert content["bnd"]["ibound"].chunksizes["y"] == (2,) + content["bnd"]["ibound"].load() + def test_open_projectfile_data__faulty_well(self, projectfile): basepath = self.basepath # Setup faulty well diff --git a/imod/util/path.py b/imod/util/path.py index 23499b255..f6cacc6e4 100644 --- a/imod/util/path.py +++ b/imod/util/path.py @@ -8,6 +8,7 @@ import pathlib import re import tempfile +from re import Pattern from typing import Any, Optional import cftime @@ -76,7 +77,7 @@ def _groupdict(stem: str, pattern: Optional[str | re.Pattern[str]]) -> dict[str, return d -def decompose(path, pattern: Optional[str] = None) -> dict[str, Any]: +def decompose(path, pattern: str | Pattern[str] | None = None) -> dict[str, Any]: r""" Parse a path, returning a dict of the parts, following the iMOD conventions. diff --git a/imod/util/time.py b/imod/util/time.py index 8525fcb6c..75b52bccb 100644 --- a/imod/util/time.py +++ b/imod/util/time.py @@ -1,5 +1,6 @@ import datetime import warnings +from typing import Any import cftime import dateutil @@ -155,7 +156,7 @@ def forcing_starts_ends(package_times: np.ndarray, globaltimes: np.ndarray): return starts_ends -def _convert_datetimes(times: np.ndarray, use_cftime: bool): +def _convert_datetimes(times: list[Any], use_cftime: bool): """ Return times as np.datetime64[ns] or cftime.DatetimeProlepticGregorian depending on whether the dates fall within the inclusive bounds of diff --git a/pyproject.toml b/pyproject.toml index 366b87efe..f74a2b29f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -118,7 +118,11 @@ warn_return_any = false module = [ "imod.data.*", "imod.evaluate.*", - "imod.formats.*", + "imod.formats.array_io.*", + "imod.formats.gen.*", + "imod.formats.prj.*", + "imod.formats.ipf", + "imod.formats.rasterio", "imod.schemata.*", "imod.templates.*", "imod.tests.*",