fix #549

[dbrakenhoff]

• find first cross method to guarantee finding the first point value is reached.
• add right/left kwargs to allow setting values for values above/below interpolation array

[OnnoEbbens]

Thanks for the PR. It works perfectly. I had already started to develop another function which does the same thing (in a totally different way). Just posting it here for reference

I think we should use your function in this PR for nlmod because it is more readable and probably easier to scale for large DataArrays. Although my function is about 5x faster (tested on the Schoonhoven example notebook).

def get_isosurface(ds,
                   da,
                   value,
                   left=np.nan,
                   right=np.nan):
    """Linear interpolation to compute the elevation of an isosurface.

    Currently only supports linear interpolation.

    Parameters
    ----------
    ds : xarray.Dataset
        Model dataset, should contain top and botm. 
    da : xarray.DataArray
        DataArray containing the values for which the isosurface is computed.
        Should have a layer dimension.
    value : float
        The value of the isosurface to compute.
    left : float, optional
        elevation value to return when value is above the maximum of da.
        The default is np.nan.
    right : float, optional
        elevation value to return when value is below the minimum of da.
        The default is np.nan.

    Returns
    -------
    xarray.DataArray
        The elevation of the isosurface. Dimensions are equal to the dimensions of da 
        except for the layer dimension, which is dropped.

    Note
    ----
    if the values in da in all layers at a location (x,y) are nan, a nan value is returned.
    Otherwise the isosurface is computed using the non-nan values.
    """
    
    # get cell centers
    thickness = nlmod.dims.calculate_thickness(ds)
    top = ds["botm"] + thickness
    zcellcenters = (top + ds["botm"]) / 2

    # deal with nan values
    if da.isnull().any():
        # dataarray contains nan values, assuming these are inactive cells with zero thickness
        zcellcenters = xr.where(da.isnull(), np.nan, zcellcenters).ffill(dim='layer')
        da = da.ffill(dim='layer')
        zcellcenters = zcellcenters.bfill(dim='layer')
        da = da.bfill(dim='layer')

    # get boolean arrays
    above_or_at_threshold = da >= value
    at_threshold = da == value

    # get the layer above the isosurface from boolean array
    ilay_isosurface = abs(above_or_at_threshold.astype(int).diff(dim='layer')).argmax(dim='layer')

    # use linear interpolation to get the isosurface elevation
    val_above = da.isel(layer=ilay_isosurface)
    val_below = da.isel(layer=ilay_isosurface+1)
    z_above = zcellcenters.isel(layer=ilay_isosurface)
    z_below = zcellcenters.isel(layer=ilay_isosurface + 1)
    isosurface = z_above + (z_below - z_above) * (value - val_above) / (val_below - val_above)

    # edge case1: all da values are above threshold
    isosurface = xr.where(above_or_at_threshold.all(dim='layer') & (~at_threshold.any(dim='layer')),
                          left,
                          isosurface)

    # edge case2: all da values are above or at threshold and at least one value matches threshold
    isosurface = xr.where(above_or_at_threshold.all(dim='layer') & (at_threshold.any(dim='layer')),
                          zcellcenters.isel(layer=at_threshold.argmax(dim='layer')),
                          isosurface)

    # edge case3: the threshold is exactly matched in the top cell of da, the top layer is the isosurface
    isosurface = xr.where(at_threshold.isel(layer=0),
                          zcellcenters.isel(layer=0),
                          isosurface)

    # edge case4: all da values are below threshold
    isosurface = xr.where((~above_or_at_threshold).all(dim='layer'), right, isosurface)

    return isosurface

[OnnoEbbens]

I also made an isosurface using 'nearest' interpolation.

