Add floating brick detection to LegoStructure

src/legogpt/data/lego_structure.py

@@ -1,8 +1,9 @@
 import re
 import warnings
-import numpy as np
 from dataclasses import dataclass
 
+import numpy as np
+
 from legogpt.data.lego_library import lego_library, dimensions_to_brick_id
 
 
@@ -84,6 +85,8 @@ class LegoStructure:
     """
 
     def __init__(self, bricks: list[LegoBrick], world_dim: int = 20):
+        self.world_dim = world_dim
+
         # Check if structure starts at ground level
         z0 = min(brick.z for brick in bricks)
         if z0 != 0:
@@ -115,6 +118,19 @@ class LegoStructure:
     def has_collisions(self) -> bool:
         return np.any(self.voxel_occupancy > 1)
 
+    @property
+    def has_floating_bricks(self) -> bool:
+        return any(self._is_floating(brick) for brick in self.bricks)
+
+    def _is_floating(self, brick: LegoBrick) -> bool:
+        if brick.z == 0:
+            return False  # Supported by ground
+        if np.any(self.voxel_occupancy[*brick.slice_2d, brick.z - 1]):
+            return False  # Supported from below
+        if brick.z != self.world_dim - 1 and np.any(self.voxel_occupancy[*brick.slice_2d, brick.z + 1]):
+            return False  # Supported from above
+        return True
+
     @classmethod
     def from_json(cls, lego_json: dict):
         bricks = [LegoBrick.from_json(v) for k, v in lego_json.items() if k.isdigit()]

tests/test_lego_structure.py

@@ -37,3 +37,13 @@ def test_lego_structure():
 def test_collision_check(brick_txt: str, has_collisions: bool):
     lego = LegoStructure.from_txt(brick_txt)
     assert lego.has_collisions == has_collisions
+
+
+@pytest.mark.parametrize(
+    'brick_txt,has_floating_bricks', [
+        ('2x6 (0,0,0)\n2x6 (2,0,0)\n', False),
+        ('2x6 (0,0,0)\n2x6 (2,0,1)\n', True),
+    ])
+def test_floating_check(brick_txt: str, has_floating_bricks: bool):
+    lego = LegoStructure.from_txt(brick_txt)
+    assert lego.has_floating_bricks == has_floating_bricks