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 88.43%. Comparing base (9d07be1) to head (d1df7ea).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1408   +/-   ##
=======================================
  Coverage   88.43%   88.43%           
=======================================
  Files          93       93           
  Lines       18641    18647    +6     
=======================================
+ Hits        16485    16491    +6     
  Misses       2156     2156           

☔ 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

[amaloney]

@dmiracle I was able to reproduce your results. Would you like to convert this from a draft PR?

[amaloney]

Great work, thanks again for the contribution