Search code examples
pythonpython-3.xpandaspython-polars

How to calculate ema in polars


So, basically, I'm trying to write an equivalent of this code but use polars instead of pandas.

def ema(close, length=None, talib=None, offset=None, **kwargs):
    """Indicator: Exponential Moving Average (EMA)"""
    # Validate Arguments
    length = int(length) if length and length > 0 else 10
    adjust = kwargs.pop("adjust", False)
    sma = kwargs.pop("sma", True)
    close = verify_series(close, length)
    offset = get_offset(offset)

    if close is None: return

    # Calculate Result
    if sma:
        close = close.copy()
        sma_nth = close[0:length].mean()
        close[:length - 1] = npNaN
        close.iloc[length - 1] = sma_nth
    ema = close.ewm(span=length, adjust=adjust).mean()

and this is my code

def polars_ema(close, length=None, offset=None, **kwargs):
    """Indicator: Exponential Moving Average (EMA)"""
    # Validate Arguments
    length = int(length) if length and length > 0 else 10
    adjust = kwargs.pop("adjust", False)
    sma = kwargs.pop("sma", True)

    if close is None: return

    # Calculate Result
    if sma:
        sma_nth = close.slice(0, length).mean()
        nans = pl.Series([npNaN] * (length - 1))
        sma_nth_series = pl.Series("sma_nth", [sma_nth])
        nans_plus_sma_nth = nans.append(sma_nth_series)
        rest_of_close = close.slice(length, close.len())
        close = nans_plus_sma_nth.append(rest_of_close)
        ema = close.ewm_mean(span=length, adjust=adjust, ignore_nulls=False, min_periods=0)
        print(ema)

The problem is close.ewm_mean is returning a series of NaN values even though the close series in both functions are identical :-(

I have tried changing the arguments in ewm_mean function, but it has been to no avail.

Runnable Example

this is some sample data that I'm passing to polars_ema().

Series: 'close' [f64]
[
    1.08086
    1.08069
    1.08077
    1.08077
    1.08052
    1.08055
    1.08068
    1.08073
    1.08077
    1.08073
    1.08068
    1.08062
    1.08052
    1.0806
    1.08063
    1.08064
    1.08063
    1.08053
    1.08067
    1.08058
]
polars_ema(first_20_items, length=10)

Solution

  • As @jqurious suggested npNaN was the problem. this works now.

    def polars_ema(close, length=None, offset=None, **kwargs):
        """Indicator: Exponential Moving Average (EMA)"""
        # Validate Arguments
        length = int(length) if length and length > 0 else 10
        adjust = kwargs.pop("adjust", False)
        sma = kwargs.pop("sma", True)
    
        if close is None: return
    
        # Calculate Result
        if sma:
            sma_nth = close.slice(0, length).mean()
            nones = pl.Series([None] * (length - 1), dtype=pl.Float64)
            sma_nth_series = pl.Series("sma_nth", [sma_nth])
            nans_plus_sma_nth = nones.append(sma_nth_series)
            rest_of_close = close.slice(length, close.len())
            close = nans_plus_sma_nth.append(rest_of_close)
            ema = close.ewm_mean(span=length, adjust=adjust, ignore_nulls=False, min_periods=0)
            print(ema)