feat: add Rust-backed discrete convolution for random variables

applpy/algebra.py

@@ -58,8 +58,6 @@ def convolution(random_variable_1, random_variable_2):
                 2. random_variable_2: A random variable
     Output:     1. The convolution of random_variable_1 and random_variable_2
     """
-    # If the two random variables are not both continuous or
-    #   both discrete, return an error
     if random_variable_1.domain_type != random_variable_2.domain_type:
         discrete_domain_types = ["discrete", "discrete_functional"]
         if (random_variable_1.domain_type not in discrete_domain_types) and (
@@ -140,44 +138,10 @@ def convolution(random_variable_1, random_variable_2):
                 raise RVError(err_string)
         random_variable_2 = Convert(random_variable_2)
 
-    # If the distributions are discrete, find and return the convolution
-    #   of the two random variables.
     if random_variable_1.is_discrete():
-        # Convert each random variable to its pdf form
-        left_pdf_rv = pdf(random_variable_1)
-        right_pdf_rv = pdf(random_variable_2)
-        # Create function and support lists for the convolution of the
-        #   two random variables
-        convolution_support_candidates = []
-        convolution_function_candidates = []
-        for i in range(len(left_pdf_rv.support)):
-            for j in range(len(right_pdf_rv.support)):
-                convolution_support_candidates.append(
-                    left_pdf_rv.support[i] + right_pdf_rv.support[j]
-                )
-                convolution_function_candidates.append(left_pdf_rv.func[i] * right_pdf_rv.func[j])
-        # Sort the function and support lists for the convolution
-        sorted_convolution_terms = list(
-            zip(convolution_support_candidates, convolution_function_candidates)
-        )
-        sorted_convolution_terms.sort()
-        sorted_convolution_support = []
-        sorted_convolution_functions = []
-        for i in range(len(sorted_convolution_terms)):
-            sorted_convolution_support.append(sorted_convolution_terms[i][0])
-            sorted_convolution_functions.append(sorted_convolution_terms[i][1])
-        # Remove redundant elements in the support list
-        unique_convolution_support = []
-        unique_convolution_functions = []
-        for i in range(len(sorted_convolution_support)):
-            if sorted_convolution_support[i] not in unique_convolution_support:
-                unique_convolution_support.append(sorted_convolution_support[i])
-                unique_convolution_functions.append(sorted_convolution_functions[i])
-            else:
-                support_index = unique_convolution_support.index(sorted_convolution_support[i])
-                unique_convolution_functions[support_index] += sorted_convolution_functions[i]
-        # Create and return the new random variable
-        return RV(unique_convolution_functions, unique_convolution_support, ["discrete", "pdf"])
+        fast_rv_1 = random_variable_1.to_fast_rv()
+        fast_rv_2 = random_variable_2.to_fast_rv()
+        return RV.from_fast_rv(fast_rv_1 + fast_rv_2)
 
 
 def product(random_variable_1, random_variable_2):

rust/src/algorithms/algebra.rs

@@ -129,6 +129,131 @@ pub fn product_discrete(
     Ok(product_rv)
 }
 
+/// Computes the sum of two independent discrete random variables
+///
+/// # Arguments
+/// * `random_variable_1` - the first random variable
+/// * `random_variable_2` - the second random variable
+///
+/// # Returns
+/// * `sum_rv` - the product of the two random variables
+///
+/// # Examples
+/// ```
+/// use applpy_rust::algorithms::algebra::convolution_discrete;
+/// use applpy_rust::algorithms::number::Number;
+/// use applpy_rust::algorithms::rv::{DomainType, FunctionalForm, RandomVariable};
+/// use num_rational::Rational64;
+///
+/// let rv1 = RandomVariable {
+///     function: vec![
+///         Number::Rational(Rational64::new(1, 2)),
+///         Number::Rational(Rational64::new(1, 2)),
+///     ],
+///     support: vec![Number::Integer(1), Number::Integer(2)],
+///     functional_form: FunctionalForm::Pdf,
+///     domain_type: DomainType::Discrete,
+/// };
+///
+/// let rv2 = RandomVariable {
+///     function: vec![
+///         Number::Rational(Rational64::new(1, 2)),
+///         Number::Rational(Rational64::new(1, 2)),
+///     ],
+///     support: vec![Number::Integer(2), Number::Integer(3)],
+///     functional_form: FunctionalForm::Pdf,
+///     domain_type: DomainType::Discrete,
+/// };
+///
+/// let sum = convolution_discrete(&rv1, &rv2).unwrap();
+///
+/// assert_eq!(
+///     sum.support,
+///     vec![Number::Integer(3), Number::Integer(4), Number::Integer(5)]
+/// );
+/// assert_eq!(
+///     sum.function,
+///     vec![
+///         Number::Rational(Rational64::new(1, 4)),
+///         Number::Rational(Rational64::new(1, 2)),
+///         Number::Rational(Rational64::new(1, 4)),
+///     ]
+/// );
+/// assert!(sum.verify_pdf(None).unwrap());
+/// ```
+pub fn convolution_discrete(
+    random_variable_1: &RandomVariable,
+    random_variable_2: &RandomVariable,
+) -> Result<RandomVariable, String> {
+    let pdf_random_variable_1 = random_variable_1.to_pdf()?;
+    let function_1 = pdf_random_variable_1.function;
+    let support_1 = pdf_random_variable_1.support;
+
+    let pdf_random_variable_2 = random_variable_2.to_pdf()?;
+    let function_2 = pdf_random_variable_2.function;
+    let support_2 = pdf_random_variable_2.support;
+
+    // Find the values and probabilities for support_1 + support_2
+    // for all combinations of support values
+    let mut raw_conv_support = Vec::new();
+    for &s1 in support_1.iter() {
+        for &s2 in support_2.iter() {
+            let support_sum = s1 + s2;
+            raw_conv_support.push(support_sum);
+        }
+    }
+
+    let mut raw_conv_function = Vec::new();
+    for &f1 in function_1.iter() {
+        for &f2 in function_2.iter() {
+            let probability = f1 * f2;
+            raw_conv_function.push(probability);
+        }
+    }
+
+    // Sorts the results by the support values
+    let mut raw_conv_pairs: Vec<_> = raw_conv_support
+        .into_iter()
+        .zip(raw_conv_function)
+        .collect();
+
+    raw_conv_pairs.sort_by(|a, b| {
+        let first_value = a.0.to_f64();
+        let second_value = b.0.to_f64();
+        first_value.total_cmp(&second_value)
+    });
+
+    let (sorted_support, sorted_function): (Vec<Number>, Vec<Number>) =
+        raw_conv_pairs.into_iter().unzip();
+
+    // Remove redundant elements from the support
+    let mut conv_support = Vec::new();
+    let mut conv_function = Vec::new();
+
+    for (&s, &f) in sorted_support.iter().zip(sorted_function.iter()) {
+        let support_index = conv_support.iter().position(|&x| x == s);
+
+        match support_index {
+            Some(index) => {
+                conv_function[index] += f;
+            }
+            None => {
+                conv_support.push(s);
+                conv_function.push(f);
+            }
+        }
+    }
+
+    let sum_rv = RandomVariable {
+        function: conv_function,
+        support: conv_support,
+        functional_form: FunctionalForm::Pdf,
+        domain_type: DomainType::Discrete,
+    };
+
+    Ok(sum_rv)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -235,4 +360,90 @@ mod tests {
         assert_eq!(product.domain_type, DomainType::Discrete);
         assert!(product.verify_pdf(None).unwrap());
     }
+
+    #[test]
+    fn convolution_discrete_combines_duplicate_support_values() {
+        let rv1 = two_point_pdf(
+            [0, 1],
+            [Rational64::new(1, 2), Rational64::new(1, 2)],
+            FunctionalForm::Pdf,
+        );
+        let rv2 = two_point_pdf(
+            [2, 3],
+            [Rational64::new(1, 5), Rational64::new(4, 5)],
+            FunctionalForm::Pdf,
+        );
+
+        let sum = convolution_discrete(&rv1, &rv2).unwrap();
+
+        assert_eq!(
+            sum.support,
+            vec![Number::Integer(2), Number::Integer(3), Number::Integer(4)]
+        );
+        assert_eq!(
+            sum.function,
+            vec![
+                Number::Rational(Rational64::new(1, 10)),
+                Number::Rational(Rational64::new(1, 2)),
+                Number::Rational(Rational64::new(2, 5))
+            ]
+        );
+        assert!(sum.verify_pdf(None).unwrap());
+    }
+
+    #[test]
+    fn convolution_discrete_accepts_cdf_inputs_by_converting_to_pdf() {
+        let rv1 = two_point_pdf(
+            [1, 2],
+            [Rational64::new(1, 4), Rational64::new(1, 1)],
+            FunctionalForm::Cdf,
+        );
+        let rv2 = two_point_pdf(
+            [3, 5],
+            [Rational64::new(1, 2), Rational64::new(1, 1)],
+            FunctionalForm::Cdf,
+        );
+
+        let sum = convolution_discrete(&rv1, &rv2).unwrap();
+
+        assert_eq!(
+            sum.support,
+            vec![
+                Number::Integer(4),
+                Number::Integer(5),
+                Number::Integer(6),
+                Number::Integer(7)
+            ]
+        );
+        assert_eq!(
+            sum.function,
+            vec![
+                Number::Rational(Rational64::new(1, 8)),
+                Number::Rational(Rational64::new(3, 8)),
+                Number::Rational(Rational64::new(1, 8)),
+                Number::Rational(Rational64::new(3, 8))
+            ]
+        );
+        assert!(sum.verify_pdf(None).unwrap());
+    }
+
+    #[test]
+    fn convolution_discrete_preserves_pdf_and_discrete_domain() {
+        let rv1 = two_point_pdf(
+            [2, 4],
+            [Rational64::new(1, 3), Rational64::new(2, 3)],
+            FunctionalForm::Pdf,
+        );
+        let rv2 = two_point_pdf(
+            [1, 3],
+            [Rational64::new(3, 10), Rational64::new(7, 10)],
+            FunctionalForm::Pdf,
+        );
+
+        let sum = convolution_discrete(&rv1, &rv2).unwrap();
+
+        assert_eq!(sum.functional_form, FunctionalForm::Pdf);
+        assert_eq!(sum.domain_type, DomainType::Discrete);
+        assert!(sum.verify_pdf(None).unwrap());
+    }
 }

rust/src/algorithms/rv.rs

@@ -1,7 +1,7 @@
 #![allow(dead_code)]
 
 use std::fmt;
-use std::ops::Mul;
+use std::ops::{Add, AddAssign, Mul};
 
 use num_rational::Rational64;
 use num_traits::cast::ToPrimitive;
@@ -63,6 +63,22 @@ pub struct RandomVariable {
     pub domain_type: DomainType,
 }
 
+impl Add for RandomVariable {
+    type Output = Result<RandomVariable, String>;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        let sum_rv = algebra::convolution_discrete(&self, &rhs)?;
+        Ok(sum_rv)
+    }
+}
+
+impl AddAssign for RandomVariable {
+    fn add_assign(&mut self, rhs: Self) {
+        *self =
+            algebra::convolution_discrete(self, &rhs).expect("failed to sum the random variables");
+    }
+}
+
 impl Mul for RandomVariable {
     type Output = Result<RandomVariable, String>;
 

rust/src/lib.rs

@@ -36,6 +36,10 @@ fn applpy_rust(_py: Python<'_>, module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_function(wrap_pyfunction!(python::api::bootstrap_rv_py, module)?)?;
 
     // random variable algebra functions
+    module.add_function(wrap_pyfunction!(
+        python::api::convolution_discrete_py,
+        module
+    )?)?;
     module.add_function(wrap_pyfunction!(python::api::product_discrete_py, module)?)?;
 
     // transformation functions