def test_obim_python(threads_1):
_ = threads_1
order = []
def metric(out, i):
return out[i]
def f(out, i, ctx):
order.append(i)
orig = out[i]
out[i] = 10 - i
if orig == 0:
ctx.push(i)
out = np.zeros(10, dtype=int)
for_each(
range(10),
partial(f, out),
worklist=OrderedByIntegerMetric(partial(metric, out)),
)
assert np.allclose(out, np.array([10, 9, 8, 7, 6, 5, 4, 3, 2, 1]))
assert order == [
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0
]
def sssp(graph: PropertyGraph, source, length_property, shift, property_name):
dists = create_distance_array(graph, source, length_property)
# Define the struct type here so it can depend on the type of the weight property
UpdateRequest = np.dtype([("src", np.uint32), ("dist", dists.dtype)])
init_bag = InsertBag[UpdateRequest]()
init_bag.push((source, 0))
t = StatTimer("Total SSSP")
t.start()
for_each(
init_bag,
sssp_operator(graph, dists, graph.get_edge_property(length_property)),
worklist=OrderedByIntegerMetric(obim_indexer(shift)),
disable_conflict_detection=True,
loop_name="SSSP",
)
t.stop()
print("Elapsed time: ", t.get(), "milliseconds.")
graph.add_node_property(pyarrow.table({property_name: dists}))