Added rolling quantile for time index

miranov25 · Nov 6, 2023 · 1d37097 · 1d37097
1 parent 0573e9a
commit 1d37097
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/RootInteractive/Tools/RDataFrame/RollingQuantile.cpp b/RootInteractive/Tools/RDataFrame/RollingQuantile.cpp
@@ -172,6 +172,43 @@ OutputIt rolling_quantile(InputIt first, InputIt last, OutputIt d_first, size_t
 
 }
 
+// TODO: If performance is bad, add a better algorithm than brute force
+template <class InputIt, class TimeIt, class OutputIt, class DistT>
+OutputIt rolling_quantile_weighted(InputIt first, InputIt last, TimeIt t_first, OutputIt d_first, DistT window, double quantile, bool center){
+	std::vector<InputIt> window_contents;
+	InputIt low = first;
+	InputIt high = first;
+	TimeIt t_low = t_first;
+	TimeIt t_high = t_first;
+	DistT off_low = center ? -window/2 : 0;
+	DistT off_high = center ? window/2 : window;
+	// Sanity checks on quantile to prevent buffer overflow
+	if(quantile >= 1.){
+		quantile = 1.;
+	}
+	if(quantile <= 0. || quantile != quantile){
+		quantile = 0.;
+	}
+	for(;first<last;++first){
+		while(high < last && *t_first + off_high > *t_high){
+			++high;
+			++t_high;
+		}	
+		while(*t_first + off_low > *t_low){
+			++low;
+			++t_low;
+		}
+		window_contents.clear();
+		for(InputIt j = low; j<high; ++j){
+			window_contents.push_back(j);
+		}
+		std::nth_element(window_contents.begin(), window_contents.begin() + (size_t)((high - low)*quantile), window_contents.end(), [](InputIt a, InputIt b){return *a < *b;});
+		*d_first = **(window_contents.begin() + (size_t)((high - low)*quantile));
+	       ++d_first;
+		++t_first;	       
+	}
+	return d_first;
+}
 
 }
 

diff --git a/RootInteractive/Tools/RDataFrame/test_RDataFrame_Array.py b/RootInteractive/Tools/RDataFrame/test_RDataFrame_Array.py
@@ -239,6 +239,8 @@ def test_define2(rdf):
     print(rdf.Sum("arrayRollingMedian1D1").GetValue())
     rdf = makeDefine("arrayRollingMedian1D2", "abs(rollingQuantile(array1D0, 7, center=True, quantile=.5)[3:-3] - array1D0[3:-3])", rdf, None, 3)
     print(rdf.Sum("arrayRollingMedian1D2").GetValue())
+    rdf = makeDefine("arrayRollingMedian1D3", "abs(rollingQuantile(array1D0, 7, center=True, quantile=.5, time=array1D0)[3:-3] - array1D0[3:-3])", rdf, None, 3)
+    print(rdf.Sum("arrayRollingMedian1D3").GetValue())
     return rdf
 
 def test_exception(rdf):