diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index ea33596..719fed0 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -36,17 +36,29 @@ jobs: name: colcon-logs-${{ matrix.ros }} path: ros_ws/log - # The pytest suites only run under colcon above, so cover the wheel path as well. - wheel: + # The pytest suites only run under colcon above, so cover the wheel path the explorer uses. + wheel_and_explorer: runs-on: ubuntu-latest - name: Wheel + name: Wheel and explorer steps: - uses: actions/checkout@v5 - uses: astral-sh/setup-uv@v7 - name: Install package with dev extra run: uv venv && uv pip install -e ".[dev]" - name: Python tests - run: uv run pytest test/test_python_bindings.py -v + run: uv run pytest test/test_python_bindings.py test/test_explorer.py -v + - name: Explorer imports cleanly + run: uv run python -c "import polymath_kinematics.explorer" + - name: Console script resolves and launches + # A healthy headless start proves the entry point doesn't run the app body on import. + run: | + uv run kinematic-explorer --server.headless=true --server.port=8501 & + for _ in $(seq 1 60); do + if curl -sf http://localhost:8501/_stcore/health; then exit 0; fi + sleep 1 + done + echo "explorer failed to become healthy" >&2 + exit 1 standalone_cmake: runs-on: ubuntu-latest diff --git a/CMakeLists.txt b/CMakeLists.txt index e654c43..1569f7d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -133,6 +133,9 @@ if(BUILD_TESTING) ament_add_pytest_test(test_python_bindings test/test_python_bindings.py WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} ) + ament_add_pytest_test(test_explorer test/test_explorer.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + ) endif() endif() diff --git a/package.xml b/package.xml index 8615ed8..1195cf1 100644 --- a/package.xml +++ b/package.xml @@ -17,6 +17,10 @@ ament_cmake_pytest ament_cmake_test catch2 + + python3-pandas + python3-matplotlib ament_cmake diff --git a/polymath_kinematics/explorer/__init__.py b/polymath_kinematics/explorer/__init__.py new file mode 100644 index 0000000..e7f73de --- /dev/null +++ b/polymath_kinematics/explorer/__init__.py @@ -0,0 +1,81 @@ +# Copyright (c) 2025-present Polymath Robotics, Inc. All rights reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kinematic Explorer - trajectory simulation and visualization. + +This subpackage provides pure functions for trajectory simulation, plotting, +and export. The streamlit UI is in kinematic_explorer_app.py. +""" + +from .config import ( + KINEMATIC_EQUATIONS, + LATTICE_CONFIG, + TRAJECTORY_EQUATIONS, + LatticeConfig, +) +from .export import trajectories_to_dataframe +from .plotting import ( + get_traj_attr, + plot_analysis, + plot_articulated_footprint, + plot_lattice, + plot_trajectory_with_footprints, + plot_vehicle_footprint, + select_symmetric_trajectories, +) +from .simulation import ( + generate_lattice_articulated, + generate_lattice_bicycle, + generate_lattice_differential, + single_articulated_trajectory, + single_bicycle_trajectory, + single_differential_trajectory, +) +from .types import ( + AnyTrajectory, + ArticulatedTrajectory, + BicycleTrajectory, + DifferentialTrajectory, + Trajectory, +) + +__all__ = [ + # Types + 'AnyTrajectory', + 'ArticulatedTrajectory', + 'BicycleTrajectory', + 'DifferentialTrajectory', + 'Trajectory', + # Config + 'KINEMATIC_EQUATIONS', + 'LATTICE_CONFIG', + 'LatticeConfig', + 'TRAJECTORY_EQUATIONS', + # Simulation + 'generate_lattice_articulated', + 'generate_lattice_bicycle', + 'generate_lattice_differential', + 'single_articulated_trajectory', + 'single_bicycle_trajectory', + 'single_differential_trajectory', + # Plotting + 'get_traj_attr', + 'plot_analysis', + 'plot_articulated_footprint', + 'plot_lattice', + 'plot_trajectory_with_footprints', + 'plot_vehicle_footprint', + 'select_symmetric_trajectories', + # Export + 'trajectories_to_dataframe', +] diff --git a/polymath_kinematics/explorer/cli.py b/polymath_kinematics/explorer/cli.py new file mode 100644 index 0000000..0d80581 --- /dev/null +++ b/polymath_kinematics/explorer/cli.py @@ -0,0 +1,67 @@ +# Copyright (c) 2025-present Polymath Robotics, Inc. All rights reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Launcher for the Kinematic Explorer Streamlit app. + +Kept out of ``kinematic_explorer_app`` because that module's body *is* the Streamlit script: +importing it from an entry point would run every widget in bare mode before the server starts. +""" + +from __future__ import annotations + +import pathlib +import subprocess +import sys + + +def app_path() -> pathlib.Path: + """Absolute path to the Streamlit app script.""" + return pathlib.Path(__file__).resolve().parent.parent / 'kinematic_explorer_app.py' + + +def main(argv: list[str] | None = None) -> int: + """Run ``streamlit run kinematic_explorer_app.py``, forwarding any extra arguments. + + Extra arguments go to Streamlit, so e.g. ``kinematic-explorer --server.port=8600`` works. + """ + args = list(sys.argv[1:] if argv is None else argv) + + script = app_path() + if not script.is_file(): + print(f'error: could not locate the explorer app at {script}', file=sys.stderr) + return 1 + + command = [ + sys.executable, + '-m', + 'streamlit', + 'run', + str(script), + '--browser.gatherUsageStats=false', + ] + # Default to loopback; pass --server.address=0.0.0.0 to serve over the network. + if not any(arg.startswith('--server.address') for arg in args): + command.append('--server.address=localhost') + command.extend(args) + try: + return subprocess.call(command) + except FileNotFoundError: + print( + 'error: streamlit is not installed. Install the explorer extra with:\n uv pip install -e ".[explorer]"', + file=sys.stderr, + ) + return 1 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/polymath_kinematics/explorer/config.py b/polymath_kinematics/explorer/config.py new file mode 100644 index 0000000..2b25541 --- /dev/null +++ b/polymath_kinematics/explorer/config.py @@ -0,0 +1,120 @@ +# Copyright (c) 2025-present Polymath Robotics, Inc. All rights reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Configuration constants and data structures for kinematic exploration.""" + +from __future__ import annotations + +from dataclasses import dataclass + +# Footprint overhang slider defaults (metres, measured beyond the reference axle/centre). + +# Bicycle: pose reference is the rear axle, so front_overhang_m = wheelbase + front overhang. +DEFAULT_BICYCLE_FRONT_OVERHANG_M = 0.6 +DEFAULT_BICYCLE_REAR_OVERHANG_M = 0.5 + +# Differential drive: pose reference is the body centre, so the overhangs pass through directly. +DEFAULT_DIFFERENTIAL_FRONT_OVERHANG_M = 0.4 +DEFAULT_DIFFERENTIAL_REAR_OVERHANG_M = 0.4 + +# Articulated: base_link is the articulation joint; joint-to-bumper distances are axle distance + +# overhang, so a positive overhang makes the body extend behind the rear axle (counterweight) and +# ahead of the front axle (bucket). +DEFAULT_ARTICULATED_FRONT_OVERHANG_M = 1.0 +DEFAULT_ARTICULATED_REAR_OVERHANG_M = 0.8 + + +@dataclass +class LatticeConfig: + """Configuration for plotting trajectory lattices.""" + + group_key: str + group_label: str + group_unit: str + color_key: str + color_label: str + color_unit: str + color_is_angle: bool = False + angle_key: str = '' # Key for angle in trajectory selection + vel_key: str = '' # Key for velocity in trajectory selection + + +LATTICE_CONFIG: dict[str, LatticeConfig] = { + 'Differential Drive': LatticeConfig( + group_key='base_wheel_velocity', + group_label='Base Wheel Velocity', + group_unit='rad/s', + color_key='angular_velocity', + color_label='Angular Velocity', + color_unit='rad/s', + color_is_angle=False, + angle_key='angular_velocity', + vel_key='base_wheel_velocity', + ), + 'Bicycle': LatticeConfig( + group_key='drive_velocity', + group_label='Drive Velocity', + group_unit='m/s', + color_key='steering_angle', + color_label='Steering Angle', + color_unit='deg', + color_is_angle=True, + angle_key='steering_angle', + vel_key='drive_velocity', + ), + 'Articulated': LatticeConfig( + group_key='drive_velocity', + group_label='Drive Velocity', + group_unit='m/s', + color_key='articulation_angle', + color_label='Articulation Angle', + color_unit='deg', + color_is_angle=True, + angle_key='articulation_angle', + vel_key='drive_velocity', + ), +} + + +KINEMATIC_EQUATIONS = { + 'Differential Drive': { + 'title': 'Differential Drive Kinematics', + 'equations': [ + r'v = \frac{r}{2}(\omega_L + \omega_R)', + r'\omega = \frac{r}{W}(\omega_R - \omega_L)', + ], + 'variables': r'$r$ = wheel radius, $W$ = track width, $\omega_L, \omega_R$ = wheel velocities', + }, + 'Bicycle': { + 'title': 'Bicycle Model Kinematics', + 'equations': [ + r'\omega = \frac{v \tan(\delta)}{L}', + r'R = \frac{L}{\tan(\delta)}', + ], + 'variables': r'$L$ = wheelbase, $\delta$ = steering angle, $R$ = turning radius', + }, + 'Articulated': { + 'title': 'Articulated Vehicle Kinematics (Corke & Ridley)', + 'equations': [ + r'\omega = \frac{v \sin\gamma + L_r \dot{\gamma}}{L_f \cos\gamma + L_r}', + r'R_f = \frac{L_f \cos\gamma + L_r}{\sin\gamma}', + ], + 'variables': r'$L_f, L_r$ = front/rear distances to articulation joint, $\gamma$ = articulation angle, $\dot{\gamma}$ = articulation rate', + 'reference': "Corke & Ridley, IEEE IO'A 2001", + }, +} + +TRAJECTORY_EQUATIONS = r""" +**Trajectory Integration (Euler method):** +$$\dot{x} = v \cos(\theta), \quad \dot{y} = v \sin(\theta), \quad \dot{\theta} = \omega$$ +""" diff --git a/polymath_kinematics/explorer/export.py b/polymath_kinematics/explorer/export.py new file mode 100644 index 0000000..9e40aeb --- /dev/null +++ b/polymath_kinematics/explorer/export.py @@ -0,0 +1,124 @@ +# Copyright (c) 2025-present Polymath Robotics, Inc. All rights reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Export functions for trajectory data.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd + +from .types import AnyTrajectory, ArticulatedTrajectory, BicycleTrajectory, DifferentialTrajectory + + +def trajectories_to_dataframe( + trajectories: list[AnyTrajectory], + model_type: str, +) -> pd.DataFrame: + """Convert trajectories to a pandas DataFrame for export. + + Uses vectorized construction for better performance. + + Args: + trajectories: List of trajectory dataclasses + model_type: One of "Differential Drive", "Bicycle", or "Articulated" + + Returns: + DataFrame with trajectory data, columns depend on model type + """ + if not trajectories: + return pd.DataFrame() + + # Calculate total size for preallocation + num_total = sum(len(trajectory) for trajectory in trajectories) + + # Preallocate arrays for common columns + trajectory_ids = np.empty(num_total, dtype=int) + time_array = np.empty(num_total) + x_array = np.empty(num_total) + y_array = np.empty(num_total) + theta_array = np.empty(num_total) + linear_velocity_array = np.empty(num_total) + angular_velocity_array = np.empty(num_total) + + # Model-specific arrays + if model_type == 'Differential Drive': + left_wheel_array = np.empty(num_total) + right_wheel_array = np.empty(num_total) + base_wheel_velocity_array = np.empty(num_total) + elif model_type == 'Bicycle': + drive_velocity_array = np.empty(num_total) + steering_angle_array = np.empty(num_total) + turning_radius_array = np.empty(num_total) + elif model_type == 'Articulated': + drive_velocity_array = np.empty(num_total) + articulation_angle_array = np.empty(num_total) + turning_radius_array = np.empty(num_total) + + # Fill arrays + current_index = 0 + for trajectory_index, trajectory in enumerate(trajectories): + num_points = len(trajectory) + end_index = current_index + num_points + + trajectory_ids[current_index:end_index] = trajectory_index + time_array[current_index:end_index] = trajectory.time + x_array[current_index:end_index] = trajectory.x + y_array[current_index:end_index] = trajectory.y + theta_array[current_index:end_index] = trajectory.theta + linear_velocity_array[current_index:end_index] = trajectory.linear_velocity + angular_velocity_array[current_index:end_index] = trajectory.angular_velocity + + if model_type == 'Differential Drive': + trajectory_differential: DifferentialTrajectory = trajectory # type: ignore[assignment] + left_wheel_array[current_index:end_index] = trajectory_differential.left_wheel + right_wheel_array[current_index:end_index] = trajectory_differential.right_wheel + base_wheel_velocity_array[current_index:end_index] = trajectory_differential.base_wheel_velocity + elif model_type == 'Bicycle': + trajectory_bicycle: BicycleTrajectory = trajectory # type: ignore[assignment] + drive_velocity_array[current_index:end_index] = trajectory_bicycle.drive_velocity + steering_angle_array[current_index:end_index] = trajectory_bicycle.steering_angle + turning_radius_array[current_index:end_index] = trajectory_bicycle.turning_radius + elif model_type == 'Articulated': + trajectory_articulated: ArticulatedTrajectory = trajectory # type: ignore[assignment] + drive_velocity_array[current_index:end_index] = trajectory_articulated.drive_velocity + articulation_angle_array[current_index:end_index] = trajectory_articulated.articulation_angle + turning_radius_array[current_index:end_index] = trajectory_articulated.turning_radius + + current_index = end_index + + # Build DataFrame from arrays (single allocation) + data: dict[str, np.ndarray] = { + 'trajectory_id': trajectory_ids, + 'time': time_array, + 'x': x_array, + 'y': y_array, + 'theta': theta_array, + 'linear_velocity': linear_velocity_array, + 'angular_velocity': angular_velocity_array, + } + + if model_type == 'Differential Drive': + data['left_wheel_velocity'] = left_wheel_array + data['right_wheel_velocity'] = right_wheel_array + data['base_wheel_velocity'] = base_wheel_velocity_array + elif model_type == 'Bicycle': + data['drive_velocity'] = drive_velocity_array + data['steering_angle'] = steering_angle_array + data['turning_radius'] = turning_radius_array + elif model_type == 'Articulated': + data['drive_velocity'] = drive_velocity_array + data['articulation_angle'] = articulation_angle_array + data['turning_radius'] = turning_radius_array + + return pd.DataFrame(data) diff --git a/polymath_kinematics/explorer/plotting.py b/polymath_kinematics/explorer/plotting.py new file mode 100644 index 0000000..eea164a --- /dev/null +++ b/polymath_kinematics/explorer/plotting.py @@ -0,0 +1,542 @@ +# Copyright (c) 2025-present Polymath Robotics, Inc. All rights reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Plotting functions for kinematic visualization. + +All functions return matplotlib Figure objects and have no streamlit dependencies. +""" + +from __future__ import annotations + +import matplotlib.pyplot as plt +import numpy as np + +from .config import LATTICE_CONFIG +from .types import AnyTrajectory, ArticulatedTrajectory, BicycleTrajectory, DifferentialTrajectory + + +def get_traj_attr(trajectory: AnyTrajectory, key: str) -> float: + """Get trajectory attribute by key name.""" + return getattr(trajectory, key) + + +def plot_vehicle_footprint( + ax: plt.Axes, + x: float, + y: float, + theta: float, + body_corners: np.ndarray, + color: str = 'blue', + alpha: float = 0.5, + steering_angle: float | None = None, + front_axle_offset_m: float | None = None, + track_width_m: float | None = None, + wheel_length_frac: float = 0.18, + wheel_width_frac: float = 0.08, +) -> None: + """Draw a single-body vehicle footprint at a given pose. + + ``body_corners`` is the projector's world-frame ``footprint`` — any vertex count, so the + heading arrow is measured from the polygon's extent along ``theta`` rather than from fixed + corner indices. + + ``steering_angle`` (radians) draws two front-wheel indicators at ``front_axle_offset_m`` + ahead of the pose reference, spaced by ``track_width_m``. Pass ``None`` for models without + steered wheels (e.g. differential drive). + """ + cos_theta, sin_theta = np.cos(theta), np.sin(theta) + rotation_matrix = np.array([[cos_theta, -sin_theta], [sin_theta, cos_theta]]) + + body = np.asarray(body_corners, dtype=float) + ax.fill(body[:, 0], body[:, 1], color=color, alpha=alpha, edgecolor=color, linewidth=0.5) + + # Heading arrow spans the body's extent along the heading, measured through the pose. + origin = np.array([x, y]) + heading = np.array([cos_theta, sin_theta]) + lateral = np.array([-sin_theta, cos_theta]) + relative = body - origin + along = relative @ heading + across = relative @ lateral + body_length = float(along.max() - along.min()) + body_width = float(across.max() - across.min()) + arrow_start = origin + heading * along.min() + ax.arrow( + arrow_start[0], + arrow_start[1], + heading[0] * body_length, + heading[1] * body_length, + head_width=body_width * 0.3, + head_length=body_length * 0.1, + fc=color, + ec=color, + alpha=min(1.0, alpha + 0.3), + length_includes_head=True, + ) + + if steering_angle is not None and front_axle_offset_m is not None and track_width_m is not None: + # Front-wheel indicators: rectangles at the front axle, offset left/right by half the + # track width, each rotated by steering_angle in the body frame. + wheel_l = body_length * wheel_length_frac + wheel_w = body_width * wheel_width_frac + wheel_corners_centered = np.array([ + [-wheel_l / 2.0, -wheel_w / 2.0], + [+wheel_l / 2.0, -wheel_w / 2.0], + [+wheel_l / 2.0, +wheel_w / 2.0], + [-wheel_l / 2.0, +wheel_w / 2.0], + [-wheel_l / 2.0, -wheel_w / 2.0], + ]) + cos_s, sin_s = np.cos(steering_angle), np.sin(steering_angle) + wheel_rotation = np.array([[cos_s, -sin_s], [sin_s, cos_s]]) + wheel_corners_steered = (wheel_rotation @ wheel_corners_centered.T).T # (5, 2) + for lateral_offset in (-track_width_m / 2.0, +track_width_m / 2.0): + corners_body = wheel_corners_steered + np.array([front_axle_offset_m, lateral_offset]) + corners_in_world = (rotation_matrix @ corners_body.T).T + np.array([x, y]) + ax.fill( + corners_in_world[:, 0], + corners_in_world[:, 1], + color=color, + alpha=min(1.0, alpha + 0.2), + edgecolor=color, + linewidth=0.5, + ) + + +def plot_articulated_footprint( + ax: plt.Axes, + front_corners: np.ndarray, + rear_corners: np.ndarray, + joint_xy: np.ndarray, + front_theta: float, + color: str = 'blue', + alpha: float = 0.5, +) -> None: + """Draw an articulated vehicle footprint: two independently posed body polygons. + + Each polygon comes from the projector already in the world frame, anchored at its own axle, + so nothing here assumes a vertex count or ordering. ``joint_xy`` and ``front_theta`` come from + the projector's ``joint_pose`` and articulation angle; they place the joint marker and the + front-body heading arrow, which cannot be inferred from an arbitrary polygon. + """ + fc = np.asarray(front_corners, dtype=float) + rc = np.asarray(rear_corners, dtype=float) + joint = np.asarray(joint_xy, dtype=float) + + ax.fill(fc[:, 0], fc[:, 1], color=color, alpha=alpha, edgecolor=color, linewidth=0.5) + ax.fill(rc[:, 0], rc[:, 1], color=color, alpha=alpha * 0.8, edgecolor=color, linewidth=0.5) + ax.plot(joint[0], joint[1], 'o', color=color, markersize=4, alpha=min(1.0, alpha + 0.3)) + + # Arrow runs from the joint to the front body's forward extent along the front-body heading. + heading = np.array([np.cos(front_theta), np.sin(front_theta)]) + lateral = np.array([-np.sin(front_theta), np.cos(front_theta)]) + relative = fc - joint + front_length = float((relative @ heading).max()) + front_width = float(np.ptp(relative @ lateral)) + if front_length > 0.0: + ax.arrow( + joint[0], + joint[1], + heading[0] * front_length, + heading[1] * front_length, + head_width=front_width * 0.2, + head_length=front_length * 0.08, + fc=color, + ec=color, + alpha=min(1.0, alpha + 0.3), + length_includes_head=True, + ) + + +def select_symmetric_trajectories( + trajectories: list[AnyTrajectory], + model_type: str, + num_angles: int = 5, + num_velocities: int = 1, +) -> list[AnyTrajectory]: + """Select a symmetric subset of trajectories for visualization. + + Args: + trajectories: Full list of trajectory dataclasses + model_type: Model type to determine which keys to use + num_angles: Number of steering/articulation angles (should be odd for symmetry) + num_velocities: Number of velocities to include + + Returns: + Filtered list of trajectories, symmetrically selected + """ + if not trajectories: + return [] + + config = LATTICE_CONFIG.get(model_type) + if config is None: + return trajectories[:5] + + angle_key = config.angle_key + vel_key = config.vel_key + + all_angles = sorted(set(get_traj_attr(trajectory, angle_key) for trajectory in trajectories)) + all_velocities = sorted(set(get_traj_attr(trajectory, vel_key) for trajectory in trajectories)) + + # Select symmetric angles + if len(all_angles) <= num_angles: + selected_angles = set(all_angles) + else: + indices = np.linspace(0, len(all_angles) - 1, num_angles, dtype=int) + selected_angles = {all_angles[i] for i in indices} + + # Select velocities + if len(all_velocities) <= num_velocities: + selected_velocities = set(all_velocities) + else: + indices = np.linspace(0, len(all_velocities) - 1, num_velocities, dtype=int) + selected_velocities = {all_velocities[i] for i in indices} + + selected = [ + trajectory + for trajectory in trajectories + if any(abs(get_traj_attr(trajectory, angle_key) - angle) < 1e-6 for angle in selected_angles) + and any(abs(get_traj_attr(trajectory, vel_key) - velocity) < 1e-6 for velocity in selected_velocities) + ] + + selected.sort(key=lambda trajectory: get_traj_attr(trajectory, angle_key)) + return selected + + +def _build_legend_label(trajectory: AnyTrajectory, model_type: str) -> str: + """Build legend label for a trajectory based on model type.""" + if model_type == 'Articulated': + articulated: ArticulatedTrajectory = trajectory # type: ignore[assignment] + return f'y={np.degrees(articulated.articulation_angle):+.0f} deg, v={articulated.drive_velocity:.1f}' + elif model_type == 'Bicycle': + bicycle: BicycleTrajectory = trajectory # type: ignore[assignment] + return f'd={np.degrees(bicycle.steering_angle):+.0f} deg, v={bicycle.drive_velocity:.1f}' + elif model_type == 'Differential Drive': + differential: DifferentialTrajectory = trajectory # type: ignore[assignment] + return f'w={differential.angular_velocity:+.1f}, base={differential.base_wheel_velocity:.0f}' + return '' + + +def plot_trajectory_with_footprints( + trajectories: list[AnyTrajectory], + model_type: str, + model_params: dict, + num_footprints: int = 5, +) -> plt.Figure: + """Plot selected trajectories with vehicle footprints at intervals. + + A trajectory with no footprint series (projector given no body dims) still has its path + drawn; only the overlay is skipped. + + Args: + trajectories: List of trajectory dataclasses (already filtered/selected) + model_type: "Differential Drive", "Bicycle", or "Articulated" + model_params: Bicycle only — ``wheelbase`` and ``track_width`` for the steered-wheel + indicators. Unused for the other models. + num_footprints: Number of footprints per trajectory + + Returns: + matplotlib Figure + """ + fig, ax = plt.subplots(figsize=(16, 9), layout='constrained') + + colormap = plt.cm.tab10 + legend_handles = [] + legend_labels = [] + + for index, trajectory in enumerate(trajectories): + color = colormap(index % 10) + + (line,) = ax.plot(trajectory.x, trajectory.y, color=color, linewidth=2, alpha=0.7) + legend_handles.append(line) + legend_labels.append(_build_legend_label(trajectory, model_type)) + + num_points = len(trajectory) + footprint_indices = np.linspace(0, num_points - 1, num_footprints, dtype=int) + + for footprint_num, footprint_index in enumerate(footprint_indices): + footprint_alpha = 0.3 + 0.5 * (footprint_num / max(1, num_footprints - 1)) + + if model_type == 'Articulated': + articulated: ArticulatedTrajectory = trajectory # type: ignore[assignment] + if ( + articulated.front_footprint_series is None + or articulated.rear_footprint_series is None + or articulated.joint_pose_series is None + ): + continue + joint_x, joint_y, joint_theta = articulated.joint_pose_series[footprint_index] + if articulated.articulation_angle_series is not None: + gamma = float(articulated.articulation_angle_series[footprint_index]) + else: + gamma = articulated.articulation_angle + plot_articulated_footprint( + ax, + articulated.front_footprint_series[footprint_index], + articulated.rear_footprint_series[footprint_index], + joint_xy=np.array([joint_x, joint_y]), + # joint_pose.theta is the rear-body heading; the front body leads it by gamma. + front_theta=joint_theta + gamma, + color=color, + alpha=footprint_alpha, + ) + elif model_type == 'Bicycle': + bicycle: BicycleTrajectory = trajectory # type: ignore[assignment] + if bicycle.footprint_series is None: + continue + if bicycle.steering_angle_series is not None: + sample_steering_angle = float(bicycle.steering_angle_series[footprint_index]) + else: + sample_steering_angle = bicycle.steering_angle + plot_vehicle_footprint( + ax, + trajectory.x[footprint_index], + trajectory.y[footprint_index], + trajectory.theta[footprint_index], + bicycle.footprint_series[footprint_index], + color=color, + alpha=footprint_alpha, + steering_angle=sample_steering_angle, + front_axle_offset_m=model_params['wheelbase'], + track_width_m=model_params['track_width'], + ) + else: # Differential Drive — no steered wheels. + if trajectory.footprint_series is None: + continue + plot_vehicle_footprint( + ax, + trajectory.x[footprint_index], + trajectory.y[footprint_index], + trajectory.theta[footprint_index], + trajectory.footprint_series[footprint_index], + color=color, + alpha=footprint_alpha, + ) + + ax.set_xlabel('X (m)') + ax.set_ylabel('Y (m)') + ax.set_title(f'Trajectory with Vehicle Footprints ({model_type})') + # Equal data scaling (circles stay circular) but force a 16:9 landscape box; the view + # limits expand to fill the box rather than distorting the geometry. + ax.set_aspect('equal') + ax.set_box_aspect(9 / 16) + ax.grid(True, alpha=0.3) + ax.axhline(y=0, color='k', linewidth=0.5) + ax.axvline(x=0, color='k', linewidth=0.5) + + if legend_handles: + title = { + 'Articulated': 'Articulation', + 'Bicycle': 'Steering', + 'Differential Drive': 'Angular Vel', + }.get(model_type, '') + ax.legend(legend_handles, legend_labels, loc='upper right', title=title) + + return fig + + +def plot_lattice( + trajectories: list[AnyTrajectory], + model_type: str, + group_values: list[float], +) -> plt.Figure: + """Plot trajectory lattice grouped by velocity, colored by steering/articulation angle. + + Args: + trajectories: List of trajectory dataclasses + model_type: "Differential Drive", "Bicycle", or "Articulated" + group_values: List of velocity values to create subplots for + + Returns: + matplotlib Figure + """ + config = LATTICE_CONFIG[model_type] + + num_columns = len(group_values) + # 16:9 landscape; grow width for the (rare) multi-column case. + fig_width = max(16, 8 * num_columns) + fig, axes = plt.subplots(1, num_columns, figsize=(fig_width, 9), layout='constrained') + + if num_columns == 1: + axes = [axes] + + color_values = [get_traj_attr(trajectory, config.color_key) for trajectory in trajectories] + if config.color_is_angle: + color_values_display = [np.degrees(value) for value in color_values] + else: + color_values_display = color_values + + colormap = plt.cm.coolwarm + normalizer = plt.Normalize(vmin=min(color_values_display), vmax=max(color_values_display)) + + for ax, group_value in zip(axes, group_values): + ax.set_title(f'{config.group_label} = {group_value:.1f} {config.group_unit}') + ax.set_xlabel('X (m)') + ax.set_ylabel('Y (m)') + # Equal data scaling (circles stay circular) with a 16:9 landscape box. + ax.set_aspect('equal') + ax.set_box_aspect(9 / 16) + ax.grid(True, alpha=0.3) + ax.axhline(y=0, color='k', linewidth=0.5) + ax.axvline(x=0, color='k', linewidth=0.5) + + for trajectory in trajectories: + if abs(get_traj_attr(trajectory, config.group_key) - group_value) < 0.001: + raw_color_value = get_traj_attr(trajectory, config.color_key) + color_value = np.degrees(raw_color_value) if config.color_is_angle else raw_color_value + color = colormap(normalizer(color_value)) + ax.plot(trajectory.x, trajectory.y, color=color, linewidth=1.5, alpha=0.8) + ax.plot(trajectory.x[-1], trajectory.y[-1], 'o', color=color, markersize=3) + + scalar_mappable = plt.cm.ScalarMappable(cmap=colormap, norm=normalizer) + scalar_mappable.set_array([]) + colorbar = fig.colorbar(scalar_mappable, ax=axes, orientation='horizontal', fraction=0.05, pad=0.12) + colorbar.set_label(f'{config.color_label} ({config.color_unit})') + + return fig + + +def plot_analysis( + trajectories: list[AnyTrajectory], + model_type: str, + group_values: list[float], +) -> plt.Figure: + """Plot kinematic analysis for all velocities in the lattice. + + Args: + trajectories: List of trajectory dataclasses + model_type: "Differential Drive", "Bicycle", or "Articulated" + group_values: List of velocity values to analyze + + Returns: + matplotlib Figure with two subplots + """ + fig, axes = plt.subplots(1, 2, figsize=(10, 4), layout='constrained') + + if model_type == 'Articulated': + ax_left = axes[0] + for velocity in group_values: + angles = [] + angular_velocities = [] + for trajectory in trajectories: + articulated: ArticulatedTrajectory = trajectory # type: ignore[assignment] + if abs(articulated.drive_velocity - velocity) < 0.001: + angles.append(np.degrees(articulated.articulation_angle)) + angular_velocities.append(articulated.angular_velocity) + if angles: + ax_left.plot(angles, angular_velocities, 'o-', label=f'v={velocity:.1f} m/s', markersize=5) + ax_left.set_xlabel('Articulation Angle (deg)') + ax_left.set_ylabel('Angular Velocity (rad/s)') + ax_left.set_title('Articulation -> Angular Velocity') + ax_left.legend() + ax_left.grid(True, alpha=0.3) + + ax_right = axes[1] + radii = [] + angles = [] + for trajectory in trajectories: + articulated: ArticulatedTrajectory = trajectory # type: ignore[assignment] + if abs(articulated.articulation_angle) > 0.00001: + radii.append(abs(articulated.turning_radius)) + angles.append(np.degrees(abs(articulated.articulation_angle))) + if angles: + ax_right.plot(angles, radii, 'o-', markersize=5) + + ax_right.set_xlabel('|Articulation Angle| (deg)') + ax_right.set_ylabel('Turning Radius (m)') + ax_right.set_title('Turning Radius vs Articulation') + # Single unlabeled series — no legend. + ax_right.grid(True, alpha=0.3) + ax_right.set_ylim(bottom=0) + + elif model_type == 'Bicycle': + ax_left = axes[0] + for velocity in group_values: + angles = [] + angular_velocities = [] + for trajectory in trajectories: + bicycle: BicycleTrajectory = trajectory # type: ignore[assignment] + if abs(bicycle.drive_velocity - velocity) < 0.001: + angles.append(np.degrees(bicycle.steering_angle)) + angular_velocities.append(bicycle.angular_velocity) + if angles: + ax_left.plot(angles, angular_velocities, 'o-', label=f'v={velocity:.1f} m/s', markersize=5) + ax_left.set_xlabel('Steering Angle (deg)') + ax_left.set_ylabel('Angular Velocity (rad/s)') + ax_left.set_title('Steering -> Angular Velocity') + ax_left.legend() + ax_left.grid(True, alpha=0.3) + + ax_right = axes[1] + + radii = [] + angles = [] + for trajectory in trajectories: + bicycle: BicycleTrajectory = trajectory # type: ignore[assignment] + if abs(bicycle.steering_angle) > 0.00001: + radii.append(abs(bicycle.turning_radius)) + angles.append(np.degrees(abs(bicycle.steering_angle))) + + if angles: + ax_right.plot(angles, radii, 'o-', markersize=5) + + ax_right.set_xlabel('|Steering Angle| (deg)') + ax_right.set_ylabel('Turning Radius (m)') + ax_right.set_title('Turning Radius vs Steering') + # Single unlabeled series — no legend. + ax_right.grid(True, alpha=0.3) + ax_right.set_ylim(bottom=0) + + elif model_type == 'Differential Drive': + ax_left = axes[0] + base_velocity = group_values[len(group_values) // 2] + left_wheel_velocities = [] + right_wheel_velocities = [] + angular_velocities = [] + for trajectory in trajectories: + differential: DifferentialTrajectory = trajectory # type: ignore[assignment] + if abs(differential.base_wheel_velocity - base_velocity) < 0.001: + left_wheel_velocities.append(differential.left_wheel) + right_wheel_velocities.append(differential.right_wheel) + angular_velocities.append(differential.angular_velocity) + if left_wheel_velocities: + ax_left.plot(angular_velocities, left_wheel_velocities, 'o-', label='Left wheel', markersize=5) + ax_left.plot(angular_velocities, right_wheel_velocities, 's-', label='Right wheel', markersize=5) + ax_left.set_xlabel('Angular Velocity (rad/s)') + ax_left.set_ylabel('Wheel Velocity (rad/s)') + ax_left.set_title(f'Wheel Velocities (base={base_velocity:.0f} rad/s)') + ax_left.legend() + ax_left.grid(True, alpha=0.3) + + ax_right = axes[1] + for base_velocity in group_values: + linear_velocities = [] + angular_velocities_list = [] + for trajectory in trajectories: + differential: DifferentialTrajectory = trajectory # type: ignore[assignment] + if abs(differential.base_wheel_velocity - base_velocity) < 0.001: + linear_velocities.append(differential.linear_velocity) + angular_velocities_list.append(differential.angular_velocity) + if linear_velocities: + ax_right.plot( + angular_velocities_list, + linear_velocities, + 'o-', + label=f'base={base_velocity:.0f} rad/s', + markersize=5, + ) + ax_right.set_xlabel('Angular Velocity (rad/s)') + ax_right.set_ylabel('Linear Velocity (m/s)') + ax_right.set_title('Body Velocities') + ax_right.legend() + ax_right.grid(True, alpha=0.3) + + return fig diff --git a/polymath_kinematics/explorer/simulation.py b/polymath_kinematics/explorer/simulation.py new file mode 100644 index 0000000..66e513d --- /dev/null +++ b/polymath_kinematics/explorer/simulation.py @@ -0,0 +1,490 @@ +# Copyright (c) 2025-present Polymath Robotics, Inc. All rights reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Trajectory simulation and lattice generation functions. + +Every trajectory comes from the matching C++ projector, so the explorer and the production +kinematics share one forward-simulation codepath. Lattice cells set initial == target to hold +a constant command; the single-trajectory helpers drive initial != target to show the ramp. + +All functions are pure (no streamlit dependency) for testability and reuse. +""" + +from __future__ import annotations + +import numpy as np + +from polymath_kinematics import ( + ArticulatedModel, + ArticulatedProjector, + AxleReference, + BicycleModel, + BicycleProjector, + DifferentialDriveModel, + DifferentialDriveProjector, + Footprint, + Pose2D, +) + +from .types import ArticulatedTrajectory, BicycleTrajectory, DifferentialTrajectory + + +def _states_to_arrays(states): + """Pull (time, x, y, theta, omega) arrays out of a list of ProjectedState samples.""" + n = len(states) + time_arr = np.empty(n) + x_arr = np.empty(n) + y_arr = np.empty(n) + theta_arr = np.empty(n) + omega_arr = np.empty(n) + for i, s in enumerate(states): + time_arr[i] = s.time_s + x_arr[i] = s.pose.x + y_arr[i] = s.pose.y + theta_arr[i] = s.pose.theta + omega_arr[i] = s.angular_velocity_rad_s + return time_arr, x_arr, y_arr, theta_arr, omega_arr + + +def _footprint_series(states, attr: str): + """Stack the per-sample footprint polygons from a projection into an (N, V, 2) array. + + ``attr`` is the projected-state attribute holding the footprint (``footprint`` for single-body + models, ``front_footprint`` / ``rear_footprint`` for articulated). V is the polygon's vertex + count. Returns ``None`` if any sample has an empty footprint (no footprint was set on the + projector), so downstream consumers can fall back cleanly. + """ + corners = [] + for s in states: + polygon = getattr(s, attr) + if not polygon: + return None + corners.append([[p.x, p.y] for p in polygon]) + return np.asarray(corners, dtype=float) + + +def _joint_pose_series(states): + """Stack the articulated projector's per-sample joint pose into an (N, 3) [x, y, theta] array.""" + return np.asarray([[s.joint_pose.x, s.joint_pose.y, s.joint_pose.theta] for s in states], dtype=float) + + +def generate_lattice_differential( + wheel_radius: float, + track_width: float, + base_wheel_velocities: tuple[float, ...], + wheel_velocity_diffs: tuple[float, ...], + duration: float, + time_step: float = 0.02, + footprint: Footprint | None = None, +) -> list[DifferentialTrajectory]: + """Generate trajectory lattice by sweeping wheel velocity differences. + + Each wheel-velocity pair becomes a body command, then a constant-command projection. + """ + model = DifferentialDriveModel(wheel_radius, track_width) + + # Resolve all cells first: the projector bounds must bracket the sweep, or outer cells clamp. + cells = [] + for base_velocity in base_wheel_velocities: + for velocity_diff in wheel_velocity_diffs: + left_wheel = base_velocity - velocity_diff / 2 + right_wheel = base_velocity + velocity_diff / 2 + body_velocity = model.wheel_velocities_to_body_velocity(left_wheel, right_wheel) + cells.append(( + base_velocity, + left_wheel, + right_wheel, + body_velocity.linear_velocity_m_s, + body_velocity.angular_velocity_rad_s, + )) + + if not cells: + return [] + + linear_velocities = [cell[3] for cell in cells] + angular_velocities = [cell[4] for cell in cells] + projector = DifferentialDriveProjector( + model, + min(linear_velocities), + max(linear_velocities), + min(angular_velocities), + max(angular_velocities), + footprint or [], + ) + + trajectories: list[DifferentialTrajectory] = [] + for base_velocity, left_wheel, right_wheel, linear_velocity, angular_velocity in cells: + states = projector.project( + horizon_s=duration, + dt_s=time_step, + initial_pose=Pose2D(), + initial_linear_velocity_m_s=linear_velocity, + initial_angular_velocity_rad_s=angular_velocity, + target_linear_velocity_m_s=linear_velocity, + target_angular_velocity_rad_s=angular_velocity, + # initial == target, so the ramp never fires and the acceleration is irrelevant. + linear_acceleration_m_s2=0.0, + angular_acceleration_rad_s2=0.0, + ) + time_arr, x_arr, y_arr, theta_arr, _omega = _states_to_arrays(states) + + trajectories.append( + DifferentialTrajectory( + time=time_arr, + x=x_arr, + y=y_arr, + theta=theta_arr, + linear_velocity=linear_velocity, + angular_velocity=angular_velocity, + left_wheel=left_wheel, + right_wheel=right_wheel, + base_wheel_velocity=base_velocity, + footprint_series=_footprint_series(states, 'footprint'), + ) + ) + + return trajectories + + +def generate_lattice_bicycle( + wheelbase: float, + track_width: float, + wheel_radius: float, + drive_velocities: tuple[float, ...], + steering_angles: tuple[float, ...], + duration: float, + time_step: float = 0.02, + min_steering_angle_rad: float | None = None, + max_steering_angle_rad: float | None = None, + steering_rate_rad_s: float = 0.0, + axle_reference: AxleReference = AxleReference.REAR, + footprint: Footprint | None = None, +) -> list[BicycleTrajectory]: + """Generate trajectory lattice by sweeping steering angles. + + Each lattice cell is generated via BicycleProjector.project() with the + initial steering angle equal to the target, so no ramp occurs and the + behavior matches a constant-steering sweep. ``steering_rate_rad_s`` is + exposed for future use; with a non-zero value and ``min/max`` set, the + projector will ramp toward the target — but a constant-input lattice + is the typical use. + """ + angles = list(steering_angles) + if not angles: + return [] + min_angle = min_steering_angle_rad if min_steering_angle_rad is not None else min(angles) + max_angle = max_steering_angle_rad if max_steering_angle_rad is not None else max(angles) + + model = BicycleModel(wheelbase, track_width, wheel_radius) + projector = BicycleProjector(model, min_angle, max_angle, axle_reference, footprint or []) + + trajectories: list[BicycleTrajectory] = [] + for drive_velocity in drive_velocities: + for steering_angle in angles: + states = projector.project( + horizon_s=duration, + dt_s=time_step, + initial_pose=Pose2D(), + initial_steering_angle_rad=steering_angle, + target_steering_angle_rad=steering_angle, + steering_rate_rad_s=steering_rate_rad_s, + linear_velocity_m_s=drive_velocity, + ) + time_arr, x_arr, y_arr, theta_arr, _omega = _states_to_arrays(states) + + trajectories.append( + BicycleTrajectory( + time=time_arr, + x=x_arr, + y=y_arr, + theta=theta_arr, + linear_velocity=drive_velocity, + angular_velocity=states[0].angular_velocity_rad_s, + drive_velocity=drive_velocity, + steering_angle=steering_angle, + turning_radius=states[0].steering_state.turning_radius_m, + footprint_series=_footprint_series(states, 'footprint'), + ) + ) + + return trajectories + + +def generate_lattice_articulated( + articulation_to_front: float, + articulation_to_rear: float, + front_track: float, + rear_track: float, + front_wheel_radius: float, + rear_wheel_radius: float, + drive_velocities: tuple[float, ...], + articulation_angles: tuple[float, ...], + duration: float, + time_step: float = 0.02, + min_articulation_angle_rad: float | None = None, + max_articulation_angle_rad: float | None = None, + articulation_rate_rad_s: float = 0.0, + axle_reference: AxleReference = AxleReference.REAR, + front_footprint: Footprint | None = None, + rear_footprint: Footprint | None = None, +) -> list[ArticulatedTrajectory]: + """Generate trajectory lattice by sweeping articulation angles. + + See ``generate_lattice_bicycle`` for the projector-based rationale. + """ + angles = list(articulation_angles) + if not angles: + return [] + min_angle = min_articulation_angle_rad if min_articulation_angle_rad is not None else min(angles) + max_angle = max_articulation_angle_rad if max_articulation_angle_rad is not None else max(angles) + + model = ArticulatedModel( + articulation_to_front, + articulation_to_rear, + front_track, + rear_track, + front_wheel_radius, + rear_wheel_radius, + ) + projector = ArticulatedProjector( + model, min_angle, max_angle, axle_reference, front_footprint or [], rear_footprint or [] + ) + + trajectories: list[ArticulatedTrajectory] = [] + for drive_velocity in drive_velocities: + for articulation_angle in angles: + states = projector.project( + horizon_s=duration, + dt_s=time_step, + initial_pose=Pose2D(), + initial_articulation_angle_rad=articulation_angle, + target_articulation_angle_rad=articulation_angle, + articulation_rate_rad_s=articulation_rate_rad_s, + linear_velocity_m_s=drive_velocity, + ) + time_arr, x_arr, y_arr, theta_arr, _omega = _states_to_arrays(states) + + trajectories.append( + ArticulatedTrajectory( + time=time_arr, + x=x_arr, + y=y_arr, + theta=theta_arr, + linear_velocity=drive_velocity, + angular_velocity=states[0].angular_velocity_rad_s, + drive_velocity=drive_velocity, + articulation_angle=articulation_angle, + turning_radius=states[0].vehicle_state.front_axle_turning_radius_m, + front_footprint_series=_footprint_series(states, 'front_footprint'), + rear_footprint_series=_footprint_series(states, 'rear_footprint'), + joint_pose_series=_joint_pose_series(states), + ) + ) + + return trajectories + + +# ---------------------------------------------------------------------------- +# Single-trajectory helpers — used by the "Single Projected Trajectory" sections +# of the Streamlit explorer. Each returns one trajectory generated by the +# corresponding C++ projector, packed into a dataclass that the existing +# plot_trajectory_with_footprints helper can consume. +# ---------------------------------------------------------------------------- + + +def single_bicycle_trajectory( + wheelbase: float, + track_width: float, + wheel_radius: float, + initial_steering_angle_rad: float, + target_steering_angle_rad: float, + steering_rate_rad_s: float, + drive_velocity: float, + duration: float, + time_step: float = 0.02, + min_steering_angle_rad: float | None = None, + max_steering_angle_rad: float | None = None, + axle_reference: AxleReference = AxleReference.REAR, + footprint: Footprint | None = None, +) -> BicycleTrajectory: + """Ramp a bicycle steering angle from `initial` toward `target` at `rate` rad/s + over `duration` seconds, returning the resulting trajectory. + + Default min/max bracket `[min(initial, target), max(initial, target)]` so the + clamp never fires unless the caller explicitly sets bounds. + """ + if min_steering_angle_rad is None: + min_steering_angle_rad = min(initial_steering_angle_rad, target_steering_angle_rad) + if max_steering_angle_rad is None: + max_steering_angle_rad = max(initial_steering_angle_rad, target_steering_angle_rad) + + model = BicycleModel(wheelbase, track_width, wheel_radius) + projector = BicycleProjector(model, min_steering_angle_rad, max_steering_angle_rad, axle_reference, footprint or []) + states = projector.project( + horizon_s=duration, + dt_s=time_step, + initial_pose=Pose2D(), + initial_steering_angle_rad=initial_steering_angle_rad, + target_steering_angle_rad=target_steering_angle_rad, + steering_rate_rad_s=steering_rate_rad_s, + linear_velocity_m_s=drive_velocity, + ) + time_arr, x_arr, y_arr, theta_arr, _omega = _states_to_arrays(states) + steering_series = np.asarray([s.steering_angle_rad for s in states]) + return BicycleTrajectory( + time=time_arr, + x=x_arr, + y=y_arr, + theta=theta_arr, + linear_velocity=drive_velocity, + angular_velocity=states[0].angular_velocity_rad_s, + drive_velocity=drive_velocity, + steering_angle=target_steering_angle_rad, + turning_radius=states[-1].steering_state.turning_radius_m, + steering_angle_series=steering_series, + footprint_series=_footprint_series(states, 'footprint'), + ) + + +def single_articulated_trajectory( + articulation_to_front: float, + articulation_to_rear: float, + front_track: float, + rear_track: float, + front_wheel_radius: float, + rear_wheel_radius: float, + initial_articulation_angle_rad: float, + target_articulation_angle_rad: float, + articulation_rate_rad_s: float, + drive_velocity: float, + duration: float, + time_step: float = 0.02, + min_articulation_angle_rad: float | None = None, + max_articulation_angle_rad: float | None = None, + axle_reference: AxleReference = AxleReference.REAR, + front_footprint: Footprint | None = None, + rear_footprint: Footprint | None = None, +) -> ArticulatedTrajectory: + """Ramp an articulation angle from `initial` toward `target` at `rate` rad/s.""" + if min_articulation_angle_rad is None: + min_articulation_angle_rad = min(initial_articulation_angle_rad, target_articulation_angle_rad) + if max_articulation_angle_rad is None: + max_articulation_angle_rad = max(initial_articulation_angle_rad, target_articulation_angle_rad) + + model = ArticulatedModel( + articulation_to_front, + articulation_to_rear, + front_track, + rear_track, + front_wheel_radius, + rear_wheel_radius, + ) + projector = ArticulatedProjector( + model, + min_articulation_angle_rad, + max_articulation_angle_rad, + axle_reference, + front_footprint or [], + rear_footprint or [], + ) + states = projector.project( + horizon_s=duration, + dt_s=time_step, + initial_pose=Pose2D(), + initial_articulation_angle_rad=initial_articulation_angle_rad, + target_articulation_angle_rad=target_articulation_angle_rad, + articulation_rate_rad_s=articulation_rate_rad_s, + linear_velocity_m_s=drive_velocity, + ) + time_arr, x_arr, y_arr, theta_arr, _omega = _states_to_arrays(states) + articulation_series = np.asarray([s.articulation_angle_rad for s in states]) + return ArticulatedTrajectory( + time=time_arr, + x=x_arr, + y=y_arr, + theta=theta_arr, + linear_velocity=drive_velocity, + angular_velocity=states[0].angular_velocity_rad_s, + drive_velocity=drive_velocity, + articulation_angle=target_articulation_angle_rad, + turning_radius=states[-1].vehicle_state.front_axle_turning_radius_m, + articulation_angle_series=articulation_series, + front_footprint_series=_footprint_series(states, 'front_footprint'), + rear_footprint_series=_footprint_series(states, 'rear_footprint'), + joint_pose_series=_joint_pose_series(states), + ) + + +def single_differential_trajectory( + wheel_radius: float, + track_width: float, + initial_linear_velocity: float, + initial_angular_velocity: float, + target_linear_velocity: float, + target_angular_velocity: float, + linear_acceleration: float, + angular_acceleration: float, + duration: float, + time_step: float = 0.02, + min_linear_velocity: float | None = None, + max_linear_velocity: float | None = None, + min_angular_velocity: float | None = None, + max_angular_velocity: float | None = None, + footprint: Footprint | None = None, +) -> DifferentialTrajectory: + """Ramp diff-drive body command (v, omega) from initial toward target at the + given accelerations, integrating pose forward over `duration`. + """ + if min_linear_velocity is None: + min_linear_velocity = min(initial_linear_velocity, target_linear_velocity) + if max_linear_velocity is None: + max_linear_velocity = max(initial_linear_velocity, target_linear_velocity) + if min_angular_velocity is None: + min_angular_velocity = min(initial_angular_velocity, target_angular_velocity) + if max_angular_velocity is None: + max_angular_velocity = max(initial_angular_velocity, target_angular_velocity) + + model = DifferentialDriveModel(wheel_radius, track_width) + projector = DifferentialDriveProjector( + model, + min_linear_velocity, + max_linear_velocity, + min_angular_velocity, + max_angular_velocity, + footprint or [], + ) + states = projector.project( + horizon_s=duration, + dt_s=time_step, + initial_pose=Pose2D(), + initial_linear_velocity_m_s=initial_linear_velocity, + initial_angular_velocity_rad_s=initial_angular_velocity, + target_linear_velocity_m_s=target_linear_velocity, + target_angular_velocity_rad_s=target_angular_velocity, + linear_acceleration_m_s2=linear_acceleration, + angular_acceleration_rad_s2=angular_acceleration, + ) + time_arr, x_arr, y_arr, theta_arr, _omega = _states_to_arrays(states) + final_wheels = states[-1].wheel_velocities + return DifferentialTrajectory( + time=time_arr, + x=x_arr, + y=y_arr, + theta=theta_arr, + linear_velocity=states[-1].linear_velocity_m_s, + angular_velocity=states[-1].angular_velocity_rad_s, + left_wheel=final_wheels.left_wheel_velocity_rad_s, + right_wheel=final_wheels.right_wheel_velocity_rad_s, + base_wheel_velocity=(final_wheels.left_wheel_velocity_rad_s + final_wheels.right_wheel_velocity_rad_s) / 2.0, + footprint_series=_footprint_series(states, 'footprint'), + ) diff --git a/polymath_kinematics/explorer/types.py b/polymath_kinematics/explorer/types.py new file mode 100644 index 0000000..cb83450 --- /dev/null +++ b/polymath_kinematics/explorer/types.py @@ -0,0 +1,95 @@ +# Copyright (c) 2025-present Polymath Robotics, Inc. All rights reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Trajectory data types for kinematic exploration.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Union + +import numpy as np + + +@dataclass +class Trajectory: + """Base trajectory data with common fields. + + ``footprint_series``, when present, is the per-sample single-body footprint produced by the + C++ projector: an ``(N, V, 2)`` array of world-frame polygon vertices (N samples, V vertices, + xy). It is ``None`` when no footprint was set on the projector (single-body models only; the + articulated model uses the front/rear series on ``ArticulatedTrajectory``). + """ + + time: np.ndarray + x: np.ndarray + y: np.ndarray + theta: np.ndarray + linear_velocity: float + angular_velocity: float + footprint_series: Optional[np.ndarray] = None + + def __len__(self) -> int: + return len(self.time) + + +@dataclass +class DifferentialTrajectory(Trajectory): + """Trajectory for differential drive model.""" + + left_wheel: float = 0.0 + right_wheel: float = 0.0 + base_wheel_velocity: float = 0.0 + + +@dataclass +class BicycleTrajectory(Trajectory): + """Trajectory for bicycle model. + + ``steering_angle`` is the steady-state / target value used for legend labels and + analysis plots. For ramped trajectories, ``steering_angle_series`` holds the + per-sample steering angle (same length as ``time``); the footprint visualizer + uses it to draw front-wheel indicators at the angle reported for each sample. + """ + + drive_velocity: float = 0.0 + steering_angle: float = 0.0 + turning_radius: float = 0.0 + steering_angle_series: Optional[np.ndarray] = None + + +@dataclass +class ArticulatedTrajectory(Trajectory): + """Trajectory for articulated vehicle model. + + ``articulation_angle`` is the steady-state / target value used for legend labels + and analysis plots. For ramped trajectories, ``articulation_angle_series`` holds + the per-sample articulation angle (same length as ``time``); the footprint + visualizer uses it to fold the rear segment progressively across the trajectory. + """ + + drive_velocity: float = 0.0 + articulation_angle: float = 0.0 + turning_radius: float = 0.0 + articulation_angle_series: Optional[np.ndarray] = None + # Per-sample front/rear body footprints from the C++ projector, each an (N, V, 2) array of + # world-frame polygon vertices. None when no footprint was set on the projector. + front_footprint_series: Optional[np.ndarray] = None + rear_footprint_series: Optional[np.ndarray] = None + # Per-sample articulation-joint pose as an (N, 3) array of [x, y, rear-body theta]. The joint + # cannot be recovered from an arbitrary polygon, so the projector reports it directly. + joint_pose_series: Optional[np.ndarray] = None + + +# Type alias for any trajectory type +AnyTrajectory = Union[DifferentialTrajectory, BicycleTrajectory, ArticulatedTrajectory] diff --git a/polymath_kinematics/kinematic_explorer_app.py b/polymath_kinematics/kinematic_explorer_app.py new file mode 100644 index 0000000..67fdc47 --- /dev/null +++ b/polymath_kinematics/kinematic_explorer_app.py @@ -0,0 +1,1045 @@ +# Copyright (c) 2025-present Polymath Robotics, Inc. All rights reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Interactive Kinematic Model Explorer - Streamlit UI. + +Run with: streamlit run kinematic_explorer_app.py +""" + +from __future__ import annotations + +import io +import json +from datetime import datetime + +import matplotlib.pyplot as plt +import numpy as np +import streamlit as st + +from polymath_kinematics import AxleReference, rectangle_footprint +from polymath_kinematics.explorer import ( + KINEMATIC_EQUATIONS, + LATTICE_CONFIG, + TRAJECTORY_EQUATIONS, + AnyTrajectory, + ArticulatedTrajectory, + BicycleTrajectory, + DifferentialTrajectory, + generate_lattice_articulated, + generate_lattice_bicycle, + generate_lattice_differential, + get_traj_attr, + plot_analysis, + plot_lattice, + plot_trajectory_with_footprints, + select_symmetric_trajectories, + single_articulated_trajectory, + single_bicycle_trajectory, + single_differential_trajectory, + trajectories_to_dataframe, +) +from polymath_kinematics.explorer.config import ( + DEFAULT_ARTICULATED_FRONT_OVERHANG_M, + DEFAULT_ARTICULATED_REAR_OVERHANG_M, + DEFAULT_BICYCLE_FRONT_OVERHANG_M, + DEFAULT_BICYCLE_REAR_OVERHANG_M, + DEFAULT_DIFFERENTIAL_FRONT_OVERHANG_M, + DEFAULT_DIFFERENTIAL_REAR_OVERHANG_M, +) + +AXLE_REFERENCES = {'Rear axle': AxleReference.REAR, 'Front axle': AxleReference.FRONT} + + +def _bicycle_footprint( + wheelbase: float, track_width: float, front_overhang: float, rear_overhang: float, axle_reference: str +): + """Rectangle for a boxy bicycle body, expressed in the selected axle's frame. + + The sliders describe the body relative to the axles it sits between: `front_overhang` ahead of + the FRONT axle, `rear_overhang` behind the REAR axle. Whichever axle is the reference, the body + lands in the same physical place. + """ + if AXLE_REFERENCES[axle_reference] == AxleReference.FRONT: + return rectangle_footprint(front_overhang, wheelbase + rear_overhang, track_width) + return rectangle_footprint(wheelbase + front_overhang, rear_overhang, track_width) + + +@st.cache_data +def _cached_lattice_differential( + wheel_radius: float, + track_width: float, + base_wheel_velocities: tuple[float, ...], + wheel_velocity_diffs: tuple[float, ...], + duration: float, + time_step: float, + front_overhang: float, + rear_overhang: float, +) -> list[DifferentialTrajectory]: + return generate_lattice_differential( + wheel_radius, + track_width, + base_wheel_velocities, + wheel_velocity_diffs, + duration, + time_step, + # Pose reference is the body centre, so the overhangs are the bumper distances directly. + footprint=rectangle_footprint(front_overhang, rear_overhang, track_width), + ) + + +@st.cache_data +def _cached_lattice_bicycle( + wheelbase: float, + track_width: float, + wheel_radius: float, + drive_velocities: tuple[float, ...], + steering_angles: tuple[float, ...], + duration: float, + time_step: float, + front_overhang: float, + rear_overhang: float, + axle_reference: str, +) -> list[BicycleTrajectory]: + return generate_lattice_bicycle( + wheelbase, + track_width, + wheel_radius, + drive_velocities, + steering_angles, + duration, + time_step, + axle_reference=AXLE_REFERENCES[axle_reference], + footprint=_bicycle_footprint(wheelbase, track_width, front_overhang, rear_overhang, axle_reference), + ) + + +@st.cache_data +def _cached_lattice_articulated( + articulation_to_front: float, + articulation_to_rear: float, + front_track: float, + rear_track: float, + front_wheel_radius: float, + rear_wheel_radius: float, + drive_velocities: tuple[float, ...], + articulation_angles: tuple[float, ...], + duration: float, + time_step: float, + front_overhang: float, + rear_overhang: float, + axle_reference: str, +) -> list[ArticulatedTrajectory]: + return generate_lattice_articulated( + articulation_to_front, + articulation_to_rear, + front_track, + rear_track, + front_wheel_radius, + rear_wheel_radius, + drive_velocities, + articulation_angles, + duration, + time_step, + axle_reference=AXLE_REFERENCES[axle_reference], + front_footprint=rectangle_footprint(front_overhang, articulation_to_front, front_track), + rear_footprint=rectangle_footprint(articulation_to_rear, rear_overhang, rear_track), + ) + + +@st.cache_data +def _cached_single_bicycle( + wheelbase: float, + track_width: float, + wheel_radius: float, + initial_steering_angle_rad: float, + target_steering_angle_rad: float, + steering_rate_rad_s: float, + drive_velocity: float, + duration: float, + time_step: float, + front_overhang: float, + rear_overhang: float, + axle_reference: str, +) -> BicycleTrajectory: + return single_bicycle_trajectory( + wheelbase, + track_width, + wheel_radius, + initial_steering_angle_rad=initial_steering_angle_rad, + target_steering_angle_rad=target_steering_angle_rad, + steering_rate_rad_s=steering_rate_rad_s, + drive_velocity=drive_velocity, + duration=duration, + time_step=time_step, + axle_reference=AXLE_REFERENCES[axle_reference], + footprint=_bicycle_footprint(wheelbase, track_width, front_overhang, rear_overhang, axle_reference), + ) + + +@st.cache_data +def _cached_single_articulated( + articulation_to_front: float, + articulation_to_rear: float, + front_track: float, + rear_track: float, + front_wheel_radius: float, + rear_wheel_radius: float, + initial_articulation_angle_rad: float, + target_articulation_angle_rad: float, + articulation_rate_rad_s: float, + drive_velocity: float, + duration: float, + time_step: float, + front_overhang: float, + rear_overhang: float, + axle_reference: str, +) -> ArticulatedTrajectory: + return single_articulated_trajectory( + articulation_to_front, + articulation_to_rear, + front_track, + rear_track, + front_wheel_radius, + rear_wheel_radius, + initial_articulation_angle_rad=initial_articulation_angle_rad, + target_articulation_angle_rad=target_articulation_angle_rad, + articulation_rate_rad_s=articulation_rate_rad_s, + drive_velocity=drive_velocity, + duration=duration, + time_step=time_step, + axle_reference=AXLE_REFERENCES[axle_reference], + # Each body polygon is measured from its own axle: the joint is L_f behind the front axle + # and L_r ahead of the rear axle. + front_footprint=rectangle_footprint(front_overhang, articulation_to_front, front_track), + rear_footprint=rectangle_footprint(articulation_to_rear, rear_overhang, rear_track), + ) + + +@st.cache_data +def _cached_single_differential( + wheel_radius: float, + track_width: float, + initial_linear_velocity: float, + initial_angular_velocity: float, + target_linear_velocity: float, + target_angular_velocity: float, + linear_acceleration: float, + angular_acceleration: float, + duration: float, + time_step: float, + front_overhang: float, + rear_overhang: float, +) -> DifferentialTrajectory: + return single_differential_trajectory( + wheel_radius, + track_width, + initial_linear_velocity=initial_linear_velocity, + initial_angular_velocity=initial_angular_velocity, + target_linear_velocity=target_linear_velocity, + target_angular_velocity=target_angular_velocity, + linear_acceleration=linear_acceleration, + angular_acceleration=angular_acceleration, + duration=duration, + time_step=time_step, + # Pose reference is the body centre; see _cached_lattice_differential. + footprint=rectangle_footprint(front_overhang, rear_overhang, track_width), + ) + + +def get_config_dict(model_type: str) -> dict: + """Get current configuration as a dictionary (depends on session state).""" + config = { + 'model_type': model_type, + 'timestamp': datetime.now().isoformat(), + 'simulation': { + 'duration': st.session_state.sim_duration, + 'dt': st.session_state.sim_dt, + }, + } + + if model_type == 'Differential Drive': + config['model_parameters'] = { + 'wheel_radius': st.session_state.diff_wheel_radius, + 'track_width': st.session_state.diff_track_width, + 'front_overhang': st.session_state.diff_front_overhang, + 'rear_overhang': st.session_state.diff_rear_overhang, + } + config['control_inputs'] = { + 'base_vel_min': st.session_state.diff_base_vel_min, + 'base_vel_max': st.session_state.diff_base_vel_max, + 'n_base_vels': st.session_state.diff_n_base_vels, + 'max_wheel_diff': st.session_state.diff_max_wheel_diff, + 'n_samples': st.session_state.diff_n_samples, + } + elif model_type == 'Bicycle': + config['model_parameters'] = { + 'wheelbase': st.session_state.bike_wheelbase, + 'track_width': st.session_state.bike_track_width, + 'wheel_radius': st.session_state.bike_wheel_radius, + 'front_overhang': st.session_state.bike_front_overhang, + 'rear_overhang': st.session_state.bike_rear_overhang, + 'axle_reference': st.session_state.bike_axle_reference, + } + config['control_inputs'] = { + 'v_min': st.session_state.bike_v_min, + 'v_max': st.session_state.bike_v_max, + 'n_velocities': st.session_state.bike_n_velocities, + 'max_steer_deg': st.session_state.bike_max_steer, + 'n_steer': st.session_state.bike_n_steer, + } + elif model_type == 'Articulated': + config['model_parameters'] = { + 'articulation_to_front_axle': st.session_state.art_to_front, + 'articulation_to_rear_axle': st.session_state.art_to_rear, + 'front_track_width': st.session_state.art_front_track, + 'rear_track_width': st.session_state.art_rear_track, + 'front_wheel_radius': st.session_state.art_front_wheel_r, + 'rear_wheel_radius': st.session_state.art_rear_wheel_r, + 'front_overhang': st.session_state.art_front_overhang, + 'rear_overhang': st.session_state.art_rear_overhang, + 'axle_reference': st.session_state.art_axle_reference, + } + config['control_inputs'] = { + 'v_min': st.session_state.art_v_min, + 'v_max': st.session_state.art_v_max, + 'n_velocities': st.session_state.art_n_velocities, + 'max_articulation_deg': st.session_state.art_max_angle, + 'n_angles': st.session_state.art_n_angles, + } + + return config + + +def init_session_state(): + """Initialize default values for all model parameters.""" + defaults = { + # Differential Drive + 'diff_wheel_radius': 0.1, + 'diff_track_width': 0.5, + 'diff_front_overhang': DEFAULT_DIFFERENTIAL_FRONT_OVERHANG_M, + 'diff_rear_overhang': DEFAULT_DIFFERENTIAL_REAR_OVERHANG_M, + 'diff_base_vel_min': 5.0, + 'diff_base_vel_max': 15.0, + 'diff_n_base_vels': 3, + 'diff_max_wheel_diff': 10.0, + 'diff_n_samples': 9, + # Bicycle + 'bike_wheelbase': 2.5, + 'bike_track_width': 1.5, + 'bike_wheel_radius': 0.3, + 'bike_axle_reference': 'Rear axle', + 'bike_front_overhang': DEFAULT_BICYCLE_FRONT_OVERHANG_M, + 'bike_rear_overhang': DEFAULT_BICYCLE_REAR_OVERHANG_M, + 'bike_v_min': 1.0, + 'bike_v_max': 3.0, + 'bike_n_velocities': 3, + 'bike_max_steer': 30.0, + 'bike_n_steer': 9, + # Articulated + 'art_to_front': 1.5, + 'art_to_rear': 1.2, + 'art_front_track': 1.8, + 'art_rear_track': 1.6, + 'art_axle_reference': 'Rear axle', + 'art_front_overhang': DEFAULT_ARTICULATED_FRONT_OVERHANG_M, + 'art_rear_overhang': DEFAULT_ARTICULATED_REAR_OVERHANG_M, + 'art_front_wheel_r': 0.4, + 'art_rear_wheel_r': 0.5, + 'art_v_min': 1.0, + 'art_v_max': 3.0, + 'art_n_velocities': 3, + 'art_max_angle': 30.0, + 'art_n_angles': 9, + # Simulation + 'sim_duration': 3.0, + 'sim_dt': 0.02, + 'export_dpi': 150, + 'export_format': 'png', + } + for key, value in defaults.items(): + if key not in st.session_state: + st.session_state[key] = value + + +st.set_page_config(page_title='Kinematic Explorer', layout='wide') +st.title('Kinematic Model Explorer') + +init_session_state() + +# Sidebar - Model Selection +st.sidebar.header('Model Configuration') +model_type = st.sidebar.selectbox('Model Type', ['Differential Drive', 'Bicycle', 'Articulated']) + +st.sidebar.subheader('Model Parameters') + +# Model-specific parameters +if model_type == 'Differential Drive': + st.session_state.diff_wheel_radius = st.sidebar.slider( + 'Wheel Radius (m)', 0.05, 0.5, st.session_state.diff_wheel_radius, 0.01 + ) + st.session_state.diff_track_width = st.sidebar.slider( + 'Track Width (m)', 0.2, 2.0, st.session_state.diff_track_width, 0.05 + ) + st.session_state.diff_front_overhang = st.sidebar.slider( + 'Front Overhang (m)', + 0.0, + 2.0, + st.session_state.diff_front_overhang, + 0.05, + help='Body length ahead of the body centre (the pose reference).', + ) + st.session_state.diff_rear_overhang = st.sidebar.slider( + 'Rear Overhang (m)', + 0.0, + 2.0, + st.session_state.diff_rear_overhang, + 0.05, + help='Body length behind the body centre (the pose reference).', + ) + st.sidebar.markdown('---') + st.sidebar.markdown(f'**Wheel Radius:** {st.session_state.diff_wheel_radius:.2f} m') + st.sidebar.markdown(f'**Track Width:** {st.session_state.diff_track_width:.2f} m') + st.sidebar.markdown( + f'**Body Length:** {st.session_state.diff_front_overhang + st.session_state.diff_rear_overhang:.2f} m' + ) + +elif model_type == 'Bicycle': + st.session_state.bike_wheelbase = st.sidebar.slider('Wheelbase (m)', 1.0, 5.0, st.session_state.bike_wheelbase, 0.1) + st.session_state.bike_track_width = st.sidebar.slider( + 'Track Width (m)', 1.0, 3.0, st.session_state.bike_track_width, 0.1 + ) + st.session_state.bike_wheel_radius = st.sidebar.slider( + 'Wheel Radius (m)', 0.1, 0.6, st.session_state.bike_wheel_radius, 0.05 + ) + st.session_state.bike_axle_reference = st.sidebar.selectbox( + 'Pose / Footprint Reference', + list(AXLE_REFERENCES), + index=list(AXLE_REFERENCES).index(st.session_state.bike_axle_reference), + help='Axle that trajectory poses and the footprint polygon are measured from.', + ) + st.session_state.bike_front_overhang = st.sidebar.slider( + 'Front Overhang (m)', + 0.0, + 2.0, + st.session_state.bike_front_overhang, + 0.05, + help='Body length ahead of the FRONT axle. The pose reference is the rear axle, so the ' + 'front bumper sits a wheelbase plus this overhang ahead of it.', + ) + st.session_state.bike_rear_overhang = st.sidebar.slider( + 'Rear Overhang (m)', + 0.0, + 2.0, + st.session_state.bike_rear_overhang, + 0.05, + help='Body length behind the rear axle (the pose reference).', + ) + st.sidebar.markdown('---') + st.sidebar.markdown(f'**Wheelbase:** {st.session_state.bike_wheelbase:.2f} m') + st.sidebar.markdown(f'**Track Width:** {st.session_state.bike_track_width:.2f} m') + st.sidebar.markdown( + f'**Body Length:** ' + f'{st.session_state.bike_wheelbase + st.session_state.bike_front_overhang + st.session_state.bike_rear_overhang:.2f} m' + ) + +elif model_type == 'Articulated': + st.session_state.art_to_front = st.sidebar.slider( + 'Articulation to Front Axle (m)', 0.5, 3.0, st.session_state.art_to_front, 0.1 + ) + st.session_state.art_to_rear = st.sidebar.slider( + 'Articulation to Rear Axle (m)', 0.5, 3.0, st.session_state.art_to_rear, 0.1 + ) + st.session_state.art_front_track = st.sidebar.slider( + 'Front Track Width (m)', 1.0, 3.0, st.session_state.art_front_track, 0.1 + ) + st.session_state.art_rear_track = st.sidebar.slider( + 'Rear Track Width (m)', 1.0, 3.0, st.session_state.art_rear_track, 0.1 + ) + st.session_state.art_axle_reference = st.sidebar.selectbox( + 'Pose Reference', + list(AXLE_REFERENCES), + index=list(AXLE_REFERENCES).index(st.session_state.art_axle_reference), + help='Axle that trajectory poses are measured from. Each body footprint is always measured from its own axle.', + ) + st.session_state.art_front_overhang = st.sidebar.slider( + 'Front Overhang (m)', + 0.0, + 3.0, + st.session_state.art_front_overhang, + 0.1, + help='Body length ahead of the front axle (e.g. bucket).', + ) + st.session_state.art_rear_overhang = st.sidebar.slider( + 'Rear Overhang (m)', + 0.0, + 3.0, + st.session_state.art_rear_overhang, + 0.1, + help='Body length behind the rear axle (e.g. counterweight).', + ) + st.session_state.art_front_wheel_r = st.sidebar.slider( + 'Front Wheel Radius (m)', 0.2, 0.8, st.session_state.art_front_wheel_r, 0.05 + ) + st.session_state.art_rear_wheel_r = st.sidebar.slider( + 'Rear Wheel Radius (m)', 0.2, 0.8, st.session_state.art_rear_wheel_r, 0.05 + ) + st.sidebar.markdown('---') + st.sidebar.markdown(f'**Total Wheelbase:** {st.session_state.art_to_front + st.session_state.art_to_rear:.2f} m') + +# Lattice Configuration +st.sidebar.header('Lattice Configuration') +st.session_state.sim_duration = st.sidebar.slider( + 'Simulation Duration (s)', 1.0, 10.0, st.session_state.sim_duration, 0.5 +) + +with st.sidebar.expander('Advanced Settings'): + st.caption('Trajectories come from the C++ projectors (Euler integration at the time step below).') + st.session_state.sim_dt = st.slider( + 'Time Step (s)', + 0.005, + 0.1, + st.session_state.sim_dt, + 0.005, + help='Smaller values are more accurate but slower', + ) + st.caption(f'Steps per trajectory: {int(st.session_state.sim_duration / st.session_state.sim_dt)}') + +# Generate trajectories based on model type +trajectories: list[AnyTrajectory] +group_values: list[float] + +if model_type == 'Differential Drive': + st.sidebar.subheader('Control Inputs') + col1, col2 = st.sidebar.columns(2) + with col1: + st.session_state.diff_base_vel_min = st.number_input( + 'Min Base Vel (rad/s)', 1.0, 20.0, st.session_state.diff_base_vel_min, 1.0 + ) + with col2: + st.session_state.diff_base_vel_max = st.number_input( + 'Max Base Vel (rad/s)', 1.0, 30.0, st.session_state.diff_base_vel_max, 1.0 + ) + st.session_state.diff_n_base_vels = st.sidebar.slider( + 'Number of Base Velocities', 1, 5, st.session_state.diff_n_base_vels + ) + base_wheel_velocities = tuple( + np.linspace( + st.session_state.diff_base_vel_min, st.session_state.diff_base_vel_max, st.session_state.diff_n_base_vels + ) + ) + + st.session_state.diff_max_wheel_diff = st.sidebar.slider( + 'Max Wheel Velocity Diff (rad/s)', 1.0, 20.0, st.session_state.diff_max_wheel_diff, 1.0 + ) + st.session_state.diff_n_samples = st.sidebar.slider('Diff Samples', 3, 15, st.session_state.diff_n_samples, 2) + wheel_diffs = tuple( + np.linspace( + -st.session_state.diff_max_wheel_diff, st.session_state.diff_max_wheel_diff, st.session_state.diff_n_samples + ) + ) + + trajectories = _cached_lattice_differential( + wheel_radius=st.session_state.diff_wheel_radius, + track_width=st.session_state.diff_track_width, + base_wheel_velocities=base_wheel_velocities, + wheel_velocity_diffs=wheel_diffs, + duration=st.session_state.sim_duration, + time_step=st.session_state.sim_dt, + front_overhang=st.session_state.diff_front_overhang, + rear_overhang=st.session_state.diff_rear_overhang, + ) + group_values = list(base_wheel_velocities) + +elif model_type == 'Bicycle': + st.sidebar.subheader('Control Inputs') + col1, col2 = st.sidebar.columns(2) + with col1: + st.session_state.bike_v_min = st.number_input('Min Drive Vel (m/s)', 0.5, 5.0, st.session_state.bike_v_min, 0.5) + with col2: + st.session_state.bike_v_max = st.number_input( + 'Max Drive Vel (m/s)', 0.5, 10.0, st.session_state.bike_v_max, 0.5 + ) + st.session_state.bike_n_velocities = st.sidebar.slider( + 'Number of Velocities', 1, 5, st.session_state.bike_n_velocities + ) + drive_velocities = tuple( + np.linspace(st.session_state.bike_v_min, st.session_state.bike_v_max, st.session_state.bike_n_velocities) + ) + + st.session_state.bike_max_steer = st.sidebar.slider( + 'Max Steering Angle (deg)', 5.0, 45.0, st.session_state.bike_max_steer, 5.0 + ) + st.session_state.bike_n_steer = st.sidebar.slider('Steering Samples', 3, 15, st.session_state.bike_n_steer, 2) + steering_angles = tuple( + np.linspace( + -np.radians(st.session_state.bike_max_steer), + np.radians(st.session_state.bike_max_steer), + st.session_state.bike_n_steer, + ) + ) + + trajectories = _cached_lattice_bicycle( + wheelbase=st.session_state.bike_wheelbase, + track_width=st.session_state.bike_track_width, + wheel_radius=st.session_state.bike_wheel_radius, + drive_velocities=drive_velocities, + steering_angles=steering_angles, + duration=st.session_state.sim_duration, + time_step=st.session_state.sim_dt, + front_overhang=st.session_state.bike_front_overhang, + rear_overhang=st.session_state.bike_rear_overhang, + axle_reference=st.session_state.bike_axle_reference, + ) + group_values = list(drive_velocities) + +else: # Articulated + st.sidebar.subheader('Control Inputs') + col1, col2 = st.sidebar.columns(2) + with col1: + st.session_state.art_v_min = st.number_input('Min Drive Vel (m/s)', 0.5, 5.0, st.session_state.art_v_min, 0.5) + with col2: + st.session_state.art_v_max = st.number_input('Max Drive Vel (m/s)', 0.5, 10.0, st.session_state.art_v_max, 0.5) + st.session_state.art_n_velocities = st.sidebar.slider( + 'Number of Velocities', 1, 5, st.session_state.art_n_velocities + ) + drive_velocities = tuple( + np.linspace(st.session_state.art_v_min, st.session_state.art_v_max, st.session_state.art_n_velocities) + ) + + st.session_state.art_max_angle = st.sidebar.slider( + 'Max Articulation Angle (deg)', 5.0, 45.0, st.session_state.art_max_angle, 5.0 + ) + st.session_state.art_n_angles = st.sidebar.slider('Articulation Samples', 3, 15, st.session_state.art_n_angles, 2) + articulation_angles = tuple( + np.linspace( + -np.radians(st.session_state.art_max_angle), + np.radians(st.session_state.art_max_angle), + st.session_state.art_n_angles, + ) + ) + + trajectories = _cached_lattice_articulated( + articulation_to_front=st.session_state.art_to_front, + articulation_to_rear=st.session_state.art_to_rear, + front_track=st.session_state.art_front_track, + rear_track=st.session_state.art_rear_track, + front_wheel_radius=st.session_state.art_front_wheel_r, + rear_wheel_radius=st.session_state.art_rear_wheel_r, + drive_velocities=drive_velocities, + articulation_angles=articulation_angles, + duration=st.session_state.sim_duration, + time_step=st.session_state.sim_dt, + front_overhang=st.session_state.art_front_overhang, + rear_overhang=st.session_state.art_rear_overhang, + axle_reference=st.session_state.art_axle_reference, + ) + group_values = list(drive_velocities) + +# Kinematic Equations Section +with st.expander('Kinematic Equations', expanded=False): + eq_info = KINEMATIC_EQUATIONS[model_type] + st.subheader(eq_info['title']) + for eq in eq_info['equations']: + st.latex(eq) + st.markdown(eq_info['variables']) + st.markdown('---') + st.markdown(TRAJECTORY_EQUATIONS) + +# Trajectory Lattice — only the longest (highest-velocity) group is shown. Lower velocities are +# the same trajectory fan at a shorter distance, so rendering all of them is redundant. +longest_group = [max(group_values)] if group_values else group_values +st.header('Trajectory Lattice') +lattice_fig = plot_lattice(trajectories, model_type, longest_group) +lattice_col, _ = st.columns([2, 1]) # constrain to 2/3 page width so the 16:9 plot isn't full-bleed +lattice_col.pyplot(lattice_fig) +plt.close(lattice_fig) + +# Kinematic Analysis +st.header('Kinematic Analysis') +analysis_fig = plot_analysis(trajectories, model_type, group_values) +analysis_col, _ = st.columns([2, 1]) # constrain to 2/3 page width so the plot isn't full-bleed +analysis_col.pyplot(analysis_fig) +plt.close(analysis_fig) + +# Vehicle Footprint Visualization +st.header('Vehicle Visualization') + +vel_options = list(group_values) +# Body outlines come from the projector footprints; model_params only carries what the corners +# can't give us — the bicycle's front-axle offset and track width. +model_params: dict = {} +if model_type == 'Differential Drive': + vel_label = 'Base Wheel Velocity (rad/s)' + body_length = st.session_state.diff_front_overhang + st.session_state.diff_rear_overhang + st.caption(f'Vehicle dimensions: {body_length:.2f}m x {st.session_state.diff_track_width:.2f}m') +elif model_type == 'Bicycle': + # front_axle_offset is measured from the pose reference, so it collapses to 0 when the pose + # already sits at the front axle. + model_params = { + 'wheelbase': 0.0 + if AXLE_REFERENCES[st.session_state.bike_axle_reference] == AxleReference.FRONT + else st.session_state.bike_wheelbase, + 'track_width': st.session_state.bike_track_width, + } + vel_label = 'Drive Velocity (m/s)' + body_length = ( + st.session_state.bike_wheelbase + st.session_state.bike_front_overhang + st.session_state.bike_rear_overhang + ) + st.caption( + f'Vehicle dimensions: {body_length:.2f}m x {st.session_state.bike_track_width:.2f}m ' + f'(wheelbase={st.session_state.bike_wheelbase:.2f}m)' + ) +else: # Articulated + vel_label = 'Drive Velocity (m/s)' + # Body lengths are measured from the articulation joint (base_link) to each bumper: + # joint-to-axle + overhang. + front_length = st.session_state.art_to_front + st.session_state.art_front_overhang + rear_length = st.session_state.art_to_rear + st.session_state.art_rear_overhang + st.caption( + f'Articulated: front={front_length:.2f}m x {st.session_state.art_front_track:.2f}m, ' + f'rear={rear_length:.2f}m x {st.session_state.art_rear_track:.2f}m' + ) + +col1, col2 = st.columns([3, 1]) +with col2: + st.markdown('**Visualization Settings**') + num_footprints = st.slider('Footprints per trajectory', 3, 8, 4, 1, key='num_footprints') + + if model_type == 'Differential Drive': + angle_label = 'Angular velocities' + elif model_type == 'Bicycle': + angle_label = 'Steering angles' + else: + angle_label = 'Articulation angles' + + num_angles_viz = st.slider( + angle_label, 1, 9, 5, 2, key='num_angles_viz', help='Number of angles to show (symmetric)' + ) + + if len(vel_options) > 1: + viz_vel = st.select_slider( + vel_label, + options=[round(v, 1) for v in vel_options], + value=round(vel_options[len(vel_options) // 2], 1), + key='viz_velocity', + ) + else: + viz_vel = vel_options[0] + st.caption(f'{vel_label}: {viz_vel:.1f}') + +with col1: + config = LATTICE_CONFIG[model_type] + # Narrow to the chosen velocity first, then let the library pick a symmetric angle subset. + velocity_filtered = [ + trajectory for trajectory in trajectories if abs(get_traj_attr(trajectory, config.vel_key) - viz_vel) < 0.05 + ] + viz_trajectories = select_symmetric_trajectories( + velocity_filtered, model_type, num_angles=num_angles_viz, num_velocities=1 + ) + + if viz_trajectories: + footprint_fig = plot_trajectory_with_footprints(viz_trajectories, model_type, model_params, num_footprints) + st.pyplot(footprint_fig) + plt.close(footprint_fig) + else: + st.warning('No trajectories selected. Adjust the visualization settings.') + +# Single Projected Trajectory — ramps the model command from initial → target at a chosen rate, +# clamped to [min, max], and plots the resulting trajectory live. Velocity/horizon are inherited. +st.header('Single Projected Trajectory') +st.caption( + 'Set the initial input, target input, and rate of change. The plot updates as you move the ' + 'sliders. Drive velocity and simulation horizon are inherited from the sections above.' +) + + +def _render_single_trajectory(single_traj): + """Render one projected trajectory with footprints (shared by all model branches below).""" + fig = plot_trajectory_with_footprints([single_traj], model_type, model_params, num_footprints=5) + plot_col, _ = st.columns([2, 1]) # constrain to 2/3 page width so the 16:9 plot isn't full-bleed + plot_col.pyplot(fig) + plt.close(fig) + + +if model_type == 'Bicycle': + col_a, col_b, col_c = st.columns(3) + with col_a: + single_initial_deg = st.slider( + 'Initial Steering (deg)', + -st.session_state.bike_max_steer, + st.session_state.bike_max_steer, + 0.0, + 1.0, + key='bike_single_initial_deg', + ) + with col_b: + single_target_deg = st.slider( + 'Target Steering (deg)', + -st.session_state.bike_max_steer, + st.session_state.bike_max_steer, + float(st.session_state.bike_max_steer / 2), + 1.0, + key='bike_single_target_deg', + ) + with col_c: + single_rate_deg_s = st.slider( + 'Steering Rate (deg/s)', + 0.0, + 180.0, + 30.0, + 1.0, + key='bike_single_rate_deg_s', + ) + + single_traj = _cached_single_bicycle( + wheelbase=st.session_state.bike_wheelbase, + track_width=st.session_state.bike_track_width, + wheel_radius=st.session_state.bike_wheel_radius, + initial_steering_angle_rad=float(np.radians(single_initial_deg)), + target_steering_angle_rad=float(np.radians(single_target_deg)), + steering_rate_rad_s=float(np.radians(single_rate_deg_s)), + drive_velocity=float(viz_vel), + duration=st.session_state.sim_duration, + time_step=st.session_state.sim_dt, + front_overhang=st.session_state.bike_front_overhang, + rear_overhang=st.session_state.bike_rear_overhang, + axle_reference=st.session_state.bike_axle_reference, + ) + _render_single_trajectory(single_traj) + st.caption( + f'Drive velocity: {viz_vel:.2f} m/s · ' + f'Horizon: {st.session_state.sim_duration:.1f}s · ' + f'dt: {st.session_state.sim_dt:.3f}s' + ) + +elif model_type == 'Articulated': + col_a, col_b, col_c = st.columns(3) + with col_a: + single_initial_deg = st.slider( + 'Initial Articulation (deg)', + -st.session_state.art_max_angle, + st.session_state.art_max_angle, + 0.0, + 1.0, + key='art_single_initial_deg', + ) + with col_b: + single_target_deg = st.slider( + 'Target Articulation (deg)', + -st.session_state.art_max_angle, + st.session_state.art_max_angle, + float(st.session_state.art_max_angle / 2), + 1.0, + key='art_single_target_deg', + ) + with col_c: + single_rate_deg_s = st.slider( + 'Articulation Rate (deg/s)', + 0.0, + 90.0, + 15.0, + 1.0, + key='art_single_rate_deg_s', + ) + + single_traj = _cached_single_articulated( + articulation_to_front=st.session_state.art_to_front, + articulation_to_rear=st.session_state.art_to_rear, + front_track=st.session_state.art_front_track, + rear_track=st.session_state.art_rear_track, + front_wheel_radius=st.session_state.art_front_wheel_r, + rear_wheel_radius=st.session_state.art_rear_wheel_r, + initial_articulation_angle_rad=float(np.radians(single_initial_deg)), + target_articulation_angle_rad=float(np.radians(single_target_deg)), + articulation_rate_rad_s=float(np.radians(single_rate_deg_s)), + drive_velocity=float(viz_vel), + duration=st.session_state.sim_duration, + time_step=st.session_state.sim_dt, + front_overhang=st.session_state.art_front_overhang, + rear_overhang=st.session_state.art_rear_overhang, + axle_reference=st.session_state.art_axle_reference, + ) + _render_single_trajectory(single_traj) + st.caption( + f'Drive velocity: {viz_vel:.2f} m/s · ' + f'Horizon: {st.session_state.sim_duration:.1f}s · ' + f'dt: {st.session_state.sim_dt:.3f}s' + ) + +else: # Differential Drive — ramp linear AND angular body command. + st.markdown('**Body command**') + col_a, col_b, col_c, col_d = st.columns(4) + with col_a: + single_initial_v = st.slider( + 'Initial Linear (m/s)', + -3.0, + 3.0, + 0.0, + 0.1, + key='diff_single_initial_v', + ) + with col_b: + single_target_v = st.slider( + 'Target Linear (m/s)', + -3.0, + 3.0, + 1.0, + 0.1, + key='diff_single_target_v', + ) + with col_c: + single_initial_omega = st.slider( + 'Initial Angular (rad/s)', + -3.0, + 3.0, + 0.0, + 0.1, + key='diff_single_initial_omega', + ) + with col_d: + single_target_omega = st.slider( + 'Target Angular (rad/s)', + -3.0, + 3.0, + 0.5, + 0.1, + key='diff_single_target_omega', + ) + st.markdown('**Acceleration limits**') + col_e, col_f = st.columns(2) + with col_e: + single_linear_accel = st.slider( + 'Linear Accel (m/s²)', + 0.0, + 5.0, + 1.0, + 0.1, + key='diff_single_linear_accel', + ) + with col_f: + single_angular_accel = st.slider( + 'Angular Accel (rad/s²)', + 0.0, + 5.0, + 1.0, + 0.1, + key='diff_single_angular_accel', + ) + + single_traj = _cached_single_differential( + wheel_radius=st.session_state.diff_wheel_radius, + track_width=st.session_state.diff_track_width, + initial_linear_velocity=float(single_initial_v), + initial_angular_velocity=float(single_initial_omega), + target_linear_velocity=float(single_target_v), + target_angular_velocity=float(single_target_omega), + linear_acceleration=float(single_linear_accel), + angular_acceleration=float(single_angular_accel), + duration=st.session_state.sim_duration, + time_step=st.session_state.sim_dt, + front_overhang=st.session_state.diff_front_overhang, + rear_overhang=st.session_state.diff_rear_overhang, + ) + _render_single_trajectory(single_traj) + st.caption(f'Horizon: {st.session_state.sim_duration:.1f}s · dt: {st.session_state.sim_dt:.3f}s') + +# Parameter Table +with st.expander('Current Configuration', expanded=False): + config_dict = get_config_dict(model_type) + + col1, col2 = st.columns(2) + with col1: + st.markdown('**Model Parameters:**') + for key, value in config_dict['model_parameters'].items(): + st.markdown(f'- {key.replace("_", " ").title()}: `{value}`') + + with col2: + st.markdown('**Control Inputs:**') + for key, value in config_dict['control_inputs'].items(): + st.markdown(f'- {key.replace("_", " ").title()}: `{value}`') + + st.markdown('**Simulation Settings:**') + st.markdown(f'- Duration: `{config_dict["simulation"]["duration"]}s`, dt: `{config_dict["simulation"]["dt"]}s`') + +# Summary stats +st.header('Summary') +col1, col2, col3, col4 = st.columns(4) + +with col1: + st.metric('Model Type', model_type) +with col2: + st.metric('Total Trajectories', len(trajectories)) +with col3: + if model_type == 'Articulated': + max_art = max(abs(trajectory.articulation_angle) for trajectory in trajectories) # type: ignore[union-attr] + st.metric('Max Articulation', f'{np.degrees(max_art):.1f} deg') + elif model_type == 'Bicycle': + max_steer = max(abs(trajectory.steering_angle) for trajectory in trajectories) # type: ignore[union-attr] + st.metric('Max Steering', f'{np.degrees(max_steer):.1f} deg') + elif model_type == 'Differential Drive': + st.metric('Track Width', f'{st.session_state.diff_track_width:.2f} m') +with col4: + max_omega = max(abs(trajectory.angular_velocity) for trajectory in trajectories) + st.metric('Max Angular Vel', f'{max_omega:.2f} rad/s') + +# Export Section +st.sidebar.header('Export') + +with st.sidebar.expander('Export Settings'): + st.session_state.export_format = st.selectbox( + 'Image Format', + ['png', 'svg', 'pdf'], + index=['png', 'svg', 'pdf'].index(st.session_state.export_format), + ) + st.session_state.export_dpi = st.select_slider( + 'DPI (for PNG)', + [72, 150, 300, 600], + value=st.session_state.export_dpi, + ) + +# CSV Export +df = trajectories_to_dataframe(trajectories, model_type) +csv_buffer = io.StringIO() +df.to_csv(csv_buffer, index=False) +st.sidebar.download_button( + label='Download Trajectory CSV', + data=csv_buffer.getvalue(), + file_name=f'trajectories_{model_type.lower().replace(" ", "_")}.csv', + mime='text/csv', +) + +# JSON Config Export +config_export = get_config_dict(model_type) +st.sidebar.download_button( + label='Download Config JSON', + data=json.dumps(config_export, indent=2), + file_name=f'config_{model_type.lower().replace(" ", "_")}.json', + mime='application/json', +) + +# Image Export — matches the displayed lattice (longest/highest-velocity group only), rendered +# into memory so the browser downloads it rather than the server writing a file. +image_format = st.session_state.export_format +image_buffer = io.BytesIO() +export_fig = plot_lattice(trajectories, model_type, longest_group) +save_kwargs = {'bbox_inches': 'tight', 'format': image_format} +if image_format == 'png': + save_kwargs['dpi'] = st.session_state.export_dpi +export_fig.savefig(image_buffer, **save_kwargs) +plt.close(export_fig) +st.sidebar.download_button( + label=f'Download Lattice Image ({image_format.upper()})', + data=image_buffer.getvalue(), + file_name=f'lattice_{model_type.lower().replace(" ", "_")}.{image_format}', + mime={'png': 'image/png', 'svg': 'image/svg+xml', 'pdf': 'application/pdf'}[image_format], +) diff --git a/test/test_explorer.py b/test/test_explorer.py new file mode 100644 index 0000000..242d57c --- /dev/null +++ b/test/test_explorer.py @@ -0,0 +1,657 @@ +# Copyright (c) 2025-present Polymath Robotics, Inc. All rights reserved +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the kinematic explorer module.""" + +import math + +import numpy as np +import pytest + +from polymath_kinematics import AxleReference, DifferentialDriveModel, rectangle_footprint +from polymath_kinematics.explorer import ( + KINEMATIC_EQUATIONS, + LATTICE_CONFIG, + generate_lattice_articulated, + generate_lattice_bicycle, + generate_lattice_differential, + plot_analysis, + plot_lattice, + plot_trajectory_with_footprints, + select_symmetric_trajectories, + single_articulated_trajectory, + single_bicycle_trajectory, + single_differential_trajectory, + trajectories_to_dataframe, +) + + +class TestLatticeGeneration: + def test_generate_lattice_differential(self): + trajectories = generate_lattice_differential( + wheel_radius=0.1, + track_width=0.5, + base_wheel_velocities=(10.0,), + wheel_velocity_diffs=(-5.0, 0.0, 5.0), + duration=1.0, + time_step=0.1, + ) + + # ceil(duration/dt) + 1 samples, same as the bicycle and articulated lattices. + assert len(trajectories) == 3 + assert len(trajectories[0].time) == 11 + assert trajectories[0].linear_velocity == pytest.approx(1.0) + + def test_generate_lattice_differential_footprints(self): + # With body dimensions supplied, the projector emits a footprint per sample. + trajectories = generate_lattice_differential( + wheel_radius=0.1, + track_width=0.5, + base_wheel_velocities=(10.0,), + wheel_velocity_diffs=(0.0, 5.0), + duration=1.0, + time_step=0.1, + footprint=rectangle_footprint(0.4, 0.4, 0.5), + ) + for trajectory in trajectories: + assert trajectory.footprint_series is not None + assert trajectory.footprint_series.shape == (len(trajectory.time), 4, 2) + + def test_generate_lattice_differential_no_dims_has_no_footprints(self): + trajectories = generate_lattice_differential( + wheel_radius=0.1, + track_width=0.5, + base_wheel_velocities=(10.0,), + wheel_velocity_diffs=(0.0,), + duration=1.0, + time_step=0.1, + ) + assert trajectories[0].footprint_series is None + + def test_generate_lattice_differential_outer_cells_are_not_clamped(self): + # Bounds must span the sweep, or the extreme cells get clamped back toward the middle. + diffs = (-10.0, 0.0, 10.0) + trajectories = generate_lattice_differential( + wheel_radius=0.1, + track_width=0.5, + base_wheel_velocities=(10.0,), + wheel_velocity_diffs=diffs, + duration=1.0, + time_step=0.1, + ) + model = DifferentialDriveModel(0.1, 0.5) + for trajectory, velocity_diff in zip(trajectories, diffs): + expected = model.wheel_velocities_to_body_velocity(10.0 - velocity_diff / 2, 10.0 + velocity_diff / 2) + assert trajectory.angular_velocity == pytest.approx(expected.angular_velocity_rad_s) + assert trajectory.linear_velocity == pytest.approx(expected.linear_velocity_m_s) + + def test_generate_lattice_differential_multiple_velocities(self): + trajectories = generate_lattice_differential( + wheel_radius=0.1, + track_width=0.5, + base_wheel_velocities=(10.0, 15.0), + wheel_velocity_diffs=(-5.0, 0.0, 5.0), + duration=1.0, + time_step=0.1, + ) + assert len(trajectories) == 6 + + def test_generate_lattice_bicycle(self): + trajectories = generate_lattice_bicycle( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + drive_velocities=(2.0,), + steering_angles=(0.0, math.radians(15)), + duration=1.0, + time_step=0.1, + ) + # Projector returns ceil(duration/dt) + 1 samples (initial state seeded as element 0). + assert len(trajectories) == 2 + assert len(trajectories[0].time) == 11 + assert trajectories[0].drive_velocity == pytest.approx(2.0) + + def test_generate_lattice_articulated(self): + trajectories = generate_lattice_articulated( + articulation_to_front=1.5, + articulation_to_rear=1.2, + front_track=1.8, + rear_track=1.6, + front_wheel_radius=0.4, + rear_wheel_radius=0.5, + drive_velocities=(2.0,), + articulation_angles=(0.0, math.radians(15)), + duration=1.0, + time_step=0.1, + ) + assert len(trajectories) == 2 + assert len(trajectories[0].time) == 11 + assert trajectories[0].drive_velocity == pytest.approx(2.0) + + +class TestDataframeExport: + def test_differential_drive_export(self): + trajectories = generate_lattice_differential( + wheel_radius=0.1, + track_width=0.5, + base_wheel_velocities=(10.0,), + wheel_velocity_diffs=(-5.0, 0.0, 5.0), + duration=1.0, + time_step=0.1, + ) + dataframe = trajectories_to_dataframe(trajectories, 'Differential Drive') + + expected_columns = [ + 'trajectory_id', + 'time', + 'x', + 'y', + 'theta', + 'linear_velocity', + 'angular_velocity', + 'left_wheel_velocity', + 'right_wheel_velocity', + 'base_wheel_velocity', + ] + assert list(dataframe.columns) == expected_columns + assert len(dataframe) == 33 # 3 trajectories * 11 samples each + + def test_bicycle_export(self): + trajectories = generate_lattice_bicycle( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + drive_velocities=(2.0,), + steering_angles=(0.0,), + duration=1.0, + time_step=0.1, + ) + dataframe = trajectories_to_dataframe(trajectories, 'Bicycle') + + assert 'drive_velocity' in dataframe.columns + assert 'steering_angle' in dataframe.columns + assert 'turning_radius' in dataframe.columns + + def test_articulated_export(self): + trajectories = generate_lattice_articulated( + articulation_to_front=1.5, + articulation_to_rear=1.2, + front_track=1.8, + rear_track=1.6, + front_wheel_radius=0.4, + rear_wheel_radius=0.5, + drive_velocities=(2.0,), + articulation_angles=(0.0,), + duration=1.0, + time_step=0.1, + ) + dataframe = trajectories_to_dataframe(trajectories, 'Articulated') + + assert 'drive_velocity' in dataframe.columns + assert 'articulation_angle' in dataframe.columns + assert 'turning_radius' in dataframe.columns + + def test_empty_trajectories(self): + dataframe = trajectories_to_dataframe([], 'Bicycle') + assert dataframe.empty + + +class TestTrajectorySelection: + def test_select_symmetric_trajectories(self): + trajectories = generate_lattice_differential( + wheel_radius=0.1, + track_width=0.5, + base_wheel_velocities=(10.0, 15.0), + wheel_velocity_diffs=(-5.0, -2.5, 0.0, 2.5, 5.0), + duration=1.0, + time_step=0.1, + ) + selected = select_symmetric_trajectories(trajectories, 'Differential Drive', num_angles=3, num_velocities=1) + assert len(selected) == 3 + + def test_select_all_when_fewer_than_requested(self): + trajectories = generate_lattice_bicycle( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + drive_velocities=(2.0,), + steering_angles=(0.0, math.radians(15)), + duration=1.0, + time_step=0.1, + ) + selected = select_symmetric_trajectories(trajectories, 'Bicycle', num_angles=5, num_velocities=1) + assert len(selected) == 2 + + def test_empty_trajectories_selection(self): + selected = select_symmetric_trajectories([], 'Bicycle', num_angles=5, num_velocities=1) + assert selected == [] + + +class TestPlotting: + @pytest.fixture + def differential_trajectories(self): + # Body dims are supplied so the footprint-overlay path is actually exercised. + return generate_lattice_differential( + wheel_radius=0.1, + track_width=0.5, + base_wheel_velocities=(10.0, 15.0), + wheel_velocity_diffs=(-5.0, 0.0, 5.0), + duration=1.0, + time_step=0.1, + footprint=rectangle_footprint(0.4, 0.4, 0.5), + ) + + @pytest.fixture + def bicycle_trajectories(self): + return generate_lattice_bicycle( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + drive_velocities=(2.0, 3.0), + steering_angles=(-math.radians(15), 0.0, math.radians(15)), + duration=1.0, + time_step=0.1, + footprint=rectangle_footprint(2.5 + 0.6, 0.5, 1.5), + ) + + @pytest.fixture + def articulated_trajectories(self): + return generate_lattice_articulated( + articulation_to_front=1.5, + articulation_to_rear=1.2, + front_track=1.8, + rear_track=1.6, + front_wheel_radius=0.4, + rear_wheel_radius=0.5, + drive_velocities=(2.0,), + articulation_angles=(-math.radians(15), 0.0, math.radians(15)), + duration=1.0, + time_step=0.1, + # Each polygon is measured from its own axle; the joint is L_f behind the front + # axle and L_r ahead of the rear axle. + front_footprint=rectangle_footprint(1.0, 1.5, 1.8), + rear_footprint=rectangle_footprint(1.2, 0.8, 1.6), + ) + + def test_plot_lattice_differential(self, differential_trajectories): + assert plot_lattice(differential_trajectories, 'Differential Drive', [10.0, 15.0]) is not None + + def test_plot_lattice_bicycle(self, bicycle_trajectories): + assert plot_lattice(bicycle_trajectories, 'Bicycle', [2.0, 3.0]) is not None + + def test_plot_lattice_articulated(self, articulated_trajectories): + assert plot_lattice(articulated_trajectories, 'Articulated', [2.0]) is not None + + def test_plot_analysis_differential(self, differential_trajectories): + assert plot_analysis(differential_trajectories, 'Differential Drive', [10.0, 15.0]) is not None + + def test_plot_analysis_bicycle(self, bicycle_trajectories): + assert plot_analysis(bicycle_trajectories, 'Bicycle', [2.0, 3.0]) is not None + + def test_plot_analysis_articulated(self, articulated_trajectories): + assert plot_analysis(articulated_trajectories, 'Articulated', [2.0]) is not None + + def test_plot_trajectory_with_footprints_differential(self, differential_trajectories): + selected = select_symmetric_trajectories( + differential_trajectories, 'Differential Drive', num_angles=3, num_velocities=1 + ) + figure = plot_trajectory_with_footprints(selected, 'Differential Drive', {}, num_footprints=3) + assert figure is not None + + def test_plot_trajectory_with_footprints_bicycle(self, bicycle_trajectories): + selected = select_symmetric_trajectories(bicycle_trajectories, 'Bicycle', num_angles=3, num_velocities=1) + figure = plot_trajectory_with_footprints( + selected, 'Bicycle', {'wheelbase': 2.5, 'track_width': 1.5}, num_footprints=3 + ) + assert figure is not None + + def test_plot_trajectory_with_footprints_articulated(self, articulated_trajectories): + selected = select_symmetric_trajectories( + articulated_trajectories, 'Articulated', num_angles=3, num_velocities=1 + ) + figure = plot_trajectory_with_footprints(selected, 'Articulated', {}, num_footprints=3) + assert figure is not None + + def test_plot_trajectory_with_footprints_skips_missing_footprints(self): + # No body dims: the path is still drawn, the overlay skipped rather than raising. + without_dims = generate_lattice_bicycle( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + drive_velocities=(2.0,), + steering_angles=(math.radians(15),), + duration=1.0, + time_step=0.1, + ) + assert without_dims[0].footprint_series is None + figure = plot_trajectory_with_footprints(without_dims, 'Bicycle', {}, num_footprints=3) + assert figure is not None + + +class TestConfig: + def test_lattice_config_keys(self): + assert 'Differential Drive' in LATTICE_CONFIG + assert 'Bicycle' in LATTICE_CONFIG + assert 'Articulated' in LATTICE_CONFIG + + def test_lattice_config_fields(self): + config = LATTICE_CONFIG['Bicycle'] + assert config.group_key == 'drive_velocity' + assert config.vel_key == 'drive_velocity' + assert config.angle_key == 'steering_angle' + + def test_kinematic_equations_present(self): + for model_type in ('Differential Drive', 'Bicycle', 'Articulated'): + assert model_type in KINEMATIC_EQUATIONS + assert 'title' in KINEMATIC_EQUATIONS[model_type] + assert 'equations' in KINEMATIC_EQUATIONS[model_type] + assert 'variables' in KINEMATIC_EQUATIONS[model_type] + + +class TestSingleTrajectory: + def test_single_bicycle_trajectory_reaches_target(self): + # Initial=0, target=0.3, rate=0.3 rad/s, duration=2s → step adds at most 0.3*dt; + # over 2s we reach the target well before the horizon ends. The trajectory's + # `steering_angle` field reflects the target (steady-state). + traj = single_bicycle_trajectory( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + initial_steering_angle_rad=0.0, + target_steering_angle_rad=0.3, + steering_rate_rad_s=0.3, + drive_velocity=1.0, + duration=2.0, + time_step=0.1, + ) + assert traj.steering_angle == pytest.approx(0.3) + assert len(traj.time) == 21 + # Pose advanced from origin (we drove forward). + assert traj.x[-1] > 0.0 + + def test_single_articulated_trajectory_reaches_target(self): + traj = single_articulated_trajectory( + articulation_to_front=1.66, + articulation_to_rear=1.44, + front_track=2.0, + rear_track=2.0, + front_wheel_radius=0.723, + rear_wheel_radius=0.723, + initial_articulation_angle_rad=0.0, + target_articulation_angle_rad=0.4, + articulation_rate_rad_s=0.5, + drive_velocity=1.0, + duration=2.0, + time_step=0.1, + ) + assert traj.articulation_angle == pytest.approx(0.4) + assert len(traj.time) == 21 + assert traj.x[-1] > 0.0 + + def test_single_differential_trajectory_ramps_both_velocities(self): + traj = single_differential_trajectory( + wheel_radius=0.1, + track_width=0.5, + initial_linear_velocity=0.0, + initial_angular_velocity=0.0, + target_linear_velocity=1.0, + target_angular_velocity=0.5, + linear_acceleration=1.0, + angular_acceleration=1.0, + duration=2.0, + time_step=0.1, + ) + # Targets reached in ~1.0s and ~0.5s respectively; final state pinned to target. + assert traj.linear_velocity == pytest.approx(1.0) + assert traj.angular_velocity == pytest.approx(0.5) + assert len(traj.time) == 21 + + def test_single_bicycle_zero_rate_keeps_initial_angle(self): + # With rate=0 the angle never advances; turning_radius reflects the (unchanging) + # initial angle, not the target. + traj = single_bicycle_trajectory( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + initial_steering_angle_rad=0.0, + target_steering_angle_rad=0.5, + steering_rate_rad_s=0.0, + drive_velocity=1.0, + duration=1.0, + time_step=0.1, + ) + # Zero steering → straight line along +x. + assert traj.x[-1] == pytest.approx(1.0) + assert traj.y[-1] == pytest.approx(0.0) + + def test_single_bicycle_steering_series_brackets_initial_and_target(self): + # Initial=0 → target=0.3 at rate=0.3 rad/s; ramp completes in 1s. Over duration=2s + # the series starts at 0 and ends pinned at 0.3. + traj = single_bicycle_trajectory( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + initial_steering_angle_rad=0.0, + target_steering_angle_rad=0.3, + steering_rate_rad_s=0.3, + drive_velocity=1.0, + duration=2.0, + time_step=0.1, + ) + assert traj.steering_angle_series is not None + assert len(traj.steering_angle_series) == len(traj.time) + assert traj.steering_angle_series[0] == pytest.approx(0.0) + assert traj.steering_angle_series[-1] == pytest.approx(0.3) + + def test_single_articulated_articulation_series_brackets_initial_and_target(self): + traj = single_articulated_trajectory( + articulation_to_front=1.66, + articulation_to_rear=1.44, + front_track=2.0, + rear_track=2.0, + front_wheel_radius=0.723, + rear_wheel_radius=0.723, + initial_articulation_angle_rad=0.0, + target_articulation_angle_rad=0.4, + articulation_rate_rad_s=0.5, + drive_velocity=1.0, + duration=2.0, + time_step=0.1, + ) + assert traj.articulation_angle_series is not None + assert len(traj.articulation_angle_series) == len(traj.time) + assert traj.articulation_angle_series[0] == pytest.approx(0.0) + assert traj.articulation_angle_series[-1] == pytest.approx(0.4) + + def test_single_bicycle_footprint_series_none_without_dims(self): + traj = single_bicycle_trajectory( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + initial_steering_angle_rad=0.0, + target_steering_angle_rad=0.2, + steering_rate_rad_s=0.0, + drive_velocity=1.0, + duration=1.0, + time_step=0.1, + ) + # No footprint dimensions passed → projector emits empty footprints → series is None. + assert traj.footprint_series is None + + def test_single_bicycle_footprint_series_present_with_dims(self): + traj = single_bicycle_trajectory( + wheelbase=2.5, + track_width=1.5, + wheel_radius=0.3, + initial_steering_angle_rad=0.0, + target_steering_angle_rad=0.2, + steering_rate_rad_s=0.0, + drive_velocity=1.0, + duration=1.0, + time_step=0.1, + footprint=rectangle_footprint(3.0, 1.0, 1.5), + ) + assert traj.footprint_series is not None + # (N samples, 4 corners, xy) + assert traj.footprint_series.shape == (len(traj.time), 4, 2) + + def test_single_articulated_footprint_series_present_with_dims(self): + traj = single_articulated_trajectory( + articulation_to_front=1.66, + articulation_to_rear=1.44, + front_track=2.0, + rear_track=2.0, + front_wheel_radius=0.723, + rear_wheel_radius=0.723, + initial_articulation_angle_rad=0.0, + target_articulation_angle_rad=0.4, + articulation_rate_rad_s=0.5, + drive_velocity=1.0, + duration=1.0, + time_step=0.1, + front_footprint=rectangle_footprint(0.54, 1.66, 2.0), + rear_footprint=rectangle_footprint(1.44, 0.56, 2.0), + ) + assert traj.front_footprint_series is not None + assert traj.rear_footprint_series is not None + assert traj.front_footprint_series.shape == (len(traj.time), 4, 2) + assert traj.rear_footprint_series.shape == (len(traj.time), 4, 2) + + def test_single_articulated_footprint_extends_behind_rear_axle(self): + # The rear polygon is measured from the rear axle, and the pose (REAR reference) IS the + # rear axle, so a rear_overhang of 0.8 must put the rearmost point at exactly -0.8. + articulation_to_rear = 1.44 + rear_overhang = 0.8 + traj = single_articulated_trajectory( + articulation_to_front=1.66, + articulation_to_rear=articulation_to_rear, + front_track=2.0, + rear_track=2.0, + front_wheel_radius=0.723, + rear_wheel_radius=0.723, + initial_articulation_angle_rad=0.0, + target_articulation_angle_rad=0.0, + articulation_rate_rad_s=0.0, + drive_velocity=0.0, # stay at the origin so the geometry is exact + duration=0.2, + time_step=0.1, + front_footprint=rectangle_footprint(1.0, 1.66, 2.0), + rear_footprint=rectangle_footprint(articulation_to_rear, rear_overhang, 2.0), + ) + assert traj.rear_footprint_series is not None + rear0 = traj.rear_footprint_series[0] # (4, 2) corners of the first sample + assert rear0[:, 0].min() == pytest.approx(-rear_overhang) + # Forward extent reaches the joint, which sits L_r ahead of the rear axle. + assert rear0[:, 0].max() == pytest.approx(articulation_to_rear) + + def test_single_bicycle_footprint_extends_ahead_of_front_axle(self): + # Pose reference is the rear axle, so front_overhang_m = wheelbase + overhang. Regression + # guard: a wheelbase fraction used to put the front bumper behind the front axle. + wheelbase = 2.5 + front_overhang = 0.6 + rear_overhang = 0.5 + traj = single_bicycle_trajectory( + wheelbase=wheelbase, + track_width=1.5, + wheel_radius=0.3, + initial_steering_angle_rad=0.0, + target_steering_angle_rad=0.0, + steering_rate_rad_s=0.0, + drive_velocity=0.0, # stay at the origin so the geometry is exact + duration=0.2, + time_step=0.1, + footprint=rectangle_footprint(wheelbase + front_overhang, rear_overhang, 1.5), + ) + assert traj.footprint_series is not None + body0 = traj.footprint_series[0] # (4, 2) corners of the first sample + max_x = body0[:, 0].max() + min_x = body0[:, 0].min() + assert max_x == pytest.approx(wheelbase + front_overhang) + assert max_x > wheelbase # front bumper is AHEAD of the front axle + assert min_x == pytest.approx(-rear_overhang) + assert max_x - min_x == pytest.approx(wheelbase + front_overhang + rear_overhang) + + def test_single_bicycle_axle_reference_describes_the_same_body(self): + # The same physical vehicle described from each axle must put the body in the same place. + wheelbase, front_overhang, rear_overhang, width = 2.5, 0.6, 0.5, 1.5 + common = dict( + track_width=width, + wheel_radius=0.3, + initial_steering_angle_rad=0.0, + target_steering_angle_rad=0.0, + steering_rate_rad_s=0.0, + drive_velocity=0.0, + duration=0.2, + time_step=0.1, + ) + from_rear = single_bicycle_trajectory( + wheelbase=wheelbase, + axle_reference=AxleReference.REAR, + footprint=rectangle_footprint(wheelbase + front_overhang, rear_overhang, width), + **common, + ) + from_front = single_bicycle_trajectory( + wheelbase=wheelbase, + axle_reference=AxleReference.FRONT, + footprint=rectangle_footprint(front_overhang, wheelbase + rear_overhang, width), + **common, + ) + # Both start at the origin *of their own reference axle*, so the front-referenced body + # sits a wheelbase back in world terms; shift it to compare the same physical placement. + rear_body = from_rear.footprint_series[0] + front_body = from_front.footprint_series[0] + np.array([wheelbase, 0.0]) + assert front_body == pytest.approx(rear_body) + + def test_single_articulated_reports_joint_pose_series(self): + articulation_to_rear = 1.44 + traj = single_articulated_trajectory( + articulation_to_front=1.66, + articulation_to_rear=articulation_to_rear, + front_track=2.0, + rear_track=2.0, + front_wheel_radius=0.723, + rear_wheel_radius=0.723, + initial_articulation_angle_rad=0.0, + target_articulation_angle_rad=0.0, + articulation_rate_rad_s=0.0, + drive_velocity=0.0, + duration=0.2, + time_step=0.1, + ) + assert traj.joint_pose_series is not None + assert traj.joint_pose_series.shape == (len(traj.time), 3) + # REAR reference at the origin: the joint sits L_r straight ahead. + assert traj.joint_pose_series[0][0] == pytest.approx(articulation_to_rear) + assert traj.joint_pose_series[0][1] == pytest.approx(0.0) + + def test_single_differential_footprint_straddles_body_centre(self): + # Pose reference is the body centre, so the overhangs are the bumper distances directly. + front_overhang = 0.4 + rear_overhang = 0.3 + traj = single_differential_trajectory( + wheel_radius=0.1, + track_width=0.5, + initial_linear_velocity=0.0, + initial_angular_velocity=0.0, + target_linear_velocity=0.0, + target_angular_velocity=0.0, + linear_acceleration=0.0, + angular_acceleration=0.0, + duration=0.2, + time_step=0.1, + footprint=rectangle_footprint(front_overhang, rear_overhang, 0.5), + ) + assert traj.footprint_series is not None + body0 = traj.footprint_series[0] + assert body0[:, 0].max() == pytest.approx(front_overhang) + assert body0[:, 0].min() == pytest.approx(-rear_overhang)