Fixed a bug in medain and kth_element visitors related to handling nans

Author: hosseinmoein

include/DataFrame/DataFrameStatsVisitors.h

@@ -953,10 +953,9 @@ struct  CorrVisitor  {
 
         const auto      &idx = *idx_begin;
         const size_type col_s =
-            std::min(std::distance(idx_begin, idx_end),
-                     std::distance(column_begin1, column_end1));
-
-        assert((col_s == size_type(std::distance(column_begin2, column_end2))));
+            std::min ({ std::distance(idx_begin, idx_end),
+                        std::distance(column_begin1, column_end1),
+                        std::distance(column_begin2, column_end2) });
 
         if (type_ == correlation_type::pearson)  {
             while (column_begin1 < column_end1 && column_begin2 < column_end2)
@@ -2532,29 +2531,44 @@ struct  KthValueVisitor  {
 
     DEFINE_VISIT_BASIC_TYPES_2
 
-    template<typename U>
-    using vec_type = std::vector<U, typename allocator_declare<U, A>::type>;
+    using vec_type = std::vector<T, typename allocator_declare<T, A>::type>;
 
     template <forward_iterator K, forward_iterator H>
     inline void
     operator() (const K &, const K &,
                 const H &values_begin, const H &values_end)  {
 
-        vec_type<value_type>    aux (values_begin, values_end);
+        vec_type        aux;
+        const size_type col_s = std::distance(values_begin, values_end);
+
+        if (skip_nan_)  {
+            aux.reserve(col_s);
+            std::copy_if(values_begin, values_end,
+                         std::back_inserter(aux),
+                         [](T x) -> bool { return (! is_nan__(x)); });
+        }
+        else
+            aux.insert(aux.begin(), values_begin, values_end);
+        compute_size_ = aux.size();
+
+        const size_type kth =
+            std::round(double(kth_element_ * compute_size_) / double(col_s));
 
-        result_ = find_kth_element_(aux, 0, aux.size() - 1, kth_element_);
+        result_ = find_kth_element_(aux, 0, compute_size_ - 1, kth);
     }
 
     inline void pre ()  { result_ = value_type(); }
     inline void post ()  {   }
     inline result_type get_result () const  { return (result_); }
+    inline size_type get_compute_size() const  { return (compute_size_); }
 
     explicit KthValueVisitor (size_type ke, bool skipnan = true)
         : kth_element_(ke), skip_nan_(skipnan)  {   }
 
 private:
 
     result_type     result_ {  };
+    size_type       compute_size_ { 0 };
     const size_type kth_element_;
     const bool      skip_nan_;
 
@@ -2612,17 +2626,21 @@ struct  MedianVisitor  {
     operator() (const K &idx_begin, const K &idx_end,
                 const H &column_begin, const H &column_end)  {
 
-        GET_COL_SIZE2
-
-        KthValueVisitor<value_type, index_type, A> kv_visitor (col_s >> 1);
-
+        const std::size_t                           col_s =
+            std::distance(column_begin, column_end);
+        const std::size_t                           half = col_s >> 1;
+        KthValueVisitor<value_type, index_type, A>  kv_visitor (half + 1);
 
         kv_visitor.pre();
         kv_visitor(idx_begin, idx_end, column_begin, column_end);
         kv_visitor.post();
         result_ = kv_visitor.get_result();
-        if (! (col_s & 0x01))  { // Even
-            KthValueVisitor<value_type, I, A>   kv_visitor2 ((col_s >> 1) + 1);
+
+        const size_type cs = kv_visitor.get_compute_size();
+
+        if (! (cs & 0x01))  { // Even
+            KthValueVisitor<value_type, I, A>   kv_visitor2 (
+                 cs < col_s ? half + 2 : half);
 
             kv_visitor2.pre();
             kv_visitor2(idx_begin, idx_end, column_begin, column_end);

test/dataframe_tester.cc

@@ -2131,7 +2131,7 @@ static void test_median()  {
     double                  result =
         df.single_act_visit<double>("dblcol_1", med_visit, true).get_result();
 
-    assert(result == 10.0);
+    assert(result == 11.0);
 
     result = df.single_act_visit<double>("dblcol_2", med_visit).get_result();
     assert(result == 10.50);
@@ -2140,10 +2140,30 @@ static void test_median()  {
     int                 result2 =
         df.single_act_visit<int>("intcol_1", med_visit2).get_result();
 
-    assert(result2 == 10);
+    assert(result2 == 11);
 
     result2 = df.single_act_visit<int>("intcol_2", med_visit2).get_result();
     assert(result2 == 10);
+
+    using TestDF = StdDataFrame<std::string>;
+
+    std::vector<std::string> syms = { "AAPL", "IBM", "TSLA", "MSFT", "CSCO" };
+    std::vector<double>      c1 = { 1.0, 2.0, 3.0, 4.0 };
+    std::vector<double>      c2 = { 0.01, 0.02, 0.03 };
+    std::vector<double>      c3 =
+        { 0.0, std::numeric_limits<double>::quiet_NaN(), 0.1 };
+    TestDF                   testDF ;
+
+    testDF.load_data(std::move(syms),
+                     std::make_pair("c1",  c1),
+                     std::make_pair("c2",  c2),
+                     std::make_pair("c3",  c3));
+
+    MedianVisitor<double, std::string>  md;
+
+    assert((testDF.single_act_visit<double>("c1", md).get_result() == 2.5));
+    assert((testDF.single_act_visit<double>("c2", md).get_result() == 0.02));
+    assert((testDF.single_act_visit<double>("c3", md).get_result() == 0.05));
 }
 
 // -----------------------------------------------------------------------------

test/dataframe_tester_2.cc

@@ -2388,18 +2388,18 @@ static void test_LowessVisitor()  {
     df.single_act_visit<double, double>("dep_var", "indep_var", l_v);
 
     auto    actual_yfit = StlVecType<double> {
-        68.1432, 119.432, 122.75, 135.633, 142.724, 165.905, 169.447, 185.617,
-        186.017, 191.865, 198.03, 202.234, 206.178, 215.053, 216.586, 220.408,
-        226.671, 229.052, 229.185, 230.023, 231.657,
+       67.988, 119.351, 122.673, 135.574, 142.677, 165.901, 169.442, 185.5469,
+       185.946, 191.751, 197.912, 202.10997, 206.052, 214.933, 216.473,
+       220.319, 226.653, 229.068, 229.203, 230.054, 231.714,
     };
 
     for (size_t idx = 0; idx < actual_yfit.size(); ++idx)
         assert(fabs(l_v.get_result()[idx] - actual_yfit[idx]) < 0.001);
 
     auto    actual_weights = StlVecType<double> {
-        0.641773, 0.653544, 0.940738, 0.865302, 0.990575, 0.971522, 0.92929,
-        0.902444, 0.918228, 0.924041, 0.855054, 0.824388, 0.586045, 0.945216,
-        0.94831, 0.998031, 0.999834, 0.991263, 0.993165, 0.972067, 0.990308,
+        0.665908, 0.674181, 0.945216, 0.873828, 0.991117, 0.973495, 0.934069,
+        0.909536, 0.923308, 0.928475, 0.863709, 0.837148, 0.612762, 0.948307,
+        0.951239, 0.998073, 0.99984, 0.991830, 0.993602, 0.974109, 0.990844,
     };
 
     for (size_t idx = 0; idx < actual_weights.size(); ++idx)