diff --git a/patsy/state.py b/patsy/state.py index 8d674ba..046f6bb 100644 --- a/patsy/state.py +++ b/patsy/state.py @@ -105,7 +105,14 @@ def __init__(self): def memorize_chunk(self, x): x = atleast_2d_column_default(x) - self._count += x.shape[0] + if safe_issubdtype(x.dtype, np.inexact): + # Missing values are left out of the mean, column by column; rows + # containing them are dropped later by the NA action. + missing = np.isnan(x) + self._count += np.sum(~missing, 0) + x = np.where(missing, 0, x) + else: + self._count += x.shape[0] this_total = np.sum(x, 0, dtype=wide_dtype_for(x)) # This is to handle potentially multi-column x's: if self._sum is None: @@ -168,12 +175,20 @@ def memorize_chunk(self, x, center=True, rescale=True, ddof=0): if self.current_mean is None: self.current_mean = np.zeros(x.shape[1], dtype=wide_dtype_for(x)) self.current_M2 = np.zeros(x.shape[1], dtype=wide_dtype_for(x)) + # Missing values are left out of the mean and variance, column by + # column; rows containing them are dropped later by the NA action. + if safe_issubdtype(x.dtype, np.inexact): + present = ~np.isnan(x) + else: + present = np.ones(x.shape, dtype=bool) # XX this can surely be vectorized but I am feeling lazy: for i in range(x.shape[0]): - self.current_n += 1 - delta = x[i, :] - self.current_mean - self.current_mean += delta / self.current_n - self.current_M2 += delta * (x[i, :] - self.current_mean) + self.current_n = self.current_n + present[i, :] + delta = np.where(present[i, :], x[i, :] - self.current_mean, 0) + self.current_mean += delta / np.maximum(self.current_n, 1) + self.current_M2 += delta * np.where( + present[i, :], x[i, :] - self.current_mean, 0 + ) def memorize_finish(self): pass diff --git a/patsy/test_state.py b/patsy/test_state.py index 2c5a8e8..e009a9d 100644 --- a/patsy/test_state.py +++ b/patsy/test_state.py @@ -135,6 +135,20 @@ def test_Center(): check_stateful(Center, True, [1.3, -10.1, 7.0, 12.0], [-1.25, -12.65, 4.45, 9.45]) +def test_Center_missing_values(): + # NaNs are skipped column by column when computing the mean, and stay NaN + data = np.array([[1.0, 10.0], [np.nan, 20.0], [3.0, np.nan], [4.0, 30.0]]) + expected = np.array( + [[-1.666666, -10.0], [np.nan, 0.0], [0.333333, np.nan], [1.333333, 10.0]] + ) + for chunks in ([data], [data[:2], data[2:]]): + t = Center() + for chunk in chunks: + t.memorize_chunk(chunk) + t.memorize_finish() + assert np.allclose(t.transform(data), expected, equal_nan=True, atol=1e-6) + + def test_stateful_transform_wrapper(): assert np.allclose(center([1, 2, 3]), [-1, 0, 1]) assert np.allclose(center([1, 2, 1, 2]), [-0.5, 0.5, -0.5, 0.5]) @@ -205,3 +219,19 @@ def test_Standardize(): ddof=1, ) check_stateful(Standardize, True, r20, r20, center=False, rescale=False, ddof=1) + + +def test_Standardize_missing_values(): + # NaNs are skipped column by column when computing the mean and variance + data = np.array([[1.0, 10.0], [np.nan, 20.0], [3.0, np.nan], [4.0, 30.0]]) + expected = np.empty_like(data) + for j in range(data.shape[1]): + col = data[:, j] + observed = col[~np.isnan(col)] + expected[:, j] = (col - observed.mean()) / observed.std(ddof=1) + for chunks in ([data], [data[:2], data[2:]]): + t = Standardize() + for chunk in chunks: + t.memorize_chunk(chunk, ddof=1) + t.memorize_finish() + assert np.allclose(t.transform(data, ddof=1), expected, equal_nan=True)