Source code for state_space_setup_assistant.yamlgen

"""Builders for the two YAML artifacts, in the *existing* schemas.

The wizard never calls the math engines with ad-hoc arguments: it writes
these dicts to disk and runs urdf_state_space.build_from_yaml / the same
controller path as ss_design on them, so the exported files are exactly what
produced the on-screen results.
"""

import os
from typing import Dict, List, Optional

import yaml




def model_yaml_dict(
    urdf_source: str,
    joints: Dict[str, Dict],
    velocities: bool = False,
    output_joints: Optional[List[str]] = None,
    floating_base: bool = False,
    export_npz: Optional[str] = None,
    export_mat: Optional[str] = None,
    dt: Optional[float] = None,
) -> Dict:
    """Build a model.yaml dict (schema of urdf_state_space.config).

    ``joints`` maps joint name -> {actuated?: bool, damping?: float,
    q_eq?: float}; empty per-joint entries are dropped.
    """
    cfg_joints = {}
    for name, entry in joints.items():
        e = {}
        if entry.get('actuated'):
            e['actuated'] = True
        if entry.get('damping') not in (None, 0, 0.0):
            e['damping'] = float(entry['damping'])
        if entry.get('q_eq') not in (None,):
            if float(entry['q_eq']) != 0.0:
                e['q_eq'] = float(entry['q_eq'])
        if e:
            cfg_joints[name] = e
    if not any(e.get('actuated') for e in cfg_joints.values()):
        raise ValueError('at least one joint must be actuated')

    cfg = {'model': {'urdf': urdf_source}}
    if floating_base:
        cfg['model']['floating_base'] = True
    cfg['joints'] = cfg_joints
    outputs = {}
    if velocities:
        outputs['velocities'] = True
    if output_joints is not None:
        outputs['joints'] = list(output_joints)
    if outputs:
        cfg['outputs'] = outputs
    export = {}
    if export_npz:
        export['npz'] = export_npz
    if export_mat:
        export['mat'] = export_mat
    if dt:
        export['dt'] = float(dt)
    if export:
        cfg['export'] = export
    return cfg





def design_yaml_dict(controller: str, params: Optional[Dict] = None) -> Dict:
    """Build a design.yaml dict (schema of state_space_control's ss_design)."""
    cfg = {'controller': controller}
    if params:
        cfg['params'] = params
    return cfg





def dump_yaml(cfg: Dict, path: str) -> str:
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with open(path, 'w') as f:
        yaml.safe_dump(cfg, f, default_flow_style=False, sort_keys=False)
    return path