Allowing a custom centering option. (#53)

For example, use --centering custom --center 30 50 20

src/ttmask/box_setup.py

@@ -2,10 +2,12 @@
 import einops
 
 
-def box_setup(sidelength: int, centering: str) -> np.ndarray:
+def box_setup(sidelength: int, centering: str, custom_center: tuple):
     c = sidelength // 2
     if centering == "visual" and sidelength % 2 == 0:
         center = np.array([c, c, c]) - 0.5
+    elif centering == "custom":
+        center = np.array([custom_center])
     else:
         center = np.array([c, c, c])
     mask = np.zeros(shape=(sidelength, sidelength, sidelength), dtype=np.float32)
@@ -14,6 +16,6 @@ def box_setup(sidelength: int, centering: str) -> np.ndarray:
     coordinates = np.indices([sidelength, sidelength, sidelength])
     coordinates = einops.rearrange(coordinates, 'zyx d h w -> d h w zyx')
 
-    #coordinates now expressed relative to center
+    # coordinates now expressed relative to center
     coordinates_centered = coordinates - center
-    return (coordinates_centered, mask)
+    return coordinates_centered, mask

src/ttmask/cone.py

@@ -1,4 +1,7 @@
 from pathlib import Path
+from typing import Tuple
+from typing_extensions import Annotated
+
 import numpy as np
 import einops
 import typer
@@ -14,10 +17,11 @@ def cone(
     cone_base_diameter: float, 
     soft_edge_width: int,
     pixel_size: float,
-    centering: str
+    centering: str,
+    center: tuple
 ) -> np.ndarray:
     # establish our coordinate system and empty mask
-    coordinates_centered, mask = box_setup(sidelength, centering)
+    coordinates_centered, mask = box_setup(sidelength, centering, center)
     # 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')
@@ -63,8 +67,9 @@ def cone_cli(
     pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("cone.mrc")),
     centering: str = typer.Option("standard"),
+    center: Annotated[Tuple[int, int, int], typer.Option()] = (50, 50, 50)
 ):
-    mask = cone(sidelength, cone_height, cone_base_diameter, soft_edge_width, pixel_size, centering)
+    mask = cone(sidelength, cone_height, cone_base_diameter, soft_edge_width, pixel_size, centering, center)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:

src/ttmask/cube.py

@@ -1,4 +1,7 @@
 from pathlib import Path
+from typing import Tuple
+from typing_extensions import Annotated
+
 import numpy as np
 import typer
 import mrcfile
@@ -13,10 +16,11 @@ def cube(
     wall_thickness: float,
     soft_edge_width: int,
     pixel_size: float,
-    centering: str
+    centering: str,
+    center: tuple
 ) -> np.ndarray:
      # establish our coordinate system and empty mask
-    coordinates_centered, mask = box_setup(sidelength, centering)
+    coordinates_centered, mask = box_setup(sidelength, centering, center)
     #converting relative coordinates to xyz distances (i.e. not a negative number) :
     xyz_distances = np.abs(coordinates_centered)
 
@@ -43,9 +47,10 @@ def cube_cli(
     soft_edge_width: int = typer.Option(0),
     pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("cube.mrc")),
-    centering: float = typer.Option("standard"),
+    centering: str = typer.Option("standard"),
+    center: Annotated[Tuple[int, int, int], typer.Option()] = (50, 50, 50)
 ):
-    mask = cube(sidelength, cube_sidelength, wall_thickness, soft_edge_width, pixel_size, centering)
+    mask = cube(sidelength, cube_sidelength, wall_thickness, soft_edge_width, pixel_size, centering, center)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:

src/ttmask/cuboid.py

@@ -1,8 +1,9 @@
 from pathlib import Path
-import numpy as np
-import typer
 from typing import Tuple
 from typing_extensions import Annotated
+
+import numpy as np
+import typer
 import mrcfile
 
 from ._cli import cli
@@ -15,10 +16,11 @@ def cuboid(
     wall_thickness: float,
     soft_edge_width: int,
     pixel_size: float,
-    centering: str
+    centering: str,
+    center: tuple
 ) -> np.ndarray:
     # establish our coordinate system and empty mask
-    coordinates_centered, mask = box_setup(sidelength, centering)
+    coordinates_centered, mask = box_setup(sidelength, centering, center)
     #converting relative coordinates to xyz distances (i.e. not a negative number) :
     xyz_distances = np.abs(coordinates_centered)
 
@@ -46,8 +48,9 @@ def cuboid_cli(
     pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("cuboid.mrc")),
     centering: str = typer.Option("standard"),
+    center: Annotated[Tuple[int, int, int], typer.Option()] = (50, 50, 50)
 ):
-    mask = cuboid(sidelength, cuboid_sidelengths, wall_thickness, soft_edge_width, pixel_size, centering)
+    mask = cuboid(sidelength, cuboid_sidelengths, wall_thickness, soft_edge_width, pixel_size, centering, center)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:

src/ttmask/cylinder.py

@@ -1,11 +1,12 @@
 from pathlib import Path
+from typing import Tuple
+from typing_extensions import Annotated
 
 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
 
