From 1b01ab9f76092c38720e3b0affeb0f9793923481 Mon Sep 17 00:00:00 2001
From: Louis Potok <louispotok@gmail.com>
Date: Wed, 15 May 2019 16:33:25 -0700
Subject: [PATCH] add function and test for rolling median

---
 python/fbprophet/diagnostics.py            | 49 ++++++++++++++++++++++
 python/fbprophet/tests/test_diagnostics.py | 23 ++++++++++
 2 files changed, 72 insertions(+)

diff --git a/python/fbprophet/diagnostics.py b/python/fbprophet/diagnostics.py
index 5f54dca..5b4a2b5 100644
--- a/python/fbprophet/diagnostics.py
+++ b/python/fbprophet/diagnostics.py
@@ -317,6 +317,55 @@ def rolling_mean_by_h(x, h, w, name):
     return pd.DataFrame({'horizon': res_h, name: res_x})
 
 
+def rolling_median_by_h(x, h, w, name):
+    """Compute a rolling median of x, after first aggregating by h.
+
+    Right-aligned. Computes a single median for each unique value of h. Each
+    median is over at least w samples.
+
+    For each h where there are fewer than w samples, we take samples from the previous h,
+    moving backwards. (In other words, we ~ assume that the x's are shuffled within each h.)
+
+    Parameters
+    ----------
+    x: Array.
+    h: Array of horizon for each value in x.
+    w: Integer window size (number of elements).
+    name: Name for metric in result dataframe
+
+    Returns
+    -------
+    Dataframe with columns horizon and name, the rolling median of x.
+    """
+    # Aggregate over h
+    df = pd.DataFrame({'x': x, 'h': h})
+    grouped = df.groupby('h')
+    df2 = grouped.size().reset_index().sort_values('h')
+    hs = df2['h'].values
+
+    res_h = []
+    res_x = []
+    # Start from the right and work backwards
+    i = len(hs) - 1
+    while i >= 0:
+        h_i = hs[i]
+        xs = grouped.get_group(h_i).x.tolist()
+        j = i - 1
+        while (len(xs) < w) and (j >= 0):
+            # Include points from the previous horizon. All of them if still
+            # less than w, otherwise just enough to get to w.
+            xs.append(x[j])
+            j -= 1
+        if len(xs) < w:
+            # Ran out of horizons before enough points.
+            break
+        res_h.append(hs[i])
+        res_x.append(np.median(xs))
+        i -= 1
+    res_h.reverse()
+    res_x.reverse()
+    return pd.DataFrame({'horizon': res_h, name: res_x})
+
 
 # The functions below specify performance metrics for cross-validation results.
 # Each takes as input the output of cross_validation, and returns the statistic
diff --git a/python/fbprophet/tests/test_diagnostics.py b/python/fbprophet/tests/test_diagnostics.py
index f682b49..51f58c7 100644
--- a/python/fbprophet/tests/test_diagnostics.py
+++ b/python/fbprophet/tests/test_diagnostics.py
@@ -174,6 +174,29 @@ class TestDiagnostics(TestCase):
         self.assertTrue(np.allclose(np.array([7.]), df['horizon'].values))
         self.assertTrue(np.allclose(np.array([4.5]), df['x'].values))
 
+    def test_rolling_median(self):
+        x = np.arange(10)
+        h = np.arange(10)
+        df = diagnostics.rolling_median_by_h(x=x, h=h, w=1, name='x')
+        self.assertTrue(np.array_equal(x, df['x'].values))
+        self.assertTrue(np.array_equal(h, df['horizon'].values))
+
+        df = diagnostics.rolling_median_by_h(x, h, w=4, name='x')
+        x_true = x[3:] - 1.5
+        self.assertTrue(np.allclose(x_true, df['x'].values))
+        self.assertTrue(np.array_equal(np.arange(3, 10), df['horizon'].values))
+
+        h = np.array([1., 2., 3., 4., 4., 4., 4., 4., 7., 7.])
+        x_true = np.array([1.0, 5.0, 8.0])
+        h_true = np.array([3., 4., 7.])
+        df = diagnostics.rolling_median_by_h(x, h, w=3, name='x')
+        self.assertTrue(np.allclose(x_true, df['x'].values))
+        self.assertTrue(np.array_equal(h_true, df['horizon'].values))
+
+        df = diagnostics.rolling_median_by_h(x, h, w=10, name='x')
+        self.assertTrue(np.allclose(np.array([7.]), df['horizon'].values))
+        self.assertTrue(np.allclose(np.array([4.5]), df['x'].values))
+
     def test_copy(self):
         df = DATA_all.copy()
         df['cap'] = 200.