Deltares / Deltares/imod-python

simulation.regrid_like seems excessively slow + a hypothesis

Open
#844 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

performance
Dominant language
Python
Stars
41
Forks
12
Avg merge
21h 8m
Merged PRs (30d)
1

Description

I was looking at a workflow made by @HendrikKok, together with @jc-hunink.
This workflow makes use of the regrid_like functionality of the Modflow6Simulation. It was taking a very long time, while the data were relatively small (approximately 300 by 300 cells). We didn't time it, but the regrid_like took in the order of minutes. This seemed very suspicious, because the data were coming from a project file, and so they were dask-backed DataArrays. In this case, no IO is happening, no actual regridding is happening, the only somewhat costly step is computing the regridding weights. Because we're only regridding in 2D, this is a few calls to numpy searchsorted followed by a broadcasting operation. Nothing big.

The regridding weights are computed for each variable. The weights are currently not re-used (#429), which is obviously not the best for performance and it will inflate memory usage somewhat (they should be shared on the model level). But this doesn't seem to be the problem: I ran a manual test on two DataArrays and it was quite quick. Even adding a .compute resulted in a running time of only 0.25 s.
Since the model has about 10 packages, with on average at most 5 variables, this should result in a regrid_like time of no more than 10 to 15 seconds. Instead, we were waiting for minutes: at least an order of magnitude slower!

The question is: if it's not IO, and it's not the regridder, what's causing the delay?
My current hypothesis is that it's caused by repeated the alignment that xarray performs. I've had a number of bad experiences with it.
It's one of the reasons that are several places where I've used dask.array.stack or dask.array.block instead of using an xarray merge or concat. In those cases, I was also looking at an order of magnitude slowdown.
Maybe this matches the profiling by @JoerivanEngelen in https://github.com/Deltares/imod-python/pull/722#issuecomment-1884907538

So instead of setting variables one by one, they could be set via xr.merge(), here: https://github.com/Deltares/imod-python/blob/b59a7700a75d10bf33ed51c6e532bf89ac625410/imod/mf6/package.py#L680

In case alignment remains a bottleneck, we could collect the .data from every regridded results, and build the xarray.Dataset directly from a dict of numpy/dask arrays. From the xarray docs:

temp = 15 + 8 * np.random.randn(2, 2, 3)
precip = 10 * np.random.rand(2, 2, 3)
lon = [[-99.83, -99.32], [-99.79, -99.23]]
lat = [[42.25, 42.21], [42.63, 42.59]]
ds = xr.Dataset(
    {
        "temperature": (["x", "y", "time"], temp),
        "precipitation": (["x", "y", "time"], precip),
    },
    coords={
        "lon": (["x", "y"], lon),
        "lat": (["x", "y"], lat),
        "time": pd.date_range("2014-09-06", periods=3),
        "reference_time": pd.Timestamp("2014-09-05"),
    },
)

Since the individual variables don't come with coordinates attached, there's only a single coordinate definition, and xarray won't have to do any alignment.

In principle, such work could be done behind the scenes as well, if the xugrid Regridders accept datasets as arguments. However, most regridding is variable specific, so it generally doesn't make that much sense to do so in with a single call. You'd need to provide another dict mapping var -> method. It would also mean you cannot cache regridding weights smartly: they would only be shared within a single dataset. And in case of simulations, you primarily expect methods to be shared across packages (datasets), not across variables within a single dataset.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at Modflow6Simulation.regrid_like and the variable-setting code in imod/mf6/package.py around line 680. Profile a representative 300-by-300 workflow and compare repeated variable assignment with the proposed merge approach, checking whether alignment dominates runtime. Done means identifying the delay and documenting a measured improvement or a confirmed cause.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.