I want to design a model that takes an array of data as input and predicts another array of data as output. One metric that I want to use is the pearsonr. In other words, I want to minimize the difference between pearsonr(x_train, y_train) and pearsonr(x_train, y_pred)

Instead of using the basic mean squared error, I want to penalize the loss function more heavily if this pearsonr difference is high. Therefore, I wanted to design a custom loss function where I can get access to each training sample x_train[i], so I can calculate the pearsonr and convert it to a weight factor that can be added to the current loss.

This is what I have for now:

def custom_loss_wrapper(x_train):
    def custom_loss(y_true, y_pred):
        mse_loss = tf.keras.losses.MSE(y_true, y_pred)
        # calculate true pearsonr
        target_pearsonr = pearsonr(x_train, y_true)
        # calculate predicted pearsonr
        predicted_pearsonr = pearsonr(x_train, y_pred)
        pearsonr_diff = tf.abs(target_pearsonr - predicted_pearsonr)
        penalty = tf.multiply(pearsonr_diff, 2.0) 
        weighted_loss = mse_loss + penalty
        return weighted_loss
    return custom_loss

In this code, I wanted to use a wrapper so I can pass in the X_train. However, I think this x_train is the the whole tensor in its entirety, and the y_true and y_pred are the individual ground truth and prediction for each sample of X_train.

How do I change the code so I can calculate a pearsonr for each sample and then use it in the custom loss function?