def get_isosurface_nearest(ds,
                            da,
                            value,
                            left=np.nan,
                            right=np.nan):
    """get the elevation of the neareast cell top/bottom to an isosurface.

    Parameters
    ----------
    ds : xarray.Dataset
        Model dataset, should contain top and botm. 
    da : xarray.DataArray
        DataArray containing the values for which the isosurface is computed.
        Should have a layer dimension.
    value : float
        The value of the isosurface to compute.
    left : float, optional
        elevation value to return when value is above the maximum of da.
        The default is np.nan.
    right : float, optional
        elevation value to return when value is below the minimum of da.
        The default is np.nan.

    Returns
    -------
    xarray.DataArray
        The elevation of the isosurface. Dimensions are equal to the dimensions of da 
        except for the layer dimension, which is dropped.

    Note
    ----
    if the values in da in all layers at a location (x,y) are nan, a nan value is returned.
    Otherwise the isosurface is computed using the non-nan values.
    """
    botm = ds['botm']
    top = ds['top']

    # deal with nan values
    if da.isnull().any():
        # dataarray contains nan values, assuming these are inactive cells with zero thickness
        botm = xr.where(da.isnull(), np.nan, botm).ffill(dim='layer')
        top = xr.where(da.isnull(), np.nan, top).ffill(dim='layer')
        da = da.ffill(dim='layer')
        botm = botm.bfill(dim='layer')
        top = top.bfill(dim='layer')
        da = da.bfill(dim='layer')

    # get boolean arrays
    above_or_at_threshold = da >= value
    at_threshold = da == value

    # get the layer above the isosurface from boolean array
    ilay_isosurface = abs(above_or_at_threshold.astype(int).diff(dim='layer')).argmax(dim='layer')

    # get the bottom of the layer above the isosurface, which is the nearest cell border to the isosurface
    isosurface = botm.isel(layer=ilay_isosurface)

    # edge case1: all da values are at or above threshold
    isosurface = xr.where(above_or_at_threshold.all(dim='layer'),
                          left,
                          isosurface)

    # edge case3: the threshold is exactly matched in the top cell of da, the top layer is the isosurface
    isosurface = xr.where(at_threshold.isel(layer=0),
                          top,
                          isosurface)

    # edge case4: all da values are below threshold
    isosurface = xr.where((~above_or_at_threshold).all(dim='layer'), right, isosurface)

    return isosurface 

[dbrakenhoff]

Good busy! I think computing the isosurface directly from the entire array object is probably faster until the arrays start getting really big.

There are some theoretical opportunities to parallelize/numba the 1d version when applying it to a DataArray but I haven't dared to attempt those optimizations. Additionally it's supposed to work on dask arrays, allowing delayed computation for very large datasets that do not fit into memory. The xr.apply_ufunc is an intimidating and confusing method however, and I haven't bothered to fully investigate all the options yet.

As for implementing different isosurface methods, alternatively, it could be possible to have 1d versions of your code as well, which can be applied in the current implementation of get_isosurface, e.g. get_isosurface(method="nearest") which then implements get_nearest_surface().

But if you prefer your implementations, I'm fine replacing what I've written, and perhaps moving that to a bigdata.py file or something. I think it's worth keeping around for when we do run into something very large, but for most models we're building now, your methods are just as convenient and faster.

[dbrakenhoff]

@OnnoEbbens, I asked copilot for some help building the optimized and parallel friendly versions of the isosurface_1d function based on my initial template. I've implemented both of these in get_isosurface() and you can select "numpy" or "numba" with the method="numpy" kwarg. Both work on in-memory arrays and dask arrays, and both work in parallel.

The numba implementation is a lot faster than the original (tested here for the entire LHM chloride dataset), note this is not the final syntax now, this was from earlier testing:

numba was not yet a dependency of nlmod, so I added it. Should we make it optional instead, or is this fine?

[rubencalje]

@OnnoEbbens, I asked copilot for some help building the optimized and parallel friendly versions of the isosurface_1d function based on my initial template. I've implemented both of these in get_isosurface() and you can select "numpy" or "numba" with the method="numpy" kwarg. Both work on in-memory arrays and dask arrays, and both work in parallel.

The numba implementation is a lot faster than the original (tested here for the entire LHM chloride dataset), note this is not the final syntax now, this was from earlier testing:

numba was not yet a dependency of nlmod, so I added it. Should we make it optional instead, or is this fine?

I am in favor of making numba optional. Numba can be hard to install, which can create problems when you want to use nlmod in a qgis-plugin for example. We can show a warning in methods that use numba that numba is preferred.