Skip to content

Commit

Permalink
Adding ewmMean, ewmStd, ewmVar (#129)
Browse files Browse the repository at this point in the history 
* Adding ewmMean, ewmStd, ewmVar

* Converting to JS/TS doc syntax

* Linting fixes

* Adding EwmOps to shared traits

* Removing duplicate code, adjusting params

* Moving tests from series to expr
  • Loading branch information
@Bidek56
Bidek56 authored Oct 15, 2023
1 parent 0387b7d commit d80f015
Show file tree
Hide file tree
Showing 4 changed files with 410 additions and 0 deletions.
146 changes: 146 additions & 0 deletions __tests__/expr.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1988,3 +1988,149 @@ describe("Round<T>", () => {
expect(exprFn).toThrow();
});
});

describe("EWM", () => {
test("ewmMean", () => {
let s = pl.Series("s", [2, 5, 3]);
let df = pl.DataFrame([s]);
let expected = pl.DataFrame({ s, ewmMean: [2.0, 4.0, 3.4285714285714284] });

let seriesActual = df.getColumn("s").ewmMean().rename("ewmMean");
let actual = df.withColumn(col("s").ewmMean().as("ewmMean"));

expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmMean"));

seriesActual = df
.getColumn("s")
.ewmMean({ alpha: 0.5, adjust: true, ignoreNulls: true })
.rename("ewmMean");
actual = df.withColumn(
col("s")
.ewmMean({ alpha: 0.5, adjust: true, ignoreNulls: true })
.as("ewmMean"),
);

expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmMean"));

seriesActual = df
.getColumn("s")
.ewmMean({ alpha: 0.5, adjust: false, ignoreNulls: true })
.rename("ewmMean");
actual = df.withColumn(
col("s")
.ewmMean({ alpha: 0.5, adjust: false, ignoreNulls: true })
.as("ewmMean"),
);

expected = pl.DataFrame({ s, ewmMean: [2.0, 3.5, 3.25] });
expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmMean"));

seriesActual = df
.getColumn("s")
.ewmMean(0.5, false, 1, true)
.rename("ewmMean");
actual = df.withColumn(col("s").ewmMean(0.5, false, 1, true).as("ewmMean"));

expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmMean"));

s = pl.Series("a", [2, 3, 5, 7, 4]);
df = pl.DataFrame([s]);

seriesActual = df
.getColumn("a")
.ewmMean({ adjust: true, minPeriods: 2, ignoreNulls: true })
.round(5)
.rename("ewmMean");
actual = df.withColumn(
col("a")
.ewmMean({ adjust: true, minPeriods: 2, ignoreNulls: true })
.round(5)
.as("ewmMean"),
);

expected = pl.DataFrame({ ewmMean: [null, 2.66667, 4, 5.6, 4.77419], s });
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmMean"));
});

test("ewmStd", () => {
const s = pl.Series("s", [2, 5, 3]);
const df = pl.DataFrame([s]);
const expected = pl.DataFrame({ s, ewmStd: [0, 2.12132, 1.38873] });

let seriesActual = df.getColumn("s").ewmStd().round(5).rename("ewmStd");
let actual = df.withColumn(col("s").ewmStd().round(5).as("ewmStd"));

expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmStd"));

seriesActual = df
.getColumn("s")
.ewmStd({ alpha: 0.5, adjust: true, ignoreNulls: true })
.round(5)
.rename("ewmStd");
actual = df.withColumn(
col("s")
.ewmStd({ alpha: 0.5, adjust: true, ignoreNulls: true })
.round(5)
.as("ewmStd"),
);

expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmStd"));

seriesActual = df
.getColumn("s")
.ewmStd(0.5, true, 1, false)
.round(5)
.rename("ewmStd");
actual = df.withColumn(
col("s").ewmStd(0.5, true, 1, false).round(5).as("ewmStd"),
);

expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmStd"));
});

test("ewmVar", () => {
const s = pl.Series("s", [2, 5, 3]);
const df = pl.DataFrame([s]);
const expected = pl.DataFrame({ s, ewmVar: [0, 4.5, 1.92857] });

let seriesActual = df.getColumn("s").ewmVar().round(5).rename("ewmVar");
let actual = df.withColumn(col("s").ewmVar().round(5).as("ewmVar"));

expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmVar"));

seriesActual = df
.getColumn("s")
.ewmVar({ alpha: 0.5, adjust: true, ignoreNulls: true })
.round(5)
.rename("ewmVar");
actual = df.withColumn(
col("s")
.ewmVar({ alpha: 0.5, adjust: true, ignoreNulls: true })
.round(5)
.as("ewmVar"),
);

expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmVar"));

seriesActual = df
.getColumn("s")
.ewmVar(0.5, true, 1, false)
.round(5)
.rename("ewmVar");
actual = df.withColumn(
col("s").ewmVar(0.5, true, 1, false).round(5).as("ewmVar"),
);

expect(actual).toFrameEqual(expected);
expect(seriesActual).toSeriesEqual(expected.getColumn("ewmVar"));
});
});
114 changes: 114 additions & 0 deletions polars/lazy/expr/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import {
Round,
Sample,
Serialize,
EwmOps,
} from "../../shared_traits";
import { InterpolationMethod, FillNullStrategy, RankMethod } from "../../types";
/**
Expand All @@ -32,6 +33,7 @@ export interface Expr
Cumulative<Expr>,
Sample<Expr>,
Round<Expr>,
EwmOps<Expr>,
Serialize {
/** @ignore */
_expr: any;
Expand Down Expand Up @@ -140,6 +142,7 @@ export interface Expr
* @param other Expression to compute dot product with
*/
dot(other: any): Expr;

