python - How to write a for loop for multiple machine learning metrics if one metric requires a parameter specification?

Question

I created a for loop for generating metrics on my train and test set. However, in order to calculated Root Mean Square Error (RMSE), I need to either 1) take the sqrt of Mean Square Error or 2) set the parameter mean_squared_error(squared = False). However, I only want a parameter for the RMSE, not for the MAE or the R2.

If I try the below I, understandably, get an error TypeError: mean_squared_error() missing 2 required positional arguments: 'y_true' and 'y_pred' because the parentheses should only come in the for loop.

from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error

#Metrics on train and test
metrics = {
    'RMSE' : mean_squared_error,
    'MAE' : mean_absolute_error,
    'R2' : r2_score
}
    #Train and Test
for key in metrics:
    i = metrics[key]
    train_score = i(y_train, train_predictions)
    test_score = i(y_test, y_pred)
    print(f'Train set {key}: {train_score:.4f}')
    print(f'Test set {key}: {test_score:.4f}')

Answer 1

It works fine for me both in Python 2.7.16 and python 3.7.4 and sklearn version 0.20.3 (does not have the "squared" argument but for that issue it can be solved in a various number of ways in the for loop, the most straightforward is a condition on the key:

for key in metrics:
   i = metrics[key]
   if key == "RMSE":
      train_score = i(y_train, train_predictions, squared=False)
      test_score = i(y_test, y_pred, squared=False)
   else:
      train_score = i(y_train, train_predictions)
      test_score = i(y_test, y_pred)

You can try with those versions, i do not see why your code should not work.