Consider the following code. I want to duplicate the figure under # Plot for all samples, and on the duplicated figure, I want to add the subset of selected peaks as markings with a cross (on y_test_subset). this will help me to visualize which peaks where selected for evaluating the model based on the subset of peaks. import pandas as pd
import numpy as np
import pickle
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import GroupShuffleSplit
from tensorflow.keras.models import Sequential, load_model
from tensorflow.keras.layers import LSTM, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping, Callback
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import matplotlib.pyplot as plt
import time

# ===========================
# 1. Data Loading and Preprocessing
# ===========================

# Set simulation data loading option (only "option1" is available now)
simulation_data = "option1"

# Target parameter: 'curvature', 'compression' or 'Top_tension'
Target_parameter = "Top_tension"

Time_step = 0.1 #timestep used in the simulations

# Set the training mode: 'standard', 'maxima', or 'minima'
training_mode = 'maxima'  

# Set the evaluation mode: 'all', 'maxima', 'minima'
evaluation_mode = 'maxima'

sequence_length = 100  # Define the length of sequences

# Initialize lists to hold features and targets
X_sequences = []
y_sequences = []
groups = []
time_indices = []


if simulation_data == "option1":
    # Define file paths for each variable
    if Target_parameter == 'curvature':
        file_data = 'curvature_max_data.txt'
    elif Target_parameter == 'compression':
        file_data = 'compression_max_data.txt'
    elif Target_parameter == 'Top_tension':
        file_data = 'Top_tension_max_data.txt' 
        
    file_paths = {
        'heave': 'heave_data.txt',
        'heaveVelocity': 'heaveVelocity_data.txt',
        'surge': 'surge_data.txt',
        'surgeVelocity': 'surgeVelocity_data.txt',
        'sway': 'sway_data.txt',
        'swayVelocity': 'swayVelocity_data.txt',
        'cable_parameter': file_data
    }

    # Load each variable's data from the respective text files
    data_dict = {}
    for variable, file_path in file_paths.items():
        data = pd.read_csv(file_path, header=None)
        data_dict[variable] = data  # Each column represents a separate simulation

    num_simulations = data_dict['heave'].shape[1]  # Number of simulations

    # Extract features and targets for each simulation and generate sequences
    for i in range(num_simulations):
        # Extract features for the i-th simulation
        heave = data_dict['heave'].iloc[:, i].values
        heaveVelocity = data_dict['heaveVelocity'].iloc[:, i].values
        surge = data_dict['surge'].iloc[:, i].values
        surgeVelocity = data_dict['surgeVelocity'].iloc[:, i].values
        sway = data_dict['sway'].iloc[:, i].values
        swayVelocity = data_dict['swayVelocity'].iloc[:, i].values

        # Stack features horizontally (shape: [time_steps, 6])
        features = np.column_stack([
            heave,
            heaveVelocity,
            surge,
            surgeVelocity,
            sway,
            swayVelocity
        ])

        # Extract target for the i-th simulation
        target = data_dict['cable_parameter'].iloc[:, i].values  # Shape: [time_steps]

        # Generate sequences
        num_sequences = len(features) - sequence_length
        for t in range(num_sequences):
            X_seq = features[t:t + sequence_length]
            y_seq = target[t + sequence_length]  # Predicting the value at t + sequence_length
            X_sequences.append(X_seq)
            y_sequences.append(y_seq)
            groups.append(i)  # Keep track of the simulation index
            time_indices.append(t + sequence_length)  # Time index of the target value

    # Convert lists to NumPy arrays
    X_sequences = np.array(X_sequences)  # Shape: [num_sequences_total, sequence_length, num_features]
    y_sequences = np.array(y_sequences)  # Shape: [num_sequences_total]
    groups = np.array(groups)  # Shape: [num_sequences_total]
    time_indices = np.array(time_indices)  # Shape: [num_sequences_total]

# ===========================
# 2. Train-Test Split with Grouping
# ===========================

# Split the data into training and testing sets using GroupShuffleSplit
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
train_idx, test_idx = next(gss.split(X_sequences, y_sequences, groups=groups))

X_train, X_test = X_sequences[train_idx], X_sequences[test_idx]
y_train, y_test = y_sequences[train_idx], y_sequences[test_idx]
groups_train = groups[train_idx]
groups_test = groups[test_idx]
time_indices_train = time_indices[train_idx]
time_indices_test = time_indices[test_idx]

# ===========================
# 3. Data Scaling
# ===========================

# Initialize scalers
scaler_features = MinMaxScaler()
scaler_target = MinMaxScaler()

# Flatten the training data to fit the scalers
X_train_flat = X_train.reshape(-1, X_train.shape[2])  # Shape: [num_samples * seq_length, n_features]
y_train_flat = y_train.reshape(-1, 1)                 # Shape: [num_samples, 1]

# Fit scalers on training data
scaler_features.fit(X_train_flat)
scaler_target.fit(y_train_flat)

# Transform training data
X_train_scaled = scaler_features.transform(X_train_flat).reshape(X_train.shape)
y_train_scaled = scaler_target.transform(y_train_flat).reshape(-1)

