tidy up of code including refactoring box creation (#41)

teamtomo · Jun 17, 2024 · 24f77fe · 24f77fe
1 parent cd79dd0
commit 24f77fe
Show file tree

Hide file tree

Showing 8 changed files with 104 additions and 125 deletions.
diff --git a/src/ttmask/box_setup.py b/src/ttmask/box_setup.py
@@ -0,0 +1,16 @@
+import numpy as np
+import einops
+
+
+def box_setup(sidelength: int) -> np.ndarray:
+    c = sidelength // 2
+    center = np.array([c, c, c])
+    mask = np.zeros(shape=(sidelength, sidelength, sidelength), dtype=np.float32)
+
+    # 3d coordinates of all voxels
+    coordinates = np.indices([sidelength, sidelength, sidelength])
+    coordinates = einops.rearrange(coordinates, 'zyx d h w -> d h w zyx')
+
+    #coordinates now expressed relative to center
+    coordinates_centered = coordinates - center
+    return (coordinates_centered, mask)
diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py
@@ -7,6 +7,7 @@
 
 from ._cli import cli
 from .soft_edge import add_soft_edge
+from .box_setup import box_setup
 
 
 @cli.command(name='cone')
@@ -18,30 +19,23 @@ def cone(
     pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("cone.mrc"))
 ):
-    c = sidelength // 2
-    center = np.array([c, c, c])
-    mask = np.zeros(shape=(sidelength, sidelength, sidelength), dtype=np.float32)
-
-    # 3d positions of all voxels
-    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
-    magnitudes = np.linalg.norm(centered, axis=-1)
-
+    # establish our coordinate system and empty mask
+    coordinates_centered, mask = box_setup(sidelength)
+    # distances between each pixel and center :
+    magnitudes = np.linalg.norm(coordinates_centered, axis=-1)
     magnitudes = einops.rearrange(magnitudes, 'd h w -> d h w 1')
 
     # Check for zeros in magnitudes and replace them with a small value to avoid Nan warning
     near_zero = 1e-8
     magnitudes = np.where(magnitudes == 0, near_zero, magnitudes)
-    normalised = centered / magnitudes
+    normalised = coordinates_centered / magnitudes
 
     principal_axis = np.array([1, 0, 0])
     dot_product = np.dot(normalised, principal_axis)
     angles_radians = np.arccos(dot_product)
     angles = np.rad2deg(angles_radians)
 
-    z_distance = centered[:, :, :, 0]  # (100, 100, 100)
+    z_distance = coordinates_centered[:, :, :, 0]  # (100, 100, 100)
 
     # Calculate the angle from the tip of the cone to the edge of the base
     cone_base_radius = (cone_base_diameter / 2) / pixel_size

diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py
@@ -1,12 +1,12 @@
 from pathlib import Path
 
 import numpy as np
-import einops
 import typer
 import mrcfile
 
 from .soft_edge import add_soft_edge
 from ._cli import cli
+from .box_setup import box_setup
 
 
 @cli.command(name='cube')
@@ -18,28 +18,23 @@ def cube(
     output: Path = typer.Option(Path("cube.mrc")),
     wall_thickness: float = typer.Option(0),
 ):
-    c = sidelength // 2
-    center = np.array([c, c, c])
-    mask = np.zeros(shape=(sidelength, sidelength, sidelength), dtype=np.float32)
+    # establish our coordinate system and empty mask
+    coordinates_centered, mask = box_setup(sidelength)
+    #converting relative coordinates to xyz distances (i.e. not a negative number) :
+    xyz_distances = np.abs(coordinates_centered)
 
-    # 3d positions of all voxels
-    positions = np.indices([sidelength, sidelength, sidelength])
-    positions = einops.rearrange(positions, 'zyx d h w -> d h w zyx')
-
-    # calculate the distance between the center and every pixel position
-    print(center.shape)
-    print(positions.shape)
-
-    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)
+    # set up criteria for which pixels are inside the cube and modify values to 1.
+    in_cube = np.all(xyz_distances < np.array(cube_sidelength) / (pixel_size * 2), axis=-1)
     mask[in_cube] = 1
 
+    # if requested, criteria set up for pixels within the hollowed area and these values changed to zero
     if wall_thickness != 0:
