diff --git a/src/ttmask/box_setup.py b/src/ttmask/box_setup.py index 5cb0e34..53ae0a4 100644 --- a/src/ttmask/box_setup.py +++ b/src/ttmask/box_setup.py @@ -2,9 +2,12 @@ import einops -def box_setup(sidelength: int) -> np.ndarray: +def box_setup(sidelength: int, centering: str) -> np.ndarray: c = sidelength // 2 - center = np.array([c, c, c]) + if centering == "visual" and sidelength % 2 == 0: + center = np.array([c, c, c]) - 0.5 + else: + center = np.array([c, c, c]) mask = np.zeros(shape=(sidelength, sidelength, sidelength), dtype=np.float32) # 3d coordinates of all voxels diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py index d017f59..9cc5100 100644 --- a/src/ttmask/cone.py +++ b/src/ttmask/cone.py @@ -13,10 +13,11 @@ def cone( cone_height: float, cone_base_diameter: float, soft_edge_width: int, - pixel_size: float + pixel_size: float, + centering: str ) -> np.ndarray: # establish our coordinate system and empty mask - coordinates_centered, mask = box_setup(sidelength) + coordinates_centered, mask = box_setup(sidelength, centering) # 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') @@ -61,8 +62,9 @@ def cone_cli( soft_edge_width: int = typer.Option(0), pixel_size: float = typer.Option(1), output: Path = typer.Option(Path("cone.mrc")), + centering: str = typer.Option("standard"), ): - mask = cone(sidelength, cone_height, cone_base_diameter, soft_edge_width, pixel_size) + mask = cone(sidelength, cone_height, cone_base_diameter, soft_edge_width, pixel_size, centering) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py index 8bdb4c0..a098a8f 100644 --- a/src/ttmask/cube.py +++ b/src/ttmask/cube.py @@ -12,10 +12,11 @@ def cube( cube_sidelength: float, soft_edge_width: float, pixel_size: float, - wall_thickness: float + wall_thickness: float, + centering: str ) -> np.ndarray: - # establish our coordinate system and empty mask - coordinates_centered, mask = box_setup(sidelength) + # establish our coordinate system and empty mask + coordinates_centered, mask = box_setup(sidelength, centering) #converting relative coordinates to xyz distances (i.e. not a negative number) : xyz_distances = np.abs(coordinates_centered) @@ -42,8 +43,9 @@ def cube_cli( pixel_size: float = typer.Option(1), output: Path = typer.Option(Path("cube.mrc")), wall_thickness: float = typer.Option(0), + centering: str = typer.Option("standard"), ): - mask = cube(sidelength, cube_sidelength, soft_edge_width, pixel_size) + mask = cube(sidelength, cube_sidelength, soft_edge_width, pixel_size, wall_thickness, centering) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: diff --git a/src/ttmask/cuboid.py b/src/ttmask/cuboid.py index c8539f0..f12caab 100644 --- a/src/ttmask/cuboid.py +++ b/src/ttmask/cuboid.py @@ -14,10 +14,11 @@ def cuboid( cuboid_sidelengths: Tuple[float, float, float], soft_edge_width: float, pixel_size: float, - wall_thickness: float + wall_thickness: float, + centering: str ) -> np.ndarray: # establish our coordinate system and empty mask - coordinates_centered, mask = box_setup(sidelength) + coordinates_centered, mask = box_setup(sidelength, centering) #converting relative coordinates to xyz distances (i.e. not a negative number) : xyz_distances = np.abs(coordinates_centered) @@ -44,8 +45,9 @@ def cuboid_cli( pixel_size: float = typer.Option(1), wall_thickness: float = typer.Option(0), output: Path = typer.Option(Path("cuboid.mrc")), + centering: str = typer.Option("standard"), ): - mask = cuboid(sidelength, cuboid_sidelengths, soft_edge_width, pixel_size,wall_thickness) + mask = cuboid(sidelength, cuboid_sidelengths, soft_edge_width, pixel_size,wall_thickness, centering) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: diff --git a/src/ttmask/curved_surface.py b/src/ttmask/curved_surface.py index f71939f..7311ebc 100644 --- a/src/ttmask/curved_surface.py +++ b/src/ttmask/curved_surface.py @@ -12,12 +12,13 @@ def curved_surface( fit_sphere_diameter: float, soft_edge_width: int, pixel_size: float, - surface_thickness: float + surface_thickness: float, + centering: str ) -> np.ndarray: sphere_radius = fit_sphere_diameter / 2 # establish our coordinate system and empty mask - coordinates_centered, mask = box_setup(sidelength) + coordinates_centered, mask = box_setup(sidelength, centering) coordinates_shifted = coordinates_centered - ([0, sphere_radius, 0]) @@ -47,8 +48,9 @@ def curved_surface_cli( pixel_size: float = typer.Option(1), output: Path = typer.Option(Path("curved_surface.mrc")), surface_thickness: float = typer.Option(...), + centering: str = typer.Option("standard"), ): - mask = curved_surface(sidelength, fit_sphere_diameter, soft_edge_width, pixel_size, surface_thickness) + mask = curved_surface(sidelength, fit_sphere_diameter, soft_edge_width, pixel_size, surface_thickness, centering) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: diff --git a/src/ttmask/cylinder.py b/src/ttmask/cylinder.py index ea412fc..268a9b1 100644 --- a/src/ttmask/cylinder.py +++ b/src/ttmask/cylinder.py @@ -15,12 +15,13 @@ def cylinder( cylinder_diameter: float, wall_thickness: float, soft_edge_width: int, - pixel_size: float + pixel_size: float, + centering: str ) -> np.ndarray: cylinder_radius = cylinder_diameter / 2 # establish our coordinate system and empty mask - coordinates_centered, mask = box_setup(sidelength) + coordinates_centered, mask = box_setup(sidelength, centering) # converting relative coordinates to xyz distances (i.e. not a negative number) : xyz_distances = np.abs(coordinates_centered) @@ -53,8 +54,9 @@ def cylinder_cli( soft_edge_width: int = typer.Option(0), pixel_size: float = typer.Option(1), output: Path = typer.Option(Path("cylinder.mrc")), + centering: str = typer.Option("standard"), ): - mask = cylinder(sidelength, cylinder_height, cylinder_diameter, wall_thickness, soft_edge_width, pixel_size) + mask = cylinder(sidelength, cylinder_height, cylinder_diameter, wall_thickness, soft_edge_width, pixel_size, centering) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: diff --git a/src/ttmask/ellipsoid.py b/src/ttmask/ellipsoid.py index 0395d6c..54a340d 100644 --- a/src/ttmask/ellipsoid.py +++ b/src/ttmask/ellipsoid.py @@ -15,10 +15,11 @@ def ellipsoid( ellipsoid_dimensions: Tuple[float, float, float], soft_edge_width: float, pixel_size: float, - wall_thickness: float + wall_thickness: float, + centering: str ) -> np.ndarray: # establish our coordinate system and empty mask - coordinates_centered, mask = box_setup(sidelength) + coordinates_centered, mask = box_setup(sidelength, centering) #converting relative coordinates to xyz distances (i.e. not a negative number) : xyz_distances = np.abs(coordinates_centered) @@ -56,8 +57,9 @@ def ellipsoid_cli( soft_edge_width: float = typer.Option(0), pixel_size: float = typer.Option(1), output: Path = typer.Option(Path("ellipsoid.mrc")), + centering: str = typer.Option("standard"), ): - mask = ellipsoid(sidelength, ellipsoid_dimensions, soft_edge_width, pixel_size) + mask = ellipsoid(sidelength, ellipsoid_dimensions, soft_edge_width, pixel_size, centering) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py index e59e577..2bf5917 100644 --- a/src/ttmask/sphere.py +++ b/src/ttmask/sphere.py @@ -12,12 +12,13 @@ def sphere( sphere_diameter: float, soft_edge_width: int, pixel_size: float, - wall_thickness: float + wall_thickness: float, + centering: str ) -> np.ndarray: sphere_radius = sphere_diameter / 2 # establish our coordinate system and empty mask - coordinates_centered, mask = box_setup(sidelength) + coordinates_centered, mask = box_setup(sidelength, centering) # determine distances of each pixel to the center distance_to_center = np.linalg.norm(coordinates_centered, axis=-1) @@ -44,8 +45,9 @@ def sphere_cli( pixel_size: float = typer.Option(1), output: Path = typer.Option(Path("sphere.mrc")), wall_thickness: float = typer.Option(0), + centering: str = typer.Option("standard"), ): - mask = sphere(sidelength, sphere_diameter, soft_edge_width, pixel_size, wall_thickness) + mask = sphere(sidelength, sphere_diameter, soft_edge_width, pixel_size, wall_thickness, centering) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: diff --git a/src/ttmask/tube.py b/src/ttmask/tube.py index db680d7..ce23573 100644 --- a/src/ttmask/tube.py +++ b/src/ttmask/tube.py @@ -15,11 +15,12 @@ def tube( wall_thickness: float, soft_edge_width: int, pixel_size: float, + centering: str ) -> np.ndarray: tube_radius = tube_diameter / 2 # establish our coordinate system and empty mask - coordinates_centered, mask = box_setup(sidelength) + coordinates_centered, mask = box_setup(sidelength, centering) #converting relative coordinates to xyz distances (i.e. not a negative number) : xyz_distances = np.abs(coordinates_centered) @@ -46,8 +47,9 @@ def tube_cli( tube_diameter: float = typer.Option(...), wall_thickness: float = typer.Option(0), output: Path = typer.Option(Path("tube.mrc")), + centering: str = typer.Option("standard"), ): - mask = tube(sidelength, tube_height, tube_diameter, wall_thickness) + mask = tube(sidelength, tube_height, tube_diameter, wall_thickness, centering) # Save the mask to an MRC file with mrcfile.new(output, overwrite=True) as mrc: diff --git a/tests/test_functions.py b/tests/test_functions.py index d96c2d9..7b73e57 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -2,20 +2,20 @@ import numpy as np def test_cone(): - mask = cone(100, 50, 50, 0, 1) + mask = cone(100, 50, 50, 0, 1, "standard") assert mask.shape == (100, 100, 100) 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(): - mask = cube(100, 50, 0, 1, 0) + mask = cube(100, 50, 0, 1, 0, "standard") assert mask.shape == (100, 100, 100) # Test against volume of cube +- center and subpixel issues assert mask.sum() > 48**3 assert mask.sum() < 52**3 def test_cuboid(): - mask = cuboid(100, (50,40,30), 0, 1, 0) + mask = cuboid(100, (50,40,30), 0, 1, 0, "standard") assert mask.shape == (100, 100, 100) # Test against volume of cuboid +- center and subpixel issues assert mask.sum() > 48 * 38 * 28 @@ -28,26 +28,26 @@ def test_cuboid(): # assert mask.sum() < 2 * np.pi * 25 * 50 # Area of cylinder def test_cylinder(): - mask = cylinder(100, 50, 50, 0, 0, 1) + mask = cylinder(100, 50, 50, 0, 0, 1, "standard") 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) + mask = ellipsoid(100, (50,40,30), 0, 1,0, "standard") 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) + mask = sphere(100, 50, 0, 1,0, "standard") 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) + mask = tube(100, 50, 50, 0, 0, 1, "standard") 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