signal_metadata ¶

def resolve_dimension_metadata(
    meta: ColumnMetadata, unit_system: UnitSystem
) -> ColumnMetadata:
    """Adds a concrete `unit` string to `meta` derived from its `dimension` and the given `unit_system`. The `dimension` is preserved so callers can still see both the physical quantity type and the concrete unit. Dimensionless signals (`dimension="[]"`) are left unchanged. All other fields (e.g. `quantity`) are preserved. Cached, since `merge_signal_metadata` runs every timestep."""
    if meta.dimension is None or meta.dimension == "[]":
        return meta
    return _resolve_dimension(meta, _base_units(unit_system))

@cache
def dim(dimension: Dimension | str) -> ColumnMetadata:
    """Builds a `dimension`-tagged `ColumnMetadata`. Accepts either a `Dimension` enum member or a Pint unit name string (e.g. `"inch"`), in which case Pint looks up the matching `Dimension` member by dimensionality."""
    if isinstance(dimension, Dimension):
        # check Dimension first: it's a StrEnum and is also an instance of str,
        # so this must come before the str branch to avoid the Pint-lookup path running on
        # a known Dimension value (which would return the wrong enum member for TORQUE vs ENERGY)
        return ColumnMetadata(dimension=str(dimension))
    from mujoco_mojo.stochas import ureg

    target = ureg.get_dimensionality(dimension)
    for member in Dimension:
        if ureg.get_dimensionality(member.value) == target:
            return ColumnMetadata(dimension=str(member))
    raise ValueError(
        f"No Dimension member matches the Pint dimensionality of unit {dimension!r}"
    )

@cache
def unit(unit: str) -> ColumnMetadata:
    """Builds a `unit`-tagged `ColumnMetadata`."""
    return ColumnMetadata(unit=unit)

def dimensionless_metadata() -> ColumnMetadata:
    """Tags a signal as known to have no units (e.g. quaternions, rotation matrices, axis vectors, enum/index values), distinguishing "known to be unitless" from "nobody tagged this column"."""
    return dim(Dimension.DIMENSIONLESS)

def angle_metadata() -> ColumnMetadata:
    """Tags a signal as an angle in radians. MuJoCo's compiled model always stores angles in radians regardless of the XML `angle=` attribute, so this is a concrete unit rather than a scale-ambiguous dimension."""
    return unit(str(Angle.RADIAN))

def angular_rate_metadata(per: str = "second") -> ColumnMetadata:
    """Tags a signal as an angular rate (e.g. angular velocity, angular acceleration) in radians per `per`."""
    return unit(f"{Angle.RADIAN} / {per}")

def torque_metadata() -> ColumnMetadata:
    """Tags a signal as torque. Includes an extra `quantity` hint since torque and energy share the same Pint dimensionality and can't be told apart by dimension alone."""
    return _torque_metadata()

@cache
def force_or_torque(jnt_type: int) -> ColumnMetadata:
    """Returns force metadata for a slide joint's generalized force, or torque metadata for any rotational joint (hinge/ball/free)."""
    import mujoco

    if jnt_type == mujoco.mjtJoint.mjJNT_SLIDE:
        return dim(Dimension.FORCE)
    return torque_metadata()

def merge_signal_metadata(
    builtin: ColumnMetadata | None,
    channel: str,
    user_metadata: MetadataOverrides | None,
    *,
    unit_system: UnitSystem | None = None,
) -> ColumnMetadata:
    """Merges a built-in default with the caller-supplied override/extension for `channel`, with user-supplied fields winning on conflict. If `unit_system` is provided and `builtin` has a `dimension`, it is resolved to a concrete unit. Returns an empty `ColumnMetadata` if both are empty. Nothing here is validated: the result is validated when it is used to build a `Column`."""
    if builtin is not None and unit_system is not None:
        builtin = resolve_dimension_metadata(builtin, unit_system)
    override = (user_metadata or {}).get(channel)
    if override is None:
        return builtin if builtin is not None else ColumnMetadata()
    override_meta = ColumnMetadata.unchecked(override)
    return override_meta if builtin is None else builtin | override_meta