store the complete KstestResult at each iteration

Author: fra-pcmgf

dyson_equalizer/dyson_equalizer.py

@@ -9,6 +9,7 @@
 import numpy as np
 from matplotlib import pyplot as plt
 from scipy.stats import ks_1samp
+from scipy.stats._stats_py import KstestResult # noqa
 
 from dyson_equalizer.algorithm import compute_scaling_factors, compute_low_rank_approximation_mp, scale_matrix, \
     marchenko_pastur_cdf
@@ -17,11 +18,39 @@
 
 @dataclass
 class IterationStatistics:
+    """
+    This class stores the results of each iteration of the Dyson Equalizer
+
+    Attributes
+    ----------
+    delta_Y_hat: float
+        The mean absolute difference in values between Y_hat of this iteration and the previous one
+
+    x_hat_mean: float
+        The mean of the x_hat values
+
+    y_hat_mean: float
+        The mean of the y_hat values
+
+    kstest_result: KstestResult
+        The result of the Kolmogorov-Smirnov test
+
+    """
     delta_Y_hat: float = np.nan
-    ks_pvalue: float = np.nan
     x_hat_mean: float = np.nan
     y_hat_mean: float = np.nan
+    kstest_result: KstestResult = None
 
+    @property
+    def ks_pvalue(self) -> float:
+        """
+        the p-value of the Kolmogorov-Smirnov test
+        Returns
+        -------
+        p-value: float
+            the p-value of the Kolmogorov-Smirnov test
+        """
+        return self.kstest_result.pvalue if self.kstest_result else np.nan
 
 class DysonEqualizer:
     """
@@ -143,7 +172,7 @@ def compute(
             self.S = svd.S
         self.S_hat = svd_hat.S
 
-        stats.ks_pvalue = self.ks_pvalue_Y_hat()
+        stats.kstest_result = self.ks_Y_hat()
         self.iteration_statistics.append(stats)
         return self
 
@@ -378,6 +407,24 @@ def plot_mp_density_Y_hat(
             ax=ax,
         )
 
+
+    def ks_Y_hat(
+            self
+    ) -> KstestResult:
+        """ Computes the Kolmogorov–Smirnov test between the density of eigenvalues of ¹⁄ₙŶŶᵀ
+            and the Marchenko-Pastur distribution
+
+        Returns
+        -------
+        result : KstestResult
+            The result of the Kolmogorov–Smirnov test
+
+        """
+        m, n = sorted(self.Y_hat.shape)
+        eigs = self.S_hat ** 2 / n
+        ksr = ks_1samp(eigs, cdf=marchenko_pastur_cdf, args=[m/n])
+        return ksr
+
     def ks_pvalue_Y_hat(
             self
     ) -> float:
@@ -386,14 +433,28 @@ def ks_pvalue_Y_hat(
 
         Returns
         -------
-        p-value : DysonEqualizer
+        p-value : float
             The p-value of the Kolmogorov–Smirnov test
 
         """
-        m, n = sorted(self.Y_hat.shape)
-        eigs = self.S_hat ** 2 / n
+        return self.ks_Y_hat().pvalue
+
+    def ks_Y(
+            self
+    ) -> KstestResult:
+        """ Computes the Kolmogorov–Smirnov test between the density of eigenvalues of ¹⁄ₙYYᵀ
+            and the Marchenko-Pastur distribution
+
+        Returns
+        -------
+        result : KstestResult
+            The result of the Kolmogorov–Smirnov test
+
+        """
+        m, n = sorted(self.Y.shape)
+        eigs = self.S ** 2 / n
         ksr = ks_1samp(eigs, cdf=marchenko_pastur_cdf, args=[m/n])
-        return ksr.pvalue
+        return ksr
 
     def ks_pvalue_Y(
             self
@@ -403,11 +464,8 @@ def ks_pvalue_Y(
 
         Returns
         -------
-        p-value : DysonEqualizer
+        p-value : float
             The p-value of the Kolmogorov–Smirnov test
 
         """
-        m, n = sorted(self.Y.shape)
-        eigs = self.S ** 2 / n
-        ksr = ks_1samp(eigs, cdf=marchenko_pastur_cdf, args=[m/n])
-        return ksr.pvalue
+        return self.ks_Y().pvalue