From c957b6f474ee77783a928cb2858db414c435e1b2 Mon Sep 17 00:00:00 2001 From: Miles Graham Date: Wed, 5 Jun 2024 15:35:46 +0100 Subject: [PATCH 1/2] soft edge written as a function (refactored) --- src/ttmask/__init__.py | 1 - src/ttmask/cone.py | 13 ++++++------- src/ttmask/cube.py | 7 ++----- src/ttmask/cuboid.py | 7 ++----- src/ttmask/cylinder.py | 7 ++----- src/ttmask/ellipsoid.py | 7 ++----- src/ttmask/soft_edge.py | 9 +++++++++ src/ttmask/sphere.py | 7 ++----- 8 files changed, 25 insertions(+), 33 deletions(-) create mode 100644 src/ttmask/soft_edge.py diff --git a/src/ttmask/__init__.py b/src/ttmask/__init__.py index 5f0d6c4..e1e99b2 100644 --- a/src/ttmask/__init__.py +++ b/src/ttmask/__init__.py @@ -17,4 +17,3 @@ from .cone import cone from .ellipsoid import ellipsoid - diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py index 9f7bb11..9ec3a89 100644 --- a/src/ttmask/cone.py +++ b/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 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]) @@ -50,16 +51,14 @@ 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) - - 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) + soft_edge(mask, soft_edge_width) mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True) diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py index 148004f..6366b59 100644 --- a/src/ttmask/cube.py +++ b/src/ttmask/cube.py @@ -4,6 +4,7 @@ from ._cli import cli from scipy.ndimage import distance_transform_edt import mrcfile +from .soft_edge import soft_edge @cli.command(name='cube') @@ -33,10 +34,6 @@ def cube( mask[in_cube] = 1 - 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) + soft_edge(mask, soft_edge_width) mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True) \ No newline at end of file diff --git a/src/ttmask/cuboid.py b/src/ttmask/cuboid.py index 59cfbf0..63d8679 100644 --- a/src/ttmask/cuboid.py +++ b/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 soft_edge @cli.command(name='cuboid') @@ -43,10 +44,6 @@ def cuboid( mask[inside_cuboid] = 1 - 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) + soft_edge(mask, soft_edge_width) mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True) diff --git a/src/ttmask/cylinder.py b/src/ttmask/cylinder.py index 2be308d..99f1763 100644 --- a/src/ttmask/cylinder.py +++ b/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 soft_edge @cli.command(name='cylinder') @@ -39,11 +40,7 @@ def cylinder( mask[np.logical_and(idx_z, idx_xy_outer)] = 1 mask[np.logical_and(idx_z, idx_xy_inner)] = 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) + soft_edge(mask, soft_edge_width) mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True) diff --git a/src/ttmask/ellipsoid.py b/src/ttmask/ellipsoid.py index 17456aa..eab50e1 100644 --- a/src/ttmask/ellipsoid.py +++ b/src/ttmask/ellipsoid.py @@ -4,6 +4,7 @@ import typer from scipy.ndimage import distance_transform_edt import mrcfile +from .soft_edge import soft_edge @cli.command(name='ellipsoid') @@ -39,10 +40,6 @@ def ellipsoid( (z_magnitude ** 2) / (z_axis_length ** 2)) <= 1 mask[in_ellipsoid] = 1 - 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) + soft_edge(mask, soft_edge_width) mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True) diff --git a/src/ttmask/soft_edge.py b/src/ttmask/soft_edge.py new file mode 100644 index 0000000..04b4ea9 --- /dev/null +++ b/src/ttmask/soft_edge.py @@ -0,0 +1,9 @@ +import numpy as np +from scipy.ndimage import distance_transform_edt + +def 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) \ No newline at end of file diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py index 9fce1eb..16b32a5 100644 --- a/src/ttmask/sphere.py +++ b/src/ttmask/sphere.py @@ -4,6 +4,7 @@ import typer from scipy.ndimage import distance_transform_edt import mrcfile +from .soft_edge import soft_edge @cli.command(name='sphere') @@ -34,11 +35,7 @@ def sphere( idx = distance < (sphere_radius / pixel_size) mask[idx] = 1 - 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) + soft_edge(mask, soft_edge_width) mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True) From 483c77adaba420eccd74e30fcdf758b374ddd17e Mon Sep 17 00:00:00 2001 From: Miles Graham Date: Wed, 5 Jun 2024 15:52:51 +0100 Subject: [PATCH 2/2] make soft edge function work on copy of the original array and returning and output instead of changing 'in place' --- src/ttmask/cone.py | 12 ++++-------- src/ttmask/cube.py | 8 ++++---- src/ttmask/cuboid.py | 7 +++---- src/ttmask/cylinder.py | 7 +++---- src/ttmask/ellipsoid.py | 6 +++--- src/ttmask/soft_edge.py | 12 +++++++----- src/ttmask/sphere.py | 7 +++---- 7 files changed, 27 insertions(+), 32 deletions(-) diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py index 9ec3a89..f8ce632 100644 --- a/src/ttmask/cone.py +++ b/src/ttmask/cone.py @@ -4,7 +4,7 @@ from scipy.ndimage import distance_transform_edt import mrcfile from ._cli import cli -from .soft_edge import soft_edge +from .soft_edge import add_soft_edge @cli.command(name='cone') @@ -24,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') @@ -57,10 +57,6 @@ def cone( # 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) - soft_edge(mask, soft_edge_width) - - - mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True) - - + mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True) diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py index 6366b59..bb4baa1 100644 --- a/src/ttmask/cube.py +++ b/src/ttmask/cube.py @@ -4,13 +4,13 @@ from ._cli import cli from scipy.ndimage import distance_transform_edt import mrcfile -from .soft_edge import soft_edge +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") @@ -34,6 +34,6 @@ def cube( mask[in_cube] = 1 - soft_edge(mask, soft_edge_width) + mask = add_soft_edge(mask, soft_edge_width) - mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True) \ No newline at end of file + mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True) diff --git a/src/ttmask/cuboid.py b/src/ttmask/cuboid.py index 63d8679..0da5279 100644 --- a/src/ttmask/cuboid.py +++ b/src/ttmask/cuboid.py @@ -6,7 +6,7 @@ from typing_extensions import Annotated from scipy.ndimage import distance_transform_edt import mrcfile -from .soft_edge import soft_edge +from .soft_edge import add_soft_edge @cli.command(name='cuboid') @@ -43,7 +43,6 @@ def cuboid( inside_cuboid = np.all(difference < (np.array(cuboid_sidelengths) / (2 * pixel_size)), axis=-1) mask[inside_cuboid] = 1 + mask = add_soft_edge(mask, soft_edge_width) - soft_edge(mask, soft_edge_width) - - mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True) + mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True) diff --git a/src/ttmask/cylinder.py b/src/ttmask/cylinder.py index 99f1763..4c873a6 100644 --- a/src/ttmask/cylinder.py +++ b/src/ttmask/cylinder.py @@ -4,7 +4,7 @@ from scipy.ndimage import distance_transform_edt import mrcfile from ._cli import cli -from .soft_edge import soft_edge +from .soft_edge import add_soft_edge @cli.command(name='cylinder') @@ -40,7 +40,6 @@ def cylinder( mask[np.logical_and(idx_z, idx_xy_outer)] = 1 mask[np.logical_and(idx_z, idx_xy_inner)] = 0 - soft_edge(mask, soft_edge_width) - - 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) diff --git a/src/ttmask/ellipsoid.py b/src/ttmask/ellipsoid.py index eab50e1..45ff935 100644 --- a/src/ttmask/ellipsoid.py +++ b/src/ttmask/ellipsoid.py @@ -4,7 +4,7 @@ import typer from scipy.ndimage import distance_transform_edt import mrcfile -from .soft_edge import soft_edge +from .soft_edge import add_soft_edge @cli.command(name='ellipsoid') @@ -37,9 +37,9 @@ 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 - soft_edge(mask, soft_edge_width) + mask = add_soft_edge(mask, soft_edge_width) mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True) diff --git a/src/ttmask/soft_edge.py b/src/ttmask/soft_edge.py index 04b4ea9..490c212 100644 --- a/src/ttmask/soft_edge.py +++ b/src/ttmask/soft_edge.py @@ -1,9 +1,11 @@ import numpy as np from scipy.ndimage import distance_transform_edt -def 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) \ No newline at end of file +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 diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py index 16b32a5..cb8ee37 100644 --- a/src/ttmask/sphere.py +++ b/src/ttmask/sphere.py @@ -4,7 +4,7 @@ import typer from scipy.ndimage import distance_transform_edt import mrcfile -from .soft_edge import soft_edge +from .soft_edge import add_soft_edge @cli.command(name='sphere') @@ -35,7 +35,6 @@ def sphere( idx = distance < (sphere_radius / pixel_size) mask[idx] = 1 - soft_edge(mask, soft_edge_width) - - 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)