Spaces:
Paused
Paused
| """ | |
| GENESIS Urban Layout Simulation Viewer | |
| Loads OBJ models and runs Traffic, Rain, and Wind simulations using PyVista | |
| """ | |
| import numpy as np | |
| import pyvista as pv | |
| import json | |
| import argparse | |
| from pathlib import Path | |
| from dataclasses import dataclass, field | |
| from typing import List, Tuple, Dict, Optional, Any | |
| from enum import Enum | |
| import time | |
| # ============================================================================ | |
| # SIMULATION MODE ENUM | |
| # ============================================================================ | |
| class SimulationMode(Enum): | |
| NONE = 0 | |
| TRAFFIC = 1 | |
| RAIN = 2 | |
| WIND = 3 | |
| # ============================================================================ | |
| # TRAFFIC SIMULATION | |
| # ============================================================================ | |
| class TrafficAgent: | |
| """Single traffic agent (vehicle/pedestrian)""" | |
| position: np.ndarray | |
| velocity: np.ndarray | |
| road_index: int | |
| progress: float # 0.0 to 1.0 along road | |
| speed: float | |
| agent_type: str # 'car', 'bike', 'pedestrian' | |
| class TrafficSimulation: | |
| """Traffic simulation with agents moving along roads""" | |
| def __init__(self, roads_data: List[Dict], num_agents: int = 50): | |
| self.roads = roads_data | |
| self.num_agents = num_agents | |
| self.agents: List[TrafficAgent] = [] | |
| self.agent_colors = { | |
| 'car': [1.0, 0.2, 0.2], # Red | |
| 'bike': [0.2, 1.0, 0.2], # Green | |
| 'pedestrian': [0.2, 0.2, 1.0] # Blue | |
| } | |
| self.paused = False | |
| self._initialize_agents() | |
| def _initialize_agents(self): | |
| """Create initial agents on roads""" | |
| self.agents = [] | |
| if not self.roads: | |
| return | |
| agent_types = ['car', 'bike', 'pedestrian'] | |
| type_speeds = {'car': 15.0, 'bike': 8.0, 'pedestrian': 2.0} | |
| type_weights = [0.5, 0.3, 0.2] | |
| for i in range(self.num_agents): | |
| # Random road selection weighted by length | |
| road_lengths = [r.get('length', 10) for r in self.roads] | |
| total_length = sum(road_lengths) | |
| if total_length == 0: | |
| continue | |
| probs = [l / total_length for l in road_lengths] | |
| road_idx = np.random.choice(len(self.roads), p=probs) | |
| road = self.roads[road_idx] | |
| # Random position along road | |
| progress = np.random.random() | |
| start = np.array(road['start_point'] + [0.5]) # Slight elevation | |
| end = np.array(road['end_point'] + [0.5]) | |
| position = start + progress * (end - start) | |
| # Direction along road (randomly forward or backward) | |
| direction = 1 if np.random.random() > 0.5 else -1 | |
| road_vec = end - start | |
| road_length = np.linalg.norm(road_vec) | |
| if road_length > 0: | |
| velocity = direction * road_vec / road_length | |
| else: | |
| velocity = np.array([1.0, 0.0, 0.0]) | |
| # Agent type | |
| agent_type = np.random.choice(agent_types, p=type_weights) | |
| speed = type_speeds[agent_type] * (0.8 + 0.4 * np.random.random()) | |
| agent = TrafficAgent( | |
| position=position, | |
| velocity=velocity * speed, | |
| road_index=road_idx, | |
| progress=progress, | |
| speed=speed, | |
| agent_type=agent_type | |
| ) | |
| self.agents.append(agent) | |
| def update(self, dt: float = 0.016): | |
| """Update agent positions""" | |
| if self.paused or not self.roads: | |
| return | |
| for agent in self.agents: | |
| road = self.roads[agent.road_index] | |
| start = np.array(road['start_point'] + [0.5]) | |
| end = np.array(road['end_point'] + [0.5]) | |
| road_vec = end - start | |
| road_length = np.linalg.norm(road_vec) | |
| if road_length == 0: | |
| continue | |
| # Update progress along road | |
| direction = np.sign(np.dot(agent.velocity, road_vec)) | |
| agent.progress += direction * agent.speed * dt / road_length | |
| # Handle road endpoints - switch to connected road or reverse | |
| if agent.progress > 1.0 or agent.progress < 0.0: | |
| # Find connected road or reverse | |
| current_end = end if agent.progress > 1.0 else start | |
| connected = self._find_connected_road(agent.road_index, current_end) | |
| if connected is not None: | |
| agent.road_index = connected | |
| new_road = self.roads[connected] | |
| new_start = np.array(new_road['start_point'] + [0.5]) | |
| new_end = np.array(new_road['end_point'] + [0.5]) | |
| # Determine entry point | |
| dist_to_start = np.linalg.norm(current_end[:2] - new_start[:2]) | |
| dist_to_end = np.linalg.norm(current_end[:2] - new_end[:2]) | |
| if dist_to_start < dist_to_end: | |
| agent.progress = 0.0 | |
| new_vec = new_end - new_start | |
| else: | |
| agent.progress = 1.0 | |
| new_vec = new_start - new_end | |
| new_length = np.linalg.norm(new_vec) | |
| if new_length > 0: | |
| agent.velocity = new_vec / new_length * agent.speed | |
| else: | |
| # Reverse direction | |
| agent.velocity = -agent.velocity | |
| agent.progress = np.clip(agent.progress, 0.0, 1.0) | |
| # Update actual position | |
| agent.position = start + agent.progress * road_vec | |
| agent.position[2] = 0.5 # Keep at road level | |
| def _find_connected_road(self, current_road_idx: int, | |
| endpoint: np.ndarray, | |
| tolerance: float = 5.0) -> Optional[int]: | |
| """Find a road connected to the given endpoint""" | |
| for i, road in enumerate(self.roads): | |
| if i == current_road_idx: | |
| continue | |
| start = np.array(road['start_point']) | |
| end = np.array(road['end_point']) | |
| if np.linalg.norm(endpoint[:2] - start) < tolerance: | |
| return i | |
| if np.linalg.norm(endpoint[:2] - end) < tolerance: | |
| return i | |
| return None | |
| def get_points_and_colors(self) -> Tuple[np.ndarray, np.ndarray]: | |
| """Get agent positions and colors for rendering""" | |
| if not self.agents: | |
| return np.array([[0, 0, 0]]), np.array([[1, 1, 1]]) | |
| positions = np.array([a.position for a in self.agents]) | |
| colors = np.array([self.agent_colors[a.agent_type] for a in self.agents]) | |
| return positions, colors | |
| def reset(self): | |
| """Reset simulation""" | |
| self._initialize_agents() | |
| # ============================================================================ | |
| # RAIN SIMULATION | |
| # ============================================================================ | |
| class RainSimulation: | |
| """Rain particle simulation""" | |
| def __init__(self, bounds: Tuple[float, float, float, float], | |
| max_height: float = 100.0, | |
| num_drops: int = 2000, | |
| rain_speed: float = 50.0): | |
| self.bounds = bounds # minx, miny, maxx, maxy | |
| self.max_height = max_height | |
| self.num_drops = num_drops | |
| self.rain_speed = rain_speed | |
| self.drops: np.ndarray = None | |
| self.velocities: np.ndarray = None | |
| self.paused = False | |
| self._initialize_drops() | |
| def _initialize_drops(self): | |
| """Initialize raindrop positions""" | |
| minx, miny, maxx, maxy = self.bounds | |
| # Random positions within bounds | |
| x = np.random.uniform(minx, maxx, self.num_drops) | |
| y = np.random.uniform(miny, maxy, self.num_drops) | |
| z = np.random.uniform(0, self.max_height, self.num_drops) | |
| self.drops = np.column_stack([x, y, z]) | |
| # Velocities - mostly downward with slight variation | |
| vx = np.random.uniform(-2, 2, self.num_drops) | |
| vy = np.random.uniform(-2, 2, self.num_drops) | |
| vz = np.full(self.num_drops, -self.rain_speed) | |
| self.velocities = np.column_stack([vx, vy, vz]) | |
| def update(self, dt: float = 0.016): | |
| """Update raindrop positions""" | |
| if self.paused: | |
| return | |
| # Move drops | |
| self.drops += self.velocities * dt | |
| # Reset drops that hit ground | |
| ground_mask = self.drops[:, 2] < 0 | |
| if np.any(ground_mask): | |
| minx, miny, maxx, maxy = self.bounds | |
| num_reset = np.sum(ground_mask) | |
| self.drops[ground_mask, 0] = np.random.uniform(minx, maxx, num_reset) | |
| self.drops[ground_mask, 1] = np.random.uniform(miny, maxy, num_reset) | |
| self.drops[ground_mask, 2] = self.max_height | |
| def get_points_and_colors(self) -> Tuple[np.ndarray, np.ndarray]: | |
| """Get drop positions and colors""" | |
| # Color based on height - lighter at top, darker near ground | |
| heights = self.drops[:, 2] / self.max_height | |
| colors = np.column_stack([ | |
| 0.3 + 0.4 * heights, # R | |
| 0.5 + 0.3 * heights, # G | |
| 0.8 + 0.2 * heights # B | |
| ]) | |
| return self.drops.copy(), colors | |
| def get_rain_lines(self) -> pv.PolyData: | |
| """Get rain as short lines for better visualization""" | |
| line_length = 2.0 | |
| # Create line endpoints | |
| starts = self.drops.copy() | |
| ends = self.drops + np.array([0, 0, line_length]) | |
| # Build lines | |
| points = [] | |
| lines = [] | |
| for i in range(len(self.drops)): | |
| idx = len(points) | |
| points.append(starts[i]) | |
| points.append(ends[i]) | |
| lines.append([2, idx, idx + 1]) | |
| if not points: | |
| return pv.PolyData() | |
| mesh = pv.PolyData(np.array(points)) | |
| mesh.lines = np.array(lines) | |
| return mesh | |
| def reset(self): | |
| """Reset simulation""" | |
| self._initialize_drops() | |
| # ============================================================================ | |
| # WIND SIMULATION | |
| # ============================================================================ | |
| class WindSimulation: | |
| """Wind particle flow simulation""" | |
| def __init__(self, bounds: Tuple[float, float, float, float], | |
| max_height: float = 50.0, | |
| num_particles: int = 1500, | |
| wind_speed: float = 20.0, | |
| wind_direction: float = 45.0): # degrees | |
| self.bounds = bounds | |
| self.max_height = max_height | |
| self.num_particles = num_particles | |
| self.base_wind_speed = wind_speed | |
| self.wind_direction = np.radians(wind_direction) | |
| self.particles: np.ndarray = None | |
| self.velocities: np.ndarray = None | |
| self.lifetimes: np.ndarray = None | |
| self.max_lifetime = 5.0 | |
| self.paused = False | |
| self.buildings: List[Dict] = [] | |
| self._initialize_particles() | |
| def set_buildings(self, plots_data: List[Dict]): | |
| """Set building data for wind interaction""" | |
| self.buildings = plots_data | |
| def set_wind_direction(self, degrees: float): | |
| """Set wind direction in degrees""" | |
| self.wind_direction = np.radians(degrees) | |
| self._update_base_velocity() | |
| def _update_base_velocity(self): | |
| """Update base velocity based on wind direction""" | |
| vx = self.base_wind_speed * np.cos(self.wind_direction) | |
| vy = self.base_wind_speed * np.sin(self.wind_direction) | |
| # Add variation to existing particles | |
| if self.velocities is not None: | |
| noise_x = np.random.uniform(-3, 3, self.num_particles) | |
| noise_y = np.random.uniform(-3, 3, self.num_particles) | |
| noise_z = np.random.uniform(-1, 1, self.num_particles) | |
| self.velocities[:, 0] = vx + noise_x | |
| self.velocities[:, 1] = vy + noise_y | |
| self.velocities[:, 2] = noise_z | |
| def _initialize_particles(self): | |
| """Initialize wind particles""" | |
| minx, miny, maxx, maxy = self.bounds | |
| # Spawn particles from upwind edge | |
| wind_dx = np.cos(self.wind_direction) | |
| wind_dy = np.sin(self.wind_direction) | |
| # Determine spawn edge based on wind direction | |
| if abs(wind_dx) > abs(wind_dy): | |
| if wind_dx > 0: # Wind from left | |
| x = np.full(self.num_particles, minx - 10) | |
| else: # Wind from right | |
| x = np.full(self.num_particles, maxx + 10) | |
| y = np.random.uniform(miny, maxy, self.num_particles) | |
| else: | |
| x = np.random.uniform(minx, maxx, self.num_particles) | |
| if wind_dy > 0: # Wind from bottom | |
| y = np.full(self.num_particles, miny - 10) | |
| else: # Wind from top | |
| y = np.full(self.num_particles, maxy + 10) | |
| z = np.random.uniform(1, self.max_height, self.num_particles) | |
| self.particles = np.column_stack([x, y, z]) | |
| # Initialize velocities | |
| vx = self.base_wind_speed * wind_dx + np.random.uniform(-3, 3, self.num_particles) | |
| vy = self.base_wind_speed * wind_dy + np.random.uniform(-3, 3, self.num_particles) | |
| vz = np.random.uniform(-1, 1, self.num_particles) | |
| self.velocities = np.column_stack([vx, vy, vz]) | |
| # Lifetimes for particle recycling | |
| self.lifetimes = np.random.uniform(0, self.max_lifetime, self.num_particles) | |
| def update(self, dt: float = 0.016): | |
| """Update particle positions with building interaction""" | |
| if self.paused: | |
| return | |
| # Update lifetimes | |
| self.lifetimes -= dt | |
| # Move particles | |
| self.particles += self.velocities * dt | |
| # Add turbulence | |
| turbulence = np.random.uniform(-0.5, 0.5, self.particles.shape) * dt * 10 | |
| self.velocities += turbulence | |
| # Clamp vertical velocity | |
| self.velocities[:, 2] = np.clip(self.velocities[:, 2], -5, 5) | |
| # Simple building interaction - deflect around buildings | |
| for building in self.buildings: | |
| bounds = building.get('bounds', [0, 0, 0, 0]) | |
| height = building.get('max_building_height', 10) | |
| bminx, bminy, bmaxx, bmaxy = bounds | |
| # Check particles near building | |
| in_x = (self.particles[:, 0] > bminx) & (self.particles[:, 0] < bmaxx) | |
| in_y = (self.particles[:, 1] > bminy) & (self.particles[:, 1] < bmaxy) | |
| below_height = self.particles[:, 2] < height | |
| affected = in_x & in_y & below_height | |
| if np.any(affected): | |
| # Push particles up and around | |
| self.velocities[affected, 2] += 10 * dt | |
| # Deflect horizontally | |
| center_x = (bminx + bmaxx) / 2 | |
| center_y = (bminy + bmaxy) / 2 | |
| dx = self.particles[affected, 0] - center_x | |
| dy = self.particles[affected, 1] - center_y | |
| self.velocities[affected, 0] += np.sign(dx) * 5 * dt | |
| self.velocities[affected, 1] += np.sign(dy) * 5 * dt | |
| # Reset expired or out-of-bounds particles | |
| minx, miny, maxx, maxy = self.bounds | |
| margin = 50 | |
| out_of_bounds = ( | |
| (self.particles[:, 0] < minx - margin) | | |
| (self.particles[:, 0] > maxx + margin) | | |
| (self.particles[:, 1] < miny - margin) | | |
| (self.particles[:, 1] > maxy + margin) | | |
| (self.particles[:, 2] < 0) | | |
| (self.particles[:, 2] > self.max_height * 2) | | |
| (self.lifetimes < 0) | |
| ) | |
| if np.any(out_of_bounds): | |
| self._respawn_particles(out_of_bounds) | |
| def _respawn_particles(self, mask: np.ndarray): | |
| """Respawn particles at upwind edge""" | |
| num_respawn = np.sum(mask) | |
| if num_respawn == 0: | |
| return | |
| minx, miny, maxx, maxy = self.bounds | |
| wind_dx = np.cos(self.wind_direction) | |
| wind_dy = np.sin(self.wind_direction) | |
| # Spawn at upwind edge | |
| if abs(wind_dx) > abs(wind_dy): | |
| if wind_dx > 0: | |
| self.particles[mask, 0] = minx - 10 | |
| else: | |
| self.particles[mask, 0] = maxx + 10 | |
| self.particles[mask, 1] = np.random.uniform(miny, maxy, num_respawn) | |
| else: | |
| self.particles[mask, 0] = np.random.uniform(minx, maxx, num_respawn) | |
| if wind_dy > 0: | |
| self.particles[mask, 1] = miny - 10 | |
| else: | |
| self.particles[mask, 1] = maxy + 10 | |
| self.particles[mask, 2] = np.random.uniform(1, self.max_height, num_respawn) | |
| # Reset velocities | |
| self.velocities[mask, 0] = self.base_wind_speed * wind_dx + np.random.uniform(-3, 3, num_respawn) | |
| self.velocities[mask, 1] = self.base_wind_speed * wind_dy + np.random.uniform(-3, 3, num_respawn) | |
| self.velocities[mask, 2] = np.random.uniform(-1, 1, num_respawn) | |
| # Reset lifetimes | |
| self.lifetimes[mask] = self.max_lifetime | |
| def get_points_and_colors(self) -> Tuple[np.ndarray, np.ndarray]: | |
| """Get particle positions and colors""" | |
| # Color based on velocity magnitude | |
| speeds = np.linalg.norm(self.velocities, axis=1) | |
| max_speed = self.base_wind_speed * 1.5 | |
| normalized_speed = np.clip(speeds / max_speed, 0, 1) | |
| # Color gradient: blue (slow) -> cyan -> white (fast) | |
| colors = np.column_stack([ | |
| 0.5 + 0.5 * normalized_speed, # R | |
| 0.8 + 0.2 * normalized_speed, # G | |
| 1.0 * np.ones_like(normalized_speed) # B | |
| ]) | |
| return self.particles.copy(), colors | |
| def get_streamlines(self) -> pv.PolyData: | |
| """Get wind as short streamlines""" | |
| line_length = 3.0 | |
| # Normalize velocities for direction | |
| speeds = np.linalg.norm(self.velocities, axis=1, keepdims=True) | |
| speeds = np.maximum(speeds, 0.001) | |
| directions = self.velocities / speeds | |
| # Create line endpoints | |
| starts = self.particles.copy() | |
| ends = self.particles + directions * line_length | |
| points = [] | |
| lines = [] | |
| for i in range(len(self.particles)): | |
| idx = len(points) | |
| points.append(starts[i]) | |
| points.append(ends[i]) | |
| lines.append([2, idx, idx + 1]) | |
| if not points: | |
| return pv.PolyData() | |
| mesh = pv.PolyData(np.array(points)) | |
| mesh.lines = np.array(lines) | |
| return mesh | |
| def reset(self): | |
| """Reset simulation""" | |
| self._initialize_particles() | |
| # ============================================================================ | |
| # MAIN SIMULATION VIEWER | |
| # ============================================================================ | |
| class UrbanSimulationViewer: | |
| """Main simulation viewer combining all simulations""" | |
| def __init__(self, obj_path: str, json_path: str = None): | |
| self.obj_path = Path(obj_path) | |
| self.json_path = Path(json_path) if json_path else None | |
| # Load data | |
| self.mesh = None | |
| self.layout_data = None | |
| self.bounds = None | |
| # Simulations | |
| self.traffic_sim = None | |
| self.rain_sim = None | |
| self.wind_sim = None | |
| # Visualization | |
| self.plotter = None | |
| self.current_mode = SimulationMode.NONE | |
| self.simulation_actors = {} | |
| self.is_paused = False | |
| self.last_update_time = time.time() | |
| # Load resources | |
| self._load_resources() | |
| self._initialize_simulations() | |
| def _load_resources(self): | |
| """Load OBJ mesh and JSON data""" | |
| print("=" * 60) | |
| print("LOADING RESOURCES") | |
| print("=" * 60) | |
| # Load OBJ | |
| if self.obj_path.exists(): | |
| print(f"Loading OBJ: {self.obj_path}") | |
| self.mesh = pv.read(str(self.obj_path)) | |
| self.bounds = self.mesh.bounds | |
| print(f" Vertices: {self.mesh.n_points}") | |
| print(f" Faces: {self.mesh.n_cells}") | |
| print(f" Bounds: {self.bounds}") | |
| else: | |
| raise FileNotFoundError(f"OBJ file not found: {self.obj_path}") | |
| # Load JSON if provided | |
| if self.json_path and self.json_path.exists(): | |
| print(f"Loading JSON: {self.json_path}") | |
| with open(self.json_path, 'r') as f: | |
| self.layout_data = json.load(f) | |
| print(f" Roads: {len(self.layout_data.get('roads', []))}") | |
| print(f" Plots: {len(self.layout_data.get('plots', []))}") | |
| else: | |
| print("No JSON data - using mesh bounds for simulation") | |
| # Create minimal data from mesh bounds | |
| minx, maxx, miny, maxy, minz, maxz = self.bounds | |
| self.layout_data = { | |
| 'roads': [], | |
| 'plots': [], | |
| 'site_info': { | |
| 'bounds': [minx, miny, maxx, maxy] | |
| } | |
| } | |
| def _initialize_simulations(self): | |
| """Initialize all simulation systems""" | |
| print("\n" + "=" * 60) | |
| print("INITIALIZING SIMULATIONS") | |
| print("=" * 60) | |
| # Get bounds | |
| if self.layout_data and 'site_info' in self.layout_data: | |
| bounds = self.layout_data['site_info']['bounds'] | |
| sim_bounds = (bounds[0], bounds[1], bounds[2], bounds[3]) | |
| else: | |
| minx, maxx, miny, maxy, _, _ = self.bounds | |
| sim_bounds = (minx, miny, maxx, maxy) | |
| # Calculate max building height | |
| max_height = 50.0 | |
| if self.layout_data and 'plots' in self.layout_data: | |
| heights = [p.get('max_building_height', 10) for p in self.layout_data['plots']] | |
| if heights: | |
| max_height = max(heights) + 20 | |
| # Traffic simulation | |
| roads = self.layout_data.get('roads', []) if self.layout_data else [] | |
| if roads: | |
| self.traffic_sim = TrafficSimulation(roads, num_agents=100) | |
| print(f"Traffic: {len(self.traffic_sim.agents)} agents on {len(roads)} roads") | |
| else: | |
| # Generate fake roads from bounds for demo | |
| fake_roads = self._generate_fake_roads(sim_bounds) | |
| self.traffic_sim = TrafficSimulation(fake_roads, num_agents=50) | |
| print(f"Traffic: {len(self.traffic_sim.agents)} agents (generated roads)") | |
| # Rain simulation | |
| self.rain_sim = RainSimulation( | |
| bounds=sim_bounds, | |
| max_height=max_height + 50, | |
| num_drops=3000, | |
| rain_speed=60.0 | |
| ) | |
| print(f"Rain: {self.rain_sim.num_drops} drops") | |
| # Wind simulation | |
| self.wind_sim = WindSimulation( | |
| bounds=sim_bounds, | |
| max_height=max_height, | |
| num_particles=2000, | |
| wind_speed=25.0, | |
| wind_direction=45.0 | |
| ) | |
| # Set buildings for wind interaction | |
| if self.layout_data and 'plots' in self.layout_data: | |
| self.wind_sim.set_buildings(self.layout_data['plots']) | |
| print(f"Wind: {self.wind_sim.num_particles} particles") | |
| def _generate_fake_roads(self, bounds: Tuple) -> List[Dict]: | |
| """Generate fake roads if JSON not provided""" | |
| minx, miny, maxx, maxy = bounds | |
| roads = [] | |
| # Grid of roads | |
| num_h = 5 | |
| num_v = 5 | |
| for i in range(num_h): | |
| y = miny + (i + 1) * (maxy - miny) / (num_h + 1) | |
| roads.append({ | |
| 'start_point': [minx, y], | |
| 'end_point': [maxx, y], | |
| 'width': 12.0, | |
| 'length': maxx - minx | |
| }) | |
| for i in range(num_v): | |
| x = minx + (i + 1) * (maxx - minx) / (num_v + 1) | |
| roads.append({ | |
| 'start_point': [x, miny], | |
| 'end_point': [x, maxy], | |
| 'width': 12.0, | |
| 'length': maxy - miny | |
| }) | |
| return roads | |
| def _setup_plotter(self): | |
| """Setup PyVista plotter with controls""" | |
| self.plotter = pv.Plotter(title="GENESIS Urban Simulation Viewer") | |
| self.plotter.set_background('black') | |
| # Add mesh | |
| self.plotter.add_mesh( | |
| self.mesh, | |
| color='lightgray', | |
| opacity=0.9, | |
| show_edges=False, | |
| name='layout_mesh' | |
| ) | |
| # Set top-down camera | |
| self._set_top_down_view() | |
| # Add text overlay | |
| self._add_ui_text() | |
| # Setup key bindings | |
| self.plotter.add_key_event('space', self._toggle_pause) | |
| self.plotter.add_key_event('r', self._reset_simulations) | |
| self.plotter.add_key_event('0', lambda: self._set_mode(SimulationMode.NONE)) | |
| self.plotter.add_key_event('1', lambda: self._set_mode(SimulationMode.TRAFFIC)) | |
| self.plotter.add_key_event('2', lambda: self._set_mode(SimulationMode.RAIN)) | |
| self.plotter.add_key_event('3', lambda: self._set_mode(SimulationMode.WIND)) | |
| self.plotter.add_key_event('Up', lambda: self._adjust_wind_direction(15)) | |
| self.plotter.add_key_event('Down', lambda: self._adjust_wind_direction(-15)) | |
| self.plotter.add_key_event('t', self._toggle_top_view) | |
| self.plotter.add_key_event('p', self._toggle_perspective) | |
| def _set_top_down_view(self): | |
| """Set camera to top-down view""" | |
| if self.bounds: | |
| minx, maxx, miny, maxy, minz, maxz = self.bounds | |
| cx = (minx + maxx) / 2 | |
| cy = (miny + maxy) / 2 | |
| # Calculate camera height based on scene size | |
| width = maxx - minx | |
| height = maxy - miny | |
| cam_height = max(width, height) * 1.5 | |
| self.plotter.camera_position = [ | |
| (cx, cy, cam_height), # Camera position | |
| (cx, cy, 0), # Focal point | |
| (0, 1, 0) # Up vector | |
| ] | |
| def _toggle_top_view(self): | |
| """Toggle between top-down and isometric view""" | |
| if self.bounds: | |
| minx, maxx, miny, maxy, minz, maxz = self.bounds | |
| cx = (minx + maxx) / 2 | |
| cy = (miny + maxy) / 2 | |
| width = maxx - minx | |
| height = maxy - miny | |
| cam_dist = max(width, height) * 1.2 | |
| self.plotter.camera_position = [ | |
| (cx, cy, cam_dist), | |
| (cx, cy, 0), | |
| (0, 1, 0) | |
| ] | |
| def _toggle_perspective(self): | |
| """Toggle to perspective/isometric view""" | |
| if self.bounds: | |
| minx, maxx, miny, maxy, minz, maxz = self.bounds | |
| cx = (minx + maxx) / 2 | |
| cy = (miny + maxy) / 2 | |
| width = maxx - minx | |
| height = maxy - miny | |
| cam_dist = max(width, height) * 0.8 | |
| self.plotter.camera_position = [ | |
| (cx - cam_dist, cy - cam_dist, cam_dist * 0.7), | |
| (cx, cy, 0), | |
| (0, 0, 1) | |
| ] | |
| def _add_ui_text(self): | |
| """Add UI text overlay""" | |
| instructions = ( | |
| "Controls:\n" | |
| "SPACE - Pause/Resume\n" | |
| "R - Reset\n" | |
| "0 - No overlay\n" | |
| "1 - Traffic\n" | |
| "2 - Rain\n" | |
| "3 - Wind\n" | |
| "T - Top view\n" | |
| "P - Perspective\n" | |
| "↑/↓ - Wind direction" | |
| ) | |
| self.plotter.add_text( | |
| instructions, | |
| position='upper_left', | |
| font_size=10, | |
| color='white', | |
| name='instructions' | |
| ) | |
| self.status_text = self.plotter.add_text( | |
| "Mode: None | Status: Running", | |
| position='upper_right', | |
| font_size=12, | |
| color='cyan', | |
| name='status' | |
| ) | |
| def _update_status_text(self): | |
| """Update status text""" | |
| mode_names = { | |
| SimulationMode.NONE: "None", | |
| SimulationMode.TRAFFIC: "Traffic", | |
| SimulationMode.RAIN: "Rain", | |
| SimulationMode.WIND: "Wind" | |
| } | |
| status = "Paused" if self.is_paused else "Running" | |
| mode = mode_names.get(self.current_mode, "Unknown") | |
| extra = "" | |
| if self.current_mode == SimulationMode.WIND: | |
| direction = np.degrees(self.wind_sim.wind_direction) | |
| extra = f" | Direction: {direction:.0f}°" | |
| text = f"Mode: {mode} | Status: {status}{extra}" | |
| self.plotter.add_text( | |
| text, | |
| position='upper_right', | |
| font_size=12, | |
| color='cyan', | |
| name='status' | |
| ) | |
| def _toggle_pause(self): | |
| """Toggle pause state""" | |
| self.is_paused = not self.is_paused | |
| if self.traffic_sim: | |
| self.traffic_sim.paused = self.is_paused | |
| if self.rain_sim: | |
| self.rain_sim.paused = self.is_paused | |
| if self.wind_sim: | |
| self.wind_sim.paused = self.is_paused | |
| self._update_status_text() | |
| print(f"{'Paused' if self.is_paused else 'Resumed'}") | |
| def _reset_simulations(self): | |
| """Reset all simulations""" | |
| if self.traffic_sim: | |
| self.traffic_sim.reset() | |
| if self.rain_sim: | |
| self.rain_sim.reset() | |
| if self.wind_sim: | |
| self.wind_sim.reset() | |
| print("Simulations reset") | |
| def _set_mode(self, mode: SimulationMode): | |
| """Set simulation display mode""" | |
| self.current_mode = mode | |
| self._clear_simulation_actors() | |
| self._update_status_text() | |
| mode_names = { | |
| SimulationMode.NONE: "None", | |
| SimulationMode.TRAFFIC: "Traffic", | |
| SimulationMode.RAIN: "Rain", | |
| SimulationMode.WIND: "Wind" | |
| } | |
| print(f"Mode: {mode_names.get(mode, 'Unknown')}") | |
| def _adjust_wind_direction(self, delta: float): | |
| """Adjust wind direction""" | |
| if self.wind_sim: | |
| new_dir = np.degrees(self.wind_sim.wind_direction) + delta | |
| self.wind_sim.set_wind_direction(new_dir) | |
| self._update_status_text() | |
| print(f"Wind direction: {new_dir:.0f}°") | |
| def _clear_simulation_actors(self): | |
| """Clear simulation visualization actors""" | |
| for name in list(self.simulation_actors.keys()): | |
| try: | |
| self.plotter.remove_actor(name) | |
| except: | |
| pass | |
| self.simulation_actors = {} | |
| def _update_callback(self): | |
| """Main update callback for animation""" | |
| current_time = time.time() | |
| dt = min(current_time - self.last_update_time, 0.1) # Cap dt | |
| self.last_update_time = current_time | |
| # Update simulations | |
| if self.current_mode == SimulationMode.TRAFFIC and self.traffic_sim: | |
| self.traffic_sim.update(dt) | |
| self._render_traffic() | |
| elif self.current_mode == SimulationMode.RAIN and self.rain_sim: | |
| self.rain_sim.update(dt) | |
| self._render_rain() | |
| elif self.current_mode == SimulationMode.WIND and self.wind_sim: | |
| self.wind_sim.update(dt) | |
| self._render_wind() | |
| def _render_traffic(self): | |
| """Render traffic agents""" | |
| if not self.traffic_sim: | |
| return | |
| points, colors = self.traffic_sim.get_points_and_colors() | |
| if len(points) > 0: | |
| point_cloud = pv.PolyData(points) | |
| point_cloud['colors'] = (colors * 255).astype(np.uint8) | |
| # Remove old actor | |
| if 'traffic_points' in self.simulation_actors: | |
| try: | |
| self.plotter.remove_actor('traffic_points') | |
| except: | |
| pass | |
| actor = self.plotter.add_mesh( | |
| point_cloud, | |
| scalars='colors', | |
| rgb=True, | |
| point_size=15, | |
| render_points_as_spheres=True, | |
| name='traffic_points' | |
| ) | |
| self.simulation_actors['traffic_points'] = actor | |
| def _render_rain(self): | |
| """Render rain particles""" | |
| if not self.rain_sim: | |
| return | |
| # Use lines for rain | |
| rain_mesh = self.rain_sim.get_rain_lines() | |
| if rain_mesh.n_points > 0: | |
| if 'rain_lines' in self.simulation_actors: | |
| try: | |
| self.plotter.remove_actor('rain_lines') | |
| except: | |
| pass | |
| actor = self.plotter.add_mesh( | |
| rain_mesh, | |
| color='lightblue', | |
| line_width=1, | |
| opacity=0.6, | |
| name='rain_lines' | |
| ) | |
| self.simulation_actors['rain_lines'] = actor | |
| def _render_wind(self): | |
| """Render wind particles""" | |
| if not self.wind_sim: | |
| return | |
| # Use streamlines for wind | |
| wind_mesh = self.wind_sim.get_streamlines() | |
| if wind_mesh.n_points > 0: | |
| if 'wind_lines' in self.simulation_actors: | |
| try: | |
| self.plotter.remove_actor('wind_lines') | |
| except: | |
| pass | |
| actor = self.plotter.add_mesh( | |
| wind_mesh, | |
| color='white', | |
| line_width=2, | |
| opacity=0.7, | |
| name='wind_lines' | |
| ) | |
| self.simulation_actors['wind_lines'] = actor | |
| def run(self): | |
| """Run the simulation viewer""" | |
| print("\n" + "=" * 60) | |
| print("STARTING SIMULATION VIEWER") | |
| print("=" * 60) | |
| print("\nControls:") | |
| print(" SPACE - Pause/Resume") | |
| print(" R - Reset simulations") | |
| print(" 0 - No overlay") | |
| print(" 1 - Traffic simulation") | |
| print(" 2 - Rain simulation") | |
| print(" 3 - Wind simulation") | |
| print(" T - Top-down view") | |
| print(" P - Perspective view") | |
| print(" ↑/↓ - Change wind direction") | |
| print("\nClose window to exit.") | |
| self._setup_plotter() | |
| # Start with traffic mode | |
| self._set_mode(SimulationMode.TRAFFIC) | |
| # Timer-driven animation (works on older PyVista) | |
| self.plotter.add_timer_event(max_steps=10**9, duration=33, callback=self._update_callback) | |
| # Show plotter | |
| self.plotter.show() | |
| # ============================================================================ | |
| # CLI INTERFACE | |
| # ============================================================================ | |
| def main(): | |
| parser = argparse.ArgumentParser( | |
| description='GENESIS Urban Layout Simulation Viewer', | |
| formatter_class=argparse.RawDescriptionHelpFormatter, | |
| epilog=""" | |
| Examples: | |
| python simulation_viewer.py layout.obj | |
| python simulation_viewer.py layout.obj --json layout.json | |
| python simulation_viewer.py layout.obj -j layout.json --wind-dir 90 | |
| """ | |
| ) | |
| parser.add_argument('obj_file', type=str, help='Path to OBJ file') | |
| parser.add_argument('--json', '-j', type=str, help='Path to JSON layout file') | |
| parser.add_argument('--wind-dir', type=float, default=45.0, | |
| help='Initial wind direction in degrees (default: 45)') | |
| parser.add_argument('--num-traffic', type=int, default=100, | |
| help='Number of traffic agents (default: 100)') | |
| parser.add_argument('--num-rain', type=int, default=3000, | |
| help='Number of rain drops (default: 3000)') | |
| parser.add_argument('--num-wind', type=int, default=2000, | |
| help='Number of wind particles (default: 2000)') | |
| args = parser.parse_args() | |
| # Validate files | |
| obj_path = Path(args.obj_file) | |
| if not obj_path.exists(): | |
| print(f"Error: OBJ file not found: {obj_path}") | |
| return 1 | |
| json_path = args.json | |
| if json_path and not Path(json_path).exists(): | |
| print(f"Warning: JSON file not found: {json_path}") | |
| print("Proceeding without layout data...") | |
| json_path = None | |
| # Create and run viewer | |
| try: | |
| viewer = UrbanSimulationViewer(str(obj_path), json_path) | |
| # Apply custom settings | |
| if viewer.wind_sim: | |
| viewer.wind_sim.set_wind_direction(args.wind_dir) | |
| viewer.run() | |
| except Exception as e: | |
| print(f"Error: {e}") | |
| import traceback | |
| traceback.print_exc() | |
| return 1 | |
| return 0 | |
| if __name__ == "__main__": | |
| exit(main()) |