# Transform testing data
X_test_flat = X_test.reshape(-1, X_test.shape[2])
y_test_flat = y_test.reshape(-1, 1)

X_test_scaled = scaler_features.transform(X_test_flat).reshape(X_test.shape)
y_test_scaled = scaler_target.transform(y_test_flat).reshape(-1)

# ===========================
# 4. Generating Sample Weights for Maxima/Minima Emphasis
# ===========================

# Initialize sample weights
sample_weights = np.ones(len(y_train))

# Set the weight factor for maxima/minima
weight_factor = 20  # Increase to give more emphasis

if training_mode in ['maxima', 'minima']:
    from scipy.signal import find_peaks
    if Target_parameter == 'curvature':
        threshold_filter = 0.1
    elif Target_parameter == 'compression':
        threshold_filter = 0.5
    elif Target_parameter == 'Top_tension':
        threshold_filter = 0.5 
        
    unique_groups = np.unique(groups_train)

    for sim in unique_groups:
        # Get the full target array for this simulation
        sim_target_full = data_dict['cable_parameter'].iloc[:, sim].values

        # Find peaks or troughs in sim_target_full
        if training_mode == 'maxima':
            peaks, _ = find_peaks(sim_target_full,threshold=threshold_filter,distance=65)
            peak_indices = peaks  # Time indices of maxima
        elif training_mode == 'minima':
            troughs, _ = find_peaks(-sim_target_full,threshold=threshold_filter,distance=65)
            peak_indices = troughs  # Time indices of minima

        # Get the indices of samples in training set that belong to this simulation
        sim_indices = np.where(groups_train == sim)[0]
        sim_time_indices = time_indices_train[sim_indices]

        # Now, for each sample in sim_indices, check if its time index is in peak_indices
        for idx, time_idx in zip(sim_indices, sim_time_indices):
            if time_idx in peak_indices:
                sample_weights[idx] = weight_factor

# ===========================
# 5. Custom Callback for Timing
# ===========================

class TimeHistory(Callback):
    def on_train_begin(self, logs=None):
        self.times = []

    def on_epoch_begin(self, epoch, logs=None):
        self.epoch_time_start = time.time()

    def on_epoch_end(self, epoch, logs=None):
        epoch_time = time.time() - self.epoch_time_start
        self.times.append(epoch_time)
        print(f" Epoch {epoch + 1} time: {epoch_time:.2f} seconds")

time_callback = TimeHistory()

# ===========================
# 6. Building the LSTM Model
# ===========================

sequence_length = X_train_scaled.shape[1]
n_features = X_train_scaled.shape[2]

model = Sequential()
model.add(LSTM(units=50, return_sequences=True, input_shape=(sequence_length, n_features)))
model.add(Dropout(0.2))
model.add(LSTM(units=50, return_sequences=False))  # Changed to return_sequences=False
model.add(Dropout(0.2))
model.add(Dense(units=1))  # Output layer for regression

# Compile the model
model.compile(optimizer='adam', loss='mean_squared_error')

# ===========================
# 7. Training the Model
# ===========================

# EarlyStopping to prevent overfitting
early_stopping = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)

# Start training timer
start_time = time.time()

# Train the model
history = model.fit(
    X_train_scaled, y_train_scaled,
    epochs=100,
    batch_size=32,
    validation_split=0.2,
    callbacks=[early_stopping, time_callback],
    shuffle=True,  # Shuffle for better generalization
    sample_weight=sample_weights  # Include sample weights
)

# End training timer
end_time = time.time()
total_time = end_time - start_time
print(f"Total Training Time: {total_time:.2f} seconds")

# ===========================
# 8. Saving the Model and Scalers
# ===========================

# Save the trained model
model.save('cable_'+Target_parameter+'_lstm_model.h5')
print("Model saved to disk as .h5")

# Save the scalers
with open('scaler_features-'+Target_parameter+'.pkl', 'wb') as f:
    pickle.dump(scaler_features, f)
with open('scaler_target-'+Target_parameter+'.pkl', 'wb') as f:
    pickle.dump(scaler_target, f)
print("Scalers saved to disk as .pkl'")

# ===========================
# 9. Loading the Model and Scalers
# ===========================

# Load the model
loaded_model = load_model('cable_'+Target_parameter+'_lstm_model.h5')

# Load the scalers
with open('scaler_features-'+Target_parameter+'.pkl', 'rb') as f:
    loaded_scaler_features = pickle.load(f)
with open('scaler_target-'+Target_parameter+'.pkl', 'rb') as f:
    loaded_scaler_target = pickle.load(f)

# ===========================
# 10. Making Predictions
# ===========================

# Predict on the test set
y_pred_scaled = loaded_model.predict(X_test_scaled)  # Shape: [num_samples, 1]

# Inverse transform the predictions and actual values to original scale
y_pred = loaded_scaler_target.inverse_transform(y_pred_scaled).flatten()
y_test_original = loaded_scaler_target.inverse_transform(y_test_scaled.reshape(-1, 1)).flatten()

