dataframe ¶

def select(self, *columns: Column | str) -> MojoDataFrame:
    """
    Keeps every column that is one of `columns` or lies beneath one, as decided by `Column.covers`. A string is parsed with `Column.parse`.

    A partial column selects a subtree, so `select("Bodies/box1")` keeps `Bodies/box1/xpos:x` and `Bodies/box1:ke_trans`, while `select(Column(category="Bodies", attr="x"))` keeps every body's `:x` components. Names are matched as parts, never as patterns (each part is escaped), so an object name containing regex characters is matched literally.
    """
    wanted = [Column.parse(c) if isinstance(c, str) else c for c in columns]
    if not wanted:
        return _MojoFrame.from_pl(self._df.select([]))
    # matched by polars' regex engine, so a frame with tens of thousands of columns is not walked in Python
    return _MojoFrame.from_pl(self._df.select(pl.col(covering_pattern(wanted))))

def pose_at(
    self,
    index: int,
    source: Column | str,
    *,
    pos_channel: str = "xpos",
    quat_channel: str = "quat",
) -> PoseQuat:
    """
    The pose of `source` (for example `Sites/A`) at row `index`.

    Reads only the seven pose columns of that one row, so the cost does not grow with the size of the frame. Raises a `ValueError` naming every missing column. For every row at once use `pose_arrays`.
    """
    from mujoco_mojo.mjcf.pose import PoseQuat

    pos_names, quat_names = pose_column_names(source, pos_channel, quat_channel)
    names = (*pos_names, *quat_names)
    try:
        # `get_column` is a lookup by name; `select` would cost time proportional to the frame's width
        row = {n: self._df.get_column(n)[index] for n in names}
    except pl.exceptions.ColumnNotFoundError:
        raise pose_missing_error(
            source, self._df.columns, pos_channel, quat_channel
        ) or ValueError(f"No pose for {str(source)!r}") from None
    return PoseQuat.from_row(
        row, source, pos_channel=pos_channel, quat_channel=quat_channel
    )

def pose_arrays(
    self,
    source: Column | str,
    *,
    pos_channel: str = "xpos",
    quat_channel: str = "quat",
) -> tuple[np.ndarray, np.ndarray]:
    """
    The pose of `source` over the whole frame as plain arrays: positions `(n_rows, 3)` and quaternions `(n_rows, 4)` in MJCF `w, x, y, z` order.

    Prefer this over calling `pose_at` per row for a trajectory: it makes one array copy per group, where a per-row `PoseQuat` would validate a pydantic model for every row. Raises a `ValueError` naming every missing column.
    """
    pos_names, quat_names = pose_column_names(source, pos_channel, quat_channel)
    try:
        pos = np.column_stack(
            [self._df.get_column(n).to_numpy() for n in pos_names]
        )
        quat = np.column_stack(
            [self._df.get_column(n).to_numpy() for n in quat_names]
        )
    except pl.exceptions.ColumnNotFoundError:
        raise pose_missing_error(
            source, self._df.columns, pos_channel, quat_channel
        ) or ValueError(f"No pose for {str(source)!r}") from None
    return pos, quat

def select_category(self, category: SignalCategory | str) -> MojoDataFrame:
    """Filter columns belonging to a specific SignalCategory (e.g., 'Bodies')."""
    return _MojoFrame.from_pl(self._df.select(pl.col(rf"^{category}/.*$")))

def select_name(self, name: str) -> MojoDataFrame:
    """General filter for columns associated with a specific object name (e.g., 'racket')."""
    # Matches Category/Name:Attr or Category/Name/Sub:Attr
    return _MojoFrame.from_pl(self._df.select(pl.col(rf"^[^/]+/{name}/.*$")))

def select_channel(self, channel: str) -> MojoDataFrame:
    """
    Selects all components of a specific channel across any category.

    Matches the logical 'folder' before the attribute separator.
    Example: 'xpos' matches 'Bodies/Hand/xpos:x' and 'Bodies/Hand/xpos:y'.
    """
    # Regex Breakdown:
    # ^.*/        -> Start and match any prefix ending in a slash (the path)
    # {channel}   -> The specific channel name (e.g., xpos)
    # (?::.*)?    -> Optionally match a colon followed by any attribute component
    # $           -> End of string
    return _MojoFrame.from_pl(self._df.select(pl.col(rf"^.*/{channel}(?::.*)?$")))

