ENH: softimpute

Author: Hông-Lan Botterman

qolmat/imputations/imputers.py

@@ -1791,7 +1791,7 @@ def __init__(
         tau: float = 0,
         max_iterations: int = 100,
         verbose: bool = False,
-        projected: bool = False,
+        projected: bool = True,
     ):
         super().__init__(
             imputer_params=(
@@ -1844,7 +1844,7 @@ def _fit_element(
         self._check_dataframe(df)
         assert col == "__all__"
         hyperparams = self.get_hyperparams()
-        model = softimpute.SoftImpute(**hyperparams)
+        model = softimpute.SoftImpute(random_state=self._rng, **hyperparams)
         model = model.fit(df.values)
         return model
 

qolmat/utils/plot.py

@@ -9,6 +9,7 @@
 import matplotlib.pyplot as plt
 import matplotlib.ticker as plticker
 import numpy as np
+from numpy.typing import NDArray
 import pandas as pd
 import scipy
 from mpl_toolkits.axes_grid1 import make_axes_locatable
@@ -155,53 +156,72 @@ def plot_images(
 
 
 def make_ellipses(
-    X: np.ndarray,
+    x: NDArray,
+    y: NDArray,
     ax: mpl.axes.Axes,
-    color: Union[str, Tuple[float, float, float]],
+    n_std: float = 2,
+    color: Union[str, Any, Tuple[float, float, float]] = "None",
 ):
-    """Draw ellipses on a figure
+    """
+    Create a plot of the covariance confidence ellipse of *x* and *y*.
 
     Parameters
     ----------
-    X : np.ndarray
-        array for the ellipse
-    ax : matplotlib.axes._subplots.AxesSubplot
-        matplotlib ax handles
-    color : Union[str, Tuple[float, float, float]]
-        ellipse's color
+    x, y : array-like, shape (n, )
+        Input data.
+
+    ax : matplotlib.axes.Axes
+        The axes object to draw the ellipse into.
+
+    n_std : float
+        The number of standard deviations to determine the ellipse's radiuses.
+
+    color : Optional[str]
+        facecolor
+
+    Returns
+    -------
+    matplotlib.patches.Ellipse
     """
-    covariances = np.cov(X)  # gmm.covariances_[0] # [n][:2, :2]
-    v, w = np.linalg.eigh(covariances)
-    u = w[0] / np.linalg.norm(w[0])
-    angle = np.arctan2(u[1], u[0])
-    angle = 180 * angle / np.pi  # convert to degrees
-    center = X.mean(axis=0)  # .means_[0]
-    v[v < 0] = np.nan
-    v = 2.0 * np.sqrt(2.0) * np.sqrt(v)
-    ell = mpl.patches.Ellipse(center, v[0], v[1], angle=180 + angle, color=color)
+    if x.size != y.size:
+        raise ValueError("x and y must be the same size")
+
+    cov = np.cov(x, y)
+    pearson = cov[0, 1] / np.sqrt(cov[0, 0] * cov[1, 1])
+    ell_radius_x = np.sqrt(1 + pearson) * 2.5
+    ell_radius_y = np.sqrt(1 - pearson) * 2.5
+    ell = mpl.patches.Ellipse((0, 0), width=ell_radius_x, height=ell_radius_y, facecolor=color)
+    scale_x = np.sqrt(cov[0, 0]) * n_std
+    mean_x = np.mean(x)
+    scale_y = np.sqrt(cov[1, 1]) * n_std
+    mean_y = np.mean(y)
+    transf = (
+        mpl.transforms.Affine2D().rotate_deg(45).scale(scale_x, scale_y).translate(mean_x, mean_y)
+    )
+    ell.set_transform(transf + ax.transData)
+    ax.add_patch(ell)
     ell.set_clip_box(ax.bbox)
-    ell.set_alpha(0.5)
-    ax.add_artist(ell)
+    ell.set_alpha(0.4)
     ax.set_aspect("equal", "datalim")
 
 
 def compare_covariances(
-    df1: pd.DataFrame,
-    df2: pd.DataFrame,
+    df_1: pd.DataFrame,
+    df_2: pd.DataFrame,
     col_x: str,
     col_y: str,
-    ax: mpl.axes._subplots.AxesSubplot,
+    ax: mpl.axes.Axes,
     label: str = "",
-    color: Optional[Union[str, Tuple[float, float, float]]] = None,
+    color: Union[None, str, Tuple[float, float, float], Tuple[float, float, float, float]] = None,
 ):
     """
     Covariance plot: scatter plot with ellipses
 
     Parameters
     ----------
-    df1 : pd.DataFrame
+    df_1 : pd.DataFrame
         dataframe with raw data
-    df2 : pd.DataFrame
+    df_2 : pd.DataFrame
         dataframe with imputations
     col_x : str
         variable x, column's name of dataframe df1 to compare with
@@ -210,12 +230,14 @@ def compare_covariances(
     ax : matplotlib.axes._subplots.AxesSubplot
         matplotlib ax handles
     """
+    df1 = df_1.dropna()
+    df2 = df_2.dropna()
     if color is None:
         color = tab10(0)
-    ax.scatter(df2[col_x], df2[col_y], marker=".", color=color, s=0.2, label=label)
-    ax.scatter(df1[col_x], df1[col_y], marker=".", color="black", s=0.2)
-    make_ellipses(df1[[col_x, col_y]], ax, "black")
-    make_ellipses(df2[[col_x, col_y]], ax, color)
+    ax.scatter(df2[col_x], df2[col_y], marker=".", color=color, s=2, alpha=0.7, label="imputed")
+    ax.scatter(df1[col_x], df1[col_y], marker=".", color="black", s=2, alpha=0.7, label="original")
+    make_ellipses(df1[col_x], df1[col_y], ax, color="black")
+    make_ellipses(df2[col_x], df2[col_y], ax, color=color)
     ax.set_xlabel(col_x)
     ax.set_ylabel(col_y)
 

tests/imputations/test_imputers.py

@@ -275,18 +275,31 @@ def test_ImputerRPCA_fit_transform(df: pd.DataFrame) -> None:
     np.testing.assert_allclose(result, expected, atol=1e-2)
 
 
+@pytest.mark.parametrize("df", [df_incomplete])
+def test_ImputerSoftImpute_fit_transform(df: pd.DataFrame) -> None:
+    imputer = imputers.ImputerSoftImpute(
+        columnwise=False, max_iterations=100, tau=0.3, random_state=4
+    )
+    result = imputer.fit_transform(df)
+    expected = pd.DataFrame(
+        {
+            "col1": [0, 1.327, 2, 3, 0.137],
+            "col2": [-1, 0.099, 0.5, 0.122, 1.5],
+        }
+    )
+    np.testing.assert_allclose(result, expected, atol=1e-2)
+
+
 @pytest.mark.parametrize("df", [df_timeseries])
 def test_ImputerEM_fit_transform(df: pd.DataFrame) -> None:
     imputer = imputers.ImputerEM(method="sample", dt=1e-3, random_state=42)
     result = imputer.fit_transform(df)
-    print(result)
     expected = pd.DataFrame(
         {
             "col1": [i for i in range(20)],
             "col2": [0, 0.773, 2, 2.621, 2] + [i for i in range(5, 20)],
         }
     )
-    print(result)
     np.testing.assert_allclose(result, expected, atol=1e-2)
 
 
@@ -307,7 +320,6 @@ def test_ImputerEM_fit_transform(df: pd.DataFrame) -> None:
     imputers.ImputerMICE(groups=("group",)),
     imputers.ImputerRegressor(groups=("group",), estimator=LinearRegression()),
     imputers.ImputerRPCA(groups=("group",)),
-    imputers.ImputerSoftImpute(groups=("group",)),
     imputers.ImputerEM(groups=("group",)),
 ]
 
@@ -335,7 +347,6 @@ def test_models_fit_transform_grouped(imputer):
         imputers.ImputerMICE(),
         imputers.ImputerRegressor(),
         imputers.ImputerRPCA(tau=0, lam=0),
-        imputers.ImputerSoftImpute(),
         imputers.ImputerEM(),
     ]
 )

tests/utils/test_plot.py

@@ -75,7 +75,7 @@ def test__utils_plot_plot_images(
 def test_utils_plot_make_ellipses(X: np.ndarray, mocker: MockerFixture):
     mocker.patch("matplotlib.pyplot.show")
     ax = plt.gca()
-    plot.make_ellipses(X, ax, color="blue")
+    plot.make_ellipses(X[1], X[2], ax, color="blue")
     assert len(plt.gcf().get_axes()) > 0
     plt.close("all")