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
49impl From<TrigonometricFunction> for FunctionExpr {
50 fn from(value: TrigonometricFunction) -> Self {
51 Self::Trigonometry(value)
52 }
53}
54
55pub(super) fn apply_trigonometric_function(
56 s: &Column,
57 trig_function: TrigonometricFunction,
58) -> PolarsResult<Column> {
59 use DataType::*;
60 match s.dtype() {
61 Float32 => {
62 let ca = s.f32().unwrap();
63 apply_trigonometric_function_to_float(ca, trig_function)
64 },
65 Float64 => {
66 let ca = s.f64().unwrap();
67 apply_trigonometric_function_to_float(ca, trig_function)
68 },
69 dt if dt.is_primitive_numeric() => {
70 let s = s.cast(&Float64)?;
71 apply_trigonometric_function(&s, trig_function)
72 },
73 dt => polars_bail!(op = "trigonometry", dt),
74 }
75}
76
77pub(super) fn apply_arctan2(s: &mut [Column]) -> PolarsResult<Option<Column>> {
78 let y = &s[0];
79 let x = &s[1];
80
81 let y_len = y.len();
82 let x_len = x.len();
83
84 match (y_len, x_len) {
85 (1, _) | (_, 1) => arctan2_on_columns(y, x),
86 (len_a, len_b) if len_a == len_b => arctan2_on_columns(y, x),
87 _ => polars_bail!(
88 ComputeError:
89 "y shape: {} in `arctan2` expression does not match that of x: {}",
90 y_len, x_len,
91 ),
92 }
93}
94
95fn arctan2_on_columns(y: &Column, x: &Column) -> PolarsResult<Option<Column>> {
96 use DataType::*;
97 match y.dtype() {
98 Float32 => {
99 let y_ca: &ChunkedArray<Float32Type> = y.f32().unwrap();
100 arctan2_on_floats(y_ca, x)
101 },
102 Float64 => {
103 let y_ca: &ChunkedArray<Float64Type> = y.f64().unwrap();
104 arctan2_on_floats(y_ca, x)
105 },
106 _ => {
107 let y = y.cast(&DataType::Float64)?;
108 arctan2_on_columns(&y, x)
109 },
110 }
111}
112
113fn arctan2_on_floats<T>(y: &ChunkedArray<T>, x: &Column) -> PolarsResult<Option<Column>>
114where
115 T: PolarsFloatType,
116 T::Native: Float,
117 ChunkedArray<T>: IntoColumn,
118{
119 let dtype = T::get_dtype();
120 let x = x.cast(&dtype)?;
121 let x = y
122 .unpack_series_matching_type(x.as_materialized_series())
123 .unwrap();
124
125 Ok(Some(
126 broadcast_binary_elementwise(y, x, |yv, xv| Some(yv?.atan2(xv?))).into_column(),
127 ))
128}
129
130fn apply_trigonometric_function_to_float<T>(
131 ca: &ChunkedArray<T>,
132 trig_function: TrigonometricFunction,
133) -> PolarsResult<Column>
134where
135 T: PolarsFloatType,
136 T::Native: Float,
137 ChunkedArray<T>: IntoColumn,
138{
139 match trig_function {
140 TrigonometricFunction::Cos => cos(ca),
141 TrigonometricFunction::Cot => cot(ca),
142 TrigonometricFunction::Sin => sin(ca),
143 TrigonometricFunction::Tan => tan(ca),
144 TrigonometricFunction::ArcCos => arccos(ca),
145 TrigonometricFunction::ArcSin => arcsin(ca),
146 TrigonometricFunction::ArcTan => arctan(ca),
147 TrigonometricFunction::Cosh => cosh(ca),
148 TrigonometricFunction::Sinh => sinh(ca),
149 TrigonometricFunction::Tanh => tanh(ca),
150 TrigonometricFunction::ArcCosh => arccosh(ca),
151 TrigonometricFunction::ArcSinh => arcsinh(ca),
152 TrigonometricFunction::ArcTanh => arctanh(ca),
153 TrigonometricFunction::Degrees => degrees(ca),
154 TrigonometricFunction::Radians => radians(ca),
155 }
156}
157
158fn cos<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
159where
160 T: PolarsFloatType,
161 T::Native: Float,
162 ChunkedArray<T>: IntoColumn,
163{
164 Ok(ca.apply_values(|v| v.cos()).into_column())
165}
166
167fn cot<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
168where
169 T: PolarsFloatType,
170 T::Native: Float,
171 ChunkedArray<T>: IntoColumn,
172{
173 Ok(ca.apply_values(|v| v.tan().powi(-1)).into_column())
174}
175
176fn sin<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
177where
178 T: PolarsFloatType,
179 T::Native: Float,
180 ChunkedArray<T>: IntoColumn,
181{
182 Ok(ca.apply_values(|v| v.sin()).into_column())
183}
184
185fn tan<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
186where
187 T: PolarsFloatType,
188 T::Native: Float,
189 ChunkedArray<T>: IntoColumn,
190{
191 Ok(ca.apply_values(|v| v.tan()).into_column())
192}
193
194fn arccos<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
195where
196 T: PolarsFloatType,
197 T::Native: Float,
198 ChunkedArray<T>: IntoColumn,
199{
200 Ok(ca.apply_values(|v| v.acos()).into_column())
201}
202
203fn arcsin<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
204where
205 T: PolarsFloatType,
206 T::Native: Float,
207 ChunkedArray<T>: IntoColumn,
208{
209 Ok(ca.apply_values(|v| v.asin()).into_column())
210}
211
212fn arctan<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
213where
214 T: PolarsFloatType,
215 T::Native: Float,
216 ChunkedArray<T>: IntoColumn,
217{
218 Ok(ca.apply_values(|v| v.atan()).into_column())
219}
220
221fn cosh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
222where
223 T: PolarsFloatType,
224 T::Native: Float,
225 ChunkedArray<T>: IntoColumn,
226{
227 Ok(ca.apply_values(|v| v.cosh()).into_column())
228}
229
230fn sinh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
231where
232 T: PolarsFloatType,
233 T::Native: Float,
234 ChunkedArray<T>: IntoColumn,
235{
236 Ok(ca.apply_values(|v| v.sinh()).into_column())
237}
238
239fn tanh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
240where
241 T: PolarsFloatType,
242 T::Native: Float,
243 ChunkedArray<T>: IntoColumn,
244{
245 Ok(ca.apply_values(|v| v.tanh()).into_column())
246}
247
248fn arccosh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
249where
250 T: PolarsFloatType,
251 T::Native: Float,
252 ChunkedArray<T>: IntoColumn,
253{
254 Ok(ca.apply_values(|v| v.acosh()).into_column())
255}
256
257fn arcsinh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
258where
259 T: PolarsFloatType,
260 T::Native: Float,
261 ChunkedArray<T>: IntoColumn,
262{
263 Ok(ca.apply_values(|v| v.asinh()).into_column())
264}
265
266fn arctanh<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
267where
268 T: PolarsFloatType,
269 T::Native: Float,
270 ChunkedArray<T>: IntoColumn,
271{
272 Ok(ca.apply_values(|v| v.atanh()).into_column())
273}
274
275fn degrees<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
276where
277 T: PolarsFloatType,
278 T::Native: Float,
279 ChunkedArray<T>: IntoColumn,
280{
281 Ok(ca.apply_values(|v| v.to_degrees()).into_column())
282}
283
284fn radians<T>(ca: &ChunkedArray<T>) -> PolarsResult<Column>
285where
286 T: PolarsFloatType,
287 T::Native: Float,
288 ChunkedArray<T>: IntoColumn,
289{
290 Ok(ca.apply_values(|v| v.to_radians()).into_column())
291}