official-code/src/utils/msnn.py

"""
Mixed Synthetic Nearest Neighbors algorithm
"""
import sys
import warnings 
import random
import numpy as np 
import networkx as nx
from networkx.algorithms.clique import find_cliques
from sklearn.utils import check_array
from .snn import SyntheticNearestNeighbors

class MixedSyntheticNearestNeighbors(SyntheticNearestNeighbors):
    def __init__(self, *args, treatment_mask=None, treatment_level_scale_dict=None, missing_treatment=None, **kwargs): 
        super().__init__(*args, **kwargs)
        self.treatment_mask = treatment_mask
        self.treatment_level_scale_dict = treatment_level_scale_dict
        self.missing_treatment = missing_treatment

    def _find_anchors(self, X, missing_pair): 
        (missing_row, missing_col) = missing_pair
        """
        now for the mixed snn: 
        obs_rows is where treatment_mask[:, missing_col] equals to missing_treatment, 
        and obs_col is where treatment_mask[missing_row, :] is not zero. 
        """
        if self.treatment_mask is not None:
            obs_rows = np.argwhere(self.treatment_mask[:, missing_col]==self.missing_treatment).flatten()
            obs_cols = np.argwhere(self.treatment_mask[missing_row, :]!="not_observed").flatten()
        else: 
            # Warning: no treatment_mask provided, using default SNN
            warnings.warn(
				"no treatment_mask provided, using default SNN"
			)
            obs_rows = np.argwhere(~np.isnan(X[:, missing_col])).flatten()
            obs_cols = np.argwhere(~np.isnan(X[missing_row, :])).flatten()
        # print("obs_rows, obs_cols, missing_row, missing_col: ", obs_rows, obs_cols, missing_row, missing_col)

        # make sure (i,j) not in (obs_rows, obs_cols)
        obs_rows = np.setdiff1d(obs_rows, missing_row)
        obs_cols = np.setdiff1d(obs_cols, missing_col)
        # print("obs_rows, obs_cols: ", obs_rows, obs_cols)

        # create bipartite incidence matrix 
        if self.treatment_mask is not None:
            # treatment of missing_row for each column (column-wise labels)
            col_treatments = self.treatment_mask[missing_row, obs_cols]  # shape (|obs_cols|,)
            # treatment for candidate rows * cols
            B = (self.treatment_mask[np.ix_(obs_rows, obs_cols)] == col_treatments)
            # print("B: ", B)
            # print("B_shape: ", B.shape)
            # print("B_nonzero: ", np.count_nonzero(B))
            # print("col_treatments: ", col_treatments)
            
            # baseline: 
            # obs_cols = np.argwhere(self.treatment_mask[missing_row, :]==self.missing_treatment).flatten()
            # obs_cols = np.setdiff1d(obs_cols, missing_col)
            # B = (self.treatment_mask[np.ix_(obs_rows, obs_cols)] == self.missing_treatment)
        else: 
            B = X[obs_rows]
            B = B[:, obs_cols]
            if not np.any(np.isnan(B)): # check if fully connected already
                return (obs_rows, obs_cols)
            B[np.isnan(B)] = 0 

		# bipartite graph 
        (n_rows, n_cols) = B.shape 
        A = np.block([[np.ones((n_rows, n_rows)), B],
		  			  [B.T, np.ones((n_cols, n_cols))]])
        G = nx.from_numpy_array(A)

		# find max clique that yields the most self.n_neighbors : 1 matrix
        cliques = list(find_cliques(G))
        d_min = 0 
        max_clique_rows_idx = False
        max_clique_cols_idx = False
        for clique in cliques: 
            clique = np.sort(clique)
            clique_rows_idx = clique[clique<n_rows]
            clique_cols_idx = clique[clique>=n_rows] - n_rows
            d = min(len(clique_rows_idx // self.n_neighbors), len(clique_cols_idx))
            if d>d_min: 
                d_min = d 
                max_clique_rows_idx = clique_rows_idx
                max_clique_cols_idx = clique_cols_idx

        # determine model learning rows & cols 
        anchor_rows = obs_rows[max_clique_rows_idx]
        anchor_cols = obs_cols[max_clique_cols_idx]
        # print("anchor_rows, anchor_cols, shape: ", anchor_rows, anchor_cols, anchor_rows.shape, anchor_cols.shape)
        return (anchor_rows, anchor_cols)

    def _synth_neighbor(self, X, missing_pair, anchor_rows, anchor_cols, covariates=None): 
        """
		construct the k-th mixed synthetic neighbor 
		"""
        # normalize & initialize
        (missing_row, missing_col) = missing_pair
        target_scale = 1
        if self.treatment_mask is not None and self.treatment_level_scale_dict is not None: 
            # scaling matrix: each element of treatment_mask, 
            # if not zero, replacing by corresponding value of treatment_level; 
            # if zero, replacing by 0
            # X: entry-wise product of X and 1/scaling_matrix
            vectorized_lookup = np.vectorize(self.treatment_level_scale_dict.get, otypes=[np.float64])
            scaling_matrix = vectorized_lookup(self.treatment_mask)
            X = X / scaling_matrix
            target_scale = vectorized_lookup([self.missing_treatment])[0]
        # print("X: ", X)
        y1 = X[missing_row, anchor_cols]
        X1 = X[anchor_rows, :] 
        X1 = X1[:, anchor_cols]
        X2 = X[anchor_rows, missing_col]

        # add covariates
        if covariates is not None: 
            y1_covariates = np.hstack([y1, covariates[missing_row]])
            X1_covariates = np.hstack([X1, covariates[anchor_rows]])
        else:
            y1_covariates = y1.copy()
            X1_covariates = X1.copy()

        # print("X1: ", X1_covariates)

        # learn k-th synthetic neighbor
        (beta, _, s_rank, v_rank) = self._pcr(X1_covariates.T, y1_covariates)

        # prediction
        pred = target_scale * self._clip(X2@beta) 

        # diagnostics 
        train_error = self._train_error(X1.T, y1, beta)
        subspace_inclusion_stat = self._subspace_inclusion(v_rank, X2) 
        feasible = self._isfeasible(train_error, subspace_inclusion_stat)

        # assign weight of k-th synthetic neighbor
        if self.weights=='uniform':
            weight = 1 
        elif self.weights=='distance':
            d = train_error + subspace_inclusion_stat
            weight = 1/d if d>0 else sys.float_info.max
        return (pred, feasible, weight)
    
    def fit_transform(self, X, covariates=None, test_set=None): 
        """
        complete missing entries in matrix 
        """ 
        # get missing entries to impute
        missing_set = test_set if test_set is not None else np.argwhere(self.treatment_mask != self.missing_treatment)
        # print("missing set: ", missing_set)
        # print the size of missing set
        # print("missing set size: ", missing_set.shape[0])
        num_missing = len(missing_set)

        # check and prepare data 
        X = self._prepare_input_data(X, missing_set)

        # check weights
        self.weights = self._check_weights(self.weights)

        # initialize 
        X_imputed = X.copy() 
        std_matrix = np.zeros(X.shape)
        self.feasible = np.empty(X.shape)
        self.feasible.fill(np.nan)

        # complete missing entries 
        for (i, missing_pair) in enumerate(missing_set): 
            if self.verbose: 
                print("[MSNN] iteration {} of {}".format(i+1, num_missing))

            # predict missing entry
            (pred, feasible) = self._predict(X, 
                                                missing_pair=missing_pair,
                                                covariates=covariates)

            # store in imputed matrices
            (missing_row, missing_col) = missing_pair
            X_imputed[missing_row, missing_col] = pred 
            self.feasible[missing_row, missing_col] = feasible
            print("i: ", i, "pred: ", pred, "feasible: ", feasible)

        if self.verbose:
            print("[MSNN] complete")
        return X_imputed