-        within_hollowing = np.all(difference < ((np.array(cube_sidelength) / (pixel_size * 2)) - wall_thickness),
+        within_hollowing = np.all(xyz_distances < ((np.array(cube_sidelength) / (pixel_size * 2)) - wall_thickness),
                                   axis=-1)
         mask[within_hollowing] = 0
 
+    #if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
+
+    #output created with desired pixel size.
     mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
diff --git a/src/ttmask/cuboid.py b/src/ttmask/cuboid.py
@@ -1,16 +1,14 @@
 from pathlib import Path
 
 import numpy as np
-import einops
 import typer
 from typing import Tuple
 from typing_extensions import Annotated
 import mrcfile
-from .soft_edge import add_soft_edge
-
-from ._cli import cli
 
 from ._cli import cli
+from .soft_edge import add_soft_edge
+from .box_setup import box_setup
 
 
 @cli.command(name='cuboid')
@@ -22,38 +20,23 @@ def cuboid(
     output: str = typer.Option(Path("cuboid.mrc")),
     wall_thickness: float = typer.Option(0),
 ):
-    c = sidelength // 2
-    center = np.array([c, c, c])
-    mask = np.zeros(shape=(sidelength, sidelength, sidelength), dtype=np.float32)
-
-    # 3d positions of all voxels
-    positions = np.indices([sidelength, sidelength, sidelength])
-    positions = einops.rearrange(positions, 'zyx d h w -> d h w zyx')
-
-    # calculate the distance between the center and every pixel position
-    print(center.shape)
-    print(positions.shape)
-
-    print('calculating distance')
-    difference = np.abs(positions - center)  # (100, 100, 100, 3)
-    # z = difference[:, :, :, 0]
-    # y = difference[:, :, :, 1]
-    # x = difference[:, :, :, 2]
-    # idx_z = z < cuboid_sidelengths[0] / 2
-    # idx_y = y < cuboid_sidelengths[1] / 2
-    # idx_x = x < cuboid_sidelengths[2] / 2
-
-    # mask[np.logical_not(idx)] = 1 #if you wanted to do opposite for whatever reason
-
-    inside_cuboid = np.all(difference < (np.array(cuboid_sidelengths) / (2 * pixel_size)), axis=-1)
+    # establish our coordinate system and empty mask
+    coordinates_centered, mask = box_setup(sidelength)
+    #converting relative coordinates to xyz distances (i.e. not a negative number) :
+    xyz_distances = np.abs(coordinates_centered)
 
+    # set up criteria for which pixels are inside the cuboid and modify values to 1.
+    inside_cuboid = np.all(xyz_distances < (np.array(cuboid_sidelengths) / (2 * pixel_size)), axis=-1)
     mask[inside_cuboid] = 1
-
+
+    # if requested, criteria set up for pixels within the hollowed area and these values changed to zero
     if wall_thickness != 0:
-        within_hollowing = np.all(difference < ((np.array(cuboid_sidelengths) / (2 * pixel_size)) - wall_thickness), axis=-1)
+        within_hollowing = np.all(xyz_distances < ((np.array(cuboid_sidelengths) / (2 * pixel_size)) - wall_thickness),
+                                  axis=-1)
         mask[within_hollowing] = 0
-
-    mask = add_soft_edge(mask, soft_edge_width)
 
-    mrcfile.write(output, mask, voxel_size= pixel_size, overwrite=True)
+    # if requested, a soft edge is added to the mask
+    mask = add_soft_edge(mask, soft_edge_width)
 
+    # output created with desired pixel size.
+    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
diff --git a/src/ttmask/cylinder.py b/src/ttmask/cylinder.py
@@ -1,12 +1,12 @@
 from pathlib import Path
 
 import numpy as np
-import einops
 import typer
 import mrcfile
 
 from ._cli import cli
 from .soft_edge import add_soft_edge
+from .box_setup import box_setup
 
 
 @cli.command(name='cylinder')
@@ -21,28 +21,27 @@ def cylinder(
 ):
     cylinder_radius = cylinder_diameter / 2
 
-    c = sidelength // 2
-    center = np.array([c, c, c])
-    mask = np.zeros(shape=(sidelength, sidelength, sidelength), dtype=np.float32)
+    # establish our coordinate system and empty mask
+    coordinates_centered, mask = box_setup(sidelength)
 
-    # 3d positions of all voxels
-    positions = np.indices([sidelength, sidelength, sidelength])
-    positions = einops.rearrange(positions, 'zyx d h w -> d h w zyx')
-
-    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)
+    # converting relative coordinates to xyz distances (i.e. not a negative number) :
+    xyz_distances = np.abs(coordinates_centered)
 
+    # set up criteria for which pixels are inside the cylinder and modify values to 1.
 
+    xy_distance = np.linalg.norm(xyz_distances[:, :, :, [1, 2]], axis=-1)
+    within_xy = xy_distance < (cylinder_radius / pixel_size)
+    within_z = xyz_distances[:, :, :, 0] < (cylinder_height / (2 * pixel_size))
     mask[np.logical_and(within_z, within_xy)] = 1
 
+    # if requested, criteria set up for pixels within the hollowed area and these values changed to zero
     if wall_thickness != 0:
