From 0e67a0d799b7da2f025c7b274e2bf205f029572f Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Sun, 16 Mar 2025 21:10:27 -0400
Subject: [PATCH 01/12] added scalar mnar simulator

---
 .../datasets/synthetic_data/loader.py         | 89 ++++++++++++++++++-
 1 file changed, 86 insertions(+), 3 deletions(-)

diff --git a/src/nearest_neighbors/datasets/synthetic_data/loader.py b/src/nearest_neighbors/datasets/synthetic_data/loader.py
index 4e08ae8..f1f637f 100644
--- a/src/nearest_neighbors/datasets/synthetic_data/loader.py
+++ b/src/nearest_neighbors/datasets/synthetic_data/loader.py
@@ -1,6 +1,7 @@
 from nearest_neighbors.dataloader_base import NNDataLoader
 from nearest_neighbors.dataloader_factory import register_dataset
 import numpy as np
+import math as math
 from typing import Any
 
 params = {
@@ -39,6 +40,7 @@
 }
 
 
+
 @register_dataset("synthetic_data", params)
 class SyntheticDataLoader(NNDataLoader):
     """Data from the Heartsteps V1 study formatted into a matrix or tensor.
@@ -123,6 +125,7 @@ def download_data(self) -> None:
         )
         pass
 
+
     def _generate_simulated_data(self) -> None:
         """Generates the simulated data with no missing values"""
         # row and column latent factors:
@@ -186,8 +189,88 @@ def _make_mcar(self) -> None:
         self.availability_mask = A
 
     def _make_mnar(self) -> None:
-        raise NotImplementedError("MNAR yet to be implemented")
-        # TODO: Aashish/Caleb/Kyuesong/Tatha: come back to this later
+        """Makes values missing not at random (MNAR), specifically staggered adoption (non-positivity & confounded)"""
+        U = self.row_latent
+
+        N = self.num_rows
+        T = self.num_cols
+
+        missing_mask = np.zeros((N, T))
+        pre_Masking = np.zeros((N, T))
+
+        # Divide units into 3 groups
+        g1_inds = np.arange(0, N // 3)
+        g2_inds = np.arange(N // 3, 2 * N // 3)
+        g3_inds = np.arange(2 * N // 3, N)
+
+        gamma_1 = [2, 0.7, 1, 0.7]
+        gamma_2 = [2, 0.2, 1, 0.2]
+
+        #TODO: make beta a parameter (currently hardcoded)
+        #first group adopts at the first 30% of the time period
+        #second group adopts at the first 70% of the time period
+        beta = [0.3, 0.7]
+
+        T1_lower = math.floor(T ** beta[0])
+        T2_lower = math.floor(T ** beta[1])
+
+        for i in range(N):
+            if i in g1_inds:
+                pre_Masking[i, :] = np.concatenate(
+                    (np.ones(T1_lower), np.zeros(T - T1_lower))
+                )
+                for t in range(T - T1_lower):
+                    pre_Masking[i, (t + T1_lower)] = np.random.binomial( #each units' adoption time probability is affected by their neighbors
+                        1,
+                        self._expit(
+                            gamma_1[0]
+                            + (0.99**t) * gamma_1[1] * U[i - 1]
+                            + gamma_1[2] * U[i]
+                            + (0.99**t) * gamma_1[3] * U[i + 1]
+                        ),
+                        1,
+                    )
+                pre_A = pre_Masking[i, :]
+                if len([i for i in range(len(pre_A)) if pre_A[i] == 0]) == 0:
+                    missing_mask[i, :] = pre_A
+                elif len([i for i in range(len(pre_A)) if pre_A[i] == 0]) > 0:
+                    adopt_time = min([i for i in range(len(pre_A)) if pre_A[i] == 0])
+                    missing_mask[i, :] = np.concatenate(
+                        (np.ones(adopt_time), np.zeros(T - adopt_time))
+                    )
+            elif i in g2_inds:
+                pre_Masking[i, :] = np.concatenate(
+                    (np.ones(T2_lower), np.zeros(T - T2_lower))
+                )
+                for t in range(T - T2_lower):
+                    pre_Masking[i, (t + T2_lower)] = np.random.binomial(
+                        1,
+                        self._expit(
+                            gamma_2[0]
+                            + (1.01**t) * gamma_2[1] * U[i - 1]
+                            + gamma_2[2] * U[i]
+                            + (1.01**t) * gamma_2[3] * U[i + 1]
+                        ),
+                        1,
+                    )
+                pre_A = pre_Masking[i, :]
+                if len([i for i in range(len(pre_A)) if pre_A[i] == 0]) == 0:
+                    missing_mask[i, :] = pre_A
+                elif len([i for i in range(len(pre_A)) if pre_A[i] == 0]) > 0:
+                    adopt_time = min([i for i in range(len(pre_A)) if pre_A[i] == 0])
+                    missing_mask[i, :] = np.concatenate(
+                        (np.ones(adopt_time), np.zeros(T - adopt_time))
+                    )
+            elif i in g3_inds:
+                missing_mask[i, :] = np.ones(T)
+
+        data_obs = self.data_noisy.copy()
+        data_obs[missing_mask] = np.nan
+        A = ~missing_mask  # A = NOT M, i.e. A_ij = 1 if Y_ij is observed, 0 if missing
+        self.data_obs = data_obs
+        self.availability_mask = A
+
+
 
     def get_full_state_as_dict(self, include_metadata: bool = False) -> dict:
         """Returns the full state of this object as a dictionary"""
@@ -257,7 +340,6 @@ def process_data_distribution(
         raise NotImplementedError(
             "Distributional setting not yet implemented for synthetic data"
         )
-
     # HELPER FUNCTIONS
     def _expit(self, x: np.ndarray) -> np.ndarray:
         """Helper function to apply the logistic sigmoid function to an array"""
@@ -283,3 +365,4 @@ def _transform_simulated_data(self, Y: np.ndarray) -> np.ndarray:
                 "non_lin must be one of '', 'expit', 'tanh', 'sin', 'cubic', or 'sinh'."
             )
         return Y
+    
\ No newline at end of file

From b52e3b3f2deccc88bc8aa510d6640e917974bcd8 Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Sun, 16 Mar 2025 22:09:52 -0400
Subject: [PATCH 02/12] fixed nnimputer fitmethods + created fit_method.py

---
 src/nearest_neighbors/nnimputer.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/nearest_neighbors/nnimputer.py b/src/nearest_neighbors/nnimputer.py
index 43976a3..be76c07 100644
--- a/src/nearest_neighbors/nnimputer.py
+++ b/src/nearest_neighbors/nnimputer.py
@@ -121,6 +121,8 @@ class FitMethod(ABC):
     @abstractmethod
     def fit(
         self,
+        row: int, 
+        column: int,
         data_array: npt.NDArray,
         mask_array: npt.NDArray,
         imputer: NearestNeighborImputer,
@@ -128,6 +130,8 @@ def fit(
         """Find the best distance threshold for the given data.
 
         Args:
