diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index 55c787846..1cbe89a3c 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -89,6 +89,10 @@ Fixed ``front`` budget was left out, even though you only need to provide one of ``front``, ``lower`` or ``right``. Leaving out ``front`` now works as described in the documentation. +- Fixed big performance degradation with :func:`imod.idf.open_subdomains` where + it would take a long time to open lots of idf files. Performance is now + significantly improved up to the same speed as before the change that caused + the performance degradation. Changed ~~~~~~~ diff --git a/docs/developing/test-plan.rst b/docs/developing/test-plan.rst index f75c86aa7..eb916da3e 100644 --- a/docs/developing/test-plan.rst +++ b/docs/developing/test-plan.rst @@ -162,6 +162,13 @@ Criteria for user acceptance tests of the 1.0 release are: - Some small differences in how rivers and drains are allocated and their conductances are distributed. This leads to local head differences in the output of the simulations in the south of Limburg. +* There additionally is a user acceptance test to test performance of the + :func:`imod.formats.idf.open_subdomains` function. This function is used to + open a large number of IDF files and load them into memory. The test will + write 32,000 IDF files to disk, and then open them and load them into memory. + The time it takes to do this should be less than 10 minutes, otherwise it will + fail. (For reference: It takes about 3 minutes on a machine with 32 GB of RAM + and an Intel i7-12800HX CPU.) Manual checks ************* diff --git a/imod/formats/idf.py b/imod/formats/idf.py index 5096623d2..68eb1e1c8 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -12,14 +12,16 @@ from collections.abc import Iterable from pathlib import Path from re import Pattern -from typing import Any +from typing import Any, DefaultDict +import dask +import dask.array import numpy as np import xarray as xr import imod from imod.formats import array_io -from imod.typing.structured import merge_partitions +from imod.typing.structured import merge_partitions_as_da_components # Make sure we can still use the built-in function... f_open = open @@ -220,6 +222,75 @@ def _more_than_one_unique_value(values: Iterable[Any]): return len(set(values)) != 1 +def _merge_subdomains( + paths_per_subdomain: DefaultDict[Any, list[str]], + use_cftime: bool, + pattern: str | Pattern, +): + """ + Open and spatially merge all subdomain IDF files for one timestep. + Returns an ``xr.DataArray`` with all coordinates computed by + ``merge_partitions``. Called directly (eagerly) once to obtain a + coordinate template, and indirectly via ``_merge_subdomains_values`` + inside ``dask.delayed`` for actual computation. + """ + das = [] + for pathlist in paths_per_subdomain.values(): + da = open(pathlist, use_cftime=use_cftime, pattern=pattern) + if "subdomain" in da.dims: + da = da.isel(subdomain=0, drop=True) + das.append(da) + return merge_partitions_as_da_components(das) + + +def _merge_subdomains_values( + paths_per_subdomain: DefaultDict[Any, list[str]], + use_cftime: bool, + pattern: str | Pattern, +): + """Wraps ``_merge_subdomains`` to return just a numpy array for ``dask.array.from_delayed``.""" + data, _, _, _ = _merge_subdomains(paths_per_subdomain, use_cftime, pattern) + return data + + +def _merge_subdomains_to_dataarray( + paths_per_subdomain: DefaultDict[Any, list[str]], + use_cftime: bool, + pattern: str | Pattern, +) -> xr.DataArray: + """Wraps ``_merge_subdomains`` to return a DataArray for coordinate template.""" + data, coords, dims, name = _merge_subdomains( + paths_per_subdomain, use_cftime, pattern + ) + return xr.DataArray( + data=data, + coords=coords, + dims=dims, + name=name, + ) + + +def check_subdomain_consistency( + parsed: list[dict[str, Any]], paths: list[str], pattern: str | Pattern +): + """Check that each subdomain has the same number of IDF files.""" + grouped = defaultdict(list) + for match, p in zip(parsed, paths): + try: + key = match["subdomain"] + except KeyError as e: + raise KeyError(f"{e} in path: {p} with pattern: {pattern}") + grouped[key].append(p) + + n_idf_per_subdomain = { + subdomain_id: len(path_ls) for subdomain_id, path_ls in grouped.items() + } + if _more_than_one_unique_value(n_idf_per_subdomain.values()): + raise ValueError( + f"Each subdomain must have the same number of IDF files, found: {n_idf_per_subdomain}" + ) + + def open_subdomains( path: str | Path, use_cftime: bool = False, pattern: str | Pattern = None ) -> xr.DataArray: @@ -243,6 +314,11 @@ def open_subdomains( xarray.DataArray """ + # This function is a wrapper around open() that groups by subdomain and + # 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))) if pattern is None: @@ -254,30 +330,70 @@ def open_subdomains( pattern = "{name}_{time}_l{layer}_p{subdomain}" parsed = [imod.util.path.decompose(path, pattern) for path in paths] - grouped = defaultdict(list) - for match, path in zip(parsed, paths): - try: - key = match["subdomain"] - except KeyError as e: - raise KeyError(f"{e} in path: {path} with pattern: {pattern}") - grouped[key].append(path) + check_subdomain_consistency(parsed, paths, pattern) - n_idf_per_subdomain = { - subdomain_id: len(path_ls) for subdomain_id, path_ls in grouped.items() - } - if _more_than_one_unique_value(n_idf_per_subdomain.values()): - raise ValueError( - f"Each subdomain must have the same number of IDF files, found: {n_idf_per_subdomain}" - ) + has_time = "time" in parsed[0] - das = [] - for pathlist in grouped.values(): - da = open(pathlist, use_cftime=use_cftime, pattern=pattern) - da = da.isel(subdomain=0, drop=True) - das.append(da) + # 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( + lambda: defaultdict(list) + ) + + for match, p in zip(parsed, paths): + if has_time: + time_key = match["time"] + else: + # Work around for files without time dimension + # (imod.util.time._convert_datetimes special-cases this string) + time_key = "steady-state" + grouped_by_time[time_key][match["subdomain"]].append(p) + + # 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 + # (which are numpy) are used; the dask data array is discarded. + first_time_key = raw_times_sorted[0] + 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 - name = das[0].name - return merge_partitions(das)[name] # as DataArray for backwards compatibility + return xr.DataArray(data, coords, dims, name=template.name, attrs=template.attrs) def open_dataset(globpath, use_cftime=False, pattern=None): diff --git a/imod/tests/test_formats/test_idf.py b/imod/tests/test_formats/test_idf.py index 9a39c0276..af19ad894 100644 --- a/imod/tests/test_formats/test_idf.py +++ b/imod/tests/test_formats/test_idf.py @@ -1,3 +1,5 @@ +import datetime + import numpy as np import pytest import xarray as xr @@ -156,6 +158,7 @@ def test_open_subdomains(subdomains, expected, equidistant, tmp_path): expected_coords = util.spatial._xycoords((0.0, 8.0, 0.0, 6.0), (dx, dy)) assert da.dims == ("time", "layer", "y", "x") + assert da.name == "subdomains" assert np.all(da.isel(time=0) == expected) assert len(da.x) == 8 @@ -186,6 +189,7 @@ def test_open_subdomains__start_nr_not_zero( expected_coords = util.spatial._xycoords((0.0, 8.0, 0.0, 6.0), (dx, dy)) assert da.dims == ("time", "layer", "y", "x") + assert da.name == "subdomains" assert np.all(da.isel(time=0) == expected) assert len(da.x) == 8 @@ -211,10 +215,31 @@ def test_open_subdomains_pattern_None(subdomains, expected, equidistant, tmp_pat da = idf.open_subdomains(tmp_path / "subdomains_*.idf").load() assert da.dims == ("time", "layer", "y", "x") + assert da.name == "subdomains" assert np.all(da.isel(time=0) == expected) +@parametrize_with_cases( + "subdomains,expected,equidistant", cases=SubdomainCases, has_tag="no_species" +) +def test_open_subdomains_pattern_no_time(subdomains, expected, equidistant, tmp_path): + """ + Read with pattern without {time}, the time in the idf filename should then + be included in the name of the DataArray + """ + _save_subdomains_no_species(subdomains, tmp_path) + # Test with pattern + pattern = r"{name}_l{layer}_p{subdomain}" + da = idf.open_subdomains(tmp_path / "subdomains_*.idf", pattern=pattern).load() + + assert da.dims == ("layer", "y", "x") + assert da.name == "subdomains_20000101" + + assert "time" not in da.coords.keys() + assert np.all(da == expected) + + @parametrize_with_cases( "subdomains,expected,equidistant", cases=SubdomainCases, has_tag="species" ) @@ -230,6 +255,7 @@ def test_open_subdomains_species(subdomains, expected, equidistant, tmp_path): expected_coords = util.spatial._xycoords((0.0, 8.0, 0.0, 6.0), (dx, dy)) assert da.dims == ("species", "time", "layer", "y", "x") + assert da.name == "subdomains" assert np.all(da.isel(time=0) == expected) assert len(da.x) == 8 @@ -457,3 +483,74 @@ def test_open_dataset(test_layerda, tmp_path): for name, da_result in result.items(): assert isinstance(da_result, xr.DataArray) assert da_result.dims == ("layer", "y", "x") + + +@pytest.mark.timeout(600, method="thread") # 10 minutes +@pytest.mark.user_acceptance +def test_open_subdomains_large_scale(tmp_path): + """ + User acceptance test: write 32,000 IDF files for 16 subdomains, 40 layers, + and 50 timesteps, then open and load with imod.idf.open_subdomains. + Target: completes within 10 minutes (when setting up this test, it took 4 + minutes on my laptop). + """ + n_subdomains = 16 # 4 x 4 grid + n_layers = 40 + n_times = 50 + nrow, ncol = 5, 5 # cells per subdomain + dx, dy = 1.0, -1.0 + grid_cols = 4 + grid_rows = 4 # grid_cols * grid_rows == n_subdomains + + assert n_subdomains * 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) + ] + + # Precompute one DataArray per subdomain (data is the same for all + # timesteps and layers; only the spatial extent differs per subdomain). + data = np.ones((nrow, ncol), dtype=np.float32) + subdomain_das = [] + for s in range(n_subdomains): + col_idx = s % grid_cols + row_idx = s // grid_cols + xmin = col_idx * ncol * dx + xmax = xmin + ncol * dx + ymax = (grid_rows - row_idx) * nrow * abs(dy) + ymin = ymax - nrow * abs(dy) + coords = util.spatial._xycoords((xmin, xmax, ymin, ymax), (dx, dy)) + da = xr.DataArray(data, dims=("y", "x"), coords=coords, name="head") + subdomain_das.append(da) + + idf_dir = tmp_path / "idf_files" + idf_dir.mkdir(exist_ok=True) + + # Write all 32,000 IDF files + for s, da in enumerate(subdomain_das): + for layer in range(1, n_layers + 1): + for date_str in date_strs: + idf.write( + idf_dir / f"head_{date_str}_l{layer}_p{s:04d}.idf", + da, + ) + + # Open subdomains and load into memory + t0 = datetime.datetime.now() + result = idf.open_subdomains(idf_dir / "head_*.idf").load() + elapsed = datetime.datetime.now() - t0 + with open(tmp_path / "open_subdomains_large_scale.log", "w") as f: + f.write( + f"Elapsed time for open_subdomains + 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"] == grid_rows * nrow + assert result.sizes["x"] == grid_cols * ncol + assert result.values.dtype == np.float32 + assert np.all(result.values == 1.0) diff --git a/imod/typing/structured.py b/imod/typing/structured.py index e429a14a8..363c9d8ad 100644 --- a/imod/typing/structured.py +++ b/imod/typing/structured.py @@ -2,7 +2,7 @@ import itertools from collections import defaultdict -from typing import DefaultDict, List, Set, Tuple +from typing import Any, DefaultDict, List, Set, Tuple, cast import dask import numpy as np @@ -155,7 +155,29 @@ def _merge_nonequidistant_coords( return out -def _merge_partitions(das: List[xr.DataArray]) -> xr.DataArray: +def merge_partitions_as_da_components( + das: List[xr.DataArray], +) -> tuple[np.ndarray, dict[str, Any], tuple[str, ...], str]: + """ + Merge a list of xarray DataArrays into components for a single DataArray. + + The returned components are not yet combined into a DataArray, to avoid + excessive overhead for imod.idf._merge_subdomains() + + Parameters + ---------- + das: list of xr.DataArray + The list of DataArrays to merge. + + Returns + ------- + data: np.ndarray + The merged data array. + coords: dict[str, Any] + The merged coordinates. + dims: tuple[str, ...] + The dimension names of the merged DataArray. + """ # Do some input checking check_dtypes(das) check_dims(das) @@ -175,7 +197,7 @@ def _merge_partitions(das: List[xr.DataArray]) -> xr.DataArray: # Collect coordinates first = das[0] - coords = dict(first.coords) + coords = cast(dict[str, Any], dict(first.coords)) coords["x"] = x coords["y"] = y[::-1] if _is_nonequidistant_coord(first, "dx"): @@ -183,7 +205,10 @@ def _merge_partitions(das: List[xr.DataArray]) -> xr.DataArray: if _is_nonequidistant_coord(first, "dy"): coords["dy"] = ("y", _merge_nonequidistant_coords(das, "dy", iys, nrow)) + dims = cast(tuple[str, ...], first.dims) + name = cast(str, first.name) arrays = [da.data for da in das] + data: np.ndarray if first.chunks is None: # If the data is in memory, merge all at once. data = merge_arrays(arrays, ixs, iys, yx_shape) @@ -228,11 +253,19 @@ def _merge_partitions(das: List[xr.DataArray]) -> xr.DataArray: # After merging, the xy chunks are always (1, 1) reshaped = merged_blocks.reshape(block_shape + (1, 1)) data = dask.array.block(reshaped.tolist()) + return data, coords, dims, name + +def _merge_partitions_da(das: List[xr.DataArray]) -> xr.DataArray: + """ + Merge a list of xarray DataArrays into a single DataArray. + """ + data, coords, dims, name = merge_partitions_as_da_components(das) return xr.DataArray( data=data, coords=coords, - dims=first.dims, + dims=dims, + name=name, ) @@ -244,12 +277,12 @@ def merge_partitions( unique_keys = {key for da in das for key in da.keys()} merged_ls = [] for key in unique_keys: - merged_ls.append(_merge_partitions([da[key] for da in das]).rename(key)) + merged_ls.append(_merge_partitions_da([da[key] for da in das]).rename(key)) return xr.merge(merged_ls, compat="no_conflicts") elif isinstance(first_item, xr.DataArray): # Store name to rename after concatenation name = first_item.name - return _merge_partitions(das).to_dataset(name=name) # type: ignore + return _merge_partitions_da(das).to_dataset(name=name) # type: ignore else: raise TypeError( f"Expected type: xr.DataArray or xr.Dataset, got {type(first_item)}"