Dataframe

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 _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.JOINTS}/{name}/.*$"))
    )

def select_site(self, name: SiteName) -> MojoDataFrame:
    """Select all signals recorded at a specific Site."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.SITES}/{name}/.*$"))
    )

def select_geom(self, name: GeomName) -> MojoDataFrame:
    """Select all signals associated with a specific Geom (contacts, etc.)."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.GEOMS}/{name}/.*$"))
    )

def select_sensor(self, name: SensorName) -> MojoDataFrame:
    """Select data from a specific named Sensor."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.SENSORS}/{name}/.*$"))
    )

def select_actuator(self, name: ActuatorName) -> MojoDataFrame:
    """Select data from a specific Actuator."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.ACTUATORS}/{name}/.*$"))
    )

def select_tendon(self, name: TendonName) -> MojoDataFrame:
    """Select data from a specific Tendon."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.TENDONS}/{name}/.*$"))
    )

def select_camera(self, name: CameraName) -> MojoDataFrame:
    """Select pose or FOV data from a specific Camera."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.CAMERAS}/{name}/.*$"))
    )

def select_light(self, name: LightName) -> MojoDataFrame:
    """Select pose or intensity data from a specific Light."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.LIGHTS}/{name}/.*$"))
    )

def select_equality(self, name: EqualityName) -> MojoDataFrame:
    """Select force/error data from an Equality constraint."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.CONSTRAINTS}/{name}/.*$"))
    )

def select_plugin(self, name: InstanceName) -> MojoDataFrame:
    """Select custom state data from a specific Plugin Instance."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.PLUGINS}/{name}/.*$"))
    )

def select_flex(self, name: FlexName) -> MojoDataFrame:
    """Select vertex/stress data from a Deformable Flex object."""
    return _MojoFrame.from_pl(
        self._df.select(pl.col(rf"^{SignalCategory.DEFORMABLES}/{name}/.*$"))
    )

def get_manifest(
    self,
    extra_columns: list[str] | None = None,
    column_metadata: dict[str, dict[str, Any]] | 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 `read_column_metadata()`) to populate `column_units`."""
    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}
    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
        ),
        "available_quats": sorted(
            b for b, s in bm.items() if {"w", "x", "y", "z"}.issubset(s)
        ),
        "column_metadata": col_meta,
    }

def with_unit_system(
    self,
    target: UnitSystem,
    *,
    path: Path | str | None = None,
    column_metadata: dict[str, dict[str, Any]] | 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 `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: dict[str, dict[str, Any]]
    if column_metadata is not None:
        meta = column_metadata
    elif path is not None:
        meta = read_column_metadata(path)
    else:
        meta = {}

    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 "unit" in col_meta:
            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 "dimension" in col_meta:
            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 with_rotation(self, quat_base: str, invert: bool = True) -> MojoDataFrame:
    """
    Rotates all 3D vectors into a new frame using the specified quaternion.

    Args:
        quat_base (str): Prefix for the [w,x,y,z] quaternion group.
        invert (bool, optional): If True, performs World to Local transformation (use False with the same `quat_base` to revert the rotation). Defaults to True.

    Returns:
        Self: DataFrame with transformed :x, :y, :z columns.

    """
    if quat_base not in self.quaternion_bases:
        logger.warning(
            f"Rotation failed: Quaternion base '{quat_base}' not found (please see the quaternion_bases property for valid columns)."
        )
        return _MojoFrame.from_pl(self._df)
    rotated = RotationFilter(quat_col=quat_base, invert=invert).apply_to_frame(
        self._df, self.rotatable_bases
    )
    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)