Fixes tech.ml.dataset#254

cnuernber · Jun 14, 2021 · 297b247 · 297b247
1 parent 165f7e3
commit 297b247
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 13 deletions.
diff --git a/src/tech/v3/datatype/rolling.clj b/src/tech/v3/datatype/rolling.clj
@@ -163,6 +163,19 @@ tech.v3.datatype.rolling> (fixed-rolling-window-ranges 10 3 :right)
         (WindowRange. (+ idx padding) window-size)))))
 
 
+(defn- fixed-reader-rolling-window
+  [item window-size relative-window-position edge-mode window-fn options]
+  (let [n-elems (dtype-base/ecount item)]
+    (->> (window-ranges->window-reader
+          item (fixed-rolling-window-ranges n-elems window-size
+                                            relative-window-position)
+          edge-mode)
+         (emap/emap window-fn
+                    (:datatype options
+                               (packing/unpack-datatype
+                                (dtype-base/elemwise-datatype item)))))))
+
+
 (defn fixed-rolling-window
   "Return a lazily evaluated rolling window of window-fn applied to each window.  The
   iterable or sequence is padded such that there are the same number of values in the
@@ -207,19 +220,11 @@ user>
    (vectorized-dispatch-1
     (fn [_] (throw (ex-info "Rolling windows aren't defined on scalars" {})))
     (fn [_res-dtype item]
-      (->> (pad-sequence (quot (long window-size) 2) item)
-           (fixed-window-sequence window-size 1)
-           (map window-fn)))
+      (fixed-reader-rolling-window (vec item) window-size relative-window-position
+                                   edge-mode window-fn options))
     (fn [_result-dtype item]
-      (let [n-elems (dtype-base/ecount item)]
-        (->> (window-ranges->window-reader
-              item (fixed-rolling-window-ranges n-elems window-size
-                                                relative-window-position)
-              edge-mode)
-             (emap/emap window-fn
-                        (:datatype options
-                                   (packing/unpack-datatype
-                                    (dtype-base/elemwise-datatype item)))))))
+      (fixed-reader-rolling-window item window-size relative-window-position
+                                   edge-mode window-fn options))
     nil
     item))
   ([item window-size window-fn]

diff --git a/test/tech/v3/datatype_test.clj b/test/tech/v3/datatype_test.clj
@@ -603,7 +603,23 @@
                                      {:relative-window-position :right})))
   (is (= [0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19]
          (rolling/fixed-rolling-window (range 20) 5 (fn [w] (last w))
-                                       {:relative-window-position :left} ))))
+                                       {:relative-window-position :left}))))
+
+
+(deftest rolling-window-position-seq
+  (is (dfn/equals [0 1 3 6 10 15 20 25 30 35 40 45 50 55 60 65 70 75 80 85]
+                  (rolling/fixed-rolling-window (apply list (range 20)) 5 dfn/sum
+                                                {:relative-window-position :left})))
+  (is (dfn/equals [3 6 10 15 20 25 30 35 40 45 50 55 60 65 70 75 80 85 89 92]
+                  (rolling/fixed-rolling-window (apply list (range 20)) 5 dfn/sum
+                                                {:relative-window-position :center})))
+  (is (dfn/equals [10 15 20 25 30 35 40 45 50 55 60 65 70 75 80 85 89 92 94 95]
+                  (rolling/fixed-rolling-window (apply list (range 20)) 5 dfn/sum
+                                                {:relative-window-position :right})))
+  (is (dfn/equals [0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19]
+                  (rolling/fixed-rolling-window (apply list (range 20)) 5
+                                                (fn [w] (last w))
+                                                {:relative-window-position :left}))))
 
 
 (deftest binary-search