b_fnn_MCdrop_wmap_R.py

# ==============================================
# 1. Import Necessary Libraries
# ==============================================
# %%
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader, random_split
import torch.optim as optim
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
import optuna
from optuna.trial import TrialState
import matplotlib.pyplot as plt
import random

# ==============================================
# 2. Set Random Seeds for Reproducibility
# ==============================================

seed = 42
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)

# ==============================================
# 3. Data Loading and Preprocessing
# ==============================================

# Load the data
saturation_maps = np.load("saturation_maps.npy")
resistivity_maps = np.load("resistivity_maps.npy")
total_water_contents = np.load("total_water_contents.npy")
data_dd_rhoa_array = np.load("data_dd.npy")
data_slm_rhoa_array = np.load("data_slm.npy")
apparent_resistivity = np.concatenate([data_dd_rhoa_array, data_slm_rhoa_array], axis=1)

print(f"Saturation Maps Shape: {saturation_maps.shape}")
print(f"Resistivity Maps Shape: {resistivity_maps.shape}")
print(f"Total Water Contents Shape: {total_water_contents.shape}")
print(f"DD Rhoa Array Shape: {data_dd_rhoa_array.shape}")
print(f"SLM Rhoa Array Shape: {data_slm_rhoa_array.shape}")
print(f"All Rhoa Array Shape: {apparent_resistivity.shape}")

# Normalize the input features
scaler = StandardScaler()
X = scaler.fit_transform(apparent_resistivity)

# Optional: Dimensionality Reduction using PCA
pca = PCA(n_components=0.95)  # Retain 95% of the variance
X_reduced = pca.fit_transform(X)
print(f"Reduced feature count: {X_reduced.shape[1]}")

# Prepare the target variable
# Flatten the saturation maps
y = saturation_maps.reshape(len(saturation_maps), -1)
print(f"Flattened Saturation Maps Shape: {y.shape}")

# Convert to PyTorch tensors
X_tensor = torch.tensor(X_reduced, dtype=torch.float32)
y_tensor = torch.tensor(y, dtype=torch.float32)

# Create Dataset
dataset = TensorDataset(X_tensor, y_tensor)

# Split into Train, Validation, and Test sets
train_size = int(0.7 * len(dataset))
val_size = int(0.15 * len(dataset))
test_size = len(dataset) - train_size - val_size

train_dataset, val_dataset, test_dataset = random_split(
    dataset, [train_size, val_size, test_size]
)

# ==============================================
# 4. Define the Model Class
# ==============================================

def build_model(trial, input_dim, output_dim):
    # Suggest hyperparameters for tuning
    num_layers = trial.suggest_int('num_layers', 2, 5)
    hidden_units = trial.suggest_categorical('hidden_units', [256, 512, 1024])
    dropout_rate = trial.suggest_uniform('dropout_rate', 0.0, 0.5)
    activation_name = trial.suggest_categorical('activation', ['ReLU', 'LeakyReLU', 'ELU'])
    
    # Choose activation function
    if activation_name == 'ReLU':
        activation = nn.ReLU()
    elif activation_name == 'LeakyReLU':
        activation = nn.LeakyReLU()
    elif activation_name == 'ELU':
        activation = nn.ELU()
    
    layers = []
    in_features = input_dim
    for i in range(num_layers):
        layers.append(nn.Linear(in_features, hidden_units))
        layers.append(activation)
        layers.append(nn.Dropout(dropout_rate))
        in_features = hidden_units
    layers.append(nn.Linear(hidden_units, output_dim))  # Output layer for saturation map
    
    model = nn.Sequential(*layers)
    return model

# ==============================================
# 5. Define the Objective Function for Optuna
# ==============================================

def objective(trial):
    # Generate the model
    input_dim = X_reduced.shape[1]
    output_dim = y.shape[1]
    model = build_model(trial, input_dim, output_dim)
    
    # Suggest hyperparameters for the optimizer
    lr = trial.suggest_loguniform('learning_rate', 1e-5, 1e-3)
    optimizer_name = trial.suggest_categorical('optimizer', ['Adam', 'RMSprop'])
    weight_decay = trial.suggest_loguniform('weight_decay', 1e-8, 1e-4)
    
    # Choose optimizer
    if optimizer_name == 'Adam':
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    elif optimizer_name == 'RMSprop':
        optimizer = optim.RMSprop(model.parameters(), lr=lr, weight_decay=weight_decay)
    
    # Create DataLoaders
    batch_size = trial.suggest_categorical('batch_size', [32, 64, 128])
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)
    
    # Define loss function
    criterion = nn.MSELoss()
    
    # Training loop
    num_epochs = 30
    patience = 3  # For early stopping
    best_val_loss = float('inf')
    trigger_times = 0
    
    for epoch in range(num_epochs):
        model.train()
        train_loss = 0.0
        for X_batch, y_batch in train_loader:
            optimizer.zero_grad()
            outputs = model(X_batch)
            loss = criterion(outputs, y_batch)
            loss.backward()
            optimizer.step()
            train_loss += loss.item() * X_batch.size(0)
        train_loss /= train_size
        
        # Validation
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for X_batch, y_batch in val_loader:
                outputs = model(X_batch)
                loss = criterion(outputs, y_batch)
                val_loss += loss.item() * X_batch.size(0)
        val_loss /= val_size
        
        trial.report(val_loss, epoch)
        
        # Handle pruning
        if trial.should_prune():
            raise optuna.exceptions.TrialPruned()
        
        # Early stopping
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            trigger_times = 0
        else:
            trigger_times += 1
            if trigger_times >= patience:
                break
        
    return best_val_loss

