45 filter graph

docs/user-guides/getting-started.md

@@ -307,7 +307,7 @@ MuJoCo does not currently ship first-party Python typing stubs. To enable proper
 ---
 
 !!! tip "Tip: Scaffold a New Project"
-    Before diving into the generate script, run `mujoco-mojo init` in a new directory to get a working project skeleton — `simulation.py`, `run.sh`, and `reloaded.sh` — pre-wired and ready to edit.
+    Before diving into the generate script, run `mujoco-mojo init` in a new directory to get a working project skeleton (`simulation.py`, `run.sh`, and `reloaded.sh`) pre-wired and ready to edit.
 
     ```bash linenums="0"
     mujoco-mojo init

docs/user-guides/runtime-script.md

@@ -49,7 +49,7 @@ Mojo provides high-level force abstractions like `PointToPointForce`, which auto
 
         ---
 
-        **6-DOF force and torque**. The "multi-tool" of loads—defines full spatial dynamics in either a relative frame or local to a site.
+        **6-DOF force and torque**. Defines full spatial dynamics in either a relative frame or local to a site.
 
         > Think of it as a `VectorForce` and `VectorTorque` combined into one robust plugin.
 

src/mujoco_mojo/utils/dataframe.py

@@ -231,18 +231,22 @@ def select_flex(self, name: FlexName) -> MojoDataFrame:
             self._df.select(pl.col(rf"^{SignalCategory.DEFORMABLES}/{name}/.*$"))
         )
 
-    def _get_base_map(self) -> dict[str, set[str]]:
+    def _get_base_map(
+        self, extra_columns: list[str] | None = None
+    ) -> dict[str, set[str]]:
         """
         Internal helper to group suffixes by their common prefixes.
 
         Example:
-            'Bodies/racket:xvelr:x', 'Bodies/racket:xvelr:y' -> {'Bodies/racket:xvelr': {'x', 'y'}}
+            'Bodies/racket/xvelr:x', 'Bodies/racket/xvelr:y' -> {'Bodies/racket/xvelr': {'x', 'y'}}
+
+        extra_columns are included in discovery but are not expected to exist in self._df.
 
         """
         base_map: dict[str, set[str]] = {}
-        for c in self._df.columns:
+        for c in list(self._df.columns) + (extra_columns or []):
             if ":" in c:
-                prefix, suffix = c.rsplit(":")  # we may have :ke_trans, or ke_rot, etc.
+                prefix, suffix = c.rsplit(":", 1)
                 base_map.setdefault(prefix, set()).add(suffix)
         return base_map
 
@@ -280,12 +284,17 @@ def quaternion_columns(self) -> list[str]:
             expanded.extend([f"{base}:w", f"{base}:x", f"{base}:y", f"{base}:z"])
         return expanded
 
-    def get_manifest(self) -> ColumnManifest:
-        """Returns the structured manifest used by the frontend."""
+    def get_manifest(self, extra_columns: list[str] | None = None) -> ColumnManifest:
+        """Returns the structured manifest used by the frontend. extra_columns are appended to 'all' and included in rotatable/quat discovery."""
+        bm = self._get_base_map(extra_columns)
         return {
-            "all": self._df.columns,
-            "rotatable_vectors": sorted(list(self.rotatable_bases)),
-            "available_quats": sorted(list(self.quaternion_bases)),
+            "all": list(self._df.columns) + (extra_columns or []),
+            "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)
+            ),
         }
 
     def with_rotation(self, quat_base: str, invert: bool = True) -> MojoDataFrame:

src/mujoco_mojo/utils/filters/__init__.py

@@ -2,18 +2,26 @@
     AbsoluteValueFilter,
     AnyFilter,
     ClipFilter,
+    ComparisonFilter,
     DeadbandFilter,
     DerivativeFilter,
+    ExpFilter,
     FilterType,
     HighPassFilter,
     IntegralFilter,
+    LogFilter,
     LowPassFilter,
     MedianFilter,
     NormalizeFilter,
+    PowerFilter,
     RollingMeanFilter,
+    RotationFilter,
+    RoundFilter,
     SavitzkyGolayFilter,
     ScaleFilter,
+    SignFilter,
     TaringFilter,
+    TrigFilter,
     UnitFilter,
     WrapFilter,
     filter_adapter,
@@ -23,18 +31,26 @@
     "AbsoluteValueFilter",
     "AnyFilter",
     "ClipFilter",
+    "ComparisonFilter",
     "DeadbandFilter",
     "DerivativeFilter",
+    "ExpFilter",
     "FilterType",
     "HighPassFilter",
     "IntegralFilter",
+    "LogFilter",
     "LowPassFilter",
     "MedianFilter",
     "NormalizeFilter",
+    "PowerFilter",
     "RollingMeanFilter",
+    "RotationFilter",
+    "RoundFilter",
     "SavitzkyGolayFilter",
     "ScaleFilter",
+    "SignFilter",
     "TaringFilter",
+    "TrigFilter",
     "UnitFilter",
     "WrapFilter",
     "filter_adapter",