Skip to content
86 changes: 86 additions & 0 deletions dimos/protocol/tf/buffer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
# Copyright 2025-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.


"""Bounded timestamp index for generated transforms; frame search stays in TF."""

from copy import deepcopy
import math
from typing import cast

from sortedcontainers import SortedDict # type: ignore[import-untyped]

from dimos.message_utils import time as message_time
from dimos.messages.geometry_msgs.msg import TransformStamped
from dimos.types.timestamped import to_human_readable


class TBuffer:
"""A bounded time index of generated transforms, keyed by integer nanoseconds."""

def __init__(self, buffer_size: float = 10.0) -> None:
if math.isnan(buffer_size) or buffer_size <= 0:
raise ValueError("buffer_size must be positive")
self.buffer_size = buffer_size
self._entries: SortedDict = SortedDict()

def __len__(self) -> int:
return len(self._entries)

def add(self, transform: TransformStamped) -> None:
stamp = message_time.to_nanoseconds(transform.header.stamp)
self._entries[stamp] = deepcopy(transform)
if math.isfinite(self.buffer_size):
oldest = self._entries.peekitem(-1)[0] - round(self.buffer_size * 1_000_000_000)
while self._entries and self._entries.peekitem(0)[0] < oldest:
self._entries.popitem(0)

def first(self) -> TransformStamped | None:
return deepcopy(self._entries.peekitem(0)[1]) if self._entries else None

def last(self) -> TransformStamped | None:
return deepcopy(self._entries.peekitem(-1)[1]) if self._entries else None

def get(
self, time_point: float | None = None, time_tolerance: float = 1.0
) -> TransformStamped | None:
if time_point is None:
return self.last()
stamp = message_time.to_nanoseconds(message_time.time_from_seconds(time_point))
if math.isnan(time_tolerance) or time_tolerance < 0:
raise ValueError("time_tolerance must be nonnegative")
index = self._entries.bisect_left(stamp)
candidates = [
self._entries.peekitem(i) for i in (index - 1, index) if 0 <= i < len(self._entries)
]
if not candidates:
return None
nearest, result = min(candidates, key=lambda item: (abs(item[0] - stamp), -item[0]))
if abs(nearest - stamp) / 1_000_000_000 > time_tolerance:
return None
return deepcopy(cast("TransformStamped", result))

def __str__(self) -> str:
first, last = self.first(), self.last()
if first is None or last is None:
return "TBuffer(empty)"
start, end = (
message_time.to_seconds(first.header.stamp),
message_time.to_seconds(last.header.stamp),
)
return (
f"TBuffer({first.header.frame_id} -> {first.child_frame_id}, "
f"{len(self)} msgs, {end - start:.2f}s "
f"[{to_human_readable(start)} - {to_human_readable(end)}])"
)
91 changes: 91 additions & 0 deletions dimos/protocol/tf/test_buffer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
# Copyright 2025-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.


"""Timestamp retention and ownership, independently of transport/graph migration."""

from dimos.messages.builtin_interfaces.msg import Time
from dimos.messages.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3
from dimos.messages.std_msgs.msg import Header
from dimos.protocol.tf.buffer import TBuffer


def test_buffer_retains_distinct_nanoseconds_and_owns_its_snapshots():
buffer = TBuffer()
first = TransformStamped(
header=Header(frame_id="a", stamp=Time(sec=1700000000, nanosec=123456789)),
child_frame_id="b",
transform=Transform(
translation=Vector3(x=0.0, y=0.0, z=0.0),
rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0),
),
)
second = TransformStamped(
header=Header(frame_id="a", stamp=Time(sec=1700000000, nanosec=123456790)),
child_frame_id="b",
transform=Transform(
translation=Vector3(x=0.0, y=0.0, z=0.0),
rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0),
),
)
buffer.add(first)
buffer.add(second)
first.header.stamp.nanosec = 0
assert len(buffer) == 2
stored_first = buffer.first()
assert stored_first is not None
assert stored_first.header.stamp.nanosec == 123456789
stored_first.header.stamp.nanosec = 1
unchanged = buffer.first()
assert unchanged is not None
assert unchanged.header.stamp.nanosec == 123456789
latest = buffer.get()
assert latest is not None
assert latest.header.stamp.nanosec == 123456790


def test_buffer_drops_out_of_order_data_older_than_retained_window():
buffer = TBuffer(buffer_size=1)
for sec in (10, 11, 1):
buffer.add(
TransformStamped(
header=Header(stamp=Time(sec=sec, nanosec=0), frame_id=""),
child_frame_id="",
transform=Transform(
translation=Vector3(x=0.0, y=0.0, z=0.0),
rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0),
),
)
)
assert len(buffer) == 2
first = buffer.first()
assert first is not None
assert first.header.stamp.sec == 10


def test_buffer_preserves_negative_stamp():
buffer = TBuffer()
buffer.add(
TransformStamped(
header=Header(stamp=Time(sec=-1, nanosec=500000000), frame_id=""),
child_frame_id="",
transform=Transform(
translation=Vector3(x=0.0, y=0.0, z=0.0),
rotation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0),
),
)
)
result = buffer.get(time_point=-0.5, time_tolerance=0)
assert result is not None
assert result.header.stamp == Time(sec=-1, nanosec=500000000)
Loading