refactor soft edge generation into a function

src/ttmask/__init__.py

@@ -17,4 +17,3 @@
 from .cone import cone
 from .ellipsoid import ellipsoid
 
-

src/ttmask/cone.py

@@ -4,6 +4,7 @@
 from scipy.ndimage import distance_transform_edt
 import mrcfile
 from ._cli import cli
+from .soft_edge import add_soft_edge
 
 
 @cli.command(name='cone')
@@ -13,7 +14,7 @@ def cone(
     cone_base_diameter: float = typer.Option(...),
     soft_edge_width: int = typer.Option(0),
     pixel_size: float = typer.Option(...),
-    output: str = typer.Option("cone.mrc")
+    output: str = typer.Option("cone.mrc"),
 ):
     c = sidelength // 2
     center = np.array([c, c, c])
@@ -23,7 +24,7 @@ def cone(
     positions = np.indices([sidelength, sidelength, sidelength])
     positions = einops.rearrange(positions, 'zyx d h w -> d h w zyx')
 
-    centered = positions - center  #pixels relative to center point
+    centered = positions - center  # pixels relative to center point
     magnitudes = np.linalg.norm(centered, axis=-1)
 
     magnitudes = einops.rearrange(magnitudes, 'd h w -> d h w 1')
@@ -50,18 +51,12 @@ def cone(
     # mask[within_cone_height] = 1
     mask[np.logical_and(within_cone_height, within_cone_angle)] = 1
 
+    #   need to adjust the center of the hollow cone otherwise the cone thins towards the apex
+    # thickness will need to decrease towards the apex anyway - it's surely not possible / realistic?
+
     # Shift the mask in the z-axis by cone_height / 2
     z_shift = -int(cone_height / 2)
     mask = np.roll(mask, z_shift, axis=0)
+    mask = add_soft_edge(mask, soft_edge_width)
 
-
-    distance_from_edge = distance_transform_edt(mask == 0)
-    boundary_pixels = (distance_from_edge <= soft_edge_width) & (distance_from_edge != 0)
-    normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / soft_edge_width) * np.pi
-
-    mask[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5)
-
-
-    mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True)
-
-
+    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)

src/ttmask/cube.py

@@ -4,12 +4,13 @@
 from ._cli import cli
 from scipy.ndimage import distance_transform_edt
 import mrcfile
+from .soft_edge import add_soft_edge
 
 
 @cli.command(name='cube')
 def cube(
     sidelength: int = typer.Option(...),
-    cube_sidelength: float =typer.Option(...),
+    cube_sidelength: float = typer.Option(...),
     soft_edge_width: float = typer.Option(0),
     pixel_size: float = typer.Option(...),
     output: str = typer.Option("cube.mrc"),
@@ -32,15 +33,10 @@ def cube(
 
     in_cube = np.all(difference < np.array(cube_sidelength) / (pixel_size * 2), axis=-1)
     mask[in_cube] = 1
-
+    
     if wall_thickness != 0:
         within_hollowing = np.all(difference < ((np.array(cube_sidelength) / (pixel_size * 2)) - wall_thickness), axis=-1)
         mask[within_hollowing] = 0
-
-    distance_from_edge = distance_transform_edt(mask == 0)
-    boundary_pixels = (distance_from_edge <= soft_edge_width) & (distance_from_edge != 0)
-    normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / soft_edge_width) * np.pi
-
-    mask[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5)
-
-    mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True)
+
+    mask = add_soft_edge(mask, soft_edge_width)
+    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)

src/ttmask/cuboid.py

@@ -6,6 +6,7 @@
 from typing_extensions import Annotated
 from scipy.ndimage import distance_transform_edt
 import mrcfile
+from .soft_edge import add_soft_edge
 
 
 @cli.command(name='cuboid')
