Create functions for mask creation that return a ndarray (#47)

* Refactor functions + Tests

* Wow Copilot was freelancing

* More copilot fixes

* Copilot fix cuboid

* Fix curved surface

* Final copilot fixes

* More wiggle room in tests

src/ttmask/cone.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-
 import numpy as np
 import einops
 import typer
@@ -9,16 +8,13 @@
 from .soft_edge import add_soft_edge
 from .box_setup import box_setup
 
-
-@cli.command(name='cone')
 def cone(
-    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"))
-):
+    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 :
@@ -55,4 +51,19 @@ def cone(
     mask = np.roll(mask, z_shift, axis=0)
     mask = add_soft_edge(mask, soft_edge_width)
 
-    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
+    return mask
+
+@cli.command(name='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")),
+):
+    mask = cone(sidelength, cone_height, cone_base_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))

src/ttmask/cube.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-
 import numpy as np
 import typer
 import mrcfile
@@ -8,17 +7,14 @@
 from ._cli import cli
 from .box_setup import box_setup
 
-
-@cli.command(name='cube')
 def cube(
-    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
+    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)
@@ -36,5 +32,19 @@ def cube(
     #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='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),
+):
+    mask = cube(sidelength, cube_sidelength, 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))

src/ttmask/cuboid.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-
 import numpy as np
 import typer
 from typing import Tuple
@@ -10,16 +9,13 @@
 from .soft_edge import add_soft_edge
 from .box_setup import box_setup
 
-
-@cli.command(name='cuboid')
 def cuboid(
-    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),
-):
+    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) :
@@ -37,6 +33,20 @@ def cuboid(
 
     # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
+
+    return mask
+
+@cli.command(name='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),
+    wall_thickness: float = typer.Option(0),
+    output: Path = typer.Option(Path("cuboid.mrc")),
+):
+    mask = cuboid(sidelength, cuboid_sidelengths, soft_edge_width, pixel_size,wall_thickness)
 
-    # 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))

src/ttmask/curved_surface.py

@@ -1,6 +1,4 @@
 from pathlib import Path
-
-
 import numpy as np
 import typer
 import mrcfile
@@ -9,15 +7,13 @@
 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(...),
-):
+    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
@@ -41,5 +37,19 @@ def curved_surface(
     # 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")),
+    surface_thickness: float = typer.Option(...),
+):
+    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:
+        mrc.set_data(mask.astype(np.float32))

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))

src/ttmask/ellipsoid.py

@@ -10,17 +10,13 @@
 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),
-):
+    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 +47,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))

src/ttmask/map2mask.py

@@ -7,18 +7,12 @@
 from .soft_edge import add_soft_edge
 from .add_padding import add_padding
 
-@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")),
-    pixel_size: float = typer.Option(...),
-    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)
+def mask_from_map(
+    map_data: np.ndarray, 
+    binarization_threshold: float, 
+    soft_edge_width: int, 
+    padding_width: int
+) -> np.ndarray:
 
     above_threshold = map_data >= binarization_threshold
     below_threshold = map_data < binarization_threshold
@@ -29,4 +23,21 @@ def map2mask(
     padded_mask = add_padding(map_data, padding_width)
     mask = add_soft_edge(padded_mask, soft_edge_width)
 
-    mrcfile.write(output_mask, mask, voxel_size=pixel_size, overwrite=True)
+    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")),
+    pixel_size: float = typer.Option(...),
+    soft_edge_width: int = typer.Option(0),
+    padding_width: int = typer.Option(0),
+):
+    with mrcfile.open(input_map, permissive=True) as mrc:
+        data = mrc.data
+    mask = mask_from_map(data, binarization_threshold, soft_edge_width, padding_width)
+
+    # Save the mask to an MRC file
+    with mrcfile.new(output_mask, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))