scripts/topology_optimization.py

import torch
import math
import numpy as np
from gpytorch.settings import cholesky_jitter
import matplotlib.pyplot as plt
import torch.optim
import matplotlib.pyplot as plt
from tqdm import tqdm
from onescience.utils.GP_TO import  plot_predictions_and_residuals,plot_loss_history,plot_density_and_velocity_fields
plt.rcParams['font.family'] = 'DejaVu Sans'
plt.rcParams['font.size'] = 15
plt.rcParams['figure.dpi'] = 150


checkpoints=[1000,10000,20000,30000,40000,50000]
lambdaa = 0.1

# Determine volume loss based on title
gamma_values = {
    'rugby': 0.9,
    'pipebend': 0.08 * torch.pi,
    'doublepipe': 1 / 3,
    'diffuser': 0.5
    }

def modified_sigmoid(x, alpha=12.0):
    """Compute the modified sigmoid function."""
    return 1 / (1 + torch.exp(-alpha * (x - 0.5)))


def infer_square_grid(num_points):
    side = int(math.sqrt(num_points))
    if side * side != num_points:
        raise ValueError(
            "Numerical differentiation expects a square collocation grid. "
            f"Got {num_points} points."
        )
    return side, side


def share_mean_module(model_list):
    model_ref = model_list[0]
    for model in model_list:
        model.mean_module_NN_All = model_ref.mean_module_NN_All
    return model_ref


def clear_cached_kernels(model_list):
    for model in model_list:
        model.k_xX = None
        model.chol_decomp = None


def predict_fields(model_list):
    model_ref = share_mean_module(model_list)
    collocation_x = model_ref.collocation_x
    m_col = model_ref.mean_module_NN_All(collocation_x)

    if model_ref.k_xX is None:
        for model in model_list:
            model.k_xX = model.covar_module(model.train_inputs[0], collocation_x).evaluate()

    if model_ref.chol_decomp is None:
        with cholesky_jitter(1e-7):
            for model in model_list:
                model.chol_decomp = model.covar_module(model.train_inputs[0]).cholesky()

    offsets = []
    for i, model in enumerate(model_list):
        target_offset = model.train_targets.unsqueeze(-1) - model.mean_module_NN_All(model.train_inputs[0])[:, i].unsqueeze(-1)
        offsets.append(model.chol_decomp._cholesky_solve(target_offset))

    u = (m_col[:, 0].unsqueeze(-1) + model_list[0].k_xX.t() @ offsets[0]).squeeze(-1)
    v = (m_col[:, 1].unsqueeze(-1) + model_list[1].k_xX.t() @ offsets[1]).squeeze(-1)
    p = (m_col[:, 2].unsqueeze(-1) + model_list[2].k_xX.t() @ offsets[2]).squeeze(-1)
    ro_tensor = (m_col[:, 3].unsqueeze(-1) + model_list[3].k_xX.t() @ offsets[3]).squeeze(-1)
    ro = modified_sigmoid(ro_tensor)
    return {"x": collocation_x, "u": u, "v": v, "p": p, "ro": ro}


def compute_dynamic_weights_ic(pde_loss, bc_loss, ic_loss, model):
    """
    Compute dynamic weights for loss functions based on gradients.
    """
    params_to_update = [param for param in model.parameters() if param.requires_grad]
    
    def compute_gradients(loss, scaling_factor):
        gradients = torch.autograd.grad(scaling_factor * loss, params_to_update, retain_graph=True, allow_unused=True)
        values = [p.reshape(-1).cpu().tolist() for p in gradients if p is not None]
        return torch.abs(torch.tensor([v for val in values for v in val]))
    
    delta_pde = compute_gradients(pde_loss, 1.0)
    delta_bc = compute_gradients(bc_loss, model.alpha)
    delta_ic = compute_gradients(ic_loss, model.beta)

    temp_bc = torch.max(delta_pde) / torch.mean(delta_bc)
    temp_ic = torch.max(delta_pde) / torch.mean(delta_ic)

    return (
        (1.0 - lambdaa) * model.alpha + lambdaa * temp_bc,
        (1.0 - lambdaa) * model.beta + lambdaa * temp_ic
    )

