fix: vectorize LineSegments3D encoding and decoding

dimos/msgs/nav_msgs/LineSegments3D.py

@@ -12,115 +12,119 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""LineSegments3D: collection of 3D line segments for graph edge visualization.
+"""LineSegments3D: 3D line segments with a per-segment weight.
 
-On the wire uses ``nav_msgs/Path`` — consecutive pose pairs form segments.
-Renders as ``rr.LineStrips3D`` with each segment as a separate strip.
+On the wire uses ``nav_msgs/Path``. Consecutive pose pairs form segments and
+``orientation.w`` carries the weight.
 """
 
 from __future__ import annotations
 
+import struct
 import time
 from typing import TYPE_CHECKING, BinaryIO
 
-from dimos_lcm.nav_msgs import Path as LCMPath
+import numpy as np
 
 from dimos.types.timestamped import Timestamped
 
 if TYPE_CHECKING:
+    from numpy.typing import ArrayLike, NDArray
     from rerun._baseclasses import Archetype
 
+# Path prefix after the 8 byte fingerprint: poses_length, header seq, stamp sec, stamp nsec, frame_id length.
+_PREFIX = struct.Struct(">iiiiI")
+# PoseStamped bytes before its frame_id text: header seq, stamp sec, stamp nsec, frame_id length.
+_POSE_HEAD = 16
+_POSE_DOUBLES = 7
+_POSE_TAIL = _POSE_DOUBLES * 8
 
-class LineSegments3D(Timestamped):
-    """Line segments for graph edge visualization.
 
-    Wire format: ``nav_msgs/Path`` — consecutive pose pairs are segments.
-    ``orientation.w`` encodes traversability: 1.0=traversable, 0.5=partial, 0.0=unreachable.
-    """
+class LineSegments3D(Timestamped):
+    """Line segments as an (N, 2, 3) array plus one weight per segment."""
 
     msg_name = "nav_msgs.LineSegments3D"
     ts: float
     frame_id: str
-    _segments: list[tuple[tuple[float, float, float], tuple[float, float, float]]]
-    _traversability: list[float]
+    segments: NDArray[np.float64]
+    weights: NDArray[np.float64]
 
     def __init__(
         self,
         ts: float | None = None,
         frame_id: str = "map",
-        segments: list[tuple[tuple[float, float, float], tuple[float, float, float]]] | None = None,
-        traversability: list[float] | None = None,
+        segments: ArrayLike | None = None,
+        weights: ArrayLike | None = None,
     ) -> None:
         self.frame_id = frame_id
         self.ts = time.time() if ts is None else ts
-        self._segments = segments or []
-        self._traversability = traversability or [1.0] * len(self._segments)
+        self.segments = np.asarray(
+            segments if segments is not None else np.empty((0, 2, 3)), dtype=np.float64
+        ).reshape(-1, 2, 3)
+        self.weights = (
+            np.ones(len(self.segments))
+            if weights is None
+            else np.asarray(weights, dtype=np.float64).reshape(-1)
+        )
 
     def lcm_encode(self) -> bytes:
         raise NotImplementedError("Encoded on C++ side")
 
     @classmethod
     def lcm_decode(cls, data: bytes | BinaryIO) -> LineSegments3D:
