Tensor of categoricals

[tommyod]

I want fit and transform over a Tensor of categoricals to behave like fitting each categorical.

import matplotlib.pyplot as plt
import numpy as np

from generalized_additive_models import GAM, Categorical, Spline, Tensor
from generalized_additive_models.datasets import load_powerlifters

# Load data and filter it
df = load_powerlifters().rename(columns=lambda s:s.removeprefix("best3").removesuffix("kg"))

# Melt dataframe so each exercise (squat, bench, deadlift) ends up in a row
df = df.melt(id_vars=["sex", "age", "bodyweight"], value_vars=["squat", "bench", "deadlift"], value_name="lifted", var_name="exercise")

# Predict total weight lifted, given age, bodyweight and sex
target = df["lifted"]
age = Spline("age", penalty=1e3, num_splines=8)
bodyweight = Spline("bodyweightkg", penalty=1e3, num_splines=8)

sex = Categorical("sex", penalty=0)
sex.fit_transform(df)
array([[1., 0.],
       [1., 0.],
       [0., 1.],
       ...,
       [1., 0.],
       [1., 0.],
       [1., 0.]])

However:

sex_ex_exercise = Tensor([Categorical(feature='sex', penalty=0), Categorical(feature='exercise', penalty=0)])

sex_ex_exercise.fit_transform(df)
Out[149]: 
array([[-1.51653018e-01, -1.51653018e-01,  1.33333333e+00,
        -5.15013649e-01, -5.15013649e-01,  4.18554080e-14],
       [-1.51653018e-01, -1.51653018e-01,  1.33333333e+00,
        -5.15013649e-01, -5.15013649e-01,  4.18554080e-14],
       [-4.84986351e-01, -4.84986351e-01,  2.50910404e-14,
        -1.81680315e-01, -1.81680315e-01,  1.33333333e+00],
       ...,
       [-1.51653018e-01,  1.33333333e+00, -1.51653018e-01,
        -5.15013649e-01, -1.36557432e-14, -5.15013649e-01],
       [-1.51653018e-01,  1.33333333e+00, -1.51653018e-01,
        -5.15013649e-01, -1.36557432e-14, -5.15013649e-01],
       [-1.51653018e-01,  1.33333333e+00, -1.51653018e-01,
        -5.15013649e-01, -1.36557432e-14, -5.15013649e-01]])