naive jit implementation for cooling

[dmiracle]

Tried out a simple numba implementation for layout cooling and got more than a 10x speedup. I made some questionable choices (matrix = matrix.toarray()) but just putting it up to say that the approach gave a good speedup.

[dmiracle]

Here is some sample code where the speedup can be seen (12 min -> 1 min on my m1 mac powerbook)

import networkx as nx
import pandas as pd
import time

from datashader.layout import forceatlas2_layout

# Make a random graph with 50000 nodes and 200000 edges and convert to a dataframe
G = nx.gnm_random_graph(50000, 200000)
nodes = pd.DataFrame({'id': list(G.nodes)}).set_index('id')
edges = pd.DataFrame({'source': [x[0] for x in list(G.edges)], 'target': [x[1] for x in list(G.edges)]})
edges.reset_index(drop=False, inplace=True)
edges.rename(columns={'index': 'id'}, inplace=True)
edges = edges.set_index('id')

# Run the forceatlas2 layout
start_time = time.perf_counter()
forcedirected = forceatlas2_layout(nodes, edges)
end_time = time.perf_counter()
print(f"Time taken: {end_time - start_time} seconds")

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 87.05%. Comparing base (12e7471) to head (d12fede).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1408   +/-   ##
=======================================
  Coverage   87.04%   87.05%           
=======================================
  Files          93       93           
  Lines       18749    18755    +6     
=======================================
+ Hits        16321    16327    +6     
  Misses       2428     2428           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[holovizbot]

This pull request has been mentioned on HoloViz Discourse. There might be relevant details there:

https://discourse.holoviz.org/t/how-can-i-get-datashader-force-directed-layout-to-eat-more-cores/8575/4