def compute_fvm_residuals_higher_order(u, v, p, f_x, f_y, nu, dx, dy):
    # Check the device of the input tensors and move other tensors to the same device
    device = u.device
    
    # Extend the fields with ghost cells to handle boundary conditions
    u_ext = torch.zeros((u.shape[0] + 4, u.shape[1] + 4), device=device)
    v_ext = torch.zeros((v.shape[0] + 4, v.shape[1] + 4), device=device)
    p_ext = torch.zeros((p.shape[0] + 4, p.shape[1] + 4), device=device)
    
    # Copy the interior
    u_ext[2:-2, 2:-2] = u
    v_ext[2:-2, 2:-2] = v
    p_ext[2:-2, 2:-2] = p
    
    # Apply ghost cells for no-slip boundary conditions
    u_ext[:2, :] = u_ext[2:4, :]
    u_ext[-2:, :] = u_ext[-4:-2, :]
    u_ext[:, :2] = u_ext[:, 2:4]
    u_ext[:, -2:] = u_ext[:, -4:-2]
    
    v_ext[:2, :] = v_ext[2:4, :]
    v_ext[-2:, :] = v_ext[-4:-2, :]
    v_ext[:, :2] = v_ext[:, 2:4]
    v_ext[:, -2:] = v_ext[:, -4:-2]
    
    # Laplacian of u and v (fourth-order central differences for second derivatives)
    u_xx = (-u_ext[4:, 2:-2] + 16 * u_ext[3:-1, 2:-2] - 30 * u_ext[2:-2, 2:-2] + 16 * u_ext[1:-3, 2:-2] - u_ext[:-4, 2:-2]) / (12 * dx**2)
    u_yy = (-u_ext[2:-2, 4:] + 16 * u_ext[2:-2, 3:-1] - 30 * u_ext[2:-2, 2:-2] + 16 * u_ext[2:-2, 1:-3] - u_ext[2:-2, :-4]) / (12 * dy**2)
    laplacian_u = u_xx + u_yy
    
    v_xx = (-v_ext[4:, 2:-2] + 16 * v_ext[3:-1, 2:-2] - 30 * v_ext[2:-2, 2:-2] + 16 * v_ext[1:-3, 2:-2] - v_ext[:-4, 2:-2]) / (12 * dx**2)
    v_yy = (-v_ext[2:-2, 4:] + 16 * v_ext[2:-2, 3:-1] - 30 * v_ext[2:-2, 2:-2] + 16 * v_ext[2:-2, 1:-3] - v_ext[2:-2, :-4]) / (12 * dy**2)
    laplacian_v = v_xx + v_yy
    
    # Gradients of p using fourth-order central difference
    grad_p_x = (-p_ext[4:, 2:-2] + 8 * p_ext[3:-1, 2:-2] - 8 * p_ext[1:-3, 2:-2] + p_ext[:-4, 2:-2]) / (12 * dx)
    grad_p_y = (-p_ext[2:-2, 4:] + 8 * p_ext[2:-2, 3:-1] - 8 * p_ext[2:-2, 1:-3] + p_ext[2:-2, :-4]) / (12 * dy)
    
    # Momentum equation residuals
    residual_u = -nu * laplacian_u + grad_p_x + f_x
    residual_v = -nu * laplacian_v + grad_p_y + f_y
    
    # Mass conservation (divergence of velocity) using central differences
    div_u = (u_ext[2:-2, 2:-2] - u_ext[1:-3, 2:-2]) / dx
    div_v = (v_ext[2:-2, 2:-2] - v_ext[2:-2, 1:-3]) / dy
    
    u_y = (u_ext[2:-2, 2:-2] - u_ext[2:-2, 1:-3]) / dy
    v_x = (v_ext[2:-2, 2:-2] - v_ext[1:-3, 2:-2]) / dx
    residual_mass = div_u + div_v
    
    return residual_u.reshape(-1), residual_v.reshape(-1), residual_mass.reshape(-1), div_u.reshape(-1), div_v.reshape(-1), u_y.reshape(-1), v_x.reshape(-1)



