nu_command/math/
variance.rs

1use crate::math::utils::run_with_function;
2use nu_engine::command_prelude::*;
3
4#[derive(Clone)]
5pub struct MathVariance;
6
7impl Command for MathVariance {
8    fn name(&self) -> &str {
9        "math variance"
10    }
11
12    fn signature(&self) -> Signature {
13        Signature::build("math variance")
14            .input_output_types(vec![
15                (Type::List(Box::new(Type::Number)), Type::Number),
16                (Type::table(), Type::record()),
17                (Type::record(), Type::record()),
18            ])
19            .switch(
20                "sample",
21                "calculate sample variance (i.e. using N-1 as the denominator)",
22                Some('s'),
23            )
24            .allow_variants_without_examples(true)
25            .category(Category::Math)
26    }
27
28    fn description(&self) -> &str {
29        "Returns the variance of a list of numbers or of each column in a table."
30    }
31
32    fn search_terms(&self) -> Vec<&str> {
33        vec!["deviation", "dispersion", "variation", "statistics"]
34    }
35
36    fn is_const(&self) -> bool {
37        true
38    }
39
40    fn run(
41        &self,
42        engine_state: &EngineState,
43        stack: &mut Stack,
44        call: &Call,
45        input: PipelineData,
46    ) -> Result<PipelineData, ShellError> {
47        let sample = call.has_flag(engine_state, stack, "sample")?;
48        run_with_function(call, input, compute_variance(sample))
49    }
50
51    fn run_const(
52        &self,
53        working_set: &StateWorkingSet,
54        call: &Call,
55        input: PipelineData,
56    ) -> Result<PipelineData, ShellError> {
57        let sample = call.has_flag_const(working_set, "sample")?;
58        run_with_function(call, input, compute_variance(sample))
59    }
60
61    fn examples(&self) -> Vec<Example> {
62        vec![
63            Example {
64                description: "Get the variance of a list of numbers",
65                example: "[1 2 3 4 5] | math variance",
66                result: Some(Value::test_float(2.0)),
67            },
68            Example {
69                description: "Get the sample variance of a list of numbers",
70                example: "[1 2 3 4 5] | math variance --sample",
71                result: Some(Value::test_float(2.5)),
72            },
73            Example {
74                description: "Compute the variance of each column in a table",
75                example: "[[a b]; [1 2] [3 4]] | math variance",
76                result: Some(Value::test_record(record! {
77                    "a" => Value::test_int(1),
78                    "b" => Value::test_int(1),
79                })),
80            },
81        ]
82    }
83}
84
85fn sum_of_squares(values: &[Value], span: Span) -> Result<Value, ShellError> {
86    let n = Value::int(values.len() as i64, span);
87    let mut sum_x = Value::int(0, span);
88    let mut sum_x2 = Value::int(0, span);
89    for value in values {
90        let v = match &value {
91            Value::Int { .. } | Value::Float { .. } => Ok(value),
92            Value::Error { error, .. } => Err(*error.clone()),
93            other => Err(ShellError::UnsupportedInput {
94                msg: format!("Attempted to compute the sum of squares of a non-int, non-float value '{}' with a type of `{}`.",
95                        other.coerce_string()?, other.get_type()),
96                input: "value originates from here".into(),
97                msg_span: span,
98                input_span: value.span(),
99            }),
100        }?;
101        let v_squared = &v.mul(span, v, span)?;
102        sum_x2 = sum_x2.add(span, v_squared, span)?;
103        sum_x = sum_x.add(span, v, span)?;
104    }
105
106    let sum_x_squared = sum_x.mul(span, &sum_x, span)?;
107    let sum_x_squared_div_n = sum_x_squared.div(span, &n, span)?;
108
109    let ss = sum_x2.sub(span, &sum_x_squared_div_n, span)?;
110
111    Ok(ss)
112}
113
114pub fn compute_variance(
115    sample: bool,
116) -> impl Fn(&[Value], Span, Span) -> Result<Value, ShellError> {
117    move |values: &[Value], span: Span, head: Span| {
118        let n = if sample {
119            values.len() - 1
120        } else {
121            values.len()
122        };
123        // sum_of_squares() needs the span of the original value, not the call head.
124        let ss = sum_of_squares(values, span)?;
125        let n = Value::int(n as i64, head);
126        ss.div(head, &n, head)
127    }
128}
129
130#[cfg(test)]
131mod test {
132    use super::*;
133
134    #[test]
135    fn test_examples() {
136        use crate::test_examples;
137
138        test_examples(MathVariance {})
139    }
140}
nu_command/math/variance.rs

nu_command/math/
variance.rs
