fix: PatternHoleGenerator

Author: anhkhoangoho

qolmat/benchmark/metrics.py

@@ -384,32 +384,6 @@ def kolmogorov_smirnov_test(
     )
 
 
-def distance_correlation_complement(
-    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame
-) -> pd.Series:
-    """Correlation distance between columns of 2 dataframes.
-
-    Parameters
-    ----------
-    df1 : pd.DataFrame
-        true dataframe
-    df2 : pd.DataFrame
-        predicted dataframe
-    df_mask : pd.DataFrame
-        Elements of the dataframes to compute on
-
-    Returns
-    -------
-    pd.Series
-        Correlation distance
-    """
-
-    df1 = df1.fillna(0.0)
-    df2 = df2[["TEMP", "PRES"]].fillna(0.0)
-    print(df1.shape, df2.shape)
-    return 1.0 - pd.Series(dcor.distance_correlation(df1.values, df2.values), index=["All"])
-
-
 def _total_variance_distance_1D(df1: pd.Series, df2: pd.Series) -> float:
     """Compute Total Variance Distance for a categorical feature
     It is based on TVComplement in https://github.com/sdv-dev/SDMetrics
@@ -859,3 +833,31 @@ def frechet_distance(
         return pd.Series((frechet_dist / df_true.shape[0]), index=["All"])
     else:
         return pd.Series(np.repeat(frechet_dist, len(df1.columns)))
+
+
+def distance_correlation_complement(
+    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame
+) -> pd.Series:
+    """Correlation distance between columns of 2 dataframes.
+
+    Parameters
+    ----------
+    df1 : pd.DataFrame
+        true dataframe
+    df2 : pd.DataFrame
+        predicted dataframe
+    df_mask : pd.DataFrame
+        Elements of the dataframes to compute on
+
+    Returns
+    -------
+    pd.Series
+        Correlation distance
+    """
+    df1[~df_mask] = np.nan
+    df2[~df_mask] = np.nan
+
+    df1 = df1.dropna(axis=0, how="all").dropna(axis=1, how="any")
+    df2 = df2.dropna(axis=0, how="all").dropna(axis=1, how="any")
+
+    return 1.0 - pd.Series(dcor.distance_correlation(df1.values, df2.values), index=["All"])

qolmat/benchmark/missing_patterns.py

@@ -4,6 +4,7 @@
 import logging
 from typing import Callable, List, Optional, Tuple, Union
 
+import random
 import numpy as np
 import pandas as pd
 from sklearn.model_selection import GroupShuffleSplit
@@ -753,8 +754,8 @@ def fit(self, X: pd.DataFrame) -> PatternHoleGenerator:
         super().fit(X)
         df_isna = X.isna().apply(lambda x: self.get_pattern(x), axis=1).to_frame(name="pattern")
         self.df_isna = df_isna["pattern"]
-
-        patterns = self.df_isna.value_counts().index.to_list()
+        self.patterns_counts = self.df_isna.value_counts()
+        patterns = self.patterns_counts.index.to_list()
         if "_ALLNAN_" in patterns:
             patterns.remove("_ALLNAN_")
         if "_EMPTY_" in patterns:
@@ -772,11 +773,29 @@ def generate_mask(self, X: pd.DataFrame) -> pd.DataFrame:
         X : pd.DataFrame
             Initial dataframe with a missing pattern to be imitated.
         """
+        if self.ngroups is not None:
+            self.fit(X)
+
         df_mask = pd.DataFrame(False, index=X.index, columns=X.columns)
-        patterns_selected = ["_EMPTY_"] + [self.patterns[0]]
-        df_mask.loc[self.df_isna[self.df_isna.isin(patterns_selected)].index] = True
-        df_mask[X.isna()] = False
+        patterns_selected = ["_EMPTY_"]
+        patterns = self.patterns
+        for k in range(len(self.patterns)):
+            pattern = random.choice(patterns)
+            patterns_selected_ = patterns_selected + [pattern]
+            patterns.remove(pattern)
+            df_mask_ = df_mask.copy()
+            X_ = X.copy()
+
+            df_mask_.loc[self.df_isna[self.df_isna.isin(patterns_selected_)].index] = True
+            X_[~df_mask_] = np.nan
+            X_ = X_.dropna(axis=0, how="all").dropna(axis=1, how="any")
+            if X_.size == 0:
+                break
+            patterns_selected.append(pattern)
+            if self.patterns_counts.loc[patterns_selected_].sum() / len(X) > self.ratio_masked:
+                break
 
+        df_mask.loc[self.df_isna[self.df_isna.isin(patterns_selected)].index] = True
         return df_mask
 
     def get_pattern(self, row: pd.Series) -> str: