MLP with data and physics loss

Question

I have the following problem: Training an MLP on 4 inputs while also having an estimation from a physical model (with some error). I now want to compute a combined loss from the physics and data loss and use this for training the MLP (1 output as a regression value). I'm using tensorflow and I'm having some errors in seperating the targets and providing them to the model. I appreciate any suggestions. Thank you! This is my model and code so far:

import tensorflow as tf
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras import regularizers
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import SGD
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import joblib
import tensorflow.keras.backend as K

# Load data
train_csv_path = "/home/qiki1832/PycharmProjects/datasets_TABULA/GroundTruth_train20250219.csv"
val_csv_path = "/home/qiki1832/PycharmProjects/datasets_TABULA/GroundTruth_val20250218.csv"

# Read and clean data
df_train = pd.read_csv(train_csv_path, delimiter=",").dropna()
df_val = pd.read_csv(val_csv_path, delimiter=",").dropna()

# Features and targets
features = ['T_se', 'T_si', 'q_si', 'd']
targets = ['R']  # Only R_true is used for actual training
pred_phys = ['R_phys,i']  # R_phys,i is used for physics loss

X_train = df_train[features].values
y_train = df_train[targets].values
X_val = df_val[features].values
y_val = df_val[targets].values

# Normalize inputs
scaler_X = StandardScaler()
X_train_scaled = scaler_X.fit_transform(X_train)
X_val_scaled = scaler_X.transform(X_val)

# Normalize targets
scaler_y = StandardScaler()
y_train_scaled = scaler_y.fit_transform(y_train)
y_val_scaled = scaler_y.transform(y_val)

# Normalize the physics targets (R_phys) using the same scaler
scaler_phys = StandardScaler()
y_train_scaled_R_phys = scaler_phys.fit_transform(df_train[pred_phys].values)
y_val_scaled_R_phys = scaler_phys.transform(df_val[pred_phys].values)

# Save scalers for later use
joblib.dump(scaler_X, "/home/qiki1832/PycharmProjects/PINN/PhysicsRegularisation/MLP_01/transient_scaler_X.pkl")
joblib.dump(scaler_y, "/home/qiki1832/PycharmProjects/PINN/PhysicsRegularisation/MLP_01/transient_scaler_y.pkl")
joblib.dump(scaler_phys, "/home/qiki1832/PycharmProjects/PINN/PhysicsRegularisation/MLP_01/transient_scaler_phys.pkl")

# Convert to tensors
X_train_tensor = tf.convert_to_tensor(X_train_scaled, dtype=tf.float32)
X_val_tensor = tf.convert_to_tensor(X_val_scaled, dtype=tf.float32)

# Stack both targets together for proper loss computation
y_train_tensor = tf.convert_to_tensor(np.hstack([y_train_scaled, y_train_scaled_R_phys]), dtype=tf.float32)
y_val_tensor = tf.convert_to_tensor(np.hstack([y_val_scaled, y_val_scaled_R_phys]), dtype=tf.float32)

print('Training data shape:', X_train_tensor.shape, y_train_tensor.shape)
print('Validation data shape:', X_val_tensor.shape, y_val_tensor.shape)

# Define MLP model
def create_model():
    model = Sequential([
        Dense(4, activation="relu", input_shape=(X_train_tensor.shape[1],), kernel_regularizer=regularizers.l2(1e-3)),
        Dense(1, activation="linear")  # Single output: R_pred
    ])
    return model

# Create and compile model
model = create_model()
optimizer = SGD(learning_rate=0.001, momentum=0.8)