# ==============================================
# 6. Run Hyperparameter Optimization with Optuna
# ==============================================

study = optuna.create_study(direction='minimize')
study.optimize(objective, n_trials=30)

# Print study statistics
print("Number of finished trials: ", len(study.trials))
print("Best trial:")
trial = study.best_trial

print(f"  Value (Best Validation Loss): {trial.value}")
print("  Params: ")
for key, value in trial.params.items():
    print(f"    {key}: {value}")

# ==============================================
# 7. Train the Final Model with Best Hyperparameters
# ==============================================

# Build the best model
input_dim = X_reduced.shape[1]
output_dim = y.shape[1]
best_model = build_model(trial, input_dim, output_dim)

# Define the optimizer with best hyperparameters
if trial.params['optimizer'] == 'Adam':
    optimizer = optim.Adam(best_model.parameters(), lr=trial.params['learning_rate'], weight_decay=trial.params['weight_decay'])
elif trial.params['optimizer'] == 'RMSprop':
    optimizer = optim.RMSprop(best_model.parameters(), lr=trial.params['learning_rate'], weight_decay=trial.params['weight_decay'])

# DataLoaders with best batch size
batch_size = trial.params['batch_size']
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
test_loader = DataLoader(test_dataset, batch_size=batch_size)

# Loss function
criterion = nn.MSELoss()

# Training the final model
num_epochs = 30
patience = 5
best_val_loss = float('inf')
trigger_times = 0

for epoch in range(num_epochs):
    best_model.train()
    train_loss = 0.0
    for X_batch, y_batch in train_loader:
        optimizer.zero_grad()
        outputs = best_model(X_batch)
        loss = criterion(outputs, y_batch)
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * X_batch.size(0)
    train_loss /= train_size
    
    # Validation
    best_model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for X_batch, y_batch in val_loader:
            outputs = best_model(X_batch)
            loss = criterion(outputs, y_batch)
            val_loss += loss.item() * X_batch.size(0)
    val_loss /= val_size
    
    # Early stopping
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        trigger_times = 0
        # Save the best model
        torch.save(best_model.state_dict(), 'best_model.pt')
    else:
        trigger_times += 1
        if trigger_times >= patience:
            print('Early stopping!')
            break
    
    print(f'Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}')

# ==============================================
# 8. Evaluate the Model on Test Data
# ==============================================

# Load the best model
best_model.load_state_dict(torch.load('best_model.pt'))

# Evaluate on Test Set
best_model.eval()
test_loss = 0.0
with torch.no_grad():
    for X_batch, y_batch in test_loader:
        outputs = best_model(X_batch)
        loss = criterion(outputs, y_batch)
        test_loss += loss.item() * X_batch.size(0)
test_loss /= test_size
print(f'Test Loss: {test_loss:.6f}')

# ==============================================
# 9. Visualize Some Predictions
# ==============================================

# Select a few samples from the test set
num_samples = 5
X_test_tensor = X_tensor[test_dataset.indices]
y_test_tensor = y_tensor[test_dataset.indices]

indices = np.random.choice(len(X_test_tensor), num_samples, replace=False)

best_model.eval()
with torch.no_grad():
    for idx in indices:
        X_sample = X_test_tensor[idx].unsqueeze(0)
        y_true = y_test_tensor[idx].numpy().reshape(15, 142)
        y_pred = best_model(X_sample).numpy().reshape(15, 142)
        
        # Plot the true and predicted saturation maps
        fig, axes = plt.subplots(1, 2, figsize=(12, 5))
        im1 = axes[0].imshow(y_true, cmap='viridis')
        axes[0].set_title('True Saturation Map')
        fig.colorbar(im1, ax=axes[0])
        
        im2 = axes[1].imshow(y_pred, cmap='viridis')
        axes[1].set_title('Predicted Saturation Map')
        fig.colorbar(im2, ax=axes[1])
        
        plt.show()

# ==============================================
# 10. Save the Scaler and PCA Model for Future Use
# ==============================================

import joblib

# Save the scaler and PCA model
joblib.dump(scaler, 'scaler.joblib')
joblib.dump(pca, 'pca.joblib')

# ==============================================
# 11. Load the Scaler and PCA Model (Example)
# ==============================================

# To load the scaler and PCA model in the future:
# scaler = joblib.load('scaler.joblib')
# pca = joblib.load('pca.joblib')

# ==============================================
# End of Script
# ==============================================