neuro-sam 0.1.0__py3-none-any.whl

neuro_sam/brightest_path_lib/algorithm/waypointastar.py

@@ -0,0 +1,449 @@
+"""Advanced optimized A* search implementation with waypoint support for finding brightest paths.
+This extends the BidirectionalAStarSearch to support user-defined auxiliary waypoints with
+performance optimizations matching the core algorithm.
+"""
+
+import heapq
+import math
+import numpy as np
+from collections import defaultdict
+from typing import List, Tuple, Dict, Set, Any, Optional
+import numba as nb
+from numba import njit, prange, jit
+
+from neuro_sam.brightest_path_lib.algorithm.astar import (
+    BidirectionalAStarSearch, array_equal, euclidean_distance_scaled,
+    find_2D_neighbors_optimized, find_3D_neighbors_optimized
+)
+from neuro_sam.brightest_path_lib.cost import Reciprocal
+from neuro_sam.brightest_path_lib.heuristic import Euclidean
+from neuro_sam.brightest_path_lib.image import ImageStats
+from neuro_sam.brightest_path_lib.input import CostFunction, HeuristicFunction
+from neuro_sam.brightest_path_lib.node import Node
+
+
+# Numba-optimized function to check if all points are on the same z-plane
+@nb.njit(cache=True)
+def check_same_z_plane(start_point, goal_point, waypoints):
+    """Check if all points are on the same z-plane for 3D optimization"""
+    # If points are 2D, return True
+    if len(start_point) == 2:
+        return True
+        
+    # For 3D points, check z-coordinates
+    if len(start_point) == 3:
+        z_val = start_point[0]
+        
+        # Check if goal has same z
+        if goal_point[0] != z_val:
+            return False
+            
+        # Check all waypoints
+        for i in range(len(waypoints)):
+            if waypoints[i][0] != z_val:
+                return False
+                
+        # All points have same z
+        return True
+        
+    # For higher dimensions, don't use the optimization
+    return False
+
+
+# Numba-optimized function to create 2D points from 3D points
+@nb.njit(cache=True)
+def convert_3d_to_2d_points(start_point, goal_point, waypoints):
+    """Convert 3D points to 2D by removing z-coordinate"""
+    start_2d = np.array([start_point[1], start_point[2]], dtype=np.int32)
+    goal_2d = np.array([goal_point[1], goal_point[2]], dtype=np.int32)
+    
+    # Convert waypoints
+    waypoints_2d = np.empty((len(waypoints), 2), dtype=np.int32)
+    for i in range(len(waypoints)):
+        waypoints_2d[i, 0] = waypoints[i][1]
+        waypoints_2d[i, 1] = waypoints[i][2]
+        
+    return start_2d, goal_2d, waypoints_2d
+
+
+# Numba-optimized function to convert 2D path back to 3D
+@nb.njit(cache=True)
+def convert_2d_path_to_3d(path_2d, z_value):
+    """Convert a 2D path back to 3D by adding z-coordinate"""
+    path_3d = np.empty((len(path_2d), 3), dtype=np.int32)
+    
+    for i in range(len(path_2d)):
+        path_3d[i, 0] = z_value
+        path_3d[i, 1] = path_2d[i][0]
+        path_3d[i, 2] = path_2d[i][1]
+        
+    return path_3d
+
+
+class WaypointBidirectionalAStarSearch:
+    """Advanced bidirectional A* search implementation with waypoint support
+    
+    This implementation allows users to specify intermediate points that
+    the path should pass through, breaking down a complex search into
+    multiple simpler searches between consecutive points.
+    
+    Performance optimization: When waypoints are on the same z-plane, 
+    the algorithm automatically uses 2D pathfinding instead of 3D.
+    """
+
+    def __init__(
+        self,
+        image: np.ndarray,
+        start_point: np.ndarray,
+        goal_point: np.ndarray,
+        waypoints: List[np.ndarray] = None,
+        scale: Tuple = (1.0, 1.0),
+        cost_function: CostFunction = CostFunction.RECIPROCAL,
+        heuristic_function: HeuristicFunction = HeuristicFunction.EUCLIDEAN,
+        open_nodes=None,
+        use_hierarchical: bool = False,
+        weight_heuristic: float = 1.0
+    ):
+        """Initialize waypoint-enabled A* search
+        
+        Parameters
+        ----------
+        image : numpy ndarray
+            The image to search
+        start_point, goal_point : numpy ndarray
+            Start and goal coordinates
+        waypoints : List[numpy ndarray], optional
+            List of intermediate points that the path must pass through,
+            in the order they should be visited
+        scale : tuple
+            Image scale factors
+        cost_function, heuristic_function : Enum
+            Functions to use for cost and heuristic
+        open_nodes : Queue, optional
+            Queue for visualization
+        use_hierarchical : bool
+            Whether to use hierarchical search for large images
+        weight_heuristic : float
+            Weight for heuristic (> 1.0 makes search faster but less optimal)
+        """
+        self._validate_inputs(image, start_point, goal_point, waypoints)
+        
+        # Basic parameters
+        self.image = image
+        self.image_stats = ImageStats(image)
+        self.start_point = np.round(start_point).astype(np.int32)
+        self.goal_point = np.round(goal_point).astype(np.int32)
+        
+        # Process and validate waypoints if provided
+        if waypoints and len(waypoints) > 0:
+            self.waypoints = np.array([np.round(wp).astype(np.int32) for wp in waypoints])
+        else:
+            # Empty array with correct shape for numba compatibility
+            self.waypoints = np.empty((0, len(self.start_point)), dtype=np.int32)
+        
+        self.scale = scale
+        self.open_nodes = open_nodes
+        self.use_hierarchical = use_hierarchical
+        self.weight_heuristic = weight_heuristic
+        
+        if cost_function == CostFunction.RECIPROCAL:
+            self.cost_function = Reciprocal(
+                min_intensity=self.image_stats.min_intensity, 
+                max_intensity=self.image_stats.max_intensity)
+        
+        if heuristic_function == HeuristicFunction.EUCLIDEAN:
+            self.heuristic_function = Euclidean(scale=self.scale)
+        
+        # State variables
+        self.is_canceled = False
+        self.found_path = False
+        self.evaluated_nodes = 0
+        self.result = []
+        self.segment_results = []  # Store individual segment paths
+        self.segment_evaluated_nodes = []  # Track nodes evaluated per segment
+        
+        # Determine if we should use 2D mode (all points have same z-coordinate)
+        self.use_2d_mode = check_same_z_plane(
+            self.start_point, self.goal_point, self.waypoints)
+        
+        # For 2D optimization when applicable
+        self.z_value = self.start_point[0] if len(self.start_point) == 3 else None
+        
+    def _validate_inputs(
+        self,
+        image: np.ndarray,
+        start_point: np.ndarray,
+        goal_point: np.ndarray,
+        waypoints: Optional[List[np.ndarray]] = None
+    ):
+        """Validate input parameters"""
+        if image is None or start_point is None or goal_point is None:
+            raise TypeError("Image, start_point, and goal_point cannot be None")
+        if len(image) == 0 or len(start_point) == 0 or len(goal_point) == 0:
+            raise ValueError("Image, start_point, and goal_point cannot be empty")
+            
+        # Validate waypoints if provided
+        if waypoints:
+            for i, wp in enumerate(waypoints):
+                if wp is None:
+                    raise TypeError(f"Waypoint {i} cannot be None")
+                if len(wp) == 0:
+                    raise ValueError(f"Waypoint {i} cannot be empty")
+                if len(wp) != len(start_point):
+                    raise ValueError(f"Waypoint {i} dimensions must match start_point dimensions")
+
+    @property
+    def found_path(self) -> bool:
+        return self._found_path
+
+    @found_path.setter
+    def found_path(self, value: bool):
+        if value is None:
+            raise TypeError
+        self._found_path = value
+
+    @property
+    def is_canceled(self) -> bool:
+        return self._is_canceled
+
+    @is_canceled.setter
+    def is_canceled(self, value: bool):
+        if value is None:
+            raise TypeError
+        self._is_canceled = value
+
+    def search(self, verbose: bool = False) -> List[np.ndarray]:
+        """Perform A* search with waypoints
+        
+        This method breaks down the search into multiple segments:
+        start→waypoint₁, waypoint₁→waypoint₂, ..., waypointₙ→goal
+        
+        Performance optimization: Uses 2D mode when all points are on the same z-plane
+        
+        Returns
+        -------
+        List[np.ndarray]
+            Complete path from start to goal through all waypoints
+        """
+        # Reset state
+        self.result = []
+        self.segment_results = []
+        self.segment_evaluated_nodes = []
+        self.evaluated_nodes = 0
+        
+        # Check if we can use 2D mode for performance optimization
+        if self.use_2d_mode and len(self.start_point) == 3:
+            if verbose:
+                print("Using 2D mode optimization (all points on same z-plane)")
+            return self._search_2d_mode(verbose)
+        else:
+            return self._search_normal_mode(verbose)
+    
+    def _search_2d_mode(self, verbose: bool = False) -> List[np.ndarray]:
+        """Perform search in 2D mode for performance optimization
+        
+        This extracts a 2D slice from the 3D image when all points are on the same z-plane
+        """
+        # Extract 2D slice and convert points to 2D
+        image_2d = self.image[self.z_value]
+        start_2d, goal_2d, waypoints_2d = convert_3d_to_2d_points(
+            self.start_point, self.goal_point, self.waypoints)
+        
+        # Create a list of all points in order, including start and goal
+        wp_list = [waypoints_2d[i] for i in range(len(waypoints_2d))]
+        all_points_2d = [start_2d] + wp_list + [goal_2d]
+        
+        # Track overall success
+        overall_success = True
+        
+        if verbose:
+            if len(wp_list) > 0:
+                print(f"Starting 2D search with {len(wp_list)} waypoints")
+            else:
+                print("Starting 2D search (no waypoints)")
+        
+        # Process each segment (start to first waypoint, waypoint to waypoint, last waypoint to goal)
+        for i in range(len(all_points_2d) - 1):
+            if self.is_canceled:
+                return []
+                
+            point_a = all_points_2d[i]
+            point_b = all_points_2d[i+1]
+            
+            if verbose:
+                if i == 0:
+                    print(f"Searching from start to {'waypoint' if i+1 < len(all_points_2d)-1 else 'goal'} {i+1}")
+                elif i+1 == len(all_points_2d)-1:
+                    print(f"Searching from waypoint {i} to goal")
+                else:
+                    print(f"Searching from waypoint {i} to waypoint {i+1}")
+            
+            # Create A* search for this segment
+            segment_search = BidirectionalAStarSearch(
+                image=image_2d,
+                start_point=point_a,
+                goal_point=point_b,
+                scale=self.scale[:2],  # Only use x,y scale
+                cost_function=CostFunction.RECIPROCAL,
+                heuristic_function=HeuristicFunction.EUCLIDEAN,
+                open_nodes=self.open_nodes,
+                use_hierarchical=self.use_hierarchical,
+                weight_heuristic=self.weight_heuristic
+            )
+            
+            # Run the search for this segment
+            segment_path_2d = segment_search.search(verbose=verbose)
+            
+            # Track stats and results
+            self.segment_evaluated_nodes.append(segment_search.evaluated_nodes)
+            self.evaluated_nodes += segment_search.evaluated_nodes
+            
+            # Check if segment search was successful
+            if segment_search.found_path and len(segment_path_2d) > 0:
+                # Convert 2D path back to 3D
+                segment_path_3d = convert_2d_path_to_3d(segment_path_2d, self.z_value)
+                # Store as list of arrays for compatibility
+                segment_path_list = [segment_path_3d[i] for i in range(len(segment_path_3d))]
+                self.segment_results.append(segment_path_list)
+                
+                if verbose:
+                    print(f"Found path for segment {i+1} ({len(segment_path_2d)} points, {segment_search.evaluated_nodes} nodes evaluated)")
+            else:
+                if verbose:
+                    print(f"Failed to find path for segment {i+1}")
+                overall_success = False
+                break
+        
+        # If all segments were successful, combine the paths
+        if overall_success:
+            self._construct_complete_path()
+            self.found_path = True
+            
+            if verbose:
+                print(f"Complete path found: {len(self.result)} points, {self.evaluated_nodes} total nodes evaluated")
+        else:
+            if verbose:
+                print("Failed to find a complete path through all waypoints")
+        
+        return self.result
+    
+    def _search_normal_mode(self, verbose: bool = False) -> List[np.ndarray]:
+        """Perform search in normal mode (without special 2D optimization)"""
+        # Create a list of all points in order, including start and goal
+        wp_list = [self.waypoints[i] for i in range(len(self.waypoints))]
+        all_points = [self.start_point] + wp_list + [self.goal_point]
+        
+        # Track overall success
+        overall_success = True
+        
+        if verbose:
+            if len(wp_list) > 0:
+                print(f"Starting search with {len(wp_list)} waypoints")
+            else:
+                print("Starting search (no waypoints)")
+        
+        # Process each segment (start to first waypoint, waypoint to waypoint, last waypoint to goal)
+        for i in range(len(all_points) - 1):
+            if self.is_canceled:
+                return []
+                
+            point_a = all_points[i]
+            point_b = all_points[i+1]
+            
+            if verbose:
+                if i == 0:
+                    print(f"Searching from start to {'waypoint' if i+1 < len(all_points)-1 else 'goal'} {i+1}")
+                elif i+1 == len(all_points)-1:
+                    print(f"Searching from waypoint {i} to goal")
+                else:
+                    print(f"Searching from waypoint {i} to waypoint {i+1}")
+            
+            # Create A* search for this segment
+            segment_search = BidirectionalAStarSearch(
+                image=self.image,
+                start_point=point_a,
+                goal_point=point_b,
+                scale=self.scale,
+                cost_function=CostFunction.RECIPROCAL,
+                heuristic_function=HeuristicFunction.EUCLIDEAN,
+                open_nodes=self.open_nodes,
+                use_hierarchical=self.use_hierarchical,
+                weight_heuristic=self.weight_heuristic
+            )
+            
+            # Run the search for this segment
+            segment_path = segment_search.search(verbose=verbose)
+            
+            # Track stats and results
+            self.segment_evaluated_nodes.append(segment_search.evaluated_nodes)
+            self.evaluated_nodes += segment_search.evaluated_nodes
+            
+            # Check if segment search was successful
+            if segment_search.found_path and len(segment_path) > 0:
+                self.segment_results.append(segment_path)
+                
+                if verbose:
+                    print(f"Found path for segment {i+1} ({len(segment_path)} points, {segment_search.evaluated_nodes} nodes evaluated)")
+            else:
+                if verbose:
+                    print(f"Failed to find path for segment {i+1}")
+                overall_success = False
+                break
+        
+        # If all segments were successful, combine the paths
+        if overall_success:
+            self._construct_complete_path()
+            self.found_path = True
+            
+            if verbose:
+                print(f"Complete path found: {len(self.result)} points, {self.evaluated_nodes} total nodes evaluated")
+        else:
+            if verbose:
+                print("Failed to find a complete path through all waypoints")
+        
+        return self.result
+    
+    def _construct_complete_path(self):
+        """Combine segment paths into a complete path
+        
+        This handles removing duplicate points at segment boundaries.
+        """
+        if not self.segment_results:
+            return
+            
+        # Start with the first segment
+        self.result = self.segment_results[0].copy()
+        
+        # Add each subsequent segment (skipping the first point to avoid duplication)
+        for i in range(1, len(self.segment_results)):
+            segment = self.segment_results[i]
+            
+            # Skip the first point of each subsequent segment as it should 
+            # be the same as the last point of the previous segment
+            # But verify they're actually the same
+            if len(segment) > 1 and np.array_equal(self.result[-1], segment[0]):
+                self.result.extend(segment[1:])
+            else:
+                # Something's wrong - just append everything
+                self.result.extend(segment)
+    
+    def get_segment_info(self):
+        """Get information about each segment of the path
+        
+        Returns
+        -------
+        dict
+            Dictionary with information about each path segment
+        """
+        if not self.segment_results:
+            return None
+            
+        info = {
+            'num_segments': len(self.segment_results),
+            'segment_lengths': [len(path) for path in self.segment_results],
+            'segment_evaluated_nodes': self.segment_evaluated_nodes,
+            'total_path_length': len(self.result),
+            'total_evaluated_nodes': self.evaluated_nodes,
+            'used_2d_optimization': self.use_2d_mode and len(self.start_point) == 3
+        }
+        return info