polars_plan/dsl/function_expr/
trigonometry.rs1use num_traits::Float;
2use polars_core::chunked_array::ops::arity::broadcast_binary_elementwise;
3
4use super::*;
5
6#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
7#[derive(Clone, Copy, PartialEq, Debug, Eq, Hash)]
8pub enum TrigonometricFunction {
9 Cos,
10 Cot,
11 Sin,
12 Tan,
13 ArcCos,
14 ArcSin,
15 ArcTan,
16 Cosh,
17 Sinh,
18 Tanh,
19 ArcCosh,
20 ArcSinh,
21 ArcTanh,
22 Degrees,
23 Radians,
24}
25
26impl Display for TrigonometricFunction {
27 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
28 use self::*;
29 match self {
30 TrigonometricFunction::Cos => write!(f, "cos"),
31 TrigonometricFunction::Cot => write!(f, "cot"),
32 TrigonometricFunction::Sin => write!(f, "sin"),
33 TrigonometricFunction::Tan => write!(f, "tan"),
34 TrigonometricFunction::ArcCos => write!(f, "arccos"),
35 TrigonometricFunction::ArcSin => write!(f, "arcsin"),
36 TrigonometricFunction::ArcTan => write!(f, "arctan"),
37 TrigonometricFunction::Cosh => write!(f, "cosh"),
38 TrigonometricFunction::Sinh => write!(f, "sinh"),
39 TrigonometricFunction::Tanh => write!(f, "tanh"),
40 TrigonometricFunction::ArcCosh => write!(f, "arccosh"),
41 TrigonometricFunction::ArcSinh => write!(f, "arcsinh"),
42 TrigonometricFunction::ArcTanh => write!(f, "arctanh"),
43 TrigonometricFunction::Degrees => write!(f, "degrees"),
44 TrigonometricFunction::Radians => write!(f, "radians"),
45 }
46 }
47}
48
49pub(super) fn apply_trigonometric_function(
50 s: &Column,
51 trig_function: TrigonometricFunction,
52) -> PolarsResult<Column> {
53 use DataType::*;
54 match s.dtype() {
55 Float32 => {
56 let ca = s.f32().unwrap();
57 apply_trigonometric_function_to_float(ca, trig_function)
58 },
59 Float64 => {
60 let ca = s.f64().unwrap();
61 apply_trigonometric_function_to_float(ca, trig_function)
62 },
63 dt if dt.is_primitive_numeric() => {
64 let s = s.cast(&Float64)?;
65 apply_trigonometric_function(&s, trig_function)
66 },
67 dt => polars_bail!(op = "trigonometry", dt),
68 }
69}
70
71pub(super) fn apply_arctan2(s: &mut [Column]) -> PolarsResult<Option<Column>> {
72 let y = &s[0];
73 let x = &s[1];
74
75 let y_len = y.len();
76 let x_len = x.len();
77
78 match (y_len, x_len) {
79 (1, _) | (_, 1) => arctan2_on_columns(y, x),
80 (len_a, len_b) if len_a == len_b => arctan2_on_columns(y, x),
81 _ => polars_bail!(
82 ComputeError:
83 "y shape: {} in `arctan2` expression does not match that of x: {}",
84 y_len, x_len,
85 ),
86 }
87}
88
89fn arctan2_on_columns(y: &Column, x: &Column) -> PolarsResult<Option<Column>> {
90 use DataType::*;
91 match y.dtype() {
92 Float32 => {
93 let y_ca: &ChunkedArray<Float32Type> = y.f32().unwrap();
94 arctan2_on_floats(y_ca, x)
95 },
96 Float64 => {
97 let y_ca: &ChunkedArray<Float64Type> = y.f64().unwrap();
98 arctan2_on_floats(y_ca, x)
99 },
100 _ => {
101 let y = y.cast(&DataType::Float64)?;
102 arctan2_on_columns(&y, x)
103 },
104 }
105}
106
107fn arctan2_on_floats<T>(y: &ChunkedArray<T>, x: &Column) -> PolarsResult<Option<Column>>
108where
109 T: PolarsFloatType,
110 T::Native: Float,
111 ChunkedArray<T>: IntoColumn,
112{
113 let dtype = T::get_dtype();
114 let x = x.cast(&dtype)?;
115 let x = y
116 .unpack_series_matching_type(x.as_materialized_series())
117 .unwrap();
118
119 Ok(Some(
120 broadcast_binary_elementwise(y, x, |yv, xv| Some(yv?.atan2(xv?))).into_column(),
121 ))
122}
123
124fn apply_trigonometric_function_to_float<T>(
125 ca: &ChunkedArray<T>,
126 trig_function: TrigonometricFunction,
127) -> PolarsResult<Column>
128where
129 T: PolarsFloatType,
130 T::Native: Float,
131 ChunkedArray<T>: IntoColumn,
132{
133 match trig_function {
134 TrigonometricFunction::Cos => cos(ca),
135 TrigonometricFunction::Cot => cot(ca),
136 TrigonometricFunction::Sin => sin(ca),
137 TrigonometricFunction::Tan => tan(ca),
138 TrigonometricFunction::ArcCos => arccos(ca),
139 TrigonometricFunction::ArcSin => arcsin(ca),
140 TrigonometricFunction::ArcTan => arctan(ca),
141 TrigonometricFunction::Cosh => cosh(ca),
142 TrigonometricFunction::Sinh => sinh(ca),
143 TrigonometricFunction::Tanh => tanh(ca),
144 TrigonometricFunction::ArcCosh => arccosh(ca),
145 TrigonometricFunction::ArcSinh => arcsinh(ca),
146 TrigonometricFunction::ArcTanh => arctanh(ca),
147 TrigonometricFunction::Degrees => degrees(ca),
148 TrigonometricFunction::Radians => radians(ca),
149 }
150}
151
152fn cos<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
153where
154 T: PolarsFloatType,
155 T::Native: Float,
156 ChunkedArray<T>: IntoColumn,
157{
158 Ok(ca.apply_values(|v| v.cos()).into_column())
159}
160
161fn cot<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
162where
163 T: PolarsFloatType,
164 T::Native: Float,
165 ChunkedArray<T>: IntoColumn,
166{
167 Ok(ca.apply_values(|v| v.tan().powi(-1)).into_column())
168}
169
170fn sin<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
171where
172 T: PolarsFloatType,
173 T::Native: Float,
174 ChunkedArray<T>: IntoColumn,
175{
176 Ok(ca.apply_values(|v| v.sin()).into_column())
177}
178
179fn tan<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
180where
181 T: PolarsFloatType,
182 T::Native: Float,
183 ChunkedArray<T>: IntoColumn,
184{
185 Ok(ca.apply_values(|v| v.tan()).into_column())
186}
187
188fn arccos<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
189where
190 T: PolarsFloatType,
191 T::Native: Float,
192 ChunkedArray<T>: IntoColumn,
193{
194 Ok(ca.apply_values(|v| v.acos()).into_column())
195}
196
197fn arcsin<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
198where
199 T: PolarsFloatType,
200 T::Native: Float,
201 ChunkedArray<T>: IntoColumn,
202{
203 Ok(ca.apply_values(|v| v.asin()).into_column())
204}
205
206fn arctan<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
207where
208 T: PolarsFloatType,
209 T::Native: Float,
210 ChunkedArray<T>: IntoColumn,
211{
212 Ok(ca.apply_values(|v| v.atan()).into_column())
213}
214
215fn cosh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
216where
217 T: PolarsFloatType,
218 T::Native: Float,
219 ChunkedArray<T>: IntoColumn,
220{
221 Ok(ca.apply_values(|v| v.cosh()).into_column())
222}
223
224fn sinh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
225where
226 T: PolarsFloatType,
227 T::Native: Float,
228 ChunkedArray<T>: IntoColumn,
229{
230 Ok(ca.apply_values(|v| v.sinh()).into_column())
231}
232
233fn tanh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
234where
235 T: PolarsFloatType,
236 T::Native: Float,
237 ChunkedArray<T>: IntoColumn,
238{
239 Ok(ca.apply_values(|v| v.tanh()).into_column())
240}
241
242fn arccosh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
243where
244 T: PolarsFloatType,
245 T::Native: Float,
246 ChunkedArray<T>: IntoColumn,
247{
248 Ok(ca.apply_values(|v| v.acosh()).into_column())
249}
250
251fn arcsinh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
252where
253 T: PolarsFloatType,
254 T::Native: Float,
255 ChunkedArray<T>: IntoColumn,
256{
257 Ok(ca.apply_values(|v| v.asinh()).into_column())
258}
259
260fn arctanh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
261where
262 T: PolarsFloatType,
263 T::Native: Float,
264 ChunkedArray<T>: IntoColumn,
265{
266 Ok(ca.apply_values(|v| v.atanh()).into_column())
267}
268
269fn degrees<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
270where
271 T: PolarsFloatType,
272 T::Native: Float,
273 ChunkedArray<T>: IntoColumn,
274{
275 Ok(ca.apply_values(|v| v.to_degrees()).into_column())
276}
277
278fn radians<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
279where
280 T: PolarsFloatType,
281 T::Native: Float,
282 ChunkedArray<T>: IntoColumn,
283{
284 Ok(ca.apply_values(|v| v.to_radians()).into_column())
285}