update deps + faster impl

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ppca_rs"
-version = "0.5.1"
+version = "0.5.2"
 edition = "2021"
 publish = false
 
@@ -14,10 +14,10 @@ crate-type = ["staticlib"]
 
 [dependencies]
 ppca = { path = "./ppca" }
-pyo3 = { version = "0.18.3", features = ["extension-module"] }
-numpy = { version = "0.18.0", features = ["nalgebra"] }
+pyo3 = { version = "0.21.2", features = ["extension-module"] }
+numpy = { version = "0.21.0", features = ["nalgebra"] }
 bincode = "1.3.3"
 rayon = "1.7.0"
-nalgebra = "0.32.2"
+nalgebra = "0.32.6"
 rand = "0.8.5"
 rand_distr = "0.4.3"

ppca/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ppca"
-version = "0.5.0"
+version = "0.5.2"
 edition = "2021"
 authors = ["Pedro Bittencourt Arruda <[email protected]>"]
 description = "Rust implementation of the Probabilistic Principal Component Analysis model"
@@ -22,12 +22,12 @@ categories = ["algorithms", "data-structures", "mathematics"]
 [dependencies]
 
 approx = "0.5.1"
-bit-vec = { version = "0.6.3", features = ["serde"] }
-nalgebra = { version = "0.32.2", features = ["serde-serialize"] }
+bit-vec = { version = "0.8.0", features = ["serde"] }
+nalgebra = { version = "0.32.6", features = ["serde-serialize"] }
 rand = "0.8.5"
 rand_distr = "0.4.3"
 rayon = "1.7.0"
-ndarray = "0.15.6"
+ndarray = "0.16.1"
 serde = { version = "1.0.160", features = ["rc"] }
 serde_derive = "1.0.160"
-ordered-float = "3.6.0"
+ordered-float = "4.2.2"

ppca/src/dataset.rs

