scripts/builtin/imputeByMean.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.
#
#-------------------------------------------------------------

# impute the data by mean value and if the feature is categorical then by mode value
# Related to [SYSTEMDS-2662] dependency function for cleaning pipelines
#
# INPUT:
# -------------------------------------------------------------------------------------
# X       Data Matrix (Recoded Matrix for categorical features)
# mask    A 0/1 row vector for identifying numeric (0) and categorical features (1)
# -------------------------------------------------------------------------------------
#
# OUTPUT:
# -----------------------------------------------------------------------------------
# X     imputed dataset
# -----------------------------------------------------------------------------------

m_imputeByMean = function(Matrix[Double] X, Matrix[Double] mask)
return(Matrix[Double] X, Matrix[Double] imputedVec)
{

  #  mean imputation
  colMean = matrix(0, rows=1, cols=ncol(X))
  parfor(i in 1:ncol(X))
  {
    if(as.scalar(mask[1, i]) == 0)
    {
      nX = removeEmpty(target=X[, i], margin="rows", select = (is.na(X[, i]) == 0))
      colMean[1, i] = mean(nX)
    }
  }

  if(sum(mask) > 0)
  {
    # mode imputation
    cX = X*mask
    [X_c, colMode] = imputeByMode(cX)
    imputedVec = colMean + colMode
  }
  else 
  {
    imputedVec = colMean
  }
  X = imputeByMeanApply(X, imputedVec)
}