/**
* Exclude certain columns from a wildcard/regex selection.
*
Expand Down Expand Up @@ -727,6 +730,117 @@ export const _Expr = (_expr: any): Expr => {

return _Expr(_expr.dot(expr));
},
ewmMean(
opts: {
alpha?: number;
adjust?: boolean;
minPeriods?: number;
bias?: boolean;
ignoreNulls?: boolean;
},
adjust?: boolean,
minPeriods?: number,
bias?: boolean,
ignoreNulls?: boolean,
) {
if (opts) {
if (typeof opts === "number") {
return wrap(
"ewmMean",
opts,
adjust ?? true,
minPeriods ?? 1,
bias ?? false,
ignoreNulls ?? true,
);
} else {
return wrap(
"ewmMean",
opts.alpha ?? 0.5,
opts.adjust ?? true,
opts.minPeriods ?? 1,
opts.bias ?? false,
opts.ignoreNulls ?? true,
);
}
} else {
return wrap("ewmMean", 0.5, true, 1, false, true);
}
},
ewmStd(
opts: {
alpha?: number;
adjust?: boolean;
minPeriods?: number;
bias?: boolean;
ignoreNulls?: boolean;
},
adjust?: boolean,
minPeriods?: number,
bias?: boolean,
ignoreNulls?: boolean,
) {
if (opts) {
if (typeof opts === "number") {
return wrap(
"ewmStd",
opts,
adjust ?? true,
minPeriods ?? 1,
bias ?? false,
ignoreNulls ?? true,
);
} else {
return wrap(
"ewmStd",
opts.alpha ?? 0.5,
opts.adjust ?? true,
opts.minPeriods ?? 1,
opts.bias ?? false,
opts.ignoreNulls ?? true,
);
}
} else {
return wrap("ewmStd", 0.5, true, 1, false, true);
}
},
ewmVar(
opts: {
alpha?: number;
adjust?: boolean;
minPeriods?: number;
bias?: boolean;
ignoreNulls?: boolean;
},
adjust?: boolean,
minPeriods?: number,
bias?: boolean,
ignoreNulls?: boolean,
) {
if (opts) {
if (typeof opts === "number") {
return wrap(
"ewmVar",
opts,
adjust ?? true,
minPeriods ?? 1,
bias ?? false,
ignoreNulls ?? true,
);
} else {
return wrap(
"ewmVar",
opts.alpha ?? 0.5,
opts.adjust ?? true,
opts.minPeriods ?? 1,
opts.bias ?? false,
opts.ignoreNulls ?? true,
);
}
} else {
return wrap("ewmVar", 0.5, true, 1, false, true);
}
},
exclude(...columns) {
return _Expr(_expr.exclude(columns.flat(2)));
},
Expand Down
11 changes: 11 additions & 0 deletions polars/series/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import {
Round,
Sample,
Serialize,
EwmOps,
} from "../shared_traits";
import { col } from "../lazy/functions";
import { InterpolationMethod, RankMethod } from "../types";
Expand All @@ -32,6 +33,7 @@ export interface Series
Cumulative<Series>,
Round<Series>,
Sample<Series>,
EwmOps<Series>,
Serialize {
inner(): any;
name: string;
Expand Down Expand Up @@ -1265,6 +1267,15 @@ export function _Series(_s: any): Series {
equals(field: Series) {
return this.eq(field);
},
ewmMean(...args) {
return expr_op("ewmMean", ...args);
},
ewmStd(...args) {
return expr_op("ewmStd", ...args);
},
ewmVar(...args) {
return expr_op("ewmVar", ...args);
},
explode() {
return wrap("explode");
},
Expand Down
Loading

0 comments on commit d80f015

Please sign in to comment.