@@ -193,7 +193,7 @@ impl Dataset {
     /// Lists the dimensions which as masked in __all__ samples in this dataset.
     pub fn empty_dimensions(&self) -> Vec<usize> {
         let Some(n_dimensions) = self.data.first().map(|sample| sample.mask().0.len()) else {
-            return vec![]
+            return vec![];
         };
         let new_mask = || BitVec::from_elem(n_dimensions, false);
         let poormans_or = |mut this: BitVec, other: &BitVec| {

ppca/src/output_covariance.rs

@@ -1,6 +1,7 @@
 use nalgebra::{DMatrix, DVector};
 use serde_derive::{Deserialize, Serialize};
 use std::borrow::Cow;
+use std::sync::OnceLock;
 
 use crate::utils::Mask;
 
@@ -21,6 +22,18 @@ pub(crate) struct OutputCovariance<'a> {
     pub(crate) isotropic_noise: f64,
     /// The matrix mapping hidden state to output state, denoted as `C`.
     pub(crate) transform: Cow<'a, DMatrix<f64>>,
+    #[serde(default)]
+    #[serde(skip_serializing)]
+    #[serde(skip_deserializing)]
+    inner_product: OnceLock<DMatrix<f64>>,
+    #[serde(default)]
+    #[serde(skip_serializing)]
+    #[serde(skip_deserializing)]
+    inner_matrix: OnceLock<DMatrix<f64>>,
+    #[serde(default)]
+    #[serde(skip_serializing)]
+    #[serde(skip_deserializing)]
+    inner_inverse: OnceLock<DMatrix<f64>>,
 }
 
 impl<'a> OutputCovariance<'a> {
@@ -31,6 +44,9 @@ impl<'a> OutputCovariance<'a> {
         OutputCovariance {
             isotropic_noise,
             transform: Cow::Owned(transform),
+            inner_product: OnceLock::new(),
+            inner_matrix: OnceLock::new(),
+            inner_inverse: OnceLock::new(),
         }
     }
 
@@ -54,21 +70,34 @@ impl<'a> OutputCovariance<'a> {
     //         + &*self.transform * self.transform.transpose()
     // }
 
-    pub(crate) fn inner_product(&self) -> DMatrix<f64> {
+    fn do_inner_product(&self) -> DMatrix<f64> {
         self.transform.transpose() * &*self.transform
     }
 
-    pub(crate) fn inner_matrix(&self) -> DMatrix<f64> {
-        DMatrix::identity(self.state_size(), self.state_size()) * self.isotropic_noise.powi(2)
-            + self.inner_product()
+    fn inner_product(&self) -> &DMatrix<f64> {
+        self.inner_product.get_or_init(|| self.do_inner_product())
     }
 
-    pub(crate) fn inner_inverse(&self) -> DMatrix<f64> {
+    fn do_inner_matrix(&self) -> DMatrix<f64> {
+        DMatrix::identity(self.state_size(), self.state_size())
+            + self.inner_product() / self.isotropic_noise.powi(2)
+    }
+
+    fn inner_matrix(&self) -> &DMatrix<f64> {
+        self.inner_matrix.get_or_init(|| self.do_inner_matrix())
+    }
+
+    fn do_inner_inverse(&self) -> DMatrix<f64> {
         self.inner_matrix()
+            .clone()
             .try_inverse()
             .expect("inner matrix is always invertible")
     }
 
+    fn inner_inverse(&self) -> &DMatrix<f64> {
+        self.inner_inverse.get_or_init(|| self.do_inner_inverse())
+    }
+
     /// Calculates the linear transformation that estimates the hidden state from the
     /// observation.
     ///
@@ -86,58 +115,51 @@ impl<'a> OutputCovariance<'a> {
     /// ```
     /// C^T/sigma^2 - C^T*C/sigma^2*(I + C^T*C/sigma^2)^-1*C^T/sigma^2
     /// ```
-    /// Which can be calculated in `O(output_length * state_length^3)`.
+    /// Which can be calculated in `O(output_length * state_length^3)`. This can be futher simplified to
+    /// ```
+    /// (I - C^T*C/sigma^2*(I + C^T*C/sigma^2)^-1) * C^T/sigma^2
+    ///     = ((I + C^T*C/sigma^2) - C^T*C/sigma^2) * (I + C^T*C/sigma^2)^-1 * C^T/sigma^2
+    ///     = (I + C^T*C/sigma^2)^-1 * C^T/sigma^2
+    /// ```
+    /// Which retains the same complexity, but uses fewer operations.
     pub(crate) fn estimator_transform(&self) -> DMatrix<f64> {
-        (self.transform.transpose()
-            - self.inner_product() * self.inner_inverse() * self.transform.transpose())
-            / self.isotropic_noise.powi(2)
+        self.inner_inverse() * self.transform.transpose() / self.isotropic_noise.powi(2)
     }
 
     /// The covariance of the estimator that estimates hidden state from the observation.
     /// See `OutputCovariance.estimator_transform` for the explanation on the derivation.
     pub(crate) fn estimator_covariance(&self) -> DMatrix<f64> {
-        DMatrix::identity(self.state_size(), self.state_size())
-            - self.estimator_transform() * &*self.transform
+        self.inner_inverse().clone()
     }
 
     /// Calculates the log of the determinant of the output covariance matrix form masked
     /// data. This uses the _Matrix Determinant Lemma_ shenanigan to speed up computation:
     /// ```
     /// det(I * sigma^2 + C * C^T) = det(I + C^T * C / sigma^2) * det(I * sigma^2)
     /// ```
-    /// This can be simplified to
-    /// ```
-    /// det(I * sigma^2 + C * C^T) = det(I * sigma^2 + C^T * C)
-    ///     * sigma^(2 * (output_size - state_size))
-    /// ```
     /// The first `det` on the right side is the determinant of
     /// `OutputCovariance.inner_matrix`.
     pub(crate) fn covariance_log_det(&self) -> f64 {
-        // NOTE: not always `output_size > state_size`.
         self.inner_matrix().determinant().ln()
-            + self.isotropic_noise.ln()
-                * 2.0
-                * (self.output_size() as f64 - self.state_size() as f64)
+            + self.isotropic_noise.ln() * 2.0 * (self.output_size() as f64)
     }
 
     pub(crate) fn masked(&self, mask: &Mask) -> OutputCovariance<'static> {
         assert_eq!(mask.0.len(), self.output_size());
-        OutputCovariance {
-            isotropic_noise: self.isotropic_noise,
-            transform: Cow::Owned(DMatrix::from_rows(
-                &mask.filter(self.transform.row_iter()).collect::<Vec<_>>(),
-            )),
-        }
+        OutputCovariance::new_owned(
+            self.isotropic_noise,
+            DMatrix::from_rows(&mask.filter(self.transform.row_iter()).collect::<Vec<_>>()),
+        )
     }
 
     pub(crate) fn quadratic_form(&self, x: &DVector<f64>) -> f64 {
         let norm_squared = x.norm_squared();
         let transpose_transformed = self.transform.transpose() * x;
 
         (norm_squared
-            // this is a scalar.
             - (transpose_transformed.transpose() * self.inner_inverse() * transpose_transformed)
-                [(0, 0)])
+                [(0, 0)]
+                / self.isotropic_noise.powi(2))
             / self.isotropic_noise.powi(2)
     }
 }

ppca/src/ppca_model.rs

@@ -5,7 +5,6 @@ use rand::Rng;
 use rand_distr::Bernoulli;
 use rayon::prelude::*;
 use serde_derive::{Deserialize, Serialize};
-use std::borrow::Cow;
 use std::sync::Arc;
 
 use crate::dataset::{Dataset, MaskedSample};
@@ -61,10 +60,7 @@ impl PPCAModel {
         }
 
         PPCAModel(Arc::new(PPCAModelInner {
-            output_covariance: OutputCovariance {
-                isotropic_noise: 1.0,
-                transform: Cow::Owned(rand_transform),
-            },
+            output_covariance: OutputCovariance::new_owned(1.0, rand_transform),
             mean: DVector::zeros(output_size),
         }))
     }
@@ -384,10 +380,10 @@ impl PPCAModel {
         }
 
         PPCAModel(Arc::new(PPCAModelInner {
-            output_covariance: OutputCovariance {
-                transform: Cow::Owned(new_transform),
-                isotropic_noise: isotropic_noise_sq.sqrt(),
-            },
+            output_covariance: OutputCovariance::new_owned(
+                isotropic_noise_sq.sqrt(),
+                new_transform,
+            ),
             mean: new_mean,
         }))
     }