Merge pull request #123 from oemof/Fix/networkx_graph

Fix nx graph methods

dhnx/graph.py

@@ -15,19 +15,24 @@
 import pandas as pd
 
 
-def thermal_network_to_nx_graph(thermal_network):
+def thermal_network_to_nx_graph(thermal_network, type_of_graph=None):
     r"""
 
     Parameters
     ----------
     thermal_network : dhnx.network.ThermalNetwork
+    type_of_graph : networkx graph class
+        Must be either Graph() or DiGraph()
 
     Returns
     -------
     nx_graph : nx:MultiDigraph
         networkx graph of thermal_network
     """
-    nx_graph = nx.DiGraph()  # TODO: Check if this line can be removed.
+    if type_of_graph is None:
+        type_of_graph = nx.Graph()
+
+    nx_graph = type_of_graph
 
     edges = thermal_network.components['pipes'].copy()
 

dhnx/network.py

@@ -13,6 +13,7 @@
 
 import os
 
+import networkx as nx
 import pandas as pd
 
 from .graph import thermal_network_to_nx_graph
@@ -112,7 +113,16 @@ def to_csv_folder(self, dirname):
         exporter.save()
 
     def to_nx_graph(self):
-        nx_graph = thermal_network_to_nx_graph(self)
+        nx_graph = thermal_network_to_nx_graph(
+            self, type_of_graph=nx.DiGraph(),
+        )
+
+        return nx_graph
+
+    def to_nx_undirected_graph(self):
+        nx_graph = thermal_network_to_nx_graph(
+            self, type_of_graph=nx.Graph(),
+        )
 
         return nx_graph