Bug/median

CHANGELOG.md

@@ -19,7 +19,8 @@
 - [#614](https://github.com/helmholtz-analytics/heat/pull/614) New feature: printing of DNDarrays and ``__repr__`` and ``__str__`` functions
 - [#615](https://github.com/helmholtz-analytics/heat/pull/615) New feature: `skew()`
 - [#615](https://github.com/helmholtz-analytics/heat/pull/615) New feature: `kurtosis()`
-- [#618](https://github.com/helmholtz-analytics/heat/pull/618) Printing of unbalnced DNDarrays added
+- [#618](https://github.com/helmholtz-analytics/heat/pull/618) Printing of unbalanced DNDarrays added
+- [#624](https://github.com/helmholtz-analytics/heat/pull/624) Bugfix: distributed median() indexing and casting
 
 # v0.4.0
 

heat/core/statistics.py

@@ -1323,17 +1323,17 @@ def local_percentile(data, axis, indices):
 
     # sanitize q
     if isinstance(q, list) or isinstance(q, tuple):
-        t_perc_dtype = torch.promote_types(type(q[0]), t_x.dtype)
+        t_perc_dtype = torch.promote_types(type(q[0]), torch.float32)
         t_q = torch.tensor(q, dtype=t_perc_dtype, device=t_x.device)
     elif np.isscalar(q):
-        t_perc_dtype = torch.promote_types(type(q), t_x.dtype)
+        t_perc_dtype = torch.promote_types(type(q), torch.float32)
         t_q = torch.tensor([q], dtype=t_perc_dtype, device=t_x.device)
     elif isinstance(q, dndarray.DNDarray):
         if x.comm.is_distributed() and q.split is not None:
             # q needs to be local
             q.resplit_(axis=None)
         t_q = q._DNDarray__array
-        t_perc_dtype = torch.promote_types(t_q.dtype, t_x.dtype)
+        t_perc_dtype = torch.promote_types(t_q.dtype, torch.float32)
     else:
         raise TypeError("DNDarray, list or tuple supported, but q was {}".format(type(q)))
 
@@ -1407,8 +1407,8 @@ def local_percentile(data, axis, indices):
 
         if split == axis:
             # map percentile location: which q on what rank
-            t_indices_map = torch.ones((size, nperc), dtype=t_q.dtype, device=t_q.device) * -1.0
-            t_local_indices = torch.ones((1, nperc), dtype=t_q.dtype, device=t_q.device) * -1.0
+            t_indices_map = torch.ones((size, nperc), dtype=t_indices.dtype, device=t_q.device) * -1
+            t_local_indices = torch.ones((1, nperc), dtype=t_indices.dtype, device=t_q.device) * -1
             offset, _, chunk = x.comm.chunk(gshape, split)
             chunk_start = chunk[split].start
             chunk_stop = chunk[split].stop
@@ -1429,15 +1429,19 @@ def local_percentile(data, axis, indices):
         data.get_halo(1)
         t_data = data.array_with_halos
         # fill out percentile
+        t_ind_on_rank -= offset
         t_map_sum = t_indices_map.sum(axis=1)
         perc_ranks = torch.where(t_map_sum > -1 * nperc)[0].tolist()
-        for r in perc_ranks:
+        for r_id, r in enumerate(perc_ranks):
             # chunk of the global percentile that will be populated by rank r
-            _, _, perc_chunk = x.comm.chunk(output_shape, join, rank=r, w_size=len(perc_ranks))
+            _, _, perc_chunk = x.comm.chunk(output_shape, join, rank=r_id, w_size=len(perc_ranks))
             perc_slice = perc_slice[:join] + (perc_chunk[join],) + perc_slice[join + 1 :]
             local_p = factories.zeros(percentile[perc_slice].shape, dtype=perc_dtype, comm=x.comm)
             if rank == r:
-                local_p = factories.array(local_percentile(t_data, axis, t_ind_on_rank - offset))
+                if rank > 0:
+                    # correct indices for halo
+                    t_ind_on_rank += 1
+                local_p = factories.array(local_percentile(t_data, axis, t_ind_on_rank))
             x.comm.Bcast(local_p, root=r)
             percentile[perc_slice] = local_p
     else:

heat/core/tests/test_statistics.py

@@ -971,7 +971,13 @@ def test_percentile(self):
         q = 100
         p_np = np.percentile(x_np, q, axis=0)
         p_ht = ht.percentile(x_ht, q, axis=0)
-        self.assertAlmostEqual(p_ht.numpy().all(), p_np.all())
+        self.assertTrue((p_ht.numpy() == p_np).all())
+
+        # test median (q = 50)
+        q = 50
+        p_np = np.percentile(x_np, q, axis=0)
+        p_ht = ht.median(x_ht_split0, axis=0)
+        self.assertTrue((p_ht.numpy() == p_np).all())
 
         # test list q and writing to output buffer
         q = [0.1, 2.3, 15.9, 50.0, 84.1, 97.7, 99.9]