Merge pull request #203 from sdmccabe/merge-cmr-rcmr

Merge regularized CM into CM reconstructor

netrd/reconstruction/__init__.py

@@ -1,7 +1,6 @@
 from .base import BaseReconstructor
 from .random import RandomReconstructor
 from .correlation_matrix import CorrelationMatrixReconstructor
-from .regularized_correlation_matrix import RegularizedCorrelationMatrixReconstructor
 from .partial_correlation_matrix import PartialCorrelationMatrixReconstructor
 from .partial_correlation_influence import PartialCorrelationInfluenceReconstructor
 from .free_energy_minimization import FreeEnergyMinimizationReconstructor
@@ -22,7 +21,6 @@
 __all__ = [
     'RandomReconstructor',
     'CorrelationMatrixReconstructor',
-    'RegularizedCorrelationMatrixReconstructor',
     'PartialCorrelationMatrixReconstructor',
     'PartialCorrelationInfluenceReconstructor',
     'FreeEnergyMinimizationReconstructor',

netrd/reconstruction/correlation_matrix.py

@@ -16,18 +16,29 @@
 
 
 class CorrelationMatrixReconstructor(BaseReconstructor):
-    def fit(self, TS, threshold_type='range', **kwargs):
-        """Uses an unregularized form of the precision matrix.
+    def fit(self, TS, num_eigs=None, threshold_type='range', **kwargs):
+        """Uses the correlation matrix.
 
-        The results dictionary also stores the raw correlation matrix as
-        `'weights_matrix'` and the thresholded version of the correlation
-        matrix as `'thresholded_matrix'`. For details see [1].
+        If ``num_eigs`` is `None`, perform the reconstruction using the
+        unregularized correlation matrix. Otherwise, construct a regularized
+        precision matrix using ``num_eigs`` eigenvectors and eigenvalues of the
+        correlation matrix. For details on the regularization method, see [1].
+
+        The results dictionary also stores the raw correlation matrix
+        (potentially regularized) as `'weights_matrix'` and the thresholded
+        version of the correlation matrix as `'thresholded_matrix'`. For
+        details see [2].
 
         Parameters
         ----------
 
         TS (np.ndarray)
-            Array consisting of $L$ observations from $N$ sensors
+            Array consisting of :math:`L` observations from :math:`N` sensors
+
+        num_eigs (int)
+            The number of eigenvalues to use. (This corresponds to the
+            amount of regularization.) The number of eigenvalues used must
+            be less than :math:`N`.
 
         threshold_type (str)
             Which thresholding function to use on the matrix of
@@ -43,16 +54,45 @@ def fit(self, TS, threshold_type='range', **kwargs):
         References
         ----------
 
-        [1] https://github.com/valeria-io/visualising_stocks_correlations/blob/master/corr_matrix_viz.ipynb
+        [1] https://bwlewis.github.io/correlation-regularization/
+
+        [2] https://github.com/valeria-io/visualising_stocks_correlations/blob/master/corr_matrix_viz.ipynb
 
         """
 
         # get the correlation matrix
         cor = np.corrcoef(TS)
-        self.results['weights_matrix'] = cor
+
+        if num_eigs:
+            N = TS.shape[0]
+            if num_eigs > N:
+                raise ValueError(
+                    "The number of eigenvalues used must be less "
+                    "than the number of sensors."
+                )
+
+            # get eigenvalues and eigenvectors of the correlation matrix
+            vals, vecs = np.linalg.eigh(cor)
+            idx = vals.argsort()[::-1]
+            vals = vals[idx]
+            vecs = vecs[:, idx]
+
+            # construct the precision matrix and store it
+            P = (vecs[:, :num_eigs]) @ (
+                1 / (vals[:num_eigs]).reshape(num_eigs, 1) * (vecs[:, :num_eigs]).T
+            )
+            P = P / (
+                np.sqrt(np.diag(P)).reshape(N, 1) @ np.sqrt(np.diag(P)).reshape(1, N)
+            )
+            mat = P
+        else:
+            mat = cor
+
+        # store the appropriate source matrix
+        self.results['weights_matrix'] = mat
 
         # threshold the correlation matrix
-        A = threshold(cor, threshold_type, **kwargs)
+        A = threshold(mat, threshold_type, **kwargs)
         self.results['thresholded_matrix'] = A
 
         # construct the network