model_fdm.py

import torch
import numpy as np
import matplotlib.pyplot as plt
from Dataset import DataAdditiveManufacturing, DataThermoforming
from model import NeuralNetwork

DEVICE = torch.device('cpu')

# Set global plotting parameters
plt.rcParams.update({'font.size': 14,
                     'figure.figsize': (10, 8),
                     'lines.linewidth':  2,
                     'lines.markersize': 6,
                     'axes.grid': True,
                     'axes.labelsize': 16,
                     'legend.fontsize': 10,
                     'xtick.labelsize': 14,
                     'ytick.labelsize': 14,
                     'figure.autolayout': True
                     })

def set_seed(seed=42):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)    

def train_neural_network(model, inputs, outputs, optimizer, epochs=1000, lr_scheduler=None):
    model.train()
    for epoch in range(epochs):
        optimizer.zero_grad()
        predictions = model(inputs)
        loss = torch.mean(torch.square(predictions - outputs))
        loss.backward()
        optimizer.step()

        if lr_scheduler:
            lr_scheduler.step()

        if epoch % 100 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item()}, Learning Rate: {optimizer.param_groups[0]["lr"]}')

def kfold_indices(n_samples, k=5, seed=42, shuffle=True):
    rng = np.random.default_rng(seed)
    indices = np.arange(n_samples)
    if shuffle:
        rng.shuffle(indices)
    fold_sizes = np.full(k, n_samples // k, dtype=int)
    fold_sizes[: n_samples % k] += 1
    current = 0
    folds = []
    for fold_size in fold_sizes:
        start, stop = current, current + fold_size
        folds.append(indices[start:stop])
        current = stop
    return folds

def ridge_fit_predict(x_train, y_train, x_test, alpha=1.0):
    # Closed-form ridge regression: W = (X^T X + alpha I)^-1 X^T Y
    x_aug = np.concatenate([x_train, np.ones((x_train.shape[0], 1))], axis=1)
    xtx = x_aug.T @ x_aug
    reg = alpha * np.eye(xtx.shape[0], dtype=x_train.dtype)
    reg[-1, -1] = 0.0  # don't regularize bias
    w = np.linalg.solve(xtx + reg, x_aug.T @ y_train)
    x_test_aug = np.concatenate([x_test, np.ones((x_test.shape[0], 1))], axis=1)
    return x_test_aug @ w

def kfold_ridge_baseline(inputs, outputs, k=5, alpha=1.0, seed=42):
    folds = kfold_indices(len(inputs), k=k, seed=seed, shuffle=True)
    mse_folds = []
    r2_folds = []
    for i in range(k):
        test_idx = folds[i]
        train_idx = np.concatenate([f for j, f in enumerate(folds) if j != i])

        x_train = inputs[train_idx]
        y_train = outputs[train_idx]
        x_test = inputs[test_idx]
        y_test = outputs[test_idx]

        # Train-only normalization
        x_mean = x_train.mean(axis=0)
        x_std = x_train.std(axis=0) + 1e-8
        y_mean = y_train.mean(axis=0)
        y_std = y_train.std(axis=0) + 1e-8

        x_train_n = (x_train - x_mean) / x_std
        x_test_n = (x_test - x_mean) / x_std
        y_train_n = (y_train - y_mean) / y_std

        y_pred_n = ridge_fit_predict(x_train_n, y_train_n, x_test_n, alpha=alpha)
        y_pred = y_pred_n * y_std + y_mean

        mse = np.mean((y_pred - y_test) ** 2, axis=0)
        ss_res = np.sum((y_test - y_pred) ** 2, axis=0)
        ss_tot = np.sum((y_test - np.mean(y_test, axis=0)) ** 2, axis=0)
        r2 = 1 - ss_res / ss_tot
        mse_folds.append(mse)
        r2_folds.append(r2)

    mse_folds = np.stack(mse_folds, axis=0)
    r2_folds = np.stack(r2_folds, axis=0)
    print("Ridge k-fold CV (alpha=%.3g, k=%d)" % (alpha, k))
    print("MSE mean:", np.mean(mse_folds, axis=0))
    print("MSE std:", np.std(mse_folds, axis=0))
    print("R2 mean:", np.mean(r2_folds, axis=0))
    print("R2 std:", np.std(r2_folds, axis=0))

def main():
    dataset = DataAdditiveManufacturing()
    inputs = dataset.get_input(normalize=False)
    outputs = dataset.get_output(normalize=False)

    idx_train = np.random.choice(len(inputs), size=int(0.95 * len(inputs)), replace=False)
    idx_test = np.setdiff1d(np.arange(len(inputs)), idx_train)

    # Normalize using train-only statistics to avoid test leakage
    x_train = inputs[idx_train]
    y_train = outputs[idx_train]
    x_test = inputs[idx_test]
    y_test = outputs[idx_test]

    x_mean = x_train.mean(axis=0)
    x_std = x_train.std(axis=0) + 1e-8
    y_mean = y_train.mean(axis=0)
    y_std = y_train.std(axis=0) + 1e-8

    x_train_n = (x_train - x_mean) / x_std
    x_test_n = (x_test - x_mean) / x_std
    y_train_n = (y_train - y_mean) / y_std
    y_test_n = (y_test - y_mean) / y_std

    inputs_train = torch.tensor(x_train_n, dtype=torch.float32).to(DEVICE)
    outputs_train = torch.tensor(y_train_n, dtype=torch.float32).to(DEVICE)

    inputs_test = torch.tensor(x_test_n, dtype=torch.float32).to(DEVICE)
    outputs_test = torch.tensor(y_test_n, dtype=torch.float32).to(DEVICE)
    
    layer_sizes = [inputs.shape[1], 64, 32, outputs.shape[1]]
    dropout_rate = 0.1
    model = NeuralNetwork(layer_sizes, dropout_rate=dropout_rate, activation=torch.nn.ReLU).to(DEVICE)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2000, gamma=0.9)

    # Create a proper dataset that keeps input-output pairs together
    train_dataset = torch.utils.data.TensorDataset(inputs_train, outputs_train)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True)

    # Train the model
    epochs = 5000
    best_test_loss = float("inf")
    patience = 400
    patience_left = patience
    for epoch in range(epochs):
        model.train()
        for inputs_batch, outputs_batch in train_loader:
            inputs_batch = inputs_batch.to(DEVICE)
            outputs_batch = outputs_batch.to(DEVICE)
            optimizer.zero_grad()
            predictions = model(inputs_batch)
            loss = torch.mean(torch.square(predictions - outputs_batch))
            loss.backward()
            optimizer.step()

        if lr_scheduler:
            lr_scheduler.step()

        if epoch % 200 == 0:
            model.eval()
            with torch.no_grad():
                train_pred = model(inputs_train, train=False)
                train_loss = torch.mean(torch.square(train_pred - outputs_train))
                test_pred = model(inputs_test, train=False)
                test_loss = torch.mean(torch.square(test_pred - outputs_test))
            print(f'Epoch {epoch}, Train Loss: {train_loss.item():.6f}, Test Loss: {test_loss.item():.6f}')
            if test_loss.item() < best_test_loss - 1e-6:
                best_test_loss = test_loss.item()
                patience_left = patience
            else:
                patience_left -= 1
                if patience_left <= 0:
                    print(f"Early stopping at epoch {epoch}")
                    break


    predictions = model.predict(inputs_test)
    test_loss = torch.mean(torch.square(predictions - outputs_test))
    print(f'Test Loss: {test_loss.item()}. Samples: {idx_test}')

    x = np.arange(0, len(idx_test))

    outputs_test = outputs_test.cpu().numpy() * y_std + y_mean
    predictions = predictions.cpu().numpy() * y_std + y_mean
    # for sample in outputs_test:
    #     print(f'Test samples: {sample}')
    plt.figure(figsize=(10, 6))
    plt.plot(x, outputs_test[:, 0], color='b', linestyle='--', label='True Phi7_Change')
    plt.plot(x, predictions[:, 0], color='b', linestyle='-', label='Predicted Phi7_Change')
    plt.plot(x, outputs_test[:, 1], color='r', linestyle='--', label='True Phi8_Change')
    plt.plot(x, predictions[:, 1], color='r', linestyle='-', label='Predicted Phi8_Change')
    plt.plot(x, outputs_test[:, 2], color='g', linestyle='--', label='True Phi9_Change')
    plt.plot(x, predictions[:, 2], color='g', linestyle='-', label='Predicted Phi9_Change')
    
    plt.gca().xaxis.set_major_locator(plt.MaxNLocator(integer=True))    
    plt.xlabel('Sample Index')
    plt.xticks(ticks=range(len(idx_test)),labels=idx_test + 1)
    plt.ylabel('Angle Change (Degrees)')
    plt.title('Angle Change Prediction')
    plt.legend(loc='lower left')
    plt.savefig('fdm_simulation.png')


    plt.figure(figsize=(10, 6))
    plt.plot(x, outputs_test[:, -1], color='m', linestyle='--', label='True Global_Max_Stress')
    plt.plot(x, predictions[:, -1], color='m', linestyle='-', label='Predicted Global_Max_Stress')
    plt.xlabel('Sample Index')
    plt.xticks(ticks=range(len(idx_test)),labels=idx_test + 1)
    plt.ylabel('Stress (MPa)')
    plt.title('Global Max Stress Prediction')
    plt.legend(loc='lower left')
    plt.savefig('fdm_stress_prediction.png')



    # MSE
    mse = np.mean((predictions - outputs_test) ** 2, axis=0)
    # print(f'Mean Squared Error for Phi1_Change: {mse[0]:.6f}, Phi2_Change: {mse[1]:.6f}, Phi3_Change: {mse[2]:.6f}, Phi7_Change: {mse[3]:.6f}, Phi8_Change: {mse[4]:.6f}, Phi9_Change: {mse[5]:.6f}, Global_Max_Stress: {mse[6]:.6f}')
    print(f'Mean Squared Error for Phi7_Change: {mse[0]:.6f}, Phi8_Change: {mse[1]:.6f}, Phi9_Change: {mse[2]:.6f}, Global_Max_Stress: {mse[3]:.6f}')


    # R 2 score
    ss_ress = np.sum((outputs_test - predictions) ** 2, axis=0)
    ss_tots = np.sum((outputs_test - np.mean(outputs_test, axis=0)) ** 2, axis=0)
    r2_scores = 1 - ss_ress / ss_tots
    # print(f'R² Score for Phi1_Change: {r2_scores[0]:.6f}, Phi2_Change: {r2_scores[1]:.6f}, Phi3_Change: {r2_scores[2]:.6f}, Phi7_Change: {r2_scores[3]:.6f}, Phi8_Change: {r2_scores[4]:.6f}, Phi9_Change: {r2_scores[5]:.6f}, Global_Max_Stress: {r2_scores[6]:.6f}')
    print(f'R² Score for Phi7_Change: {r2_scores[0]:.6f}, Phi8_Change: {r2_scores[1]:.6f}, Phi9_Change: {r2_scores[2]:.6f}, Global_Max_Stress: {r2_scores[3]:.6f}')
  
    # Error

    # Save the model
    model_save_path = './model_fdm_ckpt.pth'
    model_config = {'layer_sizes': layer_sizes,
                    'dropout_rate': dropout_rate
                    }
    checkpoint = {
        'model_state_dict': model.state_dict(),
        'model_config': model_config
    }
    torch.save(checkpoint, model_save_path)

