scripts/builtin/imputeByKNN.dml

#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#-------------------------------------------------------------

# Imputes missing values, indicated by NaNs, using KNN-based methods
# (k-nearest neighbors by euclidean distance). In order to avoid NaNs in
# distance computation and meaningful nearest neighbor search, we initialize
# the missing values by column means. Currently, only the column with the most
# missing values is actually imputed.
#
# INPUT:
# ------------------------------------------------------------------------------
# X             Matrix with missing values, which are represented as NaNs
# method        Method used for imputing missing values with different performance and accuracy tradeoffs:\n
#               - 'dist' (default): Compute all-pairs distances and impute the missing values by closest. O(N^2 * #features)
#               - 'dist_missing': Compute distances between data and records with missing values. O(N*M * #features), assuming that the number of records with MV is M<<N.
#               - 'dist_sample': Compute distances between sample of data and records with missing values. O(S*M * #features) with M<<N and S<<N, but suboptimal imputation.
#
# seed          Root seed value for random/sample calls for deterministic behavior. -1 for true randomization
# sample_frac   Sample fraction for 'dist_sample' (value between 0 and 1)
# ------------------------------------------------------------------------------
#
# OUTPUT:
# ------------------------------------------------------------------------------
# result     Imputed dataset
# ------------------------------------------------------------------------------

m_imputeByKNN = function(Matrix[Double] X, String method="dist", Int seed=-1, Double sample_frac=0.1)
  return(Matrix[Double] result)
{
  #KNN-Imputation Script

  imputedValue = X

  #Impute NaN value with temporary mean value of the column
  filled_matrix = imputeByMean(X, matrix(0, cols = ncol(X), rows = 1))

  if(method == "dist") {
    #Calculate the distance using dist method after imputation with mean
    distance_matrix = dist(filled_matrix)

    #Change 0 value so rowIndexMin will ignore that diagonal value
    distance_matrix = replace(target=distance_matrix, pattern=0, replacement=999)

    #Get the minimum distance row-wise computation
    minimum_index = rowIndexMin(distance_matrix)

    #Create aligned matrix from minimum index
    aligned = table(minimum_index, seq(1, nrow(X)), odim1=nrow(X), odim2=nrow(X))

    #Get the X records that need to be imputed
    imputedValue = t(filled_matrix) %*% aligned
    imputedValue = t(imputedValue)
  }
  else if(method == "dist_missing") {
    #assuming small missing values
    imputedValue = compute_missing_values(X, filled_matrix, seed, 1.0)
  }
  else if(method == "dist_sample"){
    #assuming large missing values
    imputedValue = compute_missing_values(X, filled_matrix, seed, sample_frac)
  }
  else {
    stop("Method is unknown or not yet implemented")
  }

  #Impute the value
  result = replace(target=X, pattern=NaN, replacement=0)
  result = result + (imputedValue * is.nan(X))
}

compute_missing_values = function (Matrix[Double] X, Matrix[Double] filled_matrix, Int seed, Double sample_frac)
    return (Matrix[Double] imputedValue)
{
    #Split the matrix into containing NaN values (missing records) and not containing NaN values (M2 records)
    maskNAN = is.nan(X)
    I = rowSums(maskNAN) != 0
    missing = removeEmpty(target=filled_matrix, margin="rows", select=I)

    Y = (rowSums(maskNAN)==0)
    M2 = removeEmpty(target=X, margin = "rows", select = Y)

    if (sample_frac != 1.0) {
        #Create permutation matrix for sampling sample_frac*nrow(X) rows
        I = rand(rows=nrow(M2), cols=1, seed=seed) <= sample_frac;
        M2 = removeEmpty(target=M2, margin="rows", select=I);
    }

    #Calculate the euclidean distance between fully records and missing records, and then find the min value row wise
    dotM2 = rowSums(M2 * M2) %*% matrix(1, rows = 1, cols = nrow(missing))
    dotMissing = t(rowSums(missing * missing) %*% matrix(1, rows = 1, cols = nrow(M2)))
    D = sqrt(dotM2 + dotMissing - 2 * (M2 %*% t(missing)))
    minD = rowIndexMin(t(D))

    #Get the index location of the missing value
    pos = rowMaxs(maskNAN)
    missing_indices = seq(1, nrow(pos)) * pos

    #Put the replacement value in the missing indices
    I2 = removeEmpty(target=missing_indices, margin="rows")
    R = table(I2, 1, minD, odim1=nrow(X), odim2=1)

    #Replace the 0 to avoid error in table()
    R = replace(target=R, pattern=0, replacement=nrow(X)+1)

    #Create aligned matrix from minimum index
    aligned = table(R, seq(1, nrow(X)), odim1=nrow(X), odim2=nrow(X))

    #Reshape the subset
    reshaped = rbind(M2, matrix(0, rows=nrow(X) - nrow(M2), cols=ncol(X)))

    #Get the subset records that need to be imputed
    imputedValue = t(reshaped) %*% aligned
    imputedValue = t(imputedValue)
}