"""
Phase space visualization for MechanicsDSL
Specialized tools for phase space and Poincaré section analysis.
"""
from typing import Tuple
try:
import matplotlib.pyplot as plt
HAS_MATPLOTLIB = True
except ImportError:
HAS_MATPLOTLIB = False
from ..utils import logger
[docs]
class PhaseSpaceVisualizer:
"""
Phase space and Poincaré section visualization.
Provides tools for analyzing dynamical systems through
phase portraits and stroboscopic maps.
"""
def __init__(self):
if not HAS_MATPLOTLIB:
raise ImportError("matplotlib required for visualization")
[docs]
def plot_phase_portrait(
self, solution: dict, coordinate_index: int = 0, title: str = "Phase Portrait"
) -> plt.Figure:
"""
Plot phase space trajectory (q vs q_dot).
Args:
solution: Simulation result
coordinate_index: Which coordinate to plot
title: Plot title
Returns:
matplotlib Figure
"""
y = solution["y"]
coords = solution.get("coordinates", [])
if coordinate_index >= len(coords):
raise ValueError(f"coordinate_index {coordinate_index} out of range")
q = y[2 * coordinate_index]
q_dot = y[2 * coordinate_index + 1]
coord_name = coords[coordinate_index]
fig, ax = plt.subplots(figsize=(8, 6))
ax.plot(q, q_dot, lw=0.5)
ax.scatter(q[0], q_dot[0], c="green", s=50, zorder=5, label="Start")
ax.scatter(q[-1], q_dot[-1], c="red", s=50, zorder=5, label="End")
ax.set_xlabel(f"{coord_name}")
ax.set_ylabel(f"{coord_name}_dot")
ax.set_title(title)
ax.legend()
ax.grid(True, alpha=0.3)
return fig
[docs]
def plot_phase_portrait_3d(
self,
solution: dict,
coords: Tuple[int, int, int] = (0, 0, 1),
title: str = "3D Phase Space",
) -> plt.Figure:
"""
Plot 3D phase space trajectory.
Args:
solution: Simulation result
coords: Tuple of (coord1_idx, coord1_type, coord2_idx)
where type 0=position, 1=velocity
title: Plot title
Returns:
matplotlib Figure
"""
y = solution["y"]
t = solution["t"]
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection="3d")
x = y[2 * coords[0] + coords[1]]
y_data = y[2 * coords[2]]
z = t
ax.plot(x, y_data, z, lw=0.5)
ax.set_xlabel("Coordinate 1")
ax.set_ylabel("Coordinate 2")
ax.set_zlabel("Time")
ax.set_title(title)
return fig
[docs]
def plot_poincare_section(
self,
solution: dict,
section_var: int = 0,
section_value: float = 0.0,
plot_vars: Tuple[int, int] = (1, 2),
title: str = "Poincaré Section",
) -> plt.Figure:
"""
Plot Poincaré section (stroboscopic map).
Args:
solution: Simulation result
section_var: State variable index for section condition
section_value: Value where section is taken
plot_vars: Which variables to plot (indices)
title: Plot title
Returns:
matplotlib Figure
"""
y = solution["y"]
# Find crossings
section_data = y[section_var]
crossings_idx = []
for i in range(len(section_data) - 1):
if (section_data[i] - section_value) * (section_data[i + 1] - section_value) < 0:
if section_data[i + 1] > section_data[i]: # Upward crossing
crossings_idx.append(i)
if len(crossings_idx) < 2:
logger.warning("Not enough crossings for Poincaré section")
return plt.figure()
# Extract values at crossings
x_points = [y[plot_vars[0], i] for i in crossings_idx]
y_points = [y[plot_vars[1], i] for i in crossings_idx]
fig, ax = plt.subplots(figsize=(8, 6))
ax.scatter(x_points, y_points, s=5, alpha=0.5)
ax.set_xlabel(f"State variable {plot_vars[0]}")
ax.set_ylabel(f"State variable {plot_vars[1]}")
ax.set_title(title)
ax.grid(True, alpha=0.3)
return fig