Skip to main content

nu_command/math/
stddev.rs

1use super::variance::{compute_variance as variance, values_unit};
2use crate::math::utils::{
3    NumericUnit, expand_range_input, run_with_function, run_with_function_and_cell_paths,
4    wrap_unit_f64,
5};
6use nu_engine::command_prelude::*;
7
8#[derive(Clone)]
9pub struct MathStddev;
10
11impl Command for MathStddev {
12    fn name(&self) -> &str {
13        "math stddev"
14    }
15
16    fn signature(&self) -> Signature {
17        Signature::build("math stddev")
18            .input_output_types(vec![
19                (Type::List(Box::new(Type::Number)), Type::Number),
20                (Type::List(Box::new(Type::Duration)), Type::Duration),
21                (Type::List(Box::new(Type::Filesize)), Type::Filesize),
22                (Type::Range, Type::Number),
23                (Type::table(), Type::record()),
24                (Type::record(), Type::record()),
25            ])
26            .switch(
27                "sample",
28                "Calculate sample standard deviation (i.e. using N-1 as the denominator).",
29                Some('s'),
30            )
31            .rest(
32                "columns",
33                SyntaxShape::CellPath,
34                "The cell-paths/columns to operate on.",
35            )
36            .allow_variants_without_examples(true)
37            .category(Category::Math)
38    }
39
40    fn description(&self) -> &str {
41        "Returns the standard deviation of a list of numbers, or of each column in a table."
42    }
43
44    fn search_terms(&self) -> Vec<&str> {
45        vec![
46            "SD",
47            "standard",
48            "deviation",
49            "dispersion",
50            "variation",
51            "statistics",
52        ]
53    }
54
55    fn is_const(&self) -> bool {
56        true
57    }
58
59    fn run(
60        &self,
61        engine_state: &EngineState,
62        stack: &mut Stack,
63        call: &Call,
64        input: PipelineData,
65    ) -> Result<PipelineData, ShellError> {
66        let cell_paths: Vec<CellPath> = call.rest(engine_state, stack, 0)?;
67        let sample = call.has_flag(engine_state, stack, "sample")?;
68        let mf = compute_stddev(sample);
69        if cell_paths.is_empty() {
70            let input = expand_range_input(input, call.head)?;
71            return run_with_function(call, input, mf);
72        }
73        run_with_function_and_cell_paths(call, input, cell_paths, engine_state.signals(), mf)
74    }
75
76    fn run_const(
77        &self,
78        working_set: &StateWorkingSet,
79        call: &Call,
80        input: PipelineData,
81    ) -> Result<PipelineData, ShellError> {
82        let cell_paths: Vec<CellPath> = call.rest_const(working_set, 0)?;
83        let sample = call.has_flag_const(working_set, "sample")?;
84        let mf = compute_stddev(sample);
85        if cell_paths.is_empty() {
86            let input = expand_range_input(input, call.head)?;
87            return run_with_function(call, input, mf);
88        }
89        run_with_function_and_cell_paths(
90            call,
91            input,
92            cell_paths,
93            working_set.permanent().signals(),
94            mf,
95        )
96    }
97
98    fn examples(&self) -> Vec<Example<'_>> {
99        vec![
100            Example {
101                description: "Compute the standard deviation of a list of numbers.",
102                example: "[1 2 3 4 5] | math stddev",
103                result: Some(Value::test_float(std::f64::consts::SQRT_2)),
104            },
105            Example {
106                description: "Compute the sample standard deviation of a list of numbers.",
107                example: "[1 2 3 4 5] | math stddev --sample",
108                result: Some(Value::test_float(1.5811388300841898)),
109            },
110            Example {
111                description: "Compute the standard deviation of each column in a table.",
112                example: "[[a b]; [1 2] [3 4]] | math stddev",
113                result: Some(Value::test_record(record! {
114                    "a" => Value::test_float(1.0),
115                    "b" => Value::test_float(1.0),
116                })),
117            },
118            Example {
119                description: "Compute the standard deviation of list-valued columns in a record.",
120                example: "{alice: [1 3], bob: [4 6]} | math stddev",
121                result: Some(Value::test_record(record! {
122                    "alice" => Value::test_float(1.0),
123                    "bob" => Value::test_float(1.0),
124                })),
125            },
126            Example {
127                description: "Compute the standard deviation of a single column using a cell path.",
128                example: "{alice: [1 3], bob: [4 6]} | math stddev alice",
129                result: Some(Value::test_record(record! {
130                    "alice" => Value::test_float(1.0),
131                    "bob" => Value::list(
132                        vec![Value::test_int(4), Value::test_int(6)],
133                        Span::test_data(),
134                    ),
135                })),
136            },
137        ]
138    }
139}
140
141pub fn compute_stddev(sample: bool) -> impl Fn(&[Value], Span, Span) -> Result<Value, ShellError> {
142    move |values: &[Value], span: Span, head: Span| {
143        let unit = values_unit(values, head)?;
144        // variance() produces its own usable error, so we can use `?` to propagate the error.
145        let variance = variance(sample)(values, span, head)?;
146        let val_span = variance.span();
147        let sqrt = match variance {
148            Value::Float { val, .. } => val.sqrt(),
149            Value::Int { val, .. } => (val as f64).sqrt(),
150            other => return Ok(other),
151        };
152        // Duration/filesize keep their unit (stddev has the same unit as the inputs).
153        // Numbers stay floats, matching the previous behavior.
154        match unit {
155            NumericUnit::Number => Ok(Value::float(sqrt, val_span)),
156            other => Ok(wrap_unit_f64(other, sqrt, val_span)),
157        }
158    }
159}
160
161#[cfg(test)]
162mod test {
163    use super::*;
164
165    #[test]
166    fn test_examples() -> nu_test_support::Result {
167        nu_test_support::test().examples(MathStddev)
168    }
169}