diff --git a/moderne_visualizations_misc/call_graph_uml.ipynb b/moderne_visualizations_misc/call_graph_uml.ipynb
index 3788a36..0cf4496 100644
--- a/moderne_visualizations_misc/call_graph_uml.ipynb
+++ b/moderne_visualizations_misc/call_graph_uml.ipynb
@@ -88,63 +88,44 @@
     "    return f\"{name}()\"\n",
     "\n",
     "\n",
-    "for i, row in df.iterrows():\n",
-    "    from_class = row[\"fromClass\"]\n",
-    "    to_class = row[\"toClass\"]\n",
-    "\n",
-    "    from_name = row[\"fromName\"]\n",
-    "    to_name = row[\"toName\"]\n",
-    "\n",
-    "    if is_valid_class(from_class):\n",
-    "        if is_valid_method(from_name):\n",
-    "            from_method = make_method(from_name)\n",
-    "            to_method = make_method(to_name)\n",
-    "            plantuml_code.append(\n",
-    "                f'\"{from_class}::{from_method}\" --> \"{to_class}::{to_method}\"'\n",
-    "            )\n",
-    "\n",
-    "            if from_class not in classes:\n",
-    "                classes[from_class] = set()\n",
-    "\n",
-    "            members = classes[from_class]\n",
-    "            method_name = make_method(from_name)\n",
-    "            if method_name not in members:\n",
-    "                members.add(make_method(from_name))\n",
-    "\n",
-    "            if to_class not in classes:\n",
-    "                classes[to_class] = set()\n",
-    "\n",
-    "            members = classes[to_class]\n",
-    "            method_name = make_method(to_name)\n",
-    "            if method_name not in members:\n",
-    "                members.add(make_method(to_name))\n",
-    "\n",
-    "    if is_valid_class(to_class):\n",
-    "        if is_valid_method(to_name):\n",
-    "            if to_class not in classes:\n",
-    "                classes[to_class] = set()\n",
-    "\n",
-    "            members = classes[to_class]\n",
-    "            method_name = make_method(to_name)\n",
-    "            if method_name not in members:\n",
-    "                members.add(make_method(to_name))\n",
-    "\n",
-    "            if from_class not in classes:\n",
-    "                classes[from_class] = set()\n",
-    "\n",
-    "            members = classes[from_class]\n",
-    "            method_name = make_method(from_name)\n",
-    "            if method_name not in members:\n",
-    "                members.add(make_method(from_name))\n",
-    "\n",
-    "    if from_class != to_class and is_valid_class(to_class):\n",
-    "        if is_valid_method(to_name):\n",
-    "            from_method = make_method(to_name)\n",
-    "            to_method = make_method(to_name)\n",
-    "            plantuml_code.append(\n",
-    "                f'\"{to_class}::{from_method}\" <-- \"{from_class}::{to_method}\"'\n",
-    "            )\n",
+    "def add_method_to_class(class_name, method_name):\n",
+    "    if class_name not in classes:\n",
+    "        classes[class_name] = set()\n",
+    "    members = classes[class_name]\n",
+    "    method = make_method(method_name)\n",
+    "    if method not in members:\n",
+    "        members.add(method)\n",
+    "\n",
     "\n",
+    "for i, row in df.iterrows():\n",
+    "    from_class, to_class = row[\"fromClass\"], row[\"toClass\"]\n",
+    "    from_name, to_name = row[\"fromName\"], row[\"toName\"]\n",
+    "\n",
+    "    if is_valid_class(from_class) and is_valid_method(from_name):\n",
+    "        add_method_to_class(from_class, from_name)\n",
+    "        add_method_to_class(to_class, to_name)\n",
+    "\n",
+    "        from_method = make_method(from_name)\n",
+    "        to_method = make_method(to_name)\n",
+    "        plantuml_code.append(\n",
+    "            f'\"{from_class}::{from_method}\" --> \"{to_class}::{to_method}\"')\n",
+    "\n",
+    "    if is_valid_class(to_class) and is_valid_method(to_name):\n",
+    "        add_method_to_class(to_class, to_name)\n",
+    "        add_method_to_class(from_class, from_name)\n",
+    "\n",
+    "        from_method = make_method(from_name)\n",
+    "        to_method = make_method(to_name)\n",
+    "        plantuml_code.append(\n",
+    "            f'\"{to_class}::{to_method}\" <-- \"{from_class}::{from_method}\"')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "# Create classes\n",
     "for _class in classes:\n",
     "    plantuml_code.append(f'class \"{_class}\"')\n",
@@ -160,7 +141,7 @@
     "\n",
     "plantuml_code_string = \"\\n\".join(plantuml_code)\n",
     "\n",
-    "# print(plantuml_code_string)"
+    "# print(plantuml_code_string)\n"
    ]
   },
   {