@@ -16,12 +17,13 @@ def cylinder(
     wall_thickness: float, 
     soft_edge_width: int,
     pixel_size: float,
-    centering: str
+    centering: str,
+    center: tuple
 ) -> np.ndarray:
     cylinder_radius = cylinder_diameter / 2
 
     # establish our coordinate system and empty mask
-    coordinates_centered, mask = box_setup(sidelength, centering)
+    coordinates_centered, mask = box_setup(sidelength, centering, center)
 
     # converting relative coordinates to xyz distances (i.e. not a negative number) :
     xyz_distances = np.abs(coordinates_centered)
@@ -55,8 +57,9 @@ def cylinder_cli(
     pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("cylinder.mrc")),
     centering: str = typer.Option("standard"),
+    center: Annotated[Tuple[int, int, int], typer.Option()] = (50, 50, 50)
 ):
-    mask = cylinder(sidelength, cylinder_height, cylinder_diameter, wall_thickness, soft_edge_width, pixel_size, centering)
+    mask = cylinder(sidelength, cylinder_height, cylinder_diameter, wall_thickness, soft_edge_width, pixel_size, centering, center)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:

src/ttmask/ellipsoid.py

@@ -1,10 +1,10 @@
 from pathlib import Path
+from typing import Tuple
+from typing_extensions import Annotated
 
 import numpy as np
 import typer
 import mrcfile
-from typing import Tuple
-from typing_extensions import Annotated
 
 from .soft_edge import add_soft_edge
 from ._cli import cli
@@ -16,11 +16,14 @@ def ellipsoid(
     wall_thickness: float,
     soft_edge_width: int,
     pixel_size: float,
-    centering: str
+    centering: str,
+    center: tuple
 ) -> np.ndarray:
+
     # 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) :
+    coordinates_centered, mask = box_setup(sidelength, centering, center)
+
+    # converting relative coordinates to xyz distances (i.e. not a negative number) :
     xyz_distances = np.abs(coordinates_centered)
 
     #extract xyz magnitudes from xyz_distances
@@ -59,8 +62,9 @@ def ellipsoid_cli(
     pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("ellipsoid.mrc")),
     centering: str = typer.Option("standard"),
+    center: Annotated[Tuple[int, int, int], typer.Option()] = (50, 50, 50)
 ):
-    mask = ellipsoid(sidelength, ellipsoid_dimensions, wall_thickness, soft_edge_width, pixel_size, centering)
+    mask = ellipsoid(sidelength, ellipsoid_dimensions, wall_thickness, soft_edge_width, pixel_size, centering, center)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:

src/ttmask/sphere.py

@@ -1,4 +1,7 @@
 from pathlib import Path
+from typing import Tuple
+from typing_extensions import Annotated
+
 import numpy as np
 import typer
 import mrcfile
@@ -13,12 +16,13 @@ def sphere(
     wall_thickness: float,
     soft_edge_width: int,
     pixel_size: float,
-    centering: str
+    centering: str,
+    center: tuple
 ) -> np.ndarray:
     sphere_radius = sphere_diameter / 2
 
     # establish our coordinate system and empty mask
-    coordinates_centered, mask = box_setup(sidelength, centering)
+    coordinates_centered, mask, _ = box_setup(sidelength, centering, center)
 
     # determine distances of each pixel to the center
     distance_to_center = np.linalg.norm(coordinates_centered, axis=-1)
@@ -46,8 +50,9 @@ def sphere_cli(
     pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("sphere.mrc")),
     centering: str = typer.Option("standard"),
+    center: Annotated[Tuple[int, int, int], typer.Option()] = (50, 50, 50)
 ):
-    mask = sphere(sidelength, sphere_diameter, wall_thickness, soft_edge_width, pixel_size, centering)
+    mask = sphere(sidelength, sphere_diameter, wall_thickness, soft_edge_width, pixel_size, centering, center)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc:

src/ttmask/tube.py

@@ -1,9 +1,12 @@
 from pathlib import Path
+from typing import Tuple
+from typing_extensions import Annotated
 
 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
@@ -15,12 +18,14 @@ def tube(
     wall_thickness: float,
     soft_edge_width: int,
     pixel_size: float,
-    centering: str
+    centering: str,
+    center: tuple
 ) -> np.ndarray:
     tube_radius = tube_diameter / 2
 
     # establish our coordinate system and empty mask
-    coordinates_centered, mask = box_setup(sidelength, centering)
+    coordinates_centered, mask, = box_setup(sidelength, centering, center)
+
     #converting relative coordinates to xyz distances (i.e. not a negative number) :
     xyz_distances = np.abs(coordinates_centered)
 
@@ -50,8 +55,9 @@ def tube_cli(
     pixel_size: float = typer.Option(1),
     output: Path = typer.Option(Path("tube.mrc")),
     centering: str = typer.Option("standard"),
+    center: Annotated[Tuple[int, int, int], typer.Option()] = (50, 50, 50)
 ):
-    mask = tube(sidelength, tube_height, tube_diameter, wall_thickness, soft_edge_width, pixel_size, centering)
+    mask = tube(sidelength, tube_height, tube_diameter, wall_thickness, soft_edge_width, pixel_size, centering, center)
 
     # Save the mask to an MRC file
     with mrcfile.new(output, overwrite=True) as mrc: