Merge pull request #188 from martinfleis/pycno_align

tobler/pycno/pycno.py

@@ -28,11 +28,6 @@
 def pycno(
     gdf, value_field, cellsize, r=0.2, handle_null=True, converge=3, verbose=True
 ):
-    try:
-        from astropy.convolution import convolve as astro_convolve
-    except (ImportError, ModuleNotFoundError):
-        raise ImportError("Pycnophylactic interpolation requires the astropy package")
-
     """Returns a smooth pycnophylactic interpolation raster for a given geodataframe
 
     Args:
@@ -115,6 +110,13 @@ def smooth2D(data):
 
     # The convolution function from astropy handles nulls.
     def astroSmooth2d(data):
+        try:
+            from astropy.convolution import convolve as astro_convolve
+        except (ImportError, ModuleNotFoundError) as err:
+            raise ImportError(
+                "Pycnophylactic interpolation with handle_null=True "
+                "requires the astropy package"
+            ) from err
         s1d = lambda s: astro_convolve(s, [0.5, 0, 0.5])
         # pad the data array with the mean value
         padarray = pad(data, 1, "constant", constant_values=nanmean(data))
@@ -125,7 +127,6 @@ def astroSmooth2d(data):
         return padarray[1:-1, 1:-1]
 
     def correct2Da(data):
-
         for idx, val in gdf[value_field].items():
             # Create zone mask from feature_array
             mask = masked_where(feature_array == idx, feature_array).mask
@@ -137,7 +138,6 @@ def correct2Da(data):
         return data
 
     def correct2Dm(data):
-
         for idx, val in gdf[value_field].items():
             # Create zone mask from feature_array
             mask = masked_where(feature_array == idx, feature_array).mask
@@ -186,7 +186,6 @@ def correct2Dm(data):
 
 
 def save_pycno(pycno_array, transform, crs, filestring, driver="GTiff"):
-
     """Saves a numpy array as a raster, largely a helper function for pycno
     Args:
         pycno_array (numpy array): 2D numpy array of pycnophylactic surface
@@ -218,7 +217,6 @@ def save_pycno(pycno_array, transform, crs, filestring, driver="GTiff"):
 
 
 def extract_values(pycno_array, gdf, transform, fieldname="Estimate"):
-
     """Extract raster value sums according to a provided polygon geodataframe
     Args:
         pycno_array (numpy array): 2D numpy array of pycnophylactic surface.
@@ -238,7 +236,7 @@ def extract_values(pycno_array, gdf, transform, fieldname="Estimate"):
             [geom], pycno_array.shape, transform=transform, invert=True
         )
         estimates.append(nansum(pycno_array[mask]))
-    out = pd.Series(estimates)
+    out = pd.Series(estimates, index=gdf.index)
     return out
 
 

tobler/tests/test_pycno.py

@@ -24,3 +24,11 @@ def test_pycno_interpolate():
         source_df=sac1, target_df=sac2, variables=["TOT_POP"], cellsize=500
     )
     assert_almost_equal(pyc.TOT_POP.sum(), 1794618.503, decimal=1)
+
+def test_custom_index():
+    sac1, sac2 = datasets()
+    sac2 = sac2.set_index("ZIP")
+    pyc = pycno_interpolate(
+        source_df=sac1, target_df=sac2, variables=["TOT_POP"], cellsize=500
+    )
+    assert_almost_equal(pyc.TOT_POP.sum(), 1794618.503, decimal=1)