Skip to content

Commit

Permalink
Expose frac_padding arg to window-based functions
Browse files Browse the repository at this point in the history
  • Loading branch information
fjclark committed Jan 25, 2024
1 parent 587a23f commit 3add1af
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 1 deletion.
8 changes: 8 additions & 0 deletions red/equilibration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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).
Expand Down
13 changes: 12 additions & 1 deletion red/sse.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down

0 comments on commit 3add1af

Please sign in to comment.