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

polars_ops/series/ops/
linear_space.rs