@@ -43,15 +44,12 @@ def cuboid(
     inside_cuboid = np.all(difference < (np.array(cuboid_sidelengths) / (2 * pixel_size)), axis=-1)
 
     mask[inside_cuboid] = 1
-
+       
     if wall_thickness != 0:
         within_hollowing = np.all(difference < ((np.array(cuboid_sidelengths) / (2 * pixel_size)) - wall_thickness), axis=-1)
         mask[within_hollowing] = 0
-
-    distance_from_edge = distance_transform_edt(mask == 0)
-    boundary_pixels = (distance_from_edge <= soft_edge_width) & (distance_from_edge != 0)
-    normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / soft_edge_width) * np.pi
-
-    mask[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5)
+
+    mask = add_soft_edge(mask, soft_edge_width)
 
     mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True)
+

src/ttmask/cylinder.py

@@ -4,6 +4,7 @@
 from scipy.ndimage import distance_transform_edt
 import mrcfile
 from ._cli import cli
+from .soft_edge import add_soft_edge
 
 
 @cli.command(name='cylinder')
@@ -40,11 +41,6 @@ def cylinder(
         within_xy_hollowing = xy_distance < ((cylinder_radius - wall_thickness) / pixel_size)
         mask[np.logical_and(within_z_hollowing, within_xy_hollowing)] = 0
 
-    distance_from_edge = distance_transform_edt(mask == 0)
-    boundary_pixels = (distance_from_edge <= soft_edge_width) & (distance_from_edge != 0)
-    normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / soft_edge_width) * np.pi
-
-    mask[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5)
-
-    mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True)
+    mask = add_soft_edge(mask, soft_edge_width)
 
+    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)

src/ttmask/ellipsoid.py

@@ -4,6 +4,7 @@
 import typer
 from scipy.ndimage import distance_transform_edt
 import mrcfile
+from .soft_edge import add_soft_edge
 
 
 @cli.command(name='ellipsoid')
@@ -37,18 +38,14 @@ def ellipsoid(
     x_axis_length = width / (2 * pixel_size)
 
     in_ellipsoid = (((x_magnitude) ** 2) / (x_axis_length ** 2)) + ((y_magnitude ** 2) / (y_axis_length ** 2)) + (
-            (z_magnitude ** 2) / (z_axis_length ** 2)) <= 1
+        (z_magnitude ** 2) / (z_axis_length ** 2)) <= 1
     mask[in_ellipsoid] = 1
-
+    
     if wall_thickness != 0:
         in_hollowing = (((x_magnitude) ** 2) / ((x_axis_length - wall_thickness) ** 2)) + ((y_magnitude ** 2) / ((y_axis_length - wall_thickness) ** 2)) + (
             (z_magnitude ** 2) / ((z_axis_length - wall_thickness) ** 2)) <= 1
         mask[in_hollowing] = 0
 
-    distance_from_edge = distance_transform_edt(mask == 0)
-    boundary_pixels = (distance_from_edge <= soft_edge_width) & (distance_from_edge != 0)
-    normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / soft_edge_width) * np.pi
-
-    mask[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5)
+    mask = add_soft_edge(mask, soft_edge_width)
 
     mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)

src/ttmask/soft_edge.py

@@ -0,0 +1,11 @@
+import numpy as np
+from scipy.ndimage import distance_transform_edt
+
+
+def add_soft_edge(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)
+    normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / width) * np.pi
+    output = np.copy(mask)
+    output[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5)
+    return output

src/ttmask/sphere.py

@@ -4,6 +4,7 @@
 import typer
 from scipy.ndimage import distance_transform_edt
 import mrcfile
+from .soft_edge import add_soft_edge
 
 
 @cli.command(name='sphere')
@@ -34,16 +35,11 @@ def sphere(
     print('calculating which pixels are in sphere')
     idx = distance < (sphere_radius / pixel_size)
     mask[idx] = 1
-
+    
     if wall_thickness != 0:
         within_hollowing = distance < ((sphere_radius - wall_thickness) / pixel_size)
         mask[within_hollowing] = 0
+
+    mask = add_soft_edge(mask, soft_edge_width)
 
-    distance_from_edge = distance_transform_edt(mask == 0)
-    boundary_pixels = (distance_from_edge <= soft_edge_width) & (distance_from_edge != 0)
-    normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / soft_edge_width) * np.pi
-
-    mask[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5)
-
-    mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True)
-
+    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)