+            row (int): Row index
+            column (int): Column index  
             data_array (npt.NDArray): Data matrix
             mask_array (npt.NDArray): Mask matrix
             imputer (NearestNeighborImputer): Imputer object

From 57087eb0b35a7b2d4b5f522788e33ee9dce32bf2 Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Tue, 18 Mar 2025 15:21:56 -0400
Subject: [PATCH 03/12] adding files

---
 src/nearest_neighbors/nnimputer.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/nearest_neighbors/nnimputer.py b/src/nearest_neighbors/nnimputer.py
index be76c07..2356bc9 100644
--- a/src/nearest_neighbors/nnimputer.py
+++ b/src/nearest_neighbors/nnimputer.py
@@ -121,8 +121,6 @@ class FitMethod(ABC):
     @abstractmethod
     def fit(
         self,
-        row: int, 
-        column: int,
         data_array: npt.NDArray,
         mask_array: npt.NDArray,
         imputer: NearestNeighborImputer,

From c1f84f4e66d2ac9cd6dd7e66b5f3e744684d0fe2 Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Tue, 18 Mar 2025 15:31:53 -0400
Subject: [PATCH 04/12] added files

---
 src/nearest_neighbors/fit_method.py | 113 ++++++++++++++++++++++++++++
 1 file changed, 113 insertions(+)
 create mode 100644 src/nearest_neighbors/fit_method.py

diff --git a/src/nearest_neighbors/fit_method.py b/src/nearest_neighbors/fit_method.py
new file mode 100644
index 0000000..8e8dd9a
--- /dev/null
+++ b/src/nearest_neighbors/fit_method.py
@@ -0,0 +1,113 @@
+from .nnimputer import FitMethod
+from .nnimputer import NearestNeighborImputer
+from .data_types import DistributionKernelMMD
+import numpy.typing as npt
+import numpy as np
+
+
+class DirectOptimization(FitMethod):
+    """Non-cross-validation fit method. Analytically optimizes the squared MMD error."""
+
+    def __init__(self, row: int, column: int, kernel: str, eta_cand: npt.NDArray, delta: float):
+        """Initialize the fit method with additional parameters.
+
+        Args:
+            kernel (str): Kernel to use for the MMD
+            eta_cand (npt.NDArray): Candidate distance thresholds
+            delta (float): Significance level
+            row (int): target row index
+            column (int): target column index
+
+        """
+        supported_kernels = ["exponential"]
+
+        if kernel not in supported_kernels:
+            raise ValueError(
+                f"Kernel {kernel} is not supported. Supported kernels are {supported_kernels}"
+            )
+
+        self.kernel = kernel
+        self.eta_cand = eta_cand
+        self.delta = delta
+        self.row = row
+        self.column = column
+
+    def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: NearestNeighborImputer) -> float:
+        """Analytically optimizes the squared MMD error.
+        
+        Args:
+            row (int): target row index
+            column (int): target column index
+            data_array (npt.NDArray): Data array
+            mask_array (npt.NDArray): Mask array
+            imputer (NearestNeighborImputer): Imputer
+
+        Returns:
+            float: Best distance threshold
+
+        """
+        if self.kernel == "exponential":
+            sup_kern = 1 #TODO: need to change for general kernels
+
+        delta = self.delta
+        eta_cand = self.eta_cand
+        row = self.row
+        column = self.column
+
+        n_rows, n_cols = data_array.shape[0], data_array.shape[1]
+        n = data_array[0, 0].shape[0] # number of samples per distribution
+        data_type = DistributionKernelMMD(self.kernel)
+
+        row_distances = np.zeros(n_rows)
+        for i in range(n_rows):
+            # Get columns observed in both row i and row
+            overlap_columns = np.logical_and(mask_array[row], mask_array[i])
+
+            if not np.any(overlap_columns):
+                row_distances[i] = np.inf
+                continue
+
+            # Calculate distance between rows
+            for j in range(n_cols):
+                if (
+                    not overlap_columns[j] or j == column
+                ):  # Skip missing values and the target column
+                    continue
+                row_distances[i] += data_type.distance(
+                    data_array[row, j], data_array[i, j]
+                )
+            row_distances[i] /= np.sum(overlap_columns)
+
+        perf = []
+
+        for eta in eta_cand:
+            neighborhood = np.where( (row_distances < eta)*(mask_array[:, column]) == 1 )[0] # Set of neighbors: (i) within eta distance (ii) observed
+
+            if sum(np.isin(neighborhood, row)) == 1: # Pretending as if (i, t) entry is missing
+                neighborhood = np.delete(neighborhood, np.where(neighborhood == i)[0])
+
+            if len(neighborhood) == 0: # Default (null) output when there is zero neighbor
+                perf.append(10**5) # Avoid selecting such eta without neighbors
+            else:
+                overlap = []
+                for neighbor in neighborhood:
+                    overlap.append(np.sum(mask_array[row, :]*mask_array[neighbor, :]))
+
+                Bias = 8*np.exp(1/np.exp(1))*sup_kern*np.log(2*n_rows/delta)/(np.sqrt(2*np.log(2)*np.min(overlap)))
+                Variance = 4*sup_kern*(np.log(n) + 1.5)/(n*len(neighborhood))
+
+                perf.append(eta + Bias + Variance)
+
+        eta_star= eta_cand[np.argmin(perf)]
+        return eta_star
+    
+
+
+
+class CrossValidation(FitMethod):
+    """Cross-validation fit method. Uses cross-validation to find the best distance threshold."""
+
+    def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: NearestNeighborImputer) -> float:
+        """Uses cross-validation to find the best distance threshold."""
+        pass
+    pass
\ No newline at end of file

From 98c0e9418504a77d23610de66b73f4c95755fc6a Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Tue, 18 Mar 2025 15:39:15 -0400
Subject: [PATCH 05/12] added (not aligned yet with the base loader) mcar
 simulator for distributional setting