-        lcm_msg = LCMPath.lcm_decode(data)
-        header_ts = lcm_msg.header.stamp.sec + lcm_msg.header.stamp.nsec / 1e9
-        frame_id = lcm_msg.header.frame_id
-
-        segments = []
-        traversability = []
-        poses = lcm_msg.poses
-        for i in range(0, len(poses) - 1, 2):
-            p1, p2 = poses[i], poses[i + 1]
-            segments.append(
-                (
-                    (p1.pose.position.x, p1.pose.position.y, p1.pose.position.z),
-                    (p2.pose.position.x, p2.pose.position.y, p2.pose.position.z),
-                )
-            )
-            traversability.append(p1.pose.orientation.w)
+        """Read the Path payload through strided array views.
+
+        Every pose header must carry the same frame_id length, which is what the planner emits.
+        """
+        raw = data if isinstance(data, bytes) else data.read()
+        count, _, sec, nsec, frame_len = _PREFIX.unpack_from(raw, 8)
+        offset = 8 + _PREFIX.size
+        frame_id = raw[offset : offset + frame_len][:-1].decode("utf-8", "replace")
+        offset += frame_len
+        ts = sec + nsec / 1e9
+        if count == 0:
+            return cls(ts=ts, frame_id=frame_id)
+        if count % 2:
+            raise ValueError(f"LineSegments3D needs pose pairs, got {count} poses")
+        (pose_frame_len,) = struct.unpack_from(">I", raw, offset + _POSE_HEAD - 4)
+        stride = _POSE_HEAD + pose_frame_len + _POSE_TAIL
+        lens = np.ndarray(
+            (count,), dtype=">u4", buffer=raw, offset=offset + _POSE_HEAD - 4, strides=(stride,)
+        )
+        if len(raw) != offset + count * stride or not np.all(lens == pose_frame_len):
+            raise ValueError("LineSegments3D poses must share one frame_id length")
+        poses = np.ndarray(
+            (count, _POSE_DOUBLES),
+            dtype=">f8",
+            buffer=raw,
+            offset=offset + _POSE_HEAD + pose_frame_len,
+            strides=(stride, 8),
+        )
         return cls(
-            ts=header_ts, frame_id=frame_id, segments=segments, traversability=traversability
+            ts=ts,
+            frame_id=frame_id,
+            segments=poses[:, :3].astype(np.float64).reshape(-1, 2, 3),
+            weights=poses[0::2, 6].astype(np.float64),
         )
 
-    def to_rerun(
-        self,
-        z_offset: float = 1.7,
-        color: tuple[int, int, int, int] = (0, 255, 150, 255),
-        radii: float = 0.04,
-    ) -> Archetype:
-        """Render as ``rr.LineStrips3D`` — color-coded by traversability.
-
-        Green = traversable (reachable from robot), red = non-traversable.
-        """
+    def to_rerun(self, z_offset: float = 0.0, radii: float = 0.04) -> Archetype:
+        """Render as ``rr.LineStrips3D``, green to red by log-scale weight."""
         import rerun as rr
 
-        if not self._segments:
+        if len(self.segments) == 0:
             return rr.LineStrips3D([])
-
-        strips = []
-        colors = []
-        for idx, (p1, p2) in enumerate(self._segments):
-            strips.append(
-                [
-                    [p1[0], p1[1], p1[2] + z_offset],
-                    [p2[0], p2[1], p2[2] + z_offset],
-                ]
-            )
-            trav = self._traversability[idx] if idx < len(self._traversability) else 1.0
-            if trav >= 0.9:
-                colors.append((0, 220, 100, 200))  # green = fully traversable
-            elif trav >= 0.4:
-                colors.append((255, 180, 0, 200))  # yellow = partially traversable
-            else:
-                colors.append((255, 50, 50, 150))  # red = non-traversable
-
-        return rr.LineStrips3D(
-            strips,
-            colors=colors,
-            radii=[radii] * len(strips),
-        )
+        strips = self.segments.astype(np.float32)
+        strips[:, :, 2] += z_offset
+        log_w = np.log10(np.maximum(self.weights, 1e-6))
+        lo, hi = float(log_w.min()), float(log_w.max())
+        norm = (log_w - lo) / (hi - lo) if hi > lo else np.zeros_like(log_w)
+        r = (255 * norm).astype(np.uint8)
+        g = (255 * (1.0 - norm)).astype(np.uint8)
+        colors = np.column_stack([r, g, np.full_like(r, 60), np.full_like(r, 220)])
+        return rr.LineStrips3D(strips, colors=colors, radii=radii)
 
     def __len__(self) -> int:
-        return len(self._segments)
+        return len(self.segments)
 
     def __str__(self) -> str:
-        return f"LineSegments3D(frame_id='{self.frame_id}', segments={len(self._segments)})"
+        return f"LineSegments3D(frame_id='{self.frame_id}', segments={len(self.segments)})"

