How to define Custom Loss Function?

Question

Hello, I am trying to make a custom Loss Function through the Keras Functional API and I am having problems with this implementation since I have never done anything like.

I want to make a Loss Function that in addition to receiving the data of y_true and y_pred, it also receives a third variable (y_teorico) that would be calculated through the input data itself. Next, how I’d like to do Oss:

Loss = MSE (y_true, y_pred) + MSE (y_pred, y_teorico)

Following my dataset, where the first 14 lines are the network input and the output is the last column’M', which is a direct measurement.

This third variable(y_teorico) I want to enter in Loss Function is calculated using a formula with the input data:

y_teorico = (np.sqrt(
                     2*(CONSTANTE**2) + 
                     2*((E1*E2)) -
                     2*((px1*px2) + (py1*py2) + (pz1*pz2))
                     ))

I do not know how to do this inside my Oss, I was reading that could be done through a custom Train step as in the following link, but even so I’m having difficulties to do so.

https://keras.io/guides/customizing_what_happens_in_fit/#a-first-simple-example

Going on like I’m trying to do:

Defining training and testing:

x_train = data.drop(columns=['M'])
X_train = x.apply(lambda x: (x - np.mean(x)) / (np.max(x) - np.min(x)))

y_train = data.M

Creating the model

class CustomModel(keras.Model):
def train_step(self, data):
    X_train, y_train = data
   
    with tf.GradientTape() as tape:
        y_pred = self(x, training=True)

        elec_mass = 0.5109989461
        y_teoric = (np.sqrt(
                      2*(elec_mass**2) + 
                      2*((self.E1*self.E2)) -
                      2*((self.px1*self.px2) + (self.py1*self.py2) + (self.pz1*self.pz2))
                      ))
                        
        # Compute the loss value
        # (the loss function is configured in `compile()`)

        GAMMA = 1
        L_data = self.compiled_loss(y, y_pred, y_teoric, regularization_losses=self.losses)
        L_phy = GAMMA * self.compiled_loss(y_pred, y_teoric)
        custom_loss = L_data + L_phy
        loss = custom_loss

    # Compute gradients
    trainable_vars = self.trainable_variables
    gradients = tape.gradient(loss, trainable_vars)
    
    # Update weights
    self.optimizer.apply_gradients(zip(gradients, trainable_vars))
    
    # Update metrics (includes the metric that tracks the loss)
    self.compiled_metrics.update_state(y, y_pred)
    
    # Return a dict mapping metric names to current value
    return {m.name: m.result() for m in self.metrics}

#Layers

#Model Layers
input = keras.Input(shape=(14,))
hidden1 = keras.layers.Dense(14, activation='relu')(input)
hidden2 = keras.layers.Dense(7, activation='relu')(hidden1)
output = keras.layers.Dense(1, activation= relu_advanced)(hidden2)

model = CustomModel(inputs=input, outputs=output)

#Finally the FIT()

model.compile(optimizer="adam", loss="mse", metrics=["mse"])
early_stop = keras.callbacks.EarlyStopping(monitor='val_loss', patience=5)

history = model.fit(X_train, y_train, 
                    epochs           = 100,
                    batch_size       = 100,
                    verbose          = 2,
                    validation_split = 0.2,
                    callbacks=[early_stop]
                    )