pytme 0.2.2__cp311-cp311-macosx_14_0_arm64.whl → 0.2.4__cp311-cp311-macosx_14_0_arm64.whl

tme/backends/cupy_backend.py

@@ -149,10 +149,10 @@ class CupyBackend(NumpyFFTWBackend):
             cache.clear()
 
         def rfftn(arr: CupyArray, out: CupyArray) -> CupyArray:
-            return cufft.rfftn(arr)
+            return cufft.rfftn(arr, s=fast_shape)
 
         def irfftn(arr: CupyArray, out: CupyArray) -> CupyArray:
-            return cufft.irfftn(arr)
+            return cufft.irfftn(arr, s=fast_shape)
 
         PLAN_CACHE[current_device] = [fast_shape, fast_ft_shape]
 
@@ -167,11 +167,6 @@ class CupyBackend(NumpyFFTWBackend):
         fast_shape = [next_fast_len(x, real=True) for x in convolution_shape]
         fast_ft_shape = list(fast_shape[:-1]) + [fast_shape[-1] // 2 + 1]
 
-        # This almost never happens but avoid cuFFT casting errors
-        is_odd = fast_shape[-1] % 2
-        fast_shape[-1] += is_odd
-        fast_ft_shape[-1] += is_odd
-
         return convolution_shape, fast_shape, fast_ft_shape
 
     def max_filter_coordinates(self, score_space, min_distance: Tuple[int]):

tme/backends/jax_backend.py

@@ -119,19 +119,6 @@ class JaxBackend(NumpyFFTWBackend):
 
         return rfftn, irfftn
 
-    def compute_convolution_shapes(
-        self, arr1_shape: Tuple[int], arr2_shape: Tuple[int]
-    ) -> Tuple[List[int], List[int], List[int]]:
-        conv_shape, fast_shape, fast_ft_shape = super().compute_convolution_shapes(
-            arr1_shape, arr2_shape
-        )
-
-        is_odd = fast_shape[-1] % 2
-        fast_shape[-1] += is_odd
-        fast_ft_shape[-1] += is_odd
-
-        return conv_shape, fast_shape, fast_ft_shape
-
     def rigid_transform(
         self,
         arr: BackendArray,
@@ -144,8 +131,8 @@ class JaxBackend(NumpyFFTWBackend):
         **kwargs,
     ) -> Tuple[BackendArray, BackendArray]:
         rotate_mask = arr_mask is not None
-        center = self.divide(self.to_backend_array(arr.shape), 2)[:, None]
 
+        center = self.divide(self.to_backend_array(arr.shape) - 1, 2)[:, None]
         indices = self._array_backend.indices(arr.shape, dtype=self._float_dtype)
         indices = indices.reshape((arr.ndim, -1))
         indices = indices.at[:].add(-center)
@@ -200,7 +187,7 @@ class JaxBackend(NumpyFFTWBackend):
         target_shape = tuple(
             (x.stop - x.start + p) for x, p in zip(splits[0][0], target_pad)
         )
-        fast_shape, fast_ft_shape, shift = matching_data._fourier_padding(
+        conv_shape, fast_shape, fast_ft_shape, shift = matching_data._fourier_padding(
             target_shape=self.to_numpy_array(target_shape),
             template_shape=self.to_numpy_array(matching_data._template.shape),
             pad_fourier=False,
@@ -210,13 +197,26 @@ class JaxBackend(NumpyFFTWBackend):
             "convolution_mode": convolution_mode,
             "fourier_shift": shift,
             "targetshape": target_shape,
-            "templateshape": matching_data._template.shape,
+            "templateshape": matching_data.template.shape,
+            "convolution_shape": conv_shape,
         }
 
         create_target_filter = matching_data.target_filter is not None
         create_template_filter = matching_data.template_filter is not None
         create_filter = create_target_filter or create_template_filter
 
+        # Applying the filter leads to more FFTs
+        fastt_shape = matching_data._template.shape
+        if create_template_filter:
+            # _, fastt_shape, _, tshift = matching_data._fourier_padding(
+            #     target_shape=self.to_numpy_array(matching_data._template.shape),
+            #     template_shape=self.to_numpy_array(
+            #         [1 for _ in matching_data._template.shape]
+            #     ),
+            #     pad_fourier=False,
+            # )
+            fastt_shape = matching_data._template.shape
+
         ret, template_filter, target_filter = [], 1, 1
         rotation_mapping = {
             self.tobytes(matching_data.rotations[i]): i
@@ -246,12 +246,12 @@ class JaxBackend(NumpyFFTWBackend):
 
             if create_template_filter:
                 template_filter = matching_data.template_filter(
-                    shape=matching_data._template.shape, **filter_args
+                    shape=fastt_shape, **filter_args
                 )["data"]
                 template_filter = template_filter.at[(0,) * template_filter.ndim].set(0)
 
             if create_target_filter:
-                target_filter = matching_data.template_filter(
+                target_filter = matching_data.target_filter(
                     shape=fast_shape, **filter_args
                 )["data"]
                 target_filter = target_filter.at[(0,) * target_filter.ndim].set(0)
@@ -260,8 +260,8 @@ class JaxBackend(NumpyFFTWBackend):
             base, targets = None, self._array_backend.stack(targets)
             scores, rotations = scan_inner(
                 targets,
-                matching_data.template,
-                matching_data.template_mask,
+                self.topleft_pad(matching_data.template, fastt_shape),
+                self.topleft_pad(matching_data.template_mask, fastt_shape),
                 matching_data.rotations,
                 template_filter,
                 target_filter,
@@ -280,3 +280,18 @@ class JaxBackend(NumpyFFTWBackend):
                 ret.append(tuple(temp._postprocess(**analyzer_args)))
 
         return ret
+
+    def get_available_memory(self) -> int:
+        import jax
+
+        _memory = {"cpu": 0, "gpu": 0}
+        for device in jax.devices():
+            if device.platform == "cpu":
+                _memory["cpu"] = super().get_available_memory()
+            else:
+                mem_stats = device.memory_stats()
+                _memory["gpu"] += mem_stats.get("bytes_limit", 0)
+
+        if _memory["gpu"] > 0:
+            return _memory["gpu"]
+        return _memory["cpu"]

tme/backends/npfftw_backend.py

@@ -186,7 +186,7 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
     def to_sharedarr(
         self, arr: NDArray, shared_memory_handler: type = None
     ) -> shm_type:
-        if type(shared_memory_handler) == SharedMemoryManager:
+        if isinstance(shared_memory_handler, SharedMemoryManager):
             shm = shared_memory_handler.SharedMemory(size=arr.nbytes)
         else:
             shm = shared_memory.SharedMemory(create=True, size=arr.nbytes)
@@ -347,7 +347,8 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         cache: bool = False,
     ) -> Tuple[NDArray, NDArray]:
         translation = self.zeros(arr.ndim) if translation is None else translation
-        center = self.divide(self.to_backend_array(arr.shape), 2)
+
+        center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
         if not use_geometric_center:
             center = self.center_of_mass(arr, cutoff=0)
 

tme/backends/pytorch_backend.py

@@ -81,13 +81,13 @@ class PytorchBackend(NumpyFFTWBackend):
 
     def max(self, *args, **kwargs) -> NDArray:
         ret = self._array_backend.amax(*args, **kwargs)
-        if type(ret) == self._array_backend.Tensor:
+        if isinstance(ret, self._array_backend.Tensor):
             return ret
         return ret[0]
 
     def min(self, *args, **kwargs) -> NDArray:
         ret = self._array_backend.amin(*args, **kwargs)
-        if type(ret) == self._array_backend.Tensor:
+        if isinstance(ret, self._array_backend.Tensor):
             return ret
         return ret[0]
 
@@ -154,7 +154,7 @@ class PytorchBackend(NumpyFFTWBackend):
             1, -1
         )
         if unraveled_coords.size(0) == 1:
-            return tuple(unraveled_coords[0, :].tolist())
+            return (unraveled_coords[0, :],)
 
         else:
             return tuple(unraveled_coords.T)
@@ -206,7 +206,9 @@ class PytorchBackend(NumpyFFTWBackend):
         else:
             raise NotImplementedError("Operation only implemented for 2 and 3D inputs.")
 
-        pool = func(kernel_size=min_distance, return_indices=True)
+        pool = func(
+            kernel_size=min_distance, padding=min_distance // 2, return_indices=True
+        )
         _, indices = pool(score_space.reshape(1, 1, *score_space.shape))
         coordinates = self.unravel_index(indices.reshape(-1), score_space.shape)
         coordinates = self.transpose(self.stack(coordinates))
@@ -217,7 +219,7 @@ class PytorchBackend(NumpyFFTWBackend):
 
     def from_sharedarr(self, args) -> TorchTensor:
         if self.device == "cuda":
-            return args[0]
+            return args
 
         shm, shape, dtype = args
         required_size = int(self._array_backend.prod(self.to_backend_array(shape)))
@@ -235,13 +237,12 @@ class PytorchBackend(NumpyFFTWBackend):
 
         nbytes = arr.numel() * arr.element_size()
 
-        if type(shared_memory_handler) == SharedMemoryManager:
+        if isinstance(shared_memory_handler, SharedMemoryManager):
             shm = shared_memory_handler.SharedMemory(size=nbytes)
         else:
             shm = shared_memory.SharedMemory(create=True, size=nbytes)
 
         shm.buf[:nbytes] = arr.numpy().tobytes()
-
         return shm, arr.shape, arr.dtype
 
     def transpose(self, arr):
@@ -415,6 +416,8 @@ class PytorchBackend(NumpyFFTWBackend):
             yield None
 
     def device_count(self) -> int:
+        if self.device == "cpu":
+            return 1
         return self._array_backend.cuda.device_count()
 
     def reverse(self, arr: TorchTensor) -> TorchTensor:

tme/density.py

@@ -116,8 +116,8 @@ class Density:
         response = "Density object at {}\nOrigin: {}, Sampling Rate: {}, Shape: {}"
         return response.format(
             hex(id(self)),
-            tuple(np.round(self.origin, 3)),
-            tuple(np.round(self.sampling_rate, 3)),
+            tuple(round(float(x), 3) for x in self.origin),
+            tuple(round(float(x), 3) for x in self.sampling_rate),
             self.shape,
         )
 
@@ -306,6 +306,10 @@ class Density:
                 "std": float(mrc.header.rms),
             }
 
+        non_standard_crs = not np.all(crs_index == (0, 1, 2))
+        if non_standard_crs:
+            warnings.warn("Non standard MAPC, MAPR, MAPS, adapting data and origin.")
+
         if is_gzipped(filename):
             if use_memmap:
                 warnings.warn(
@@ -315,6 +319,10 @@ class Density:
             use_memmap = False
 
         if subset is not None:
+            subset = tuple(
+                subset[i] if i < len(subset) else slice(0, data_shape[i])
+                for i in crs_index
+            )
             subset_shape = tuple(x.stop - x.start for x in subset)
             if np.allclose(subset_shape, data_shape):
                 return cls._load_mrc(
@@ -328,18 +336,16 @@ class Density:
                 dtype=data_type,
                 header_size=1024 + extended_header,
             )
-            return data, origin, sampling_rate, metadata
-
-        if not use_memmap:
+        elif subset is None and not use_memmap:
             with mrcfile.open(filename, header_only=False) as mrc:
                 data = mrc.data.astype(np.float32, copy=False)
         else:
             with mrcfile.mrcmemmap.MrcMemmap(filename, header_only=False) as mrc:
                 data = mrc.data
 
-        if not np.all(crs_index == (0, 1, 2)):
+        if non_standard_crs:
             data = np.transpose(data, crs_index)
-            start = np.take(start, crs_index)
+            origin = np.take(origin, crs_index)
 
         return data, origin, sampling_rate, metadata
 
@@ -738,9 +744,16 @@ class Density:
         >>> )
 
         :py:meth:`Density.from_structure` supports a variety of methods to convert
-        atoms into densities. In addition to 'atomic_weight', 'atomic_number',
-        and 'van_der_waals_radius', its possible to use experimentally determined
-        scattering factors from various sources:
+        atoms into densities
+
+        >>> density = Density.from_structure(
+        >>>    filename_or_structure = path_to_structure,
+        >>>    weight_type = "gaussian",
+        >>>    weight_type_args={"resolution": "20"}
+        >>> )
+
+        In addition its possible to use experimentally determined scattering factors
+        from various sources:
 
         >>> density = Density.from_structure(
         >>>    filename_or_structure = path_to_structure,
@@ -748,7 +761,7 @@ class Density:
         >>>    weight_type_args={"source": "dt1969"}
         >>> )
 
-        or a lowpass filtered representation introduced in [1]_:
+        and their lowpass filtered representation introduced in [1]_:
 
         >>> density = Density.from_structure(
         >>>    filename_or_structure = path_to_structure,
@@ -873,6 +886,7 @@ class Density:
             mrc.header.nzstart, mrc.header.nystart, mrc.header.nxstart = np.rint(
                 np.divide(self.origin, self.sampling_rate)
             )
+            mrc.header.origin = tuple(x for x in self.origin)
             # mrcfile library expects origin to be in xyz format
             mrc.header.mapc, mrc.header.mapr, mrc.header.maps = (1, 2, 3)
             mrc.header["origin"] = tuple(self.origin[::-1])
@@ -1594,7 +1608,7 @@ class Density:
         rotation_matrix: NDArray,
         translation: NDArray = None,
         order: int = 3,
-        use_geometric_center: bool = False,
+        use_geometric_center: bool = True,
     ) -> "Density":
         """
         Performs a rigid transform of the class instance.

tme/external/bindings.cpp

@@ -0,0 +1,332 @@
+/*  Pybind extensions for template matching score space analyzers.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+*/
+
+#include <vector>
+#include <iostream>
+#include <limits>
+
+#include <pybind11/stl.h>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+template <typename T>
+void absolute_minimum_deviation(
+    py::array_t<T, py::array::c_style> coordinates,
+    py::array_t<T, py::array::c_style> output) {
+    auto coordinates_data = coordinates.data();
+    auto output_data = output.mutable_data();
+    int n = coordinates.shape(0);
+    int k = coordinates.shape(1);
+    int ik, jk, in, jn;
+
+    for (int i = 0; i < n; ++i) {
+        ik = i * k;
+        in = i * n;
+        for (int j = i + 1; j < n; ++j) {
+            jk = j * k;
+            jn = j * n;
+            T min_distance = std::abs(coordinates_data[ik] - coordinates_data[jk]);
+            for (int p = 1; p < k; ++p) {
+                min_distance = std::min(min_distance,
+                    std::abs(coordinates_data[ik + p] - coordinates_data[jk + p]));
+            }
+            output_data[in + j] = min_distance;
+            output_data[jn + i] = min_distance;
+        }
+        output_data[in + i] = 0;
+    }
+}
+
+template <typename T>
+std::pair<double, std::pair<int, int>> max_euclidean_distance(
+    py::array_t<T, py::array::c_style> coordinates) {
+    auto coordinates_data = coordinates.data();
+    int n = coordinates.shape(0);
+    int k = coordinates.shape(1);
+
+    double distance = 0.0;
+    double difference = 0.0;
+    double max_distance = -1;
+    double squared_distances = 0.0;
+
+    int ik, jk;
+    int max_i = -1, max_j = -1;
+
+    for (int i = 0; i < n; ++i) {
+        ik = i * k;
+        for (int j = i + 1; j < n; ++j) {
+            jk = j * k;
+            squared_distances = 0.0;
+            for (int p = 0; p < k; ++p) {
+                difference = static_cast<double>(
+                    coordinates_data[ik + p] - coordinates_data[jk + p]
+                );
+                squared_distances += (difference * difference);
+            }
+            distance = std::sqrt(squared_distances);
+            if (distance > max_distance) {
+                max_distance = distance;
+                max_i = i;
+                max_j = j;
+            }
+        }
+    }
+
+    return std::make_pair(max_distance, std::make_pair(max_i, max_j));
+}
+
+
+template <typename T>
+inline py::array_t<int, py::array::c_style> find_candidate_indices(
+    py::array_t<T, py::array::c_style> coordinates,
+    T min_distance) {
+    auto coordinates_data = coordinates.data();
+    int n = coordinates.shape(0);
+    int k = coordinates.shape(1);
+    int ik, jk;
+
+    std::vector<int> candidate_indices;
+    candidate_indices.reserve(n / 2);
+    candidate_indices.push_back(0);
+
+    for (int i = 1; i < n; ++i) {
+        bool is_candidate = true;
+        ik = i * k;
+        for (int candidate_index : candidate_indices) {
+            jk = candidate_index * k;
+            T distance = std::pow(coordinates_data[ik] - coordinates_data[jk], 2);
+            for (int p = 1; p < k; ++p) {
+                distance += std::pow(coordinates_data[ik + p] - coordinates_data[jk + p], 2);
+            }
+            distance = std::sqrt(distance);
+            if (distance <= min_distance) {
+                is_candidate = false;
+                break;
+            }
+        }
+        if (is_candidate) {
+            candidate_indices.push_back(i);
+        }
+    }
+
+    py::array_t<int, py::array::c_style> output({(int)candidate_indices.size()});
+    auto output_data = output.mutable_data();
+
+    for (size_t i = 0; i < candidate_indices.size(); ++i) {
+        output_data[i] = candidate_indices[i];
+    }
+
+    return output;
+}
+
+template <typename T>
+py::array_t<T, py::array::c_style> find_candidate_coordinates(
+    py::array_t<T, py::array::c_style> coordinates,
+    T min_distance) {
+
+    py::array_t<int, py::array::c_style> candidate_indices_array = find_candidate_indices(
+        coordinates, min_distance);
+    auto candidate_indices_data = candidate_indices_array.data();
+    int num_candidates = candidate_indices_array.shape(0);
+    int k = coordinates.shape(1);
+    auto coordinates_data = coordinates.data();
+
+    py::array_t<T, py::array::c_style> output({num_candidates, k});
+    auto output_data = output.mutable_data();
+
+    for (int i = 0; i < num_candidates; ++i) {
+        int candidate_index = candidate_indices_data[i] * k;
+        std::copy(
+            coordinates_data + candidate_index,
+            coordinates_data + candidate_index + k,
+            output_data + i * k
+        );
+    }
+
+    return output;
+}
+
+template <typename U, typename T>
+py::dict max_index_by_label(
+    py::array_t<U, py::array::c_style> labels,
+    py::array_t<T, py::array::c_style> scores
+    ) {
+
+    const U* labels_ptr = labels.data();
+    const T* scores_ptr = scores.data();
+
+    std::unordered_map<U, std::pair<T, ssize_t>> max_scores;
+
+    U label;
+    T score;
+    for (ssize_t i = 0; i < labels.size(); ++i) {
+        label = labels_ptr[i];
+        score = scores_ptr[i];
+
+        auto it = max_scores.insert({label, {score, i}});
+
+        if (score > it.first->second.first) {
+            it.first->second = {score, i};
+        }
+    }
+
+    py::dict ret;
+    for (auto& item: max_scores) {
+        ret[py::cast(item.first)] = py::cast(item.second.second);
+    }
+
+    return ret;
+}
+
+
+template <typename T>
+py::tuple online_statistics(
+    py::array_t<T, py::array::c_style> arr,
+    unsigned long long int n = 0,
+    double rmean = 0,
+    double ssqd = 0,
+    T reference = 0) {
+
+    auto in = arr.data();
+    int size = arr.size();
+
+    T max_value = std::numeric_limits<T>::lowest();
+    T min_value = std::numeric_limits<T>::max();
+    double delta, delta_prime;
+
+    unsigned long long int nbetter_or_equal = 0;
+
+    for(int i = 0; i < size; i++){
+        n++;
+        delta = in[i] - rmean;
+        rmean += delta / n;
+        delta_prime = in[i] - rmean;
+        ssqd += delta * delta_prime;
+
+        max_value = std::max(in[i], max_value);
+        min_value = std::min(in[i], min_value);
+        if (in[i] >= reference)
+            nbetter_or_equal++;
+    }
+
+    return py::make_tuple(n, rmean, ssqd, nbetter_or_equal, max_value, min_value);
+}
+
+PYBIND11_MODULE(extensions, m) {
+
+    m.def("absolute_minimum_deviation", absolute_minimum_deviation<double>,
+        "Compute pairwise absolute minimum deviation for a set of coordinates (float64).",
+        py::arg("coordinates"), py::arg("output"));
+    m.def("absolute_minimum_deviation", absolute_minimum_deviation<float>,
+        "Compute pairwise absolute minimum deviation for a set of coordinates (float32).",
+        py::arg("coordinates"), py::arg("output"));
+    m.def("absolute_minimum_deviation", absolute_minimum_deviation<int64_t>,
+        "Compute pairwise absolute minimum deviation for a set of coordinates (int64).",
+        py::arg("coordinates"), py::arg("output"));
+    m.def("absolute_minimum_deviation", absolute_minimum_deviation<int32_t>,
+        "Compute pairwise absolute minimum deviation for a set of coordinates (int32).",
+        py::arg("coordinates"), py::arg("output"));
+
+
+    m.def("max_euclidean_distance", max_euclidean_distance<double>,
+        "Identify pair of points with maximal euclidean distance (float64).",
+        py::arg("coordinates"));
+    m.def("max_euclidean_distance", max_euclidean_distance<float>,
+        "Identify pair of points with maximal euclidean distance (float32).",
+          py::arg("coordinates"));
+    m.def("max_euclidean_distance", max_euclidean_distance<int64_t>,
+        "Identify pair of points with maximal euclidean distance (int64).",
+          py::arg("coordinates"));
+    m.def("max_euclidean_distance", max_euclidean_distance<int32_t>,
+        "Identify pair of points with maximal euclidean distance (int32).",
+          py::arg("coordinates"));
+
+
+    m.def("find_candidate_indices", &find_candidate_indices<double>,
+          "Finds candidate indices with minimum distance (float64).",
+          py::arg("coordinates"), py::arg("min_distance"));
+    m.def("find_candidate_indices", &find_candidate_indices<float>,
+          "Finds candidate indices with minimum distance (float32).",
+          py::arg("coordinates"), py::arg("min_distance"));
+    m.def("find_candidate_indices", &find_candidate_indices<int64_t>,
+          "Finds candidate indices with minimum distance (int64).",
+          py::arg("coordinates"), py::arg("min_distance"));
+    m.def("find_candidate_indices", &find_candidate_indices<int32_t>,
+          "Finds candidate indices with minimum distance (int32).",
+          py::arg("coordinates"), py::arg("min_distance"));
+
+
+    m.def("find_candidate_coordinates", &find_candidate_coordinates<double>,
+          "Finds candidate coordinates with minimum distance (float64).",
+          py::arg("coordinates"), py::arg("min_distance"));
+    m.def("find_candidate_coordinates", &find_candidate_coordinates<float>,
+          "Finds candidate coordinates with minimum distance (float32).",
+          py::arg("coordinates"), py::arg("min_distance"));
+    m.def("find_candidate_coordinates", &find_candidate_coordinates<int64_t>,
+          "Finds candidate coordinates with minimum distance (int64).",
+          py::arg("coordinates"), py::arg("min_distance"));
+    m.def("find_candidate_coordinates", &find_candidate_coordinates<int32_t>,
+          "Finds candidate coordinates with minimum distance (int32).",
+          py::arg("coordinates"), py::arg("min_distance"));
+
+
+    m.def("max_index_by_label", &max_index_by_label<double, double>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<double, float>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<double, int64_t>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<double, int32_t>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+
+    m.def("max_index_by_label", &max_index_by_label<float, double>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<float, float>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<float, int64_t>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<float, int32_t>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+
+    m.def("max_index_by_label", &max_index_by_label<int64_t, double>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<int64_t, float>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<int64_t, int64_t>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<int64_t, int32_t>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+
+    m.def("max_index_by_label", &max_index_by_label<int32_t, double>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<int32_t, float>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<int32_t, int64_t>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+    m.def("max_index_by_label", &max_index_by_label<int32_t, int32_t>,
+          "Maximum value by label", py::arg("labels"), py::arg("scores"));
+
+
+    m.def("online_statistics", &online_statistics<double>, py::arg("arr"),
+          py::arg("n") = 0, py::arg("rmean") = 0,
+          py::arg("ssqd") = 0, py::arg("reference") = 0,
+          "Compute running online statistics on a numpy array.");
+    m.def("online_statistics", &online_statistics<float>, py::arg("arr"),
+          py::arg("n") = 0, py::arg("rmean") = 0,
+          py::arg("ssqd") = 0, py::arg("reference") = 0,
+          "Compute running online statistics on a numpy array.");
+    m.def("online_statistics", &online_statistics<int64_t>, py::arg("arr"),
+          py::arg("n") = 0, py::arg("rmean") = 0,
+          py::arg("ssqd") = 0, py::arg("reference") = 0,
+          "Compute running online statistics on a numpy array.");
+    m.def("online_statistics", &online_statistics<int32_t>, py::arg("arr"),
+          py::arg("n") = 0, py::arg("rmean") = 0,
+          py::arg("ssqd") = 0, py::arg("reference") = 0,
+          "Compute running online statistics on a numpy array.");
+}