diff --git a/src/ttmask/__init__.py b/src/ttmask/__init__.py index fc7d461..f4ee12b 100644 --- a/src/ttmask/__init__.py +++ b/src/ttmask/__init__.py @@ -18,4 +18,5 @@ from .ellipsoid import ellipsoid from .map2mask import map2mask from .tube import tube +from .curved_surface import curved_surface diff --git a/src/ttmask/curved_surface.py b/src/ttmask/curved_surface.py new file mode 100644 index 0000000..b19edb4 --- /dev/null +++ b/src/ttmask/curved_surface.py @@ -0,0 +1,45 @@ +from pathlib import Path + + +import numpy as np +import typer +import mrcfile + +from ._cli import cli +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 + + # 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) + + + # 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) + + # 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 3fb2243..f7217e6 100644 --- a/src/ttmask/sphere.py +++ b/src/ttmask/sphere.py @@ -9,7 +9,6 @@ from .soft_edge import add_soft_edge from .box_setup import box_setup - @cli.command(name='sphere') def sphere( sidelength: int = typer.Option(...), @@ -25,7 +24,7 @@ def sphere( coordinates_centered, mask = box_setup(sidelength) #determine distances of each pixel to the center - distance_to_center = np.linalg.norm(coordinates_centered) + 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)