Addition of raster functionality

.gitignore

@@ -20,6 +20,7 @@ temp*
 
 examples/*.inp
 wntr/tests/*.png
+wntr/tests/*.tif
 
 documentation/_local
 documentation/apidoc

documentation/gis.rst

@@ -31,6 +31,12 @@
     >>> hydrant_data = gpd.read_file(examples_dir+'/data/Net1_hydrant_data.geojson')
     >>> valve_data = gpd.read_file(examples_dir+'/data/Net1_valve_data.geojson')
 
+.. doctest::
+    :hide:
+    :skipif: gpd is None or rasterio is None
+
+    >>> elevation_data_path = examples_dir+'/data/Net1_elevation_data.tif'
+
 .. _geospatial:
 
 Geospatial capabilities
@@ -47,7 +53,8 @@ The following section describes capabilities in WTNR that use GeoPandas GeoDataF
 
 .. note:: 
    Functions that use GeoDataFrames require the Python package **geopandas** :cite:p:`jvfm21` 
-   and **rtree** :cite:p:`rtree`. Both are optional dependencies of WNTR.
+   and **rtree** :cite:p:`rtree`, and functions that use raster files require the
+   Python package **rasterio**. All three are optional dependencies of WNTR.
    Note that **shapely** is installed with geopandas.
 
 The following examples use a water network generated from Net1.inp.
@@ -822,3 +829,100 @@ the census tracts (polygons) is different than the junction and pipe attributes.
    :alt: Intersection of junctions and pipes with mean income demographic data in EPANET example Net1
 
    Net1 with mean income demographic data intersected with junctions and pipes.
+
+Sample raster at points geometries
+--------------------------------------
+
+The :class:`~wntr.gis.sample_raster` function can be used to sample a raster file at point geometries,
+such as the nodes of a water network. A common use case for this function is to assign elevation to the 
+nodes of a water network, however other geospatial information such as climate or hazard data could be sampled 
+using this function.
+
+The network file, Net1.inp, in EPSG:4326 CRS is used in the example below. 
+The raster data in the GeoTIFF format is also in EPSG:4326 CRS.
+See :ref:`crs` for more information.
+
+.. doctest::
+    :skipif: gpd is None
+
+    >>> wn = wntr.network.WaterNetworkModel('networks/Net1.inp') # doctest: +SKIP
+    >>> wn_gis = wntr.network.to_gis(wn, crs='EPSG:4326')
+
+Sample elevations at junctions
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Elevation is an essential attribute for accurate simulation of pressure in a water network and is
+commonly provided in GeoTIFF (.tif) files. The following example shows how such files can be sampled 
+and assigned to the junctions and tanks of a network. 
+
+.. doctest::
+    :skipif: gpd is None or rasterio is None
+
+    >>> elevation_data_path = 'data/Net1_elevation_data.tif' # doctest: +SKIP
+    >>> junctions = wn_gis.junctions
+    >>> junction_elevations = wntr.gis.sample_raster(junctions, elevation_data_path)
+    >>> print(junction_elevations)
+    name
+    10    1400.0
+    11    2100.0
+    12    3500.0
+    13    4900.0
+    21    1200.0
+    22    2000.0
+    23    2800.0
+    31     300.0
+    32     500.0
+    dtype: float64
+
+.. doctest::
+    :skipif: gpd is None or rasterio is None
+
+    >>> tanks = wn_gis.tanks
+    >>> tank_elevations = wntr.gis.sample_raster(tanks, elevation_data_path)
+    >>> print(tank_elevations)
+    name
+    2    4500.0
+    dtype: float64
+
+To use these elevations for hydraulic simulations, 
+they need to be added to the water network object.
+
+.. doctest::
+    :skipif: gpd is None or rasterio is None
+
+    >>> for junction_name in wn.junction_name_list:
+    ...     junction = wn.get_node(junction_name)
+    ...     junction.elevation = junction_elevations[junction_name]
+
+.. doctest::
+    :skipif: gpd is None or rasterio is None
+
+    >>> for tank_name in wn.tank_name_list:
+    ...     tank = wn.get_node(tank_name)
+    ...     tank.elevation = tank_elevations[tank_name]
+
+The sampled elevations can be plotted as follows. The 
+resulting :numref:`fig-sample-elevations` illustrates Net1 with the elevations 
+sampled from the raster file.
+
+.. doctest::
+    :skipif: gpd is None or rasterio is None
+
+    >>> ax = wntr.graphics.plot_network(wn, node_attribute="elevation", link_width=1.5, 
+    ...     node_size=40, node_colorbar_label='Raster Elevation')
+
+.. doctest::
+    :skipif: gpd is None or rasterio is None
+    :hide:
+
+    >>> bounds = ax.axis('equal')
+    >>> plt.tight_layout()
+    >>> plt.savefig('sample_elevations.png', dpi=300)
+    >>> plt.close()
+
+.. _fig-sample-elevations:
+.. figure:: figures/sample_elevations.png
+   :width: 640
+   :alt: Net1 with elevations sampled from raster.
+
+   Net1 with elevations sampled from raster.

documentation/installation.rst

@@ -272,6 +272,8 @@ The following Python packages are optional:
   https://pypi.org/project/utm/
 * geopandas :cite:p:`jvfm21`: used to work with geospatial data,
   https://geopandas.org/
+* rasterio :cite:p:`rasterio`: used to work with raster data,
+  https://rasterio.readthedocs.io/
 * rtree :cite:p:`rtree`: used for overlay operations in geopandas,
   https://rtree.readthedocs.io/
 * openpyxl :cite:p:`gacl18`: used to read/write to Microsoft® Excel® spreadsheets,

documentation/references.bib

@@ -128,6 +128,14 @@ @misc{jvfm21
     year = "2021"
 }
 
+@misc{rasterio,
+  author =    {Sean Gillies and others},
+  organization = {Mapbox},
+  title =     {{Rasterio: geospatial raster I/O for {Python} programmers}},
+  year =      {2013--},
+  url = {"https://github.com/rasterio/rasterio"}
+}
+
 @article{jcmg11,
     author = "Joyner, David and \v{C}ert\'{i}k, Ond\v{r}ej and Meurer, Aaron and Granger, Brian E.",
     address = "New York, NY, USA",

requirements.txt

@@ -13,6 +13,7 @@ utm
 openpyxl
 geopandas<1.0
 fiona<1.10
+rasterio
 rtree
 
 # Documentation

wntr/gis/__init__.py

@@ -3,5 +3,5 @@
 and GIS formatted data and geospatial functions to snap data and find intersections.
 """
 from wntr.gis.network import WaterNetworkGIS
-from wntr.gis.geospatial import snap, intersect
+from wntr.gis.geospatial import snap, intersect, sample_raster
 

wntr/gis/geospatial.py

@@ -6,6 +6,7 @@
 import pandas as pd
 import numpy as np
 
+
 try:
     from shapely.geometry import MultiPoint, LineString, Point, shape
     has_shapely = True
@@ -18,6 +19,13 @@
 except ModuleNotFoundError:
     gpd = None
     has_geopandas = False
+
+try:
+    import rasterio
+    has_rasterio = True
+except ModuleNotFoundError:
+    rasterio = None
+    has_rasterio = False
 
 
 def snap(A, B, tolerance):  
@@ -57,9 +65,9 @@ def snap(A, B, tolerance):
     if not has_shapely or not has_geopandas:
         raise ModuleNotFoundError('shapley and geopandas are required')
 
-    isinstance(A, gpd.GeoDataFrame)
+    assert isinstance(A, gpd.GeoDataFrame)
     assert(A['geometry'].geom_type).isin(['Point']).all()
-    isinstance(B, gpd.GeoDataFrame)
+    assert isinstance(B, gpd.GeoDataFrame)
     assert (B['geometry'].geom_type).isin(['Point', 'LineString', 'MultiLineString']).all()
     assert A.crs == B.crs
 
@@ -197,18 +205,18 @@ def intersect(A, B, B_value=None, include_background=False, background_value=0):
     if not has_shapely or not has_geopandas:
         raise ModuleNotFoundError('shapley and geopandas are required')
 
-    isinstance(A, gpd.GeoDataFrame)
+    assert isinstance(A, gpd.GeoDataFrame)
     assert (A['geometry'].geom_type).isin(['Point', 'LineString', 
                                            'MultiLineString', 'Polygon', 
                                            'MultiPolygon']).all()
-    isinstance(B, gpd.GeoDataFrame)
+    assert isinstance(B, gpd.GeoDataFrame)
     assert (B['geometry'].geom_type).isin(['Point', 'LineString', 
                                            'MultiLineString', 'Polygon', 
                                            'MultiPolygon']).all()
     if isinstance(B_value, str):
         assert B_value in B.columns
-    isinstance(include_background, bool)
-    isinstance(background_value, (int, float))
+    assert isinstance(include_background, bool)
+    assert isinstance(background_value, (int, float))
     assert A.crs == B.crs, "A and B must have the same crs."
 
     if include_background:
@@ -280,4 +288,50 @@ def intersect(A, B, B_value=None, include_background=False, background_value=0):
 
     stats.index.name = None
 
-    return stats
+    return stats
+
+
+def sample_raster(A, filepath, bands=1):
+    """Sample a raster (e.g., GeoTIFF file) using Points in GeoDataFrame A. 
+
+    This function can take either a filepath to a raster or a virtual raster 
+    (VRT), which combines multiple raster tiles into a single object. The 
+    function opens the raster(s) and samples it at the coordinates of the point 
+    geometries in A. This function assigns nan to values that match the 
+    raster's `nodata` attribute. These sampled values are returned as a Series 
+    which has an index matching A.
+
+    Parameters
+    ----------
+    A : GeoDataFrame
+        GeoDataFrame containing Point geometries
+    filepath : str
+        Path to raster or alternatively a virtual raster (VRT)
+    bands : int or list[int] (optional, default = 1)
+        Index or indices of the bands to sample (using 1-based indexing)
+
+    Returns
+    -------
+    Series
+        Pandas Series containing the sampled values for each geometry in gdf
+    """
+    # further functionality could include other geometries (Line, Polygon),
+    # and use of multiprocessing to speed up querying.
+    if not has_rasterio:
+        raise ModuleNotFoundError('rasterio is required')
+
+    assert (A['geometry'].geom_type == "Point").all()
+    assert isinstance(filepath, str)
+    assert isinstance(bands, (int, list))
+
+    with rasterio.open(filepath) as raster:
+        xys = zip(A.geometry.x, A.geometry.y)
+
+        values = np.array(
+            tuple(raster.sample(xys, bands)), dtype=float # force to float to allow for conversion of nodata to nan
+        ).squeeze()
+
+    values[values == raster.nodata] = np.nan
+    values = pd.Series(values, index=A.index)
+
+    return values

wntr/tests/test_gis.py

@@ -23,6 +23,13 @@
     gpd = None
     has_geopandas = False
 
+try:
+    import rasterio
+    has_rasterio = True
+except ModuleNotFoundError:
+    rasterio = None
+    has_rasterio = False
+
 testdir = dirname(abspath(str(__file__)))
 datadir = join(testdir, "networks_for_testing")
 ex_datadir = join(testdir, "..", "..", "examples", "networks")
@@ -69,7 +76,7 @@ def setUpClass(self):
 
         df = pd.DataFrame(point_data)
         self.points = gpd.GeoDataFrame(df, crs=None)
-        
+
     @classmethod
     def tearDownClass(self):
         pass
@@ -311,5 +318,48 @@ def test_snap_points_to_lines(self):
 
         assert_frame_equal(pd.DataFrame(snapped_points), expected, check_dtype=False)
 
+@unittest.skipIf(not has_rasterio,
+                 "Cannot test raster capabilities: rasterio is missing")
+class TestRaster(unittest.TestCase):
+    @classmethod
+    def setUpClass(self):
+        # use net1 junctions as example points
+        inp_file = join(ex_datadir, "Net1.inp")
+        wn = wntr.network.WaterNetworkModel(inp_file)
+        wn_gis = wn.to_gis(crs="EPSG:4326")
+        points = pd.concat((wn_gis.junctions, wn_gis.tanks))
+        self.points = points
+
+        min_lon, min_lat, max_lon, max_lat = self.points.total_bounds
+
+        resolution = 1.0
+
+        # adjust to include boundary
+        max_lon += resolution
+        min_lat -= resolution
+
+        lon_values = np.arange(min_lon, max_lon, resolution)
+        lat_values = np.arange(max_lat, min_lat, -resolution)  # Decreasing order for latitudes
+        raster_data = np.outer(lat_values,lon_values) # value is product of coordinate
+
+        transform = rasterio.transform.from_origin(min_lon, max_lat, resolution, resolution)
+        with rasterio.open(
+            "test_raster.tif", "w", driver="GTiff", height=raster_data.shape[0], width=raster_data.shape[1], 
+            count=1, dtype=raster_data.dtype, crs="EPSG:4326", transform=transform) as file:
+            file.write(raster_data, 1) 
+
+    @classmethod
+    def tearDownClass(self):
+        pass
+
+    def test_sample_raster(self):
+        raster_values = wntr.gis.sample_raster(self.points, "test_raster.tif")
+        assert (raster_values.index == self.points.index).all()
+
+        # values should be product of coordinates
+        expected_values = self.points.apply(lambda row: row.geometry.x * row.geometry.y, axis=1)
+        assert np.isclose(raster_values.values, expected_values, atol=1e-5).all()
+
+
 if __name__ == "__main__":
     unittest.main()