refactor: support generated occupancy values in existing grid algorithms

.github/workflows/message-codegen.yml

@@ -19,6 +19,7 @@ on:
       - dimos/msgs/**
       - dimos/protocol/tf/**
       - dimos/memory/**
+      - dimos/mapping/occupancy/**
       - dimos/protocol/cdr_mcap.py
       - docs/usage/transforms.md
       - dimos/hardware/sensors/camera/**
@@ -221,3 +222,8 @@ jobs:
           .native-review/bin/python -I -c 'import dimos.msgs.occupancy as values; import pytest; import sys; assert "site-packages" in str(values.__file__); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/msgs/test_occupancy.py"
           .native-review/bin/python -I -c 'import dimos.msgs.nav_msgs.OccupancyGrid; import pytest; import sys; raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/msgs/nav_msgs/test_OccupancyGrid.py" -k block_max_reduce
           MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/msgs/occupancy.py
+      - name: Verify occupancy algorithms preserve generated and existing value contracts
+        if: runner.os == 'Linux'
+        run: |
+          .native-review/bin/python -I -c 'from dimos.mapping.occupancy import operations, gradient, inflation, extrude_occupancy; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (operations, gradient, inflation, extrude_occupancy)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/mapping/occupancy/test_generated_operations.py"
+          MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/mapping/occupancy/operations.py dimos/mapping/occupancy/gradient.py dimos/mapping/occupancy/inflation.py dimos/mapping/occupancy/extrude_occupancy.py

dimos/mapping/occupancy/extrude_occupancy.py

@@ -16,13 +16,14 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid
+from dimos.mapping.occupancy.operations import _grid_values, _GridT
+from dimos.msgs.occupancy import CostValues
 
 # Rectangle type: (x, y, width, height)
 Rect = tuple[int, int, int, int]
 
 
-def identify_convex_shapes(occupancy_grid: OccupancyGrid) -> list[Rect]:
+def identify_convex_shapes(occupancy_grid: _GridT) -> list[Rect]:
     """Identify occupied zones and decompose them into convex rectangles.
 
     This function finds all occupied cells in the occupancy grid and
@@ -36,7 +37,7 @@ def identify_convex_shapes(occupancy_grid: OccupancyGrid) -> list[Rect]:
     Returns:
         List of rectangles as (x, y, width, height) tuples in grid coords.
     """
-    grid = occupancy_grid.grid
+    grid = _grid_values(occupancy_grid)
 
     # Create binary mask of occupied cells (treat UNKNOWN as OCCUPIED)
     occupied_mask = ((grid == CostValues.OCCUPIED) | (grid == CostValues.UNKNOWN)).astype(
@@ -166,7 +167,7 @@ def _largest_rect_in_histogram(
 
 
 def generate_mujoco_scene(
-    occupancy_grid: OccupancyGrid,
+    occupancy_grid: _GridT,
 ) -> str:
     """Generate a MuJoCo scene XML from an occupancy grid.
 
@@ -184,9 +185,9 @@ def generate_mujoco_scene(
     # Get rectangles from the occupancy grid
     rectangles = identify_convex_shapes(occupancy_grid)
 
-    resolution = occupancy_grid.resolution
-    origin_x = occupancy_grid.origin.position.x
-    origin_y = occupancy_grid.origin.position.y
+    resolution = occupancy_grid.info.resolution
+    origin_x = occupancy_grid.info.origin.position.x
+    origin_y = occupancy_grid.info.origin.position.y
 
     # Build XML
     xml_lines = [
@@ -208,7 +209,7 @@ def generate_mujoco_scene(
     # Add each rectangle as a box geom
     for i, (gx, gy, gw, gh) in enumerate(rectangles):
         # Convert grid coordinates to world coordinates
-        # Grid origin is top-left, world origin is at occupancy_grid.origin
+        # Grid origin is top-left, world origin is at occupancy_grid.info.origin
         # gx, gy are in grid cells, need to convert to meters
         world_x = origin_x + (gx + gw / 2) * resolution
         world_y = origin_y + (gy + gh / 2) * resolution

dimos/mapping/occupancy/gradient.py

@@ -17,7 +17,8 @@
 import numpy as np
 from scipy import ndimage
 
-from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid
+from dimos.mapping.occupancy.operations import _grid_values, _GridT, _with_grid
+from dimos.msgs.occupancy import CostValues
 
 if TYPE_CHECKING:
     from numpy.typing import NDArray
@@ -27,8 +28,8 @@
 
 
 def gradient(
-    occupancy_grid: OccupancyGrid, obstacle_threshold: int = 50, max_distance: float = 2.0
-) -> OccupancyGrid:
+    occupancy_grid: _GridT, obstacle_threshold: int = 50, max_distance: float = 2.0
+) -> _GridT:
     """Create a gradient OccupancyGrid for path planning.
 
     Creates a gradient where free space has value 0 and values increase near obstacles.
@@ -49,19 +50,19 @@ def gradient(
     """
 
     # Remember which cells are unknown
-    unknown_mask = occupancy_grid.grid == CostValues.UNKNOWN
+    unknown_mask = _grid_values(occupancy_grid) == CostValues.UNKNOWN
 
     # Create binary obstacle map
     # Consider cells >= threshold as obstacles (1), everything else as free (0)
     # Unknown cells are not considered obstacles for distance calculation
-    obstacle_map = (occupancy_grid.grid >= obstacle_threshold).astype(np.float32)
+    obstacle_map = (_grid_values(occupancy_grid) >= obstacle_threshold).astype(np.float32)
 
     # Compute distance transform (distance to nearest obstacle in cells)
     # Unknown cells are treated as if they don't exist for distance calculation
     distance_cells = cast("NDArray[np.float64]", ndimage.distance_transform_edt(1 - obstacle_map))
 
     # Convert to meters and clip to max distance
-    distance_meters = np.clip(distance_cells * occupancy_grid.resolution, 0, max_distance)
+    distance_meters = np.clip(distance_cells * occupancy_grid.info.resolution, 0, max_distance)
 
     # Invert and scale to 0-100 range
     # Far from obstacles (max_distance) -> 0
@@ -78,20 +79,14 @@ def gradient(
     gradient_data[unknown_mask] = CostValues.UNKNOWN
 
     # Create new OccupancyGrid with gradient
-    gradient_grid = OccupancyGrid(
-        grid=gradient_data,
-        resolution=occupancy_grid.resolution,
-        origin=occupancy_grid.origin,
-        frame_id=occupancy_grid.frame_id,
-        ts=occupancy_grid.ts,
-    )
+    gradient_grid = _with_grid(occupancy_grid, gradient_data)
 
     return gradient_grid
 
 
 def voronoi_gradient(
-    occupancy_grid: OccupancyGrid, obstacle_threshold: int = 50, max_distance: float = 2.0
-) -> OccupancyGrid:
+    occupancy_grid: _GridT, obstacle_threshold: int = 50, max_distance: float = 2.0
+) -> _GridT:
     """Create a Voronoi-based gradient OccupancyGrid for path planning.
 
     Unlike the regular gradient which can result in suboptimal paths in narrow
@@ -122,23 +117,17 @@ def voronoi_gradient(
         - 100: At obstacles
     """
     # Remember which cells are unknown
-    unknown_mask = occupancy_grid.grid == CostValues.UNKNOWN
+    unknown_mask = _grid_values(occupancy_grid) == CostValues.UNKNOWN
 
     # Create binary obstacle map
-    obstacle_map = (occupancy_grid.grid >= obstacle_threshold).astype(np.float32)
+    obstacle_map = (_grid_values(occupancy_grid) >= obstacle_threshold).astype(np.float32)
 
     # Check if there are any obstacles
     if not np.any(obstacle_map):
         # No obstacles - everything is free
-        gradient_data = np.zeros_like(occupancy_grid.grid, dtype=np.int8)
+        gradient_data = np.zeros_like(_grid_values(occupancy_grid), dtype=np.int8)
         gradient_data[unknown_mask] = CostValues.UNKNOWN
-        return OccupancyGrid(
-            grid=gradient_data,
-            resolution=occupancy_grid.resolution,
-            origin=occupancy_grid.origin,
-            frame_id=occupancy_grid.frame_id,
-            ts=occupancy_grid.ts,
-        )
+        return _with_grid(occupancy_grid, gradient_data)
 
     # Label connected obstacle regions (clusters)
     # This groups all cells of the same wall/obstacle together
@@ -192,7 +181,7 @@ def voronoi_gradient(
     gradient_values[obstacle_map > 0] = CostValues.OCCUPIED
 
     # Apply max_distance clipping - cells beyond max_distance from obstacles get cost 0
-    max_distance_cells = max_distance / occupancy_grid.resolution
+    max_distance_cells = max_distance / occupancy_grid.info.resolution
     gradient_values[distance_cells > max_distance_cells] = 0
 
     # Convert to int8
@@ -201,10 +190,4 @@ def voronoi_gradient(
     # Preserve unknown cells
     gradient_data[unknown_mask] = CostValues.UNKNOWN
 
-    return OccupancyGrid(
-        grid=gradient_data,
-        resolution=occupancy_grid.resolution,
-        origin=occupancy_grid.origin,
-        frame_id=occupancy_grid.frame_id,
-        ts=occupancy_grid.ts,
-    )
+    return _with_grid(occupancy_grid, gradient_data)

dimos/mapping/occupancy/inflation.py

@@ -15,21 +15,22 @@
 import numpy as np
 from scipy import ndimage
 
-from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid
+from dimos.mapping.occupancy.operations import _grid_values, _GridT, _with_grid
+from dimos.msgs.occupancy import CostValues
 
 
-def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGrid:
+def simple_inflate(occupancy_grid: _GridT, radius: float) -> _GridT:
     """Inflate obstacles by a given radius (binary inflation).
     Args:
         radius: Inflation radius in meters
     Returns:
         New OccupancyGrid with inflated obstacles
     """
     # Convert radius to grid cells
-    cell_radius = int(np.ceil(radius / occupancy_grid.resolution))
+    cell_radius = int(np.ceil(radius / occupancy_grid.info.resolution))
 
     # Get grid as numpy array
-    grid_array = occupancy_grid.grid
+    grid_array = _grid_values(occupancy_grid)
 
     # Create circular kernel for binary inflation
     y, x = np.ogrid[-cell_radius : cell_radius + 1, -cell_radius : cell_radius + 1]
@@ -44,10 +45,4 @@ def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGri
     result_grid[inflated] = CostValues.OCCUPIED
 
     # Create new OccupancyGrid with inflated data using numpy constructor
-    return OccupancyGrid(
-        grid=result_grid,
-        resolution=occupancy_grid.resolution,
-        origin=occupancy_grid.origin,
-        frame_id=occupancy_grid.frame_id,
-        ts=occupancy_grid.ts,
-    )
+    return _with_grid(occupancy_grid, result_grid)

dimos/mapping/occupancy/operations.py

@@ -12,15 +12,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from copy import deepcopy
+from typing import TypeVar
+
 import numpy as np
+from numpy.typing import NDArray
 from scipy import ndimage
 
-from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid
+from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid
+from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid
+from dimos.msgs.occupancy import CostValues, occupancy_view
+
+_GridT = TypeVar("_GridT", OccupancyGrid, CdrOccupancyGrid)
+
+
+def _grid_values(message: _GridT) -> NDArray[np.int8]:
+    if isinstance(message, CdrOccupancyGrid):
+        return occupancy_view(message)
+    return message.grid
+
+
+def _with_grid(message: _GridT, cells: NDArray[np.int8]) -> _GridT:
+    if isinstance(message, CdrOccupancyGrid):
+        if cells.shape != (message.info.height, message.info.width):
+            raise ValueError("Replacement occupancy cells must retain map dimensions")
+        result = deepcopy(message)
+        result.data = np.array(cells, dtype=np.int8, copy=True).reshape(-1)
+        return result
+    return OccupancyGrid(
+        grid=cells,
+        resolution=message.resolution,
+        origin=message.origin,
+        frame_id=message.frame_id,
+        ts=message.ts,
+    )
 
 
-def smooth_occupied(
-    occupancy_grid: OccupancyGrid, min_neighbor_fraction: float = 0.4
-) -> OccupancyGrid:
+def smooth_occupied(occupancy_grid: _GridT, min_neighbor_fraction: float = 0.4) -> _GridT:
     """Smooth occupied zones by removing unsupported protrusions.
 
     Removes occupied cells that don't have sufficient neighboring occupied
@@ -33,7 +61,7 @@ def smooth_occupied(
     Returns:
         New OccupancyGrid with smoothed occupied zones
     """
-    grid_array = occupancy_grid.grid
+    grid_array = _grid_values(occupancy_grid)
     occupied_mask = grid_array >= CostValues.OCCUPIED
 
     # Count occupied neighbors for each cell (8-connectivity).
@@ -49,16 +77,10 @@ def smooth_occupied(
     result_grid = grid_array.copy()
     result_grid[unsupported] = CostValues.FREE
 
-    return OccupancyGrid(
-        grid=result_grid,
-        resolution=occupancy_grid.resolution,
-        origin=occupancy_grid.origin,
-        frame_id=occupancy_grid.frame_id,
-        ts=occupancy_grid.ts,
-    )
+    return _with_grid(occupancy_grid, result_grid)
 
 
-def overlay_occupied(base: OccupancyGrid, overlay: OccupancyGrid) -> OccupancyGrid:
+def overlay_occupied(base: _GridT, overlay: _GridT) -> _GridT:
     """Overlay occupied zones from one grid onto another.
 
     Marks cells as occupied in the base grid wherever they are occupied
@@ -70,19 +92,13 @@ def overlay_occupied(base: OccupancyGrid, overlay: OccupancyGrid) -> OccupancyGr
     Returns:
         New OccupancyGrid with combined occupied zones
     """
-    if base.grid.shape != overlay.grid.shape:
+    if _grid_values(base).shape != _grid_values(overlay).shape:
         raise ValueError(
-            f"Grid shapes must match: base {base.grid.shape} vs overlay {overlay.grid.shape}"
+            f"Grid shapes must match: base {_grid_values(base).shape} vs overlay {_grid_values(overlay).shape}"
         )
 
-    result_grid = base.grid.copy()
-    overlay_occupied_mask = overlay.grid >= CostValues.OCCUPIED
+    result_grid = _grid_values(base).copy()
+    overlay_occupied_mask = _grid_values(overlay) >= CostValues.OCCUPIED
     result_grid[overlay_occupied_mask] = CostValues.OCCUPIED
 
-    return OccupancyGrid(
-        grid=result_grid,
-        resolution=base.resolution,
-        origin=base.origin,
-        frame_id=base.frame_id,
-        ts=base.ts,
-    )
+    return _with_grid(base, result_grid)