Consumo de memoria de la función NumPy para desviación estándar

Actualmente estoy usando los enlaces de Python de GDAL para trabajar en conjuntos de datos ráster bastante grandes (> 4 GB). Como cargarlos en la memoria de inmediato no es una solución factible para mí, los leo en bloques más pequeños y hago los cálculos pieza por pieza. Para evitar una nueva asignación para cada bloque leído, estoy usando elbuf_obj argumento (aquí) para leer los valores en una matriz NumPy preasignada. En un punto, tengo que calcular la media y la desviación estándar de todo el ráster. Naturalmente he usadonp.std para el cómputo Sin embargo, al perfilar el consumo de memoria de mi programa, me di cuenta de que con cada invocación denp.std adicionalmente se asigna y libera memoria.

Un ejemplo de trabajo mínimo que demuestra este comportamiento:

In [1]  import numpy as np
In [2]  a = np.random.rand(20e6)  # Approx. 150 MiB of memory
In [3]  %memit np.mean(a)
peak memory: 187.30 MiB, increment: 0.48 MiB
In [4]  %memit np.std(a)
peak memory: 340.24 MiB, increment: 152.91 MiB

Una búsqueda dentro del árbol fuente de NumPy en GitHub reveló que elnp.std la función invoca internamente_var funcionar desde_methods.py (aquí) En un punto_var calcula las desviaciones de la media y las resume. Por lo tanto, se crea una copia temporal de la matriz de entrada. La función esencialmente calcula la desviación estándar de la siguiente manera:

mu = sum(arr) / len(arr)
tmp = arr - mu
tmp = tmp * tmp
sd = np.sum(tmp) / len(arr)

Si bien este enfoque está bien para matrices de entrada más pequeñas, definitivamente no hay forma de ir para las más grandes. Como estoy usando bloques de memoria más pequeños como se mencionó anteriormente, esta copia adicional no es un problema innovador desde el punto de vista de la memoria en mi programa. Sin embargo, lo que me molesta es que para cada bloque se realiza y se libera una nueva asignación antes de leer el siguiente bloque.

¿Hay alguna otra función dentro de NumPy o SciPy que utilice un enfoque con consumo de memoria constante como el algoritmo Welford (Wikipedia) para el cálculo de una pasada de la media y la desviación estándar?

Otra forma de hacerlo sería implementar una versión personalizada del_var funcionar con un opcionalout argumento para un búfer preasignado (como NumPy ufuncs). Con este enfoque, la copia adicional no se eliminaría, pero al menos el consumo de memoria sería constante y se guardaría el tiempo de ejecución para las asignaciones en cada bloque.

EDITAR: Probó la implementación de Cython del algoritmo Welford como lo sugiere kezzos.

Implementación de Cython (modificado de kezzos):

cimport cython
cimport numpy as np
from libc.math cimport sqrt

@cython.boundscheck(False)
def iterative_approach(np.ndarray[np.float32_t, ndim=1] a):
    cdef long n = 0
    cdef float mean = 0
    cdef float M2 = 0
    cdef long i
    cdef float delta
    cdef float a_min = 10000000  # Must be set to Inf and -Inf for real cases
    cdef float a_max = -10000000
    for i in range(len(a)):
        n += 1
        delta = a[i] - mean
        mean += delta / n
        M2 += delta * (a[i] - mean)
        if a[i] < a_min:
            a_min = a[i]
        if a[i] > a_max:
            a_max = a[i]
    return a_min, a_max, mean, sqrt(M2 / (n - 1))

Implementación de NumPy (media y estándar pueden posiblemente calcularse en una función):

def vector_approach(a):
    return np.min(a), np.max(a), np.mean(a), np.std(a, ddof=1)

Resultados de la prueba utilizando un conjunto de datos aleatorios (tiempos en milisegundos, el mejor de 25):

----------------------------------
| Size |  Iterative |     Vector |
----------------------------------
|  1e2 |    0.00529 |    0.17149 |
|  1e3 |    0.02027 |    0.16856 |
|  1e4 |    0.17850 |    0.23069 |
|  1e5 |    1.93980 |    0.77727 |
|  1e6 |   18.78207 |    8.83245 |
|  1e7 |  180.04069 |  101.14722 |
|  1e8 | 1789.60228 | 1086.66737 |
----------------------------------

Parece que el enfoque iterativo que usa Cython es más rápido con conjuntos de datos más pequeños y el enfoque de vector NumPy (posiblemente acelerado SIMD) para conjuntos de datos más grandes con más de 10000 elementos. Todas las pruebas se realizaron con Python 2.7.9 y NumPy versión 1.9.2.

Tenga en cuenta que, en el caso real, las funciones superiores se utilizarían para calcular las estadísticas de un solo bloque del ráster. Las desviaciones estándar y los medios para todos los bloques deben combinarse con la metodología sugerida en Wikipedia (aquí) Tiene la ventaja de que no todos los elementos del ráster deben resumirse y, por lo tanto, evita el problema de desbordamiento del flotador (al menos hasta cierto punto).

Respuestas a la pregunta(2)

Su respuesta a la pregunta