# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:.conda-jme]
#     language: python
#     name: conda-env-.conda-jme-py
# ---

# ## Firstly we set up 8 dask workers

import numpy as np
import dask
from dask.distributed import Client, LocalCluster
import dask.bag as db

cluster = LocalCluster(n_workers=8)
# explicitly connect to the cluster we just created
client = Client(cluster)

client


# ## And we define a function
#
# $$
# g(x) = \sqrt{|x|} ^{\ln |x|}
# $$
#
# This is a test function for each pixel. In the real case we use an optimization workflow for each pixel.

# +
def g(x):
    return np.sqrt(np.abs(x)) ** np.log(np.abs(x))

g(3)
# -

# ## ...and a large array
#
# with 2 million pixels. Pixel values are normally distributed.

test_array = np.random.randn(2000000)

# Now we are going to apply $g(x)$ to each pixel in the test array.
#
# ## Serial

# +
# %%time

results_serial = np.zeros((2000000, 1))
for i in range(len(test_array)):
    results_serial[i] = g(test_array[i])
# -

results_serial[[0, 1, 2, 3, 10, -10, -3, -2, -1]].flatten()

# ## Parallel
#
# Using `dask.bag`

b = db.from_sequence(test_array, npartitions=24)
b = b.map(g)
b.visualize()

# %%time
results_parallel = b.compute()

results_parallel = np.array(results_parallel)
results_parallel[[0, 1, 2, 3, 10, -10, -3, -2, -1]]

# - Serial: **~11 s**
# - Parallel (8 workers): **~1.5 min**
#
# Why?

StarcoderdataJupyterScriptsDedupFiltered