-        within_z_hollowing = difference[:, :, :, 0] < (cylinder_height - wall_thickness) / (2 * pixel_size)
+        within_z_hollowing = xyz_distances[:, :, :, 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
 
+    # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
 
+    # output created with desired pixel size.
     mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
diff --git a/src/ttmask/ellipsoid.py b/src/ttmask/ellipsoid.py
@@ -1,14 +1,14 @@
 from pathlib import Path
 
 import numpy as np
-import einops
 import typer
 import mrcfile
 from typing import Tuple
 from typing_extensions import Annotated
 
 from .soft_edge import add_soft_edge
 from ._cli import cli
+from .box_setup import box_setup
 
 
 @cli.command(name='ellipsoid')
@@ -21,35 +21,35 @@ def ellipsoid(
     output: Path = typer.Option(Path("ellipsoid.mrc")),
     wall_thickness: float = typer.Option(0),
 ):
-    c = sidelength // 2
-    center = np.array([c, c, c])
-    mask = np.zeros(shape=(sidelength, sidelength, sidelength), dtype=np.float32)
+    # establish our coordinate system and empty mask
+    coordinates_centered, mask = box_setup(sidelength)
+    #converting relative coordinates to xyz distances (i.e. not a negative number) :
+    xyz_distances = np.abs(coordinates_centered)
 
-    # positions of all pixels
-    positions = np.indices([sidelength, sidelength, sidelength])
-    positions = einops.rearrange(positions, 'zyx d h w -> d h w zyx')
-
-    difference = np.abs(positions - center)  # (100, 100, 100, 3)
-    distances = np.linalg.norm(difference, axis=-1)
-
-    x_magnitude = difference[:, :, :, 2]
-    y_magnitude = difference[:, :, :, 1]
-    z_magnitude = difference[:, :, :, 0]
+    #extract xyz magnitudes from xyz_distances
+    x_magnitude = xyz_distances[:, :, :, 2]
+    y_magnitude = xyz_distances[:, :, :, 1]
+    z_magnitude = xyz_distances[:, :, :, 0]
 
+    #define desired dimensions in pixels (converting from angstrom)
     z_axis_length = ellipsoid_dimensions[0] / (2 * pixel_size)
     y_axis_length = ellipsoid_dimensions[1] / (2 * pixel_size)
     x_axis_length = ellipsoid_dimensions[2] / (2 * pixel_size)
 
+    # set up criteria for which pixels are inside the ellipsoid and modify values to 1.
     in_ellipsoid = (((x_magnitude) ** 2) / (x_axis_length ** 2)) + ((y_magnitude ** 2) / (y_axis_length ** 2)) + (
         (z_magnitude ** 2) / (z_axis_length ** 2)) <= 1
     mask[in_ellipsoid] = 1
 
+    # if requested, criteria set up for pixels within the hollowed area and these values changed to zero
     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
 
+    # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
 
+    # output created with desired pixel size.
     mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py
@@ -1,13 +1,13 @@
 from pathlib import Path
 
-from ._cli import cli
+
 import numpy as np
-import einops
 import typer
 import mrcfile
 
-from .soft_edge import add_soft_edge
 from ._cli import cli
+from .soft_edge import add_soft_edge
+from .box_setup import box_setup
 
 
 @cli.command(name='sphere')
@@ -20,29 +20,24 @@ def sphere(
     wall_thickness: float = typer.Option(0),
 ):
     sphere_radius = sphere_diameter / 2
-    c = sidelength // 2
-    center = np.array([c, c, c])
-    mask = np.zeros(shape=(sidelength, sidelength, sidelength), dtype=np.float32)
-
-    print(mask.dtype)
 
-    # 2d positions of all pixels
-    positions = np.indices([sidelength, sidelength, sidelength])
-    positions = einops.rearrange(positions, 'zyx d h w -> d h w zyx')
+    # establish our coordinate system and empty mask
+    coordinates_centered, mask = box_setup(sidelength)
 
-    print('calculating distance')
-    difference = np.abs(positions - center)  # (100, 100, 100, 3)
-    distance = np.sum(difference ** 2, axis=-1) ** 0.5
+    #determine distances of each pixel to the center
+    distance_to_center = np.linalg.norm(coordinates_centered)
 
-    # calculate whether each pixel is inside or outside the circle
-    print('calculating which pixels are in sphere')
-    idx = distance < (sphere_radius / pixel_size)
-    mask[idx] = 1
+    # set up criteria for which pixels are inside the sphere and modify values to 1.
+    inside_sphere = distance_to_center < (sphere_radius / pixel_size)
+    mask[inside_sphere] = 1
 
+    # if requested, criteria set up for pixels within the hollowed area and these values changed to zero
     if wall_thickness != 0:
-        within_hollowing = distance < ((sphere_radius - wall_thickness) / pixel_size)
+        within_hollowing = distance_to_center < ((sphere_radius - wall_thickness) / pixel_size)
         mask[within_hollowing] = 0
 
+    # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
 
+    # output created with desired pixel size.
     mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)