Configurar el API de Scikit-learn para generar DataFrames

La nueva versión 1.3 de Scikit-learn trae importantes novedades. Una de ellas es la posibilidad de configurar el API de las herramientas de transformación de datos o transformers. Hasta ahora en Scikit-learn al utilizar un transformer el resultado que se obtiene es un ndarray de NumPy. A partir de la versión 1.3, es posible configurar el tipo de objeto resultante mediante el uso del método set_config(). Pudiendo configurar Scikit-learn para generar DataFrames de Pandas en lugar de los ndarray.

Transformación de variables con Scikit-learn

Antes de entrenar un modelo una de las tareas más habituales es la normalización de las características. Una tarea que facilita el proceso de entrenamiento. Algunas de las clases que suelen utilizar habitualmente para ellos son StandardScaler y MinMaxScaler. Aunque estas clases suelen devolver un ndarray como resultado independientemente de que los datos originales se encuentren en un DataFrame.

Esto se puede ver en un ejemplo. Si se importan los datos de ejercicio físico de Linnerud disponibles en Scikit-learn dentro de un DataFrame y se normalizan con StandardScaler el resultado es un ndarray. Como se muestra en el siguiente código.

from sklearn.datasets import load_linnerud
from sklearn.preprocessing import StandardScaler

x_linnerud, y_linnerud = load_linnerud(return_X_y=True, as_frame=True)

print(f'Original (tipo: {x_linnerud.__class__.__name__})')
print(x_linnerud.head())

scaler = StandardScaler()
x_linnerud_scaler = scaler.fit_transform(x_linnerud)

print(f'\nNormalizado (tipo: {x_linnerud_scaler.__class__.__name__})')
print(x_linnerud_scaler[:5,])

Original (tipo: DataFrame)
   Chins  Situps  Jumps
0    5.0   162.0   60.0
1    2.0   110.0   60.0
2   12.0   101.0  101.0
3   12.0   105.0   37.0
4   13.0   155.0   58.0

Normalizado (tipo: ndarray)
[[-0.86367072  0.26975016 -0.20608616]
 [-1.44592064 -0.58295552 -0.20608616]
 [ 0.49491243 -0.7305392   0.61425681]
 [ 0.49491243 -0.66494645 -0.66627856]
 [ 0.68899574  0.15496286 -0.24610289]]

Modificar el API de Scikit-learn para generar DataFrames

A partir de la versión 1.3 de Scikit-learn las clases que transforman datos cuentan con el método set_output() mediante el cual se le pueden indicar el tipo de dato que se desea obtener como resultado: ndarray o DataFrame. Para que el objeto devuelva un DataFrame se tendrá que indicar la opción transform="pandas".

En el siguiente ejemplo se puede ver cómo se puede modificar el código anterior para que el resultado sea un DataFrame. (Importante: el siguiente código solo funcionará a partir de la versión 1.3 de Scikit-learn).

scaler = StandardScaler().set_output(transform="pandas")
x_linnerud_scaler = scaler.fit_transform(x_linnerud)

print(f'\nNormalizado (tipo: {x_linnerud_scaler.__class__.__name__})')
print(x_linnerud_scaler.head())

Normalizado (tipo: DataFrame)
      Chins    Situps     Jumps
0 -0.863671  0.269750 -0.206086
1 -1.445921 -0.582956 -0.206086
2  0.494912 -0.730539  0.614257
3  0.494912 -0.664946 -0.666279
4  0.688996  0.154963 -0.246103

Conclusiones

La posibilidad de cambiar el API de Scikit-learn para generar DataFrames en lugar de ndarray es una buena noticia. A partir de ahora será posible utilizar las herramientas que ofrece Pandas para el tratamiento de datos durante todo el proceso de preparación antes de entrenar un modelo. Algo que muchos usuarios agradecerán en el día a día.