def select_attribute(self, attr: str) -> MojoDataFrame:
    """
    Selects a specific attribute across all categories and objects.
    Matches columns ending in ':attr' (e.g. attr='x' matches 'Bodies/racket/xpos:x' and 'Sensors/gyro/data:x').
    """
    return _MojoFrame.from_pl(self._df.select(pl.col(rf"^.*:{attr}$")))

def select_path_part(self, part: str) -> MojoDataFrame:
    """
    Selects any column whose full path contains `part` as a substring, anywhere.

    Unlike the other `select_*` methods, this does not anchor to a specific position in the path (category, name, channel, or attribute) - it matches `part` wherever it appears.

    Example:
        `select_path_part("racket")` matches `Bodies/racket/xpos:x`, `Custom/MyGroup:racket`, and `Bodies/racket_arm/xpos:x`.

    """
    return _MojoFrame.from_pl(
        self._df.select([c for c in self._df.columns if part in c])
    )

def select_joint(self, name: JointName) -> MojoDataFrame:
    """Select all signals belonging to a specific Joint (qpos, qvel, etc.)."""
    return self._select_object(SignalCategory.JOINTS, name)

def select_site(self, name: SiteName) -> MojoDataFrame:
    """Select all signals recorded at a specific Site."""
    return self._select_object(SignalCategory.SITES, name)

def select_geom(self, name: GeomName) -> MojoDataFrame:
    """Select all signals associated with a specific Geom (contacts, etc.)."""
    return self._select_object(SignalCategory.GEOMS, name)

def select_sensor(self, name: SensorName) -> MojoDataFrame:
    """Select data from a specific named Sensor."""
    return self._select_object(SignalCategory.SENSORS, name)

def select_actuator(self, name: ActuatorName) -> MojoDataFrame:
    """Select data from a specific Actuator."""
    return self._select_object(SignalCategory.ACTUATORS, name)

def select_tendon(self, name: TendonName) -> MojoDataFrame:
    """Select data from a specific Tendon."""
    return self._select_object(SignalCategory.TENDONS, name)

def select_camera(self, name: CameraName) -> MojoDataFrame:
    """Select pose or FOV data from a specific Camera."""
    return self._select_object(SignalCategory.CAMERAS, name)

def select_light(self, name: LightName) -> MojoDataFrame:
    """Select pose or intensity data from a specific Light."""
    return self._select_object(SignalCategory.LIGHTS, name)

def select_equality(self, name: EqualityName) -> MojoDataFrame:
    """Select force/error data from an Equality constraint."""
    return self._select_object(SignalCategory.CONSTRAINTS, name)

def select_plugin(self, name: InstanceName) -> MojoDataFrame:
    """Select custom state data from a specific Plugin Instance."""
    return self._select_object(SignalCategory.PLUGINS, name)

def select_flex(self, name: FlexName) -> MojoDataFrame:
    """Select vertex/stress data from a Deformable Flex object."""
    return self._select_object(SignalCategory.DEFORMABLES, name)

def get_manifest(
    self,
    extra_columns: list[str] | None = None,
    column_metadata: Mapping[str, ColumnMetadata] | None = None,
) -> ColumnManifest:
    """Returns the structured manifest used by the frontend. `extra_columns` are appended to `all` and included in rotatable/quat discovery. Pass `column_metadata` (from `MojoDataFrame.read_column_metadata()`) to populate `column_metadata`, and to leave out any group tagged `scalar`, which is not a vector or a quaternion and so is never offered for rotation (matching `change_frame`)."""
    bm = self._get_base_map(extra_columns)
    all_cols = list(self._df.columns) + (extra_columns or [])
    meta = column_metadata or {}
    col_meta = {col: m for col in all_cols if (m := meta.get(col)) is not None}

    def is_scalar(base: str, first: str) -> bool:
        tagged = meta.get(f"{base}:{first}")
        return tagged is not None and tagged.transform_type == TransformType.SCALAR

    return {
        "all": all_cols,
        "rotatable_vectors": sorted(
            b
            for b, s in bm.items()
            if {"x", "y", "z"}.issubset(s)
            and "w" not in s
            and not is_scalar(b, "x")
        ),
        "available_quats": sorted(
            b
            for b, s in bm.items()
            if {"w", "x", "y", "z"}.issubset(s) and not is_scalar(b, "w")
        ),
        "column_metadata": col_meta,
    }

