In one of my personal python libraries I have a custom class used for computing the running average and variance of a stream of numbers:
import numpy as np
class RunningStatsVariable:
def __init__(self, ddof=0, parallel=None):
self.mean = 0
self.var = 0
self.std = 0
self._n = 0
self._s = 0
self._ddof = ddof
if parallel == 'multiprocessing':
from multiprocessing import Lock
self._lock = Lock()
elif parallel == 'threading':
from threading import Lock
self._lock = Lock()
else:
self._lock = None
def update(self, values):
if self._lock:
self._lock.acquire()
values = np.array(values, ndmin=1)
n = len(values)
self._n += n
delta = values - self.mean
self.mean += (delta / self._n).sum()
self._s += (delta * (values - self.mean)).sum()
self.var = self._s / (self._n - self._ddof) if self._n > self._ddof else 0
self.std = np.sqrt(self.var)
if self._lock:
self._lock.release()
def update_single(self, value):
if self._lock:
self._lock.acquire()
self._n += 1
old_mean = self.mean
self.mean += (value - old_mean) / self._n
self._s += (value - old_mean) * (value - self.mean)
self.var = self._s / (self._n - self._ddof) if self._n > self._ddof else 0
self.std = np.sqrt(self.var)
if self._lock:
self._lock.release()
def __str__(self):
if self.std:
return f"(μ ± σ): {self.mean} ± {self.std}"
else:
return f"{self.name}: {self.mean}"
def __len__(self):
return self._n
The idea is to make the class safe for parallel updates, regardless of the type of parallelism used (threading or multiprocessing), and regardless of which method of updating (single-value or list) is used. The code is also available on Replit. The below test seems to show this code works as expected:
data = [(np.random.random(), np.random.rand(20)) for _ in range(1000)]
def _update(var, value, arr):
var.update_single(value)
var.update(arr)
v1 = RunningStatsVariable()
for v, a in data:
_update(v1, v, a)
from joblib import Parallel, delayed
v2 = RunningStatsVariable(parallel='multiprocessing')
Parallel(n_jobs=-1)(
delayed(_update)(v2, v, a)
for v, a in data
)
v3 = RunningStasVariable(parallel='threading')
Parallel(n_jobs=-1, prefer='threads')(
delayed(_update)(v3, v, a)
for v, a in data
)
# Can't use x == y == z for float comparison, so we use np.allclose instead
print(np.allclose(v1.mean, [v2.mean, v3.mean])) # True
print(np.allclose(v1.var, [v2.var, v3.var])) # True
print(np.allclose(v1.std, [v2.std, v3.std]))])) # True
However, in parallelism testing is usually not enough to guarantee correctness. So are the updates really parallel-safe? I don't necessarily care about the parallel-safety of the __len__ and __str__ methods.
Also, is locking at the start of the update method and unlocking at the end good practice or should I use more specific locks?
Additionally, I would like this code to work for other people, who may not be using joblib for their parallelism needs. Are there other lock types I should consider?
