diff --git a/heat/core/statistics.py b/heat/core/statistics.py index 3ce8dd4791..58f627cec5 100644 --- a/heat/core/statistics.py +++ b/heat/core/statistics.py @@ -1091,13 +1091,12 @@ def median( DNDarray.median.__doc__ = median.__doc__ -def __merge_moments( - m1: torch.Tensor, m2: torch.Tensor, correction: bool = True -) -> Tuple[torch.Tensor, ...]: +def __merge_moments(m1: Tuple, m2: Tuple) -> Tuple[torch.Tensor, ...]: """ Merge two statistical moments. - If the length of ``m1`` and ``m2`` (must be equal) is ``==3`` then the second moment (variance) - is merged. This function can be expanded to merge other moments according to Reference [1] as well. + If the length of ``m1`` and ``m2`` (must be equal) is ``==3`` then the second moment (sum of squared + differences between sample data and estimated mean) is merged. This function can be expanded to merge + other moments according to Reference [1] as well. Note: all arrays must be either the same size or individual values Parameters @@ -1108,8 +1107,6 @@ def __merge_moments( m2 : Tuple Tuple of the moments to merge together, the 0th element is the moment to be merged. The tuple must be sorted in descending order of moments - correction : bool - Flag for the use of unbiased estimators (when available) References ---------- @@ -1120,6 +1117,8 @@ def __merge_moments( if len(m1) != len(m2): raise ValueError(f"m1 and m2 must be same length, currently {len(m1)} and {len(m2)}") n1, n2 = m1[-1], m2[-1] + if n2.sum() == 0: + return m1 mu1, mu2 = m1[-2], m2[-2] n = n1 + n2 delta = mu2 - mu1 @@ -1127,14 +1126,12 @@ def __merge_moments( if len(m1) == 2: # merge means return mu, n - var1, var2 = m1[-3], m2[-3] - if correction: - var_m = (var1 * (n1 - 1) + var2 * (n2 - 1) + (delta**2) * n1 * n2 / n) / (n - 1) - else: - var_m = (var1 * n1 + var2 * n2 + (delta**2) * n1 * n2 / n) / n + M2_1, M2_2 = m1[-3], m2[-3] + # this is formula (II.4) in [1]: + M2 = M2_1 + M2_2 + (delta**2) * n1 * n2 / n - if len(m1) == 3: # merge vars - return var_m, mu, n + if len(m1) == 3: # merge second moments + return M2, mu, n # TODO: This code block can be added if skew or kurtosis support multiple axes: # sk1, sk2 = m1[-4], m2[-4] @@ -2190,28 +2187,30 @@ def reduce_vars_elementwise(output_shape_i: torch.Tensor) -> DNDarray: output_shape_i : iterable Iterable with the dimensions of the output of the var function. """ + n = float(x.lshape[x.split]) + if x.lshape[x.split] != 0: mu = torch.mean(x.larray, dim=axis) - var = torch.var(x.larray, dim=axis, correction=correction) + M2 = torch.var(x.larray, dim=axis, correction=False) * n else: mu = factories.zeros(output_shape_i, dtype=x.dtype, device=x.device) - var = factories.zeros(output_shape_i, dtype=x.dtype, device=x.device) + M2 = factories.zeros(output_shape_i, dtype=x.dtype, device=x.device) - var_shape = list(var.shape) if list(var.shape) else [1] + var_shape = list(M2.shape) if list(M2.shape) else [1] - var_tot = factories.zeros(([x.comm.size, 3] + var_shape), dtype=x.dtype, device=x.device) - var_tot[x.comm.rank, 0, :] = var - var_tot[x.comm.rank, 1, :] = mu - var_tot[x.comm.rank, 2, :] = float(x.lshape[x.split]) - x.comm.Allreduce(MPI.IN_PLACE, var_tot, MPI.SUM) + M2_tot = factories.zeros(([x.comm.size, 3] + var_shape), dtype=x.dtype, device=x.device) + M2_tot[x.comm.rank, 0, :] = M2 + M2_tot[x.comm.rank, 1, :] = mu + M2_tot[x.comm.rank, 2, :] = n + x.comm.Allreduce(MPI.IN_PLACE, M2_tot, MPI.SUM) for i in range(1, x.comm.size): - var_tot[0, 0, :], var_tot[0, 1, :], var_tot[0, 2, :] = __merge_moments( - (var_tot[0, 0, :], var_tot[0, 1, :], var_tot[0, 2, :]), - (var_tot[i, 0, :], var_tot[i, 1, :], var_tot[i, 2, :]), - correction=correction, + M2_tot[0, 0, :], M2_tot[0, 1, :], M2_tot[0, 2, :] = __merge_moments( + (M2_tot[0, 0, :], M2_tot[0, 1, :], M2_tot[0, 2, :]), + (M2_tot[i, 0, :], M2_tot[i, 1, :], M2_tot[i, 2, :]), ) - return var_tot[0, 0, :][0] if var_tot[0, 0, :].size == 1 else var_tot[0, 0, :] + var = M2_tot[0, 0, :] / (M2_tot[0, 2, :] - int(correction)) + return var[0] if var.size == 1 else var # ---------------------------------------------------------------------------------------------- if axis is None: # no axis given @@ -2222,27 +2221,26 @@ def reduce_vars_elementwise(output_shape_i: torch.Tensor) -> DNDarray: ) else: # case for full matrix calculation (axis is None) + n = x.lnumel mu_in = torch.mean(x.larray) - var_in = torch.var(x.larray, correction=correction) - # Nan is returned when local tensor is empty - if torch.isnan(var_in): - var_in = 0.0 + M2_in = torch.var(x.larray, correction=False) * n + # NaN is returned when local tensor is empty + if torch.isnan(M2_in): + M2_in = 0.0 if torch.isnan(mu_in): mu_in = 0.0 - n = x.lnumel - var_tot = factories.zeros((x.comm.size, 3), dtype=x.dtype, device=x.device) + M2_tot = factories.zeros((x.comm.size, 3), dtype=x.dtype, device=x.device) var_proc = factories.zeros((x.comm.size, 3), dtype=x.dtype, device=x.device) - var_proc[x.comm.rank] = var_in, mu_in, float(n) - x.comm.Allreduce(var_proc, var_tot, MPI.SUM) + var_proc[x.comm.rank] = M2_in, mu_in, float(n) + x.comm.Allreduce(var_proc, M2_tot, MPI.SUM) for i in range(1, x.comm.size): - var_tot[0, 0], var_tot[0, 1], var_tot[0, 2] = __merge_moments( - (var_tot[0, 0], var_tot[0, 1], var_tot[0, 2]), - (var_tot[i, 0], var_tot[i, 1], var_tot[i, 2]), - correction=correction, + M2_tot[0, 0], M2_tot[0, 1], M2_tot[0, 2] = __merge_moments( + (M2_tot[0, 0], M2_tot[0, 1], M2_tot[0, 2]), + (M2_tot[i, 0], M2_tot[i, 1], M2_tot[i, 2]), ) - return var_tot[0][0] + return M2_tot[0][0] / (M2_tot[0][2] - int(correction)) else: # axis is given return __moment_w_axis(torch.var, x, axis, reduce_vars_elementwise, correction) diff --git a/tests/core/test_statistics.py b/tests/core/test_statistics.py index fb43d92b93..c8b5e03b5b 100644 --- a/tests/core/test_statistics.py +++ b/tests/core/test_statistics.py @@ -1,6 +1,9 @@ +import unittest + import numpy as np import torch import os +import unittest from itertools import combinations from scipy import stats as ss @@ -1599,3 +1602,44 @@ def test_var(self): # edge case from #2374 self.assertEqual(ht.var(ht.array([0.], split=None), axis=0, ddof=0), 0) + + @unittest.skipUnless(ht.communication.MPI_WORLD.size >= 3, "Test requires at least 3 tasks") + def test_first_two_leading_ranks_empty(self): + comm = self.comm + local_data = torch.tensor([], dtype=torch.float32) if comm.rank < 2 else torch.tensor([comm.rank], dtype=torch.float32) + data = ht.DNDarray(local_data, gshape=(comm.size-2,), dtype=ht.float32, split=0, device=ht.devices.cpu, comm=comm, balanced=False) + + self.assertEqual(ht.mean(data), np.mean(data.numpy())) + self.assertEqual(ht.var(data), np.var(data.numpy())) + + @unittest.skipUnless(ht.communication.MPI_WORLD.size == 2, "Test for two tasks") + def test_corrected_var_single_element(self): + comm = self.comm + ltensor_empty = torch.tensor([], dtype=torch.float32) + ltensor_data = torch.tensor([0.], dtype=torch.float32) + + ldata_on_first_rank = ltensor_data if comm.rank == 0 else ltensor_empty + data_on_first_rank = ht.DNDarray(ldata_on_first_rank, gshape=(1,), dtype=ht.float32, split=0, device=ht.devices.cpu, comm=comm, balanced=None) + + ldata_on_second_rank = ltensor_empty if comm.rank == 0 else ltensor_data + data_on_second_rank = ht.DNDarray(ldata_on_second_rank, gshape=(1,), dtype=ht.float32, split=0, device=ht.devices.cpu, comm=comm, balanced=False) + + self.assertTrue(ht.isnan(ht.var(data_on_first_rank, ddof=1))) + self.assertTrue(ht.isnan(ht.var(data_on_first_rank, axis=0, ddof=1))) + + self.assertTrue(ht.isnan(ht.var(data_on_second_rank, ddof=1))) + self.assertTrue(ht.isnan(ht.var(data_on_second_rank, axis=0, ddof=1))) + + def test_corrected_var_ranks_with_single_elements(self): + comm = self.comm + + local_data = torch.tensor([], dtype=torch.float32) + if ht.communication.MPI_WORLD.size == 1: + local_data = torch.tensor([0., 1.], dtype=torch.float32) + elif comm.rank <= 1: + local_data = torch.tensor([comm.rank], dtype=torch.float32) + + data = ht.DNDarray(local_data, gshape=(2,), dtype=ht.float32, split=0, device=ht.devices.cpu, comm=comm, balanced=None) + + self.assertEqual(ht.var(data, ddof=1), np.var(data.numpy(), ddof=1)) + self.assertEqual(ht.var(data, axis=0, ddof=1), np.var(data.numpy(), axis=0, ddof=1))