polars_ops/series/ops/
linear_space.rs

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

polars_ops/series/ops/
linear_space.rs