def compute_autograd_derivatives(model, z_column, collocation_x, grad_order=1):
    """
    Compute first and second-order derivatives using PyTorch autograd.
    :param model: Model to train
    :param z_column: Column of `z_all` for differentiation
    :param collocation_x: Input coordinates
    :param grad_order: Order of gradients (1 for first, 2 for second)
    :return: First and second-order derivatives as tensors
    """
    model.train()
    grad_1 = torch.autograd.grad(z_column, collocation_x, torch.ones_like(z_column), create_graph=True)[0]
    if grad_order > 1:
        grad_2_x = torch.autograd.grad(grad_1[:, 0], collocation_x, torch.ones_like(grad_1[:, 0]), create_graph=True)[0][:, 0]
        grad_2_y = torch.autograd.grad(grad_1[:, 1], collocation_x, torch.ones_like(grad_1[:, 1]), create_graph=True)[0][:, 1]
        return grad_1[:, 0], grad_1[:, 1], grad_2_x, grad_2_y
    return grad_1[:, 0], grad_1[:, 1], None, None

w_1,w_2,w_3, w_4, w_5 =0.01,0.01,1e2,1e5,1e4

def loss_volume(y, gamma=0.5):
    """
    Compute the volume loss as the squared difference between the mean of y and gamma.
    """
    mean_y = torch.mean(y)
    return torch.square(mean_y - gamma)


def dissipated_power(u, u_x, v_y, u_y, v_x):
    """
    Compute the total dissipated power based on the input velocity gradients and displacements.
    """
    # First part of the dissipated power
    p1 = (u_x**2 + v_y**2 + u_y**2 + v_x**2).sum(dim=1, keepdim=True)
    
    # Second part of the dissipated power
    u2 = (u[:, :2]**2).sum(dim=1, keepdim=True)
    p2 = alpha(u[:, 3:]) * u2
    
    return 0.5 * (p1 + p2)


def alpha(rho):
    """
    Compute the alpha parameter based on rho.
    """
    alpha_max = 2.5e4
    alpha_min = 2.5e-4
    q = 0.1
    return alpha_max + (alpha_min - alpha_max) * rho * (1 + q) / (rho + q)

