makmanalp · September 12, 2018 19:39 · Sep 12, 2018
diff --git a/network.py b/network.py
@@ -0,0 +1,82 @@
+import pandas as pd
+import json
+
+def read_network(file_name, nodes_field="nodes", edges_field="edges"):
+    network = None
+    with open(file_name, "r") as f:
+        network = json.loads(f.read())
+    nodes = network[nodes_field]
+    edges = network[edges_field]
+    other_fields = {x:network[x] for x in network.keys()
+                    if x not in [nodes_field, edges_field]}
+
+    return pd.DataFrame.from_records(nodes), pd.DataFrame.from_records(edges), other_fields
+
+
+def to_records(df):
+    """Replacement for pandas' to_dict(orient="records") which has issues with
+    upcasting ints to floats in the case of other floats being there.
+
+    https://github.com/pandas-dev/pandas/issues/12859
+    """
+    return json.loads(df.to_json(orient="records"))
+
+
+def write_network(file_name, nodes, edges, other_fields=None, nodes_field="nodes", edges_field="edges"):
+    with open(file_name, "w") as f:
+        network = {}
+        network[nodes_field] = to_records(nodes)
+        network[edges_field] = to_records(edges)
+        if other_fields is not None:
+            network.update(other_fields)
+        f.write(json.dumps(network, indent=4, separators=(',', ': ')))
+
+
+def remap_network_ids(nodes, edges, conversion_mapping, id_field="id", source_field="source", target_field="target"):
+    """conversion_mapping is a series where the index is the old ids, and the
+    values are the new ones."""
+
+    conversion_mapping = conversion_mapping.copy()
+    conversion_mapping.columns = [id_field]
+
+    nodes = nodes.merge(conversion_mapping,
+                        left_on=id_field, right_index=True,
+                        how="left", suffixes=("_old", ""))
+
+    nodes = nodes.drop([x for x in nodes.columns if x.endswith("_old")], axis=1)
+
+    # merge-and-replace
+    # Merge column renaming semantics
+
+    conversion_mapping.columns = [source_field]
+    edges = edges\
+        .merge(conversion_mapping,
+               left_on=source_field, right_index=True,
+               how="left", suffixes=("_old", ""))
+
+    conversion_mapping.columns = [target_field]
+    edges = edges\
+        .merge(conversion_mapping,
+               left_on=target_field, right_index=True,
+               how="left", suffixes=("_old", ""))
+
+    edges = edges.drop([x for x in edges.columns if x.endswith("_old")], axis=1)
+
+    return nodes, edges
+
+
+def find_neighbors(nodes, edges, this_node, extra_fields=[]):
+    connected_edges = edges[(edges.source == this_node) | (edges.target == this_node)].copy()
+
+    def get_other_node(this_node):
+        def inner(df):
+            if df.source == this_node:
+                return int(df.target)
+            else:
+                return int(df.source)
+        return inner
+
+    connected_edges["id"] = connected_edges.apply(get_other_node(this_node), axis=1)
+    connected_edges = connected_edges[["id"] + extra_fields]
+
+    return to_records(connected_edges)