diff --git a/capellambse_context_diagrams/collectors/class_tree.py b/capellambse_context_diagrams/collectors/class_tree.py
index eab5e98e..64b93653 100644
--- a/capellambse_context_diagrams/collectors/class_tree.py
+++ b/capellambse_context_diagrams/collectors/class_tree.py
@@ -25,23 +25,29 @@ def collector(
     diagram: context.ContextDiagram, params: dict[str, t.Any] | None = None
 ) -> _elkjs.ELKInputData:
     """Return the class tree data for ELK."""
+    params = params or {}
     assert isinstance(diagram.target, information.Class)
     data = generic.collector(diagram, no_symbol=True)
-    data["children"][0]["layoutOptions"] = {}
-    data["children"][0]["layoutOptions"]["elk.partitioning.partition"] = 0
+    if params.get("partitioning", False):
+        data["children"][0]["layoutOptions"] = {}
+        data["children"][0]["layoutOptions"]["elk.partitioning.partition"] = 0
     data["layoutOptions"] = LAYOUT_OPTIONS
     data["layoutOptions"]["algorithm"] = (params or {})["algorithm"]
     data["layoutOptions"]["elk.direction"] = (params or {})["direction"]
     data["layoutOptions"]["edgeRouting"] = (params or {})["edgeRouting"]
 
     made_boxes: set[str] = set()
-    for uid, (source, prop, target) in get_all_classes(diagram.target):
-        partition = uid.split(" ")[-1]
+    for _, (source, prop, target, partition) in get_all_classes(
+        diagram.target
+    ):
         if target.uuid not in made_boxes:
             made_boxes.add(target.uuid)
             box = makers.make_box(target)
-            box["layoutOptions"] = {}
-            box["layoutOptions"]["elk.partitioning.partition"] = int(partition)
+            if params.get("partitioning", False):
+                box["layoutOptions"] = {}
+                box["layoutOptions"]["elk.partitioning.partition"] = int(
+                    partition
+                )
             data["children"].append(box)
 
         width, height = helpers.extent_func(prop.name)
@@ -62,7 +68,7 @@ def collector(
 
 
 ClassContext = tuple[
-    information.Class, information.Property, information.Class
+    information.Class, information.Property, information.Class, int
 ]
 
 
@@ -74,9 +80,9 @@ def get_all_classes(
     classes: dict[str, ClassContext] = {}
     for prop in root.properties:
         if prop.type.xtype.endswith("Class"):
-            edge_id = f"{root.name} {prop.type.name} {partition}"
+            edge_id = f"{root.uuid} {prop.uuid} {prop.type.uuid}"
             if edge_id not in classes:
-                classes[edge_id] = (root, prop, prop.type)
+                classes[edge_id] = (root, prop, prop.type, partition)
                 classes.update(dict(get_all_classes(prop.type, partition)))
 
     yield from classes.items()
diff --git a/capellambse_context_diagrams/context.py b/capellambse_context_diagrams/context.py
index 64d1d2f4..53ca7ce8 100644
--- a/capellambse_context_diagrams/context.py
+++ b/capellambse_context_diagrams/context.py
@@ -350,5 +350,6 @@ def _create_diagram(self, params: dict[str, t.Any]) -> cdiagram.Diagram:
         params.setdefault("algorithm", "layered")
         params.setdefault("direction", "DOWN")
         params.setdefault("edgeRouting", "POLYLINE")
+        params.setdefault("partitioning", True)
         params["elkdata"] = class_tree.collector(self, params)
         return super()._create_diagram(params)
diff --git a/tests/test_class_tree_diagrams.py b/tests/test_class_tree_diagrams.py
index 55200fcb..98a2b48f 100644
--- a/tests/test_class_tree_diagrams.py
+++ b/tests/test_class_tree_diagrams.py
@@ -18,11 +18,22 @@ def test_class_tree_diagram_gets_rendered_successfully(
     assert diag.render(fmt)
 
 
+@pytest.mark.parametrize("edgeRouting", ["SPLINE", "ORTHOGONAL", "POLYLINE"])
+@pytest.mark.parametrize("direction", ["DOWN", "RIGHT"])
+@pytest.mark.parametrize("partitioning", [True, False])
 def test_class_tree_diagram_renders_with_additional_params(
     model: capellambse.MelodyModel,
+    edgeRouting: str,
+    direction: str,
+    partitioning: bool,
 ) -> None:
     obj = model.by_uuid(CLASS_UUID)
 
     diag = obj.class_tree_diagram
 
-    assert diag.render("svgdiagram", edgeRouting="POLYLINE", direction="RIGHT")
+    assert diag.render(
+        "svgdiagram",
+        edgeRouting=edgeRouting,
+        direction=direction,
+        partitioning=partitioning,
+    )