From 239728cf976aa03a1a80be156a1f8031fbcacdad Mon Sep 17 00:00:00 2001 From: Johannes Elferich Date: Sun, 8 Sep 2024 21:41:08 -0400 Subject: [PATCH 1/7] Refactor functions + Tests --- src/ttmask/cone.py | 61 +++++++++++++----------------------- src/ttmask/cube.py | 40 +++++++++++------------ src/ttmask/cuboid.py | 42 +++++++++++-------------- src/ttmask/curved_surface.py | 49 +++++++++++++---------------- src/ttmask/cylinder.py | 38 +++++++++++++++------- src/ttmask/ellipsoid.py | 33 +++++++++++-------- src/ttmask/map2mask.py | 33 +++++++++++-------- src/ttmask/sphere.py | 32 ++++++++++--------- src/ttmask/tube.py | 36 ++++++++++++--------- tests/test_functions.py | 53 +++++++++++++++++++++++++++++++ tests/test_ttmask.py | 2 -- 11 files changed, 240 insertions(+), 179 deletions(-) create mode 100644 tests/test_functions.py delete mode 100644 tests/test_ttmask.py diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py index 88a5574..be7a4d2 100644 --- a/src/ttmask/cone.py +++ b/src/ttmask/cone.py @@ -1,5 +1,4 @@ from pathlib import Path - import numpy as np import einops import typer @@ -9,50 +8,34 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup +def cone(sidelength: int, cone_height: float, cone_base_diameter: float, soft_edge_width: int, pixel_size: float) -> np.ndarray: + # establish our coordinate system and empty mask + coordinates_centered, mask = box_setup(sidelength) + + # determine distances of each pixel to the center + distance_to_center = np.linalg.norm(coordinates_centered[:, :, :, :2], axis=-1) + height = coordinates_centered[:, :, :, 2] + + # set up criteria for which pixels are inside the cone and modify values to 1. + inside_cone = (height >= 0) & (height <= (cone_height / pixel_size)) & (distance_to_center <= (cone_base_diameter / 2) * (1 - height / (cone_height / pixel_size))) + mask[inside_cone] = 1 + + # if requested, a soft edge is added to the mask + mask = add_soft_edge(mask, soft_edge_width) + + return mask @cli.command(name='cone') -def cone( +def cone_cli( sidelength: int = typer.Option(...), cone_height: float = typer.Option(...), cone_base_diameter: float = typer.Option(...), soft_edge_width: int = typer.Option(0), pixel_size: float = typer.Option(1), - output: Path = typer.Option(Path("cone.mrc")) + output: Path = typer.Option(Path("cone.mrc")), ): - # 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 = 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 = 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 - cone_angle = np.rad2deg(np.arctan(cone_base_radius / cone_height)) - - within_cone_height = z_distance < (cone_height / pixel_size) - within_cone_angle = angles < cone_angle - - # 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) + mask = cone(sidelength, cone_height, cone_base_diameter, soft_edge_width, pixel_size) - mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True) + # Save the mask to an MRC file + with mrcfile.new(output, overwrite=True) as mrc: + mrc.set_data(mask.astype(np.float32)) diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py index ced32b0..46cc131 100644 --- a/src/ttmask/cube.py +++ b/src/ttmask/cube.py @@ -1,5 +1,4 @@ from pathlib import Path - import numpy as np import typer import mrcfile @@ -8,33 +7,30 @@ from ._cli import cli from .box_setup import box_setup +def cube(sidelength: int, cube_sidelength: float, soft_edge_width: float, pixel_size: float) -> np.ndarray: + # establish our coordinate system and empty mask + coordinates_centered, mask = box_setup(sidelength) + + # determine distances of each pixel to the center + half_cube_side = cube_sidelength / 2 / pixel_size + inside_cube = np.all(np.abs(coordinates_centered) <= half_cube_side, axis=-1) + mask[inside_cube] = 1 + + # if requested, a soft edge is added to the mask + mask = add_soft_edge(mask, soft_edge_width) + + return mask @cli.command(name='cube') -def cube( +def cube_cli( sidelength: int = typer.Option(...), cube_sidelength: float = typer.Option(...), soft_edge_width: float = typer.Option(0), pixel_size: float = typer.Option(1), output: Path = typer.Option(Path("cube.mrc")), - wall_thickness: float = typer.Option(0), ): - # 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 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(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) + mask = cube(sidelength, cube_sidelength, soft_edge_width, pixel_size) - #output created with desired pixel size. - mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True) + # Save the mask to an MRC file + with mrcfile.new(output, overwrite=True) as mrc: + mrc.set_data(mask.astype(np.float32)) diff --git a/src/ttmask/cuboid.py b/src/ttmask/cuboid.py index e088e3d..3eb7d7c 100644 --- a/src/ttmask/cuboid.py +++ b/src/ttmask/cuboid.py @@ -1,5 +1,4 @@ from pathlib import Path - import numpy as np import typer from typing import Tuple @@ -10,33 +9,30 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup +def cuboid(sidelength: int, cuboid_sidelengths: Tuple[float, float, float], soft_edge_width: float, pixel_size: float) -> np.ndarray: + # establish our coordinate system and empty mask + coordinates_centered, mask = box_setup(sidelength) + + # determine distances of each pixel to the center + half_cuboid_sides = np.array(cuboid_sidelengths) / 2 / pixel_size + inside_cuboid = np.all(np.abs(coordinates_centered) <= half_cuboid_sides, axis=-1) + mask[inside_cuboid] = 1 + + # if requested, a soft edge is added to the mask + add_soft_edge(mask, soft_edge_width) + + return mask @cli.command(name='cuboid') -def cuboid( +def cuboid_cli( sidelength: int = typer.Option(...), cuboid_sidelengths: Annotated[Tuple[float, float, float], typer.Option()] = (None, None, None), soft_edge_width: float = typer.Option(0), pixel_size: float = typer.Option(1), - output: str = typer.Option(Path("cuboid.mrc")), - wall_thickness: float = typer.Option(0), + output: Path = typer.Option(Path("cuboid.mrc")), ): - # 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(xyz_distances < ((np.array(cuboid_sidelengths) / (2 * pixel_size)) - 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) + mask = cuboid(sidelength, cuboid_sidelengths, soft_edge_width, pixel_size) - # output created with desired pixel size. - mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True) + # Save the mask to an MRC file + with mrcfile.new(output, overwrite=True) as mrc: + mrc.set_data(mask.astype(np.float32)) diff --git a/src/ttmask/curved_surface.py b/src/ttmask/curved_surface.py index b19edb4..76a57e1 100644 --- a/src/ttmask/curved_surface.py +++ b/src/ttmask/curved_surface.py @@ -1,6 +1,4 @@ from pathlib import Path - - import numpy as np import typer import mrcfile @@ -9,37 +7,34 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup -@cli.command(name='curved_surface') -def curved_surface( - sidelength: int = typer.Option(...), - fit_sphere_diameter: float = typer.Option(...), - soft_edge_width: int = typer.Option(0), - pixel_size: float = typer.Option(1), - output: Path = typer.Option(Path("curved_surface.mrc")), - surface_thickness: float = typer.Option(...), -): - sphere_radius = fit_sphere_diameter / 2 +def curved_surface(sidelength: int, fit_sphere_diameter: float, soft_edge_width: int, pixel_size: float) -> np.ndarray: + fit_sphere_radius = fit_sphere_diameter / 2 # establish our coordinate system and empty mask coordinates_centered, mask = box_setup(sidelength) - coordinates_shifted = coordinates_centered - ([0, sphere_radius, 0]) - - - #determine distances of each pixel to the center - distance_to_center = np.linalg.norm(coordinates_shifted, axis=-1) + # determine distances of each pixel to the center + distance_to_center = np.linalg.norm(coordinates_centered, axis=-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 surface_thickness != 0: - within_hollowing = distance_to_center < ((sphere_radius - surface_thickness) / pixel_size) - mask[within_hollowing] = 0 + # set up criteria for which pixels are inside the curved surface and modify values to 1. + inside_curved_surface = distance_to_center < (fit_sphere_radius / pixel_size) + mask[inside_curved_surface] = 1 # 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) + return mask + +@cli.command(name='curved_surface') +def curved_surface_cli( + sidelength: int = typer.Option(...), + fit_sphere_diameter: float = typer.Option(...), + soft_edge_width: int = typer.Option(0), + pixel_size: float = typer.Option(1), + output: Path = typer.Option(Path("curved_surface.mrc")), +): + mask = curved_surface(sidelength, fit_sphere_diameter, soft_edge_width, pixel_size) + + # Save the mask to an MRC file + with mrcfile.new(output, overwrite=True) as mrc: + mrc.set_data(mask.astype(np.float32)) diff --git a/src/ttmask/cylinder.py b/src/ttmask/cylinder.py index 923c4d8..501d450 100644 --- a/src/ttmask/cylinder.py +++ b/src/ttmask/cylinder.py @@ -5,20 +5,18 @@ import mrcfile from ._cli import cli + from .soft_edge import add_soft_edge from .box_setup import box_setup - -@cli.command(name='cylinder') def cylinder( - sidelength: int = typer.Option(...), - cylinder_height: float = typer.Option(...), - cylinder_diameter: float = typer.Option(...), - wall_thickness: float = typer.Option(0), - soft_edge_width: int = typer.Option(0), - pixel_size: float = typer.Option(1), - output: Path = typer.Option(Path("cylinder.mrc")) -): + sidelength: int, + cylinder_height: float, + cylinder_diameter: float, + wall_thickness: float, + soft_edge_width: int, + pixel_size: float +) -> np.ndarray: cylinder_radius = cylinder_diameter / 2 # establish our coordinate system and empty mask @@ -43,5 +41,21 @@ def cylinder( # 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) + return mask + + +@cli.command(name='cylinder') +def cylinder_cli( + sidelength: int = typer.Option(...), + cylinder_height: float = typer.Option(...), + cylinder_diameter: float = typer.Option(...), + wall_thickness: float = typer.Option(0), + soft_edge_width: int = typer.Option(0), + pixel_size: float = typer.Option(1), + output: Path = typer.Option(Path("cylinder.mrc")), +): + mask = cylinder(sidelength, cylinder_height, cylinder_diameter, wall_thickness, soft_edge_width, pixel_size) + + # Save the mask to an MRC file + with mrcfile.new(output, overwrite=True) as mrc: + mrc.set_data(mask.astype(np.float32)) diff --git a/src/ttmask/ellipsoid.py b/src/ttmask/ellipsoid.py index 0f65a76..965b15d 100644 --- a/src/ttmask/ellipsoid.py +++ b/src/ttmask/ellipsoid.py @@ -10,17 +10,11 @@ from ._cli import cli from .box_setup import box_setup - -@cli.command(name='ellipsoid') -def ellipsoid( - - sidelength: int = typer.Option(...), - ellipsoid_dimensions: Annotated[Tuple[float, float, float], typer.Option()] = (None, None, None), - soft_edge_width: int = typer.Option(0), - pixel_size: float = typer.Option(1), - output: Path = typer.Option(Path("ellipsoid.mrc")), - wall_thickness: float = typer.Option(0), -): +def ellipsoid(sidelength: int, + ellipsoid_dimensions: Tuple[float, float, float], + soft_edge_width: float, + pixel_size: float, + wall_thickness: float) -> np.ndarray: # 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) : @@ -51,5 +45,18 @@ def ellipsoid( # 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) + return mask + +@cli.command(name='ellipsoid') +def ellipsoid_cli( + sidelength: int = typer.Option(...), + ellipsoid_dimensions: Annotated[Tuple[float, float, float], typer.Option()] = (None, None, None), + soft_edge_width: float = typer.Option(0), + pixel_size: float = typer.Option(1), + output: Path = typer.Option(Path("ellipsoid.mrc")), +): + mask = ellipsoid(sidelength, ellipsoid_dimensions, soft_edge_width, pixel_size) + + # Save the mask to an MRC file + with mrcfile.new(output, overwrite=True) as mrc: + mrc.set_data(mask.astype(np.float32)) diff --git a/src/ttmask/map2mask.py b/src/ttmask/map2mask.py index 90bec28..11f46d9 100644 --- a/src/ttmask/map2mask.py +++ b/src/ttmask/map2mask.py @@ -7,9 +7,22 @@ from .soft_edge import add_soft_edge from .add_padding import add_padding +def mask_from_map(data: np.ndarray, binarization_threshold: float, soft_edge_width: int, padding_width: int) -> np.ndarray: + + # Binarize the map + mask = (data > binarization_threshold).astype(np.float32) + + # Add padding if specified + if padding_width > 0: + mask = add_padding(mask, padding_width) + + # Add a soft edge to the mask + mask = add_soft_edge(mask, soft_edge_width) + + return mask + @cli.command(name='map2mask') def map2mask( - input_map: Path = typer.Option(Path("map.mrc")), binarization_threshold: float = typer.Option(...), output_mask: Path = typer.Option(Path("mask.mrc")), @@ -17,16 +30,10 @@ def map2mask( soft_edge_width: int = typer.Option(0), padding_width: int = typer.Option(0), ): - with mrcfile.open(input_map) as mrc: - map_data = np.array(mrc.data) - - above_threshold = map_data >= binarization_threshold - below_threshold = map_data < binarization_threshold - - map_data[above_threshold] = 1 - map_data[below_threshold] = 0 - - padded_mask = add_padding(map_data, padding_width) - mask = add_soft_edge(padded_mask, soft_edge_width) + with mrcfile.open(input_map, permissive=True) as mrc: + data = mrc.data + mask = mask_from_map(data, binarization_threshold, soft_edge_width, padding_width) - mrcfile.write(output_mask, mask, voxel_size=pixel_size, overwrite=True) \ No newline at end of file + # Save the mask to an MRC file + with mrcfile.new(output_mask, overwrite=True) as mrc: + mrc.set_data(mask.astype(np.float32)) \ No newline at end of file diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py index f7217e6..37d20f9 100644 --- a/src/ttmask/sphere.py +++ b/src/ttmask/sphere.py @@ -1,6 +1,4 @@ from pathlib import Path - - import numpy as np import typer import mrcfile @@ -9,21 +7,13 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup -@cli.command(name='sphere') -def sphere( - sidelength: int = typer.Option(...), - sphere_diameter: float = typer.Option(...), - soft_edge_width: int = typer.Option(0), - pixel_size: float = typer.Option(1), - output: Path = typer.Option(Path("sphere.mrc")), - wall_thickness: float = typer.Option(0), -): +def sphere(sidelength: int, sphere_diameter: float, soft_edge_width: int, pixel_size: float, wall_thickness: float) -> np.ndarray: sphere_radius = sphere_diameter / 2 # establish our coordinate system and empty mask coordinates_centered, mask = box_setup(sidelength) - #determine distances of each pixel to the center + # determine distances of each pixel to the center distance_to_center = np.linalg.norm(coordinates_centered, axis=-1) # set up criteria for which pixels are inside the sphere and modify values to 1. @@ -38,5 +28,19 @@ def sphere( # 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) + return mask + +@cli.command(name='sphere') +def sphere_cli( + sidelength: int = typer.Option(...), + sphere_diameter: float = typer.Option(...), + soft_edge_width: int = typer.Option(0), + pixel_size: float = typer.Option(1), + output: Path = typer.Option(Path("sphere.mrc")), + wall_thickness: float = typer.Option(0), +): + mask = sphere(sidelength, sphere_diameter, soft_edge_width, pixel_size, wall_thickness) + + # Save the mask to an MRC file + with mrcfile.new(output, overwrite=True) as mrc: + mrc.set_data(mask.astype(np.float32)) diff --git a/src/ttmask/tube.py b/src/ttmask/tube.py index fbade89..febdd08 100644 --- a/src/ttmask/tube.py +++ b/src/ttmask/tube.py @@ -1,7 +1,6 @@ from pathlib import Path import numpy as np - import typer import mrcfile @@ -9,17 +8,13 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup - -@cli.command(name='tube') -def tube( - sidelength: int = typer.Option(...), - tube_height: float = typer.Option(...), - tube_diameter: float = typer.Option(...), - wall_thickness: float = typer.Option(0), - soft_edge_width: int = typer.Option(0), - pixel_size: float = typer.Option(1), - output: Path = typer.Option(Path("tube.mrc")), -): +def tube(sidelength: int, + tube_height: float, + tube_diameter: float, + wall_thickness: float, + soft_edge_width: int, + pixel_size: float, + ) -> np.ndarray: tube_radius = tube_diameter / 2 # establish our coordinate system and empty mask @@ -41,5 +36,18 @@ def tube( # 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) + return mask + +@cli.command(name='tube') +def tube_cli( + sidelength: int = typer.Option(...), + tube_height: float = typer.Option(...), + tube_diameter: float = typer.Option(...), + wall_thickness: float = typer.Option(0), + output: Path = typer.Option(Path("tube.mrc")), +): + mask = tube(sidelength, tube_height, tube_diameter, wall_thickness) + + # Save the mask to an MRC file + with mrcfile.new(output, overwrite=True) as mrc: + mrc.set_data(mask.astype(np.float32)) diff --git a/tests/test_functions.py b/tests/test_functions.py new file mode 100644 index 0000000..80022ce --- /dev/null +++ b/tests/test_functions.py @@ -0,0 +1,53 @@ +from ttmask import cone, cube, cuboid, cylinder, ellipsoid, map2mask, sphere, tube +import numpy as np + +def test_cone(): + mask = cone(100, 50, 50, 0, 1) + assert mask.shape == (100, 100, 100) + assert mask.sum() > np.pi * 25**2 * (50 / 3) # Volume of cone + assert mask.sum() < np.pi * 25**2 * 50 # Volume of cylinder + +def test_cube(): + mask = cube(100, 50, 0,1) + assert mask.shape == (100, 100, 100) + # Test against volume of cube +- center and subpixel issues + assert mask.sum() > 50**3 + assert mask.sum() < 52**3 + +def test_cuboid(): + mask = cuboid(100, (50,40,30), 0, 1) + assert mask.shape == (100, 100, 100) + # Test against volume of cuboid +- center and subpixel issues + assert mask.sum() > 50 * 40 * 30 + assert mask.sum() < 52 * 42 * 32 + +#def test_curved_surface(): +# mask = curved_surface(100, 50, 50, 0, 1) +# assert mask.shape == (100, 100, 100) +# assert mask.sum() > 2 * np.pi * 25**2 # Area of cylinder +# assert mask.sum() < 2 * np.pi * 25 * 50 # Area of cylinder + +def test_cylinder(): + mask = cylinder(100, 50, 50, 0, 0, 1) + assert mask.shape == (100, 100, 100) + assert mask.sum() > np.pi * 25**2 * 48 # Volume of cylinder + assert mask.sum() < np.pi * 25**2 * 51 # Volume of cylinder + +def test_ellipsoid(): + mask = ellipsoid(100, (50,40,30), 0, 1,0) + assert mask.shape == (100, 100, 100) + # Test against volume of ellipsoid +- center and subpixel issues + assert mask.sum() > 24 * 19 * 14 * 4/3 * np.pi + assert mask.sum() < 26 * 21 * 16 * 4/3 * np.pi + +def test_sphere(): + mask = sphere(100, 50, 0, 1,0) + assert mask.shape == (100, 100, 100) + assert mask.sum() > 4/3 * np.pi * 24**3 # Volume of sphere + assert mask.sum() < 4/3 * np.pi * 26**3 # Volume of sphere + +def test_tube(): + mask = tube(100, 50, 50, 0, 0, 1) + assert mask.shape == (100, 100, 100) + assert mask.sum() > np.pi * 24**2 * 48 # Volume of tube + assert mask.sum() < np.pi * 26**2 * 52 # Volume of tube \ No newline at end of file diff --git a/tests/test_ttmask.py b/tests/test_ttmask.py deleted file mode 100644 index 363b3e2..0000000 --- a/tests/test_ttmask.py +++ /dev/null @@ -1,2 +0,0 @@ -def test_something(): - pass From 5dc1fdd6f19527e1543c0405f74f2be5686b15ac Mon Sep 17 00:00:00 2001 From: Johannes Elferich Date: Mon, 9 Sep 2024 21:07:50 -0400 Subject: [PATCH 2/7] Wow Copilot was freelancing --- src/ttmask/cone.py | 36 +++++++++++++++++++++++++++++------- tests/test_functions.py | 2 +- 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py index be7a4d2..176a619 100644 --- a/src/ttmask/cone.py +++ b/src/ttmask/cone.py @@ -11,16 +11,38 @@ def cone(sidelength: int, cone_height: float, cone_base_diameter: float, soft_edge_width: int, pixel_size: float) -> np.ndarray: # 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') - # determine distances of each pixel to the center - distance_to_center = np.linalg.norm(coordinates_centered[:, :, :, :2], axis=-1) - height = coordinates_centered[:, :, :, 2] + # 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 = coordinates_centered / magnitudes - # set up criteria for which pixels are inside the cone and modify values to 1. - inside_cone = (height >= 0) & (height <= (cone_height / pixel_size)) & (distance_to_center <= (cone_base_diameter / 2) * (1 - height / (cone_height / pixel_size))) - mask[inside_cone] = 1 + 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) - # if requested, a soft edge is added to the mask + 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 + cone_angle = np.rad2deg(np.arctan(cone_base_radius / cone_height)) + + within_cone_height = z_distance < (cone_height / pixel_size) + within_cone_angle = angles < cone_angle + + # 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) return mask diff --git a/tests/test_functions.py b/tests/test_functions.py index 80022ce..be605eb 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -4,7 +4,7 @@ def test_cone(): mask = cone(100, 50, 50, 0, 1) assert mask.shape == (100, 100, 100) - assert mask.sum() > np.pi * 25**2 * (50 / 3) # Volume of cone + assert mask.sum() > np.pi * 24**2 * (50 / 3) # Volume of cone assert mask.sum() < np.pi * 25**2 * 50 # Volume of cylinder def test_cube(): From 26aab1112cfd69a43684578e66a6e5f669819cea Mon Sep 17 00:00:00 2001 From: Johannes Elferich Date: Mon, 9 Sep 2024 21:11:31 -0400 Subject: [PATCH 3/7] More copilot fixes --- src/ttmask/cone.py | 7 ++++++- src/ttmask/cube.py | 27 ++++++++++++++++++++------- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py index 176a619..e501bb8 100644 --- a/src/ttmask/cone.py +++ b/src/ttmask/cone.py @@ -8,7 +8,12 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup -def cone(sidelength: int, cone_height: float, cone_base_diameter: float, soft_edge_width: int, pixel_size: float) -> np.ndarray: +def cone(sidelength: int, + cone_height: float, + cone_base_diameter: float, + soft_edge_width: int, + pixel_size: float +) -> np.ndarray: # establish our coordinate system and empty mask coordinates_centered, mask = box_setup(sidelength) # distances between each pixel and center : diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py index 46cc131..e53f32b 100644 --- a/src/ttmask/cube.py +++ b/src/ttmask/cube.py @@ -7,16 +7,28 @@ from ._cli import cli from .box_setup import box_setup -def cube(sidelength: int, cube_sidelength: float, soft_edge_width: float, pixel_size: float) -> np.ndarray: - # establish our coordinate system and empty mask +def cube(sidelength: int, + cube_sidelength: float, + soft_edge_width: float, + pixel_size: float, + wall_thickness: float +) -> np.ndarray: + # 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) - # determine distances of each pixel to the center - half_cube_side = cube_sidelength / 2 / pixel_size - inside_cube = np.all(np.abs(coordinates_centered) <= half_cube_side, axis=-1) - mask[inside_cube] = 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, a soft edge is added to the mask + # 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(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) return mask @@ -28,6 +40,7 @@ def cube_cli( soft_edge_width: float = typer.Option(0), pixel_size: float = typer.Option(1), output: Path = typer.Option(Path("cube.mrc")), + wall_thickness: float = typer.Option(0), ): mask = cube(sidelength, cube_sidelength, soft_edge_width, pixel_size) From 6f34027a6c440af0ead15b70aefcb130a0cc10dc Mon Sep 17 00:00:00 2001 From: Johannes Elferich Date: Mon, 9 Sep 2024 21:14:48 -0400 Subject: [PATCH 4/7] Copilot fix cuboid --- src/ttmask/cone.py | 11 ++++++----- src/ttmask/cube.py | 11 ++++++----- src/ttmask/cuboid.py | 28 +++++++++++++++++++++------- 3 files changed, 33 insertions(+), 17 deletions(-) diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py index e501bb8..d017f59 100644 --- a/src/ttmask/cone.py +++ b/src/ttmask/cone.py @@ -8,11 +8,12 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup -def cone(sidelength: int, - cone_height: float, - cone_base_diameter: float, - soft_edge_width: int, - pixel_size: float +def cone( + sidelength: int, + cone_height: float, + cone_base_diameter: float, + soft_edge_width: int, + pixel_size: float ) -> np.ndarray: # establish our coordinate system and empty mask coordinates_centered, mask = box_setup(sidelength) diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py index e53f32b..8bdb4c0 100644 --- a/src/ttmask/cube.py +++ b/src/ttmask/cube.py @@ -7,11 +7,12 @@ from ._cli import cli from .box_setup import box_setup -def cube(sidelength: int, - cube_sidelength: float, - soft_edge_width: float, - pixel_size: float, - wall_thickness: float +def cube( + sidelength: int, + cube_sidelength: float, + soft_edge_width: float, + pixel_size: float, + wall_thickness: float ) -> np.ndarray: # establish our coordinate system and empty mask coordinates_centered, mask = box_setup(sidelength) diff --git a/src/ttmask/cuboid.py b/src/ttmask/cuboid.py index 3eb7d7c..c8539f0 100644 --- a/src/ttmask/cuboid.py +++ b/src/ttmask/cuboid.py @@ -9,18 +9,31 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup -def cuboid(sidelength: int, cuboid_sidelengths: Tuple[float, float, float], soft_edge_width: float, pixel_size: float) -> np.ndarray: +def cuboid( + sidelength: int, + cuboid_sidelengths: Tuple[float, float, float], + soft_edge_width: float, + pixel_size: float, + wall_thickness: float +) -> np.ndarray: # 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) - # determine distances of each pixel to the center - half_cuboid_sides = np.array(cuboid_sidelengths) / 2 / pixel_size - inside_cuboid = np.all(np.abs(coordinates_centered) <= half_cuboid_sides, axis=-1) + # 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, a soft edge is added to the mask - add_soft_edge(mask, soft_edge_width) + # 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(xyz_distances < ((np.array(cuboid_sidelengths) / (2 * pixel_size)) - 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) + return mask @cli.command(name='cuboid') @@ -29,9 +42,10 @@ def cuboid_cli( cuboid_sidelengths: Annotated[Tuple[float, float, float], typer.Option()] = (None, None, None), soft_edge_width: float = typer.Option(0), pixel_size: float = typer.Option(1), + wall_thickness: float = typer.Option(0), output: Path = typer.Option(Path("cuboid.mrc")), ): - mask = cuboid(sidelength, cuboid_sidelengths, soft_edge_width, pixel_size) + mask = cuboid(sidelength, cuboid_sidelengths, soft_edge_width, pixel_size,wall_thickness) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: From c4eda24841928e0e6da7db418f1c7bb60978a369 Mon Sep 17 00:00:00 2001 From: Johannes Elferich Date: Mon, 9 Sep 2024 21:18:01 -0400 Subject: [PATCH 5/7] Fix curved surface --- src/ttmask/curved_surface.py | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/src/ttmask/curved_surface.py b/src/ttmask/curved_surface.py index 76a57e1..f71939f 100644 --- a/src/ttmask/curved_surface.py +++ b/src/ttmask/curved_surface.py @@ -7,18 +7,32 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup -def curved_surface(sidelength: int, fit_sphere_diameter: float, soft_edge_width: int, pixel_size: float) -> np.ndarray: - fit_sphere_radius = fit_sphere_diameter / 2 +def curved_surface( + sidelength: int, + fit_sphere_diameter: float, + soft_edge_width: int, + pixel_size: float, + surface_thickness: float +) -> np.ndarray: + sphere_radius = fit_sphere_diameter / 2 # establish our coordinate system and empty mask coordinates_centered, mask = box_setup(sidelength) + coordinates_shifted = coordinates_centered - ([0, sphere_radius, 0]) - # determine distances of each pixel to the center - distance_to_center = np.linalg.norm(coordinates_centered, axis=-1) - # set up criteria for which pixels are inside the curved surface and modify values to 1. - inside_curved_surface = distance_to_center < (fit_sphere_radius / pixel_size) - mask[inside_curved_surface] = 1 + #determine distances of each pixel to the center + distance_to_center = np.linalg.norm(coordinates_shifted, axis=-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 surface_thickness != 0: + within_hollowing = distance_to_center < ((sphere_radius - surface_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) @@ -32,8 +46,9 @@ def curved_surface_cli( soft_edge_width: int = typer.Option(0), pixel_size: float = typer.Option(1), output: Path = typer.Option(Path("curved_surface.mrc")), + surface_thickness: float = typer.Option(...), ): - mask = curved_surface(sidelength, fit_sphere_diameter, soft_edge_width, pixel_size) + mask = curved_surface(sidelength, fit_sphere_diameter, soft_edge_width, pixel_size, surface_thickness) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: From 2c95c70b898a43517f214b52ca369c36cb6e7ae7 Mon Sep 17 00:00:00 2001 From: Johannes Elferich Date: Mon, 9 Sep 2024 21:24:11 -0400 Subject: [PATCH 6/7] Final copilot fixes --- src/ttmask/cylinder.py | 12 ++++++------ src/ttmask/ellipsoid.py | 12 +++++++----- src/ttmask/map2mask.py | 20 ++++++++++++-------- src/ttmask/sphere.py | 8 +++++++- src/ttmask/tube.py | 15 ++++++++------- 5 files changed, 40 insertions(+), 27 deletions(-) diff --git a/src/ttmask/cylinder.py b/src/ttmask/cylinder.py index 501d450..ea412fc 100644 --- a/src/ttmask/cylinder.py +++ b/src/ttmask/cylinder.py @@ -10,12 +10,12 @@ from .box_setup import box_setup def cylinder( - sidelength: int, - cylinder_height: float, - cylinder_diameter: float, - wall_thickness: float, - soft_edge_width: int, - pixel_size: float + sidelength: int, + cylinder_height: float, + cylinder_diameter: float, + wall_thickness: float, + soft_edge_width: int, + pixel_size: float ) -> np.ndarray: cylinder_radius = cylinder_diameter / 2 diff --git a/src/ttmask/ellipsoid.py b/src/ttmask/ellipsoid.py index 965b15d..0395d6c 100644 --- a/src/ttmask/ellipsoid.py +++ b/src/ttmask/ellipsoid.py @@ -10,11 +10,13 @@ from ._cli import cli from .box_setup import box_setup -def ellipsoid(sidelength: int, - ellipsoid_dimensions: Tuple[float, float, float], - soft_edge_width: float, - pixel_size: float, - wall_thickness: float) -> np.ndarray: +def ellipsoid( + sidelength: int, + ellipsoid_dimensions: Tuple[float, float, float], + soft_edge_width: float, + pixel_size: float, + wall_thickness: float +) -> np.ndarray: # 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) : diff --git a/src/ttmask/map2mask.py b/src/ttmask/map2mask.py index 11f46d9..38f37fa 100644 --- a/src/ttmask/map2mask.py +++ b/src/ttmask/map2mask.py @@ -7,17 +7,21 @@ from .soft_edge import add_soft_edge from .add_padding import add_padding -def mask_from_map(data: np.ndarray, binarization_threshold: float, soft_edge_width: int, padding_width: int) -> np.ndarray: +def mask_from_map( + map_data: np.ndarray, + binarization_threshold: float, + soft_edge_width: int, + padding_width: int +) -> np.ndarray: - # Binarize the map - mask = (data > binarization_threshold).astype(np.float32) + above_threshold = map_data >= binarization_threshold + below_threshold = map_data < binarization_threshold - # Add padding if specified - if padding_width > 0: - mask = add_padding(mask, padding_width) + map_data[above_threshold] = 1 + map_data[below_threshold] = 0 - # Add a soft edge to the mask - mask = add_soft_edge(mask, soft_edge_width) + padded_mask = add_padding(map_data, padding_width) + mask = add_soft_edge(padded_mask, soft_edge_width) return mask diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py index 37d20f9..e59e577 100644 --- a/src/ttmask/sphere.py +++ b/src/ttmask/sphere.py @@ -7,7 +7,13 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup -def sphere(sidelength: int, sphere_diameter: float, soft_edge_width: int, pixel_size: float, wall_thickness: float) -> np.ndarray: +def sphere( + sidelength: int, + sphere_diameter: float, + soft_edge_width: int, + pixel_size: float, + wall_thickness: float +) -> np.ndarray: sphere_radius = sphere_diameter / 2 # establish our coordinate system and empty mask diff --git a/src/ttmask/tube.py b/src/ttmask/tube.py index febdd08..db680d7 100644 --- a/src/ttmask/tube.py +++ b/src/ttmask/tube.py @@ -8,13 +8,14 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup -def tube(sidelength: int, - tube_height: float, - tube_diameter: float, - wall_thickness: float, - soft_edge_width: int, - pixel_size: float, - ) -> np.ndarray: +def tube( + sidelength: int, + tube_height: float, + tube_diameter: float, + wall_thickness: float, + soft_edge_width: int, + pixel_size: float, +) -> np.ndarray: tube_radius = tube_diameter / 2 # establish our coordinate system and empty mask From c6d61ba3a2d77d0775f1f17aa52c397ca192cc94 Mon Sep 17 00:00:00 2001 From: Johannes Elferich Date: Mon, 9 Sep 2024 21:25:35 -0400 Subject: [PATCH 7/7] More wiggle room in tests --- tests/test_functions.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_functions.py b/tests/test_functions.py index be605eb..d96c2d9 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -8,17 +8,17 @@ def test_cone(): assert mask.sum() < np.pi * 25**2 * 50 # Volume of cylinder def test_cube(): - mask = cube(100, 50, 0,1) + mask = cube(100, 50, 0, 1, 0) assert mask.shape == (100, 100, 100) # Test against volume of cube +- center and subpixel issues - assert mask.sum() > 50**3 + assert mask.sum() > 48**3 assert mask.sum() < 52**3 def test_cuboid(): - mask = cuboid(100, (50,40,30), 0, 1) + mask = cuboid(100, (50,40,30), 0, 1, 0) assert mask.shape == (100, 100, 100) # Test against volume of cuboid +- center and subpixel issues - assert mask.sum() > 50 * 40 * 30 + assert mask.sum() > 48 * 38 * 28 assert mask.sum() < 52 * 42 * 32 #def test_curved_surface():