def with_unit_system(
    self,
    target: UnitSystem,
    *,
    path: Path | str | None = None,
    column_metadata: Mapping[str, ColumnMetadata] | None = None,
    assume_source: UnitSystem | None = None,
) -> MojoDataFrame:
    """
    Converts all columns with known units into the target unit system.

    For each column whose metadata carries a concrete `"unit"` key, applies a unit conversion from that stored unit into the equivalent unit in `target`. If `assume_source` is given, columns that only carry a `"dimension"` key (no concrete unit) are treated as if their source unit is the corresponding unit in `assume_source`, so they are converted too.

    All conversions are collected and applied in a single `with_columns()` call (Polars handles the vectorised execution across columns), so this is equivalent in cost to one pass over the data regardless of how many unit groups there are.

    Args:
        target: The target unit system (e.g. `UnitSystem.si()`).
        path: Path to the parquet file to read column metadata from. Used when `column_metadata` is not passed directly.
        column_metadata: Pre-loaded metadata dict (from `MojoDataFrame.read_column_metadata()`). Takes precedence over `path`.
        assume_source: When set, columns with only a `"dimension"` tag (no `"unit"`) are converted as if they were expressed in the corresponding unit from this system.

    """
    from mujoco_mojo.stochas import UnitSystem as _US
    from mujoco_mojo.stochas import ureg
    from mujoco_mojo.utils.filters.filters import UnitFilter

    meta = _resolve_column_metadata(path, column_metadata)

    def _base_map(us: _US) -> dict[str, str]:
        return {
            k: v
            for k, v in {
                "[length]": us.length,
                "[mass]": us.mass,
                "[time]": us.time,
                "[temperature]": us.temperature,
                "[current]": us.current,
                "[substance]": us.amount,
                "[luminosity]": us.luminosity,
            }.items()
            if v is not None
        }

    target_map = _base_map(target)
    source_map = _base_map(assume_source) if assume_source is not None else None

    def _unit_str_for(
        dim_dict: dict[str, Any], unit_map: dict[str, str]
    ) -> str | None:
        """Build a unit string from a dimensionality dict + base-unit map. Returns None if any required dimension is unconfigured."""
        u = ureg.dimensionless
        for dim_key, exp in dim_dict.items():
            base = unit_map.get(dim_key)
            if base is None:
                return None
            u = u * ureg.parse_units(base) ** exp
        return str(u)

    exprs = []
    for col in self._df.columns:
        col_meta = meta.get(col)
        src_str: str | None = None

        if col_meta is not None and col_meta.unit is not None:
            try:
                dim_dict = dict(ureg.get_dimensionality(col_meta.unit))
            except Exception:
                continue
            src_str = col_meta.unit
        elif (
            source_map is not None
            and col_meta is not None
            and col_meta.dimension is not None
        ):
            dim_str = col_meta.dimension
            if dim_str == "[]":
                continue
            try:
                dim_dict = dict(ureg.get_dimensionality(dim_str))
            except Exception:
                continue
            src_str = _unit_str_for(dim_dict, source_map)
            if src_str is None:
                continue
        else:
            continue

        tgt_str = _unit_str_for(dim_dict, target_map)
        if tgt_str is None:
            continue

        try:
            exprs.append(
                UnitFilter(from_unit=src_str, to_unit=tgt_str)
                .apply(pl.col(col))
                .alias(col)
            )
        except Exception:
            continue

    if not exprs:
        return _MojoFrame.from_pl(self._df)

    return _MojoFrame.from_pl(self._df.with_columns(exprs))

