b_fnn_MCdrop_Vtot_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
y = total_water_contents  # For predicting total water content

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

# 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):
    # Suggest hyperparameters for tuning
    num_layers = trial.suggest_int('num_layers', 2, 4)
    hidden_units = trial.suggest_categorical('hidden_units', [64, 128, 256, 512])
    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, 1))  # Output layer for regression
    
    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]
    model = build_model(trial, input_dim)
    
    # Suggest hyperparameters for the optimizer
    lr = trial.suggest_loguniform('learning_rate', 1e-5, 1e-2)
    optimizer_name = trial.suggest_categorical('optimizer', ['Adam', 'RMSprop', 'SGD'])
    weight_decay = trial.suggest_loguniform('weight_decay', 1e-8, 1e-3)
    
    # 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)
    elif optimizer_name == 'SGD':
        optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=weight_decay)
    
    # Create DataLoaders
    batch_size = trial.suggest_categorical('batch_size', [64, 128, 256])
    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 = 50
    patience = 5  # 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=50)

# 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]
best_model = build_model(trial, input_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'])
elif trial.params['optimizer'] == 'SGD':
    optimizer = optim.SGD(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 = 50
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:.4f}, Val Loss: {val_loss:.4f}')

# ==============================================
# 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:.4f}')

# ==============================================
# 9. Predict with Uncertainty Estimation using MC Dropout
# ==============================================

def predict_with_uncertainty(model, X, n_samples=100):
    model.train()  # Enable dropout layers during inference
    predictions = []
    for _ in range(n_samples):
        outputs = model(X)
        predictions.append(outputs.detach().cpu().numpy())
    predictions = np.array(predictions)
    mean_preds = predictions.mean(axis=0)
    std_preds = predictions.std(axis=0)
    return mean_preds, std_preds

# Prepare test data
X_test_tensor = X_tensor[test_dataset.indices]
y_test_tensor = y_tensor[test_dataset.indices]

# Predict with uncertainty
mean_preds, std_preds = predict_with_uncertainty(best_model, X_test_tensor)

# Convert to numpy arrays
y_test = y_test_tensor.cpu().numpy()
mean_preds = mean_preds.squeeze()
std_preds = std_preds.squeeze()

# ==============================================
# 10. Visualize the Results
# ==============================================

# Sort the data for better visualization
sorted_indices = np.argsort(y_test.squeeze())
y_test_sorted = y_test.squeeze()[sorted_indices]
mean_preds_sorted = mean_preds[sorted_indices]
std_preds_sorted = std_preds[sorted_indices]

plt.figure(figsize=(12, 6))
plt.plot(y_test_sorted, label='True Values', color='red')
plt.errorbar(range(len(y_test_sorted)), mean_preds_sorted, yerr=std_preds_sorted, fmt='o', label='Predictions with Uncertainty', ecolor='gray', alpha=0.01)
plt.legend()
plt.xlabel('Sample Index (sorted by true value)')
plt.ylabel('Total Water Content')
plt.title('Predictions with Uncertainty')
plt.show()

# ==============================================
# 11. 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')

# ==============================================
# 12. 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
# ==============================================