From 503e3126428e61f56bf5214dde9993365364b38d Mon Sep 17 00:00:00 2001 From: milesagraham Date: Wed, 5 Jun 2024 15:14:09 +0100 Subject: [PATCH] hollowness for cube, cuboid, cylinder, ellipsoid and sphere (#27) --- src/ttmask/cube.py | 10 +++++++--- src/ttmask/cuboid.py | 5 +++++ src/ttmask/cylinder.py | 21 +++++++++++---------- src/ttmask/ellipsoid.py | 6 ++++++ src/ttmask/sphere.py | 5 +++++ 5 files changed, 34 insertions(+), 13 deletions(-) diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py index 148004f..16481f8 100644 --- a/src/ttmask/cube.py +++ b/src/ttmask/cube.py @@ -12,7 +12,8 @@ def cube( cube_sidelength: float =typer.Option(...), soft_edge_width: float = typer.Option(0), pixel_size: float = typer.Option(...), - output: str = typer.Option("cube.mrc") + output: str = typer.Option("cube.mrc"), + wall_thickness: float = typer.Option(0), ): c = sidelength // 2 center = np.array([c, c, c]) @@ -29,10 +30,13 @@ def cube( print('calculating distance') difference = np.abs(positions - center) # (100, 100, 100, 3) - in_cube = np.all(difference < (np.array(cube_sidelength / pixel_size) / 2), axis=-1) - + 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 diff --git a/src/ttmask/cuboid.py b/src/ttmask/cuboid.py index 59cfbf0..9904fe7 100644 --- a/src/ttmask/cuboid.py +++ b/src/ttmask/cuboid.py @@ -15,6 +15,7 @@ def cuboid( soft_edge_width: float = typer.Option(0), pixel_size: float = typer.Option(...), output: str = typer.Option("cuboid.mrc"), + wall_thickness: float = typer.Option(0), ): c = sidelength // 2 center = np.array([c, c, c]) @@ -43,6 +44,10 @@ def cuboid( 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 diff --git a/src/ttmask/cylinder.py b/src/ttmask/cylinder.py index 2be308d..8372a10 100644 --- a/src/ttmask/cylinder.py +++ b/src/ttmask/cylinder.py @@ -10,14 +10,13 @@ def cylinder( sidelength: int = typer.Option(...), cylinder_height: float = typer.Option(...), - cylinder_outer_diameter: float = typer.Option(...), - cylinder_inner_diameter: float = typer.Option(0), + cylinder_diameter: float = typer.Option(...), + wall_thickness: float = typer.Option(0), soft_edge_width: int = typer.Option(0), pixel_size: float = typer.Option(...), output: str = typer.Option("cylinder.mrc") ): - cylinder_outer_radius = cylinder_outer_diameter / 2 - cylinder_inner_radius = cylinder_inner_diameter / 2 + cylinder_radius = cylinder_diameter / 2 c = sidelength // 2 center = np.array([c, c, c]) @@ -29,15 +28,17 @@ def cylinder( print('calculating distance') difference = np.abs(positions - center) # (100, 100, 100, 3) - xy_distance = np.sum(difference[:, :, :, [1, 2]] ** 2, axis=-1) ** 0.5 + within_z = difference[:, :, :, 0] < (cylinder_height / (2 * pixel_size)) + within_xy = xy_distance < (cylinder_radius / pixel_size) + - idx_z = difference[:, :, :, 0] < (cylinder_height / (2 * pixel_size)) - idx_xy_outer = xy_distance < (cylinder_outer_radius / pixel_size) - idx_xy_inner = xy_distance < (cylinder_inner_radius / pixel_size) + mask[np.logical_and(within_z, within_xy)] = 1 - mask[np.logical_and(idx_z, idx_xy_outer)] = 1 - mask[np.logical_and(idx_z, idx_xy_inner)] = 0 + if wall_thickness != 0: + within_z_hollowing = difference[:, :, :, 0] < (cylinder_height - wall_thickness) / (2 * pixel_size) + 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) diff --git a/src/ttmask/ellipsoid.py b/src/ttmask/ellipsoid.py index 17456aa..6d84c62 100644 --- a/src/ttmask/ellipsoid.py +++ b/src/ttmask/ellipsoid.py @@ -15,6 +15,7 @@ def ellipsoid( soft_edge_width: int = typer.Option(0), pixel_size: float = typer.Option(...), output: str = typer.Option("ellipsoid.mrc"), + wall_thickness: float = typer.Option(0), ): c = sidelength // 2 center = np.array([c, c, c]) @@ -39,6 +40,11 @@ def ellipsoid( (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 diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py index 9fce1eb..e9a519f 100644 --- a/src/ttmask/sphere.py +++ b/src/ttmask/sphere.py @@ -13,6 +13,7 @@ def sphere( soft_edge_width: int = typer.Option(0), pixel_size: float = typer.Option(...), output: str = typer.Option("sphere.mrc"), + wall_thickness: float = typer.Option(0), ): sphere_radius = sphere_diameter / 2 c = sidelength // 2 @@ -34,6 +35,10 @@ def 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 + 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