correct sample function

Author: noellehmcheng

map2loop/sampler.py

@@ -11,6 +11,7 @@
 import shapely
 import numpy
 from typing import Optional
+from osgeo import gdal
 
 
 class Sampler(ABC):
@@ -74,7 +75,7 @@ def __init__(self, decimation: int = 1):
 
     @beartype.beartype
     def sample(
-        self, spatial_data: geopandas.GeoDataFrame, dtm_data: Optional[geopandas.GeoDataFrame] = None, geology_data: Optional[geopandas.GeoDataFrame] = None
+        self, spatial_data: geopandas.GeoDataFrame, dtm_data: Optional[gdal.Dataset] = None, geology_data: Optional[geopandas.GeoDataFrame] = None
     ) -> pandas.DataFrame:
         """
         Execute sample method takes full point data, samples the data and returns the decimated points
@@ -88,10 +89,16 @@ def sample(
         data = spatial_data.copy()
         data["X"] = data.geometry.x
         data["Y"] = data.geometry.y
-        data["Z"] = set_z_values_from_raster_df(dtm_data, data)["Z"]
-        data["layerID"] = geopandas.sjoin(
-            data, geology_data, how='left'
-        )['index_right']
+        if dtm_data is not None:
+            data["Z"] = set_z_values_from_raster_df(dtm_data, data)["Z"]
+        else:
+            data["Z"] = None
+        if geology_data is not None:
+            data["layerID"] = geopandas.sjoin(
+                data, geology_data, how='left'
+            )['index_right']
+        else:
+            data["layerID"] = None
         data.reset_index(drop=True, inplace=True)
 
         return pandas.DataFrame(data[:: self.decimation].drop(columns="geometry"))