Upload qads/planner/graph.py

qads/planner/graph.py

@@ -0,0 +1,194 @@
+"""World State Graph Builder."""
+import numpy as np
+from typing import Dict, Any, List, Tuple, Optional
+import networkx as nx
+
+
+class WorldGraph:
+    """
+    Probabilistic graph representing the environment.
+    Each node stores: position, risk, traversal cost, energy cost, uncertainty, obstacle probability
+    """
+    
+    def __init__(self, resolution: float = 0.5):
+        self.resolution = resolution
+        self.graph = nx.DiGraph()
+        self.node_positions: Dict[int, Tuple[float, ...]] = {}
+        self.node_metadata: Dict[int, Dict[str, Any]] = {}
+        self.next_id = 0
+    
+    def add_node(self, 
+                 position: Tuple[float, ...],
+                 risk: float = 0.0,
+                 cost: float = 1.0,
+                 energy: float = 1.0,
+                 uncertainty: float = 0.0,
+                 obstacle_prob: float = 0.0) -> int:
+        """Add a node to the graph."""
+        node_id = self.next_id
+        self.next_id += 1
+        
+        self.graph.add_node(node_id)
+        self.node_positions[node_id] = position
+        self.node_metadata[node_id] = {
+            'risk': float(risk),
+            'cost': float(cost),
+            'energy': float(energy),
+            'uncertainty': float(uncertainty),
+            'obstacle_prob': float(obstacle_prob),
+            'traversal_prob': float(1.0 - obstacle_prob),
+            'entropy': float(-uncertainty * np.log2(uncertainty + 1e-10) if uncertainty > 0 else 0.0)
+        }
+        return node_id
+    
+    def add_edge(self, 
+                 node_a: int, 
+                 node_b: int,
+                 weight: Optional[float] = None,
+                 risk: float = 0.0):
+        """Add an edge between nodes."""
+        if weight is None:
+            pos_a = self.node_positions[node_a]
+            pos_b = self.node_positions[node_b]
+            weight = np.linalg.norm(np.array(pos_a) - np.array(pos_b))
+        
+        meta_a = self.node_metadata[node_a]
+        meta_b = self.node_metadata[node_b]
+        
+        # Composite edge cost
+        composite_cost = (
+            0.3 * weight +
+            0.2 * (meta_a['risk'] + meta_b['risk']) / 2 +
+            0.2 * (meta_a['cost'] + meta_b['cost']) / 2 +
+            0.15 * (meta_a['uncertainty'] + meta_b['uncertainty']) / 2 +
+            0.15 * (meta_a['obstacle_prob'] + meta_b['obstacle_prob']) / 2
+        )
+        
+        self.graph.add_edge(node_a, node_b, 
+                           weight=composite_cost,
+                           distance=weight,
+                           risk=risk)
+    
+    def build_grid(self, 
+                   bounds: Tuple[Tuple[float, float], ...],
+                   obstacle_map: Optional[np.ndarray] = None,
+                   uncertainty_map: Optional[np.ndarray] = None):
+        """Build grid graph from bounds and maps."""
+        if len(bounds) == 2:
+            (x_min, x_max), (y_min, y_max) = bounds
+            nx_nodes = int((x_max - x_min) / self.resolution)
+            ny_nodes = int((y_max - y_min) / self.resolution)
+            
+            # Create nodes
+            for i in range(nx_nodes):
+                for j in range(ny_nodes):
+                    x = x_min + i * self.resolution
+                    y = y_min + j * self.resolution
+                    
+                    # Get map values
+                    obs_prob = 0.0
+                    unc = 0.0
+                    if obstacle_map is not None:
+                        mi = min(int(i * obstacle_map.shape[0] / nx_nodes), obstacle_map.shape[0]-1)
+                        mj = min(int(j * obstacle_map.shape[1] / ny_nodes), obstacle_map.shape[1]-1)
+                        obs_prob = obstacle_map[mi, mj]
+                    if uncertainty_map is not None:
+                        mi = min(int(i * uncertainty_map.shape[0] / nx_nodes), uncertainty_map.shape[0]-1)
+                        mj = min(int(j * uncertainty_map.shape[1] / ny_nodes), uncertainty_map.shape[1]-1)
+                        unc = uncertainty_map[mi, mj]
+                    
+                    self.add_node(
+                        position=(x, y),
+                        risk=obs_prob * 0.5 + unc * 0.3,
+                        cost=1.0 + obs_prob,
+                        energy=1.0,
+                        uncertainty=unc,
+                        obstacle_prob=obs_prob
+                    )
+            
+            # Add edges (4-connectivity)
+            for i in range(nx_nodes):
+                for j in range(ny_nodes):
+                    idx = i * ny_nodes + j
+                    # Right
+                    if i < nx_nodes - 1:
+                        self.add_edge(idx, idx + ny_nodes)
+                        self.add_edge(idx + ny_nodes, idx)
+                    # Up
+                    if j < ny_nodes - 1:
+                        self.add_edge(idx, idx + 1)
+                        self.add_edge(idx + 1, idx)
+    
+    def get_entropy(self) -> float:
+        """Compute graph-level entropy."""
+        probs = []
+        for node_id, meta in self.node_metadata.items():
+            p = meta.get('traversal_prob', 1.0)
+            if p > 0:
+                probs.append(p)
+        
+        if not probs:
+            return 0.0
+            
+        probs = np.array(probs)
+        probs = probs / probs.sum()
+        entropy = -np.sum(probs * np.log2(probs + 1e-10))
+        return float(entropy)
+    
+    def get_uncertainty(self) -> float:
+        """Compute average node uncertainty."""
+        uncertainties = [meta['uncertainty'] for meta in self.node_metadata.values()]
+        return float(np.mean(uncertainties)) if uncertainties else 0.0
+    
+    def get_obstacle_density(self) -> float:
+        """Compute obstacle density."""
+        obs = [meta['obstacle_prob'] for meta in self.node_metadata.values()]
+        return float(np.mean(obs)) if obs else 0.0
+    
+    def find_node_at(self, position: Tuple[float, ...], tolerance: float = None) -> Optional[int]:
+        """Find node closest to position."""
+        if tolerance is None:
+            tolerance = self.resolution * 2
+            
+        pos = np.array(position)
+        best_id = None
+        best_dist = float('inf')
+        
+        for node_id, node_pos in self.node_positions.items():
+            dist = np.linalg.norm(pos - np.array(node_pos))
+            if dist < tolerance and dist < best_dist:
+                best_id = node_id
+                best_dist = dist
+                
+        return best_id
+    
+    def to_cost_matrix(self) -> np.ndarray:
+        """Convert graph to cost matrix for QAOA."""
+        n = len(self.graph.nodes)
+        cost = np.zeros((n, n))
+        
+        for u, v, data in self.graph.edges(data=True):
+            cost[u, v] = data.get('weight', 1.0)
+        
+        # Add diagonal costs
+        for node_id, meta in self.node_metadata.items():
+            cost[node_id, node_id] = meta['cost']
+        
+        return cost
+    
+    def update_node(self, node_id: int, **kwargs):
+        """Update node metadata."""
+        if node_id in self.node_metadata:
+            self.node_metadata[node_id].update(kwargs)
+    
+    def get_state_dict(self) -> Dict[str, Any]:
+        """Export graph state."""
+        return {
+            'n_nodes': len(self.graph.nodes),
+            'n_edges': len(self.graph.edges),
+            'entropy': self.get_entropy(),
+            'uncertainty': self.get_uncertainty(),
+            'obstacle_density': self.get_obstacle_density(),
+            'positions': self.node_positions.copy(),
+            'metadata': self.node_metadata.copy()
+        }