# ===========================
# 11. Evaluating the Model
# ===========================

# Evaluate model on all test samples
mse = mean_squared_error(y_test_original, y_pred)
rmse = np.sqrt(mse)
mae = mean_absolute_error(y_test_original, y_pred)
# Avoid division by zero in MAPE
non_zero = y_test_original != 0
mape = np.mean(np.abs((y_test_original[non_zero] - y_pred[non_zero]) / y_test_original[non_zero])) * 100
r2 = r2_score(y_test_original, y_pred)

print("Performance Metrics on All Test Samples:")
print(f"Mean Squared Error (MSE): {mse:.4f}")
print(f"Root Mean Squared Error (RMSE): {rmse:.4f}")
print(f"Mean Absolute Error (MAE): {mae:.4f}")
# print(f"Mean Absolute Percentage Error (MAPE): {mape:.2f}%")
print(f"R2 Score: {r2:.4f}")

# ===========================
# 11.a Evaluating and Visualizing Peaks
# ===========================

if evaluation_mode != 'all':
    peak_type = evaluation_mode  # 'maxima' or 'minima'

    is_peak_test = np.zeros(len(y_test_original), dtype=bool)

    from scipy.signal import find_peaks
    if Target_parameter == 'curvature':
        threshold_filter = 0.1
    elif Target_parameter == 'compression':
        threshold_filter = 0.5
    elif Target_parameter == 'Top_tension':
        threshold_filter = 0.5 
        
    unique_groups_test = np.unique(groups_test)

    for sim in unique_groups_test:
        # Get the full target array for this simulation
        sim_target_full = data_dict['cable_parameter'].iloc[:, sim].values

        # Find peaks or troughs in sim_target_full
        if peak_type == 'maxima':
            peaks, _ = find_peaks(sim_target_full,threshold=threshold_filter,distance=60)
            peak_indices = peaks  # Time indices of maxima
        elif peak_type == 'minima':
            troughs, _ = find_peaks(-sim_target_full,threshold=threshold_filter,distance=60)
            peak_indices = troughs  # Time indices of minima

        # Get the indices of samples in test set that belong to this simulation
        sim_indices = np.where(groups_test == sim)[0]
        sim_time_indices = time_indices_test[sim_indices]

        # Now, for each sample in sim_indices, check if its time ind☻ex is in peak_indices
        for idx, time_idx in zip(sim_indices, sim_time_indices):
            if time_idx in peak_indices:
                is_peak_test[idx] = True
else:
    is_peak_test = np.ones(len(y_test_original), dtype=bool)

# Select the samples corresponding to peaks/minima or all
y_test_subset = y_test_original[is_peak_test]
y_pred_subset = y_pred[is_peak_test]

# Calculate performance metrics for subset
mse_subset = mean_squared_error(y_test_subset, y_pred_subset)
rmse_subset = np.sqrt(mse_subset)
mae_subset = mean_absolute_error(y_test_subset, y_pred_subset)
non_zero_subset = y_test_subset != 0
mape_subset = np.mean(np.abs((y_test_subset[non_zero_subset] - y_pred_subset[non_zero_subset]) / y_test_subset[non_zero_subset])) * 100
r2_subset = r2_score(y_test_subset, y_pred_subset)

print(f"\nPerformance Metrics on {evaluation_mode.capitalize()} Samples:")
print(f"Mean Squared Error (MSE): {mse_subset:.4f}")
print(f"Root Mean Squared Error (RMSE): {rmse_subset:.4f}")
print(f"Mean Absolute Error (MAE): {mae_subset:.4f}")
# print(f"Mean Absolute Percentage Error (MAPE): {mape_subset:.2f}%")
print(f"R2 Score: {r2_subset:.4f}")

# ===========================
# 11.b Visualizing the Results
# ===========================
if Target_parameter == 'curvature':
    Units = '(rad/m)'
    Name = 'Max curvature '
elif Target_parameter == 'compression':
    Units = '(kN)'
    Name = 'Min tension '
elif Target_parameter == 'Top_tension':
    Units = '(kN)'
    Name = 'Max tension '
    
# Plot for all test samples
num_samples = len(y_test_original)
time = np.arange(0, num_samples) * Time_step  

plt.figure(figsize=(12, 6))
plt.plot(time, y_test_original, label='Actual Cable Response', alpha=0.7)
plt.plot(time, y_pred, label='Predicted Cable Response', alpha=0.7)
plt.title('Actual vs Predicted - All Samples')
plt.xlabel('Time (s)')
plt.ylabel(Name+Units)
plt.legend()
plt.show()

# Plot for peaks or minima
plt.figure(figsize=(8, 6))
plt.scatter(y_test_subset, y_pred_subset, alpha=0.7)
plt.plot([y_test_subset.min(), y_test_subset.max()], [y_test_subset.min(), y_test_subset.max()], 'r--')
plt.title(f'Actual vs Predicted at {evaluation_mode.capitalize()}')
plt.xlabel('Actual '+ Name+Units)
plt.ylabel('Predicted '+ Name+Units)
plt.show()