nfdata 0.2.0__tar.gz → 0.3.0__tar.gz

{nfdata-0.2.0 → nfdata-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nfdata
-Version: 0.2.0
+Version: 0.3.0
 Summary: Compile and edit input data for the NanoFASE model
 License: BSD-3-Clause
 Author: Sam Harrison
@@ -18,9 +18,11 @@ Requires-Dist: geopandas (>=1.0.0,<2.0.0)
 Requires-Dist: netcdf4 (>=1.7.1.post1,<2.0.0)
 Requires-Dist: numpy (>=2.0.0,<3.0.0)
 Requires-Dist: pandas (>=2.2.2,<3.0.0)
+Requires-Dist: pytest (>=8.3.5,<9.0.0)
 Requires-Dist: rasterio (>=1.3.10,<2.0.0)
 Requires-Dist: ruamel.yaml (>=0.18.6,<0.19.0)
 Requires-Dist: shapely (>=2.0.4,<3.0.0)
+Requires-Dist: whitebox (>=2.3.6,<3.0.0)
 Project-URL: Documentation, https://nerc-ceh.github.io/nanofase
 Description-Content-Type: text/markdown
 

{nfdata-0.2.0 → nfdata-0.3.0}/config.create.example.yaml

@@ -11,16 +11,38 @@ time:
   n: 365                                              # Number of timesteps to run the model for
   dt: 86400                                           # Length of each timestep in seconds
   start_date: 2015-01-01                              # Start date for the model run
+# If the `condition_dem` config option is present, the compiler will attempt to remove
+# pits, depressions and flat areas from the DEM using WhiteboxTools'
+# `breach_depressions_least_cost`.
+condition_dem:
+  # If `save_dem_to_path` is set, the conditioned DEM will be saved to this path
+  # (even if the conditioning hasn't fully worked). If not, a temporary directory
+  # is used and the resulting DEM is deleted after the processing.
+  save_dem_to_path: ./dem.tif
+  # Additional config options specified here will be passed to `breach_depressions_least_cost`
+  # See https://www.whiteboxgeo.com/manual/wbt_book/available_tools/hydrological_analysis.html#BreachDepressionsLeastCost
+  # If a value for `callback` is provided, it is ignored as we implement our own callback to
+  # deal with errors. The only required (by Whitebox) parameter is `dist`, which specifies
+  # the maximum search distance for breach paths in cells. The default value used in nfdata is 10000.
+  dist: 10000
 
 # Meteorological and hydrological
-flow_dir:                                             # Flow direction is responsible for defining the grid system
-  type: raster                                        # Type can be raster or csv (see docs)
+# -------------------------------
+# The model requires either a flow direction or DEM raster. If no flow direction
+# is provided, the DEM will be used to calculate flow direction. If both are provided,
+# the flow direction will be used. The DEM must either be conditioned (contain no pits,
+# depressions or flat areas), or if the config option `condition_dem` is set to `True`,
+# the compiler will attempt to condition the DEM using WhiteboxTools'
+# `breach_depressions_least_cost`. This is not guaranteed to work, and therefore users
+# might wish to condition the DEM themselves.
+# `flow_dir` and `dem` must be rasters that are openable by Rasterio.
+flow_dir:
   path: <root_dir>flow_dir.tif                        # Path to file, <root_dir> variable can be used from above
   crs: EPSG:27700                                     # Manually specify CRS if it is absent or ill-defined in file
-dem:                                                  # Digital Elevation Model, which defines the height above sea level
-  type: raster                                        # of each grid cell, in metres
+dem:
+  type: raster
   units: dm
-  path: <root_dir>dem.grd
+  path: <root_dir>dem.tif
 runoff:                                               # Runoff, usually provided by a hydrological model
   type: csv
   path: <root_dir>runoff.csv
@@ -41,7 +63,7 @@ precip:                                               # Precipitation
 # Soil
 # Most soil properties can be obtained from the European Soil Data Center (ESDAC): https://esdac.jrc.ec.europa.eu/
 soil_bulk_density:
-  type: raster
+  type: raster                                        # For spatial variables, only rasters are supported at the moment. Spatiotemporal variables can be CSVs
   units: t/m**3
   path: <root_dir>soil_bulk_density.tif
   # Source example:

{nfdata-0.2.0 → nfdata-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nfdata"
-version = "0.2.0"
+version = "0.3.0"
 description = "Compile and edit input data for the NanoFASE model"
 authors = ["Sam Harrison <samharrison.xg@gmail.com>"]
 license = "BSD-3-Clause"
@@ -24,8 +24,13 @@ netcdf4 = "^1.7.1.post1"
 shapely = "^2.0.4"
 Pint = "^0.24.1"
 f90nml = "^1.4.4"
+pytest = "^8.3.5"
+whitebox = "^2.3.6"
 
 
+[tool.poetry.group.dev.dependencies]
+flake8 = "^7.2.0"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

nfdata-0.3.0/src/nfdata/_errors.py

@@ -0,0 +1,11 @@
+class RoutingError(Exception):
+    """Exception raised for errors while routing waterbodies."""
+
+
+def whitebox_callback(msg):
+    """
+    Callback function to catch Whitebox errors
+    """
+    if 'Error' in msg:
+        raise RoutingError(f'Whitebox raised an error: '
+                           f'{msg}')

nfdata-0.3.0/src/nfdata/_utility.py

@@ -0,0 +1,52 @@
+"""
+Utility functions for nfdata.
+"""
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import rasterio
+from affine import Affine
+from rasterio.coords import BoundingBox
+from rasterio.crs import CRS
+from shapely.geometry import Polygon, box
+
+
+@dataclass
+class RasterInfo:
+    bounds: BoundingBox
+    crs: CRS
+    transform: Affine
+    res: Tuple[float, float]
+    width: int
+    height: int
+    dtype: str
+    count: int
+    shape: tuple
+    nodata: Optional[float]
+    bbox: Polygon
+    arr: Union[np.ndarray, np.ma.MaskedArray]
+
+
+def extract_raster_info(path: Union[str, Path],
+                        masked: bool = True) -> RasterInfo:
+    """
+    Extract data and info from a raster file and store
+    it in a RasterInfo dataclass so that it can be closed
+    """
+    with rasterio.open(path) as src:
+        return RasterInfo(
+            bounds=src.bounds,
+            crs=src.crs,
+            transform=src.transform,
+            res=src.res,
+            width=src.width,
+            height=src.height,
+            dtype=src.dtypes[0],
+            count=src.count,
+            shape=src.shape,
+            nodata=src.nodata,
+            bbox=box(*src.bounds),
+            arr=src.read(1, masked=masked)
+        )

{nfdata-0.2.0 → nfdata-0.3.0}/src/nfdata/compiler.py

@@ -1,26 +1,34 @@
+import importlib.resources as pkg_resources
+import os
 import shutil
 import sys
-import os
-import importlib.resources as pkg_resources
-from netCDF4 import Dataset
+from pathlib import Path
+
+import f90nml
+import geopandas as gpd
 import numpy as np
 import pandas as pd
-import geopandas as gpd
 import rasterio
+from netCDF4 import Dataset
+from pint import UnitRegistry
 from rasterio.crs import CRS
 from rasterio.mask import mask
-from rasterio.warp import reproject, Resampling
-from shapely.geometry import box
+from rasterio.warp import Resampling, reproject
 from ruamel.yaml import YAML
-from pint import UnitRegistry
-import f90nml
-from .router import Router
+from shapely.geometry import box
+
+from . import routing
+from ._utility import extract_raster_info
 
 
 class Compiler:
 
-    def __init__(self, task, config_path, model_vars_path):
+    def __init__(self, task, config_path, model_vars_path=None):
         """Initialise the compiler, either to compile or edit the dataset."""
+        # If we haven't been provided a model vars file, use the one in the package
+        if model_vars_path is None:
+            model_vars_path = pkg_resources.files(__package__)\
+                .joinpath('model_vars.yaml')
         # Open the config and model_vars files
         self.yaml = YAML(typ='safe')
         with open(config_path, 'r') as config_file, open(model_vars_path, 'r') as model_vars_file:
@@ -56,12 +64,15 @@ class Compiler:
                 self.vars_spatiotemporal.append(k)
             elif ('dims' in v) and (v['dims'] == ['p', 't', 'y', 'x']):
                 self.vars_spatial_point.append(k)
-            elif ('dims' in v) and (all(x in v['dims'] for x in ['y', 'x'])) and (len(v['dims']) == 3):
+            elif ('dims' in v) and \
+                    (all(x in v['dims'] for x in ['y', 'x'])) and \
+                    (len(v['dims']) == 3):
                 self.vars_spatial_1d.append(k)
         # Setup the unit registry
         self.ureg = UnitRegistry()
         # Was a root directory specified?
-        self.root_dir = self.config['root_dir'] if 'root_dir' in self.config else ''
+        self.root_dir = self.config['root_dir'] \
+            if 'root_dir' in self.config else ''
         # Create empty dict ready to save vars to memory
         self.saved_vars = {}
         # Do we want to compile or edit a dataset?
@@ -101,22 +112,29 @@ class Compiler:
         # Define the timestep as a unit, based on that given in config file
         self.ureg.define('timestep = {0} * second'.format(self.config['time']['dt']))
         # Set grid properties from the NetCDF file
-        self.grid = rasterio.open(f'netcdf:{self.config["output_nc_file"]}:flow_dir')
+        self.grid = extract_raster_info(f'netcdf:{self.config["output_nc_file"]}:flow_dir',
+                                        masked=True)
         self.grid_crs = CRS.from_wkt(self.nc['crs'].crs_wkt)
         self.grid_bbox = box(*self.grid.bounds)
-        self.grid_mask = np.ma.getmask(self.grid.read(1, masked=True))
+        self.grid_mask = np.ma.getmask(self.grid.arr)
 
     def create(self):
-        """Compile the data from multiple input files, specified in the config file, to a NetCDF dataset for
-        spatial and/or temporal data, and a Fortran namelist file for constants."""
+        """
+        Compile the data from multiple input files, specified in the config
+        file, to a NetCDF dataset for spatial and/or temporal data, and a
+        Fortran namelist file for constants.
+        """
         # Set up the dataset:
-        # - Use flow direction raster to give us the grid system (bounds, CRS, etc)
+        # - Use flow direction or DEM rasters to give us the grid system (bounds, CRS, etc)
         # - Creating the NetCDF dataset with this grid information
         # - Set tidal bounds in lieu of routing
         # - Route water bodies (e.g. set inflows, outflows and headwaters) using the flow direction and tidal bounds
-        print('Setting up dataset...\n'
-              '\t...parsing flow_dir')
-        self.parse_flow_dir()
+        print('Setting up dataset...\n')
+        # Set the flow direction either directly from the flow_dir raster,
+        # or if not provided, calculated from the dem raster. Also set the
+        # grid, CRS and grid bounding box
+        self.flow_dir, self.grid = self.parse_flow_dir()
+        print('\t...parsing flow_dir')
         print('\t...creating NetCDF file')
         self.setup_netcdf_dataset()
         self.parse_spatial_var('is_estuary', save=True)
@@ -184,72 +202,126 @@ class Compiler:
                         = os.path.join(self.config['root_dir'], path.split('<root_dir>')[1])
 
     def parse_flow_dir(self):
-        """Parse the flow direction variable and define the CRS and mask for other
-        variables based on this. This step must be performed before NetCDF file is generated."""
-        if self.config['flow_dir']['type'] == 'raster':
-            self.grid = rasterio.open(self.config['flow_dir']['path'].replace('<root_dir>', self.config['root_dir']))
-            # If a CRS has been specified for the flowdir raster, use this instead of the raster's
+        """
+        Either use the provided flow_dir raster, or if not presented,
+        calculate flow direction from the provided digital elevation
+        model
+        """
+        # If flow_dir has been provided, use it directly. Else, check that
+        # a dem has been provided and calculate the flow direction from that.
+        # If neither provided, error out
+        if 'flow_dir' in self.config:
+            # Open and read the flow_dir raster
+            flow_dir_path = self.config['flow_dir']['path'] \
+                .replace('<root_dir>', self.config['root_dir'])
+            grid = extract_raster_info(flow_dir_path, masked=True)
+            flow_dir = grid.arr
+            # If a CRS has been specified, use this instead of the raster's
             # internal CRS. This is useful if a raster has an ill-defined CRS
-            self.grid_crs = CRS.from_user_input(self.config['flow_dir']['crs']) \
-                if 'crs' in self.config['flow_dir'] else self.grid.crs
-            # Only projected rasters allowed for the moment
-            if self.grid_crs.is_geographic:
-                print('Sorry, the flow_dir raster must be projected, not geographic. '
-                      'I got a geographic CRS: \'{0}\'.'.format(self.grid_crs))
-                sys.exit()
-            # Create Shapely box from bounds, to clip other rasters to
-            self.grid_bbox = box(*self.grid.bounds)
-        elif self.config['flow_dir']['type'] == 'csv':
-            print("CSV flow direction not currently supported. Please provide a raster file.")
-            sys.exit()
-
+            if 'crs' in self.config['flow_dir']:
+                grid.crs = CRS.from_user_input(self.config['flow_dir']['crs'])
+            # Set a flag to say we used the flow_dir raster
+            var_used = 'flow_dir'
+        elif 'dem' in self.config:
+            # Path to the DEM
+            dem_path = Path(
+                self.config['dem']['path'].replace('<root_dir>',
+                                                   self.config['root_dir'])
+            ).resolve()
+            # If we have been asked to condition the DEM, do this first
+            if 'condition_dem' in self.config:
+                dem_path_conditioned, temp_dir = routing.condition_dem(
+                    dem_path,
+                    self.config['condition_dem'],
+                )
+            else:
+                # If no conditioning is required, just use the original DEM
+                dem_path_conditioned = dem_path
+            # Open and read the DEM, unmasked so that the flow direction
+            # calculation routine can use the nodata values
+            grid = extract_raster_info(dem_path_conditioned, masked=False)
+            # Clean the temporary directory if we created one
+            if temp_dir is not None:
+                temp_dir.cleanup()
+            # If a CRS has been specified, use this instead of the raster's
+            # internal CRS. This is useful if a raster has an ill-defined CRS
+            if 'crs' in self.config['dem']:
+                grid.crs = CRS.from_user_input(self.config['dem']['crs'])
+            # Calculate the flow direction, which returns a masked array
+            flow_dir = routing.calculate_flow_dir(dem=grid.arr,
+                                                  res=grid.res,
+                                                  nodata=grid.nodata)
+            # Set a flag to say we used the dem raster from grid properties
+            var_used = 'dem'
         else:
-            print("Unrecognised file type {0} for variable flow_dir. "
-                  "Type should be rs, csv or nc.".format(self.config['flow_dir']['type']))
-            sys.exit()
+            # If neither flow_dir or dem have been provided, error out
+            raise Exception('Neither `flow_dir` nor `dem` rasters '
+                            'have been provided.')
+        # Only projected rasters allowed for the moment
+        if grid.crs.is_geographic:
+            raise Exception(f'The `{var_used}` raster must be projected, '
+                            'not geographic. I got a geographic CRS: '
+                            f'{grid.crs}')
+        return flow_dir, grid
 
     def setup_netcdf_dataset(self):
-        """Create NetCDF file, add required dimensions, coordinate variables and the
-        flow direction variable."""
-        self.nc = Dataset(self.config['output']['nc_file'], 'w', format='NETCDF4')
+        """
+        Create NetCDF file, add required dimensions, coordinate variables
+        and the flow direction variable.
+        """
+        self.nc = Dataset(self.config['output']['nc_file'], 'w',
+                          format='NETCDF4')
         self.nc.title = "Input data for NanoFASE model"
-        self.nc.nanomaterial = self.config['nanomaterial'] if 'nanomaterial' in self.config else 'Unknown'
+        self.nc.nanomaterial = self.config['nanomaterial'] \
+            if 'nanomaterial' in self.config else 'Unknown'
         self.nc.Conventions = 'CF-1.8'
-        self.nc.coordinates = 'crs'                             # Needed for xarray to distinguish crs as a coordinate rather than a variable
+        # Setting the coordinates attribute is needed for xarray to
+        # distinguish crs as a coordinate rather than a variable
+        self.nc.coordinates = 'crs'
         crs_var = self.nc.createVariable('crs', 'i4')
-        crs_var.spatial_ref = self.grid_crs.to_wkt()            # QGIS/ArcGIS recognises spatial_ref to define CRS
-        crs_var.crs_wkt = self.grid_crs.to_wkt()                # Latest CF conventions say crs_wkt can be used
-        crs_var.epsg_code = self.grid_crs.to_epsg()             # Not standardised, but might be useful instead of having to decipher WKT
+        # QGIS/ArcGIS recognises spatial_ref to define the CRS
+        crs_var.spatial_ref = self.grid.crs.to_wkt()
+        # Latest CF conventions say crs_wkt can be used
+        crs_var.crs_wkt = self.grid.crs.to_wkt()
+        # Not standardised, but might be useful instead of having to
+        # decipher WKT
+        crs_var.epsg_code = self.grid.crs.to_epsg()
         # Time dimensions and coordinate variable
-        t_dim = self.nc.createDimension('t', None)
+        _ = self.nc.createDimension('t', None)
         t = self.nc.createVariable('t', 'i4', ('t',))
-        t.units = "seconds since {0} 00:00:00".format(self.config['time']['start_date'])
+        t.units = "seconds since {0} 00:00:00" \
+            .format(self.config['time']['start_date'])
         t.standard_name = 'time'
         t.calendar = 'gregorian'
-        t[:] = [i*int(self.config['time']['dt']) for i in range(int(self.config['time']['n']))]
+        t[:] = [i*int(self.config['time']['dt'])
+                for i in range(int(self.config['time']['n']))]
         # x dimension and coordinate variable
-        x_dim = self.nc.createDimension('x', self.grid.width)
+        _ = self.nc.createDimension('x', self.grid.width)
         x = self.nc.createVariable('x', 'f4', ('x',))
         x.units = 'm'
         x.standard_name = 'projection_x_coordinate'
         x.axis = 'X'
-        x[:] = [self.grid.bounds.left + i * self.grid.res[0] + 0.5 * self.grid.res[0] for i in range(self.grid.width)]
+        x[:] = [self.grid.bounds.left + i * self.grid.res[0]
+                + 0.5 * self.grid.res[0] for i in range(self.grid.width)]
         # y dimension and coordinate variable
-        y_dim = self.nc.createDimension('y', self.grid.height)
+        _ = self.nc.createDimension('y', self.grid.height)
         y = self.nc.createVariable('y', 'f4', ('y',))
         y.units = 'm'
         y.standard_name = 'projection_y_coordinate'
         y.axis = 'Y'
-        y[:] = [self.grid.bounds.top - i * self.grid.res[1] - 0.5 * self.grid.res[1] for i in range(self.grid.height)]
-        # Grid dimension (2D), max number of waterbodies per cell and grid bounds
-        d_dim = self.nc.createDimension('d', 2)
-        w_dim = self.nc.createDimension('w', 7)
-        box_dim = self.nc.createDimension('box', 4)
-        p_dim = self.nc.createDimension('p')
+        y[:] = [self.grid.bounds.top - i * self.grid.res[1]
+                - 0.5 * self.grid.res[1] for i in range(self.grid.height)]
+        # Other useful dimensions (grid, number of waterbodies, bounding
+        # box and point sources)
+        _ = self.nc.createDimension('d', 2)
+        _ = self.nc.createDimension('w', 7)
+        _ = self.nc.createDimension('box', 4)
+        _ = self.nc.createDimension('p')
         # Grid properties - shape
         grid_shape = self.nc.createVariable('grid_shape', 'i4', ('d',))
         grid_shape.units = ''
-        grid_shape.long_name = 'number of grid cells along each (x,y) grid axis'
+        grid_shape.long_name = \
+            'number of grid cells along each (x,y) grid axis'
         grid_shape[:] = self.grid.shape[::-1]
         # Grid resolution
         grid_res = self.nc.createVariable('grid_res', 'f4', ('d',))
@@ -261,9 +333,8 @@ class Compiler:
         grid_bounds.units = ''
         grid_bounds.long_name = 'bounding box of the grid'
         grid_bounds[:] = self.grid.bounds
-        # Get the flow direction array from the raster
-        self.flow_dir = self.grid.read(1, masked=True)
-        # Use the extent of the flow direction array to create a mask for all other data
+        # Use the extent of the flow direction array to create a mask for all
+        # other data
         self.grid_mask = np.ma.getmask(self.flow_dir)
         nc_var = self.nc.createVariable('flow_dir', 'i4', ('y', 'x'))
         nc_var.long_name = 'flow direction of water in grid cell'
@@ -329,7 +400,7 @@ class Compiler:
         with rasterio.open(path) as rs:
             out_bounds = rs.bounds
             # Mask clips the raster to the grid bounding box or smaller
-            out_img, _ = mask(rs, [self.grid_bbox], crop=True, filled=False)
+            out_img, _ = mask(rs, [self.grid.bbox], crop=True, filled=False)
         # The raster might still be smaller than the grid box, so let's check and 
         # construct mask array of the correct size if so
         if out_bounds.left > self.grid.bounds.left or out_bounds.right < self.grid.bounds.right \
@@ -459,8 +530,11 @@ class Compiler:
                     # Add this time step to the NetCDF file as a masked array
                     nc_var[t-1, :, :] = values.magnitude
             else:
-                print("Spatiotemporal variable ({0}) in raster format must be provided by one raster file per time step, with path denoted by /{t/}".format(var_name))
-           
+                raise Exception(f'Spatiotemporal variable ({var_name}) in '
+                                'raster format must be provided by one raster '
+                                'file per time step, with the time step '
+                                'index denoted by /{t/} in the path.')
+
         elif var_dict['type'] == 'csv':
             df = pd.read_csv(var_dict['path'], header=0)
             # Loop through the timesteps and create pivot table to obtain spatial array for each
@@ -470,8 +544,10 @@ class Compiler:
                 values = np.ma.masked_where(self.grid_mask, pt.values)
                 # Check the pivot table's shape is that of the grid we're using
                 if pt.shape != self.grid.shape:
-                    print("Inconsistent shape between {0} csv file and overall grid system ({1} and {2}). Check indices set correctly.".format(var, pt.shape, grid.shape))
-                    sys.exit()
+                    raise Exception(f'Inconsistent shape between {var} csv '
+                                    f'file and overall grid system '
+                                    f'({pt.shape} and {self.grid.shape}). '
+                                    'Check indices are set correctly.')
                 # Should the array be clipped?
                 if 'clip' in var_dict:
                     try:
@@ -561,7 +637,7 @@ class Compiler:
         # Open the supplied land use raster
         # Open the raster and clip to extent of grid (defined by flowdir raster)
         with rasterio.open(self.vars['land_use']['path']) as rs:
-            out_img, out_transform = mask(rs, [self.grid_bbox], crop=True, filled=False)
+            out_img, out_transform = mask(rs, [self.grid.bbox], crop=True, filled=False)
             src_arr = out_img[0]
 
             # Prepare a dict of lists (empty for the moment) to store arrays to be summed
@@ -597,7 +673,7 @@ class Compiler:
             for name, old_arr in nf_final_arrs.items():
                 if old_arr.shape == src_arr.shape:
                     # Re-read flowdir as reproject fills new_arr
-                    new_arr = self.grid.read(1, masked=True)
+                    new_arr = self.flow_dir.copy()
                     # Reproject. Remember we're not converting CRS here, so clc_rs
                     # CRS can be used as src and dst
                     reproject(
@@ -623,52 +699,75 @@ class Compiler:
     def routing(self):
         """Use the flow direction to route the waterbody network."""
         # Create the empty arrays to begin with a mask ready to be filled
-        router = Router(self.flow_dir)              # Set up the router to deal with common routing tasks
-        outflow_arr = np.ma.zeros((*self.flow_dir.shape, 2), dtype=np.dtype('i2'))
-        outflow_arr.mask = self.grid_mask           # Set the grid mask
-        inflows_arr = np.ma.array(np.ma.empty((*self.flow_dir.shape, 7, 2), dtype=np.dtype('i2')), mask=True)      # Max of seven inflows
-        n_waterbodies = np.ma.array(np.ma.empty(self.flow_dir.shape, dtype=np.dtype('i2')), mask=True)
-        is_headwater = np.ma.array(np.ma.empty(self.flow_dir.shape, dtype=np.dtype('u1')), mask=True)
+        outflow_arr = np.ma.zeros((*self.flow_dir.shape, 2),
+                                  dtype=np.dtype('i2'))
+        # Set the grid mask
+        outflow_arr.mask = self.grid_mask
+        # Create the relevant empty arrays
+        inflows_arr = np.ma.array(np.ma.empty((*self.flow_dir.shape, 7, 2),
+                                              dtype=np.dtype('i2')),
+                                  mask=True)
+        n_waterbodies = np.ma.array(np.ma.empty(self.flow_dir.shape,
+                                                dtype=np.dtype('i2')),
+                                    mask=True)
+        is_headwater = np.ma.array(np.ma.empty(self.flow_dir.shape,
+                                               dtype=np.dtype('u1')),
+                                   mask=True)
 
         # Use the flow direction to set outflow and inflows to each cell
-        for index, cell in np.ndenumerate(self.flow_dir):
+        for index, _ in np.ndenumerate(self.flow_dir):
             y, x = index[0] + 1, index[1] + 1
-            if not self.grid_mask[index]:       # Only for non-masked elements
-                outflow_arr[index] = router.outflow_from_flow_dir(x, y)
-                inflows_arr[index] = router.inflows_from_flow_dir(x, y)
-                n_waterbodies[index], is_headwater[index]\
-                    = router.n_waterbodies_from_inflows(x, y, outflow_arr[index], inflows_arr[index])
-        
+            # Only for non-masked elements
+            if not self.grid_mask[index]:
+                outflow_arr[index] = \
+                    routing.outflow_from_flow_dir(self.flow_dir, x, y)
+                inflows_arr[index] = \
+                    routing.inflows_from_flow_dir(self.flow_dir, x, y)
+                n_waterbodies[index], is_headwater[index] = \
+                    routing.n_waterbodies_from_inflows(self.flow_dir,
+                                                       outflow_arr[index],
+                                                       inflows_arr[index])
+
         # Create NetCDF vars for these arrays. Firstly, outflow
-        nc_var = self.nc.createVariable('outflow', np.dtype('i2'), ('y', 'x', 'd'))
+        nc_var = self.nc.createVariable('outflow',
+                                        np.dtype('i2'),
+                                        ('y', 'x', 'd'))
         nc_var.long_name = 'index of grid cell outflow'
         nc_var.units = ''
         nc_var.grid_mapping = 'crs'
         nc_var[:] = outflow_arr
 
         # Inflows
-        nc_var = self.nc.createVariable('inflows', np.dtype('i2'), ('y', 'x', 'w', 'd'))
+        nc_var = self.nc.createVariable('inflows',
+                                        np.dtype('i2'),
+                                        ('y', 'x', 'w', 'd'))
         nc_var.long_name = 'indices of grid cell inflows'
         nc_var.units = ''
         nc_var.grid_mapping = 'crs'
         nc_var[:] = inflows_arr
 
         # Number of waterbodies per cell
-        nc_var = self.nc.createVariable('n_waterbodies', np.dtype('i2'), ('y', 'x'))
+        nc_var = self.nc.createVariable('n_waterbodies',
+                                        np.dtype('i2'),
+                                        ('y', 'x'))
         nc_var.long_name = 'number of waterbodies in grid cell'
         nc_var.units = ''
         nc_var.grid_mapping = 'crs'
         nc_var[:] = n_waterbodies
 
         # Is cell a headwater?
-        nc_var = self.nc.createVariable('is_headwater', np.dtype('u1'), ('y', 'x'))
+        nc_var = self.nc.createVariable('is_headwater',
+                                        np.dtype('u1'),
+                                        ('y', 'x'))
         nc_var.long_name = 'is this cell a headwater?'
         nc_var.units = ''
         nc_var.grid_mapping = 'crs'
         nc_var[:] = is_headwater
 
     def in_model_domain(self, point):
-        """Check if a point is in the model domain."""
+        """
+        Check if a point is in the model domain.
+        """
         if (point.x >= self.grid.bounds.left)\
                 and (point.x < self.grid.bounds.right)\
                 and (point.y > self.grid.bounds.bottom)\
@@ -676,7 +775,7 @@ class Compiler:
             return True
         else:
             return False
-    
+
     @staticmethod
     def create_constants(constants_yaml, output_path):
         """Method to just create a constants file, bypassing

{nfdata-0.2.0 → nfdata-0.3.0}/src/nfdata/console.py

@@ -1,9 +1,11 @@
 #!/usr/bin/env python
-"""NanoFASE data is responsible for compiling and editing data for use in NanoFASE model."""
-import sys
-import os
+"""
+NanoFASE data is responsible for compiling and editing data for use in
+the NanoFASE model.
+"""
 import argparse
 import importlib.resources as pkg_resources
+
 from .compiler import Compiler
 
 

nfdata-0.3.0/src/nfdata/routing.py

@@ -0,0 +1,292 @@
+"""
+Utilities relating to routing water bodies, such as creating
+a flow direction grid from a digital elevation model
+"""
+import os
+import tempfile
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+from whitebox.whitebox_tools import WhiteboxTools
+
+from ._errors import RoutingError, whitebox_callback
+
+
+def in_model_domain(flow_dir, i, j):
+    """
+    Check if index [j,i] is in the model domain.
+    """
+    i_in_domain = 0 <= i < flow_dir.mask.shape[1]
+    j_in_domain = 0 <= j < flow_dir.mask.shape[0]
+    not_masked = flow_dir[j, i] is not np.ma.masked \
+        if (i_in_domain and j_in_domain) else False
+    return i_in_domain and j_in_domain and not_masked
+
+
+def outflow_from_flow_dir(flow_dir, x, y):
+    """
+    Get the outflow cell reference given the current cell
+    reference and a flow direction.
+    """
+    xy_out = {
+        1: [x+1, y],
+        2: [x+1, y+1],
+        4: [x, y+1],
+        8: [x-1, y+1],
+        16: [x-1, y],
+        32: [x-1, y-1],
+        64: [x, y-1],
+        128: [x+1, y-1]
+    }
+    return xy_out[flow_dir[y-1, x-1]]
+
+
+def inflows_from_flow_dir(flow_dir, x, y):
+    """Get the inflow cell references given the current cell
+    reference and a flow direction."""
+    j, i = y - 1, x - 1
+    # Flow direction that means the cell with indices inflows[j_n-j, i_n-i]
+    # flows into this cell
+    inflow_flow_dir = {
+        (-1, -1): 2,
+        (0, -1): 1,
+        (1, -1): 128,
+        (-1, 0): 4,
+        (0, 0): 0,
+        (1, 0): 64,
+        (-1, 1): 8,
+        (0, 1): 16,
+        (1, 1): 32
+    }
+    inflow_cells = []
+    # Loop through the neighbours and check if they're inflows to this cell
+    for j_n in range(j-1, j+2):
+        for i_n in range(i-1, i+2):
+            if in_model_domain(flow_dir, i_n, j_n):
+                if flow_dir[j_n, i_n] == inflow_flow_dir[(j_n-j, i_n-i)]:
+                    inflow_cells.append([i_n+1, j_n+1])
+    # Create masked array from the inflow_cells list
+    inflow_cells_ma = np.ma.array(np.ma.empty((7, 2), dtype=int), mask=True)
+    if len(inflow_cells) > 0:       # Only fill if there are inflows
+        inflow_cells_ma[0:len(inflow_cells)] = inflow_cells
+    return inflow_cells_ma
+
+
+def n_waterbodies_from_inflows(flow_dir, outflow, inflows):
+    """
+    Calculate the number of waterbodies from the inflows to the cell.
+    """
+    j_out, i_out = outflow[1] - 1, outflow[0] - 1
+    # Count the unmasked elements to get n_inflows
+    n_inflows = inflows.count(axis=0)[0]
+    # If there are no inflows but the outflow is to the model domain, it
+    # must be a headwater. Else, number of waterbodies is same as number of
+    # inflows
+    if n_inflows == 0 and in_model_domain(flow_dir, i_out, j_out):
+        n_waterbodies = 1
+        is_headwater = 1
+    else:
+        n_waterbodies = n_inflows
+        is_headwater = 0
+    return n_waterbodies, is_headwater
+
+
+def calculate_flow_dir(dem,
+                       res,
+                       nodata: Optional[float] = np.nan,
+                       treat_nodata_as_drain=True,
+                       attempt_to_resolve=False,
+                       lonely_cell_flowdir=1,
+                       error_if_unresolved=True):
+    """
+    Calculate the D8 flow direction from a digital elevation model
+
+    Parameters
+    ----------
+    dem : array-like with dtype=float
+        The digital elevation model to calculate flow direction from
+    res : length-2 tuple
+        The (y, x) resolution of the DEM, used to calculate distances
+        for non-square grid cells
+    nodata : float, default=np.nan
+        The value in `dem` that should be treated as nodata. Defaults
+        to NaN
+    treat_nodata_as_drain : bool, default=True
+        Should nodata cells be treated as a drain, i.e. with lower
+        elevation than all other cells? If so, this forces the flow
+        direction towards nodata cells by setting their elevation to
+        -9999. This is particularly useful where nodata represents
+        waterbodies. The `dem` is padded with one pixel of nodata on
+        each side in order to calculate flow directions from edge cells,
+        and so this flag also affects whether water drains from `dem`
+        boundaries. If this is set to False, you may end up with
+        unresolved cells
+    attempt_to_resolve : bool, default=False
+        If the resulting flow direction raster has unresolved cells
+        (flow direction = 0), setting `attempt_to_resolve` to True uses
+        same flow direction as that from the neighbouring cell with the
+        steepest sloe to this cell, effectively presuming the slope carries
+        on in the same direction. Warning: this might result in
+        neighbouring cells that flow into each other, effectively creating
+        "pools" from which flows can't drain
+    lonely_cell_flowdir : int or np.nan, default=1
+        What flow direction to assign to individual cells surrounded by
+        nodata. By default, flow direction is assumed to be eastwards (1).
+        If `nodata` or `np.nan`, the cell is removed from the output
+        raster. Note that setting to 0 and specifying
+        `error_if_unresolved=True` will result in an error if your `dem`
+        has lonely cells. There are no further checks on the validity of
+        this value, so you can use this to flag lonely cells by setting to
+        an arbitrary value
+    error_if_unresolved : bool, default=True
+        Whether to raise an error if unresolved (flow direction = 0) cells
+        are encountered
+
+    Returns
+    -------
+    np.ma.masked_array(np.int16)
+        Masked flow direction array
+    """
+
+    # Mappings between D8 codes and directions
+    directions = np.array([
+        (-1, 1),    # 128 - NE
+        (-1, 0),    # 64  - N
+        (-1, -1),   # 32  - NW
+        (0, -1),    # 16  - W
+        (1, -1),    # 8   - SW
+        (1, 0),     # 4   - S
+        (1, 1),     # 2   - SE
+        (0, 1),     # 1   - E
+    ])
+    d8_codes = np.array([128, 64, 32, 16, 8, 4, 2, 1])
+    d8_codes_opposite = np.roll(d8_codes, 4)
+    # To make processing easier, we convert nodata values to NaNs for
+    # the DEM and -32768 for the flowdir array
+    dem = np.where(dem == nodata, np.nan, dem)
+    if (lonely_cell_flowdir == nodata) or np.isnan(lonely_cell_flowdir):
+        lonely_cell_flowdir = -32768
+    # In case there are data right to the edges, we pad with one pixel
+    # of NaNs
+    dem_padded = np.pad(dem, ((1, 1), (1, 1)), constant_values=np.nan)
+    # Get the number of (padded) rows and cols
+    nrows, ncols = dem_padded.shape
+    # Create a new array to store the (unpadded) flowdir in
+    flow_dir = np.full_like(dem, -32768, dtype=np.int16)
+    # If `treat_nodata_as_drain` is True, we presume that nodata cells
+    # (including the padded region) are water and should be at some
+    # negative elevation to force flow to drain into them. If False,
+    # they are simply ignored in the flow direction calculation
+    dem_adj = np.nan_to_num(dem_padded, nan=-np.inf) \
+        if treat_nodata_as_drain else dem_padded
+
+    # Loop over the cells (except the pad)
+    for row in range(1, nrows - 1):
+        for col in range(1, ncols - 1):
+            # Get the value for this cell
+            this_elevation = dem_padded[row, col]
+            # If this cell is nodata, then we're not interested in the flow
+            # direction
+            if np.isnan(this_elevation):
+                continue
+            # Init vars to be update in the loop through neighbouring cells
+            max_slope = -np.inf
+            min_slope = 0
+            heighest_neighbour_direction = 0
+            best_direction = 0
+            lonely_cell = True
+            # Loop through neighbouring cells
+            for i, (dr, dc) in enumerate(directions):
+                # Get the elevation of this neighbouring cell
+                neighbour_elevation = dem_adj[row + dr, col + dc]
+                # If the neighbouring cell is nodata and we're not treating
+                # nodata as drains, then ignore the influence of this cell
+                if np.isnan(neighbour_elevation):
+                    continue
+                # If neighbouring cell has data, then set the lonely_cell
+                # flag to False to say the central cell isn't a single cell
+                # surrounded by nodata (including nodata cells treated as
+                # drains)
+                if neighbour_elevation != -np.inf:
+                    lonely_cell = False
+                # Difference in elevation and slope
+                dz = this_elevation - neighbour_elevation
+                distance = np.sqrt((res[0]*dr)**2 + (res[1]*dc)**2)
+                slope = dz / distance
+                # If slope is steeper than the current maximum, and
+                # it's positive (flow *away* from this cell), then
+                # set this as the best direction
+                if (slope > max_slope) and (slope > 0.0):
+                    max_slope = slope
+                    best_direction = d8_codes[i]
+                elif (slope < min_slope):
+                    # If it is a negative slope (this neighbour is higher
+                    # than the central cell), store this in case we need it
+                    # to try and resolve cells later
+                    min_slope = slope
+                    heighest_neighbour_direction = d8_codes_opposite[i]
+
+            # Fill the flow direction array. If no candidate has been
+            # found then the best_direction is 0 and the cell is unresolved
+            flow_dir_ = best_direction
+
+            # If this is a lonely cell, set it to the specified value
+            if lonely_cell:
+                flow_dir_ = lonely_cell_flowdir
+            # If we're meant to, try and resolve this cell and raise an
+            # error for unresolved cells
+            if (flow_dir_ == 0):
+                if attempt_to_resolve:
+                    flow_dir_ = heighest_neighbour_direction
+                if error_if_unresolved:
+                    neighbouring_dem = dem_adj[row-1:row+2, col-1:col+2]
+                    raise RoutingError(
+                        f"Could not resolve flow direction for cell ({col-1}, {row-1}). "
+                        f"DEM for surrounding cells:\n{neighbouring_dem}"
+                    )
+
+            # Set the best value for the flow direction
+            flow_dir[row - 1, col - 1] = flow_dir_
+
+    # Return the flow direction array as a masked array
+    return np.ma.masked_equal(flow_dir, -32768)
+
+
+def condition_dem(dem_path, config):
+    """
+    Condition a DEM using Whitebox to remove pits, depressions
+    and flat areas
+    """
+    # If we have been passed a dict of parameters, use these
+    if isinstance(config, dict):
+        # If a path to save the conditioned DEM has been provided,
+        # use this
+        if 'save_dem_to_path' in config:
+            # Get the path and remove it from the dict so that we
+            # don't pass it to the WhiteboxTools function
+            dem_path_conditioned = config.pop('save_dem_to_path')
+            dem_path_conditioned = Path(dem_path_conditioned).resolve()
+            temp_dir = None
+        else:
+            # If no path has been provided, create a temporary file
+            temp_dir = tempfile.TemporaryDirectory()
+            dem_path_conditioned = os.path.join(temp_dir.name,
+                                                'dem.tif')
+        # Use any remainning parameters in the dict as kwargs
+        kwargs = config
+        # Set dist=10000 if no value given
+        kwargs['dist'] = config.get('dist', 10000)
+        # Don't use any callback specified as we need to use our
+        # own to catch errors
+        if 'callback' in kwargs:
+            del kwargs['callback']
+
+    wbt = WhiteboxTools()
+    # Run the conditioning function and save the resulting DEM
+    # to the specified path
+    wbt.breach_depressions_least_cost(dem=dem_path,
+                                      output=dem_path_conditioned,
+                                      callback=whitebox_callback,
+                                      **kwargs)
+    return dem_path_conditioned, temp_dir

nfdata-0.2.0/src/nfdata/router.py

@@ -1,111 +0,0 @@
-import numpy as np
-
-
-class Router:
-    """The Router class is responsible for using the flow direction raster to route water bodies within the
-    model domain, setting variables required by the NanoFASE model, such as inflows, outflows and headwaters."""
-
-    def __init__(self, flow_dir):
-        """Initiliase the routing class with the flow direction array."""
-        self.flow_dir = flow_dir
-        self.grid_mask = self.flow_dir.mask
-
-    def outflow_from_flow_dir(self, x, y):
-        """Get the outflow cell reference given the current cell
-        reference and a flow direction."""
-        xy_out = {
-            1: [x+1, y],
-            2: [x+1, y+1],
-            4: [x, y+1],
-            8: [x-1, y+1],
-            16: [x-1, y],
-            32: [x-1, y-1],
-            64: [x, y-1],
-            128: [x+1, y-1]
-        }
-        return xy_out[self.flow_dir[y-1, x-1]]
-
-    def inflows_from_flow_dir(self, x, y):
-        """Get the inflow cell references given the current cell
-        reference and a flow direction."""
-        j, i = y - 1, x - 1
-        # Flow direction that means the cell with indices inflows[j_n-j, i_n-i]
-        # flows into this cell
-        inflow_flow_dir = {
-            (-1, -1): 2,
-            (0, -1): 1,
-            (1, -1): 128,
-            (-1, 0): 4,
-            (0, 0): 0,
-            (1, 0): 64,
-            (-1, 1): 8,
-            (0, 1): 16,
-            (1, 1): 32
-        }
-        inflow_cells = []
-        # Loop through the neighbours and check if they're inflows to this cell
-        for j_n in range(j-1, j+2):
-            for i_n in range(i-1, i+2):
-                if self.in_model_domain(i_n, j_n):
-                    if self.flow_dir[j_n, i_n] == inflow_flow_dir[(j_n-j, i_n-i)]:
-                        inflow_cells.append([i_n+1, j_n+1])
-        # Create masked array from the inflow_cells list
-        inflow_cells_ma = np.ma.array(np.ma.empty((7,2), dtype=int), mask=True)
-        if len(inflow_cells) > 0:       # Only fill if there are inflows
-            inflow_cells_ma[0:len(inflow_cells)] = inflow_cells
-        return inflow_cells_ma
-
-    def in_model_domain(self, i, j):
-        """Check if index [j,i] is in model domain."""
-        i_in_domain = 0 <= i < self.grid_mask.shape[1]
-        j_in_domain = 0 <= j < self.grid_mask.shape[0]
-        not_masked = self.flow_dir[j, i] is not np.ma.masked if (i_in_domain and j_in_domain) else False
-        return i_in_domain and j_in_domain and not_masked
-
-    def n_waterbodies_from_inflows(self, x, y, outflow, inflows):
-        """Calculate the number of waterbodies from the inflows to the cell."""
-        # j, i = y - 1, x - 1
-        j_out, i_out = outflow[1] - 1, outflow[0] - 1
-        n_inflows = inflows.count(axis=0)[0]        # Count the unmasked elements to get n_inflows
-        # If there are no inflows but the outflow is to the model domain, it
-        # must be a headwater. Else, number of waterbodies is same as number of inflows
-        if n_inflows == 0 and self.in_model_domain(i_out, j_out):
-            n_waterbodies = 1
-            is_headwater = 1
-        else:
-            n_waterbodies = n_inflows
-            is_headwater = 0
-        return n_waterbodies, is_headwater
-
-    # def generate_waterbody_code(self, x, y, outflow, inflows, is_estuary, is_headwater):
-    #     """Generates character code of the waterbodies of this cell, of the format
-    #     <type><index>(<inflow_point>:<outflow_point>/<n_waterbodies_along_branch>)...
-    #     where type = r,e,l,s (river, estuary, lake, sea), index is 1-indexed, inflow
-    #     and outflow_side are integers representing the geometrical point of inflow
-    #     (0 for centre, 1 for top-left, 2 for top centre, etc.) and n_waterbodies_along_branch
-    #     is the number of waterbodies between inflow and outflow."""
-    #     outflow_point = self.point_index(outflow[1]-x, outflow[0]-y)
-    #     wb_type = 'e' if is_estuary > 0 else 'r'
-    #     wb_char = ''
-    #     for i, inflow in enumerate(inflows):
-    #         if not inflow.mask.any():       # Only if this cell isn't masked
-    #             inflow_point = self.point_index(inflow[1]-x, inflow[0]-y)
-    #             wb_char = wb_char + '{0}{1}({2}:{3}/01)'.format(wb_type, i+1, inflow_point, outflow_point)
-    #     # If this is a headwater, there must be a waterbody with inflow at centre of cell
-    #     if is_headwater:
-    #         wb_char = '{0}1(0:{1}/01)'.format(wb_type, outflow_point)
-    #     return wb_char
-
-    # def point_index(self, di, dj):
-    #     point_index = {
-    #         (-1, -1): 1,
-    #         (0, -1): 8,
-    #         (1, -1): 7,
-    #         (-1, 0): 2,
-    #         (0, 0): 0,
-    #         (1, 0): 6,
-    #         (-1, 1): 3,
-    #         (0, 1): 4,
-    #         (1, 1): 5
-    #     }
-    #     return point_index[dj, di]