Fix NaN handling bug in adjusted, unbiased ema_cov with a finite horizon and ignore_na=True (#550)

Signed-off-by: Adam Glustein <adam.glustein@point72.com>

Author: AdamGlustein

cpp/csp/cppnodes/statsimpl.h

@@ -1597,38 +1597,43 @@ class AlphaDebiasEMA
 
         AlphaDebiasEMA & operator=( AlphaDebiasEMA && rhs ) = default;
 
+        
         void add( double x )
         {
-            if( m_first && likely( !isnan( x ) ) )
-            {
-                m_wsum = 1;
-                m_sqsum = 1;
-                m_first = false;
-            }
-            else if( likely( !isnan( x ) ) )
+            if( likely( !isnan( x ) ) )
             {
-                double decay_factor = ( m_ignore_na ? m_decay : pow( m_decay, m_offset ) );
-                m_wsum *= decay_factor;
-                m_sqsum *= decay_factor * decay_factor;
-                m_offset = 1;
-
-                double w0;
-                if( m_adjust )
-                    w0 = 1.0;
+                if( unlikely( m_first ) )
+                {
+                    m_wsum = 1;
+                    m_sqsum = 1;
+                    m_first = false;
+                }
                 else
-                    w0 = 1 - m_decay;
-                m_sqsum += w0 * w0;
-                m_wsum += w0;
-                if( !m_adjust )
                 {
-                    double correction = decay_factor + w0;
-                    m_wsum /= correction;
-                    m_sqsum /= ( correction * correction );
+                    double decay_factor = ( m_ignore_na ? m_decay : pow( m_decay, m_offset ) );
+                    m_wsum *= decay_factor;
+                    m_sqsum *= decay_factor * decay_factor;
+                    m_offset = 1;
+
+                    double w0;
+                    if( m_adjust )
+                        w0 = 1.0;
+                    else
+                        w0 = 1 - m_decay;
+                    m_sqsum += w0 * w0;
+                    m_wsum += w0;
+                    if( !m_adjust )
+                    {
+                        double correction = decay_factor + w0;
+                        m_wsum /= correction;
+                        m_sqsum /= ( correction * correction );
+                    }
                 }
             }
-            else if ( likely( !m_first ) )
+            else
             {
-                m_offset++;
+                if( likely( !m_first ) )
+                    m_offset++;
                 m_nan_count++;
             }
         }

csp/tests/test_stats.py

@@ -3640,6 +3640,36 @@ def g(allow_non_overlapping: bool):
         # Additional sanity check is that correlation should be 1 as middle ticks are ignored
         np.testing.assert_allclose(res["corr"][1], np.array([np.nan, 1.0, 1.0]), equal_nan=True)
 
+    def test_ema_cov_horizon_bug(self):
+        # Bug in finite horizon, adjusted, unbiased EMA covariance with ignore_na=True
+        # Also applies to ema_var/ema_std as well, as they use cov
+        # When the first data points are NaN, after the initial NaN is removed the next value that is removed has the wrong lookback weight applied to it
+
+        st = datetime(2020, 1, 1)
+        N = 15
+        K = 3
+        horizon = 10
+        alpha = 0.1
+        values = [np.nan if i < K else float(i) for i in range(N)]
+
+        @csp.graph
+        def g():
+            x = csp.curve(
+                typ=float, data=[(st + timedelta(seconds=i), values[i]) for i in range(N)]
+            )  # start with some NaNs
+            ema_std = csp.stats.ema_std(x, alpha=alpha, adjust=True, bias=False, ignore_na=True, horizon=horizon)
+            csp.add_graph_output("ema_std", ema_std)
+
+        res = csp.run(g, starttime=st, endtime=timedelta(seconds=N), output_numpy=True)
+
+        golden_ema_std = np.array(
+            [
+                pd.Series(values[max(0, j - horizon + 1) : j + 1]).ewm(alpha=alpha, ignore_na=True).std().iloc[-1]
+                for j in range(N)
+            ]
+        )
+        np.testing.assert_allclose(res["ema_std"][1], golden_ema_std, atol=1e-10)
+
 
 if __name__ == "__main__":
     unittest.main()