adding output_normalizer

[fermisea]

adds the possibility of normalizing the outputs of a symbolic fn as well.

import torch
import numpy
import pandas

n = 10000
w = numpy.random.normal(0, 1, n)
x = 10 * w + numpy.random.normal(0, 1, n)
z = 15 * w + numpy.random.normal(0, 1, n)
y = 4 * x + 4 * z + numpy.random.normal(0, 1, n)

data = pandas.DataFrame({"x": x, "y": y, "z": z, "w": w})

means = data.mean()
stds = data.std()

node = "y"
parents = ["x", "z"]

scaled_data = (data - means) / stds
X, y = scaled_data[parents], scaled_data[[node]]
X_train, X_test, y_train, y_test = X[:8000], X[8000:], y[:8000], y[8000:]

dataset = {}
dataset["train_input"] = torch.from_numpy(X_train.values).float()
dataset["test_input"] = torch.from_numpy(X_test.values).float()
dataset["train_label"] = torch.from_numpy(y_train.values).float()
dataset["test_label"] = torch.from_numpy(y_test.values).float()

model = KAN(width=[2, 1], k=3, grid=5)
model.train(dataset, steps=20)
model.auto_symbolic(lib=["x"])

node = "y"
parents = ["x", "z"]


means_list = [means[p] for p in parents]
stds_list = [stds[p] for p in parents]

output_mean = means[node]
output_std = stds[node]

model.symbolic_formula(
    normalizer=[means_list, stds_list], output_normalizer=[[output_mean], [output_std]]
)

with output normalizer: ([3.98*x_1 + 4.02*x_2 - 0.02], [x_1, x_2]) without: ([0.04*x_1 + 0.04*x_2 + 0.01], [x_1, x_2])