How can I eliminate the loop when computing adaptive EWMA for large streamflow data in Python?

If you have only few alpha values and/or have some condition on which alpha should be used for which row, you could use pl.coalesce(), pl.when() and pl.Expr.ewm_mean():

PYTHONCopied
# Creating a DataFrame with adaptive alpha values
import polars as pl
import numpy as np

df = pl.DataFrame({
    "adaptive_alpha": np.concatenate([np.repeat(0.3, 10), np.repeat(0.6, 10)]),
    "streamflow": np.arange(0, 20, 1)
})

df.with_columns(
    pl.coalesce(
        pl.when(pl.col.adaptive_alpha == alpha)
        .then(pl.col.streamflow.ewm_mean(alpha = alpha))
        for alpha in df["adaptive_alpha"].unique()
    ).alias("ewma")
).with_columns(ewma_int = pl.col.ewma.cast(pl.Int32))

PLAINTEXTCopied
shape: (20, 4)
┌────────────────┬────────────┬───────────┬──────────┐
│ adaptive_alpha ┆ streamflow ┆ ewma      ┆ ewma_int │
│ ---            ┆ ---        ┆ ---       ┆ ---      │
│ f64            ┆ i64        ┆ f64       ┆ i32      │
╞════════════════╪════════════╪═══════════╪══════════╡
│ 0.3            ┆ 0          ┆ 0.0       ┆ 0        │
│ 0.3            ┆ 1          ┆ 0.588235  ┆ 0        │
│ 0.3            ┆ 2          ┆ 1.232877  ┆ 1        │
│ 0.3            ┆ 3          ┆ 1.930517  ┆ 1        │
│ 0.3            ┆ 4          ┆ 2.676788  ┆ 2        │
│ …              ┆ …          ┆ …         ┆ …        │
│ 0.6            ┆ 15         ┆ 14.33334  ┆ 14       │
│ 0.6            ┆ 16         ┆ 15.333336 ┆ 15       │
│ 0.6            ┆ 17         ┆ 16.333335 ┆ 16       │
│ 0.6            ┆ 18         ┆ 17.333334 ┆ 17       │
│ 0.6            ┆ 19         ┆ 18.333334 ┆ 18       │
└────────────────┴────────────┴───────────┴──────────┘

There are a couple of issues in the code given:

• The code is doing a lot of computations already made in the previous loop indexes. One correct algorithm would be the one you got from DeepSeek. But the devil is in the details.

• You get integers instead of float because np.zeros_like(X) will take the same type as X so integers here and it is not what you want when you compute exponential moving average. So should use:

  ewma_data = np.zeros_like(streamflow_data, dtype='float32')

• In the polars documentation for ewm_mean one can see that there are several options to compute the EWM (see adjust parameter):
  
  I don't know which one you want but notice that by default adjust=True. If you use adjust=False in your code (and use float as above) you will get the same results as DeepSeek.

You can calculate the EWMA iteratively using a single pass over the data without recalculating intermediate values

PYTHONCopied
import numpy as np  
  
streamflow_data = np.arange(0, 20, 1)  
adaptive_alphas = np.repeat(0.3, len(streamflow_data))  
  
# Initialize the EWMA array  
ewma_data = np.zeros_like(streamflow_data, dtype=float)  
  
# Set the initial value  
ewma_data[0] = streamflow_data[0]  
  
# Compute EWMA in a vectorized way  
for i in range(1, len(streamflow_data)):
    ewma_data[i] = (
        adaptive_alphas[i] * streamflow_data[i] + (1 - adaptive_alphas[i]) * ewma_data[i - 1]
    )  
  
print(ewma_data)  

I didn't try the code but you can play around with it.
also you can always use this library
TA-Lib

PYTHONCopied
ewma_data = np.zeros_like(streamflow_data)
alpha_high, alpha_low = 0.6, 0.3
ewma_high = streamflow_series.ewm_mean(alpha=alpha_high)
ewma_low = streamflow_series.ewm_mean(alpha=alpha_low)
ewma_high_index = np.argwhere(adaptive_alphas == alpha_high)
ewma_low_index = np.argwhere(adaptive_alphas == alpha_low)
alpha_zones_high = find_continuous_intervals_vectorized(np.concatenate(ewma_high_index)) if len(ewma_high_index>0) else []
alpha_zones_low = find_continuous_intervals_vectorized(np.concatenate(ewma_low_index)) if len(ewma_low_index>0) else []
alpha_zones = alpha_zones_high + alpha_zones_low
# np.array_equal(streamflow_series.ewm_mean(alpha=current_alpha)[:i+1].to_numpy(), streamflow_series[:i+1].ewm_mean(alpha=current_alpha).to_numpy()) = True
for a_zone in alpha_zones:
    ewma_all = ewma_high if adaptive_alphas[a_zone[0]] == alpha_high else ewma_low
    ewma_data[a_zone[0]: a_zone[-1]+1] = ewma_all[a_zone[0]: a_zone[-1]+1]

PYTHONCopied
def find_continuous_intervals_vectorized(arr):
    if len(arr) == 0:
        return []
    diffs = np.diff(arr)
    boundaries = np.where(diffs != 1)[0]
    boundaries = np.concatenate(([-1], boundaries, [len(arr) - 1]))
    intervals = np.split(arr, boundaries + 1)
    intervals = [interval for interval in intervals if len(interval) > 1]
    return intervals