How do I use numpy arctan2 within polars dataframe?

Question

I am trying to use numpy arctan2 in polars dataframe. The code works outside of polars

import polars as pl
import numpy as np 
data = pl.from_dict(
{'v': [-4.293,-2.4659,-1.8378,-0.2821,-4.5649,-3.8128,-7.4274,3.3443,3.8604,-4.2200],
'u': [-11.2268,6.3478,7.1681,3.4986,2.7320,-1.0695,-10.1408,11.2327,6.6623,-8.1412]})

this works

v = data ['v'].to_list()
u = data ['u'].to_list()
wd = np.round(np.degrees(np.arctan2(v,u).tolist())+180,3)
print(wd)

I tried dozens of variations of these ideas

data.with_columns([ ( np.degrees( np.arctan2( pl.col('v'), pl.col('u'),None )  ) + 180  ).alias('wd_ck')  ]).head()
data['wd']=data.select([pl.col('v'),pl.col('u')]).apply(np.arctan2,return_dtype=pl.Float64)

I am trying to calculate wd from v,u using arctans2 inside of the polars dataframe

I’m using windows 11, python 3.9.15, numpy 1.22.3, polars 0.16.2

Asked By: bgk

||

Source

Answer 1

EDIT 2: Performance Plot (the lower the better)

EDIT 1: answer expanded after @bgk feedback in the comments

maybe by using .with_columns() or .apply()

To create a columd wd within a dataframe:

In [23]: data.with_columns([
    ...:     pl.struct(['v', 'u']).apply(
    ...:         lambda x: np.round(np.degrees(np.arctan2(x['v'], x['u'])) + 180, 3)
    ...:     ).alias('wd'),
    ...: ])
    ...:
Out[23]:
shape: (10, 3)
┌─────────┬──────────┬─────────┐
│ v       ┆ u        ┆ wd      │
│ ---     ┆ ---      ┆ ---     │
│ f64     ┆ f64      ┆ f64     │
╞═════════╪══════════╪═════════╡
│ -4.293  ┆ -11.2268 ┆ 20.926  │
│ -2.4659 ┆ 6.3478   ┆ 158.771 │
│ -1.8378 ┆ 7.1681   ┆ 165.62  │
│ -0.2821 ┆ 3.4986   ┆ 175.39  │
│ ...     ┆ ...      ┆ ...     │
│ -7.4274 ┆ -10.1408 ┆ 36.22   │
│ 3.3443  ┆ 11.2327  ┆ 196.58  │
│ 3.8604  ┆ 6.6623   ┆ 210.09  │
│ -4.22   ┆ -8.1412  ┆ 27.4    │
└─────────┴──────────┴─────────┘

To get the same result without converting to list:

wd = np.round(np.degrees(np.arctan2(data['v'], data['u'])) + 180, 3)

Where the arctan is calculated as v / u:

                         np.arctan2(data['v'], data['u'])

Then the np.degrees:

              np.degrees(np.arctan2(data['v'], data['u'])) + 180

And the round:

     np.round(np.degrees(np.arctan2(data['v'], data['u'])) + 180, 3)

A quick test to check the result against your example:

In [11]: all(
    ...:     np.round(np.degrees(np.arctan2(data['v'], data['u'])) + 180, 3)
    ...:     == np.round(np.degrees(np.arctan2(data['v'].to_list(), data['u'].to_list()).tolist()) + 180, 3)
    ...: )
    ...:
Out[11]: True

Answered By: Filippo Vitale

Answer 2

Try using map:

data.with_columns(
    [
        pl.map(
            ["v", "u"],
            lambda s: np.degrees(np.arctan2(s[0], s[1], None)) + 180)
        .round(3)
        .alias("wd_ck")
    ]
)

shape: (10, 3)
┌─────────┬──────────┬─────────┐
│ v       ┆ u        ┆ wd_ck   │
│ ---     ┆ ---      ┆ ---     │
│ f64     ┆ f64      ┆ f64     │
╞═════════╪══════════╪═════════╡
│ -4.293  ┆ -11.2268 ┆ 20.926  │
│ -2.4659 ┆ 6.3478   ┆ 158.771 │
│ -1.8378 ┆ 7.1681   ┆ 165.62  │
│ -0.2821 ┆ 3.4986   ┆ 175.39  │
│ ...     ┆ ...      ┆ ...     │
│ -7.4274 ┆ -10.1408 ┆ 36.22   │
│ 3.3443  ┆ 11.2327  ┆ 196.58  │
│ 3.8604  ┆ 6.6623   ┆ 210.09  │
│ -4.22   ┆ -8.1412  ┆ 27.4    │
└─────────┴──────────┴─────────┘

Answered By: ΩΠΟΚΕΚΡΥΜΜΕΝΟΣ

Answer 3

With respect to the other answers, they aren’t taking advantage of the fact that arctan2 and degrees are ufuncs which you can execute directly as an expression.

The somewhat confusing bit is that arctan2 takes two arguments and it isn’t obvious how to get polars to operate on a function that takes two arguments. The answer to that question is to use reduce.

For example,

df.select(pl.reduce(np.arctan2, [pl.col('v'), pl.col('u')]))

shape: (10, 1)
┌───────────┐
│ v         │
│ ---       │
│ f64       │
╞═══════════╡
│ -2.77636  │
│ -0.370523 │
│ -0.25098  │
│ -0.080458 │
│ ...       │
│ -2.509433 │
│ 0.289372  │
│ 0.525164  │
│ -2.663372 │
└───────────┘

For degrees, since it just takes one argument, you can use it directly and still have both functions in the same context, as well as adding 180 and rounding…

df.select((np.degrees(pl.reduce(np.arctan2, [pl.col('v'), pl.col('u')]))+180).round(3))

shape: (10, 1)
┌─────────┐
│ v       │
│ ---     │
│ f64     │
╞═════════╡
│ 20.926  │
│ 158.771 │
│ 165.62  │
│ 175.39  │
│ ...     │
│ 36.22   │
│ 196.58  │
│ 210.09  │
│ 27.4    │
└─────────┘

Answered By: Dean MacGregor

How do I use numpy arctan2 within polars dataframe?

Question:

Answers: