Mnar and directopt

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 = {
@@ -186,8 +187,90 @@ 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.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
 
     def get_full_state_as_dict(self, include_metadata: bool = False) -> dict:
         """Returns the full state of this object as a dictionary"""

src/nearest_neighbors/fit_method.py

@@ -0,0 +1,149 @@
+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:
+            data_array (npt.NDArray): Data array
+            mask_array (npt.NDArray): Mask array
+            imputer (NearestNeighborImputer): Imputer
+
+        Returns:
+            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
+
+        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 (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
+            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)
+
+        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
+
+
+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.
+
+        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

src/nearest_neighbors/nnimputer.py

@@ -128,6 +128,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

src/nearest_neighbors/simulations/mcar.py

@@ -139,3 +139,54 @@ 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], 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))
+            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