# Define hybrid loss function
def hybrid_loss(y_true, y_pred):
    R_true = y_true[:, 0:1]  # Extract R_true from y_true
    R_phys = y_true[:, 1:2]  # Extract R_phys from y_true

    # Compute MSE for data loss (R_pred vs R_true)
    L_data = tf.reduce_mean(tf.square(y_pred - R_true))  # MSE Loss (Data)

    # Convert R_phys to numpy array for physics loss calculation
    R_phys_numpy = R_phys.numpy() if tf.executing_eagerly() else tf.make_ndarray(R_phys)

    # Normalize R_phys for physics loss
    R_phys_normalized = scaler_phys.transform(R_phys_numpy)

    # Convert normalized R_phys back to tensor
    R_phys_normalized_tensor = tf.convert_to_tensor(R_phys_normalized, dtype=tf.float32)

    # Compute MSE for physics loss (R_pred vs normalized R_phys)
    L_physics = tf.reduce_mean(tf.square(y_pred - R_phys_normalized_tensor))  # MSE Loss (Physics)

    # Weighted combination of data loss and physics loss
    physics_weight = 0.05
    data_weight = 0.95
    L_total = data_weight * L_data + physics_weight * L_physics
    return L_total

# Custom Model Checkpoint callback
class CustomModelCheckpoint(tf.keras.callbacks.Callback):
    def __init__(self, filepath, save_freq=1):
        self.filepath = filepath
        self.save_freq = save_freq

    def on_epoch_end(self, epoch, logs=None):
        if epoch == 0 or (epoch + 1) % self.save_freq == 0:
            print(f"\nSaving model at epoch {epoch + 1}...")
            self.model.save(self.filepath.format(epoch=epoch + 1))

checkpoint_callback = CustomModelCheckpoint("/home/qiki1832/PycharmProjects/PINN/PhysicsRegularisation/MLP_01/model_epoch_{epoch}.h5", save_freq=1)

# Custom callback to track and save metrics
class MetricsLogger(tf.keras.callbacks.Callback):
    def __init__(self, X_train, y_train, X_test, y_test, filepath):
        self.X_train = X_train
        self.y_train = y_train
        self.X_test = X_test
        self.y_test = y_test
        self.filepath = filepath
        self.metrics_data = []

    def on_epoch_end(self, epoch, logs=None):
        y_train_pred = self.model.predict(self.X_train, verbose=0)
        y_test_pred = self.model.predict(self.X_test, verbose=0)

        train_mae = mean_absolute_error(self.y_train, y_train_pred)
        train_mse = mean_squared_error(self.y_train, y_train_pred)
        train_rmse = np.sqrt(train_mse)
        train_r2 = r2_score(self.y_train, y_train_pred)

        test_mae = mean_absolute_error(self.y_test, y_test_pred)
        test_mse = mean_squared_error(self.y_test, y_test_pred)
        test_rmse = np.sqrt(test_mse)
        test_r2 = r2_score(self.y_test, y_test_pred)

        # Get current learning rate
        current_lr = float(self.model.optimizer.lr.numpy()) if tf.executing_eagerly() else self.model.optimizer.lr

        # Store metrics
        self.metrics_data.append([epoch + 1, current_lr, train_mae, train_mse, train_rmse, train_r2,
                                  test_mae, test_mse, test_rmse, test_r2])

        # Save to CSV at every epoch
        metrics_df = pd.DataFrame(self.metrics_data, columns=[
            "Epoch", "Learning Rate",
            "Train MAE", "Train MSE", "Train RMSE", "Train R²",
            "Val MAE", "Val MSE", "Val RMSE", "Val R2"
        ])
        metrics_df.to_csv(self.filepath, index=False)

# Define log file path
metrics_csv_path = "/home/qiki1832/PycharmProjects/PINN/PhysicsRegularisation/MLP_01/training_epoch_metrics.csv"

# Initialize custom callback
metrics_logger = MetricsLogger(X_train_tensor, y_train_tensor, X_val_tensor, y_val_tensor, metrics_csv_path)

# Compile model
model.compile(optimizer=optimizer, loss=hybrid_loss, metrics=['mae'])

# Train the model
num_epochs = 50
batch_size = 256

model.fit(X_train_tensor, y_train_tensor,  # Now y_train_tensor includes both R_true and R_phys
          validation_data=(X_val_tensor, y_val_tensor),  # Pass both targets for validation
          epochs=num_epochs, batch_size=batch_size, verbose=1,
          callbacks=[checkpoint_callback, metrics_logger])

I tried several things: seperating the target files to have an easier normalization or a combined target file, but I encounter always errors when trying to read it.