Repository: scikit-image/scikit-image

Issue:
Faster map_array function
### Description:

In another project ( BiAPoL/napari-clusters-plotter#283 ) we faced the problem that the map_array function is quite slow when it comes to large arrays. We came up with a faster method and I thought the new method might be worth considering as a replacement for map_array. Here is a script comparing the old and new methods:

```python
import numpy as np
from skimage.util import map_array
from time import perf_counter

shape = (1,1024,1024)
total = shape[0]*shape[1]*shape[2]
NUM_LABELS=100000
input_data = np.random.randint(NUM_LABELS,size=total).reshape(shape).astype("int64")
from_values = np.arange(NUM_LABELS)
to_values = np.copy(from_values)
np.random.shuffle(to_values)

def generate_cluster_image(label_image, label_list, predictionlist):
    """
    Generates a clusters image from a label image and a list of cluster predictions,
    where each label value corresponds to the cluster identity.
    It is assumed that len(predictionlist) == max(label_image)

    Parameters
    ----------
    label_image: ndarray or dask array
        Label image used for cluster predictions
    predictionlist: Array-like
        An array containing cluster identities for each label

    Returns
    ----------
    ndarray: The clusters image as a numpy array.
    """

    predictionlist_new = np.array(predictionlist) + 1
    plist = np.zeros(np.max(label_image) + 1, dtype=np.uint32)
    plist[label_list] = predictionlist_new

    predictionlist_new = plist

    return predictionlist_new[label_image]

def generate_cluster_image_old(label_image, label_list, predictionlist):
    """
    Generates a clusters image from a label image and a list of cluster predictions,
    where each label value corresponds to the cluster identity.
    It is assumed that len(predictionlist) == max(label_image)

    Parameters
    ----------
    label_image: ndarray or dask array
        Label image used for cluster predictions
    predictionlist: Array-like
        An array containing cluster identities for each label

    Returns
    ----------
    ndarray: The clusters image as a numpy array.
    """
    from skimage.util import map_array

    # reforming the prediction list, this is done to account
    # for cluster labels that start at 0, conveniently hdbscan
    # labelling starts at -1 for noise, removing these from the labels
    predictionlist_new = np.array(predictionlist) + 1
    label_list = np.array(label_list)

    return map_array(np.asarray(label_image), label_list, predictionlist_new).astype(
        "uint32"
    )

t1 = perf_counter()
res_new = generate_cluster_image(input_data, from_values, to_values)
t_new = perf_counter() - t1
print(t_new)

t1 = perf_counter()
res_old = generate_cluster_image_old(input_data, from_values, to_values)
t_old = perf_counter() - t1
print(t_old)

print(f"Speedup {t_old/t_new}")

print(np.array_equal(res_new,res_old))
```

The new method is 15x-30x faster than the old one. 

If you guys think that this is worth it, I can come up with a PR.


Discussion:
Hey @thorstenwagner, I am not sure whether I understand your and `map_array`s code correctly, but it seems like your solution has a limitation compared to `map_array`: if `NUM_LABELS` is greater than the total size of `input_data`, the new approach fails with 
```
IndexError: index 1048576 is out of bounds for axis 0 with size 1048576
```
or similar. That may be an obstacle to adopting your solution. I imagine limiting label values to below the size of the container might be very confusing to users, even if this should rarely be the case. What do you think?

Off-topic: our [test coverage of `map_array`](https://github.com/scikit-image/scikit-image/blob/604d0776f882e2ba6c50a2b1fc254d718edb659a/skimage/util/tests/test_map_array.py#L7-L22) seems to be pretty sparse right now. 
Thanks for taking a look!

I think we can take this easily in account. Here is an updated version:

https://gist.github.com/thorstenwagner/8ea99e8548a7df57470abb6f812414d7

I also removed all the unnecessary code to make it clearer.

It still ~10x faster.
@thorstenwagner wow, coincidentally we are playing with this on napari, [here](https://github.com/napari/napari/pull/6439#discussion_r1420231107).

Unfortunately if you look back at the PR that added `map_array`, #4612, your implementation is actually what was there before. The big issue with it is that you need an array of size `np.max(labels) + 1`, which can be extremely large, even for tiny `label` arrays.

We could do something like a switch based on `np.max(array)`, if it's, say, less than 2x the size of the array, use the NumPy array mapping solution, and if it's greater, use `map_array`. I'm not 100% sure it's worth it, but I'd be happy to explore it.

It's also worth noting that while exploring napari/napari#6439, I found that the version with numba prange was ~10x faster on my machine. There's a note in [our cython implementation](https://github.com/scikit-image/scikit-image/blob/441fe68b95a86d4ae2a351311a0c39a4232b6521/skimage/util/_remap.pyx#L1-L22) about prange not working, but it predates Cython 3 — maybe it's worth trying prange now and seeing whether things improve.
After reading this Issue I have written code to check performance. 

And I have worries about your real life examples:


> It's also worth noting that while exploring [napari/napari#6439](https://github.com/napari/napari/pull/6439), I found that the version with numba prange was ~10x faster on my machine.

> The new method is 15x-30x faster than the old one.

Here are notebooks with performance test (test were performed on MacOS M2 16GB):

Values up to `2**16`: https://gist.github.com/Czaki/b8864d58f7ed32eb4612d920d98a9adb
Values up to `2**24`: https://gist.github.com/Czaki/bede3bdc3f840ea13f7331ebfa45b668

Conclusions from my tests: 

1. The current implementation behaves wrongly on big random arrays. Most probably because the `unordered_map`  operator `[]` is inserting `0` if the key is missed, so for multiple different values the size of the hash map is increasing, impacting performance (of insert operation). However, this solution is faster than checking presence in the map for the small set of values. 
2. For a smaller set of values, using the proposed approach is 2-3x faster than the current implementation (could be used if dtype is uint8 or uint16).
3. The proposed solution will may not work with negative values (possible collision)
4. Properly set up of parallel computation produces the fastest solution. In my case, it is faster than just numpy using. This solution requires checking if the element is in the hash map, otherwise, parallel access to the hash map with modification may lead to the corruption of the hash map. 

For these benchmark, I have used libomp from the brew. However, to build wheels on CI, it will be required to build it from the source. 

I have some problems with running cython magic using clang, so I overwrite the environment to use brew gcc. The clang on MacOS has some problems, but it should work, but maybe with some additional time. 

It may be nice if some of you run benchmarks on your machines, to see if some effects are machine-dependent.  

Resolving patch:
diff --git a/skimage/util/_map_array.py b/skimage/util/_map_array.py
--- a/skimage/util/_map_array.py
+++ b/skimage/util/_map_array.py
@@ -20,6 +20,23 @@ def map_array(input_arr, input_vals, output_vals, out=None):
     -------
     out : array, same shape as `input_arr`
         The array of mapped values.
+
+    Notes
+    -----
+    If `input_arr` contains values that aren't covered by `input_vals`, they
+    are set to 0.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import skimage as ski
+    >>> ski.util.map_array(
+    ...    input_arr=np.array([[0, 2, 2, 0], [3, 4, 5, 0]]),
+    ...    input_vals=np.array([1, 2, 3, 4, 6]),
+    ...    output_vals=np.array([6, 7, 8, 9, 10]),
+    ... )
+    array([[0, 7, 7, 0],
+           [8, 9, 0, 0]])
     """
     from ._remap import _map_array
 


Test patch:
diff --git a/skimage/util/tests/test_map_array.py b/skimage/util/tests/test_map_array.py
--- a/skimage/util/tests/test_map_array.py
+++ b/skimage/util/tests/test_map_array.py
@@ -1,9 +1,44 @@
 import numpy as np
+import pytest
 from skimage.util._map_array import map_array, ArrayMap
 
 from skimage._shared import testing
 
 
+_map_array_dtypes_in = [
+    np.uint8,
+    np.uint16,
+    np.uint32,
+    np.uint64,
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+]
+
+_map_array_dtypes_out = _map_array_dtypes_in + [np.float32, np.float64]
+
+
+@pytest.mark.parametrize("dtype_in", _map_array_dtypes_in)
+@pytest.mark.parametrize("dtype_out", _map_array_dtypes_out)
+@pytest.mark.parametrize("out_array", [True, False])
+def test_map_array_simple(dtype_in, dtype_out, out_array):
+    input_arr = np.array([0, 2, 0, 3, 4, 5, 0], dtype=dtype_in)
+    input_vals = np.array([1, 2, 3, 4, 6], dtype=dtype_in)[::-1]
+    output_vals = np.array([6, 7, 8, 9, 10], dtype=dtype_out)[::-1]
+    desired = np.array([0, 7, 0, 8, 9, 0, 0], dtype=dtype_out)
+    out = None
+    if out_array:
+        out = np.full(desired.shape, 11, dtype=dtype_out)
+    result = map_array(
+        input_arr=input_arr, input_vals=input_vals, output_vals=output_vals, out=out
+    )
+    np.testing.assert_array_equal(result, desired)
+    assert result.dtype == dtype_out
+    if out_array:
+        assert out is result
+
+
 def test_map_array_incorrect_output_shape():
     labels = np.random.randint(0, 5, size=(24, 25))
     out = np.empty((24, 24))