Issue #1844 performance idf open subdomains

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
 ~~~~~~~

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
 *************

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):

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)