---
 src/nearest_neighbors/simulations/mcar.py | 50 +++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/src/nearest_neighbors/simulations/mcar.py b/src/nearest_neighbors/simulations/mcar.py
index 4d0c7d1..eee6807 100644
--- a/src/nearest_neighbors/simulations/mcar.py
+++ b/src/nearest_neighbors/simulations/mcar.py
@@ -139,3 +139,53 @@ def expit(x: np.ndarray) -> np.ndarray:
     # Data[Masking == 0] = Y0[Masking == 0]
     Data: np.ndarray = np.array(Y)
     return Data, Theta, Masking
+
+
+def gendata_dist_mcar(N: int, T: int, n: int, d: int, p: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Generates multivariate Gaussian data of multiple measurements with latent dimension r = 2.
+
+    Args:
+        N (int): Number of users.
+        T (int): Number of time periods.
+        n (int): Number of samples per distribution.
+        d (int): Dimension of the data.
+        p (float): Probability of an entry being observed.
+        seed (int): Random seed for reproducibility.
+
+    Returns:
+        Data (np.ndarray): Generated data matrix of shape (N, T, n, d).
+        Masking (np.ndarray): Masking matrix indicating observed entries of shape (N, T).
+        True Mean (np.ndarray): True mean of the data of shape (N, T, d).
+        True Covariance (np.ndarray): True covariance of the data of shape (N, T, d, d).
+
+    """
+    np.random.seed(seed = seed)
+
+    ## Data Matrix (N * T * n * d)
+    Data = np.zeros( (N, T, n, d) )
+    true_Mean = np.zeros( (N, T, d) )
+    true_Cov = np.zeros( (N, T, d, d) )
+
+    u_1 = np.random.uniform(-1, 1, N)
+    u_2 = np.random.uniform(0.2, 1, N)
+
+    v_1 = np.random.uniform(-2, 2, T)
+    v_2 = np.random.uniform(0.5, 2, T)
+
+    even_ones = np.repeat([0, 1], d/2)
+    odd_ones = np.repeat([1, 0], d/2)
+
+    for i in range(N) : 
+        for t in range(T) : 
+            m_it = u_1[i]*v_1[t]*(even_ones - odd_ones)
+            c_it = np.diag(u_2[i]*v_2[t]*(0.5*even_ones + odd_ones))
+            true_Mean[i, t, :] = m_it
+            true_Cov[i, t, :, :] = c_it
+            dat_mat = np.random.multivariate_normal(m_it, c_it, size = n)
+            Data[i, t, :, :] = dat_mat
+    
+    Masking = np.zeros( (N, T) )
+
+    Masking = np.reshape(np.random.binomial(1, p, (N*T)), (N, T))
+    
+    return(Data, Masking, true_Mean, true_Cov)
\ No newline at end of file

From 1e3b48cdaaf7e5cf7eb7bcc47461620b0405d394 Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Tue, 18 Mar 2025 16:07:45 -0400
Subject: [PATCH 06/12] fixed ruff

---
 src/nearest_neighbors/fit_method.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/nearest_neighbors/fit_method.py b/src/nearest_neighbors/fit_method.py
index 8e8dd9a..ab8f23a 100644
--- a/src/nearest_neighbors/fit_method.py
+++ b/src/nearest_neighbors/fit_method.py
@@ -36,8 +36,6 @@ def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: Nearest
         """Analytically optimizes the squared MMD error.
         
         Args:
-            row (int): target row index
-            column (int): target column index
             data_array (npt.NDArray): Data array
             mask_array (npt.NDArray): Mask array
             imputer (NearestNeighborImputer): Imputer
@@ -47,7 +45,7 @@ def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: Nearest
 
         """
         if self.kernel == "exponential":
-            sup_kern = 1 #TODO: need to change for general kernels
+            sup_kern = 1  # TODO: need to change for general kernels
 
         delta = self.delta
         eta_cand = self.eta_cand

From 9233c79ed0f39ad3139c9cead17164c145cd6a83 Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Tue, 18 Mar 2025 16:15:33 -0400
Subject: [PATCH 07/12] fixed errors

---
 .../datasets/synthetic_data/loader.py         |  4 +-
 src/nearest_neighbors/fit_method.py           | 42 ++++++++++++-------
 src/nearest_neighbors/simulations/mcar.py     |  4 +-
 3 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/src/nearest_neighbors/datasets/synthetic_data/loader.py b/src/nearest_neighbors/datasets/synthetic_data/loader.py
index f1f637f..c8adf1e 100644
--- a/src/nearest_neighbors/datasets/synthetic_data/loader.py
+++ b/src/nearest_neighbors/datasets/synthetic_data/loader.py
@@ -265,8 +265,8 @@ def _make_mnar(self) -> None:
                 missing_mask[i, :] = np.ones(T)
 
         data_obs = self.data_noisy.copy()
-        data_obs[missing_mask] = np.nan
-        A = ~missing_mask  # A = NOT M, i.e. A_ij = 1 if Y_ij is observed, 0 if missing
+        data_obs[missing_mask.astype(bool)] = np.nan
+        A = ~missing_mask.astype(bool)  # A = NOT M, i.e. A_ij = 1 if Y_ij is observed, 0 if missing
         self.data_obs = data_obs
         self.availability_mask = A
 
diff --git a/src/nearest_neighbors/fit_method.py b/src/nearest_neighbors/fit_method.py
index ab8f23a..d5238f3 100644
--- a/src/nearest_neighbors/fit_method.py
+++ b/src/nearest_neighbors/fit_method.py
@@ -34,7 +34,7 @@ def __init__(self, row: int, column: int, kernel: str, eta_cand: npt.NDArray, de
 
     def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: NearestNeighborImputer) -> float:
         """Analytically optimizes the squared MMD error.
-        
+
         Args:
             data_array (npt.NDArray): Data array
             mask_array (npt.NDArray): Mask array
@@ -44,6 +44,8 @@ def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: Nearest
             float: Best distance threshold
 
         """
+        # Initialize sup_kern outside conditional
+        sup_kern = 1  # Default value
         if self.kernel == "exponential":
             sup_kern = 1  # TODO: need to change for general kernels
 
@@ -53,7 +55,7 @@ def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: Nearest
         column = self.column
 
         n_rows, n_cols = data_array.shape[0], data_array.shape[1]
-        n = data_array[0, 0].shape[0] # number of samples per distribution
+        n = data_array[0, 0].shape[0]  # number of samples per distribution
         data_type = DistributionKernelMMD(self.kernel)
 
         row_distances = np.zeros(n_rows)
@@ -79,24 +81,27 @@ def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: Nearest
         perf = []
 
         for eta in eta_cand:
-            neighborhood = np.where( (row_distances < eta)*(mask_array[:, column]) == 1 )[0] # Set of neighbors: (i) within eta distance (ii) observed
+            neighborhood = np.where((row_distances < eta) * (mask_array[:, column]) == 1)[0]  # Set of neighbors: (i) within eta distance (ii) observed
 
-            if sum(np.isin(neighborhood, row)) == 1: # Pretending as if (i, t) entry is missing
-                neighborhood = np.delete(neighborhood, np.where(neighborhood == i)[0])
+            if sum(np.isin(neighborhood, row)) == 1:  # Pretending as if (row, column) entry is missing
+                neighborhood = np.delete(neighborhood, np.where(neighborhood == row)[0])
 
-            if len(neighborhood) == 0: # Default (null) output when there is zero neighbor
-                perf.append(10**5) # Avoid selecting such eta without neighbors
+            if len(neighborhood) == 0:  # Default (null) output when there is zero neighbor
+                perf.append(10**5)  # Avoid selecting such eta without neighbors
             else:
                 overlap = []
                 for neighbor in neighborhood:
-                    overlap.append(np.sum(mask_array[row, :]*mask_array[neighbor, :]))
+                    overlap.append(np.sum(mask_array[row, :] * mask_array[neighbor, :]))
 
-                Bias = 8*np.exp(1/np.exp(1))*sup_kern*np.log(2*n_rows/delta)/(np.sqrt(2*np.log(2)*np.min(overlap)))
-                Variance = 4*sup_kern*(np.log(n) + 1.5)/(n*len(neighborhood))
+                Bias = 8 * np.exp(1/np.exp(1)) * sup_kern * np.log(2*n_rows/delta) / (np.sqrt(2*np.log(2)*np.min(overlap)))
+                Variance = 4 * sup_kern * (np.log(n) + 1.5) / (n*len(neighborhood))
 
                 perf.append(eta + Bias + Variance)
 
-        eta_star= eta_cand[np.argmin(perf)]
+        if not perf:  # Handle case when perf list is empty
+            return float('inf')  # Return infinity as a default value when no valid threshold is found
+            
+        eta_star = eta_cand[np.argmin(perf)]
         return eta_star
     
 
@@ -106,6 +111,15 @@ class CrossValidation(FitMethod):
     """Cross-validation fit method. Uses cross-validation to find the best distance threshold."""
 
     def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: NearestNeighborImputer) -> float:
-        """Uses cross-validation to find the best distance threshold."""
-        pass
-    pass
\ No newline at end of file
+        """Uses cross-validation to find the best distance threshold.
+        
+        Args:
+            data_array (npt.NDArray): Data array
+            mask_array (npt.NDArray): Mask array
+            imputer (NearestNeighborImputer): Imputer
+
+        Returns:
+            float: Best distance threshold
+            
+        """
+        return 0.0  # TODO: Implement cross-validation
\ No newline at end of file
diff --git a/src/nearest_neighbors/simulations/mcar.py b/src/nearest_neighbors/simulations/mcar.py
index eee6807..b287c5e 100644
--- a/src/nearest_neighbors/simulations/mcar.py
+++ b/src/nearest_neighbors/simulations/mcar.py
@@ -172,8 +172,8 @@ def gendata_dist_mcar(N: int, T: int, n: int, d: int, p: float, seed: int) -> Tu
     v_1 = np.random.uniform(-2, 2, T)
     v_2 = np.random.uniform(0.5, 2, T)
 
-    even_ones = np.repeat([0, 1], d/2)
-    odd_ones = np.repeat([1, 0], d/2)
+    even_ones = np.repeat([0, 1], int(d/2))
+    odd_ones = np.repeat([1, 0], int(d/2))
 
     for i in range(N) : 
         for t in range(T) : 

From 2d536628273c7dfc93df4fb033fc0178ed3b1954 Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Tue, 18 Mar 2025 16:21:14 -0400
Subject: [PATCH 08/12] fixed errors

---
 .../datasets/synthetic_data/loader.py         |  2 +-
 src/nearest_neighbors/fit_method.py           |  5 ++---
 src/nearest_neighbors/simulations/mcar.py     | 21 +++++++++----------
 3 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/src/nearest_neighbors/datasets/synthetic_data/loader.py b/src/nearest_neighbors/datasets/synthetic_data/loader.py
index c8adf1e..157c80e 100644
--- a/src/nearest_neighbors/datasets/synthetic_data/loader.py
+++ b/src/nearest_neighbors/datasets/synthetic_data/loader.py
@@ -220,7 +220,7 @@ def _make_mnar(self) -> None:
                     (np.ones(T1_lower), np.zeros(T - T1_lower))
                 )
                 for t in range(T - T1_lower):
-                    pre_Masking[i, (t + T1_lower)] = np.random.binomial( #each units' adoption time probability is affected by their neighbors
+                    pre_Masking[i, (t + T1_lower)] = np.random.binomial(  # each units' adoption time probability is affected by their neighbors
                         1,
                         self._expit(
                             gamma_1[0]
diff --git a/src/nearest_neighbors/fit_method.py b/src/nearest_neighbors/fit_method.py
index d5238f3..5ee8954 100644
--- a/src/nearest_neighbors/fit_method.py
+++ b/src/nearest_neighbors/fit_method.py
@@ -100,10 +100,9 @@ def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: Nearest
 
         if not perf:  # Handle case when perf list is empty
             return float('inf')  # Return infinity as a default value when no valid threshold is found
-            
+
         eta_star = eta_cand[np.argmin(perf)]
         return eta_star
-    
 
 
 
@@ -112,7 +111,7 @@ class CrossValidation(FitMethod):
 
     def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: NearestNeighborImputer) -> float:
         """Uses cross-validation to find the best distance threshold.
-        
+
         Args:
             data_array (npt.NDArray): Data array
             mask_array (npt.NDArray): Mask array
diff --git a/src/nearest_neighbors/simulations/mcar.py b/src/nearest_neighbors/simulations/mcar.py
index b287c5e..8c56537 100644
--- a/src/nearest_neighbors/simulations/mcar.py
+++ b/src/nearest_neighbors/simulations/mcar.py
@@ -159,12 +159,12 @@ def gendata_dist_mcar(N: int, T: int, n: int, d: int, p: float, seed: int) -> Tu
         True Covariance (np.ndarray): True covariance of the data of shape (N, T, d, d).
 
     """
-    np.random.seed(seed = seed)
+    np.random.seed(seed=seed)
 
     ## Data Matrix (N * T * n * d)
-    Data = np.zeros( (N, T, n, d) )
-    true_Mean = np.zeros( (N, T, d) )
-    true_Cov = np.zeros( (N, T, d, d) )
+    Data = np.zeros((N, T, n, d))
+    true_Mean = np.zeros((N, T, d))
+    true_Cov = np.zeros((N, T, d, d))
 
     u_1 = np.random.uniform(-1, 1, N)
     u_2 = np.random.uniform(0.2, 1, N)
@@ -175,17 +175,16 @@ def gendata_dist_mcar(N: int, T: int, n: int, d: int, p: float, seed: int) -> Tu
     even_ones = np.repeat([0, 1], int(d/2))
     odd_ones = np.repeat([1, 0], int(d/2))
 
-    for i in range(N) : 
-        for t in range(T) : 
+    for i in range(N):
+        for t in range(T):
             m_it = u_1[i]*v_1[t]*(even_ones - odd_ones)
             c_it = np.diag(u_2[i]*v_2[t]*(0.5*even_ones + odd_ones))
             true_Mean[i, t, :] = m_it
             true_Cov[i, t, :, :] = c_it
-            dat_mat = np.random.multivariate_normal(m_it, c_it, size = n)
+            dat_mat = np.random.multivariate_normal(m_it, c_it, size=n)
             Data[i, t, :, :] = dat_mat
-    
-    Masking = np.zeros( (N, T) )
 
+    Masking = np.zeros((N, T))
     Masking = np.reshape(np.random.binomial(1, p, (N*T)), (N, T))
-    
-    return(Data, Masking, true_Mean, true_Cov)
\ No newline at end of file
+
+    return Data, Masking, true_Mean, true_Cov
\ No newline at end of file

From 7f3c3c48134b54dd8710caaa1e4ea6dfdd263d1c Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Tue, 18 Mar 2025 16:39:19 -0400
Subject: [PATCH 09/12] formatting fit_method.py

---
 src/nearest_neighbors/fit_method.py | 49 ++++++++++++++++++++++-------
 1 file changed, 37 insertions(+), 12 deletions(-)

diff --git a/src/nearest_neighbors/fit_method.py b/src/nearest_neighbors/fit_method.py
index 5ee8954..0246c3d 100644
--- a/src/nearest_neighbors/fit_method.py
+++ b/src/nearest_neighbors/fit_method.py
@@ -8,7 +8,9 @@
 class DirectOptimization(FitMethod):
     """Non-cross-validation fit method. Analytically optimizes the squared MMD error."""
 
-    def __init__(self, row: int, column: int, kernel: str, eta_cand: npt.NDArray, delta: float):
+    def __init__(
+        self, row: int, column: int, kernel: str, eta_cand: npt.NDArray, delta: float
+    ):
         """Initialize the fit method with additional parameters.
 
         Args:
@@ -32,7 +34,12 @@ def __init__(self, row: int, column: int, kernel: str, eta_cand: npt.NDArray, de
         self.row = row
         self.column = column
 
-    def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: NearestNeighborImputer) -> float:
+    def fit(
+        self,
+        data_array: npt.NDArray,
+        mask_array: npt.NDArray,
+        imputer: NearestNeighborImputer,
+    ) -> float:
         """Analytically optimizes the squared MMD error.
 
         Args:
@@ -81,35 +88,53 @@ def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: Nearest
         perf = []
 
         for eta in eta_cand:
-            neighborhood = np.where((row_distances < eta) * (mask_array[:, column]) == 1)[0]  # Set of neighbors: (i) within eta distance (ii) observed
+            neighborhood = np.where(
+                (row_distances < eta) * (mask_array[:, column]) == 1
+            )[0]  # Set of neighbors: (i) within eta distance (ii) observed
 
-            if sum(np.isin(neighborhood, row)) == 1:  # Pretending as if (row, column) entry is missing
+            if (
+                sum(np.isin(neighborhood, row)) == 1
+            ):  # Pretending as if (row, column) entry is missing
                 neighborhood = np.delete(neighborhood, np.where(neighborhood == row)[0])
 
-            if len(neighborhood) == 0:  # Default (null) output when there is zero neighbor
+            if (
+                len(neighborhood) == 0
+            ):  # Default (null) output when there is zero neighbor
                 perf.append(10**5)  # Avoid selecting such eta without neighbors
             else:
                 overlap = []
                 for neighbor in neighborhood:
                     overlap.append(np.sum(mask_array[row, :] * mask_array[neighbor, :]))
 
-                Bias = 8 * np.exp(1/np.exp(1)) * sup_kern * np.log(2*n_rows/delta) / (np.sqrt(2*np.log(2)*np.min(overlap)))
-                Variance = 4 * sup_kern * (np.log(n) + 1.5) / (n*len(neighborhood))
+                Bias = (
+                    8
+                    * np.exp(1 / np.exp(1))
+                    * sup_kern
+                    * np.log(2 * n_rows / delta)
+                    / (np.sqrt(2 * np.log(2) * np.min(overlap)))
+                )
+                Variance = 4 * sup_kern * (np.log(n) + 1.5) / (n * len(neighborhood))
 
                 perf.append(eta + Bias + Variance)
 
         if not perf:  # Handle case when perf list is empty
-            return float('inf')  # Return infinity as a default value when no valid threshold is found
+            return float(
+                "inf"
+            )  # Return infinity as a default value when no valid threshold is found
 
         eta_star = eta_cand[np.argmin(perf)]
         return eta_star
 
 
-
 class CrossValidation(FitMethod):
     """Cross-validation fit method. Uses cross-validation to find the best distance threshold."""
 
-    def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: NearestNeighborImputer) -> float:
+    def fit(
+        self,
+        data_array: npt.NDArray,
+        mask_array: npt.NDArray,
+        imputer: NearestNeighborImputer,
+    ) -> float:
         """Uses cross-validation to find the best distance threshold.
 
         Args:
@@ -119,6 +144,6 @@ def fit(self, data_array: npt.NDArray, mask_array: npt.NDArray, imputer: Nearest
 
         Returns:
             float: Best distance threshold
-            
+
         """
-        return 0.0  # TODO: Implement cross-validation
\ No newline at end of file
+        return 0.0  # TODO: Implement cross-validation

From 2924a4e1c0df3fcbeddc728b07f6fea1016544a3 Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Tue, 18 Mar 2025 16:39:58 -0400
Subject: [PATCH 10/12] formatting loader.py

---
 .../datasets/synthetic_data/loader.py         | 36 +++++++++----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/nearest_neighbors/datasets/synthetic_data/loader.py b/src/nearest_neighbors/datasets/synthetic_data/loader.py
index 157c80e..edcd3dc 100644
--- a/src/nearest_neighbors/datasets/synthetic_data/loader.py
+++ b/src/nearest_neighbors/datasets/synthetic_data/loader.py
@@ -40,7 +40,6 @@
 }
 
 
-
 @register_dataset("synthetic_data", params)
 class SyntheticDataLoader(NNDataLoader):
     """Data from the Heartsteps V1 study formatted into a matrix or tensor.
@@ -125,7 +124,6 @@ def download_data(self) -> None:
         )
         pass
 
-
     def _generate_simulated_data(self) -> None:
         """Generates the simulated data with no missing values"""
         # row and column latent factors:
@@ -206,9 +204,9 @@ def _make_mnar(self) -> None:
         gamma_1 = [2, 0.7, 1, 0.7]
         gamma_2 = [2, 0.2, 1, 0.2]
 
-        #TODO: make beta a parameter (currently hardcoded)
-        #first group adopts at the first 30% of the time period
-        #second group adopts at the first 70% of the time period
+        # TODO: make beta a parameter (currently hardcoded)
+        # first group adopts at the first 30% of the time period
+        # second group adopts at the first 70% of the time period
         beta = [0.3, 0.7]
 
         T1_lower = math.floor(T ** beta[0])
@@ -220,15 +218,17 @@ def _make_mnar(self) -> None:
                     (np.ones(T1_lower), np.zeros(T - T1_lower))
                 )
                 for t in range(T - T1_lower):
