Padding mask2map (#42)

* added padding

src/ttmask/add_padding.py

@@ -0,0 +1,10 @@
+import numpy as np
+from scipy.ndimage import distance_transform_edt
+
+
+def add_padding(mask: np.ndarray, width: float) -> np.ndarray:
+    distance_from_edge = distance_transform_edt(mask == 0)
+    boundary_pixels = (distance_from_edge <= width) & (distance_from_edge != 0)
+    output = np.copy(mask)
+    output[boundary_pixels] = 1
+    return output

src/ttmask/map2mask.py

@@ -5,6 +5,7 @@
 
 from ._cli import cli
 from .soft_edge import add_soft_edge
+from .add_padding import add_padding
 
 @cli.command(name='map2mask')
 def map2mask(
@@ -14,7 +15,7 @@ def map2mask(
     output_mask: Path = typer.Option(Path("mask.mrc")),
     pixel_size: float = typer.Option(...),
     soft_edge_width: int = typer.Option(0),
-
+    padding_width: int = typer.Option(0),
 ):
     with mrcfile.open(input_map) as mrc:
         map_data = np.array(mrc.data)
@@ -25,6 +26,7 @@ def map2mask(
     map_data[above_threshold] = 1
     map_data[below_threshold] = 0
 
-    mask = add_soft_edge(map_data, soft_edge_width)
+    padded_mask = add_padding(map_data, padding_width)
+    mask = add_soft_edge(padded_mask, soft_edge_width)
 
-    mrcfile.write(output_mask, mask, voxel_size=pixel_size, overwrite=True)
+    mrcfile.write(output_mask, mask, voxel_size=pixel_size, overwrite=True)