def load_model(model_path):
    checkpoint = torch.load(model_path, map_location=DEVICE)
    model_config = checkpoint['model_config']
    model = NeuralNetwork(model_config['layer_sizes'], dropout_rate=model_config['dropout_rate'], activation=torch.nn.ReLU).to(DEVICE)
    model.load_state_dict(checkpoint['model_state_dict'])
    print(f"Model loaded from {model_path}")
    return model


def inverse_design(material_base, fiber, fiber_vf, y_target, n_restarts=5, epochs=100, use_lbfgs=True, model=None, data=None):
    if model is None:
        model = load_model('./model_fdm_ckpt.pth').to(torch.device('cpu'))

    if data is None:
        data = DataAdditiveManufacturing()
    mat_type = data.material_base_map.get(material_base, 0.0)
    fiber_type = data.fiber_type_map.get(fiber, 0.0)
    build_direction = data.build_direction_map.get("Vertical", 0.0)
    y_target_norm = torch.tensor(data.normalize_output(y_target), dtype=torch.float32)
    y_target_tensor = torch.tensor(y_target, dtype=torch.float32)
    input_mean = torch.tensor(data.input_mean)
    input_std = torch.tensor(data.input_std)
    output_mean = torch.tensor(data.output_mean)
    output_std = torch.tensor(data.output_std)
    
    weights = torch.tensor([1.0, 1.0, 1.0, 0.001], dtype=torch.float32)
    bounds = torch.tensor([[100., 300.], [50., 300.], [10., 200.]], dtype=torch.float32) # Extruder_Temp, Velocity, Bed_Temp
    best = {"loss": float('inf'), "input": None, "output": None}

    for restart in range(n_restarts):
        z = torch.randn(3, requires_grad=True)

        if use_lbfgs:
            optimizer = torch.optim.LBFGS([z], lr=0.1, max_iter=epochs, line_search_fn="strong_wolfe")
            steps = 1
        else:
            optimizer = torch.optim.Adam([z], lr=0.001)
            steps = epochs
        
        for step in range(steps):
            def closure():
                var =  bounds[:, 0] + (bounds[:, 1] - bounds[:, 0]) * torch.sigmoid(z)
                optimizer.zero_grad()
                input_raw = torch.cat([torch.tensor([mat_type, fiber_type, fiber_vf, build_direction]), var]).unsqueeze(0)
                input_norm = (input_raw - input_mean) / input_std
                output_pred = model(input_norm)
                output_pred = (output_pred * output_std) + output_mean
                loss = torch.sum(weights * (output_pred - y_target_tensor) ** 2)
                loss.backward()
                return loss

            if use_lbfgs:
                loss = optimizer.step(closure)
            else:
                loss = closure()
                optimizer.step()

            if (step + 1) % 200 == 0:
                print(f'Restart {restart + 1}, Step {step + 1}, Loss: {loss.item():.6f}, grad: {z.grad.norm().item():.6f}')
            
        with torch.no_grad():
            var =  bounds[:, 0] + (bounds[:, 1] - bounds[:, 0]) * torch.sigmoid(z)
            input_raw = torch.cat([torch.tensor([mat_type, fiber_type, fiber_vf, build_direction]), var])
            input_norm = (input_raw - input_mean) / input_std
            output_pred = model(input_norm)
            output_pred = data.denormalize_output(output_pred.numpy())
            final_loss = np.sum(weights.numpy() * (output_pred - y_target) ** 2).item()
            if final_loss < best["loss"]:
                best["loss"] = final_loss
                best["input"] = var.detach().cpu().numpy()
                best["output"] = output_pred

    return best


if __name__ == "__main__":
    set_seed(51)
    # dataset = DataAdditiveManufacturing()
    # inputs = dataset.get_input(normalize=False)
    # outputs = dataset.get_output(normalize=False)
    # kfold_ridge_baseline(inputs, outputs, k=5, alpha=1.0, seed=51)
    # main()
    
    
    best = inverse_design(material_base="HDPE", fiber="CF", fiber_vf=45.0,
                          y_target=np.array([-0.22, 0.11, -0.004, 185.2]), n_restarts=20, epochs=100, use_lbfgs=True)
    print("Best design found:")
    print(f"Extruder_Temp: {best['input'][0]:.2f}, Velocity: {best['input'][1]:.2f}, Bed_Temp: {best['input'][2]:.2f}")
    print(f"Predicted Outputs: Phi7_Change: {best['output'][0]:.4f}, Phi8_Change: {best['output'][1]:.4f}, Phi9_Change: {best['output'][2]:.4f}, Global_Max_Stress: {best['output'][3]:.4f}")