Switched cluster timestamp selection to user-defined function

JuliaStats · Jan 14, 2016 · 0c0ea0d · 0c0ea0d
1 parent 49b45af
commit 0c0ea0d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 5 deletions.
diff --git a/src/combine.jl b/src/combine.jl
@@ -44,7 +44,7 @@ end
 
 # collapse ######################
 
-function collapse{T,N,D}(ta::TimeArray{T,N,D}, f::Function; period::Function=week)
+function collapse{T,N,D}(ta::TimeArray{T,N,D}, collapse_valuecluster::Function; period::Function=week, collapse_timestampcluster::Function=last)
 
     length(ta) == 0 && return ta
 
@@ -62,8 +62,8 @@ function collapse{T,N,D}(ta::TimeArray{T,N,D}, f::Function; period::Function=wee
         next_mapped_tstamp = period(next_tstamp)
 
         if mapped_tstamp != next_mapped_tstamp
-          push!(collapsed_tstamps, tstamp)
-          collapsed_values = [collapsed_values; T[f(ta.values[cluster_startrow:i, j]) for j in 1:ncols]']
+          push!(collapsed_tstamps, collapse_timestampcluster(ta.timestamp[cluster_startrow:i]))
+          collapsed_values = [collapsed_values; T[collapse_valuecluster(ta.values[cluster_startrow:i, j]) for j in 1:ncols]']
           cluster_startrow = i+1
         end #if
 
@@ -72,8 +72,8 @@ function collapse{T,N,D}(ta::TimeArray{T,N,D}, f::Function; period::Function=wee
 
     end #for
 
-    push!(collapsed_tstamps, tstamp)
-    collapsed_values = [collapsed_values; T[f(ta.values[cluster_startrow:end, j]) for j in 1:ncols]']
+    push!(collapsed_tstamps, collapse_timestampcluster(ta.timestamp[cluster_startrow:end]))
+    collapsed_values = [collapsed_values; T[collapse_valuecluster(ta.values[cluster_startrow:end, j]) for j in 1:ncols]']
 
     N == 1 && (collapsed_values = vec(collapsed_values))
     return TimeArray(collapsed_tstamps, collapsed_values, ta.colnames, ta.meta)

diff --git a/test/combine.jl b/test/combine.jl
@@ -16,6 +16,16 @@ facts("collapse operations") do
         @fact collapse(ohlc, last).timestamp[1]               --> Date(2000,1,7)
         @fact collapse(ohlc, last, period=month).values[1, :] --> [101.0 103.88 94.5 103.75]
         @fact collapse(ohlc, last, period=month).timestamp[1] --> Date(2000,1,31)
+
+        @fact collapse(cl, first, collapse_timestampcluster=first).values[2]                  --> 97.75
+        @fact collapse(cl, first, collapse_timestampcluster=first).timestamp[2]               --> Date(2000,1,10)
+        @fact collapse(cl, first, period=month, collapse_timestampcluster=first).values[2]    --> 100.25
+        @fact collapse(cl, first, period=month, collapse_timestampcluster=first).timestamp[2] --> Date(2000,2,1)
+
+        @fact collapse(ohlc, first, collapse_timestampcluster=first).values[2, :]               --> [102.0 102.25 94.75 97.75]
+        @fact collapse(ohlc, first, collapse_timestampcluster=first).timestamp[2]               --> Date(2000,1,10)
+        @fact collapse(ohlc, first, period=month, collapse_timestampcluster=first).values[2, :] --> [104.0 105.0 100.0 100.25]
+        @fact collapse(ohlc, first, period=month, collapse_timestampcluster=first).timestamp[2] --> Date(2000,2,1)
     end
 end