polars_ops/series/ops/
linear_space.rs

1use polars_core::prelude::*;
2use polars_core::series::IsSorted;
3#[cfg(feature = "serde")]
4use serde::{Deserialize, Serialize};
5use strum_macros::IntoStaticStr;
6
7#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Default, IntoStaticStr)]
8#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
9#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
10#[strum(serialize_all = "snake_case")]
11pub enum ClosedInterval {
12    #[default]
13    Both,
14    Left,
15    Right,
16    None,
17}
18
19pub fn new_linear_space_f32(
20    start: f32,
21    end: f32,
22    n: u64,
23    closed: ClosedInterval,
24    name: PlSmallStr,
25) -> PolarsResult<Float32Chunked> {
26    let mut ca = match n {
27        0 => Float32Chunked::full_null(name, 0),
28        1 => match closed {
29            ClosedInterval::None => Float32Chunked::from_slice(name, &[(end + start) * 0.5]),
30            ClosedInterval::Left | ClosedInterval::Both => {
31                Float32Chunked::from_slice(name, &[start])
32            },
33            ClosedInterval::Right => Float32Chunked::from_slice(name, &[end]),
34        },
35        _ => Float32Chunked::from_iter_values(name, {
36            let span = end - start;
37
38            let (start, d, end) = match closed {
39                ClosedInterval::None => {
40                    let d = span / (n + 1) as f32;
41                    (start + d, d, end - d)
42                },
43                ClosedInterval::Left => (start, span / n as f32, end - span / n as f32),
44                ClosedInterval::Right => (start + span / n as f32, span / n as f32, end),
45                ClosedInterval::Both => (start, span / (n - 1) as f32, end),
46            };
47            (0..n - 1)
48                .map(move |v| (v as f32 * d) + start)
49                .chain(std::iter::once(end)) // ensures floating point accuracy of final value
50        }),
51    };
52
53    let is_sorted = if end < start {
54        IsSorted::Descending
55    } else {
56        IsSorted::Ascending
57    };
58    ca.set_sorted_flag(is_sorted);
59    Ok(ca)
60}
61
62pub fn new_linear_space_f64(
63    start: f64,
64    end: f64,
65    n: u64,
66    closed: ClosedInterval,
67    name: PlSmallStr,
68) -> PolarsResult<Float64Chunked> {
69    let mut ca = match n {
70        0 => Float64Chunked::full_null(name, 0),
71        1 => match closed {
72            ClosedInterval::None => Float64Chunked::from_slice(name, &[(end + start) * 0.5]),
73            ClosedInterval::Left | ClosedInterval::Both => {
74                Float64Chunked::from_slice(name, &[start])
75            },
76            ClosedInterval::Right => Float64Chunked::from_slice(name, &[end]),
77        },
78        _ => Float64Chunked::from_iter_values(name, {
79            let span = end - start;
80
81            let (start, d, end) = match closed {
82                ClosedInterval::None => {
83                    let d = span / (n + 1) as f64;
84                    (start + d, d, end - d)
85                },
86                ClosedInterval::Left => (start, span / n as f64, end - span / n as f64),
87                ClosedInterval::Right => (start + span / n as f64, span / n as f64, end),
88                ClosedInterval::Both => (start, span / (n - 1) as f64, end),
89            };
90            (0..n - 1)
91                .map(move |v| (v as f64 * d) + start)
92                .chain(std::iter::once(end)) // ensures floating point accuracy of final value
93        }),
94    };
95
96    let is_sorted = if end < start {
97        IsSorted::Descending
98    } else {
99        IsSorted::Ascending
100    };
101    ca.set_sorted_flag(is_sorted);
102    Ok(ca)
103}
polars_ops/series/ops/linear_space.rs

polars_ops/series/ops/
linear_space.rs