-                    pre_Masking[i, (t + T1_lower)] = np.random.binomial(  # each units' adoption time probability is affected by their neighbors
-                        1,
-                        self._expit(
-                            gamma_1[0]
-                            + (0.99**t) * gamma_1[1] * U[i - 1]
-                            + gamma_1[2] * U[i]
-                            + (0.99**t) * gamma_1[3] * U[i + 1]
-                        ),
-                        1,
+                    pre_Masking[i, (t + T1_lower)] = (
+                        np.random.binomial(  # each units' adoption time probability is affected by their neighbors
+                            1,
+                            self._expit(
+                                gamma_1[0]
+                                + (0.99**t) * gamma_1[1] * U[i - 1]
+                                + gamma_1[2] * U[i]
+                                + (0.99**t) * gamma_1[3] * U[i + 1]
+                            ),
+                            1,
+                        )
                     )
                 pre_A = pre_Masking[i, :]
                 if len([i for i in range(len(pre_A)) if pre_A[i] == 0]) == 0:
@@ -266,12 +266,12 @@ def _make_mnar(self) -> None:
 
         data_obs = self.data_noisy.copy()
         data_obs[missing_mask.astype(bool)] = np.nan
-        A = ~missing_mask.astype(bool)  # A = NOT M, i.e. A_ij = 1 if Y_ij is observed, 0 if missing
+        A = ~missing_mask.astype(
+            bool
+        )  # A = NOT M, i.e. A_ij = 1 if Y_ij is observed, 0 if missing
         self.data_obs = data_obs
         self.availability_mask = A
 
-
-
     def get_full_state_as_dict(self, include_metadata: bool = False) -> dict:
         """Returns the full state of this object as a dictionary"""
         return_dict = {
@@ -340,6 +340,7 @@ def process_data_distribution(
         raise NotImplementedError(
             "Distributional setting not yet implemented for synthetic data"
         )
+
     # HELPER FUNCTIONS
     def _expit(self, x: np.ndarray) -> np.ndarray:
         """Helper function to apply the logistic sigmoid function to an array"""
@@ -365,4 +366,3 @@ def _transform_simulated_data(self, Y: np.ndarray) -> np.ndarray:
                 "non_lin must be one of '', 'expit', 'tanh', 'sin', 'cubic', or 'sinh'."
             )
         return Y
-    
\ No newline at end of file

From c3c54096cbb747aa58578195d01c11a00a63b33a Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Tue, 18 Mar 2025 16:40:22 -0400
Subject: [PATCH 11/12] formatting fixes

---
 src/nearest_neighbors/nnimputer.py        |  2 +-
 src/nearest_neighbors/simulations/mcar.py | 16 +++++++++-------
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/nearest_neighbors/nnimputer.py b/src/nearest_neighbors/nnimputer.py
index 2356bc9..b0b29c2 100644
--- a/src/nearest_neighbors/nnimputer.py
+++ b/src/nearest_neighbors/nnimputer.py
@@ -129,7 +129,7 @@ def fit(
 
         Args:
             row (int): Row index
-            column (int): Column index  
+            column (int): Column index
             data_array (npt.NDArray): Data matrix
             mask_array (npt.NDArray): Mask matrix
             imputer (NearestNeighborImputer): Imputer object
diff --git a/src/nearest_neighbors/simulations/mcar.py b/src/nearest_neighbors/simulations/mcar.py
index 8c56537..62fda56 100644
--- a/src/nearest_neighbors/simulations/mcar.py
+++ b/src/nearest_neighbors/simulations/mcar.py
@@ -141,7 +141,9 @@ def expit(x: np.ndarray) -> np.ndarray:
     return Data, Theta, Masking
 
 
-def gendata_dist_mcar(N: int, T: int, n: int, d: int, p: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+def gendata_dist_mcar(
+    N: int, T: int, n: int, d: int, p: float, seed: int
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """Generates multivariate Gaussian data of multiple measurements with latent dimension r = 2.
 
     Args:
@@ -172,19 +174,19 @@ def gendata_dist_mcar(N: int, T: int, n: int, d: int, p: float, seed: int) -> Tu
     v_1 = np.random.uniform(-2, 2, T)
     v_2 = np.random.uniform(0.5, 2, T)
 
-    even_ones = np.repeat([0, 1], int(d/2))
-    odd_ones = np.repeat([1, 0], int(d/2))
+    even_ones = np.repeat([0, 1], int(d / 2))
+    odd_ones = np.repeat([1, 0], int(d / 2))
 
     for i in range(N):
         for t in range(T):
-            m_it = u_1[i]*v_1[t]*(even_ones - odd_ones)
-            c_it = np.diag(u_2[i]*v_2[t]*(0.5*even_ones + odd_ones))
+            m_it = u_1[i] * v_1[t] * (even_ones - odd_ones)
+            c_it = np.diag(u_2[i] * v_2[t] * (0.5 * even_ones + odd_ones))
             true_Mean[i, t, :] = m_it
             true_Cov[i, t, :, :] = c_it
             dat_mat = np.random.multivariate_normal(m_it, c_it, size=n)
             Data[i, t, :, :] = dat_mat
 
     Masking = np.zeros((N, T))
-    Masking = np.reshape(np.random.binomial(1, p, (N*T)), (N, T))
+    Masking = np.reshape(np.random.binomial(1, p, (N * T)), (N, T))
 
-    return Data, Masking, true_Mean, true_Cov
\ No newline at end of file
+    return Data, Masking, true_Mean, true_Cov

From 438a4ff938c52aa3880248763aef300247e1f382 Mon Sep 17 00:00:00 2001
From: Kyuseong Choi <97979812+kyuseongchoi5@users.noreply.github.com>
Date: Mon, 28 Apr 2025 17:27:51 -0400
Subject: [PATCH 12/12] Adding data loader and demo file

---
 .../prompteval/demo_prompteval_dataloader.py  |  41 +++
 .../datasets/prompteval/__init__.py           |   1 +
 .../datasets/prompteval/loader.py             | 255 ++++++++++++++++++
 3 files changed, 297 insertions(+)
 create mode 100644 examples/prompteval/demo_prompteval_dataloader.py
 create mode 100644 src/nearest_neighbors/datasets/prompteval/__init__.py
 create mode 100644 src/nearest_neighbors/datasets/prompteval/loader.py

diff --git a/examples/prompteval/demo_prompteval_dataloader.py b/examples/prompteval/demo_prompteval_dataloader.py
new file mode 100644
index 0000000..8ac1f58
--- /dev/null
+++ b/examples/prompteval/demo_prompteval_dataloader.py
@@ -0,0 +1,41 @@
+"""Demo: Loading and inspecting the PromptEvaldataset using the nearest_neighbors package.
+
+Run this script from the root of the repository with:
+    python examples/demo_prompteval_dataloader.py
+"""
+
+from nearest_neighbors.datasets.dataloader_factory import NNData
+
+print("====== PromptEval Data Demo ======")
+
+# Show dataset help info
+print("\nDataset Help Information:")
+NNData.help("prompteval")
+
+# Load with default settings
+print("\n====== Example 1: Default Parameters ======")
+loader = NNData.create("prompteval", seed=42)
+data, mask = loader.process_data_scalar()
+
+print("Observed Ratings Matrix (with NaNs):")
+print(data)
+print("\nMask Matrix (True = rating present, False = missing):")
+print(mask)
+
+# Load with custom sampling
+print("\n====== Example 2: Custom Parameters ======")
+loader_custom = NNData.create("movielens", seed=42, sample_users=100, sample_movies=50)
+data_custom, mask_custom = loader_custom.process_data_scalar()
+
+print("Custom-Sized Ratings Matrix:")
+print(data_custom)
+print("\nCustom Mask Matrix:")
+print(mask_custom)
+
+# View saved internal state
+print("\n====== Full State Dictionary Keys ======")
+state = loader.get_full_state_as_dict(include_metadata=True)
+print("Top-level keys:")
+print(list(state.keys()))
+print("\nCustom parameters used:")
+print(state.get("custom_params", {}))
diff --git a/src/nearest_neighbors/datasets/prompteval/__init__.py b/src/nearest_neighbors/datasets/prompteval/__init__.py
new file mode 100644
index 0000000..342ca37
--- /dev/null
+++ b/src/nearest_neighbors/datasets/prompteval/__init__.py
@@ -0,0 +1 @@
+from .loader import PromptEvalDataLoader  # noqa: F401
diff --git a/src/nearest_neighbors/datasets/prompteval/loader.py b/src/nearest_neighbors/datasets/prompteval/loader.py
new file mode 100644
index 0000000..30cc019
--- /dev/null
+++ b/src/nearest_neighbors/datasets/prompteval/loader.py
@@ -0,0 +1,255 @@
+"""Dataset loader for the PromptEval (MMLU) dataset.
+
+Source1: https://huggingface.co/datasets/PromptEval
+Source2: https://github.com/kyuseongchoi5/EfficientEval_BayesOpt
+Paper Reference for transformation implementation:
+    Felipe Maia Polo et al.
+    "Efficient multi-prompt evaluation of LLMs."
+    Neurips, 2024.
+    https://arxiv.org/pdf/2405.17202
+"""
+
+from nearest_neighbors.datasets.dataloader_base import NNDataLoader
+from nearest_neighbors.datasets.dataloader_factory import register_dataset
+import numpy as np
+import pandas as pd
+import pickle
+from typing import Any
+import logging
+from joblib import Memory
+
+
+memory = Memory(".joblib_cache", verbose=2)
+logger = logging.getLogger(__name__)
+
+params = {
+    "tasks": (
+        list[str],
+        None,
+        "List of tasks to evaluate on. By default, returns all.",
+    ),
+    "models": (
+        list[str],
+        None,
+        "List of models to evaluate by. By default, returns all.",
+    ),
+    "seed": (int, None, "Random seed for reproducibility"),
+    "propensity": (float, None, "Proportion of data to keep"),
+}
+
+
+@register_dataset("prompteval", params)
+class PromptEvalDataLoader(NNDataLoader):
+    """Data from the PromptEval study formatted into a matrix or tensor.
+    To initialize with default settings, use: NNData.create("prompteval").
+
+    """
+
+    urls = {
+        "full_data": "https://github.com/kyuseongchoi5/EfficientEval_BayesOpt/tree/main/data/MMLU/data_all.pkl"
+    }
+
+    def __init__(
+        self,
+        tasks: list[str] | None = None,
+        models: list[str] | None = None,
+        seed: int | None = None,
+        propensity: float = 1.0,  # Default to 1.0 (keeping all data)
+        **kwargs: Any,
+    ):
+        """Initializes the PromptEval data loader.
+
+        Args:
+        ----
+            tasks: benchmark tasks to evaluate on. Default: None (use all tasks).
+            models: models that are evaluated on for each tasks. Default: None (use all models).
+            seed: Random seed for reproducibility. Default: None
+            propensity: Proportion of data to keep. Default: 1.0
+            kwargs: Additional keyword arguments.
+
+        """
+        super().__init__(
+            **kwargs,
+        )
+        self.tasks = tasks
+        self.models = models
+        self.propensity = propensity
+        if seed is not None:
+            np.random.seed(
+                seed=seed
+            )  # instantiate random seed if provided but do it only once here
+
+    def process_data_scalar(self) -> tuple[np.ndarray, np.ndarray]:
+        """Processes the data into scalar setting. This implementation is applicable when generating (template * example) matrix, while fixing model and task.
+
+        Returns
+        -------
+            data: 2d processed data matrix of floats (in this case, each entry is boolean as the metric is correctness)
+            mask: Mask for processed data
+
+        """
+        if not self.tasks or not self.models:
+            raise ValueError("Tasks and models must be specified")
+
+        model = self.models[0]  # Use the first model
+        task = self.tasks[0]  # Use the first task
+        propensity = self.propensity
+
+        assert len(self.models) == 1, "Only one model is supported in scalar mode"
+        assert len(self.tasks) == 1, "Only one task is supported in scalar mode"
+
+        full_data = self._load_data()
+
+        df = pd.DataFrame(full_data[0][model])
+        df_subject = df[df["subject"] == task]
+        num_examples = len(
+            df_subject["correctness"]
+        )  # decide the column dimension of the data matrix
+
+        temp_example = np.zeros(
+            (len(full_data), num_examples)
+        )  # num_templates * num_examples
+
+        for j in range(len(full_data)):
+            # j : per prompt template
+            df = pd.DataFrame(full_data[j][model])
+            print(f"Loading {model} for {j}th prompt template with {task} subject")
+            df_subject = df[df["subject"] == task]
+            temp_example[j, :] = df_subject["correctness"]
+
+        temp_example_df = pd.DataFrame(temp_example)
+
+        # Create a mask of the original data according to the propensity
+        original_mask = (
+            temp_example_df.notna()
+        )  # This just tells us what's naturally present
+
+        n_rows = temp_example_df.shape[0]
+        n_cols = temp_example_df.shape[1]
+        n_rows_keep = int(n_rows * propensity)
+        n_cols_keep = int(n_cols * propensity)
+
+        # Randomly select which rows/columns to keep
+        rows_keep_indices = np.random.choice(n_rows, n_rows_keep, replace=False)
+        cols_keep_indices = np.random.choice(n_cols, n_cols_keep, replace=False)
+
+        # Create a new mask starting with all False
+        propensity_mask = pd.DataFrame(
+            False, index=original_mask.index, columns=original_mask.columns
+        )
+
+        # Set the randomly selected rows/columns to True
+        for i in rows_keep_indices:
+            for j in cols_keep_indices:
+                propensity_mask.iloc[i, j] = True
+
+        data = temp_example_df.to_numpy()
+        mask = propensity_mask.to_numpy()
+        self.data = data
+        self.mask = mask
+        return data, mask
+
+    def process_data_distribution(self) -> tuple[np.ndarray, np.ndarray]:
+        """Process the data into distributional setting.
+
+        Returns
+        -------
+            data: task * model * template matrix of floats (in this case, each entry average of correctness across examples)
+            mask: Mask for processed data
+
+        """
+        if not self.tasks or not self.models:
+            raise ValueError("Tasks and models must be specified")
+
+        models = self.models
+        tasks = self.tasks
+        propensity = self.propensity
+
+        full_data = self._load_data()
+
+        task_model_temp = np.zeros((len(tasks), len(models), len(full_data)))
+
+        for j in range(len(full_data)):
+            # j : per prompt template
+            for k, model in enumerate(models):
+                # model : per model
+                df = pd.DataFrame(full_data[j][model])
+                for l, subject in enumerate(tasks):
+                    print(
+                        f"Loading {model} for {j}th prompt template with {subject} subject"
+                    )
+                    df_subject = df[df["subject"] == subject]
+                    task_model_temp[l, k, j] = np.mean(df_subject["correctness"])
+
+        task_model_temp_df = pd.DataFrame(task_model_temp)
+
+        # Create a mask of the original data according to the propensity
+        original_mask = (
+            task_model_temp_df.notna()
+        )  # This just tells us what's naturally present
+
+        n_rows = task_model_temp_df.shape[0]
+        n_cols = task_model_temp_df.shape[1]
+        n_rows_keep = int(n_rows * propensity)
+        n_cols_keep = int(n_cols * propensity)
+
+        # Randomly select which rows/columns to keep
+        rows_keep_indices = np.random.choice(n_rows, n_rows_keep, replace=False)
+        cols_keep_indices = np.random.choice(n_cols, n_cols_keep, replace=False)
+
+        # Create a new mask starting with all False
+        propensity_mask = pd.DataFrame(
+            False, index=original_mask.index, columns=original_mask.columns
+        )
+
+        # Set the randomly selected rows/columns to True
+        for i in rows_keep_indices:
+            for j in cols_keep_indices:
+                propensity_mask.iloc[i, j] = True
+
+        data = task_model_temp_df.to_numpy()
+        mask = propensity_mask.to_numpy()
+        self.data = data
+        self.mask = mask
+        return data, mask
+
+    def get_full_state_as_dict(self, include_metadata: bool = False) -> dict:
+        """Returns the full state as a dictionary. For HeartSteps, this includes the data, masking matrix, and the custom parameters (if include_metadata == True
+
+        If the data and mask are None, then the data has not been processed yet. Call process_data_scalar() or process_data_distribution() to process the data first.
+
+        Args:
+            include_metadata (bool): Whether to include metadata in the dictionary. Default: False. The metadata for HeartSteps is currently empty.
+
+        """
+        full_state = {
+            "data": self.data,
+            "mask": self.mask,
+        }
+        return full_state
+
+    @classmethod
+    @memory.cache
+    def _load_data(cls) -> Any:
+        """Load the MMLU full dataset.
+
+        Returns:
+            full_data: Any Python object stored in the pickle file
+
+        """
+        logger.info("Retrieving MMLU full dataset from url...")
+        full_data_path = cls.urls["full_data"]
+
+        # GitHub URLs can't be directly downloaded; you would need to use raw content
+        # or download from releases. Here we're assuming the file is downloaded locally.
+        try:
+            # Using standard pickle module instead of pandas
+            with open(full_data_path, "rb") as f:
+                full_data = pickle.load(f)
+        except Exception as e:
+            logger.error(f"Error loading pickle file: {e}")
+            raise ValueError(
+                f"Could not load data from {full_data_path}. Make sure the file exists."
+            )
+
+        return full_data