From a09e88a3ac13ed594e9f7a27312b87ad684a5a36 Mon Sep 17 00:00:00 2001
From: Miles Graham <miles.graham98@gmail.com>
Date: Mon, 20 May 2024 11:08:32 +0100
Subject: [PATCH] added cube

---
 src/ttmask/__init__.py |  2 ++
 src/ttmask/cube.py     | 44 ++++++++++++++++++++++++++++++++++++++++++
 src/ttmask/sphere.py   |  1 +
 3 files changed, 47 insertions(+)
 create mode 100644 src/ttmask/cube.py

diff --git a/src/ttmask/__init__.py b/src/ttmask/__init__.py
index e3e1304..1c9b5d5 100644
--- a/src/ttmask/__init__.py
+++ b/src/ttmask/__init__.py
@@ -13,4 +13,6 @@
 from .sphere import sphere
 from .cylinder import cylinder
 from .cuboid import cuboid
+from .cube import cube
+
 
diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py
new file mode 100644
index 0000000..4c11468
--- /dev/null
+++ b/src/ttmask/cube.py
@@ -0,0 +1,44 @@
+import numpy as np
+import einops
+import typer
+# from ._cli import cli
+from scipy.ndimage import distance_transform_edt
+import mrcfile
+import napari
+
+
+# @cli.command(name='cube')
+def cube(
+    boxsize: int = typer.Option(...),
+    cube_sidelength: float =typer.Option(...),
+    soft_edge_size: float = typer.Option(...),
+    mrc_voxel_size: float = typer.Option(...),
+):
+    c = boxsize // 2
+    center = np.array([c, c, c])
+    mask = np.zeros(shape=(boxsize, boxsize, boxsize), dtype=np.float32)
+
+    # 3d positions of all voxels
+    positions = np.indices([boxsize, boxsize, boxsize])
+    positions = einops.rearrange(positions, 'zyx d h w -> d h w zyx')
+
+    # calculate the distance between the center and every pixel position
+    print(center.shape)
+    print(positions.shape)
+
+    print('calculating distance')
+    difference = np.abs(positions - center)  # (100, 100, 100, 3)
+
+    idx = np.all(difference < (np.array(cube_sidelength) / 2), axis=-1)
+
+    mask[idx] = 1
+
+    distance_from_edge = distance_transform_edt(mask == 0)
+    boundary_pixels = (distance_from_edge <= soft_edge_size) & (distance_from_edge != 0)
+    normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / soft_edge_size) * np.pi
+
+    mask[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5)
+
+    mrcfile.write("cube.mrc", mask, voxel_size= mrc_voxel_size, overwrite=True)
+
+cube(100, 20, 3, 1.321)
\ No newline at end of file
diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py
index 8e7eda9..f41fa2d 100644
--- a/src/ttmask/sphere.py
+++ b/src/ttmask/sphere.py
@@ -40,3 +40,4 @@ def sphere(
     mask[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5)
 
     mrcfile.write("sphere.mrc", mask, voxel_size= mrc_voxel_size, overwrite=True)
+