diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py
index 9cc5100..1bbc4cd 100644
--- a/src/ttmask/cone.py
+++ b/src/ttmask/cone.py
@@ -12,7 +12,7 @@ def cone(
     sidelength: int, 
     cone_height: float, 
     cone_base_diameter: float, 
-    soft_edge_width: int, 
+    soft_edge_width: int,
     pixel_size: float,
     centering: str
 ) -> np.ndarray:
@@ -69,3 +69,4 @@ def cone_cli(
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:
         mrc.set_data(mask.astype(np.float32))
+        mrc.voxel_size = pixel_size
diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py
index a098a8f..0b063d4 100644
--- a/src/ttmask/cube.py
+++ b/src/ttmask/cube.py
@@ -9,13 +9,13 @@
 
 def cube(
     sidelength: int, 
-    cube_sidelength: float, 
-    soft_edge_width: float, 
-    pixel_size: float,
+    cube_sidelength: float,
     wall_thickness: float,
+    soft_edge_width: int,
+    pixel_size: float,
     centering: str
 ) -> np.ndarray:
-    # establish our coordinate system and empty mask
+     # 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)
@@ -39,14 +39,15 @@ def cube(
 def cube_cli(
     sidelength: int = typer.Option(...),
     cube_sidelength: float = typer.Option(...),
-    soft_edge_width: float = typer.Option(0),
+    wall_thickness: float = typer.Option(0),
+    soft_edge_width: int = typer.Option(0),
     pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("cube.mrc")),
-    wall_thickness: float = typer.Option(0),
-    centering: str = typer.Option("standard"),
+    centering: float = typer.Option("standard"),
 ):
-    mask = cube(sidelength, cube_sidelength, soft_edge_width, pixel_size, wall_thickness, centering)
+    mask = cube(sidelength, cube_sidelength, wall_thickness, soft_edge_width, pixel_size, centering)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:
         mrc.set_data(mask.astype(np.float32))
+        mrc.voxel_size = pixel_size
diff --git a/src/ttmask/cuboid.py b/src/ttmask/cuboid.py
index f12caab..a9dee50 100644
--- a/src/ttmask/cuboid.py
+++ b/src/ttmask/cuboid.py
@@ -11,10 +11,10 @@
 
 def cuboid(
     sidelength: int, 
-    cuboid_sidelengths: Tuple[float, float, float], 
-    soft_edge_width: float, 
-    pixel_size: float,
+    cuboid_sidelengths: Tuple[float, float, float],
     wall_thickness: float,
+    soft_edge_width: int,
+    pixel_size: float,
     centering: str
 ) -> np.ndarray:
     # establish our coordinate system and empty mask
@@ -41,14 +41,15 @@ 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),
     wall_thickness: float = typer.Option(0),
+    soft_edge_width: int = typer.Option(0),
+    pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("cuboid.mrc")),
     centering: str = typer.Option("standard"),
 ):
-    mask = cuboid(sidelength, cuboid_sidelengths, soft_edge_width, pixel_size,wall_thickness, centering)
+    mask = cuboid(sidelength, cuboid_sidelengths, wall_thickness, soft_edge_width, pixel_size, centering)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:
         mrc.set_data(mask.astype(np.float32))
+        mrc.voxel_size = pixel_size
diff --git a/src/ttmask/curved_surface.py b/src/ttmask/curved_surface.py
index 7311ebc..0ab4537 100644
--- a/src/ttmask/curved_surface.py
+++ b/src/ttmask/curved_surface.py
@@ -9,10 +9,10 @@
 
 def curved_surface(
     sidelength: int, 
-    fit_sphere_diameter: float, 
-    soft_edge_width: int, 
-    pixel_size: float,
+    fit_sphere_diameter: float,
     surface_thickness: float,
+    soft_edge_width: int,
+    pixel_size: float,
     centering: str
 ) -> np.ndarray:
     sphere_radius = fit_sphere_diameter / 2
@@ -44,14 +44,15 @@ def curved_surface(
 def curved_surface_cli(
     sidelength: int = typer.Option(...),
     fit_sphere_diameter: float = typer.Option(...),
+    surface_thickness: 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(...),
     centering: str = typer.Option("standard"),
 ):
-    mask = curved_surface(sidelength, fit_sphere_diameter, soft_edge_width, pixel_size, surface_thickness, centering)
+    mask = curved_surface(sidelength, fit_sphere_diameter, surface_thickness, soft_edge_width, pixel_size, centering)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:
         mrc.set_data(mask.astype(np.float32))
+        mrc.voxel_size = pixel_size
diff --git a/src/ttmask/cylinder.py b/src/ttmask/cylinder.py
index 268a9b1..7030646 100644
--- a/src/ttmask/cylinder.py
+++ b/src/ttmask/cylinder.py
@@ -61,3 +61,4 @@ def cylinder_cli(
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:
         mrc.set_data(mask.astype(np.float32))
+        mrc.voxel_size = pixel_size
\ No newline at end of file
diff --git a/src/ttmask/ellipsoid.py b/src/ttmask/ellipsoid.py
index 54a340d..2f03222 100644
--- a/src/ttmask/ellipsoid.py
+++ b/src/ttmask/ellipsoid.py
@@ -12,10 +12,10 @@
 
 def ellipsoid(
     sidelength: int, 
-    ellipsoid_dimensions: Tuple[float, float, float], 
-    soft_edge_width: float,
-    pixel_size: float,
+    ellipsoid_dimensions: Tuple[float, float, float],
     wall_thickness: float,
+    soft_edge_width: int,
+    pixel_size: float,
     centering: str
 ) -> np.ndarray:
     # establish our coordinate system and empty mask
@@ -54,13 +54,15 @@ def 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),
+    wall_thickness: float = typer.Option(0),
+    soft_edge_width: int = 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, centering)
+    mask = ellipsoid(sidelength, ellipsoid_dimensions, wall_thickness, soft_edge_width, pixel_size, centering)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:
         mrc.set_data(mask.astype(np.float32))
+        mrc.voxel_size = pixel_size
diff --git a/src/ttmask/map2mask.py b/src/ttmask/map2mask.py
index 38f37fa..1ad5fab 100644
--- a/src/ttmask/map2mask.py
+++ b/src/ttmask/map2mask.py
@@ -9,9 +9,9 @@
 
 def mask_from_map(
     map_data: np.ndarray, 
-    binarization_threshold: float, 
-    soft_edge_width: int, 
-    padding_width: int
+    binarization_threshold: float,
+    padding_width: int,
+    soft_edge_width: int
 ) -> np.ndarray:
 
     above_threshold = map_data >= binarization_threshold
@@ -29,15 +29,15 @@ def mask_from_map(
 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),
+    soft_edge_width: int = typer.Option(0),
+    pixel_size: float = typer.Option(...),
+    output_mask: Path = typer.Option(Path("mask.mrc")),
 ):
     with mrcfile.open(input_map, permissive=True) as mrc:
         data = mrc.data
-    mask = mask_from_map(data, binarization_threshold, soft_edge_width, padding_width)
+    mask = mask_from_map(data, binarization_threshold, padding_width, soft_edge_width)
 
-    # 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
+        mrc.set_data(mask.astype(np.float32))
+        mrc.voxel_size = pixel_size
\ No newline at end of file
diff --git a/src/ttmask/soft_edge.py b/src/ttmask/soft_edge.py
index 490c212..9b25cfb 100644
--- a/src/ttmask/soft_edge.py
+++ b/src/ttmask/soft_edge.py
@@ -2,7 +2,7 @@
 from scipy.ndimage import distance_transform_edt
 
 
-def add_soft_edge(mask: np.ndarray, width: float) -> np.ndarray:
+def add_soft_edge(mask: np.ndarray, width: int) -> np.ndarray:
     distance_from_edge = distance_transform_edt(mask == 0)
     boundary_pixels = (distance_from_edge <= width) & (distance_from_edge != 0)
     normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / width) * np.pi
diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py
index 2bf5917..cf64b48 100644
--- a/src/ttmask/sphere.py
+++ b/src/ttmask/sphere.py
@@ -9,10 +9,10 @@
 
 def sphere(
     sidelength: int, 
-    sphere_diameter: float, 
-    soft_edge_width: int, 
-    pixel_size: float, 
+    sphere_diameter: float,
     wall_thickness: float,
+    soft_edge_width: int,
+    pixel_size: float,
     centering: str
 ) -> np.ndarray:
     sphere_radius = sphere_diameter / 2
@@ -41,14 +41,15 @@ def sphere(
 def sphere_cli(
     sidelength: int = typer.Option(...),
     sphere_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("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, centering)
+    mask = sphere(sidelength, sphere_diameter, wall_thickness, soft_edge_width, pixel_size, centering)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:
         mrc.set_data(mask.astype(np.float32))
+        mrc.voxel_size = pixel_size
diff --git a/src/ttmask/tube.py b/src/ttmask/tube.py
index ce23573..fc20108 100644
--- a/src/ttmask/tube.py
+++ b/src/ttmask/tube.py
@@ -46,11 +46,14 @@ def tube_cli(
     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")),
     centering: str = typer.Option("standard"),
 ):
-    mask = tube(sidelength, tube_height, tube_diameter, wall_thickness, centering)
+    mask = tube(sidelength, tube_height, tube_diameter, wall_thickness, soft_edge_width, pixel_size, centering)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:
         mrc.set_data(mask.astype(np.float32))
+        mrc.voxel_size = pixel_size
diff --git a/tests/test_functions.py b/tests/test_functions.py
index 7b73e57..6dd9e6e 100644
--- a/tests/test_functions.py
+++ b/tests/test_functions.py
@@ -8,14 +8,14 @@ def test_cone():
     assert mask.sum() < np.pi * 25**2 * 50 # Volume of cylinder
 
 def test_cube():
-    mask = cube(100, 50, 0, 1, 0, "standard")
+    mask = cube(100, 50, 0, 0, 1, "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, "standard")
+    mask = cuboid(100, (50,40,30), 0, 0, 1, "standard")
     assert mask.shape == (100, 100, 100)
     # Test against volume of cuboid +- center and subpixel issues
     assert mask.sum() > 48 * 38 * 28
@@ -34,14 +34,14 @@ def test_cylinder():
     assert mask.sum() < np.pi * 25**2 * 51 # Volume of cylinder
 
 def test_ellipsoid():
-    mask = ellipsoid(100, (50,40,30), 0, 1,0, "standard")
+    mask = ellipsoid(100, (50,40,30), 0, 0, 1, "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, "standard")
+    mask = sphere(100, 50, 0, 0, 1, "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