Add to_ldr function to LegoStructure

src/legogpt/data/lego_library.json

@@ -3,66 +3,56 @@
     "height": 2,
     "width": 4,
     "mass": 0.00216,
-    "inventory": 100000
+    "inventory": 100000,
+    "partID": "3001.DAT"
   },
   "3": {
     "height": 2,
     "width": 6,
     "mass": 0.00323,
-    "inventory": 100000
+    "inventory": 100000,
+    "partID": "2456.DAT"
   },
   "4": {
     "height": 1,
     "width": 8,
     "mass": 0.00303,
-    "inventory": 100000
+    "inventory": 100000,
+    "partID": "3008.DAT"
   },
   "5": {
     "height": 1,
     "width": 4,
     "mass": 0.00157,
-    "inventory": 100000
+    "inventory": 100000,
+    "partID": "3010.DAT"
   },
   "6": {
     "height": 1,
     "width": 6,
     "mass": 0.00228,
-    "inventory": 100000
-  },
-  "7": {
-    "height": 1,
-    "width": 4,
-    "mass": 0.00157,
-    "inventory": 100000
-  },
-  "8": {
-    "height": 1,
-    "width": 4,
-    "mass": 0.00157,
-    "inventory": 100000
+    "inventory": 100000,
+    "partID": "3009.DAT"
   },
   "9": {
     "height": 1,
     "width": 2,
     "mass": 0.00081,
-    "inventory": 100000
+    "inventory": 100000,
+    "partID": "3004.DAT"
   },
   "10": {
     "height": 1,
     "width": 1,
     "mass": 0.00043,
-    "inventory": 100000
-  },
-  "11": {
-    "height": 1,
-    "width": 2,
-    "mass": 0.00081,
-    "inventory": 100000
+    "inventory": 100000,
+    "partID": "3005.DAT"
   },
   "12": {
     "height": 2,
     "width": 2,
     "mass": 0.00115,
-    "inventory": 100000
+    "inventory": 100000,
+    "partID": "3003.DAT"
   }
 }

src/legogpt/data/lego_library.py

@@ -6,11 +6,17 @@ with open(Path(__file__).parent / 'lego_library.json') as f:
 
 max_brick_dimension = max(max(properties['height'], properties['width']) for properties in lego_library.values())
 
-_dimensions_to_brick_id_dict = {}
-for brick_id, properties in lego_library.items():
-    key = (properties['height'], properties['width'])
-    if key not in _dimensions_to_brick_id_dict.keys():
-        _dimensions_to_brick_id_dict[key] = int(brick_id)
+
+def _make_dimensions_to_brick_id_dict() -> dict:
+    result = {}
+    for brick_id, properties in lego_library.items():
+        key = (properties['height'], properties['width'])
+        if key not in result.keys():
+            result[key] = int(brick_id)
+    return result
+
+
+_dimensions_to_brick_id_dict = _make_dimensions_to_brick_id_dict()
 
 
 def dimensions_to_brick_id(h: int, w: int):
@@ -20,3 +26,10 @@ def dimensions_to_brick_id(h: int, w: int):
         return _dimensions_to_brick_id_dict[(h, w)]
     except KeyError:
         raise ValueError(f'No brick ID for brick of dimensions: {h}x{w}')
+
+
+def brick_id_to_part_id(brick_id: int) -> str:
+    """
+    Returns the part ID of the given brick, which is the ID of the brick model used in LDraw files.
+    """
+    return lego_library[str(brick_id)]['partID']

src/legogpt/data/lego_structure.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 import numpy as np
 
-from legogpt.data.lego_library import lego_library, dimensions_to_brick_id
+from legogpt.data.lego_library import lego_library, dimensions_to_brick_id, brick_id_to_part_id
 from legogpt.stability_analysis.stability_analysis import stability_score
 
 
@@ -23,6 +23,10 @@ class LegoBrick:
     def brick_id(self) -> int:
         return dimensions_to_brick_id(self.h, self.w)
 
+    @property
+    def part_id(self) -> str:
+        return brick_id_to_part_id(self.brick_id)
+
     @property
     def ori(self) -> int:
         return 1 if self.h > self.w else 0
@@ -54,6 +58,14 @@ class LegoBrick:
     def to_txt(self) -> str:
         return f'{self.h}x{self.w} ({self.x},{self.y},{self.z})\n'
 
+    def to_ldr(self, base_height: float = 0) -> str:
+        x = (self.x + self.h * 0.5) * 20
+        z = (self.y + self.w * 0.5) * 20
+        y = (self.z + base_height) * -24
+        line = f'1 115 {x} {y} {z} 0 0 1 0 1 0 -1 0 0 {self.part_id}\n'
+        step_line = '0 STEP\n'
+        return line + step_line
+
     @classmethod
     def from_json(cls, brick_json: dict):
         properties = lego_library[str(brick_json['brick_id'])]
@@ -111,6 +123,9 @@ class LegoStructure:
     def to_txt(self) -> str:
         return ''.join([brick.to_txt() for brick in self.bricks])
 
+    def to_ldr(self) -> str:
+        return ''.join([brick.to_ldr() for brick in self.bricks])
+
     def add_brick(self, brick: LegoBrick) -> None:
         self.bricks.append(brick)
         self.voxel_occupancy[brick.slice] += 1

tests/test_lego_structure.py

@@ -22,11 +22,14 @@ def test_lego_structure():
         '1': {'brick_id': 3, 'x': 0, 'y': 0, 'z': 0, 'ori': 0},
         '2': {'brick_id': 3, 'x': 2, 'y': 0, 'z': 0, 'ori': 0},
     }
+    lego_ldr = '1 115 20.0 0 60.0 0 0 1 0 1 0 -1 0 0 2456.DAT\n0 STEP\n' \
+               '1 115 60.0 0 60.0 0 0 1 0 1 0 -1 0 0 2456.DAT\n0 STEP\n'
 
     for lego in [LegoStructure.from_json(lego_json), LegoStructure.from_txt(lego_txt)]:
         assert len(lego) == 2
         assert lego.to_json() == lego_json
         assert lego.to_txt() == lego_txt
+        assert lego.to_ldr() == lego_ldr
 
 
 @pytest.mark.parametrize(