Improve progressbar in docs

docs/examples/02_surface_water.ipynb

@@ -27,10 +27,7 @@
     "import flopy\n",
     "import rioxarray\n",
     "import matplotlib.pyplot as plt\n",
-    "import nlmod\n",
-    "from geocube.api.core import make_geocube\n",
-    "from functools import partial\n",
-    "from geocube.rasterize import rasterize_image"
+    "import nlmod"
    ]
   },
   {

nlmod/read/ahn.py

@@ -1,6 +1,9 @@
 import datetime as dt
 import logging
 
+import matplotlib.pyplot as plt
+import numpy as np
+import geopandas as gpd
 import rasterio
 import rioxarray
 import xarray as xr
@@ -86,6 +89,35 @@ def get_ahn_at_point(
     res=0.5,
     **kwargs,
 ):
+    """
+    Get the height of the surface level at a certain point, defined by x and y.
+
+    Parameters
+    ----------
+    x : float
+        The x-coordinate fo the point.
+    y : float
+        The y-coordinate fo the point..
+    buffer : float, optional
+        The buffer around x and y that is downloaded. The default is 0.75.
+    return_da : bool, optional
+        Return the downloaded DataArray when True. The default is False.
+    return_mean : bool, optional
+        Resturn the mean of all non-nan pixels within buffer. Return the center pixel
+        when False. The default is False.
+    identifier : str, optional
+        The identifier passed onto get_latest_ahn_from_wcs. The default is "dsm_05m".
+    res : float, optional
+        The resolution that is passed onto get_latest_ahn_from_wcs. The default is 0.5.
+    **kwargs : dict
+        kwargs are passed onto the method get_latest_ahn_from_wcs.
+
+    Returns
+    -------
+    float
+        The surface level value at the requested point.
+
+    """
     extent = [x - buffer, x + buffer, y - buffer, y + buffer]
     ahn = get_latest_ahn_from_wcs(extent, identifier=identifier, res=res, **kwargs)
     if return_da:
@@ -99,6 +131,67 @@ def get_ahn_at_point(
         return ahn.data[int((ahn.shape[0] - 1) / 2), int((ahn.shape[1] - 1) / 2)]
 
 
+def get_ahn_along_line(line, ahn=None, dx=None, num=None, method="linear", plot=False):
+    """
+    Get the height of the surface level along a line.
+
+    Parameters
+    ----------
+    line : shapely.LineString
+        The line along which the surface level is calculated.
+    ahn : xr.DataArray, optional
+        The 2d DataArray containing surface level values. If None, ahn4-values are
+        downloaded from the web. The default is None.
+    dx : float, optional
+        The distance between the points along the line at which the surface level is
+        calculated. Only used when num is None. When dx is None, it is set to the
+        resolution of ahn. The default is None.
+    num : int, optional
+        If not None, the surface level is calculated at num equally spaced points along
+        the line. The default is None.
+    method : string, optional
+        The method to interpolate the 2d surface level values to the points along the
+        line. The default is "linear".
+    plot : bool, optional
+        if True, plot the 2d surface level, the line and the calculated heights. The
+        default is False.
+
+    Returns
+    -------
+    z : xr.DataArray
+        A DataArray with dimension s, containing surface level values along the line.
+
+    """
+    if ahn is None:
+        bbox = line.bounds
+        extent = [bbox[0], bbox[2], bbox[1], bbox[3]]
+        ahn = get_ahn4(extent)
+    if num is not None:
+        s = np.linspace(0.0, line.length, num)
+    else:
+        if dx is None:
+            dx = float(ahn.x[1] - ahn.x[0])
+        s = np.arange(0.0, line.length, dx)
+
+    x, y = zip(*[p.xy for p in line.interpolate(s)])
+
+    x = np.array(x)[:, 0]
+    y = np.array(y)[:, 0]
+
+    x = xr.DataArray(x, dims="s", coords={"s": s})
+    y = xr.DataArray(y, dims="s", coords={"s": s})
+    z = ahn.interp(x=x, y=y, method=method)
+
+    if plot:
+        _, ax = plt.subplots(figsize=(10, 10))
+        ahn.plot(ax=ax)
+        gpd.GeoDataFrame(geometry=[line]).plot(ax=ax)
+
+        _, ax = plt.subplots(figsize=(10, 10))
+        z.plot(ax=ax)
+    return z
+
+
 @cache.cache_netcdf
 def get_latest_ahn_from_wcs(
     extent=None,

pyproject.toml

@@ -74,9 +74,10 @@ rtd = [
     "nlmod[full]",
     "ipython",
     "ipykernel",
+    "ipywidgets",
     "nbsphinx",
     "sphinx_rtd_theme==1.0.0",
-    "nbconvert>6.4.5",
+    "nbconvert==7.13.0",
     "netCDF4>=1.6.3",
 ]
 

tests/test_005_external_data.py

@@ -1,4 +1,6 @@
 import pandas as pd
+import xarray as xr
+from shapely.geometry import LineString
 import pytest
 import test_001_model
 
@@ -65,9 +67,13 @@ def test_get_ahn3():
 
 def test_get_ahn4():
     extent = [98000.0, 100000.0, 494000.0, 496000.0]
-    da = nlmod.read.ahn.get_ahn4(extent)
+    ahn = nlmod.read.ahn.get_ahn4(extent)
+    assert isinstance(ahn, xr.DataArray)
+    assert not ahn.isnull().all(), "AHN only has nan values"
 
-    assert not da.isnull().all(), "AHN only has nan values"
+    line = LineString([(99000, 495000), (100000, 496000)])
+    ahn_line = nlmod.read.ahn.get_ahn_along_line(line, ahn=ahn)
+    assert isinstance(ahn_line, xr.DataArray)
 
 
 def test_get_ahn():
@@ -80,6 +86,10 @@ def test_get_ahn():
     assert not ahn_ds["ahn"].isnull().all(), "AHN only has nan values"
 
 
+def test_get_ahn_at_point():
+    nlmod.read.ahn.get_ahn_at_point(100010, 400010)
+
+
 def test_get_surface_water_ghb():
     # model with sea
     ds = test_001_model.get_ds_from_cache("basic_sea_model")
@@ -88,13 +98,13 @@ def test_get_surface_water_ghb():
     sim = nlmod.sim.sim(ds)
 
     # create time discretisation
-    tdis = nlmod.sim.tdis(ds, sim)
+    nlmod.sim.tdis(ds, sim)
 
     # create groundwater flow model
     gwf = nlmod.gwf.gwf(ds, sim)
 
     # create ims
-    ims = nlmod.sim.ims(sim)
+    nlmod.sim.ims(sim)
 
     nlmod.gwf.dis(ds, gwf)