refactor and add tests to urbanraster

src/GOSTurban/UrbanRaster.py

@@ -89,6 +89,38 @@ def geocode_cities(urban_extents):
     return urban_extents
 
 
+def _burnValue(mask, val, final_raster, allFeatures, idx, pop, cShape):
+    """Private function to burn value into final mask."""
+    mask = (mask ^ 1) * val
+    yy = np.dstack([final_raster, mask])
+    final_raster = np.amax(yy, axis=2)
+    allFeatures.append([idx, pop, val, shape(geojson.loads(json.dumps(cShape)))])
+    return final_raster, allFeatures
+
+
+def _create_urban_raster(data, comp_arr, comp_val):
+    """Private function to create urban_raster array."""
+    urban_raster = data * 0
+    urban_raster[np.where(data > comp_arr)] = comp_val
+    return urban_raster.astype("int16")
+
+
+def _smooth_urban_clusters(urban_raster):
+    """Private function to perform smoothing on an urban raster array."""
+
+    def modal(P):
+        mode = stats.mode(P)
+        if isinstance(mode.mode, float):
+            return mode.mode
+        else:
+            return mode.mode[0]
+
+    smooth_urban = generic_filter(urban_raster[0, :, :], modal, (3, 3))
+    yy = np.dstack([smooth_urban, urban_raster[0, :, :]])
+    urban_raster[0, :, :] = np.amax(yy, axis=2)
+    return urban_raster
+
+
 class urbanGriddedPop(object):
     def __init__(self, inRaster):
         """
@@ -115,14 +147,6 @@ def __init__(self, inRaster):
                 )
             )
 
-    def _burnValue(self, mask, val, final_raster, allFeatures, idx, pop, cShape):
-        """Private function to burn value into final mask."""
-        mask = (mask ^ 1) * val
-        yy = np.dstack([final_raster, mask])
-        final_raster = np.amax(yy, axis=2)
-        allFeatures.append([idx, pop, val, shape(geojson.loads(json.dumps(cShape)))])
-        return final_raster, allFeatures
-
     def calculateDegurba(
         self,
         urbDens=300,
@@ -170,32 +194,26 @@ def calculateDegurba(
             dictionary containing the final raster, the high density raster, the urban raster, and the shapes of the urban areas
 
         """
+        # read population data
         popRaster = self.inR
         data = popRaster.read()
-        urban_raster = data * 0
-        final_raster = data[0, :, :] * 0 + 11
 
-        urban_raster[np.where(data > hdDens)] = 30
-        idx = 0
-        urban_raster = urban_raster.astype("int16")
-        allFeatures = []
+        # create rasters to manipulate/populate
+        final_raster = data[0, :, :] * 0 + 11
+        urban_raster = _create_urban_raster(data, hdDens, 30)
 
+        # Smooth the HD urban clusters
         if verbose:
             tPrint(f"{print_message}: Smoothing Urban Clusters")
-
-        # Smooth the HD urban clusters
-        def modal(P):
-            mode = stats.mode(P)
-            return mode.mode[0]
-
-        smooth_urban = generic_filter(urban_raster[0, :, :], modal, (3, 3))
-        yy = np.dstack([smooth_urban, urban_raster[0, :, :]])
-        urban_raster[0, :, :] = np.amax(yy, axis=2)
+        urban_raster = _smooth_urban_clusters(urban_raster)
 
         # Analyze the high density shapes
         if verbose:
             tPrint(f"{print_message}: extracting HD clusters")
 
+        idx = 0  # init index at 0
+        allFeatures = []  # init empty list to hold features
+
         for cShape, value in features.shapes(
             urban_raster, transform=popRaster.transform
         ):
@@ -226,16 +244,14 @@ def modal(P):
                     val = 30
 
                 # Burn value into the final raster
-                final_raster, allFeatures = self._burnValue(
+                final_raster, allFeatures = _burnValue(
                     mask, val, final_raster, allFeatures, idx, pop, cShape
                 )
 
         HD_raster = final_raster
-
-        urban_raster = data * 0
+        urban_raster = _create_urban_raster(data, urbDens, 22)
         final_raster = data[0, :, :] * 0 + 11
-        urban_raster[np.where(data > urbDens)] = 22
-        urban_raster = urban_raster.astype("int16")
+
         # Analyze the high density shapes
         if verbose:
             tPrint(f"{print_message}: extracting URBAN clusters")
@@ -262,7 +278,7 @@ def modal(P):
                 if pop > urbThresh:
                     val = 21
                 # Burn value into the final raster
-                final_raster, allFeatures = self._burnValue(
+                final_raster, allFeatures = _burnValue(
                     mask, val, final_raster, allFeatures, idx, pop, cShape
                 )
 
@@ -372,6 +388,7 @@ def calculateUrban(
             geopandas dataframe of urban extents
 
         """
+        # read in urban data from the class attribute
         popRaster = self.inR
         data = popRaster.read()
         urbanData = (data > densVal) * 1

tests/test_UrbanRaster.py

@@ -71,10 +71,9 @@ def test_init_value_error(self):
     def test_burn_value(self):
         """Testing the private burn value function."""
         # make the object
-        ugp = UrbanRaster.urbanGriddedPop("str")
         with patch("shapely.geometry") as mock_shape:
             mock_shape = {"type": "Point", "coordinates": [0, 1]}
-            final_raster, allFeatures = ugp._burnValue(
+            final_raster, allFeatures = UrbanRaster._burnValue(
                 np.ones((5, 5), dtype=bool),
                 1,
                 np.zeros((5, 5)),
@@ -90,3 +89,25 @@ def test_burn_value(self):
         assert len(allFeatures) == 1
         assert len(allFeatures[0]) == 4
         assert allFeatures[0][0] == "a"
+
+
+def test_create_urban_raster():
+    """Testing the _create_urban_raster function."""
+    data = np.ones((5, 5))
+    comp_arr = np.zeros((5, 5))
+    urban_raster = UrbanRaster._create_urban_raster(data, comp_arr, 10)
+    assert urban_raster.shape == (5, 5)
+    assert urban_raster[0, 0] == 10
+    assert urban_raster.dtype == "int16"
+
+
+def test_smooth_urban_clusters():
+    """Testing the _smooth_urban_clusters function."""
+    arr = np.zeros((1, 10, 10))
+    arr[0, 3:7, 3] = 10.0
+    arr[0, 3:7, 4] = 20.0
+    arr[0, 3:7, 5] = 30.0
+    arr[0, 3:7, 6] = -10.0
+    old_arr = arr.copy()
+    arr = UrbanRaster._smooth_urban_clusters(arr)
+    assert old_arr.sum() != arr.sum()