From 62bd495a9efb9a348bdaa14a1f59597d4c53df00 Mon Sep 17 00:00:00 2001
From: Bruno Postle <bruno@postle.net>
Date: Sun, 23 Feb 2025 13:29:33 +0000
Subject: [PATCH] Create transform orientation slots #6128

Blender supports named UCS coordinate systems. Refreshing the spatial
decomposition panel now creates orientation slots for each unique
spatial element orientation.

Additionally, setting the default spatial container selects an
appropriate orientation slot. eg. if a building is rotated, selecting to
work in a storey will set the orientation to the building orientation.

These slots are only created/updated in the spatial decomposition panel,
maybe they should also be created when loading the file and other
places.

The blender API doesn't currently allow listing custom orientation slots
for deletion, so if a spatial element is deleted or renamed, the old
orientation slot will persist for the rest of the session.
---
 .../bonsai/bim/module/spatial/operator.py     |  2 +
 src/bonsai/bonsai/core/spatial.py             |  9 ++++
 src/bonsai/bonsai/tool/spatial.py             | 46 +++++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/src/bonsai/bonsai/bim/module/spatial/operator.py b/src/bonsai/bonsai/bim/module/spatial/operator.py
index dc68d95a8a..c3016f4cee 100644
--- a/src/bonsai/bonsai/bim/module/spatial/operator.py
+++ b/src/bonsai/bonsai/bim/module/spatial/operator.py
@@ -239,6 +239,7 @@ class ImportSpatialDecomposition(bpy.types.Operator):
 
     def execute(self, context):
         core.import_spatial_decomposition(tool.Spatial)
+        core.create_orientation_slots(ifc=tool.Ifc, spatial=tool.Spatial)
         return {"FINISHED"}
 
 
@@ -345,6 +346,7 @@ class SetDefaultContainer(bpy.types.Operator):
 
     def execute(self, context):
         core.set_default_container(tool.Spatial, container=tool.Ifc.get().by_id(self.container))
+        core.set_orientation_slot(spatial=tool.Spatial, element=tool.Ifc.get().by_id(self.container))
         return {"FINISHED"}
 
 
diff --git a/src/bonsai/bonsai/core/spatial.py b/src/bonsai/bonsai/core/spatial.py
index eff50ec2a3..04ae08ad6f 100644
--- a/src/bonsai/bonsai/core/spatial.py
+++ b/src/bonsai/bonsai/core/spatial.py
@@ -121,6 +121,15 @@ def import_spatial_decomposition(spatial: tool.Spatial) -> None:
     spatial.import_spatial_decomposition()
 
 
+def create_orientation_slots(ifc: tool.Ifc, spatial: tool.Spatial) -> None:
+    elements = ifc.get().by_type("IfcSpatialElement")
+    spatial.create_orientation_slots(elements, use=False)
+
+
+def set_orientation_slot(spatial: tool.Spatial, element: ifcopenshell.entity_instance) -> None:
+    spatial.create_orientation_slots([element], use=True)
+
+
 def edit_container_attributes(spatial: tool.Spatial, entity: ifcopenshell.entity_instance) -> None:
     spatial.edit_container_attributes(entity)
     spatial.import_spatial_decomposition()
diff --git a/src/bonsai/bonsai/tool/spatial.py b/src/bonsai/bonsai/tool/spatial.py
index f2c784c464..9c0c3b758d 100644
--- a/src/bonsai/bonsai/tool/spatial.py
+++ b/src/bonsai/bonsai/tool/spatial.py
@@ -38,6 +38,7 @@ import bonsai.core.geometry
 import bonsai.core.unit
 import bonsai.tool as tool
 import json
+import numpy as np
 from math import pi
 from mathutils import Vector, Matrix
 from shapely import Polygon
@@ -479,6 +480,51 @@ class Spatial(bonsai.core.tool.Spatial):
             for child in children or []:
                 cls.import_spatial_element(child, level_index + 1)
 
+    @classmethod
+    def create_orientation_slots(cls, elements: list[ifcopenshell.entity_instance], use: bool = False) -> None:
+        active_slot = bpy.context.scene.transform_orientation_slots[0]
+        initial_orientation = active_slot.type
+
+        # stash selected objects
+        active_object = bpy.context.view_layer.objects.active
+        selected_objects = list(bpy.context.view_layer.objects.selected)
+
+        # bpy.ops.transform.create_orientation() requires a dummy object
+        bpy.ops.object.empty_add(type="PLAIN_AXES")
+
+        for element in elements:
+            placement = element.ObjectPlacement
+            combined_matrix = ifcopenshell.util.placement.get_local_placement(placement)[:3, :3]
+            if np.allclose(combined_matrix, np.eye(3), atol=1e-6):
+                # this element has global orientation
+                active_slot.type = "GLOBAL"
+                continue
+            elif element.Decomposes and hasattr(element.Decomposes[0].RelatingObject, "ObjectPlacement"):
+                # this element is part of a decomposition
+                parent_placement = element.Decomposes[0].RelatingObject.ObjectPlacement
+                parent_matrix = ifcopenshell.util.placement.get_local_placement(parent_placement)[:3, :3]
+                if np.allclose(combined_matrix, parent_matrix, atol=1e-6):
+                    # this element has the same orientation as its parent
+                    cls.create_orientation_slots(elements=[element.Decomposes[0].RelatingObject], use=use)
+                    continue
+            # this element has a unique orientation
+            orientation_name = element.is_a() + "/" + element.Name
+            bpy.ops.transform.create_orientation(name=orientation_name, overwrite=True)
+            active_slot.type = orientation_name
+            active_slot.custom_orientation.matrix = np.linalg.inv(combined_matrix)
+
+        # delete dummy object
+        bpy.ops.object.delete()
+
+        # reinstate selected objects
+        for obj in selected_objects:
+            obj.select_set(True)
+        if active_object:
+            bpy.context.view_layer.objects.active = active_object
+
+        if not use:
+            active_slot.type = initial_orientation
+
     @classmethod
     def edit_container_attributes(cls, entity: ifcopenshell.entity_instance) -> None:
         # TODO