Mediana de movimento mais rápido em numpy

A implementação do Scipy não parece tão eficiente, então escrevi uma em Cython para verificar se eu poderia escrever uma mais rápida. Minha implementação em Cython é muito mais rápida do que o código inicial e também supera o Scipy.

Aqui está o código Cython ( median.pyxarquivo):

# cython: language_level=3
# cython: boundscheck=False
# cython: wraparound=False
# distutils: language=c++
import numpy as np
from numpy cimport int64_t, float64_t
from libcpp.set cimport multiset
import cython
from cython.operator import dereference as deref
from cython.operator import preincrement as inc
from cython.operator import predecrement as dec
from libcpp cimport bool as bool_t


@cython.cfunc
def value_of(it : multiset[float64_t].iterator, n : int64_t) -> float64_t:
    if n % 2 == 0:
        it2 : multiset[float64_t].iterator = it
        inc(it2)
        val1 : float64_t = deref(it)
        val2 : float64_t = deref(it2)
        return (val1 + val2) * 0.5
    return deref(it) * 0.5


def sliding_median(arr : float64_t[:], n : int64_t):
    assert n >= 3
    assert arr.size >= n

    out = np.empty(arr.size - n + 1, dtype=np.float64)
    out_view : float64_t[::1] = out
    sorted_items : multiset[float64_t] = multiset[float64_t]()
    i : int64_t = 0

    for i in range(n):
        sorted_items.insert(arr[i])

    mid1 : int64_t = (n - 1) // 2
    mid2 : int64_t = n // 2
    it : multiset[float64_t].iterator = sorted_items.begin()

    # Find the middle. Note we cannot index the ith item directly even 
    # though the implementation of std::multiset can certainly support it...
    for i in range(mid1):
        inc(it)

    out_view[0] = value_of(it, n)
    arr_size : int64_t = arr.size

    i = n
    while i < arr_size:
        old_value : float64_t = arr[i-n]
        new_value : float64_t = arr[i]
        it_value : float64_t = deref(it)
        old_it : multiset[float64_t].iterator = sorted_items.find(old_value)
        delayed_inc : bool_t = False
        delayed_dec : bool_t = False
        delayed_inc = old_value < it_value < new_value
        # Note: the median might not be accurate if there are multiple equal FP value equal to it
        delayed_dec = new_value < it_value < old_value or old_it == it and new_value < it_value
        assert new_value != it_value # Pathological cases not supported yet (multiple FP items equal to the median)
        # If the current median item is must be removed
        if old_it == it:
            inc(it) # avoid iterator invalidation
            if new_value < it_value or it_value < new_value < deref(it):
                delayed_dec = True
        sorted_items.erase(old_it)
        sorted_items.insert(new_value)
        # Must be done after the insert
        if delayed_inc: inc(it)
        if delayed_dec: dec(it)
        out_view[i-n+1] = value_of(it, n)
        i += 1

    return out

Aqui está o código para construí-lo ( setup.pyarquivo):

from setuptools import setup
from Cython.Build import cythonize
import numpy as np

setup(ext_modules=cythonize('median.pyx'), include_dirs=[np.get_include()])

Aqui está a função modificada para usar o código Cython muito mais rápido:

def moving_median_fast(signal, n=150):
    b = median.sliding_median(np.array(signal), n)
    b_full = np.concatenate([[np.nanmedian(signal)]*(n-1), b]) # Prepend first `n-1` values unchanged
    return signal - b_full

Benchmark e notas

Aqui estão os resultados de desempenho da minha CPU i5-9600KF com Numpy 2.1.3 e Scipy 1.15.2:

moving_median:           6.11 s
moving_median_fast:      0.45 s   <-----

SciPy solution:          2.02 s   (see note)
median.sliding_median:   0.20 s   (see note)

Esta implementação baseada em Cython é cerca de 13,6 vezes mais rápida que o código inicial e 4,5 vezes mais rápida que o código Scipy!

Na verdade, a solução SciPy não opera diretamente no dataframe, não precisa converter a lista para arrays Numpy e não computa np.nanmedian(signal)nenhum dos dois (sem mencionar o np.concatenate). Portanto, não deve ser comparado a moving_mediannem moving_median_fast, mas apenas a median.sliding_median.

median.sliding_mediané 10 vezes mais rápido que o código SciPy. A maior parte do tempo de moving_median_fasté, na verdade, overheads. Se você quiser reduzir tais overheads, então você deve usar matrizes Numpy no dataframe. Se todas as listas forem do mesmo tamanho (suposição feita pelo código SciPy), você deve simplesmente usar uma matriz 2D em vez de um dataframe contendo matrizes Numpy ou listas Python puras ineficientes. Você também deve usar Cython para otimizar ainda mais a linha np.concatenate([[np.nanmedian(signal)]*(n-1), b])(pré-alocando a matriz de saída para o tamanho certo e escrevendo diretamente o resultado de np.nanmediannela, e também multiplicando valores com um loop no local, etc.). Tudo isso certamente tornaria moving_median_fast>20 vezes mais rápido que moving_median.

Se isso não for suficiente, você pode usar vários threads com joblib. Observe que o GIL precisa ser liberado na função Cython (usando with nogilor @cython.nogil) para que ele possa escalar bem. Observe que a paralelização só torna o código significativamente mais rápido se você usar matrizes 2D (já que conversões de lista exigem que o GIL esteja habilitado).

Você não menciona ter NaNs em seus dados, então vou assumir que não. Nesse caso, acho que isso é o melhor que o SciPy tem a oferecer:

import numpy as np
from scipy.ndimage import median_filter
rng = np.random.default_rng(49825498549428354)
data = rng.random(size=(2000, 750))  # 2000 signals, each of length 750 
res = median_filter(data, size=(150,), axes=(-1,))  # moving window of size 150 

Acredito que houve algum trabalho recente feito para tornar isso mais rápido na versão de desenvolvimento do SciPy (nightly wheels aqui ) do que antes. Estou supondo que, em vez de reordenar cada janela do zero, ele atualiza uma estrutura de dados classificada ou particionada com base nos valores de entrada e saída, mas eu realmente não pesquisei sobre isso.

Observe as várias modeopções na documentação que controlam o que acontece no limite. Se você estiver satisfeito em obter de volta um array menor que o original em vez da condição de limite "reflect" padrão, você pode querer usar o padrão modee aparar as bordas depois.

Se você tiver NaNs, o SciPy tem um novo vectorized_filterque funcionará com o np.nanmedian, mas ele só usa stride_tricks.sliding_window_viewrecursos ocultos, então é improvável que seja mais rápido do que o que você tem.

Se CuPy for uma opção, me avise, e eu posso sugerir algo muito mais rápido. O SciPy median_filterainda não é tão rápido quanto deveria ser.

Tenho trabalhado em um projeto chamado ndfilters, que foi criado para calcular filtros multidimensionais mais rápido que o Scipy em máquinas multicore usando o compilador Numba com paralelismo habilitado.

Você pode instalar ndiltersusando pip,

pip install ndfilters

Implementei o filtro mediano ndfilters.median_filter()e você pode usá-lo de uma forma muito semelhante ao Scipy,

import numpy as np
import ndfilters

data = np.random.random(size=(2000, 750))

res = ndfilters.median_filter(data, size=150, axis=-1)

Como o Numba é um compilador just-in-time, essa função será muito lenta na primeira vez, pois precisa ser compilada.