def calculate_loss_multioutput(
    model_list,
    iteration,
    diff_method='Numerical',
    title="default",
    problem="doublepipe",
    checkpoint_steps=None,
    plot_outputs=True,
):
    """
    Calculate loss for multi-output models based on PDE, BC, and IC residuals.
    """
    # Clone collocation points
    collocation_x = model_list[0].collocation_x.clone()

    # Compute mean values at collocation points
    m_col = model_list[0].mean_module_NN_All(collocation_x)

    # Evaluate covariance matrices if not already done
    if model_list[0].k_xX is None:
        for i, model in enumerate(model_list):
            model.k_xX = model.covar_module(model.train_inputs[0], collocation_x).evaluate()

    # Perform Cholesky decomposition if not already done
    if model_list[0].chol_decomp is None:
        with cholesky_jitter(1e-7):
            for model in model_list:
                model.chol_decomp = model.covar_module(model.train_inputs[0]).cholesky()

    # Solve for offsets using Cholesky decomposition
    offsets = []
    for i, model in enumerate(model_list):
        target_offset = model.train_targets.unsqueeze(-1) - model.mean_module_NN_All(model.train_inputs[0])[:, i].unsqueeze(-1)
        offsets.append(model.chol_decomp._cholesky_solve(target_offset))

    # Compute values for velocity, pressure, and density
    ro_tensor = (m_col[:, 3].unsqueeze(-1) + model_list[3].k_xX.t() @ offsets[3]).squeeze(-1)
    ro = modified_sigmoid(ro_tensor)

    u = (m_col[:, 0].unsqueeze(-1) + model_list[0].k_xX.t() @ offsets[0]).squeeze(-1)
    v = (m_col[:, 1].unsqueeze(-1) + model_list[1].k_xX.t() @ offsets[1]).squeeze(-1)
    p = (m_col[:, 2].unsqueeze(-1) + model_list[2].k_xX.t() @ offsets[2]).squeeze(-1)

    z_all = torch.cat((u.unsqueeze(1), v.unsqueeze(1), p.unsqueeze(1), ro.unsqueeze(1)), dim=1)
    f = alpha(z_all[:, 3:]) * z_all[:, :2]
    fx, fy = f[:, :1], f[:, 1:]

    # Compute residuals
    if diff_method == 'Numerical':
        nx, ny = infer_square_grid(collocation_x.shape[0])
        dx, dy = 1.0 / nx, 1.0 / ny
        u, v, p = [arr.reshape(nx, ny) for arr in (u, v, p)]
        fx, fy = [arr.reshape(nx, ny) for arr in (fx, fy)]
        residuals = compute_fvm_residuals_higher_order(u, v, p, fx, fy, nu=1, dx=dx, dy=dy)
        residual_pde1, residual_pde2, residual_pde3, u_x, v_y, u_y, v_x = residuals
        residual_pde1, residual_pde2, residual_pde3 = [res * w for res, w in zip(residuals[:3], (w_1, w_2, w_3))]
        u, v, p = [arr.reshape(-1) for arr in (u, v, p)]
        fx, fy = [arr.reshape(-1) for arr in (fx, fy)]
    else:
        u_x, u_y, u_xx, u_yy = compute_autograd_derivatives(model_list[0], z_all[:, 0], collocation_x, grad_order=2)
        v_x, v_y, v_xx, v_yy = compute_autograd_derivatives(model_list[1], z_all[:, 1], collocation_x, grad_order=2)
        p_x, p_y, _, _ = compute_autograd_derivatives(model_list[2], z_all[:, 2], collocation_x, grad_order=1)
        residual_pde1 = (- (u_xx + u_yy) + p_x + fx[:, 0]) * w_1
        residual_pde2 = (- (v_xx + v_yy) + p_y + fy[:, 0]) * w_2
        residual_pde3 = (u_x + v_y) * w_3

    # Plot fields at checkpoints
    active_checkpoints = checkpoints if checkpoint_steps is None else checkpoint_steps
    if plot_outputs and iteration in active_checkpoints:
        plot_density_and_velocity_fields(u, v, p, ro, collocation_x, iteration, problem)
        plot_predictions_and_residuals(u, v, p, ro, collocation_x, iteration, residual_pde1, residual_pde2, residual_pde3, 1, 1, 1, problem)
    # Calculate losses
    loss_pde1 = torch.mean(residual_pde1**2)
    loss_pde2 = torch.mean(residual_pde2**2)
    loss_pde3 = torch.mean(residual_pde3**2)
    dp_loss = dissipated_power(z_all, u_x.reshape(-1, 1), v_y.reshape(-1, 1), u_y.reshape(-1, 1), v_x.reshape(-1, 1))

    gamma = next((val for key, val in gamma_values.items() if key in title), 0.5)
    vol_loss = loss_volume(z_all[:, 3:], gamma=gamma)* w_4

    # Return final losses
    return loss_pde1, loss_pde2, loss_pde3, torch.sum(dp_loss), vol_loss 




