Optimize loading meshes using numpy arrays

Andrej730 · Andrej730 · commit 299b5bcf6b32 · 2025-02-12T18:39:27.000+05:00
1) avoid calling geometry.xxx multiple times as each time it creates a new copy
2) use numpy arrays with correct data type, so Blender could perform buffer copy significantly reducing overhead
diff --git a/src/bonsai/bonsai/bim/import_ifc.py b/src/bonsai/bonsai/bim/import_ifc.py
@@ -24,6 +24,7 @@
 import traceback
 import mathutils
 import numpy as np
+import numpy.typing as npt
 import multiprocessing
 import ifcopenshell
 import ifcopenshell.geom
@@ -36,7 +37,7 @@
 import bonsai.tool as tool
 from bonsai.bim.ifc import IfcStore, IFC_CONNECTED_TYPE
 from bonsai.tool.loader import OBJECT_DATA_TYPE
-from typing import Dict, Union, Optional, Any
+from typing import Dict, Union, Optional, Any, Literal
 
 
 class MaterialCreator:
@@ -167,7 +168,7 @@ def get_empty_slot_index() -> int:
             if -1 in self.mesh["ios_material_ids"]:
                 material_to_slot[-1] = get_empty_slot_index()
 
-            material_index = [material_to_slot[mat_id] for mat_id in self.mesh["ios_material_ids"]]
+            material_index = np.array([material_to_slot[mat_id] for mat_id in self.mesh["ios_material_ids"]], dtype="I")
             self.mesh.polygons.foreach_set("material_index", material_index)
 
     def resolve_all_stylable_representation_items(
@@ -1030,7 +1031,7 @@ def create_mesh(
         self,
         element: ifcopenshell.entity_instance,
         shape: Union[ifcopenshell.geom.ShapeElementType, ifcopenshell.geom.ShapeType],
-        cartesian_point_offset=None,
+        cartesian_point_offset: Union[npt.NDArray[np.float64], Literal[False]] = None,
     ) -> bpy.types.Mesh:
         try:
             if hasattr(shape, "geometry"):
@@ -1046,32 +1047,25 @@ def create_mesh(
                 old_mesh.name = mesh_name + ".old"
             mesh = bpy.data.meshes.new(mesh_name)
 
+            verts = ifcopenshell.util.shape.get_vertices(geometry)
             if cartesian_point_offset is False:
-                verts = geometry.verts
                 mesh["has_cartesian_point_offset"] = False
             elif cartesian_point_offset is not None:
-                verts_array = np.array(geometry.verts)
-                offset = np.array([-cartesian_point_offset[0], -cartesian_point_offset[1], -cartesian_point_offset[2]])
-                offset_verts = verts_array + np.tile(offset, len(verts_array) // 3)
-                verts = offset_verts.tolist()
+                verts -= cartesian_point_offset
 
                 mesh["has_cartesian_point_offset"] = True
                 mesh["cartesian_point_offset"] = (
                     f"{cartesian_point_offset[0]},{cartesian_point_offset[1]},{cartesian_point_offset[2]}"
                 )
-            elif geometry.verts and tool.Loader.is_point_far_away(
-                (geometry.verts[0], geometry.verts[1], geometry.verts[2]), is_meters=True
-            ):
+            elif verts.size and tool.Loader.is_point_far_away(verts[0], is_meters=True):
                 # Shift geometry close to the origin based off that first vert it found
                 verts_array = np.array(geometry.verts)
-                offset = np.array([-geometry.verts[0], -geometry.verts[1], -geometry.verts[2]])
-                offset_verts = verts_array + np.tile(offset, len(verts_array) // 3)
-                verts = offset_verts.tolist()
+                offset = verts[0]
+                verts -= offset
 
                 mesh["has_cartesian_point_offset"] = True
-                mesh["cartesian_point_offset"] = f"{geometry.verts[0]},{geometry.verts[1]},{geometry.verts[2]}"
+                mesh["cartesian_point_offset"] = f"{offset[0]},{offset[1]},{offset[2]}"
             else:
-                verts = geometry.verts
                 mesh["has_cartesian_point_offset"] = False
 
             return tool.Loader.convert_geometry_to_mesh(
diff --git a/src/bonsai/bonsai/bim/module/project/operator.py b/src/bonsai/bonsai/bim/module/project/operator.py
@@ -1846,7 +1846,7 @@ def process_occurrence(self, shape: ShapeElementType) -> None:
                     max_slot_index += 1
                 material_to_slot[i] = slot_index
 
-            material_index = [(material_to_slot[i] if i != -1 else 0) for i in material_ids]
+            material_index = np.array([(material_to_slot[i] if i != -1 else 0) for i in material_ids], dtype="I")
 
             num_vertices = len(verts) // 3
             total_faces = len(faces)
diff --git a/src/bonsai/bonsai/tool/geometry.py b/src/bonsai/bonsai/tool/geometry.py
@@ -35,6 +35,7 @@
 import ifcopenshell.util.element
 import ifcopenshell.util.placement
 import ifcopenshell.util.representation
+import ifcopenshell.util.shape
 import ifcopenshell.util.shape_builder
 import ifcopenshell.util.system
 import ifcopenshell.util.unit
@@ -1756,13 +1757,9 @@ def import_item(cls, obj: bpy.types.Object) -> None:
             obj.matrix_world = rep_obj.matrix_world @ position
         else:
             geometry = tool.Loader.create_generic_shape(item)
+            verts = ifcopenshell.util.shape.get_vertices(geometry)
             if (cartesian_point_offset := cls.get_cartesian_point_offset(rep_obj)) is not None:
-                verts_array = np.array(geometry.verts)
-                offset = np.array([-cartesian_point_offset[0], -cartesian_point_offset[1], -cartesian_point_offset[2]])
-                offset_verts = verts_array + np.tile(offset, len(verts_array) // 3)
-                verts = offset_verts.tolist()
-            else:
-                verts = geometry.verts
+                verts -= cartesian_point_offset
             tool.Loader.convert_geometry_to_mesh(geometry, obj.data, verts=verts)
 
             if ios_materials := list(obj.data["ios_materials"]):
diff --git a/src/bonsai/bonsai/tool/loader.py b/src/bonsai/bonsai/tool/loader.py
@@ -35,6 +35,7 @@
 import bonsai.bim.import_ifc
 import numpy as np
 import numpy.typing as npt
+from ifcopenshell.util.shape_builder import np_to_4d
 from mathutils import Vector, Matrix
 from pathlib import Path
 from typing import Union, Any, Optional
@@ -634,7 +635,7 @@ def is_point_far_away(
         limit = cls.settings.distance_limit
         limit = limit if is_meters else (limit / cls.unit_scale)
         coords = getattr(point, "Coordinates", point)
-        return abs(coords[0]) > limit or abs(coords[1]) > limit or abs(coords[2]) > limit
+        return any(abs(c) > limit for c in coords)
 
     @classmethod
     def is_element_far_away(cls, element: ifcopenshell.entity_instance) -> bool:
@@ -879,11 +880,12 @@ def get_offset_point(cls, ifc_file: ifcopenshell.file) -> Union[npt.NDArray[np.f
             if not shape:
                 continue
             mat = ifcopenshell.util.shape.get_shape_matrix(shape)
-            point = mat @ np.array((shape.geometry.verts[0], shape.geometry.verts[1], shape.geometry.verts[2], 1.0))
+            verts = ifcopenshell.util.shape.get_vertices(shape.geometry)
+            point = (mat @ np_to_4d(verts[0]))[:3]
             if cls.is_point_far_away(point, is_meters=True):
                 # Arbitrary origins should be to the nearest millimeter.
                 # Anything more precise is just ridiculous from a practical surveying perspective.
-                return [round(float(p), 3) / cls.unit_scale for p in point[:3]]
+                return np.array([round(float(p), 3) / cls.unit_scale for p in point])
             break
 
     @classmethod
@@ -960,14 +962,19 @@ def convert_geometry_to_mesh(
         cls,
         geometry: ifcopenshell.geom.ShapeType,
         mesh: bpy.types.Mesh,
-        verts: Optional[list[float]] = None,
+        verts: Optional[npt.NDArray[np.float64]] = None,
         *,
         load_indexed_maps=True,
     ) -> bpy.types.Mesh:
+        """
+        :param verts: Numpy array of shape (n, 3).
+        """
         if verts is None:
-            verts = geometry.verts
-        if geometry.faces:
-            num_vertices = len(verts) // 3
+            verts = ifcopenshell.util.shape.get_vertices(geometry)
+        faces = ifcopenshell.util.shape.get_faces(geometry)
+        total_faces: int
+        if total_faces := faces.shape[0]:
+            num_vertices: int = verts.shape[0]
 
             # See bug 3546
             # ios_edges holds true edges that aren't triangulated.
@@ -978,34 +985,35 @@ def convert_geometry_to_mesh(
             mesh["ios_item_ids"] = ios_item_ids
 
             mesh.vertices.add(num_vertices)
-            mesh.vertices.foreach_set("co", verts)
+            mesh.vertices.foreach_set("co", verts.ravel().astype("f"))
 
             is_triangulated = True
+            num_vertex_indices = faces.size
             if is_triangulated:
-                total_faces = len(geometry.faces)
-                num_vertex_indices = len(geometry.faces)
-                loop_start = range(0, total_faces, 3)
-                num_loops = total_faces // 3
-                loop_total = [3] * num_loops
+                loop_start = np.arange(0, num_vertex_indices, 3, dtype="I")
+                loop_total = np.full(total_faces, 3, dtype="I")
+                use_smooth = np.zeros(num_vertex_indices, dtype="?")
 
                 mesh.loops.add(num_vertex_indices)
-                mesh.loops.foreach_set("vertex_index", geometry.faces)
-                mesh.polygons.add(num_loops)
+                mesh.loops.foreach_set("vertex_index", faces.ravel().astype("I"))
+                mesh.polygons.add(total_faces)
                 mesh.polygons.foreach_set("loop_start", loop_start)
                 mesh.polygons.foreach_set("loop_total", loop_total)
-                mesh.polygons.foreach_set("use_smooth", [0] * total_faces)
+                mesh.polygons.foreach_set("use_smooth", use_smooth)
             else:
+                # TODO: optimize using correct numpy array types.
                 faces_array = np.array(geometry.faces, dtype=object)
-                loop_total = tuple(len(face) for face in faces_array)
+                loop_total = np.array(tuple(len(face) for face in faces_array), dtype="I")
                 loop_start = np.cumsum((0,) + loop_total)[:-1]
                 vertex_index = np.concatenate(faces_array)
+                use_smooth = np.zeros(num_vertex_indices, dtype="?")
 
                 mesh.loops.add(len(vertex_index))
                 mesh.loops.foreach_set("vertex_index", vertex_index)
                 mesh.polygons.add(len(loop_start))
                 mesh.polygons.foreach_set("loop_start", loop_start)
                 mesh.polygons.foreach_set("loop_total", loop_total)
-                mesh.polygons.foreach_set("use_smooth", [0] * len(geometry.faces))
+                mesh.polygons.foreach_set("use_smooth", use_smooth)
 
             mesh.update()
 
@@ -1020,11 +1028,8 @@ def convert_geometry_to_mesh(
             tool.Blender.Attribute.fill_attribute(mesh, "ios_item_ids", "FACE", "INT", ios_item_ids)
             tool.Blender.Attribute.fill_attribute(mesh, "ios_material_ids", "FACE", "INT", geometry.material_ids)
         else:
-            e = geometry.edges
-            v = verts
-            vertices = [[v[i], v[i + 1], v[i + 2]] for i in range(0, len(v), 3)]
-            edges = [[e[i], e[i + 1]] for i in range(0, len(e), 2)]
-            mesh.from_pydata(vertices, edges, [])
+            edges = ifcopenshell.util.shape.get_edges(geometry)
+            mesh.from_pydata(verts.tolist(), edges.tolist(), [])
             # TODO: remove error handling after we update build in Bonsai.
             try:
                 edges_item_ids = ifcopenshell.util.shape.get_edges_representation_item_ids(geometry).tolist()
@@ -1034,8 +1039,8 @@ def convert_geometry_to_mesh(
             tool.Blender.Attribute.fill_attribute(mesh, "ios_edges_item_ids", "EDGE", "INT", edges_item_ids)
             tool.Blender.Attribute.fill_attribute(mesh, "ios_material_ids", "EDGE", "INT", geometry.material_ids)
 
-        mesh["ios_materials"] = [m.instance_id() for m in geometry.materials]
-        mesh["ios_material_ids"] = geometry.material_ids
+        mesh["ios_materials"] = [m.instance_id() for m in ifcopenshell.util.shape.get_shape_material_styles(geometry)]
+        mesh["ios_material_ids"] = ifcopenshell.util.shape.get_faces_material_style_ids(geometry).tolist()
         return mesh
 
     @classmethod
diff --git a/src/bonsai/test/tool/test_model.py b/src/bonsai/test/tool/test_model.py
@@ -614,7 +614,7 @@ def test_assign_material_to_representation_that_has_2_items_and_1_item_has_a_sty
         tool.Geometry._reload_representation(obj)
 
         def get_material_indices(mesh: bpy.types.Mesh) -> np.ndarray:
-            buffer = np.empty(len(mesh.polygons), dtype=np.int32)
+            buffer = np.empty(len(mesh.polygons), dtype="I")
             mesh.polygons.foreach_get("material_index", buffer)
             return buffer
 
