Source code for state_space_setup_assistant.export

"""Write the artifact bundle for a finished wizard run.

Everything in the bundle reproduces through the existing CLIs:

    ros2 run urdf_state_space urdf2ss model.yaml
    ros2 run state_space_control ss_design plant.npz design.yaml -o controller.npz
"""

import copy
import os
from typing import Dict, List

import numpy as np

from state_space_control.analysis import damping_report
from state_space_control.base import Plant

from .ros2_control_export import ros2_control_yaml_dict
from .yamlgen import dump_yaml

REPRO_COMMANDS = [
    'ros2 run urdf_state_space urdf2ss model.yaml',
    'ros2 run state_space_control ss_design plant.npz design.yaml '
    '-o controller.npz',
]


def _validation_md(findings: List[Dict]) -> str:
    icon = {'ok': '✓', 'warn': '⚠', 'error': '✗'}
    lines = ['# URDF validation report', '']
    for f in findings:
        lines.append(f"- {icon.get(f['level'], '?')} **{f['code']}** "
                     f"— {f['message']}")
    return '\n'.join(lines) + '\n'


def _benchmark_md(bench: Dict) -> str:
    lines = [
        '# Controller benchmark', '',
        '| Controller | Stable | Settling [s] | Overshoot [%] '
        '| Effort ∫u²dt | Ms | GM [dB] | PM [deg] |',
        '|---|---|---|---|---|---|---|---|',
    ]

    def fmt(v):
        if v is None:
            return 'n/a'
        if isinstance(v, float):
            return 'n/a' if not np.isfinite(v) else f'{v:.4g}'
        return str(v)

    for row in bench.get('table', []):
        if 'error' in row:
            lines.append(f"| {row['controller']} | — | — | — | — | — | — | "
                         f"error: {row['error']} |")
            continue
        lines.append(
            f"| {row['controller']} | {row['stable']} "
            f"| {fmt(row['settling_time'])} | {fmt(row['overshoot_pct'])} "
            f"| {fmt(row['control_effort'])} | {fmt(row['Ms'])} "
            f"| {fmt(row['gain_margin_db'])} "
            f"| {fmt(row['phase_margin_deg'])} |")
    return '\n'.join(lines) + '\n'


[docs] def write_bundle(session, output_dir: str, include_mat: bool = False, dt: float = None) -> Dict: """Write all artifacts; returns {files: [...], commands: [...]}.""" if session.model is None or session.model_cfg is None: raise ValueError('nothing to export: linearize a model first') os.makedirs(output_dir, exist_ok=True) files = [] def emit(name): files.append(name) return os.path.join(output_dir, name) cfg = copy.deepcopy(session.model_cfg) export = {'npz': 'plant.npz'} if include_mat: export['mat'] = 'plant.mat' if dt: export['dt'] = float(dt) cfg['export'] = export dump_yaml(cfg, emit('model.yaml')) session.model.save_npz(os.path.join(output_dir, 'plant.npz')) files.append('plant.npz') if include_mat: session.model.save_mat(os.path.join(output_dir, 'plant.mat')) files.append('plant.mat') summary_parts = ['# plant', session.model.summary()] if session.design_cfg is not None and session.result is not None: dump_yaml(session.design_cfg, emit('design.yaml')) session.result.save_npz(os.path.join(output_dir, 'controller.npz')) files.append('controller.npz') dump_yaml(ros2_control_yaml_dict(session.model, session.result), emit(f'{session.result.name}_ros2_control.yaml')) summary_parts += ['', '# controller', session.result.summary(), '', '# closed-loop damping', damping_report(session.result.closed_loop())] if session.validation: with open(emit('validation.md'), 'w') as f: f.write(_validation_md(session.validation)) if session.benchmark: with open(emit('benchmark.md'), 'w') as f: f.write(_benchmark_md(session.benchmark)) traces = session.benchmark.get('traces', {}) arrays = {'t': np.asarray(traces.get('t', []))} for name, y in traces.get('y', {}).items(): arrays[f'y_{name}'] = np.asarray(y) for name, u in traces.get('u', {}).items(): arrays[f'u_{name}'] = np.asarray(u) np.savez(os.path.join(output_dir, 'benchmark_traces.npz'), **arrays) files.append('benchmark_traces.npz') summary_parts += ['', '# reproduce', *REPRO_COMMANDS] with open(emit('summary.txt'), 'w') as f: f.write('\n'.join(summary_parts) + '\n') return {'output_dir': os.path.abspath(output_dir), 'files': files, 'commands': REPRO_COMMANDS}