Update bamt/display/display.py

Co-authored-by: Yury Kaminsky <[email protected]>
aimclub · Aug 11, 2023 · 48f8a84 · 48f8a84
1 parent 35b8c97
commit 48f8a84
Showing 1 changed file with 7 additions and 10 deletions.
diff --git a/bamt/display/display.py b/bamt/display/display.py
@@ -41,16 +41,13 @@ def _make_network(nodes, edges, **kwargs):
         :param kwargs: a dict passed to pyvis.network.Network
         """
         g = nx.DiGraph()
-        if not kwargs:
-            network_params = dict(
-                height="800px",
-                width="100%",
-                notebook=True,
-                directed=nx.is_directed(g),
-                layout="hierarchical",
-            )
-        else:
-            network_params = kwargs
+        network_params = dict(
+            height=kwargs.get('height', "800px"),
+            width=kwargs.get('width', "100%"),
+            notebook=kwargs.get('notebook', True),
+            directed=kwargs.get('directed', nx.is_directed(g)),
+            layout=kwargs.get('layout', "hierarchical")
+        )
 
         g.add_nodes_from(nodes)
         g.add_edges_from(edges)