From 9fbbbdf9779e38286f5a411cf1265f8914fa28d5 Mon Sep 17 00:00:00 2001
From: dilpath <dilan.private+github@outlook.com>
Date: Wed, 15 Mar 2023 00:09:38 +1100
Subject: [PATCH] model IDs default labels

---
 petab_select/plot.py | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/petab_select/plot.py b/petab_select/plot.py
index 0ca2ba10..31ae4e0b 100644
--- a/petab_select/plot.py
+++ b/petab_select/plot.py
@@ -476,16 +476,20 @@ def graph_iteration_layers(
         VIRTUAL_INITIAL_MODEL, "Virtual\nInitial\nModel"
     )
 
-    missing_labels = [
-        model.get_hash() for model in models if model.get_hash() not in labels
-    ]
-    missing_labels += [
-        model.predecessor_model_hash
+    missing_labels = {
+        model.get_hash(): model.model_id
+        for model in models
+        if model.get_hash() not in labels
+    }
+    missing_labels2 = {
+        model.predecessor_model_hash: model.predecessor_model_hash
         for model in models
         if model.predecessor_model_hash not in labels
-    ]
-    for label in missing_labels:
-        labels[label] = label
+    }
+    labels.update(missing_labels2)
+    labels.update(missing_labels)
+    # for label in missing_labels:
+    #    labels[label] = label
 
     default_draw_networkx_kwargs = {
         'node_color': NORMAL_NODE_COLOR,