add network support

Author: wybert

datashader_cli/cli.py

@@ -188,9 +188,127 @@ def points(data_path, x, y, w=600, h=600, x_range=None, y_range=None,
     t6 = datetime.datetime.now()
     print("Time to create image: ", t6-t5)
     print("Total time: ", t6-t1)
-main.add_command(points)
 
 
+import datashader as ds
+import datashader.transfer_functions as tf
+
+
+def nodesplot(nodes, name=None, canvas=None, cat=None, cvsopts=None):
+    canvas = ds.Canvas(**cvsopts) if canvas is None else canvas
+    aggregator=None if cat is None else ds.count_cat(cat)
+    agg=canvas.points(nodes,'x','y',aggregator)
+    return tf.spread(tf.shade(agg, cmap=["#FF3333"]), px=1, name=name)
+
+def edgesplot(edges, name=None, canvas=None, cvsopts=None, edge_cmap=None):
+    canvas = ds.Canvas(**cvsopts) if canvas is None else canvas
+
+    if edge_cmap:
+        return tf.shade(canvas.line(edges, 'x','y', agg=ds.count()), cmap=cc.palette[edge_cmap], name=name)
+    return tf.shade(canvas.line(edges, 'x','y', agg=ds.count()), name=name)
+    
+def graphplot(nodes, edges, name="", canvas=None, cat=None, cvsopts=None, edge_cmap=None):
+    if canvas is None:
+        xr = nodes.x.min(), nodes.x.max()
+        yr = nodes.y.min(), nodes.y.max()
+        canvas = ds.Canvas(x_range=xr, y_range=yr, **cvsopts)
+        
+    np = nodesplot(nodes, name + " nodes", canvas, cat, cvsopts)
+    ep = edgesplot(edges, name + " edges", canvas, cvsopts, edge_cmap)
+    return tf.stack(ep, np, how="over", name=name)
+
+
+@click.command()
+@click.argument('nodes_file', type=click.Path(exists=True), required=True)
+@click.argument('edges_file', type=click.Path(exists=True), required=True)
+@click.argument('output_path',default = None, type = str, required=True)
+
+@click.option('--w', default=600,type = int, required=False, help = 'How many pixels wide to make the image')
+@click.option('--h', default=600, type = int,required=False, help = 'How many pixels high to make the image')
+@click.option('--x',default='x',type= str, required=False, help = """Column name for x coordinates, e.g. "x", if use layout="geo", x is required """)
+@click.option('--y',default='y',type= str, required=False, help = """Column name for y coordinates, e.g. "y", if use layout="geo", y is required """)
+@click.option('--source',default='source',type= str, required=False, help = """Column name for source node, e.g. "source" """)
+@click.option('--target',default='target',type= str, required=False, help = """Column name for target node, e.g. "target" """)
+@click.option('--id',default=None,type= str, required=False, help = """Column name for node id, e.g. "id" """)
+@click.option('--layout',default='random',type= str, required=False, help = """Layout algorithm, e.g. "random", "forceatlas2", "geo" """)
+@click.option('--cat',default=None, type = str, required=False, help = """Column to group by, e.g. "category", see datashader docs (https://datashader.org/api.html#reductions) for more options""")
+@click.option('--background',default=None, required=False, type = str, help = """Background color, e.g. "black", "white", "#000000", "#ffffff" """)
+@click.option('--bundle',default=False, required=False, type = bool, help = """Whether to bundle the edges""")
+@click.option('--bw', default=None,type = float, required=False, help = 'initial_bandwidth for bundling')
+@click.option('--decay', default=None, type = float, required=False, help = 'decay for bundling')
+@click.option('--edge_cmap',default= None, required=False, type = str, help = 'Name of the colormap for edges, see https://colorcet.holoviz.org for more options')
+def network(nodes_file, edges_file, output_path, w=600, h=600, x="x",y="y", source="source", target="target", id=None,
+            layout="forceatlas2",cat=None, 
+            background="white", bundle=False, bw=None, decay=None, edge_cmap=None):
+
+    # read data
+    import pandas as pd
+    from datashader.layout import random_layout, circular_layout, forceatlas2_layout
+    from datashader.bundling import connect_edges, hammer_bundle
+
+    if nodes_file.endswith('.parquet'):
+        cnodes = pd.read_parquet(nodes_file)
+    elif nodes_file.endswith('.csv'):
+        cnodes = pd.read_csv(nodes_file)
+    else:
+        raise ValueError('Unsupported data format')
+    
+    cnodes = cnodes.rename(columns={x:"x", y:"y"})
+
+
+        
+    # cnodes.drop(columns=['AirportID'], inplace=True)
+    # this script can't handle id 
+
+    if edges_file.endswith('.parquet'):
+        cedges = pd.read_parquet(edges_file)
+    elif edges_file.endswith('.csv'):
+        cedges = pd.read_csv(edges_file)
+    else:
+        raise ValueError('Unsupported data format')
+    
+    cedges = cedges.rename(columns={source:"source", target:"target"})
+
+    if id:
+        cnodes["order"] = range(len(cnodes))
+        mapper = cnodes.set_index(id)["order"].to_dict()
+        cedges["source"] = cedges["source"].map(mapper)
+        cedges["target"] = cedges["target"].map(mapper)
+    # print(cedges.head())
+    # layout
+    layouts = {
+        "random":random_layout,
+        "circular":circular_layout,
+        "forceatlas2":forceatlas2_layout
+    }
+    
+    if layout == "geo":
+        layout = cnodes.rename(columns={x:"x", y:"y"})
+        # print(layout.head())
+    else:
+
+        layout = layouts[layout](cnodes, cedges)
+        print(layout.head())
+
+    cvsopts = dict(plot_height=h, plot_width=w)
+    if bundle:
+        bundle_keywords = {
+            }
+        if bw:
+            bundle_keywords["initial_bandwidth"] = bw
+        if decay:
+            bundle_keywords["decay"] = decay        
+        plot = graphplot(layout, hammer_bundle(layout,cedges, **bundle_keywords), "title",  cat=cat, cvsopts=cvsopts, edge_cmap=edge_cmap)
+
+    else:
+        plot = graphplot(layout, connect_edges(layout,cedges), "title",  cat=cat, cvsopts=cvsopts, edge_cmap=edge_cmap)
+    img = tf.set_background(plot, background)
+    img.to_pil().save(output_path)
+
+
+
+main.add_command(points)
+main.add_command(network)
 
 
 if __name__ == "__main__":

requirements.txt

@@ -3,4 +3,5 @@ datashader
 pandas
 matplotlib
 geopandas
-pyarrow
+pyarrow
+scikit-image