feat: idx

fvi.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 
-def fast_voxel_intersect(start, end, origin, step) -> tuple[list, list]:
+def fast_voxel_intersect(start, end, origin, step, shape) -> tuple[list, list, list]:
     """Compute the voxels intersected by a line segment.
 
     Args:
@@ -13,6 +13,7 @@ def fast_voxel_intersect(start, end, origin, step) -> tuple[list, list]:
     Returns:
         list: list of intersection points
         list: list of intersected voxels
+        list: list of voxels idx
     """
 
     # Convert to numpy arrays
@@ -28,17 +29,18 @@ def fast_voxel_intersect(start, end, origin, step) -> tuple[list, list]:
     # Initialize list of intersected voxels
     intersections = []
     voxels = []
+    voxels_idx = []
 
     # Compute direction of line segment
     direction = end - start
     global_distance = np.linalg.norm(direction, axis=0)
     if global_distance == 0:
-        return intersections
+        return intersections, voxels, voxels_idx
     direction = direction / global_distance
+    non_zero_direction = (direction != 0).argmax()
 
     # Compute the sign of the direction
     direction_signs = np.sign(direction)
-    is_positive = direction_signs > 0
     is_negative = direction_signs < 0
 
     # Initialize current position to start
@@ -48,7 +50,7 @@ def fast_voxel_intersect(start, end, origin, step) -> tuple[list, list]:
     while True:
         # Compute the distance to the next boundaries
         next_boundaries = np.divide(position + step * direction_signs, step)
-        distances = (is_positive * np.floor(next_boundaries) +
+        distances = ((1 - is_negative) * np.floor(next_boundaries) +
                      is_negative * np.ceil(next_boundaries)) * step - position
 
         # Determine the nearest boundary to be reached
@@ -57,7 +59,7 @@ def fast_voxel_intersect(start, end, origin, step) -> tuple[list, list]:
         clothest_boundary_distance = boundary_distances[clothest_boundary]
 
         # Check if we are done
-        distance_to_end = abs((end[0] - position[0]) / direction[0])
+        distance_to_end = abs((end[non_zero_direction] - position[non_zero_direction]) / direction[non_zero_direction])
         if clothest_boundary_distance > distance_to_end:
             break
 
@@ -69,22 +71,25 @@ def fast_voxel_intersect(start, end, origin, step) -> tuple[list, list]:
             position[clothest_boundary] / step[clothest_boundary]) * step[clothest_boundary]
 
         # Get corresponding voxel
-        boundary_reached_negativly = np.zeros(start.shape, dtype=bool)
-        boundary_reached_negativly[clothest_boundary] = is_negative[clothest_boundary]
-        voxel = np.floor(position) - boundary_reached_negativly * step
+        on_boundary = np.mod(position, step) == 0
+        voxel = np.floor(position) - is_negative * on_boundary * step
 
         # Add voxel to list
+        idx = np.floor_divide(voxel, step).astype(int)
+        if np.any(idx < 0) or np.any(idx >= shape):
+            continue
         intersections.append(position + origin)
         voxels.append(voxel + origin)
+        voxels_idx.append(idx)
 
-    return intersections, voxels
+    return intersections, voxels, voxels_idx
 
 
 if __name__ == '__main__':
     import matplotlib.pyplot as plt
 
     def update_figure():
-        positions, voxels = fast_voxel_intersect(start, end, origin, step)
+        positions, voxels, voxels_idx = fast_voxel_intersect(start, end, origin, step, shape)
 
         plt.clf()
 
@@ -94,6 +99,15 @@ if __name__ == '__main__':
                      [voxel[1], voxel[1], voxel[1] + step[1], voxel[1] + step[1]],
                      color='#e25', alpha=0.5)
         
+        for voxel_id in voxels_idx:
+            plt.fill([
+                        origin[0] + voxel_id[0] * step[0], origin[0] + (voxel_id[0] + 1) * step[0],
+                        origin[0] + (voxel_id[0] + 1) * step[0], origin[0] + voxel_id[0] * step[0]
+                    ], [
+                        origin[1] + voxel_id[1] * step[1], origin[1] + voxel_id[1] * step[1],
+                        origin[1] + (voxel_id[1] + 1) * step[1], origin[1] + (voxel_id[1] + 1) * step[1]
+                    ], color='#2e3', alpha=0.5)
+
         # Plot line segment
         plt.plot([start[0], end[0]], [start[1], end[1]], 'k-')
         plt.plot(start[0], start[1], 'go')
@@ -107,35 +121,31 @@ if __name__ == '__main__':
         plt.axis('equal')
         plt.xlim((-10, 10))
         plt.ylim((-10, 10))
-        xmin, xmax = plt.xlim()
-        ymin, ymax = plt.ylim()
-        plt.xticks(np.arange(xmin + origin[0],
-                             xmax + origin[0] + step[0], step[0]))
-        plt.yticks(np.arange(ymin + origin[1],
-                             ymax + origin[1] + step[1], step[1]))
+        plt.xticks(origin[0] + step[0] * np.arange(shape[0] + 1))
+        plt.yticks(origin[1] + step[1] * np.arange(shape[1] + 1))
         plt.grid()
         plt.draw()
 
-    def onclick(event):
+    def onkey(event):
         global start, end
-        # if event.button == 1:
-        #     start = np.array([event.xdata, event.ydata])
-        # elif event.button == 3:
-        #     end = np.array([event.xdata, event.ydata])
-        start = np.random.rand(2) * 10 - 5
-        end = np.random.rand(2) * 10 - 5
+        if event.key == ' ':
+            start = np.random.rand(2) * 20 - 10
+            end = np.random.rand(2) * 20 - 10
             update_figure()
 
     # Define voxel grid
-    origin = np.array([.1, -.3])
+    origin = np.array([-5., -5.])
     step = np.array([1.0, 1.0])
+    shape = (10, 10)
 
     # Define segment
-    start = np.random.rand(2) * 10 - 5
-    end = np.random.rand(2) * 10 - 5
+    # start = np.random.rand(2) * 20 - 10
+    # end = np.random.rand(2) * 20 - 10
+    start = np.array([0., 0.])
+    end = np.array([4., -4.])
 
     # Plot
     fig = plt.figure()
-    fig.canvas.mpl_connect('button_press_event', onclick)
+    fig.canvas.mpl_connect('key_press_event', onkey)
     update_figure()
     plt.show()