Source code for mechanics_dsl.visualization.phase_space

"""
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