From 3add1af628ad4d99e88ff42361caa51271c6225e Mon Sep 17 00:00:00 2001
From: finlayclark <s1629668@ed.ac.uk>
Date: Thu, 25 Jan 2024 16:19:31 +0000
Subject: [PATCH] Expose frac_padding arg to window-based functions

---
 red/equilibration.py |  8 ++++++++
 red/sse.py           | 13 ++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/red/equilibration.py b/red/equilibration.py
index 9231f9a..9bf1ea9 100644
--- a/red/equilibration.py
+++ b/red/equilibration.py
@@ -195,6 +195,7 @@ def detect_equilibration_window(
     kernel: _Callable[[int], _np.ndarray] = _np.bartlett,  # type: ignore
     window_size_fn: _Optional[_Callable[[int], int]] = lambda x: round(x**0.5),
     window_size: _Optional[int] = None,
+    frac_padding: float = 0.1,
     plot: bool = False,
     plot_name: _Union[str, _Path] = "equilibration_sse_window.png",
     time_units: str = "ns",
@@ -236,6 +237,12 @@ def detect_equilibration_window(
         The size of the window to use, defined in terms of time lags in the
         forwards direction. If this is not None, window_size_fn must be None.
 
+    frac_padding : float, optional, default=0.1
+        The fraction of the end of the timeseries to avoid calculating the variance
+        for. For example, if frac_padding = 0.1, the variance will be calculated
+        for the first 90% of the time series. This helps to avoid noise in the
+        variance when there are few data points.
+
     plot : bool, optional
         Whether to plot the ESS curve. The default is False.
 
@@ -283,6 +290,7 @@ def detect_equilibration_window(
         kernel=kernel,
         window_size_fn=window_size_fn,
         window_size=window_size,
+        frac_padding=frac_padding,
     )
 
     # Get the corresponding times (or indices).
diff --git a/red/sse.py b/red/sse.py
index e6f9c3b..60ee854 100644
--- a/red/sse.py
+++ b/red/sse.py
@@ -88,6 +88,7 @@ def get_sse_series_window(
     kernel: _Callable[[int], _np.ndarray] = _np.bartlett,  # type: ignore
     window_size_fn: _Optional[_Callable[[int], int]] = lambda x: round(x**0.5),
     window_size: _Optional[int] = None,
+    frac_padding: float = 0.1,
 ) -> _Tuple[_np.ndarray, _np.ndarray]:
     """
     Compute a series of squared standard errors for a time series as data
@@ -110,6 +111,12 @@ def get_sse_series_window(
         The size of the window to use, defined in terms of time lags in the
         forwards direction. If this is not None, window_size_fn must be None.
 
+    frac_padding : float, optional, default=0.1
+        The fraction of the end of the timeseries to avoid calculating the variance
+        for. For example, if frac_padding = 0.1, the variance will be calculated
+        for the first 90% of the time series. This helps to avoid noise in the
+        variance when there are few data points.
+
     Returns
     -------
     np.ndarray
@@ -124,7 +131,11 @@ def get_sse_series_window(
 
     # Compute the variance estimates.
     var_series, window_sizes = get_variance_series_window(
-        data, kernel=kernel, window_size_fn=window_size_fn, window_size=window_size
+        data,
+        kernel=kernel,
+        window_size_fn=window_size_fn,
+        window_size=window_size,
+        frac_padding=frac_padding,
     )
 
     # Compute the squared standard error series by dividing the variance series by