classify DamageDem

hhnk_threedi_tools/core/modelbuilder/create_calculation_rasters.py

@@ -1,17 +1,17 @@
 # %%
 """Rasters required to do postprocessing."""
 
+from dataclasses import dataclass
+from pathlib import Path
 from typing import Union
 
 import hhnk_research_tools as hrt
 
 from hhnk_threedi_tools import Folders
 
 
-def create_damage_dem(
-    folder: Folders,
-    dem: hrt.Raster,
-):
+@dataclass
+class DamageDem:
     """Create the damage dem. This dem differs from the original model dem in that the
     ground level of buildings are increased by an extra 10cm for a total of 15cm. This will ensure that
     no damage is calculated with low inundation depths. In the dem the floor
@@ -25,75 +25,100 @@ def create_damage_dem(
 
     Parameters
     ----------
-    folder : Folders
-        Model folders structure, including 'source_data' and 'model' sub-folders. To run this method
-        the following files should exist:
-         - folder.model.schema_base.rasters.dem tif
-         - folder.source_data.panden gpkg
-
     dem: hrt.Raster
         path to the original dem.
         default dem = folder.model.schema_base.rasters.dem
     """
 
-    damage_dem = folder.model.calculation_rasters.full_path(f"damage{dem.stem}_50cm.tif")
-    panden_gpkg = folder.source_data.panden
-    panden_raster = folder.model.calculation_rasters.panden
-
-    # Check if we should start creating the output
-    overwrite = hrt.check_create_new_file(
-        output_file=damage_dem,
-        input_files=[dem.base, panden_gpkg.base],  # If any of these are newer than output, update them.
-        overwrite=False,
-    )
-
-    if overwrite:
-        # create highres dem if it doesn't exist
-        if dem.metadata.x_res == 0.5:
-            highres_dem = dem
+    dem: hrt.Raster
+    panden_gpkg: Union[hrt.FileGDB, hrt.File, Path, str]
+    panden_raster: hrt.Raster
+    damage_dem: hrt.Raster
+
+    def __post_init__(self):
+        if self.dem.metadata.x_res == 0.5:
+            self.highres_dem = self.dem
         else:
-            highres_dem = folder.model.calculation_rasters.full_path(f"{dem.stem}_50cm.tif")
-
-        if not highres_dem.exists():
-            hrt.reproject(src=dem, target_res=0.5, output_path=highres_dem.path)
-
-        # create panden_raster if it doesn't exist
-        if not panden_raster.exists():
-            panden_gdf = panden_gpkg.load()
-            panden_gdf["base_height"] = 0.1
-            hrt.gdf_to_raster(
-                gdf=panden_gdf,
-                value_field="base_height",
-                raster_out=panden_raster,
-                nodata=0,
-                metadata=highres_dem.metadata,
-            )
+            self.highres_dem = hrt.Folder(self.damage_dem.parent).full_path(f"{self.dem.stem}_50cm.tif")
+
+    @classmethod
+    def from_folder(cls, folder: Folders, dem=None, panden_gpkg=None, panden_raster=None, damage_dem=None):
+        """Call from folder, if any of input is None take it from folder otherwise use the input.
+
+        Parameters
+        ----------
+        folder : Folders
+            Model folders structure, including 'source_data' and 'model' sub-folders. To run this method
+            the following files should exist:
+                - folder.model.schema_base.rasters.dem tif
+                - folder.source_data.panden gpkg
+        """
+        if dem is None:
+            dem = folder.model.schema_base.rasters.dem
+        if panden_gpkg is None:
+            panden_gpkg = folder.source_data.panden
+        if panden_raster is None:
+            panden_raster = folder.model.calculation_rasters.panden
+        if damage_dem is None:
+            damage_dem = folder.model.calculation_rasters.full_path(f"damage{dem.stem}_50cm.tif")
+
+        return cls(
+            dem=dem,
+            panden_gpkg=panden_gpkg,
+            panden_raster=panden_raster,
+            damage_dem=damage_dem,
+        )
 
-        # Create damage dem
-        def elevate_dem_block(block):
-            block_out = block.blocks["dem"] + block.blocks["panden"]
-
-            block_out[block.masks_all] = highres_dem.nodata
-            return block_out
-
-        calc = hrt.RasterCalculatorV2(
-            raster_out=damage_dem,
-            raster_paths_dict={
-                "dem": highres_dem,
-                "panden": panden_raster,
-            },
-            nodata_keys=["dem"],
-            mask_keys=["dem"],
-            metadata_key="dem",
-            custom_run_window_function=elevate_dem_block,
-            output_nodata=highres_dem.nodata,
-            min_block_size=4096,
-            verbose=True,
+    def create(self, overwrite=False):
+        # Check if we should start creating the output
+        overwrite = hrt.check_create_new_file(
+            output_file=self.damage_dem,
+            input_files=[self.dem.base, self.panden_gpkg.base],  # If any of these are newer than output, update them.
+            overwrite=overwrite,
         )
 
-        calc.run(overwrite=True)
-    else:
-        print(f"{damage_dem.view_name_with_parents(2)} already exists")
+        if overwrite:
+            # create panden_raster if it doesn't exist
+            if not self.panden_raster.exists():
+                panden_gdf = self.panden_gpkg.load()
+                panden_gdf["base_height"] = 0.1
+                hrt.gdf_to_raster(
+                    gdf=panden_gdf,
+                    value_field="base_height",
+                    raster_out=self.panden_raster,
+                    nodata=0,
+                    metadata=self.highres_dem.metadata,
+                )
+
+            # create highres dem if it doesn't exist
+            if not self.highres_dem.exists():
+                hrt.reproject(src=self.dem, target_res=0.5, output_path=self.highres_dem.path)
+
+            # Create damage dem
+            def elevate_dem_block(block):
+                block_out = block.blocks["dem"] + block.blocks["panden"]
+
+                block_out[block.masks_all] = self.highres_dem.nodata
+                return block_out
+
+            calc = hrt.RasterCalculatorV2(
+                raster_out=self.damage_dem,
+                raster_paths_dict={
+                    "dem": self.highres_dem,
+                    "panden": self.panden_raster,
+                },
+                nodata_keys=["dem"],
+                mask_keys=["dem"],
+                metadata_key="dem",
+                custom_run_window_function=elevate_dem_block,
+                output_nodata=self.highres_dem.nodata,
+                min_block_size=4096,
+                verbose=True,
+            )
+
+            calc.run(overwrite=True)
+        else:
+            print(f"{self.damage_dem.view_name_with_parents(2)} already exists")
 
 
 # %%