From 4fd124491dab5c80c90a5e3e5049f99b6c696f40 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Thu, 23 Jul 2026 15:18:55 +0200 Subject: [PATCH 01/13] Let Claude merge old implemenation into new one, while calling merge_partitions --- imod/formats/idf.py | 120 +++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 113 insertions(+), 7 deletions(-) diff --git a/imod/formats/idf.py b/imod/formats/idf.py index f14fdb604..debd8dd26 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -6,6 +6,7 @@ """ import glob +import itertools import pathlib import struct from collections import defaultdict @@ -14,6 +15,8 @@ from re import Pattern from typing import Any +import dask +import dask.array import numpy as np import xarray as xr @@ -220,6 +223,23 @@ def _more_than_one_unique_value(values: Iterable[Any]): return len(set(values)) != 1 +def _merge_subdomains(paths_per_subdomain, use_cftime, pattern): + """ + Open and spatially merge all subdomain IDF files for a single timestep. + Intended to be called inside ``dask.delayed``. The inner dask graphs + created by ``open()`` are computed synchronously within the task, so the + outer graph has only O(n_time) nodes — matching the original performance. + """ + 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) + name = das[0].name + return merge_partitions(das)[name].values + + def open_subdomains( path: str | Path, use_cftime: bool = False, pattern: str | Pattern = None ) -> xr.DataArray: @@ -270,14 +290,100 @@ def open_subdomains( f"Each subdomain must have the same number of IDF files, found: {n_idf_per_subdomain}" ) - 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 = defaultdict(lambda: defaultdict(list)) + all_layers = [] + all_species = [] + for match, p in zip(parsed, paths): + grouped_by_time[match["time"]][match["subdomain"]].append(p) + all_layers.append(match["layer"]) + if "species" in match: + all_species.append(match["species"]) + + # Read headers from the first timestep to determine the output shape and + # coordinates without loading any data. + first_time_key = next(iter(grouped_by_time)) + first_group = grouped_by_time[first_time_key] + unique_subdomains = sorted(first_group.keys()) + + all_x_coords = [] + all_y_coords = [] + for subdomain in unique_subdomains: + hdr = header(first_group[subdomain][0], pattern) + bounds = (hdr["xmin"], hdr["xmax"], hdr["ymin"], hdr["ymax"]) + sub_coords = imod.util.spatial._xycoords(bounds, (hdr["dx"], hdr["dy"])) + all_x_coords.append(sub_coords["x"]) + all_y_coords.append(sub_coords["y"]) + + global_x = np.unique(np.concatenate(all_x_coords)) + global_y = np.unique(np.concatenate(all_y_coords))[::-1] # descending + nrow = global_y.size + ncol = global_x.size + global_layers = np.array(sorted(set(all_layers))) + nlayer = global_layers.size + global_species = np.array(sorted(set(all_species))) if all_species else None + + first_hdr = header(first_group[unique_subdomains[0]][0], pattern) + dtype = first_hdr["dtype"] + + if global_species is not None: + nspecies = global_species.size + shape = (nspecies, 1, nlayer, nrow, ncol) + dims = ("species", "time", "layer", "y", "x") + time_axis = 1 + else: + shape = (1, nlayer, nrow, ncol) + dims = ("time", "layer", "y", "x") + time_axis = 0 + + # Sort times; datetime.datetime objects are always comparable + raw_times_sorted = sorted(grouped_by_time.keys()) + converted_times, use_cftime = imod.util.time._convert_datetimes( + raw_times_sorted, use_cftime + ) + + # 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)(group, use_cftime, pattern) + merged.append(dask.array.from_delayed(timestep_data, shape=shape, dtype=dtype)) + data = dask.array.concatenate(merged, axis=time_axis) + + if use_cftime: + time_coord = xr.CFTimeIndex(converted_times) + else: + time_coord = np.array(converted_times, dtype="datetime64[ns]") - name = das[0].name - return merge_partitions(das)[name] # as DataArray for backwards compatibility + coords = { + "y": global_y, + "x": global_x, + "layer": global_layers, + "time": time_coord, + } + if global_species is not None: + coords["species"] = global_species + + # Add z/dz coordinates if top/bottom data is present in the IDF headers + sample_paths = first_group[unique_subdomains[0]] + hdrs = [header(p, pattern) for p in sample_paths] + tops = [h.get("top") for h in hdrs] + bots = [h.get("bot") for h in hdrs] + layer_nums = [h.get("layer") for h in hdrs] + _, unique_indices = np.unique(layer_nums, return_index=True) + all_have_z = all(v is not None for v in itertools.chain(tops, bots)) + if all_have_z: + if nlayer > 1: + coords = array_io.reading._array_z_coord(coords, tops, bots, unique_indices) + else: + coords = array_io.reading._scalar_z_coord(coords, tops, bots) + + name = parsed[0]["name"] + attrs = {} + if "nodata" in first_hdr: + attrs["nodata"] = first_hdr["nodata"] + return xr.DataArray(data, coords, dims, name=name, attrs=attrs) def open_dataset(globpath, use_cftime=False, pattern=None): From 32097eb77f40809b589359185b054e39a0522940 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Thu, 23 Jul 2026 16:40:20 +0200 Subject: [PATCH 02/13] Avoid duplicating coordinate construction from headers by eagerly calling merge_subdomains for one timestep. --- imod/formats/idf.py | 111 +++++++++++++------------------------------- 1 file changed, 32 insertions(+), 79 deletions(-) diff --git a/imod/formats/idf.py b/imod/formats/idf.py index debd8dd26..558b2f980 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -6,7 +6,6 @@ """ import glob -import itertools import pathlib import struct from collections import defaultdict @@ -225,10 +224,11 @@ def _more_than_one_unique_value(values: Iterable[Any]): def _merge_subdomains(paths_per_subdomain, use_cftime, pattern): """ - Open and spatially merge all subdomain IDF files for a single timestep. - Intended to be called inside ``dask.delayed``. The inner dask graphs - created by ``open()`` are computed synchronously within the task, so the - outer graph has only O(n_time) nodes — matching the original performance. + 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(): @@ -237,7 +237,12 @@ def _merge_subdomains(paths_per_subdomain, use_cftime, pattern): da = da.isel(subdomain=0, drop=True) das.append(da) name = das[0].name - return merge_partitions(das)[name].values + return merge_partitions(das)[name] + + +def _merge_subdomains_values(paths_per_subdomain, use_cftime, pattern): + """Wraps ``_merge_subdomains`` to return a numpy array for ``dask.array.from_delayed``.""" + return _merge_subdomains(paths_per_subdomain, use_cftime, pattern).values def open_subdomains( @@ -275,12 +280,12 @@ def open_subdomains( parsed = [imod.util.path.decompose(path, pattern) for path in paths] grouped = defaultdict(list) - for match, path in zip(parsed, paths): + for match, p 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) + 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() @@ -293,97 +298,45 @@ 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 = defaultdict(lambda: defaultdict(list)) - all_layers = [] - all_species = [] for match, p in zip(parsed, paths): grouped_by_time[match["time"]][match["subdomain"]].append(p) - all_layers.append(match["layer"]) - if "species" in match: - all_species.append(match["species"]) - - # Read headers from the first timestep to determine the output shape and - # coordinates without loading any data. - first_time_key = next(iter(grouped_by_time)) - first_group = grouped_by_time[first_time_key] - unique_subdomains = sorted(first_group.keys()) - - all_x_coords = [] - all_y_coords = [] - for subdomain in unique_subdomains: - hdr = header(first_group[subdomain][0], pattern) - bounds = (hdr["xmin"], hdr["xmax"], hdr["ymin"], hdr["ymax"]) - sub_coords = imod.util.spatial._xycoords(bounds, (hdr["dx"], hdr["dy"])) - all_x_coords.append(sub_coords["x"]) - all_y_coords.append(sub_coords["y"]) - - global_x = np.unique(np.concatenate(all_x_coords)) - global_y = np.unique(np.concatenate(all_y_coords))[::-1] # descending - nrow = global_y.size - ncol = global_x.size - global_layers = np.array(sorted(set(all_layers))) - nlayer = global_layers.size - global_species = np.array(sorted(set(all_species))) if all_species else None - - first_hdr = header(first_group[unique_subdomains[0]][0], pattern) - dtype = first_hdr["dtype"] - - if global_species is not None: - nspecies = global_species.size - shape = (nspecies, 1, nlayer, nrow, ncol) - dims = ("species", "time", "layer", "y", "x") - time_axis = 1 - else: - shape = (1, nlayer, nrow, ncol) - dims = ("time", "layer", "y", "x") - time_axis = 0 - # Sort times; datetime.datetime objects are always comparable + # 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(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 + time_axis = list(dims).index("time") + # 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)(group, use_cftime, pattern) + 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 and replace the single-timestep one from the template. if use_cftime: time_coord = xr.CFTimeIndex(converted_times) else: time_coord = np.array(converted_times, dtype="datetime64[ns]") - coords = { - "y": global_y, - "x": global_x, - "layer": global_layers, - "time": time_coord, - } - if global_species is not None: - coords["species"] = global_species - - # Add z/dz coordinates if top/bottom data is present in the IDF headers - sample_paths = first_group[unique_subdomains[0]] - hdrs = [header(p, pattern) for p in sample_paths] - tops = [h.get("top") for h in hdrs] - bots = [h.get("bot") for h in hdrs] - layer_nums = [h.get("layer") for h in hdrs] - _, unique_indices = np.unique(layer_nums, return_index=True) - all_have_z = all(v is not None for v in itertools.chain(tops, bots)) - if all_have_z: - if nlayer > 1: - coords = array_io.reading._array_z_coord(coords, tops, bots, unique_indices) - else: - coords = array_io.reading._scalar_z_coord(coords, tops, bots) + coords = dict(template.coords) + coords["time"] = time_coord - name = parsed[0]["name"] - attrs = {} - if "nodata" in first_hdr: - attrs["nodata"] = first_hdr["nodata"] - return xr.DataArray(data, coords, dims, name=name, attrs=attrs) + return xr.DataArray(data, coords, dims, name=template.name, attrs=template.attrs) def open_dataset(globpath, use_cftime=False, pattern=None): From 58e92584be17068fbaca087072f0959166303662 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Thu, 23 Jul 2026 17:11:17 +0200 Subject: [PATCH 03/13] Attempt at type annotation --- imod/formats/idf.py | 58 +++++++++++++++++++++++++++++---------------- 1 file changed, 38 insertions(+), 20 deletions(-) diff --git a/imod/formats/idf.py b/imod/formats/idf.py index 558b2f980..78b61c0f2 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -12,7 +12,7 @@ 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 @@ -222,7 +222,11 @@ def _more_than_one_unique_value(values: Iterable[Any]): return len(set(values)) != 1 -def _merge_subdomains(paths_per_subdomain, use_cftime, pattern): +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 @@ -240,11 +244,35 @@ def _merge_subdomains(paths_per_subdomain, use_cftime, pattern): return merge_partitions(das)[name] -def _merge_subdomains_values(paths_per_subdomain, use_cftime, pattern): +def _merge_subdomains_values( + paths_per_subdomain: DefaultDict[Any, list[str]], + use_cftime: bool, + pattern: str | Pattern, +): """Wraps ``_merge_subdomains`` to return a numpy array for ``dask.array.from_delayed``.""" return _merge_subdomains(paths_per_subdomain, use_cftime, pattern).values +def check_subdomain_consistency( + parsed: list[dict[str, Any]], paths: list[str], pattern: str | Pattern +): + 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: @@ -279,25 +307,13 @@ 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, 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}" - ) + check_subdomain_consistency(parsed, paths, pattern) # 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 = defaultdict(lambda: defaultdict(list)) + grouped_by_time: DefaultDict[Any, DefaultDict[Any, list]] = defaultdict( + lambda: defaultdict(list) + ) for match, p in zip(parsed, paths): grouped_by_time[match["time"]][match["subdomain"]].append(p) @@ -323,7 +339,9 @@ def open_subdomains( 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) + 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) From a04dd2067b03b46b3acffcdb8f0a76cbe204dabf Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Mon, 27 Jul 2026 17:28:30 +0200 Subject: [PATCH 04/13] Work around performance issue: avoid excessively creating xr.DataArrays to get equal performance to oplossing_Huite.py --- imod/formats/idf.py | 25 ++++++++++++++++++----- imod/typing/structured.py | 43 +++++++++++++++++++++++++++++++++------ 2 files changed, 57 insertions(+), 11 deletions(-) diff --git a/imod/formats/idf.py b/imod/formats/idf.py index 78b61c0f2..a6c7c7591 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -21,7 +21,7 @@ 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 @@ -240,8 +240,7 @@ def _merge_subdomains( if "subdomain" in da.dims: da = da.isel(subdomain=0, drop=True) das.append(da) - name = das[0].name - return merge_partitions(das)[name] + return merge_partitions_as_da_components(das) def _merge_subdomains_values( @@ -250,7 +249,21 @@ def _merge_subdomains_values( pattern: str | Pattern, ): """Wraps ``_merge_subdomains`` to return a numpy array for ``dask.array.from_delayed``.""" - return _merge_subdomains(paths_per_subdomain, use_cftime, pattern).values + 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: + data, coords, dims = _merge_subdomains(paths_per_subdomain, use_cftime, pattern) + return xr.DataArray( + data=data, + coords=coords, + dims=dims, + ) def check_subdomain_consistency( @@ -328,7 +341,9 @@ def open_subdomains( # 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(grouped_by_time[first_time_key], use_cftime, pattern) + 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") diff --git a/imod/typing/structured.py b/imod/typing/structured.py index e429a14a8..401b3ed0c 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, ...]]: + """ + 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,9 @@ 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) 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 +252,18 @@ 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 + +def _merge_partitions_da(das: List[xr.DataArray]) -> xr.DataArray: + """ + Merge a list of xarray DataArrays into a single DataArray. + """ + data, coords, dims = merge_partitions_as_da_components(das) return xr.DataArray( data=data, coords=coords, - dims=first.dims, + dims=dims, ) @@ -244,12 +275,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)}" From 812e2642502bc147f459f1e7ad4bb89c23f741a7 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Mon, 27 Jul 2026 17:43:05 +0200 Subject: [PATCH 05/13] attempt supporting no time coord as well. --- imod/formats/idf.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/imod/formats/idf.py b/imod/formats/idf.py index a6c7c7591..266cd0688 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -322,13 +322,19 @@ def open_subdomains( parsed = [imod.util.path.decompose(path, pattern) for path in paths] check_subdomain_consistency(parsed, paths, pattern) + has_time = "time" in parsed[0] + # 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): - grouped_by_time[match["time"]][match["subdomain"]].append(p) + if has_time: + for match, p in zip(parsed, paths): + grouped_by_time[match["time"]][match["subdomain"]].append(p) + else: + for match, p in zip(parsed, paths): + grouped_by_time["steady-state"][match["subdomain"]].append(p) # Sort and convert times before calling _merge_subdomains so that # use_cftime is already correct when the template is built. @@ -336,6 +342,7 @@ def open_subdomains( converted_times, use_cftime = imod.util.time._convert_datetimes( raw_times_sorted, use_cftime ) + is_steady_state = all(time == "steady-state" for time in converted_times) # Call _merge_subdomains eagerly for the first timestep to obtain a # coordinate template. No data is computed — only the coordinate arrays @@ -363,6 +370,8 @@ def open_subdomains( # Build the full time coordinate and replace the single-timestep one from the template. if use_cftime: time_coord = xr.CFTimeIndex(converted_times) + elif is_steady_state: + time_coord = np.array(converted_times, dtype=str) else: time_coord = np.array(converted_times, dtype="datetime64[ns]") From 8b9b20ddaa1703bba46699ba27b9045d9da8618f Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Tue, 28 Jul 2026 08:18:41 +0200 Subject: [PATCH 06/13] preserve name as well --- imod/formats/idf.py | 5 +++-- imod/typing/structured.py | 6 ++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/imod/formats/idf.py b/imod/formats/idf.py index 266cd0688..380a9ccdb 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -249,7 +249,7 @@ def _merge_subdomains_values( pattern: str | Pattern, ): """Wraps ``_merge_subdomains`` to return a numpy array for ``dask.array.from_delayed``.""" - data, _, _ = _merge_subdomains(paths_per_subdomain, use_cftime, pattern) + data, _, _, _ = _merge_subdomains(paths_per_subdomain, use_cftime, pattern) return data @@ -258,11 +258,12 @@ def merge_subdomains_to_dataarray( use_cftime: bool, pattern: str | Pattern, ) -> xr.DataArray: - data, coords, dims = _merge_subdomains(paths_per_subdomain, use_cftime, pattern) + data, coords, dims, name = _merge_subdomains(paths_per_subdomain, use_cftime, pattern) return xr.DataArray( data=data, coords=coords, dims=dims, + name=name, ) diff --git a/imod/typing/structured.py b/imod/typing/structured.py index 401b3ed0c..a5c273297 100644 --- a/imod/typing/structured.py +++ b/imod/typing/structured.py @@ -206,6 +206,7 @@ def merge_partitions_as_da_components( 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: @@ -252,18 +253,19 @@ def merge_partitions_as_da_components( # 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 + 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 = merge_partitions_as_da_components(das) + data, coords, dims, name = merge_partitions_as_da_components(das) return xr.DataArray( data=data, coords=coords, dims=dims, + name=name, ) From 83db3fea1e6a5605fff4a127aa92288342af11b6 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Tue, 28 Jul 2026 08:19:05 +0200 Subject: [PATCH 07/13] Test for name and add test without time --- imod/tests/test_formats/test_idf.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/imod/tests/test_formats/test_idf.py b/imod/tests/test_formats/test_idf.py index 009a34461..f2d9e7efe 100644 --- a/imod/tests/test_formats/test_idf.py +++ b/imod/tests/test_formats/test_idf.py @@ -156,6 +156,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 +187,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 +213,27 @@ 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}""" + _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 == ("time", "layer", "y", "x") + assert da.name == "subdomains_20000101" + + assert da.coords["time"] == "steady-state" + assert np.all(da.isel(time=0) == expected) + @parametrize_with_cases( "subdomains,expected,equidistant", cases=SubdomainCases, has_tag="species" ) @@ -230,6 +249,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 From 083c24433bebfcef4d1bf285b8000d0c50cfd972 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Tue, 28 Jul 2026 10:48:25 +0200 Subject: [PATCH 08/13] Fix type annotation --- imod/typing/structured.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imod/typing/structured.py b/imod/typing/structured.py index a5c273297..363c9d8ad 100644 --- a/imod/typing/structured.py +++ b/imod/typing/structured.py @@ -157,7 +157,7 @@ def _merge_nonequidistant_coords( def merge_partitions_as_da_components( das: List[xr.DataArray], -) -> tuple[np.ndarray, dict[str, Any], tuple[str, ...]]: +) -> tuple[np.ndarray, dict[str, Any], tuple[str, ...], str]: """ Merge a list of xarray DataArrays into components for a single DataArray. From c65415d4b85dc31e4c181a4d2055967ede0e3505 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Tue, 28 Jul 2026 10:49:27 +0200 Subject: [PATCH 09/13] Only do time handling and include time coord when necessary --- imod/formats/idf.py | 39 ++++++++++++++++------------- imod/tests/test_formats/test_idf.py | 7 +++--- 2 files changed, 26 insertions(+), 20 deletions(-) diff --git a/imod/formats/idf.py b/imod/formats/idf.py index 380a9ccdb..ac03966c2 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -258,7 +258,9 @@ def merge_subdomains_to_dataarray( use_cftime: bool, pattern: str | Pattern, ) -> xr.DataArray: - data, coords, dims, name = _merge_subdomains(paths_per_subdomain, use_cftime, pattern) + data, coords, dims, name = _merge_subdomains( + paths_per_subdomain, use_cftime, pattern + ) return xr.DataArray( data=data, coords=coords, @@ -330,12 +332,15 @@ def open_subdomains( grouped_by_time: DefaultDict[Any, DefaultDict[Any, list]] = defaultdict( lambda: defaultdict(list) ) - if has_time: - for match, p in zip(parsed, paths): - grouped_by_time[match["time"]][match["subdomain"]].append(p) - else: - for match, p in zip(parsed, paths): - grouped_by_time["steady-state"][match["subdomain"]].append(p) + + 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. @@ -343,7 +348,6 @@ def open_subdomains( converted_times, use_cftime = imod.util.time._convert_datetimes( raw_times_sorted, use_cftime ) - is_steady_state = all(time == "steady-state" for time in converted_times) # Call _merge_subdomains eagerly for the first timestep to obtain a # coordinate template. No data is computed — only the coordinate arrays @@ -356,7 +360,10 @@ def open_subdomains( shape = template.shape # e.g. (1, nlayer, nrow, ncol) dims = template.dims # e.g. ("time", "layer", "y", "x") dtype = template.dtype - time_axis = list(dims).index("time") + if has_time: + time_axis = list(dims).index("time") + else: + time_axis = -1 # steady-state, no time dimension # One delayed task per timestep → outer graph depth 3, O(n_time) tasks merged = [] @@ -369,15 +376,13 @@ def open_subdomains( data = dask.array.concatenate(merged, axis=time_axis) # Build the full time coordinate and replace the single-timestep one from the template. - if use_cftime: - time_coord = xr.CFTimeIndex(converted_times) - elif is_steady_state: - time_coord = np.array(converted_times, dtype=str) - else: - time_coord = np.array(converted_times, dtype="datetime64[ns]") - coords = dict(template.coords) - coords["time"] = time_coord + 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) diff --git a/imod/tests/test_formats/test_idf.py b/imod/tests/test_formats/test_idf.py index f2d9e7efe..9f3e04734 100644 --- a/imod/tests/test_formats/test_idf.py +++ b/imod/tests/test_formats/test_idf.py @@ -228,11 +228,12 @@ def test_open_subdomains_pattern_no_time(subdomains, expected, equidistant, tmp_ pattern = r"{name}_l{layer}_p{subdomain}" da = idf.open_subdomains(tmp_path / "subdomains_*.idf", pattern=pattern).load() - assert da.dims == ("time", "layer", "y", "x") + assert da.dims == ("layer", "y", "x") assert da.name == "subdomains_20000101" - assert da.coords["time"] == "steady-state" - assert np.all(da.isel(time=0) == expected) + assert "time" not in da.coords.keys() + assert np.all(da == expected) + @parametrize_with_cases( "subdomains,expected,equidistant", cases=SubdomainCases, has_tag="species" From 8a3e237a46bd3728eecdd08355cc8ba5c12757fc Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Tue, 28 Jul 2026 11:01:42 +0200 Subject: [PATCH 10/13] Add comments for clarification --- imod/formats/idf.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/imod/formats/idf.py b/imod/formats/idf.py index ac03966c2..b3ed785dd 100644 --- a/imod/formats/idf.py +++ b/imod/formats/idf.py @@ -248,16 +248,17 @@ def _merge_subdomains_values( use_cftime: bool, pattern: str | Pattern, ): - """Wraps ``_merge_subdomains`` to return a numpy array for ``dask.array.from_delayed``.""" + """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( +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 ) @@ -272,6 +273,7 @@ def merge_subdomains_to_dataarray( 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: @@ -312,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: @@ -353,7 +360,7 @@ def open_subdomains( # 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( + template = _merge_subdomains_to_dataarray( grouped_by_time[first_time_key], use_cftime, pattern ) @@ -365,7 +372,9 @@ def open_subdomains( else: time_axis = -1 # steady-state, no time dimension - # One delayed task per timestep → outer graph depth 3, O(n_time) tasks + # 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] @@ -375,7 +384,7 @@ def open_subdomains( 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 and replace the single-timestep one from the template. + # Build the full time coordinate coords = dict(template.coords) if has_time: if use_cftime: From 38e9a933099789ad35914d81fc0a2690768a3d62 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Tue, 28 Jul 2026 11:21:29 +0200 Subject: [PATCH 11/13] Clarify docstring --- imod/tests/test_formats/test_idf.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/imod/tests/test_formats/test_idf.py b/imod/tests/test_formats/test_idf.py index 9f3e04734..4bb45080f 100644 --- a/imod/tests/test_formats/test_idf.py +++ b/imod/tests/test_formats/test_idf.py @@ -222,7 +222,10 @@ def test_open_subdomains_pattern_None(subdomains, expected, equidistant, tmp_pat "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}""" + """ + 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}" From 0b117c9335bc442ebbcb792fb74a780ac81a099d Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Tue, 28 Jul 2026 11:50:27 +0200 Subject: [PATCH 12/13] Update changelog --- docs/api/changelog.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index bb47dc1b2..c4169ce8b 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 ~~~~~~~ From 533d621f7d9df96a7b0e69b05378659912d619f6 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Tue, 28 Jul 2026 13:57:50 +0200 Subject: [PATCH 13/13] Add user acceptance test --- docs/developing/test-plan.rst | 7 +++ imod/tests/test_formats/test_idf.py | 73 +++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+) 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/tests/test_formats/test_idf.py b/imod/tests/test_formats/test_idf.py index 4bb45080f..3d83106c1 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 @@ -465,3 +467,74 @@ def test_save_open_arbitrary_4D(tmp_path): assert isinstance(back, xr.DataArray) # Might get shuffled dimensions for the same reason. assert set(da.dims) == set(back.dims) + + +@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)