dimos/msgs/nav_msgs/test_LineSegments3D.py

@@ -0,0 +1,65 @@
+# Copyright 2026 Dimensional Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dimos_lcm.geometry_msgs.Point import Point
+from dimos_lcm.geometry_msgs.Pose import Pose
+from dimos_lcm.geometry_msgs.PoseStamped import PoseStamped
+from dimos_lcm.geometry_msgs.Quaternion import Quaternion
+from dimos_lcm.nav_msgs.Path import Path
+from dimos_lcm.std_msgs.Header import Header
+from dimos_lcm.std_msgs.Time import Time
+import numpy as np
+import pytest
+
+from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D
+
+
+def encode_edges(n_segments: int, frame_ids: list[str] | None = None) -> bytes:
+    """A Path of consecutive pose pairs with orientation.w carrying the weight."""
+    poses = [
+        PoseStamped(
+            header=Header(frame_id=frame_ids[i % len(frame_ids)] if frame_ids else "odom"),
+            pose=Pose(
+                position=Point(float(i), float(i) * 0.5, -1.0),
+                orientation=Quaternion(w=float(i // 2) * 0.1),
+            ),
+        )
+        for i in range(2 * n_segments)
+    ]
+    header = Header(stamp=Time(12, 500_000_000), frame_id="odom")
+    return Path(poses_length=len(poses), header=header, poses=poses).lcm_encode()
+
+
+def expected_segments(n_segments: int) -> np.ndarray:
+    i = np.arange(2 * n_segments, dtype=np.float64)
+    return np.stack([i, i * 0.5, np.full_like(i, -1.0)], axis=1).reshape(-1, 2, 3)
+
+
+def test_decode_matches_the_wire_layout() -> None:
+    msg = LineSegments3D.lcm_decode(encode_edges(50))
+    assert msg.frame_id == "odom"
+    assert msg.ts == 12.5
+    np.testing.assert_array_equal(msg.segments, expected_segments(50))
+    np.testing.assert_allclose(msg.weights, np.arange(50) * 0.1)
+
+    empty = LineSegments3D.lcm_decode(encode_edges(0))
+    assert len(empty) == 0
+    assert empty.segments.shape == (0, 2, 3)
+    assert empty.weights.shape == (0,)
+
+
+def test_mixed_frame_id_lengths_are_rejected() -> None:
+    raw = encode_edges(20, frame_ids=["odom", "map", "base_link"])
+    with pytest.raises(ValueError, match="frame_id length"):
+        LineSegments3D.lcm_decode(raw)

dimos/navigation/nav_3d/mls_planner/viz.py

@@ -88,28 +88,7 @@ def render_nodes(msg: PointCloud2) -> Archetype:
 
 
 def render_node_edges(msg: LineSegments3D) -> Archetype:
-    """Color each segment by its safe-adj weight on a log-scale green->red gradient."""
-    import rerun as rr
-
-    if not msg._segments:
-        return rr.LineStrips3D([])
-    weights = np.asarray(msg._traversability, dtype=np.float64)
-    log_w = np.log10(np.maximum(weights, 1e-6))
-    lo, hi = float(log_w.min()), float(log_w.max())
-    norm = (log_w - lo) / (hi - lo) if hi > lo else np.zeros_like(log_w)
-    r = (255 * norm).astype(np.uint8)
-    g = (255 * (1.0 - norm)).astype(np.uint8)
-    b = np.full_like(r, 60)
-    a = np.full_like(r, 220)
-    colors = np.column_stack([r, g, b, a])
-    strips = [
-        [
-            [p1[0], p1[1], p1[2] + _GRAPH_Z_LIFT],
-            [p2[0], p2[1], p2[2] + _GRAPH_Z_LIFT],
-        ]
-        for p1, p2 in msg._segments
-    ]
-    return rr.LineStrips3D(strips, colors=colors, radii=[0.01] * len(strips))
+    return msg.to_rerun(z_offset=_GRAPH_Z_LIFT, radii=0.01)
 
 
 def planner_visual_override(