diff --git a/src/ttmask/box_setup.py b/src/ttmask/box_setup.py new file mode 100644 index 0000000..5cb0e34 --- /dev/null +++ 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 index 78fe096..88a5574 100644 --- 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 index 80dd9cf..ced32b0 100644 --- 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 index 7de249c..e088e3d 100644 --- 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 index 3d16489..923c4d8 100644 --- 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 index b7989c9..0f65a76 100644 --- a/src/ttmask/ellipsoid.py +++ b/src/ttmask/ellipsoid.py @@ -1,7 +1,6 @@ from pathlib import Path import numpy as np -import einops import typer import mrcfile from typing import Tuple @@ -9,6 +8,7 @@ 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 index 74a117d..3fb2243 100644 --- 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) diff --git a/src/ttmask/tube.py b/src/ttmask/tube.py index 5d151b2..fbade89 100644 --- a/src/ttmask/tube.py +++ b/src/ttmask/tube.py @@ -1,14 +1,13 @@ from pathlib import Path import numpy as np -import einops + import typer import mrcfile -from typing_extensions import Annotated - from ._cli import cli from .soft_edge import add_soft_edge +from .box_setup import box_setup @cli.command(name='tube') @@ -23,26 +22,24 @@ def tube( ): tube_radius = tube_diameter / 2 - 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') + # 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) - 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] < (tube_height / (2 * pixel_size)) + # set up criteria for which pixels are inside the tube and modify values to 1. + xy_distance = np.linalg.norm(xyz_distances[:, :, :, [1, 2]], axis=-1) + within_z = xyz_distances[:, :, :, 0] < (tube_height / (2 * pixel_size)) within_xy = xy_distance < (tube_radius / 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_xy_hollowing = xy_distance < ((tube_radius - wall_thickness) / pixel_size) mask[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)