Bug description
@jdelsman found this.
There is a check for if isinstance(xr.DataArray) and "layer" is pointwaterhead.dims,
which should be turned into:
if isinstance(xr.DataArray, xu.UgridDataArray) and "layer" is pointwaterhead.dims
I don't think we have any unittests for this function; so it should be considered experimental at the moment.