Closes Bears-R-Us#3283: histogram2d between different dtypes (Bears-R…

…-Us#3763) * Adding mixed dtypes to Histogram2D, matching numpy outputs, and testing * adding atoimic changes and check testing for multidim --------- Co-authored-by: jaketrookman <jaketrookman@users.noreply.github.com>
stress-tess · Oct 4, 2024 · 1ed6d6b · 1ed6d6b
1 parent 2dffd06
commit 1ed6d6b
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 18 deletions.
diff --git a/src/Histogram.chpl b/src/Histogram.chpl
@@ -61,10 +61,10 @@ module Histogram
         return hist;
     }
 
-    proc histogramGlobalAtomic(x: [?aD] ?etype, y: [aD] etype, xMin: etype, xMax: etype, yMin: etype, yMax: etype, numXBins: int, numYBins: int, xBinWidth: real, yBinWidth: real) throws {
+    proc histogramGlobalAtomic(x: [?aD] ?etype1, y: [aD] ?etype2, xMin: etype1, xMax: etype1, yMin: etype2, yMax: etype2, numXBins: int, numYBins: int, xBinWidth: real, yBinWidth: real) throws {
         const totNumBins = numXBins * numYBins;
         var hD = makeDistDom(totNumBins);
-        var atomicHist: [hD] atomic int;
+        var atomicHist: [hD] atomic real;
 
         // count into atomic histogram
         forall (xi, yi) in zip(x, y) {
@@ -78,7 +78,7 @@ module Histogram
             atomicHist[(xiBin * numYBins) + yiBin].add(1);
         }
 
-        var hist = makeDistArray(totNumBins,int);
+        var hist = makeDistArray(totNumBins,real);
         // copy from atomic histogram to normal histogram
         [(e,ae) in zip(hist, atomicHist)] e = ae.read();
         try! hgLogger.debug(getModuleName(),getRoutineName(),getLineNumber(),
@@ -132,10 +132,10 @@ module Histogram
         return hist;
     }
 
-    proc histogramLocalAtomic(x: [?aD] ?etype, y: [aD] etype, xMin: etype, xMax: etype, yMin: etype, yMax: etype, numXBins: int, numYBins: int, xBinWidth: real, yBinWidth: real) throws {
+    proc histogramLocalAtomic(x: [?aD] ?etype1, y: [aD] ?etype2, xMin: etype1, xMax: etype1, yMin: etype2, yMax: etype2, numXBins: int, numYBins: int, xBinWidth: real, yBinWidth: real) throws {
         const totNumBins = numXBins * numYBins;
         // allocate per-locale atomic histogram
-        var atomicHist: [PrivateSpace] [0..#totNumBins] atomic int;
+        var atomicHist: [PrivateSpace] [0..#totNumBins] atomic real;
 
         // count into per-locale private atomic histogram
         forall (xi, yi) in zip(x, y) {
@@ -147,11 +147,11 @@ module Histogram
         }
 
         // +reduce across per-locale histograms to get counts
-        var lHist: [0..#totNumBins] int;
+        var lHist: [0..#totNumBins] real;
         forall i in PrivateSpace with (+ reduce lHist) do
           lHist reduce= atomicHist[i].read();
 
-        var hist = makeDistArray(totNumBins,int);
+        var hist = makeDistArray(totNumBins,real);
         hist = lHist;
         return hist;
     }
@@ -195,9 +195,9 @@ module Histogram
         return hist;
     }
 
-    proc histogramReduceIntent(x: [?aD] ?etype, y: [aD] etype, xMin: etype, xMax: etype, yMin: etype, yMax: etype, numXBins: int, numYBins: int, xBinWidth: real, yBinWidth: real) throws {
+    proc histogramReduceIntent(x: [?aD] ?etype1, y: [aD] ?etype2, xMin: etype1, xMax: etype1, yMin: etype2, yMax: etype2, numXBins: int, numYBins: int, xBinWidth: real, yBinWidth: real) throws {
         const totNumBins = numXBins * numYBins;
-        var gHist: [0..#totNumBins] int;
+        var gHist: [0..#totNumBins] real;
 
         // count into per-task/per-locale histogram and then reduce as tasks complete
         forall (xi, yi) in zip(x, y) with (+ reduce gHist) {
@@ -208,7 +208,7 @@ module Histogram
             gHist[(xiBin * numYBins) + yiBin] += 1;
         }
 
-        var hist = makeDistArray(totNumBins,int);
+        var hist = makeDistArray(totNumBins,real);
         hist = gHist;
         return hist;
     }

diff --git a/src/HistogramMsg.chpl b/src/HistogramMsg.chpl
@@ -99,9 +99,9 @@ module HistogramMsg
         var yGenEnt: borrowed GenSymEntry = getGenericTypedArrayEntry(yName, st);
 
         // helper nested procedure
-        proc histogramHelper(type t) throws {
-            var x = toSymEntry(xGenEnt,t);
-            var y = toSymEntry(yGenEnt,t);
+        proc histogramHelper(type t1, type t2) throws {
+            var x = toSymEntry(xGenEnt,t1);
+            var y = toSymEntry(yGenEnt,t2);            
             var xMin = min reduce x.a;
             var xMax = max reduce x.a;
             var yMin = min reduce y.a;
@@ -132,9 +132,15 @@ module HistogramMsg
             }
         }
         select (xGenEnt.dtype, yGenEnt.dtype) {
-            when (DType.Int64, DType.Int64)   {histogramHelper(int);}
-            when (DType.UInt64, DType.UInt64)  {histogramHelper(uint);}
-            when (DType.Float64, DType.Float64) {histogramHelper(real);}
+            when (DType.Int64, DType.Int64) {histogramHelper(int, int);}
+            when (DType.Int64, DType.UInt64) {histogramHelper(int, uint);}
+            when (DType.Int64, DType.Float64) {histogramHelper(int, real);}
+            when (DType.UInt64, DType.Int64) {histogramHelper(uint, int);}
+            when (DType.UInt64, DType.UInt64) {histogramHelper(uint, uint);}
+            when (DType.UInt64, DType.Float64) {histogramHelper(uint, real);}
+            when (DType.Float64, DType.Int64) {histogramHelper(real, int);}
+            when (DType.Float64, DType.UInt64) {histogramHelper(real, uint);}
+            when (DType.Float64, DType.Float64) {histogramHelper(real, real);}
             otherwise {
                 var errorMsg = notImplementedError(pn,"("+dtype2str(xGenEnt.dtype)+","+dtype2str(yGenEnt.dtype)+")");
                 hgmLogger.error(getModuleName(),getRoutineName(),getLineNumber(),errorMsg);

diff --git a/tests/numpy/numeric_test.py b/tests/numpy/numeric_test.py
@@ -314,10 +314,12 @@ def test_histogram(self, num_type):
 
     @pytest.mark.skipif(host == "horizon", reason="Fails on horizon")
     @pytest.mark.skip_if_max_rank_less_than(2)
-    def test_histogram_multidim(self):
+    @pytest.mark.parametrize("num_type1", NO_BOOL)
+    @pytest.mark.parametrize("num_type2", NO_BOOL)
+    def test_histogram_multidim(self, num_type1, num_type2):
         # test 2d histogram
         seed = 1
-        ak_x, ak_y = ak.randint(1, 100, 1000, seed=seed), ak.randint(1, 100, 1000, seed=seed + 1)
+        ak_x, ak_y = ak.randint(1, 100, 1000, seed=seed, dtype=num_type1), ak.randint(1, 100, 1000, seed=seed + 1, dtype=num_type2)
         np_x, np_y = ak_x.to_ndarray(), ak_y.to_ndarray()
         np_hist, np_x_edges, np_y_edges = np.histogram2d(np_x, np_y)
         ak_hist, ak_x_edges, ak_y_edges = ak.histogram2d(ak_x, ak_y)
@@ -431,6 +433,7 @@ def test_arctan2(self, num_type, denom_type):
             ak.arctan2(pda_num, pda_denom, where=True).to_ndarray(),
             equal_nan=True,
         )
+
         assert np.allclose(
             np.arctan2(na_num[0], na_denom, where=True),
             ak.arctan2(pda_num[0], pda_denom, where=True).to_ndarray(),