def find_TO(
    model_list,
    lr_default: float = 0.01,
    num_iter: int = 500,
    title: str = 'default',
    problem: str = 'doublepipe',   # 设置默认值
    diff_method: str = 'Numerical',
    checkpoint_steps=None,
    plot_outputs=True,
) -> float:
    """
    Train models to minimize total loss using dynamic weights and record progress.

    Args:
        model_list: List of models to be trained.
        lr_default: Default learning rate for the optimizer.
        num_iter: Number of training iterations.
        title: Title for specific settings.
        diff_method: Differentiation method ('Numerical' or 'Autograd').

    Returns:
        loss history
    """
    # Use the first model as a reference for shared NN
    model_ref = share_mean_module(model_list)
    for model in model_list:
        model.train()

    # Initialize variables
    loss_total, loss_hist, loss_hist_total = [], [], []
    loss_pde1_hist, loss_pde2_hist, loss_pde3_hist = [], [], []
    loss_dp_hist, vol_loss_hist = [], []
    scaled_loss_pde1_hist, scaled_loss_pde2_hist = [], []
    scaled_loss_pde3_hist, scaled_loss_dp_hist, scaled_vol_loss_hist = [], [], []
    weights = {'alpha': [], 'beta': [], 'mu_p': []}
    result_loss_thetas = {'largest_eigval_hist': [], 'condition_number_hist': [], 'Hessian': [], 'grads': []}
    GP_NN_hist, NN_hist, time_hist = [], [], []
    dynamic_weights = True
    sigma_1, beta_1, alph_1 = 10, 10, 0.1
    f_inc, mu_F = math.inf, 1

    # Initialize model-specific parameters
    model1 = model_list[0]
    model1.alpha, model1.beta, model1.theta = 1, 1, 1

    # Set up optimizer and scheduler
    optimizer = torch.optim.Adam(model_ref.parameters(), lr=lr_default)
    milestones = sorted({int(x) for x in np.linspace(0, max(num_iter, 1), 4).tolist() if int(x) > 0})
    scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer,
        milestones=milestones,
        gamma=0.75
    )

    # Training loop
    with tqdm(range(num_iter + 1), desc='Epoch', position=0, leave=True) as pbar:
        for j in pbar:
            optimizer.zero_grad()

            # Calculate losses
            loss_pde1, loss_pde2, loss_pde3, loss_dp, vol_loss = calculate_loss_multioutput(
                model_list,
                j,
                diff_method=diff_method,
                title=title,
                problem=problem,
                checkpoint_steps=checkpoint_steps,
                plot_outputs=plot_outputs,
            )
            loss_pde = loss_pde1 + loss_pde3

            # Dynamic weight adjustments
            if dynamic_weights:
                alpha, beta = compute_dynamic_weights_ic(
                    loss_pde1 + loss_pde2, loss_pde3, vol_loss, model1
                )
                if all(torch.is_tensor(val) and not (torch.isnan(val).any() or torch.isinf(val).any()) for val in [alpha, beta]):
                    model1.alpha, model1.beta = alpha, beta

                weights['alpha'].append(alpha.detach().cpu().item())
                weights['beta'].append(beta.detach().cpu().item())
                weights['mu_p'].append(mu_F)

                loss = loss_dp + mu_F * (
                    loss_pde1 + loss_pde2 + model1.alpha * loss_pde3 + model1.beta * vol_loss
                )

                if (j + 1) % 50 == 0:
                    mu_F = min(mu_F * 1.05, 5e2)
            else:
                alph_1 = 1e-6 / (j + 1)
                beta_1, sigma_1 = 2 * (j + 1), 5 * (j + 1)
                loss = sigma_1 * loss_pde + alph_1 * loss_dp + beta_1 * vol_loss

            # Record loss history
            loss_total.append(loss.item())
            loss_pde1_hist.append((1/w_1) * loss_pde1.item())
            loss_pde2_hist.append((1/w_2) * loss_pde2.item())
            loss_pde3_hist.append((1/w_3) * loss_pde3.item())
            loss_dp_hist.append((1/w_5) * loss_dp.item())
            vol_loss_hist.append((1/w_4) * vol_loss.item())
            scaled_loss_pde1_hist.append(loss_pde1.item())
            scaled_loss_pde2_hist.append(loss_pde2.item())
            scaled_loss_pde3_hist.append(model1.alpha * loss_pde3.item())
            scaled_loss_dp_hist.append(loss_dp.item())
            scaled_vol_loss_hist.append(model1.beta * vol_loss.item())

            # Plot loss history at checkpoints
            active_checkpoints = checkpoints if checkpoint_steps is None else checkpoint_steps
            if plot_outputs and j in active_checkpoints:
                plot_loss_history(
                    loss_total,
                    [loss_pde1_hist, scaled_loss_pde1_hist],
                    [loss_pde2_hist, scaled_loss_pde2_hist],
                    [loss_pde3_hist, scaled_loss_pde3_hist],
                    [loss_dp_hist, scaled_loss_dp_hist],
                    [vol_loss_hist, scaled_vol_loss_hist],
                    j,
                    problem
                )

            # Backpropagation and optimizer step
            loss.backward(retain_graph=True)
            optimizer.step()
            scheduler.step()

            # Update progress description
            pbar.set_postfix(loss=f"{loss.item():.6f}")

            loss_hist.append(loss.item())

    loss_hist_total = loss_hist

    return loss_hist_total