Skip to content

Commit 6df604d

Browse files
Bidek56Bidek56
committed
adding trigonometry fn
1 parent 70204ab commit 6df604d

File tree

3 files changed

+291
-25
lines changed

3 files changed

+291
-25
lines changed

‎__tests__/expr.test.ts

+41-7
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,12 @@ describe("expr", () => {
9696
const fn = () => df.select(col("a").cast(pl.Int16, true));
9797
expect(fn).toThrow();
9898
});
99+
test("cos", () => {
100+
const df = pl.DataFrame({ a: [1, 2, 3] });
101+
const expected = pl.DataFrame({ cos: [0.540302, -0.416147, -0.989992] });
102+
const actual = df.select(col("a").cos().round(6).as("cos"));
103+
expect(actual).toFrameEqual(expected);
104+
});
99105
test("count", () => {
100106
const df = pl.DataFrame({ a: [1, 0, 3, 4, 6, 0] });
101107
const expected = pl.DataFrame({ a: [6] });
@@ -208,6 +214,12 @@ describe("expr", () => {
208214
const actual = df.select(col("*").exclude("b", "c"));
209215
expect(actual).toFrameEqual(expected);
210216
});
217+
test("exp", () => {
218+
const df = pl.DataFrame({ a: [1.0] });
219+
const actual = df.select(pl.col("a").exp());
220+
const expected = pl.DataFrame({ a: [Math.E] });
221+
expect(actual).toFrameEqual(expected);
222+
});
211223
test("explode", () => {
212224
const df = pl.DataFrame({
213225
letters: ["c", "a"],
@@ -329,6 +341,15 @@ describe("expr", () => {
329341
const actual = df.select(col("a").gtEq(0));
330342
expect(actual).toFrameEqual(expected);
331343
});
344+
test("gatherEvery", () => {
345+
const df = pl.DataFrame({ a: [1, 1, 2, 2, 3, 3, 8, null, 1] });
346+
let expected = pl.DataFrame({ everyother: [1, 2, 3, 8, 1] });
347+
let actual = df.select(col("a").gatherEvery(2).as("everyother"));
348+
expect(actual).toFrameEqual(expected);
349+
expected = pl.DataFrame({ everyother: [2, 3, 8, 1] });
350+
actual = df.select(col("a").gatherEvery(2, 2).as("everyother"));
351+
expect(actual).toFrameEqual(expected);
352+
});
332353
test.each`
333354
args | hashValue
334355
${[0]} | ${7355865757046787768n}
@@ -514,6 +535,16 @@ describe("expr", () => {
514535
const actual = df.select(col("a").ltEq(2).as("lt"));
515536
expect(actual).toFrameEqual(expected);
516537
});
538+
test("log", () => {
539+
let df = pl.DataFrame({ a: [1, 2, 3] });
540+
let actual = df.select(col("a").log(2).round(6).as("log"));
541+
let expected = pl.DataFrame({ log: [0.0, 1.0, 1.584963] });
542+
expect(actual).toFrameEqual(expected);
543+
df = pl.DataFrame({ a: [2] });
544+
actual = df.select(col("a").log().as("log"));
545+
expected = pl.DataFrame({ log: [Math.LN2] });
546+
expect(actual).toFrameEqual(expected);
547+
});
517548
test("max", () => {
518549
const df = pl.DataFrame({ a: [1, 5, 3] });
519550
const expected = pl.DataFrame({ max: [5] });
@@ -709,6 +740,12 @@ describe("expr", () => {
709740
const actual = df.select(col("a"), ...shifts);
710741
expect(actual).toFrameStrictEqual(expected);
711742
});
743+
test("sin", () => {
744+
const df = pl.DataFrame({ a: [1, 2, 3] });
745+
const expected = pl.DataFrame({ sin: [0.841471, 0.909297, 0.14112] });
746+
const actual = df.select(col("a").sin().round(6).as("sin"));
747+
expect(actual).toFrameEqual(expected);
748+
});
712749
test("skew", () => {
713750
const df = pl.DataFrame({ a: [1, 2, 3, 3] });
714751
const expected = pl.DataFrame({
@@ -870,13 +907,10 @@ describe("expr", () => {
870907
);
871908
expect(actual).toFrameEqual(expected);
872909
});
873-
test("gatherEvery", () => {
874-
const df = pl.DataFrame({ a: [1, 1, 2, 2, 3, 3, 8, null, 1] });
875-
let expected = pl.DataFrame({ everyother: [1, 2, 3, 8, 1] });
876-
let actual = df.select(col("a").gatherEvery(2).as("everyother"));
877-
expect(actual).toFrameEqual(expected);
878-
expected = pl.DataFrame({ everyother: [2, 3, 8, 1] });
879-
actual = df.select(col("a").gatherEvery(2, 2).as("everyother"));
910+
test("tan", () => {
911+
const df = pl.DataFrame({ a: [1, 2, 3] });
912+
const expected = pl.DataFrame({ tan: [1.557408, -2.18504, -0.142547] });
913+
const actual = df.select(col("a").tan().round(6).as("tan"));
880914
expect(actual).toFrameEqual(expected);
881915
});
882916
test("unique", () => {

‎polars/lazy/expr/index.ts

+178-18
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,36 @@ export interface Expr
6969
toJSON(): string;
7070
/** Take absolute values */
7171
abs(): Expr;
72+
/**
73+
* Get the group indexes of the group by operation.
74+
* Should be used in aggregation context only.
75+
* @example
76+
* ```
77+
>>> const df = pl.DataFrame(
78+
... {
79+
... "group": [
80+
... "one",
81+
... "one",
82+
... "one",
83+
... "two",
84+
... "two",
85+
... "two",
86+
... ],
87+
... "value": [94, 95, 96, 97, 97, 99],
88+
... }
89+
... )
90+
>>> df.group_by("group", maintain_order=True).agg(pl.col("value").aggGroups())
91+
shape: (2, 2)
92+
┌───────┬───────────┐
93+
│ group ┆ value │
94+
│ --- ┆ --- │
95+
│ str ┆ list[u32] │
96+
╞═══════╪═══════════╡
97+
│ one ┆ [0, 1, 2] │
98+
│ two ┆ [3, 4, 5] │
99+
└───────┴───────────┘
100+
*```
101+
*/
72102
aggGroups(): Expr;
73103
/**
74104
* Rename the output of an expression.
@@ -137,6 +167,24 @@ export interface Expr
137167
backwardFill(): Expr;
138168
/** Cast between data types. */
139169
cast(dtype: DataType, strict?: boolean): Expr;
170+
/**
171+
* Compute the element-wise value for the cosine.
172+
* @returns Expression of data type :class:`Float64`.
173+
* @example
174+
* ```
175+
>>> const df = pl.DataFrame({"a": [0.0]})
176+
>>> df.select(pl.col("a").cos())
177+
shape: (1, 1)
178+
┌─────┐
179+
│ a │
180+
│ --- │
181+
│ f64 │
182+
╞═════╡
183+
│ 1.0 │
184+
└─────┘
185+
* ```
186+
*/
187+
cos(): Expr;
140188
/** Count the number of values in this expression */
141189
count(): Expr;
142190
/** Calculate the n-th discrete difference.
@@ -151,7 +199,6 @@ export interface Expr
151199
* @param other Expression to compute dot product with
152200
*/
153201
dot(other: any): Expr;
154-
155202
/**
156203
* Exclude certain columns from a wildcard/regex selection.
157204
*
@@ -196,6 +243,25 @@ export interface Expr
196243
* ```
197244
*/
198245
exclude(column: string, ...columns: string[]): Expr;
246+
/**
247+
* Compute the exponential, element-wise.
248+
* @example
249+
* ```
250+
>>> const df = pl.DataFrame({"values": [1.0, 2.0, 4.0]})
251+
>>> df.select(pl.col("values").exp())
252+
shape: (3, 1)
253+
┌──────────┐
254+
│ values │
255+
│ --- │
256+
│ f64 │
257+
╞══════════╡
258+
│ 2.718282 │
259+
│ 7.389056 │
260+
│ 54.59815 │
261+
└──────────┘
262+
* ```
263+
*/
264+
exp(): Expr;
199265
/**
200266
* Explode a list or utf8 Series.
201267
*
@@ -236,6 +302,14 @@ export interface Expr
236302
flatten(): Expr;
237303
/** Fill missing values with the latest seen values */
238304
forwardFill(): Expr;
305+
/**
306+
* Take values by index.
307+
* @param index An expression that leads to a UInt32 dtyped Series.
308+
*/
309+
gather(index: Expr | number[] | Series): Expr;
310+
gather({ index }: { index: Expr | number[] | Series }): Expr;
311+
/** Take every nth value in the Series and return as a new Series. */
312+
gatherEvery(n: number, offset?: number): Expr;
239313
/** Hash the Series. */
240314
hash(k0?: number, k1?: number, k2?: number, k3?: number): Expr;
241315
hash({
@@ -356,6 +430,46 @@ export interface Expr
356430
last(): Expr;
357431
/** Aggregate to list. */
358432
list(): Expr;
433+
/***
434+
* Compute the natural logarithm of each element plus one.
435+
* This computes `log(1 + x)` but is more numerically stable for `x` close to zero.
436+
* @example
437+
* ```
438+
>>> const df = pl.DataFrame({"a": [1, 2, 3]})
439+
>>> df.select(pl.col("a").log1p())
440+
shape: (3, 1)
441+
┌──────────┐
442+
│ a │
443+
│ --- │
444+
│ f64 │
445+
╞══════════╡
446+
│ 0.693147 │
447+
│ 1.098612 │
448+
│ 1.386294 │
449+
└──────────┘
450+
* ```
451+
*/
452+
log1p(): Expr;
453+
/**
454+
* Compute the logarithm to a given base.
455+
* @param base - Given base, defaults to `e`
456+
* @example
457+
* ```
458+
>>> const df = pl.DataFrame({"a": [1, 2, 3]})
459+
>>> df.select(pl.col("a").log(base=2))
460+
shape: (3, 1)
461+
┌──────────┐
462+
│ a │
463+
│ --- │
464+
│ f64 │
465+
╞══════════╡
466+
│ 0.0 │
467+
│ 1.0 │
468+
│ 1.584963 │
469+
└──────────┘
470+
* ```
471+
*/
472+
log(base?: number): Expr;
359473
/** Returns a unit Series with the lowest value possible for the dtype of this expression. */
360474
lowerBound(): Expr;
361475
peakMax(): Expr;
@@ -503,6 +617,24 @@ export interface Expr
503617
periods,
504618
fillValue,
505619
}: { periods: number; fillValue: number }): Expr;
620+
/**
621+
* Compute the element-wise value for the sine.
622+
* @returns Expression of data type :class:`Float64`.
623+
* @example
624+
* ```
625+
>>> const df = pl.DataFrame({"a": [0.0]})
626+
>>> df.select(pl.col("a").sin())
627+
shape: (1, 1)
628+
┌─────┐
629+
│ a │
630+
│ --- │
631+
│ f64 │
632+
╞═════╡
633+
│ 0.0 │
634+
└─────┘
635+
*```
636+
*/
637+
sin(): Expr;
506638
/**
507639
* Compute the sample skewness of a data set.
508640
* For normally distributed data, the skewness should be about zero. For
@@ -565,13 +697,23 @@ export interface Expr
565697
tail(length?: number): Expr;
566698
tail({ length }: { length: number }): Expr;
567699
/**
568-
* Take values by index.
569-
* @param index An expression that leads to a UInt32 dtyped Series.
700+
* Compute the element-wise value for the tangent.
701+
* @returns Expression of data type :class:`Float64`.
702+
* @example
703+
*```
704+
>>> const df = pl.DataFrame({"a": [1.0]})
705+
>>> df.select(pl.col("a").tan().round(2))
706+
shape: (1, 1)
707+
┌──────┐
708+
│ a │
709+
│ --- │
710+
│ f64 │
711+
╞══════╡
712+
│ 1.56 │
713+
└──────┘
714+
*```
570715
*/
571-
gather(index: Expr | number[] | Series): Expr;
572-
gather({ index }: { index: Expr | number[] | Series }): Expr;
573-
/** Take every nth value in the Series and return as a new Series. */
574-
gatherEvery(n: number, offset?: number): Expr;
716+
tan(): Expr;
575717
/**
576718
* Get the unique values of this expression;
577719
* @param maintainOrder Maintain order of data. This requires more work.
@@ -709,6 +851,9 @@ export const _Expr = (_expr: any): Expr => {
709851
),
710852
);
711853
},
854+
cos() {
855+
return _Expr(_expr.cos());
856+
},
712857
count() {
713858
return _Expr(_expr.count());
714859
},
@@ -859,6 +1004,9 @@ export const _Expr = (_expr: any): Expr => {
8591004
explode() {
8601005
return _Expr(_expr.explode());
8611006
},
1007+
exp() {
1008+
return _Expr(_expr.exp());
1009+
},
8621010
extend(o, n?) {
8631011
if (n !== null && typeof n === "number") {
8641012
return _Expr(_expr.extendConstant(o, n));
@@ -908,6 +1056,17 @@ export const _Expr = (_expr: any): Expr => {
9081056
forwardFill() {
9091057
return _Expr(_expr.forwardFill());
9101058
},
1059+
gather(indices) {
1060+
if (Array.isArray(indices)) {
1061+
indices = pli.lit(Series(indices).inner());
1062+
} else {
1063+
indices = indices.inner();
1064+
}
1065+
return wrap("gather", indices);
1066+
},
1067+
gatherEvery(n, offset = 0) {
1068+
return _Expr(_expr.gatherEvery(n, offset));
1069+
},
9111070
hash(obj: any = 0, k1 = 1, k2 = 2, k3 = 3) {
9121071
if (typeof obj === "number" || typeof obj === "bigint") {
9131072
return wrap("hash", BigInt(obj), BigInt(k1), BigInt(k2), BigInt(k3));
@@ -974,7 +1133,6 @@ export const _Expr = (_expr: any): Expr => {
9741133
kurtosis(obj?, bias = true) {
9751134
const fisher = obj?.["fisher"] ?? (typeof obj === "boolean" ? obj : true);
9761135
bias = obj?.["bias"] ?? bias;
977-
9781136
return _Expr(_expr.kurtosis(fisher, bias));
9791137
},
9801138
last() {
@@ -983,6 +1141,13 @@ export const _Expr = (_expr: any): Expr => {
9831141
list() {
9841142
return _Expr(_expr.list());
9851143
},
1144+
log1p() {
1145+
console.log(_expr.log1p);
1146+
return _Expr(_expr.log1p());
1147+
},
1148+
log(base?: number) {
1149+
return _Expr(_expr.log(base ?? Math.E));
1150+
},
9861151
lowerBound() {
9871152
return _Expr(_expr.lowerBound());
9881153
},
@@ -1139,6 +1304,9 @@ export const _Expr = (_expr: any): Expr => {
11391304
skew(bias) {
11401305
return wrap("skew", bias?.bias ?? bias ?? true);
11411306
},
1307+
sin() {
1308+
return _Expr(_expr.sin());
1309+
},
11421310
slice(arg, len?) {
11431311
if (typeof arg === "number") {
11441312
return wrap("slice", pli.lit(arg), pli.lit(len));
@@ -1181,16 +1349,8 @@ export const _Expr = (_expr: any): Expr => {
11811349
tail(length) {
11821350
return _Expr(_expr.tail(length));
11831351
},
1184-
gather(indices) {
1185-
if (Array.isArray(indices)) {
1186-
indices = pli.lit(Series(indices).inner());
1187-
} else {
1188-
indices = indices.inner();
1189-
}
1190-
return wrap("gather", indices);
1191-
},
1192-
gatherEvery(n, offset = 0) {
1193-
return _Expr(_expr.gatherEvery(n, offset));
1352+
tan() {
1353+
return _Expr(_expr.tan());
11941354
},
11951355
unique(opt?) {
11961356
if (opt || opt?.maintainOrder) {

0 commit comments

Comments
 (0)