From fe22bde99873458dcbbe6e08e4ff1a9233e12c76 Mon Sep 17 00:00:00 2001
From: Scott Sanderson <ssanderson@quantopian.com>
Date: Mon, 28 Mar 2016 11:34:58 -0400
Subject: [PATCH] TEST: Test uneven buckets in quantiles.

---
 tests/pipeline/test_factor.py | 41 +++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/tests/pipeline/test_factor.py b/tests/pipeline/test_factor.py
index 2b9783aa..de30db62 100644
--- a/tests/pipeline/test_factor.py
+++ b/tests/pipeline/test_factor.py
@@ -811,6 +811,47 @@ class FactorTestCase(BasePipelineTestCase):
         for key, (res, exp) in dzip_exact(results, expected).items():
             check_arrays(res, exp)
 
+    def test_quantiles_uneven_buckets(self):
+        permute = partial(permute_rows, 5)
+        shape = (5, 5)
+
+        factor_data = permute(log1p(arange(25, dtype=float).reshape(shape)))
+        mask_data = permute(self.eye_mask(shape=shape))
+
+        f = F()
+        m = Mask()
+
+        terms = {
+            '3_masked': f.quantiles(bins=3, mask=m),
+            '7_masked': f.quantiles(bins=20, mask=m),
+        }
+
+        expected = {
+            '3_masked': [[-1, 0,  0,  1,  2],
+                         [0, -1,  0,  1,  2],
+                         [0,  0, -1,  1,  2],
+                         [0,  0,  1, -1,  2],
+                         [0,  0,  1,  2, -1]],
+            '7_masked': [[-1, 0,  2,  4,  6],
+                         [0, -1,  2,  4,  6],
+                         [0,  2, -1,  4,  6],
+                         [0,  2,  4, -1,  6],
+                         [0,  2,  4,  6, -1]],
+        }
+
+        graph = TermGraph(terms)
+        results = self.run_graph(
+            graph,
+            initial_workspace={
+                f: factor_data,
+                m: mask_data,
+            },
+            mask=self.build_mask(self.ones_mask(shape=shape)),
+        )
+
+        for key, (res, exp) in dzip_exact(results, expected).items():
+            check_arrays(res, exp)
+
     def test_quantile_helpers(self):
         f = self.f
         m = Mask()