ewm.rs

pub use polars_compute::ewm::EWMOptions;
use polars_compute::ewm::mean::ewm_mean as kernel_ewm_mean;
use polars_compute::ewm::{ewm_std as kernel_ewm_std, ewm_var as kernel_ewm_var};
use polars_core::prelude::*;

fn check_alpha(alpha: f64) -> PolarsResult<()> {
    polars_ensure!((0.0..=1.0).contains(&alpha), ComputeError: "alpha must be in [0; 1]");
    Ok(())
}

pub fn ewm_mean(s: &Series, options: EWMOptions) -> PolarsResult<Series> {
    check_alpha(options.alpha).inspect_err(|_| {
        if cfg!(debug_assertions) {
            panic!()
        }
    })?;
    match s.dtype() {
        #[cfg(feature = "dtype-f16")]
        DataType::Float16 => {
            use num_traits::AsPrimitive;

            let xs = s.f16().unwrap();
            let result = kernel_ewm_mean(
                xs,
                options.alpha.as_(),
                options.adjust,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        DataType::Float32 => {
            let xs = s.f32().unwrap();
            let result = kernel_ewm_mean(
                xs,
                options.alpha as f32,
                options.adjust,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        DataType::Float64 => {
            let xs = s.f64().unwrap();
            let result = kernel_ewm_mean(
                xs,
                options.alpha,
                options.adjust,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        dt if cfg!(debug_assertions) => panic!("{:?}", dt),
        _ => ewm_mean(&s.cast(&DataType::Float64)?, options),
    }
}

pub fn ewm_std(s: &Series, options: EWMOptions) -> PolarsResult<Series> {
    check_alpha(options.alpha)?;
    match s.dtype() {
        #[cfg(feature = "dtype-f16")]
        DataType::Float16 => {
            use num_traits::AsPrimitive;

            let xs = s.f16().unwrap();
            let result = kernel_ewm_std(
                xs,
                options.alpha.as_(),
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        DataType::Float32 => {
            let xs = s.f32().unwrap();
            let result = kernel_ewm_std(
                xs,
                options.alpha as f32,
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        DataType::Float64 => {
            let xs = s.f64().unwrap();
            let result = kernel_ewm_std(
                xs,
                options.alpha,
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        _ => ewm_std(&s.cast(&DataType::Float64)?, options),
    }
}

pub fn ewm_var(s: &Series, options: EWMOptions) -> PolarsResult<Series> {
    check_alpha(options.alpha)?;
    match s.dtype() {
        #[cfg(feature = "dtype-f16")]
        DataType::Float16 => {
            use num_traits::AsPrimitive;

            let xs = s.f16().unwrap();
            let result = kernel_ewm_var(
                xs,
                options.alpha.as_(),
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        DataType::Float32 => {
            let xs = s.f32().unwrap();
            let result = kernel_ewm_var(
                xs,
                options.alpha as f32,
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        DataType::Float64 => {
            let xs = s.f64().unwrap();
            let result = kernel_ewm_var(
                xs,
                options.alpha,
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        _ => ewm_var(&s.cast(&DataType::Float64)?, options),
    }
}