def change_frame(
    self,
    quat_base: str,
    origin_base: str,
    *,
    path: Path | str | None = None,
    column_metadata: Mapping[str, ColumnMetadata] | None = None,
    invert: bool = True,
) -> MojoDataFrame:
    """
    Re-expresses every tagged point/vector/quaternion column in this frame relative to the frame defined by `origin_base` (position) and `quat_base` (orientation).

    Applies the correct transform per column by its declared `transform_type`: points are translated then rotated, vectors are rotated only, and quaternions are composed (`q' = q_b^-1 (x) q`). Every rotatable/quaternion column in the frame must be tagged (see `mujoco_mojo.utils.signal_metadata.TransformType`) or this raises - there is no untagged fallback, since guessing the wrong kind is exactly the silent-corruption bug this method exists to prevent.

    Args:
        quat_base: Prefix for the [x,y,z,w] quaternion group giving the target frame's orientation in world coordinates (e.g. 'Bodies/chassis/xquat').
        origin_base: Prefix for the [x,y,z] position group giving the target frame's origin in world coordinates (e.g. 'Bodies/chassis/xpos').
        path: Path to the parquet file to read column metadata from. Used when `column_metadata` is not passed directly.
        column_metadata: Pre-loaded metadata dict (from `MojoDataFrame.read_column_metadata()`). Takes precedence over `path`.
        invert: If True (default), performs a world-to-local transform into the target frame; if False, reverses it (local-to-world).

    Returns:
        Self: DataFrame with every point/vector/quaternion column re-expressed in the target frame.

    Raises:
        ValueError: If `quat_base`/`origin_base` aren't valid column groups in this frame, or if any rotatable/quaternion column lacks a `transform_type` tag.

    """
    if quat_base not in self.quaternion_bases:
        msg = f"change_frame failed: quaternion base '{quat_base}' not found (see the quaternion_bases property for valid columns)."
        logger.error(msg)
        raise ValueError(msg)
    if origin_base not in self.rotatable_bases:
        msg = f"change_frame failed: origin base '{origin_base}' not found (see the rotatable_bases property for valid columns)."
        logger.error(msg)
        raise ValueError(msg)

    meta = _resolve_column_metadata(path, column_metadata)

    points, vectors = _classify_transform_bases(self.rotatable_bases, meta)
    quats = _classify_quaternion_bases(self.quaternion_bases, meta)

    rf = RotationFilter(quat_col=quat_base, origin_col=origin_base, invert=invert)
    rotated = rf.apply_to_frame(self._df, vector_bases=vectors, point_bases=points)
    rotated = rf.compose_to_frame(rotated, quaternion_bases=quats)
    return _MojoFrame.from_pl(rotated)

def with_filter_map(
    self, filter_map: dict[str, list[AnyFilter]], omit_time: bool = True
) -> MojoDataFrame:
    """
    Applies specific filter stacks to mapped columns.

    Args:
        filter_map (dict[str, list[AnyFilter]]): Dictionary mapping column names to a list of filters.
        omit_time (bool, optional): If True, skips the 'time' column even if present in the map. Defaults to True.

    Returns:
        Self: DataFrame with the transformed columns overwritten.

    """
    exprs = []
    for col_name, filters in filter_map.items():
        if col_name not in self._df.columns:
            continue

        if omit_time and col_name == TIME_COLUMN_NAME:
            continue

        expr = pl.col(col_name)
        for f in filters:
            expr = f.apply(expr)

        exprs.append(expr.alias(col_name))

    return _MojoFrame.from_pl(self._df.with_columns(exprs))

def with_filters(
    self,
    filters: list[AnyFilter],
    columns: list[str] | None = None,
    omit_time: bool = True,
) -> MojoDataFrame:
    """
    Applies a sequential stack of filters to the specified or all numeric columns.

    Args:
        filters (list[AnyFilter]): List of Filter objects to apply in order (e.g., LowPass -> Derivative).
        columns (list[str] | None, optional): Specific columns to transform. If None, applies to all available columns. Defaults to None.
        omit_time (bool, optional): If True, prevents filters from being applied to the 'time' column. Defaults to True.

    Returns:
        Self: DataFrame with the transformed columns overwritten.

    """
    target_cols = columns or self._df.columns
    filter_map = {col: filters for col in target_cols}
    return self.with_filter_map(filter_map, omit_time=omit_time)