|
| 1 | +from typing import Union |
| 2 | +import math |
| 3 | +import numpy as np |
| 4 | +from pandas import DataFrame as PdDataFrame |
| 5 | +from polars import DataFrame as PlDataFrame |
| 6 | +import polars as pl |
| 7 | +from pyindicators.exceptions import PyIndicatorException |
| 8 | + |
| 9 | + |
| 10 | +def _gauss(x: float, h: float) -> float: |
| 11 | + """Gaussian kernel function.""" |
| 12 | + return math.exp(-(x * x) / (h * h * 2)) |
| 13 | + |
| 14 | + |
| 15 | +def nadaraya_watson_envelope( |
| 16 | + data: Union[PdDataFrame, PlDataFrame], |
| 17 | + source_column: str = 'Close', |
| 18 | + bandwidth: float = 8.0, |
| 19 | + mult: float = 3.0, |
| 20 | + lookback: int = 500, |
| 21 | + upper_column: str = 'nwe_upper', |
| 22 | + lower_column: str = 'nwe_lower', |
| 23 | + middle_column: str = 'nwe_middle', |
| 24 | +) -> Union[PdDataFrame, PlDataFrame]: |
| 25 | + """ |
| 26 | + Calculate the Nadaraya-Watson Envelope indicator. |
| 27 | +
|
| 28 | + Uses Gaussian kernel regression to create a smoothed price estimate, |
| 29 | + then adds an envelope based on the mean absolute error (MAE) scaled |
| 30 | + by a multiplier. This is a non-repainting (endpoint) implementation. |
| 31 | +
|
| 32 | + Based on the TradingView "Nadaraya-Watson Envelope [LuxAlgo]" indicator. |
| 33 | +
|
| 34 | + Calculation: |
| 35 | + - Kernel weights: w(i) = exp(-i² / (2 * h²)) for i = 0..lookback-1 |
| 36 | + - Smoothed value: sum(src[t-i] * w(i)) / sum(w(i)) |
| 37 | + - MAE: SMA of |src - smoothed| over lookback period |
| 38 | + - Upper: smoothed + mult * MAE |
| 39 | + - Lower: smoothed - mult * MAE |
| 40 | +
|
| 41 | + Args: |
| 42 | + data: pandas or polars DataFrame with price data |
| 43 | + source_column: Column name for the source prices |
| 44 | + (default: 'Close') |
| 45 | + bandwidth: Gaussian kernel bandwidth / smoothing factor |
| 46 | + (default: 8.0). Higher values produce smoother curves. |
| 47 | + mult: Multiplier for the MAE envelope width (default: 3.0) |
| 48 | + lookback: Number of bars to use for kernel regression |
| 49 | + (default: 500) |
| 50 | + upper_column: Result column name for upper envelope |
| 51 | + (default: 'nwe_upper') |
| 52 | + lower_column: Result column name for lower envelope |
| 53 | + (default: 'nwe_lower') |
| 54 | + middle_column: Result column name for the smoothed line |
| 55 | + (default: 'nwe_middle') |
| 56 | +
|
| 57 | + Returns: |
| 58 | + DataFrame with added columns: |
| 59 | + - {upper_column}: Upper boundary of the envelope |
| 60 | + - {lower_column}: Lower boundary of the envelope |
| 61 | + - {middle_column}: Nadaraya-Watson smoothed estimate |
| 62 | +
|
| 63 | + Example: |
| 64 | + >>> import pandas as pd |
| 65 | + >>> from pyindicators import nadaraya_watson_envelope |
| 66 | + >>> df = pd.DataFrame({ |
| 67 | + ... 'Close': [100, 102, 101, 103, 105, 104, 106, 108, 107, 109] |
| 68 | + ... }) |
| 69 | + >>> result = nadaraya_watson_envelope(df, bandwidth=3, lookback=5) |
| 70 | + >>> print(result[['nwe_upper', 'nwe_middle', 'nwe_lower']].tail()) |
| 71 | + """ |
| 72 | + if bandwidth <= 0: |
| 73 | + raise PyIndicatorException("Bandwidth must be greater than 0") |
| 74 | + |
| 75 | + if mult < 0: |
| 76 | + raise PyIndicatorException("Multiplier must be non-negative") |
| 77 | + |
| 78 | + if lookback < 1: |
| 79 | + raise PyIndicatorException("Lookback must be at least 1") |
| 80 | + |
| 81 | + if source_column not in data.columns: |
| 82 | + raise PyIndicatorException( |
| 83 | + f"The source column '{source_column}' does not " |
| 84 | + "exist in the DataFrame." |
| 85 | + ) |
| 86 | + |
| 87 | + is_polars = isinstance(data, PlDataFrame) |
| 88 | + |
| 89 | + if is_polars: |
| 90 | + src = data[source_column].to_numpy().astype(float) |
| 91 | + else: |
| 92 | + src = data[source_column].values.astype(float) |
| 93 | + |
| 94 | + n = len(src) |
| 95 | + |
| 96 | + # Pre-compute Gaussian kernel weights |
| 97 | + max_k = min(lookback, n) |
| 98 | + weights = np.array([_gauss(i, bandwidth) for i in range(max_k)]) |
| 99 | + |
| 100 | + # Compute the smoothed (endpoint) estimate for each bar |
| 101 | + smoothed = np.full(n, np.nan) |
| 102 | + |
| 103 | + for t in range(n): |
| 104 | + # Number of available past bars (including current) |
| 105 | + available = min(t + 1, max_k) |
| 106 | + w = weights[:available] |
| 107 | + s = src[t - available + 1: t + 1][::-1] # most recent first |
| 108 | + smoothed[t] = np.nansum(s * w) / np.nansum(w) |
| 109 | + |
| 110 | + # Compute MAE (mean absolute error) using a rolling window |
| 111 | + abs_err = np.abs(src - smoothed) |
| 112 | + mae = np.full(n, np.nan) |
| 113 | + |
| 114 | + for t in range(n): |
| 115 | + window = min(t + 1, max_k) |
| 116 | + mae[t] = np.nanmean(abs_err[t - window + 1: t + 1]) |
| 117 | + |
| 118 | + upper = smoothed + mult * mae |
| 119 | + lower = smoothed - mult * mae |
| 120 | + |
| 121 | + if is_polars: |
| 122 | + data = data.with_columns([ |
| 123 | + pl.Series(name=middle_column, values=smoothed), |
| 124 | + pl.Series(name=upper_column, values=upper), |
| 125 | + pl.Series(name=lower_column, values=lower), |
| 126 | + ]) |
| 127 | + else: |
| 128 | + data[middle_column] = smoothed |
| 129 | + data[upper_column] = upper |
| 130 | + data[lower_column] = lower |
| 131 | + |
| 132 | + return data |
0 commit comments