<jupyter_start><jupyter_text>SSSP Example<jupyter_code>import dask
import numpy as np
import grblas as gb
import dask_grblas as dgb
from grblas import op
from dask.distributed import Client

client = Client()
client
# Create random data
N = 1000
num_chunks = 4
r = np.random.rand(N, N) < 0.001
r = r | r.T  # symmetric
r = r & ~np.diag(np.ones(N, dtype=bool))  # no self edges
# Option 1: create distributed Matrix from local data
def to_matrix(chunk):
    rows, cols = np.nonzero(chunk)
    values = np.random.rand(rows.size)
    return dgb.Matrix.from_values(
        rows, cols, values, nrows=chunk.shape[0], ncols=chunk.shape[1]
    )


chunks = np.array_split(r, num_chunks, axis=0)
delayed_chunks = [to_matrix(chunk) for chunk in chunks]
A = dgb.row_stack(delayed_chunks)
sources = dgb.Vector.from_values(np.random.randint(N), 0, size=N, dtype=A.dtype)
# Option 2: create distributed Matrix from distributed (delayed) data
chunks = np.array_split(r, num_chunks, axis=0)
ncols = chunks[0].shape[1]
row_lengths = np.array([chunk.shape[0] for chunk in chunks])
row_offsets = np.roll(row_lengths.cumsum(), 1)
row_offsets[0] = 0

chunked_rows = []
chunked_cols = []
chunked_vals = []
for chunk, row_offset in zip(chunks, row_offsets):
    rows, cols = np.nonzero(chunk)
    chunked_rows.append(rows + row_offset)
    chunked_cols.append(cols)
    chunked_vals.append(np.random.rand(rows.size))

delayed_rows = [dask.delayed(rows) for rows in chunked_rows]
delayed_cols = [dask.delayed(cols) for cols in chunked_cols]
delayed_vals = [dask.delayed(cols) for cols in chunked_vals]


@dask.delayed
def to_matrix(rows, cols, vals, nrows, ncols):
    # Can also use e.g. gb.Matrix.ss.import_csr
    return gb.Matrix.from_values(rows, cols, vals, nrows=nrows, ncols=ncols)


delayed_matrices = [
    to_matrix(
        delayed_rows[i] - row_offsets[i],
        delayed_cols[i],
        delayed_vals[i],
        row_lengths[i],
        ncols,
    )
    for i in range(num_chunks)
]

delayed_chunks = [
    dgb.Matrix.from_delayed(
        delayed_matrices[i],
        gb.dtypes.FP64,
        row_lengths[i],
        ncols,
    )
    for i in range(num_chunks)
]

A = dgb.row_stack(delayed_chunks)
sources = dgb.Vector.from_values(np.random.randint(N), 0, size=N, dtype=A.dtype)
# Calculate expected with grblas
B = A.compute()
v = sources.dup().compute()
v_dup = gb.Vector.new(v.dtype, size=N)
i = 0
while True:
    i += 1
    v_dup << v
    v(op.min) << B.mxv(v, op.min_plus)
    if v.isequal(v_dup):
        break
expected = v
i
# Calculate with dask-grblas
i = 0
v = sources.dup()
while True:
    i += 1
    v_dup = v.dup()
    v(op.min) << A.mxv(v, op.min_plus)
    # persist so we don't recompute every iteration
    v._delayed = v._delayed.persist()  # scheduler='synchronous')
    if v.isequal(v_dup):
        break
i
assert expected.isequal(v.compute())
expected<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12685793294158312, 'fraction_numerical': 0.010024196335983409, 'mean_word_length': 3.705691056910569, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '665731', 'license': 'Apache-2.0', 'n_tokens_mistral': 1151, 'n_tokens_neox': 1128, 'n_words': 255, 'path': 'notebooks/SSSP.ipynb', 'repo_name': 'jim22k/dask-grblas', 'size': '2893'}