Issue #1844 performance idf open subdomains

[JoerivanEngelen]

Fixes #1844

Description

Reduce dask task graph considerably by only chunking per time (and optionally per species). To avoid having to construct the coords manually again from the headers like the previous implementation, we eagerly call merge_partitions on the first timestep to just get the coords (but do not load them into memory). This implementation has similar performance to the original code. I had to modify the return statements in imod.typing.structured.py somewhat to avoid unnecessary creation of xr.DataArrays and xr.Datasets. This might look a bit confusing at first, I added some docstrings to clarify.

Other things I tried:

I also tried rechunking to layers after calling merge_partitions, so that not so much code had to be changed, but this still resulted in poor performance (5 minutes versus 6 seconds with 37400 idfs) as dask still had to optimize the task graph. (Still better than the 30+ minutes it cost before)

TODOs:

• Make time dimension optional
• See if grouping dicts twice is necessary
• Add a user acceptance test to deal with many IDFs, hopefully teamcity can handle this.

Checklist

• Links to correct issue
• Update changelog, if changes affect users
• PR title starts with Issue #nr, e.g. Issue #737
• Unit tests were added
• If feature added: Added/extended example
• If feature added: Added feature to API documentation
• If pixi.lock was changed: Ran pixi run generate-sbom and committed changes

[Huite]

Some notes:

• There may be slight overhead in constructing DataArrays in the merge_subdomains function; this requires alignment checks etc.
• This will only show up in a .compute operation.
• If relevant, a workaround is to essentially define/modify an internal functions which returns three arguments (data, dims, coords) rather than a DataArray.
• The obviates the .values wrapper here -- eliminating that might be nice anyway!
• The template can still be constructed using the higher level function which does return a DataArray: the validation xarray performs is worthwhile.

EDIT: overhead is minimal; for 3000 time steps with 50 layers, 1000 by 1000 rows & columns, it's about 1 ms of overhead.
So the only reason to do so is to avoid the .values wrapper.

[JoerivanEngelen]

Update: I managed to further improve performance by a factor 2, so that it is equal to the original solution by avoiding excessive unnecessary casting to xr.DataArray. It complicates the code somewhat, as there is now a separate function to return the components (data, coords, dims) separately, as well as a function which wraps this into a xr.DataArray.

[sonarqubecloud]

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[ClaireDons]

Nice, it was quite interesting to go through the changes. I just have a couple of questions but otherwise looks good to me :)