Fall back to memory efficient replace_values in relabel

.gitignore

@@ -4,5 +4,11 @@
 *.ipynb
 *.db
 *.py[cod]
+*.c
+*.cpp
+*.so
+lsd.egg-info
+build
 docker/lsd
 docker/requirements.txt
+docker/setup.py

Makefile

@@ -1,6 +1,5 @@
 default:
-	pip install .
-	-rm -rf dist build gunpowder.egg-info
+	python setup.py install
 
 .PHONY: install-full
 install-full:

docker/Dockerfile

@@ -6,6 +6,7 @@ LABEL maintainer [email protected]
 RUN apt-get update && apt-get install -y --no-install-recommends \
         wget \
         git \
+        python-numpy \
         libmlpack-dev && \
     rm -rf /var/lib/apt/lists/*
 
@@ -116,5 +117,7 @@ ENV PYTHONPATH ${DAISY_ROOT}:$PYTHONPATH
 # Makefile ensures that)
 ADD lsd /src/lsd/lsd
 ADD requirements.txt /src/lsd/requirements.txt
+ADD setup.py /src/lsd/setup.py
 WORKDIR /src/lsd
+RUN python setup.py build_ext --inplace
 ENV PYTHONPATH /src/lsd:$PYTHONPATH

docker/Makefile

@@ -1,11 +1,12 @@
-TAG="funkey/lsd:v0.3"
+TAG="funkey/lsd:v0.4"
 
 .PHONY: default push
 default:
 	#-docker rmi -f $(TAG)
 	-rm -rf lsd
 	cp -r ../lsd .
 	cp ../requirements.txt .
+	cp ../setup.py .
 	docker build -t $(TAG) .
 
 push: default

lsd/labels.py

@@ -1,11 +1,26 @@
+from __future__ import absolute_import
+from .replace_values import replace_values as cpp_replace_values
 import numpy as np
 
 def replace_values(array, old_values, new_values):
 
-    values_map = np.arange(old_values.max() + 1, dtype=new_values.dtype)
-    values_map[old_values] = new_values
+    max_old_value = old_values.max()
 
-    return values_map[array]
+    if max_old_value < 1024**3:
+
+        values_map = np.arange(old_values.max() + 1, dtype=new_values.dtype)
+        values_map[old_values] = new_values
+
+        return values_map[array]
+
+    else:
+
+        values_map = {
+            old_value: new_value
+            for old_value, new_value in zip(old_values, new_values)
+        }
+
+        return cpp_replace_values(array, values_map)
 
 def relabel(array, return_backwards_map=False):
     '''Relabel array, such that IDs are consecutive. Excludes 0.'''

lsd/replace_values.pyx

@@ -0,0 +1,22 @@
+from libcpp.map cimport map
+cimport cython
+import logging
+import numpy as np
+cimport numpy as np
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def replace_values(array_np, values_map):
+
+    cdef Py_ssize_t i = 0
+    cdef Py_ssize_t n = array_np.size()
+
+    replaced_np = np.zeros_like(array_np)
+    cdef np.npy_uint64[:] array = array_np
+    cdef np.npy_uint64[:] replaced = replaced_np
+    cdef map[np.npy_uint64, np.npy_uint64] cmap = values_map
+
+    for i in range(n):
+        replaced[i] = cmap[array[i]]
+
+    return replaced_np

setup.py

@@ -1,5 +1,6 @@
-from setuptools import setup
-import subprocess
+from Cython.Distutils import build_ext
+from distutils.core import setup
+from distutils.extension import Extension
 
 setup(
         name='lsd',
@@ -13,12 +14,12 @@
             'lsd',
             'lsd.gp',
         ],
-        install_requires=[
-            "numpy",
-            "scipy",
-            "h5py",
-            "scikit-image",
-            "requests",
-            "gunpowder"
-        ]
+        ext_modules=[
+            Extension(
+                'lsd.replace_values',
+                sources=['lsd/replace_values.pyx'],
+                extra_compile_args=['-O3'],
+                language='c++')
+        ],
+        cmdclass={'build_ext': build_ext}
 )