nodes: 0, module; 1, function_definition; 2, function_name:knn_impute_few_observed; 3, parameters; 4, block; 5, identifier:X; 6, identifier:missing_mask; 7, identifier:k; 8, default_parameter; 9, default_parameter; 10, expression_statement; 11, expression_statement; 12, expression_statement; 13, comment:# put the missing mask in column major order since it's accessed; 14, comment:# one column at a time; 15, expression_statement; 16, expression_statement; 17, expression_statement; 18, expression_statement; 19, comment:# get rid of infinities, replace them with a very large number; 20, expression_statement; 21, expression_statement; 22, expression_statement; 23, expression_statement; 24, comment:# trim the number of other rows we consider to exclude those; 25, comment:# with infinite distances; 26, expression_statement; 27, expression_statement; 28, for_statement; 29, return_statement; 30, identifier:verbose; 31, False; 32, identifier:print_interval; 33, integer:100; 34, comment:"""
      Seems to be the fastest kNN implementation. Pre-sorts each rows neighbors
      and then filters these sorted indices using each columns mask of
      observed values.
  
      Important detail: If k observed values are not available then uses fewer
      than k neighboring rows.
  
      Parameters
      ----------
      X : np.ndarray
          Matrix to fill of shape (n_samples, n_features)
  
      missing_mask : np.ndarray
          Boolean array of same shape as X
  
      k : int
  
      verbose : bool
      """; 35, assignment; 36, assignment; 37, assignment; 38, assignment; 39, assignment; 40, assignment; 41, assignment; 42, assignment; 43, assignment; 44, assignment; 45, assignment; 46, assignment; 47, identifier:i; 48, call; 49, block; 50, identifier:X_row_major; 51, identifier:start_t; 52, call; 53, pattern_list; 54, attribute; 55, identifier:missing_mask_column_major; 56, call; 57, identifier:observed_mask_column_major; 58, unary_operator; 59, identifier:X_column_major; 60, call; 61, pattern_list; 62, line_continuation:\; 63, call; 64, identifier:D_sorted; 65, call; 66, identifier:inv_D; 67, binary_operator:1.0 / D; 68, identifier:D_valid_mask; 69, comparison_operator:D < effective_infinity; 70, identifier:valid_distances_per_row; 71, call; 72, identifier:D_sorted; 73, list_comprehension; 74, identifier:dot; 75, attribute; 76, identifier:range; 77, argument_list; 78, expression_statement; 79, expression_statement; 80, expression_statement; 81, if_statement; 82, expression_statement; 83, for_statement; 84, attribute; 85, argument_list; 86, identifier:n_rows; 87, identifier:n_cols; 88, identifier:X; 89, identifier:shape; 90, attribute; 91, argument_list; 92, identifier:missing_mask_column_major; 93, attribute; 94, argument_list; 95, identifier:X_row_major; 96, identifier:D; 97, identifier:effective_infinity; 98, identifier:knn_initialize; 99, argument_list; 100, attribute; 101, argument_list; 102, float:1.0; 103, identifier:D; 104, identifier:D; 105, identifier:effective_infinity; 106, attribute; 107, argument_list; 108, subscript; 109, for_in_clause; 110, identifier:np; 111, identifier:dot; 112, identifier:n_rows; 113, assignment; 114, assignment; 115, assignment; 116, boolean_operator; 117, block; 118, assignment; 119, identifier:j; 120, identifier:missing_indices; 121, block; 122, identifier:time; 123, identifier:time; 124, identifier:np; 125, identifier:asarray; 126, identifier:missing_mask; 127, keyword_argument; 128, identifier:X; 129, identifier:copy; 130, keyword_argument; 131, identifier:X; 132, identifier:missing_mask; 133, keyword_argument; 134, identifier:np; 135, identifier:argsort; 136, identifier:D; 137, keyword_argument; 138, identifier:D_valid_mask; 139, identifier:sum; 140, keyword_argument; 141, identifier:D_sorted; 142, identifier:i; 143, slice; 144, pattern_list; 145, call; 146, identifier:missing_row; 147, subscript; 148, identifier:missing_indices; 149, subscript; 150, identifier:row_weights; 151, subscript; 152, identifier:verbose; 153, comparison_operator:i % print_interval == 0; 154, expression_statement; 155, identifier:candidate_neighbor_indices; 156, subscript; 157, expression_statement; 158, expression_statement; 159, expression_statement; 160, expression_statement; 161, expression_statement; 162, expression_statement; 163, if_statement; 164, identifier:order; 165, string:"F"; 166, identifier:order; 167, string:"F"; 168, identifier:verbose; 169, identifier:verbose; 170, identifier:axis; 171, integer:1; 172, identifier:axis; 173, integer:1; 174, identifier:count; 175, identifier:i; 176, identifier:count; 177, identifier:enumerate; 178, argument_list; 179, identifier:missing_mask; 180, identifier:i; 181, slice; 182, call; 183, integer:0; 184, identifier:inv_D; 185, identifier:i; 186, slice; 187, binary_operator:i % print_interval; 188, integer:0; 189, call; 190, identifier:D_sorted; 191, identifier:i; 192, assignment; 193, assignment; 194, assignment; 195, assignment; 196, assignment; 197, assignment; 198, comparison_operator:weight_sum > 0; 199, block; 200, identifier:valid_distances_per_row; 201, attribute; 202, argument_list; 203, identifier:i; 204, identifier:print_interval; 205, identifier:print; 206, argument_list; 207, identifier:observed; 208, subscript; 209, identifier:sorted_observed; 210, subscript; 211, identifier:observed_neighbor_indices; 212, subscript; 213, identifier:k_nearest_indices; 214, subscript; 215, identifier:weights; 216, subscript; 217, identifier:weight_sum; 218, call; 219, identifier:weight_sum; 220, integer:0; 221, expression_statement; 222, expression_statement; 223, expression_statement; 224, identifier:np; 225, identifier:where; 226, identifier:missing_row; 227, binary_operator:"Imputing row %d/%d with %d missing, elapsed time: %0.3f" % (
                      i + 1,
                      n_rows,
                      len(missing_indices),
                      time.time() - start_t); 228, identifier:observed_mask_column_major; 229, slice; 230, identifier:j; 231, identifier:observed; 232, identifier:candidate_neighbor_indices; 233, identifier:candidate_neighbor_indices; 234, identifier:sorted_observed; 235, identifier:observed_neighbor_indices; 236, slice; 237, identifier:row_weights; 238, identifier:k_nearest_indices; 239, attribute; 240, argument_list; 241, assignment; 242, assignment; 243, assignment; 244, string:"Imputing row %d/%d with %d missing, elapsed time: %0.3f"; 245, tuple; 246, identifier:k; 247, identifier:weights; 248, identifier:sum; 249, identifier:column; 250, subscript; 251, identifier:values; 252, subscript; 253, subscript; 254, binary_operator:dot(values, weights) / weight_sum; 255, binary_operator:i + 1; 256, identifier:n_rows; 257, call; 258, binary_operator:time.time() - start_t; 259, identifier:X_column_major; 260, slice; 261, identifier:j; 262, identifier:column; 263, identifier:k_nearest_indices; 264, identifier:X_row_major; 265, identifier:i; 266, identifier:j; 267, call; 268, identifier:weight_sum; 269, identifier:i; 270, integer:1; 271, identifier:len; 272, argument_list; 273, call; 274, identifier:start_t; 275, identifier:dot; 276, argument_list; 277, identifier:missing_indices; 278, attribute; 279, argument_list; 280, identifier:values; 281, identifier:weights; 282, identifier:time; 283, identifier:time
code: def knn_impute_few_observed(
          X, missing_mask, k, verbose=False, print_interval=100):
      """
      Seems to be the fastest kNN implementation. Pre-sorts each rows neighbors
      and then filters these sorted indices using each columns mask of
      observed values.
  
      Important detail: If k observed values are not available then uses fewer
      than k neighboring rows.
  
      Parameters
      ----------
      X : np.ndarray
          Matrix to fill of shape (n_samples, n_features)
  
      missing_mask : np.ndarray
          Boolean array of same shape as X
  
      k : int
  
      verbose : bool
      """
      start_t = time.time()
      n_rows, n_cols = X.shape
      # put the missing mask in column major order since it's accessed
      # one column at a time
      missing_mask_column_major = np.asarray(missing_mask, order="F")
      observed_mask_column_major = ~missing_mask_column_major
      X_column_major = X.copy(order="F")
      X_row_major, D, effective_infinity = \
          knn_initialize(X, missing_mask, verbose=verbose)
      # get rid of infinities, replace them with a very large number
      D_sorted = np.argsort(D, axis=1)
      inv_D = 1.0 / D
      D_valid_mask = D < effective_infinity
      valid_distances_per_row = D_valid_mask.sum(axis=1)
  
      # trim the number of other rows we consider to exclude those
      # with infinite distances
      D_sorted = [
          D_sorted[i, :count]
          for i, count in enumerate(valid_distances_per_row)
      ]
  
      dot = np.dot
      for i in range(n_rows):
          missing_row = missing_mask[i, :]
          missing_indices = np.where(missing_row)[0]
          row_weights = inv_D[i, :]
          if verbose and i % print_interval == 0:
              print(
                  "Imputing row %d/%d with %d missing, elapsed time: %0.3f" % (
                      i + 1,
                      n_rows,
                      len(missing_indices),
                      time.time() - start_t))
          candidate_neighbor_indices = D_sorted[i]
  
          for j in missing_indices:
              observed = observed_mask_column_major[:, j]
              sorted_observed = observed[candidate_neighbor_indices]
              observed_neighbor_indices = candidate_neighbor_indices[sorted_observed]
              k_nearest_indices = observed_neighbor_indices[:k]
              weights = row_weights[k_nearest_indices]
              weight_sum = weights.sum()
              if weight_sum > 0:
                  column = X_column_major[:, j]
                  values = column[k_nearest_indices]
                  X_row_major[i, j] = dot(values, weights) / weight_sum
      return X_row_major