docker run -it -u root -d --net=host \
--privileged \
--ipc=host \
--device=/dev/davinci_manager \
--device=/dev/devmm_svm \
--device=/dev/hisi_hdc \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/sbin:/usr/local/sbin \
-v /usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool \
--name pid-gan \
quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11 \
/bin/bash

mkdir -p /root/gan/
cd /root/gan/
git clone https://github.com/arkadaw9/PID-GAN

apt update
apt install libxcb1 libx11-xcb1
apt install -y libgl1-mesa-glx
apt install -y libglib2.0-0 libsm6 libxext6 libxrender-dev libgomp1
apt-get install -y zlib1g-dev build-essential

pip install torch==2.10.0 torch-npu==2.10.0
pip install scipy numpy pyDOE2 matplotlib seaborn

git apply ascend_npu_adaptation.patch

1、数值稳定性修复：epsilon 1e-8 → 1e-6
根因： NPU FP16精度下，1e-8 低于 float16 最小精度（~6e-8），torch.log(sigmoid + 1e-8) 可能产生 log(0) 导致 NaN。
实际影响评估： 当前训练使用纯 FP32（.float()），epsilon 1e-6 vs 1e-8 对判别器损失梯度的影响很小。

2、inplace操作修复：LeakyReLU inplace=True → inplace=False（1个文件，4处）
根因： NPU上inplace操作可能与autograd计算图冲突，create_graph=True 的二阶导数场景下可能导致梯度计算错误。
实际影响评估： inplace=True vs inplace=False 的数学运算完全相同，仅内存操作方式不同（原地修改 vs 新建张量），不影响精度和收敛。

3、Generator 内循环 return Bug 修复（关键修改）
根因： train_generator() 中 for gen_epoch in range(5): ... return 的 return 语句位于循环内部，导致生成器每次只更新1步就退出循环，而非论文设计的5步更新。这是原作者代码的 Bug。
影响： 这是导致 NPU 训练 Err u/Err h 偏高的主要原因。修复后与论文结果统一。

import argparse
import os
import sys
import warnings

import torch_npu
from torch_npu.contrib import transfer_to_npu

import numpy as np
import torch
import scipy.io
from pyDOE2 import lhs

warnings.filterwarnings('ignore')

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(BASE_DIR, 'Scripts'))
sys.path.insert(0, os.path.join(BASE_DIR, 'Utilities'))
sys.path.insert(0, os.path.join(BASE_DIR, 'datasets'))

DATASET_DIR = os.path.join(BASE_DIR, 'datasets')


def save_checkpoint(model, method, path):
    state = {}
    if method in ['PID-GAN', 'PIG-GAN', 'cGAN']:
        state['G_state_dict'] = model.G.state_dict()
        state['D_state_dict'] = model.D.state_dict()
        state['Q_state_dict'] = model.Q.state_dict()
        if hasattr(model, 'kG'):
            state['kG_state_dict'] = model.kG.state_dict()
    elif method in ['PINN', 'APINN']:
        net = model.net if hasattr(model, 'net') else model.pinn
        state['net_state_dict'] = net.state_dict()
    if hasattr(model, 'Xmean'):
        state['Xmean'] = model.Xmean
        state['Xstd'] = model.Xstd
    if hasattr(model, 'lambda_val'):
        state['lambda_val'] = model.lambda_val
    if hasattr(model, 'lambda_q'):
        state['lambda_q'] = model.lambda_q
    torch.save(state, path)


def get_device(use_npu, device_id):
    if use_npu and torch.npu.is_available():
        device = torch.device(f'npu:{device_id}')
        torch.npu.set_device(device)
        print(f'Using NPU device: {device}, count: {torch.npu.device_count()}', flush=True)
    elif torch.cuda.is_available():
        device = torch.device(f'cuda:{device_id}')
        print(f'Using CUDA device: {device}', flush=True)
    else:
        device = torch.device('cpu')
        print('Using CPU', flush=True)
    return device


def load_burgers_data(N_u=100, N_i=50, N_f=10000):
    data = scipy.io.loadmat(os.path.join(DATASET_DIR, 'burgers_shock.mat'))
    t = data['t'].flatten()[:, None]
    x = data['x'].flatten()[:, None]
    Exact = np.real(data['usol']).T
    X, T = np.meshgrid(x, t)
    X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
    u_star = Exact.flatten()[:, None]
    lb = X_star.min(0)
    ub = X_star.max(0)
    xx1 = np.hstack((X[0:1, :].T, T[0:1, :].T))
    uu1 = Exact[0:1, :].T
    xx2 = np.hstack((X[:, 0:1], T[:, 0:1]))
    uu2 = Exact[:, 0:1]
    xx3 = np.hstack((X[:,-1:], T[:,-1:]))
    uu3 = Exact[:,-1:]
    X_u_train = np.vstack([xx2, xx3])
    u_train = np.vstack([uu2, uu3])
    X_f_train = lb + (ub - lb) * lhs(2, N_f)
    X_f_train = np.vstack([X_f_train, X_u_train, xx1])
    idx = np.random.choice(X_u_train.shape[0], N_u, replace=False)
    X_u_train = X_u_train[idx, :]
    u_train = u_train[idx, :]
    idx = np.random.choice(xx1.shape[0], N_i, replace=False)
    X_i_train = xx1[idx, :]
    u_i_train = uu1[idx, :]
    X_u_train = np.vstack([X_u_train, X_i_train])
    u_train = np.vstack([u_train, u_i_train])
    return X_u_train, u_train, X_f_train, X_star, u_star


def load_schrodinger_data(N0=50, N_b=50, N_f=20000):
    data = scipy.io.loadmat(os.path.join(DATASET_DIR, 'NLS.mat'))
    t = data['tt'].flatten()[:, None]
    x = data['x'].flatten()[:, None]
    Exact = data['uu']
    Exact_u = np.real(Exact)
    Exact_v = np.imag(Exact)
    Exact_h = np.sqrt(Exact_u**2 + Exact_v**2)

    X, T = np.meshgrid(x, t)
    X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
    h_star = Exact_h.T.flatten()[:, None]

    lb = np.array([-5.0, 0.0])
    ub = np.array([5.0, np.pi / 2])

    idx_x = np.random.choice(x.shape[0], N0, replace=False)
    x0 = x[idx_x, :]
    u0 = Exact_u[idx_x, 0:1]
    v0 = Exact_v[idx_x, 0:1]

    idx_t = np.random.choice(t.shape[0], N_b, replace=False)
    tb = t[idx_t, :]

    X_f = lb + (ub - lb) * lhs(2, N_f)

    X0 = np.concatenate((x0, 0 * x0), 1)
    Y0 = np.concatenate((u0, v0), 1)
    X_lb = np.concatenate((0 * tb + lb[0], tb), 1)
    X_ub = np.concatenate((0 * tb + ub[0], tb), 1)

    return X0, Y0, X_f, X_lb, X_ub, X_star, h_star


def load_darcy_data(N_u=200, N_b=100, N_f=10000):
    data = np.load(os.path.join(DATASET_DIR, 'nonlinear2d_data.npz'))
    X = data['X']
    U = data['u']
    K = data['k']

    L1 = 10.
    L2 = 10.

    x1_b1 = np.zeros(N_b)[:, None]
    x2_b1 = L2 * np.random.random(N_b)[:, None]
    X_b1 = np.hstack((x1_b1, x2_b1))

    x1_b2 = L1 * np.random.random(N_b)[:, None]
    x2_b2 = np.zeros(N_b)[:, None]
    X_b2 = np.hstack((x1_b2, x2_b2))

    x1_b3 = L1 * np.ones(N_b)[:, None]
    x2_b3 = L2 * np.random.random(N_b)[:, None]
    X_b3 = np.hstack((x1_b3, x2_b3))

    x1_b4 = L1 * np.random.random(N_b)[:, None]
    x2_b4 = L2 * np.ones(N_b)[:, None]
    X_b4 = np.hstack((x1_b4, x2_b4))

    X_b = np.hstack((X_b1, X_b2, X_b3, X_b4))

    X1_f = L1 * np.random.random(N_f)[:, None]
    X2_f = L2 * np.random.random(N_f)[:, None]
    X_f = np.hstack((X1_f, X2_f))

    X_u = np.zeros((N_u, 2))
    Y_u = np.zeros((N_u, 1))
    idx_u = np.random.choice(X.shape[0], N_u, replace=False)
    for i in range(N_u):
        X_u[i, :] = X[idx_u[i], :]
        Y_u[i, :] = U[idx_u[i]]

    lb = np.array([0.0, 0.0])
    ub = np.array([10.0, 10.0])
    lbb = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    ubb = np.array([10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0])
    X_u = (X_u - lb) - 0.5 * (ub - lb)
    X_b = (X_b - lbb) - 0.5 * (ubb - lbb)
    X_f = (X_f - lb) - 0.5 * (ub - lb)

    return X_u, Y_u, X_b, X_f, N_u, U


def load_tossing_data():
    data = np.loadtxt(os.path.join(DATASET_DIR, 'tossing_trajectories.txt'))
    train_x = data[:200, :6]
    train_y = data[:200, 6:]
    test_x = data[200:, :6]
    test_y = data[200:, 6:]
    return train_x, train_y, test_x, test_y


def train_burgers(method, device, num_epochs, lambda_val, noise, hid_dim, num_layers):
    print(f'[Config] method={method}, num_epochs={num_epochs}, lambda_val={lambda_val}, noise={noise}, hid_dim={hid_dim}, num_layers={num_layers}, device={device}', flush=True)
    X_u_train, u_train, X_f_train, X_star, u_star = load_burgers_data()
    if method in ['PID-GAN', 'PIG-GAN']:
        from models_pde import Generator, Discriminator, Q_Net
        D = Discriminator(in_dim=3 if method == 'PIG-GAN' else 4, out_dim=1, hid_dim=hid_dim, num_layers=2).to(device)
        G = Generator(in_dim=3, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
        Q = Q_Net(in_dim=3, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
        if method == 'PID-GAN':
            sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Burgers', 'PID-GAN'))
            from pid import Burgers_PID
            model = Burgers_PID(X_u_train, u_train, X_f_train, X_star, u_star, G, D, Q, device, num_epochs, lambda_val, noise)
        else:
            sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Burgers', 'PIG-GAN'))
            from pig import Burgers_PIG
            model = Burgers_PIG(X_u_train, u_train, X_f_train, X_star, u_star, G, D, Q, device, num_epochs, lambda_val, noise)
    elif method == 'PINN':
        from models_pde import Net
        net = Net(in_dim=2, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
        sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Burgers', 'PINN'))
        from pinn import Burgers_PINN
        model = Burgers_PINN(X_u_train, u_train, X_f_train, X_star, u_star, net, device, num_epochs, lambda_val, noise)
    elif method == 'APINN':
        from models_pde import Net
        net = Net(in_dim=2, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
        sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Burgers', 'APINN'))
        from apinn import Burgers_APINN
        model = Burgers_APINN(X_u_train, u_train, X_f_train, X_star, u_star, net, device, num_epochs, lambda_val, noise)
    else:
        raise ValueError(f'Unknown method: {method}')
    model.train()
    print(f'Training completed: Burgers / {method} / {num_epochs} epochs on {device}', flush=True)

    checkpoint_dir = os.path.join(BASE_DIR, 'checkpoints')
    os.makedirs(checkpoint_dir, exist_ok=True)
    checkpoint_path = os.path.join(checkpoint_dir, f'burgers_{method}_epochs{num_epochs}_noise{noise}.pth')
    save_checkpoint(model, method, checkpoint_path)
    print(f'Checkpoint saved: {checkpoint_path}', flush=True)
    return model


def train_schrodinger(method, device, num_epochs, lambda_val, noise, hid_dim, num_layers):
    print(f'[Config] method={method}, num_epochs={num_epochs}, lambda_val={lambda_val}, noise={noise}, hid_dim={hid_dim}, num_layers={num_layers}, device={device}', flush=True)
    X_u_train, y_train, X_f_train, x_lb, x_ub, X_star, h_star = load_schrodinger_data()
    if method in ['PID-GAN', 'PIG-GAN']:
        from models_pde import Generator, Discriminator, Q_Net
        if method == 'PID-GAN':
            d_in_dim = 6
        else:
            d_in_dim = 4
        D = Discriminator(in_dim=d_in_dim, out_dim=1, hid_dim=hid_dim, num_layers=2).to(device)
        G = Generator(in_dim=3, out_dim=2, hid_dim=hid_dim, num_layers=num_layers).to(device)
        Q = Q_Net(in_dim=4, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
        if method == 'PID-GAN':
            sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Schrodinger', 'PID-GAN'))
            from pid import Schrodinger_PID
            model = Schrodinger_PID(X_u_train, y_train, X_f_train, x_lb, x_ub, X_star, h_star, G, D, Q, device, num_epochs, lambda_val, noise)
        else:
            sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Schrodinger', 'PIG-GAN'))
            from pig import Schrodinger_PIG
            model = Schrodinger_PIG(X_u_train, y_train, X_f_train, x_lb, x_ub, X_star, h_star, G, D, Q, device, num_epochs, lambda_val, noise)
    elif method == 'PINN':
        from models_pde import Net
        net = Net(in_dim=2, out_dim=2, hid_dim=hid_dim, num_layers=num_layers).to(device)
        sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Schrodinger', 'PINN'))
        from pinn import Schrodinger_PINN
        model = Schrodinger_PINN(X_u_train, y_train, X_f_train, x_lb, x_ub, X_star, h_star, net, device, num_epochs, lambda_val, noise)
    elif method == 'APINN':
        from models_pde import Net
        net = Net(in_dim=2, out_dim=2, hid_dim=hid_dim, num_layers=num_layers).to(device)
        sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Schrodinger', 'APINN'))
        from apinn import Schrodinger_APINN
        model = Schrodinger_APINN(X_u_train, y_train, X_f_train, x_lb, x_ub, X_star, h_star, net, device, num_epochs, noise)
    else:
        raise ValueError(f'Unknown method: {method}')
    model.train()
    print(f'Training completed: Schrodinger / {method} / {num_epochs} epochs on {device}', flush=True)
    checkpoint_dir = os.path.join(BASE_DIR, 'checkpoints')
    os.makedirs(checkpoint_dir, exist_ok=True)
    checkpoint_path = os.path.join(checkpoint_dir, f'schrodinger_{method}_epochs{num_epochs}_noise{noise}.pth')
    save_checkpoint(model, method, checkpoint_path)
    print(f'Checkpoint saved: {checkpoint_path}', flush=True)
    return model


def train_darcy(method, device, num_epochs, lambdas, noise, hid_dim, num_layers):
    print(f'[Config] method={method}, num_epochs={num_epochs}, lambdas={lambdas}, noise={noise}, hid_dim={hid_dim}, num_layers={num_layers}, device={device}', flush=True)
    x_u, y_u, x_b, x_f, N_u, u_star = load_darcy_data()
    if method in ['PID-GAN', 'PIG-GAN']:
        from models_pde import Generator, Discriminator, Q_Net
        d_in_dim = 4
        D = Discriminator(in_dim=d_in_dim, out_dim=1, hid_dim=hid_dim, num_layers=3).to(device)
        G = Generator(in_dim=4, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
        Q = Q_Net(in_dim=3, out_dim=2, hid_dim=hid_dim, num_layers=num_layers).to(device)
        kG = Generator(in_dim=1, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
        if method == 'PID-GAN':
            sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Darcy', 'PID-GAN'))
            from pid import Darcy_PID
            model = Darcy_PID(x_u, y_u, x_b, x_f, N_u, G, kG, D, Q, device, num_epochs, lambdas, noise)
        else:
            sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Darcy', 'PIG-GAN'))
            from pig import Darcy_PIG
            model = Darcy_PIG(x_u, y_u, x_b, x_f, N_u, G, kG, D, Q, device, num_epochs, lambdas, noise)
    else:
        raise ValueError(f'Darcy only supports PID-GAN and PIG-GAN, got: {method}')
    model.train()
    print(f'Training completed: Darcy / {method} / {num_epochs} epochs on {device}', flush=True)
    checkpoint_dir = os.path.join(BASE_DIR, 'checkpoints')
    os.makedirs(checkpoint_dir, exist_ok=True)
    checkpoint_path = os.path.join(checkpoint_dir, f'darcy_{method}_epochs{num_epochs}_noise{noise}.pth')
    save_checkpoint(model, method, checkpoint_path)
    print(f'Checkpoint saved: {checkpoint_path}', flush=True)
    return model


def train_tossing(method, device, num_epochs, lambda_val, noise_dim, hid_dim, num_layers):
    print(f'[Config] method={method}, num_epochs={num_epochs}, lambda_val={lambda_val}, noise_dim={noise_dim}, hid_dim={hid_dim}, num_layers={num_layers}, device={device}', flush=True)
    train_x, train_y, test_x, test_y = load_tossing_data()
    if method in ['PID-GAN', 'PIG-GAN', 'cGAN']:
        from models_imperfect import Generator, Discriminator, Q_Net
        if method == 'PID-GAN':
            d_in_dim = train_x.shape[1] + train_y.shape[1] + 1
        elif method == 'PIG-GAN':
            d_in_dim = train_x.shape[1] + train_y.shape[1]
        else:
            d_in_dim = train_x.shape[1] + train_y.shape[1]
        D = Discriminator(in_dim=d_in_dim, out_dim=1, hid_dim=hid_dim, num_layers=2).to(device)
        G = Generator(in_dim=train_x.shape[1] + noise_dim, out_dim=train_y.shape[1], hid_dim=hid_dim, num_layers=num_layers).to(device)
        Q = Q_Net(in_dim=train_x.shape[1] + train_y.shape[1], out_dim=noise_dim, hid_dim=hid_dim, num_layers=num_layers).to(device)
        if method == 'PID-GAN':
            sys.path.insert(0, os.path.join(BASE_DIR, 'Imperfect_Physics', 'Tossing', 'PID'))
            from pid import Tossing_PID
            model = Tossing_PID(train_x, train_y, test_x, test_y, G, D, Q, device, num_epochs, lambda_val, noise_dim)
        elif method == 'PIG-GAN':
            sys.path.insert(0, os.path.join(BASE_DIR, 'Imperfect_Physics', 'Tossing', 'PIG'))
            from pig import Tossing_PIG
            model = Tossing_PIG(train_x, train_y, test_x, test_y, G, D, Q, device, num_epochs, lambda_val, noise_dim)
        else:
            sys.path.insert(0, os.path.join(BASE_DIR, 'Imperfect_Physics', 'Tossing', 'cGAN'))
            from cGAN import Tossing_cGAN
            model = Tossing_cGAN(train_x, train_y, test_x, test_y, G, D, Q, device, num_epochs, lambda_val, noise_dim)
    elif method == 'PINN':
        from models_imperfect import PGNN
        net = PGNN(in_dim=train_x.shape[1], out_dim=train_y.shape[1], hid_dim=hid_dim, num_layers=num_layers).to(device)
        sys.path.insert(0, os.path.join(BASE_DIR, 'Imperfect_Physics', 'Tossing', 'PINN'))
        from pinn import Tossing_PINN
        model = Tossing_PINN(train_x, train_y, test_x, test_y, net, device, num_epochs, lambda_val)
    elif method == 'APINN':
        from models_imperfect import PGNN
        net = PGNN(in_dim=train_x.shape[1], out_dim=train_y.shape[1], hid_dim=hid_dim, num_layers=num_layers).to(device)
        sys.path.insert(0, os.path.join(BASE_DIR, 'Imperfect_Physics', 'Tossing', 'APINN'))
        from apinn import Tossing_APINN
        model = Tossing_APINN(train_x, train_y, test_x, test_y, net, device, num_epochs, lambda_val)
    else:
        raise ValueError(f'Unknown method: {method}')
    model.train()
    print(f'Training completed: Tossing / {method} / {num_epochs} epochs on {device}', flush=True)
    checkpoint_dir = os.path.join(BASE_DIR, 'checkpoints')
    os.makedirs(checkpoint_dir, exist_ok=True)
    checkpoint_path = os.path.join(checkpoint_dir, f'tossing_{method}_epochs{num_epochs}.pth')
    save_checkpoint(model, method, checkpoint_path)
    print(f'Checkpoint saved: {checkpoint_path}', flush=True)
    return model


def main():
    parser = argparse.ArgumentParser(description='PID-GAN Training on Ascend NPU')
    parser.add_argument('--problem', type=str, required=True,
                        choices=['burgers', 'schrodinger', 'darcy', 'tossing'],
                        help='PDE problem or imperfect physics problem')
    parser.add_argument('--method', type=str, required=True,
                        choices=['PID-GAN', 'PIG-GAN', 'PINN', 'APINN', 'cGAN'],
                        help='Model variant')
    parser.add_argument('--use-npu', type=int, default=1,
                        help='Use NPU device (1=yes, 0=no)')
    parser.add_argument('--device-id', type=int, default=0,
                        help='Device index (NPU or CUDA)')
    parser.add_argument('--num-epochs', type=int, default=30000,
                        help='Number of training epochs')
    parser.add_argument('--lambda-val', type=float, default=0.05,
                        help='Lambda value for physics probability')
    parser.add_argument('--noise', type=float, default=0.1,
                        help='Noise level for labeled data')
    parser.add_argument('--noise-dim', type=int, default=2,
                        help='Noise dimension (for imperfect physics)')
    parser.add_argument('--hid-dim', type=int, default=50,
                        help='Hidden dimension for networks')
    parser.add_argument('--num-layers', type=int, default=4,
                        help='Number of hidden layers for networks')
    parser.add_argument('--seed', type=int, default=1234,
                        help='Random seed')
    args = parser.parse_args()

    np.random.seed(args.seed)

    device = get_device(args.use_npu, args.device_id)

    if args.problem == 'burgers':
        train_burgers(args.method, device, args.num_epochs, args.lambda_val,
                      args.noise, args.hid_dim, args.num_layers)
    elif args.problem == 'schrodinger':
        train_schrodinger(args.method, device, args.num_epochs, args.lambda_val,
                          args.noise, args.hid_dim, args.num_layers)
    elif args.problem == 'darcy':
        lambdas = [args.lambda_val, 0.5]
        train_darcy(args.method, device, args.num_epochs, lambdas,
                    args.noise, args.hid_dim, args.num_layers)
    elif args.problem == 'tossing':
        train_tossing(args.method, device, args.num_epochs, args.lambda_val,
                      args.noise_dim, args.hid_dim, args.num_layers)


if __name__ == '__main__':
    main()

# Burgers PID-GAN
python -u train.py --problem burgers --method PID-GAN --use-npu 1 --device-id 0 --num-epochs 30000 --lambda-val 0.05 2>&1 | tee train_Burgers.log

# Schrodinger PID-GAN 
python -u train.py --problem schrodinger --method PID-GAN --use-npu 1 --device-id 0 --num-epochs 50000 --lambda-val 0.1 2>&1 | tee train_schrodinger.log

# Darcy PID-GAN 
python -u train.py --problem darcy --method PID-GAN --use-npu 1 --device-id 0 --num-epochs 30000 --lambda-val 0.1 2>&1 | tee train_darcy.log

import argparse
import os
import sys
import warnings

import torch_npu
from torch_npu.contrib import transfer_to_npu

import numpy as np
import torch
import scipy.io
from scipy.interpolate import griddata
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

warnings.filterwarnings('ignore')

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATASET_DIR = os.path.join(BASE_DIR, 'datasets')


def get_device(use_npu, device_id):
    if use_npu and torch.npu.is_available():
        device = torch.device(f'npu:{device_id}')
        torch.npu.set_device(device)
        print(f'Using NPU device: {device}, count: {torch.npu.device_count()}', flush=True)
    else:
        device = torch.device('cpu')
        print('Using CPU', flush=True)
    return device


def load_burgers_data():
    data = scipy.io.loadmat(os.path.join(DATASET_DIR, 'burgers_shock.mat'))
    t = data['t'].flatten()[:, None]
    x = data['x'].flatten()[:, None]
    Exact = np.real(data['usol']).T
    X, T = np.meshgrid(x, t)
    X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
    u_star = Exact.flatten()[:, None]
    return X, T, x, t, Exact, X_star, u_star


def load_schrodinger_inference_data():
    data = scipy.io.loadmat(os.path.join(DATASET_DIR, 'NLS.mat'))
    t = data['tt'].flatten()[:, None]
    x = data['x'].flatten()[:, None]
    Exact = data['uu']
    Exact_u = np.real(Exact)
    Exact_v = np.imag(Exact)
    Exact_h = np.sqrt(Exact_u**2 + Exact_v**2)
    X, T = np.meshgrid(x, t)
    X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
    u_star = Exact_u.T.flatten()[:, None]
    v_star = Exact_v.T.flatten()[:, None]
    h_star = Exact_h.T.flatten()[:, None]
    return X, T, x, t, Exact_u, Exact_v, Exact_h, X_star, u_star, v_star, h_star


def run_schrodinger_inference(method, device, num_epochs, noise, nsamples, checkpoint_path, output_dir):
    sys.path.insert(0, os.path.join(BASE_DIR, 'Scripts'))
    from models_pde import Generator, Discriminator, Q_Net

    X, T, x, t, Exact_u, Exact_v, Exact_h, X_star, u_star, v_star, h_star = load_schrodinger_inference_data()

    state = torch.load(checkpoint_path, map_location=device, weights_only=False)

    hid_dim = 50
    num_layers = 4
    d_num_layers = 2
    d_in_dim = 6 if method == 'PID-GAN' else 4

    G = Generator(in_dim=3, out_dim=2, hid_dim=hid_dim, num_layers=num_layers).to(device)
    D = Discriminator(in_dim=d_in_dim, out_dim=1, hid_dim=hid_dim, num_layers=d_num_layers).to(device)
    Q = Q_Net(in_dim=4, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)

    G.load_state_dict(state['G_state_dict'])
    D.load_state_dict(state['D_state_dict'])
    Q.load_state_dict(state['Q_state_dict'])
    G.eval()
    D.eval()
    Q.eval()

    print(f'[Inference] Schrodinger / {method} / noise={noise} / nsamples={nsamples} / device={device}', flush=True)

    u_pred_list = []
    v_pred_list = []
    h_pred_list = []
    f_u_pred_list = []
    f_v_pred_list = []
    for run in range(nsamples):
        x_t = torch.tensor(X_star[:, 0:1], requires_grad=True).float().to(device)
        t_t = torch.tensor(X_star[:, 1:2], requires_grad=True).float().to(device)
        noise_z = torch.randn((X_star.shape[0], 1)).float().to(device)

        y_pred = G(torch.cat([x_t, t_t, noise_z], dim=1))
        u = y_pred[:, 0:1]
        v = y_pred[:, 1:2]

        u_t = torch.autograd.grad(u, t_t, grad_outputs=torch.ones_like(u), retain_graph=True, create_graph=True)[0]
        u_x = torch.autograd.grad(u, x_t, grad_outputs=torch.ones_like(u), retain_graph=True, create_graph=True)[0]
        u_xx = torch.autograd.grad(u_x, x_t, grad_outputs=torch.ones_like(u_x), retain_graph=True, create_graph=True)[0]
        v_t = torch.autograd.grad(v, t_t, grad_outputs=torch.ones_like(v), retain_graph=True, create_graph=True)[0]
        v_x = torch.autograd.grad(v, x_t, grad_outputs=torch.ones_like(v), retain_graph=True, create_graph=True)[0]
        v_xx = torch.autograd.grad(v_x, x_t, grad_outputs=torch.ones_like(v_x), retain_graph=False, create_graph=False)[0]

        f_u = u_t + 0.5 * v_xx + (u**2 + v**2) * v
        f_v = v_t - 0.5 * u_xx - (u**2 + v**2) * u

        u_pred_list.append(u.detach().cpu().numpy())
        v_pred_list.append(v.detach().cpu().numpy())
        h_pred_list.append(np.sqrt(u.detach().cpu().numpy()**2 + v.detach().cpu().numpy()**2))
        f_u_pred_list.append(f_u.detach().cpu().numpy())
        f_v_pred_list.append(f_v.detach().cpu().numpy())

    u_pred_arr = np.array(u_pred_list)
    v_pred_arr = np.array(v_pred_list)
    h_pred_arr = np.array(h_pred_list)
    f_u_pred_arr = np.array(f_u_pred_list)
    f_v_pred_arr = np.array(f_v_pred_list)
    u_pred = u_pred_arr.mean(axis=0)
    v_pred = v_pred_arr.mean(axis=0)
    h_pred = h_pred_arr.mean(axis=0)
    f_u_pred = f_u_pred_arr.mean(axis=0)
    f_v_pred = f_v_pred_arr.mean(axis=0)
    h_dev = h_pred_arr.var(axis=0)

    error_u = np.linalg.norm(u_star - u_pred, 2) / np.linalg.norm(u_star, 2)
    error_v = np.linalg.norm(v_star - v_pred, 2) / np.linalg.norm(v_star, 2)
    error_h = np.linalg.norm(h_star - h_pred, 2) / np.linalg.norm(h_star, 2)
    residual = (f_u_pred**2).mean() + (f_v_pred**2).mean()
    print(f'[Result] Error u: {error_u:.6e} | Error v: {error_v:.6e} | Error h: {error_h:.6e} | Residual: {residual:.6e}', flush=True)

    os.makedirs(output_dir, exist_ok=True)

    lb = np.array([-5.0, 0.0])
    ub = np.array([5.0, np.pi / 2])

    H_pred = griddata(X_star, h_pred.flatten(), (X, T), method='cubic')
    H_dev_grid = griddata(X_star, h_dev.flatten(), (X, T), method='cubic')

    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    titles = ['Exact |h(x,t)|', 'Predicted |h(x,t)|', 'Uncertainty (var)']
    data_list = [Exact_h.T, H_pred, H_dev_grid]
    for ax, title, data in zip(axes, titles, data_list):
        im = ax.imshow(data, interpolation='nearest', cmap='YlGnBu',
                       extent=[lb[1], ub[1], lb[0], ub[0]],
                       origin='lower', aspect='auto')
        ax.set_title(title)
        ax.set_xlabel('$t$')
        ax.set_ylabel('$x$')
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.10)
        plt.colorbar(im, cax=cax)
    plt.tight_layout()
    fig_path = os.path.join(output_dir, f'schrodinger_{method}_heatmap.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'Heatmap saved: {fig_path}', flush=True)

    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    for idx, t_idx in enumerate([75, 100, 125]):
        ax = axes[idx]
        ax.plot(x.flatten(), Exact_h.T[t_idx, :], 'b-', label='Exact')
        ax.plot(x.flatten(), H_pred[t_idx, :], 'r--', label='Prediction')
        lower = H_pred[t_idx, :] - 2.0 * np.sqrt(np.abs(H_dev_grid[t_idx, :]))
        upper = H_pred[t_idx, :] + 2.0 * np.sqrt(np.abs(H_dev_grid[t_idx, :]))
        ax.fill_between(x.flatten(), lower, upper, facecolor='orange', alpha=0.5, label='2 std band')
        ax.set_xlabel('$x$')
        ax.set_ylabel('$|h(x,t={})|$'.format(t.flatten()[t_idx]))
        ax.set_title('$t = {}$'.format(t.flatten()[t_idx]))
        ax.legend(loc='upper center')
    plt.tight_layout()
    fig_path = os.path.join(output_dir, f'schrodinger_{method}_slices.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'Slice plot saved: {fig_path}', flush=True)

    results = {
        'error_u': error_u,
        'error_v': error_v,
        'error_h': error_h,
        'residual': residual,
        'u_pred': u_pred,
        'v_pred': v_pred,
        'h_pred': h_pred,
        'h_dev': h_dev,
        'f_u_pred': f_u_pred,
        'f_v_pred': f_v_pred,
        'X_star': X_star,
        'u_star': u_star,
        'v_star': v_star,
        'h_star': h_star,
    }
    results_path = os.path.join(output_dir, f'schrodinger_{method}_results.npz')
    np.savez(results_path, **results)
    print(f'Results saved: {results_path}', flush=True)


def load_darcy_inference_data():
    data = np.load(os.path.join(DATASET_DIR, 'nonlinear2d_data.npz'))
    X = data['X']
    U = data['u']
    K = data['k']
    u_star = U.T
    k_star = K.T
    lb = np.array([0.0, 0.0])
    ub = np.array([10.0, 10.0])
    X_star = X
    X_star_norm = (X_star - lb) - 0.5 * (ub - lb)
    return X_star, X_star_norm, u_star, k_star


def run_burgers_inference(method, device, num_epochs, noise, nsamples, checkpoint_path, output_dir):
    sys.path.insert(0, os.path.join(BASE_DIR, 'Scripts'))
    from models_pde import Generator, Discriminator, Q_Net

    X, T, x, t, Exact, X_star, u_star = load_burgers_data()

    state = torch.load(checkpoint_path, map_location=device, weights_only=False)

    hid_dim = 50
    num_layers = 4
    d_num_layers = 2
    g_in_dim = 3
    d_in_dim = 4 if method == 'PID-GAN' else 3

    G = Generator(in_dim=g_in_dim, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
    D = Discriminator(in_dim=d_in_dim, out_dim=1, hid_dim=hid_dim, num_layers=d_num_layers).to(device)
    Q = Q_Net(in_dim=3, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)

    G.load_state_dict(state['G_state_dict'])
    D.load_state_dict(state['D_state_dict'])
    Q.load_state_dict(state['Q_state_dict'])
    G.eval()
    D.eval()
    Q.eval()

    Xmean = state['Xmean']
    Xstd = state['Xstd']
    lambda_val = state['lambda_val']
    X_star_norm = (X_star - Xmean) / Xstd

    sys.path.insert(0, os.path.join(BASE_DIR, 'PDEs', 'Burgers', method))
    if method == 'PID-GAN':
        from pid import Burgers_PID
        nu = 0.01 / np.pi
        Jacobian_X = 1 / Xstd[0]
        Jacobian_T = 1 / Xstd[1]
    else:
        from pig import Burgers_PIG
        nu = 0.01 / np.pi
        Jacobian_X = 1 / Xstd[0]
        Jacobian_T = 1 / Xstd[1]

    print(f'[Inference] Burgers / {method} / noise={noise} / nsamples={nsamples} / device={device}', flush=True)

    u_pred_list = []
    f_pred_list = []
    for run in range(nsamples):
        x_t = torch.tensor(X_star_norm[:, 0:1], requires_grad=True).float().to(device)
        t_t = torch.tensor(X_star_norm[:, 1:2], requires_grad=True).float().to(device)
        noise_z = torch.randn((X_star_norm.shape[0], 1)).float().to(device)
        u = G(torch.cat([x_t, t_t, noise_z], dim=1))

        u_t = torch.autograd.grad(u, t_t, grad_outputs=torch.ones_like(u), retain_graph=True, create_graph=True)[0]
        u_x = torch.autograd.grad(u, x_t, grad_outputs=torch.ones_like(u), retain_graph=True, create_graph=True)[0]
        u_xx = torch.autograd.grad(u_x, x_t, grad_outputs=torch.ones_like(u_x), retain_graph=False, create_graph=False)[0]

        f = Jacobian_T * u_t + Jacobian_X * u * u_x - nu * (Jacobian_X ** 2) * u_xx

        u_pred_list.append(u.detach().cpu().numpy())
        f_pred_list.append(f.detach().cpu().numpy())

    u_pred_arr = np.array(u_pred_list)
    f_pred_arr = np.array(f_pred_list)
    u_pred = u_pred_arr.mean(axis=0)
    f_pred = f_pred_arr.mean(axis=0)
    u_dev = u_pred_arr.var(axis=0)

    error_u = np.linalg.norm(u_star - u_pred, 2) / np.linalg.norm(u_star, 2)
    residual = (f_pred ** 2).mean()
    print(f'[Result] Error u: {error_u:.6e} | Residual: {residual:.6e}', flush=True)

    os.makedirs(output_dir, exist_ok=True)

    U_pred = griddata(X_star, u_pred.flatten(), (X, T), method='cubic')
    U_dev = griddata(X_star, u_dev.flatten(), (X, T), method='cubic')
    Error_field = np.abs(Exact - U_pred)

    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    titles = ['Exact', 'Prediction', 'Uncertainty (variance)']
    data_list = [Exact.T, U_pred.T, U_dev.T]
    for ax, title, data in zip(axes, titles, data_list):
        im = ax.imshow(data, interpolation='nearest', cmap='rainbow',
                       extent=[t.min(), t.max(), x.min(), x.max()],
                       origin='lower', aspect='auto')
        ax.set_title(title)
        ax.set_xlabel('$t$')
        ax.set_ylabel('$x$')
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.10)
        plt.colorbar(im, cax=cax)
    plt.tight_layout()
    fig_path = os.path.join(output_dir, f'burgers_{method}_heatmap.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'Heatmap saved: {fig_path}', flush=True)

    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    for idx, t_idx in enumerate([25, 50, 75]):
        ax = axes[idx]
        ax.plot(x.flatten(), Exact[t_idx, :], 'b-', label='Exact')
        ax.plot(x.flatten(), U_pred[t_idx, :], 'r--', label='Prediction')
        lower = U_pred[t_idx, :] - 2.0 * np.sqrt(U_dev[t_idx, :])
        upper = U_pred[t_idx, :] + 2.0 * np.sqrt(U_dev[t_idx, :])
        ax.fill_between(x.flatten(), lower, upper, facecolor='orange', alpha=0.5, label='2 std band')
        ax.set_xlim([-1.1, 1.1])
        ax.set_xlabel('$x$')
        ax.set_ylabel('$u(x,t={})$'.format(t.flatten()[t_idx]))
        ax.set_title('$t = {}$'.format(t.flatten()[t_idx]))
        ax.legend(loc='upper center')
    plt.tight_layout()
    fig_path = os.path.join(output_dir, f'burgers_{method}_slices.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'Slice plot saved: {fig_path}', flush=True)

    results = {
        'error_u': error_u,
        'residual': residual,
        'u_pred': u_pred,
        'u_dev': u_dev,
        'f_pred': f_pred,
        'X_star': X_star,
        'u_star': u_star,
    }
    results_path = os.path.join(output_dir, f'burgers_{method}_results.npz')
    np.savez(results_path, **results)
    print(f'Results saved: {results_path}', flush=True)


def run_darcy_inference(method, device, num_epochs, noise, nsamples, checkpoint_path, output_dir):
    sys.path.insert(0, os.path.join(BASE_DIR, 'Scripts'))
    from models_pde import Generator, Discriminator, Q_Net

    X_star, X_star_norm, u_star, k_star = load_darcy_inference_data()
    ksat = 10.0

    state = torch.load(checkpoint_path, map_location=device, weights_only=False)

    hid_dim = 50
    num_layers = 4
    d_num_layers = 3

    G = Generator(in_dim=4, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
    kG = Generator(in_dim=1, out_dim=1, hid_dim=hid_dim, num_layers=num_layers).to(device)
    D = Discriminator(in_dim=4, out_dim=1, hid_dim=hid_dim, num_layers=d_num_layers).to(device)
    Q = Q_Net(in_dim=3, out_dim=2, hid_dim=hid_dim, num_layers=num_layers).to(device)

    G.load_state_dict(state['G_state_dict'])
    kG.load_state_dict(state['kG_state_dict'])
    D.load_state_dict(state['D_state_dict'])
    Q.load_state_dict(state['Q_state_dict'])
    G.eval()
    kG.eval()
    D.eval()
    Q.eval()

    print(f'[Inference] Darcy / {method} / noise={noise} / nsamples={nsamples} / device={device}', flush=True)

    u_pred_list = []
    f_pred_list = []
    k_pred_list = []
    for run in range(nsamples):
        x1 = torch.tensor(X_star_norm[:, 0:1], requires_grad=True).float().to(device)
        x2 = torch.tensor(X_star_norm[:, 1:2], requires_grad=True).float().to(device)
        noise_z = torch.randn((X_star_norm.shape[0], 2)).float().to(device)

        u = G(torch.cat([x1, x2, noise_z], dim=1))
        k = kG(u)
        k = k / ksat

        u_x1 = torch.autograd.grad(u, x1, grad_outputs=torch.ones_like(u), retain_graph=True, create_graph=True)[0]
        u_x2 = torch.autograd.grad(u, x2, grad_outputs=torch.ones_like(u), retain_graph=True, create_graph=True)[0]
        f_1 = torch.autograd.grad(k * u_x1, x1, grad_outputs=torch.ones_like(u_x1), retain_graph=True, create_graph=False)[0]
        f_2 = torch.autograd.grad(k * u_x2, x2, grad_outputs=torch.ones_like(u_x2), retain_graph=True, create_graph=False)[0]
        f = f_1 + f_2

        u_pred_list.append(u.detach().cpu().numpy())
        f_pred_list.append(f.detach().cpu().numpy())
        k_pred_list.append(k.detach().cpu().numpy())

    u_pred_arr = np.array(u_pred_list)
    f_pred_arr = np.array(f_pred_list)
    k_pred_arr = np.array(k_pred_list)
    u_pred = u_pred_arr.mean(axis=0)
    f_pred = f_pred_arr.mean(axis=0)
    k_pred = k_pred_arr.mean(axis=0)
    u_dev = u_pred_arr.var(axis=0)

    error_u = np.linalg.norm(u_star.T - u_pred, 2) / np.linalg.norm(u_star.T, 2)
    residual = (f_pred ** 2).mean()
    error_k = np.linalg.norm(k_star.T - k_pred, 2) / np.linalg.norm(k_star.T, 2)
    print(f'[Result] Error u: {error_u:.6e} | Residual: {residual:.6e} | Error k: {error_k:.6e}', flush=True)

    os.makedirs(output_dir, exist_ok=True)

    N_grid = 100
    x1_grid = np.linspace(0, 10, N_grid)
    x2_grid = np.linspace(0, 10, N_grid)
    X1_mesh, X2_mesh = np.meshgrid(x1_grid, x2_grid)
    X_grid = np.hstack((X1_mesh.flatten()[:, None], X2_mesh.flatten()[:, None]))
    lb = np.array([0.0, 0.0])
    ub = np.array([10.0, 10.0])
    X_grid_norm = (X_grid - lb) - 0.5 * (ub - lb)

    u_grid_list = []
    k_grid_list = []
    for run in range(nsamples):
        x1_t = torch.tensor(X_grid_norm[:, 0:1], requires_grad=False).float().to(device)
        x2_t = torch.tensor(X_grid_norm[:, 1:2], requires_grad=False).float().to(device)
        noise_z = torch.randn((X_grid_norm.shape[0], 2)).float().to(device)
        u_g = G(torch.cat([x1_t, x2_t, noise_z], dim=1))
        k_g = kG(u_g) / ksat
        u_grid_list.append(u_g.detach().cpu().numpy())
        k_grid_list.append(k_g.detach().cpu().numpy())

    u_grid_arr = np.array(u_grid_list).mean(axis=0)
    k_grid_arr = np.array(k_grid_list).mean(axis=0)
    u_grid_dev_arr = np.array(u_grid_list).var(axis=0)
    U_grid = griddata(X_star, u_pred.flatten(), (X1_mesh, X2_mesh), method='cubic')
    K_grid = griddata(X_star, k_pred.flatten(), (X1_mesh, X2_mesh), method='cubic')
    U_dev_grid = griddata(X_star, u_dev.flatten(), (X1_mesh, X2_mesh), method='cubic')

    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    titles_row1 = ['Exact u', 'Predicted u', 'Uncertainty u (var)']
    data_row1 = [u_star.T.reshape(N_grid, N_grid) if u_star.T.shape == (N_grid, N_grid) else u_star, U_grid, U_dev_grid]
    titles_row2 = ['Exact k', 'Predicted k', 'Error u']
    data_row2 = [k_star.T.reshape(N_grid, N_grid) if k_star.T.shape == (N_grid, N_grid) else k_star, K_grid, np.abs(u_star.T - u_pred)]
    for row, (titles, datas) in enumerate(zip([titles_row1, titles_row2], [data_row1, data_row2])):
        for col, (title, data) in enumerate(zip(titles, datas)):
            ax = axes[row, col]
            if isinstance(data, np.ndarray) and data.ndim <= 2:
                try:
                    im = ax.imshow(data, interpolation='nearest', cmap='rainbow',
                                   extent=[0, 10, 0, 10], origin='lower', aspect='auto')
                except Exception:
                    im = ax.imshow(data.flatten().reshape(int(np.sqrt(data.flatten().shape[0])), -1),
                                   interpolation='nearest', cmap='rainbow',
                                   extent=[0, 10, 0, 10], origin='lower', aspect='auto')
            else:
                im = ax.imshow(np.zeros((10, 10)), cmap='rainbow', extent=[0, 10, 0, 10], origin='lower', aspect='auto')
            ax.set_title(title)
            divider = make_axes_locatable(ax)
            cax = divider.append_axes("right", size="5%", pad=0.10)
            plt.colorbar(im, cax=cax)
    plt.tight_layout()
    fig_path = os.path.join(output_dir, f'darcy_{method}_heatmap.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'Heatmap saved: {fig_path}', flush=True)

    results = {
        'error_u': error_u,
        'residual': residual,
        'error_k': error_k,
        'u_pred': u_pred,
        'u_dev': u_dev,
        'f_pred': f_pred,
        'k_pred': k_pred,
        'X_star': X_star,
        'u_star': u_star.T,
        'k_star': k_star.T,
    }
    results_path = os.path.join(output_dir, f'darcy_{method}_results.npz')

    np.savez(results_path, **results)
    print(f'Results saved: {results_path}', flush=True)


def main():
    parser = argparse.ArgumentParser(description='PID-GAN Inference on Ascend NPU')
    parser.add_argument('--problem', type=str, required=True,
                        choices=['burgers', 'schrodinger', 'darcy', 'tossing'])
    parser.add_argument('--method', type=str, required=True,
                        choices=['PID-GAN', 'PIG-GAN', 'PINN', 'APINN', 'cGAN'])
    parser.add_argument('--checkpoint', type=str, required=True,
                        help='Path to checkpoint .pth file')
    parser.add_argument('--use-npu', type=int, default=1)
    parser.add_argument('--device-id', type=int, default=0)
    parser.add_argument('--noise', type=float, default=0.1)
    parser.add_argument('--num-epochs', type=int, default=30000,
                        help='Used to locate checkpoint if --checkpoint not given')
    parser.add_argument('--nsamples', type=int, default=500,
                        help='Number of stochastic forward passes for prediction')
    parser.add_argument('--output-dir', type=str, default=None,
                        help='Output directory for plots and results')

    args = parser.parse_args()
    device = get_device(args.use_npu, args.device_id)

    if args.output_dir is None:
        args.output_dir = os.path.join(BASE_DIR, 'results')

    if args.problem == 'burgers':
        run_burgers_inference(args.method, device, args.num_epochs, args.noise,
                              args.nsamples, args.checkpoint, args.output_dir)
    elif args.problem == 'darcy':
        run_darcy_inference(args.method, device, args.num_epochs, args.noise,
                            args.nsamples, args.checkpoint, args.output_dir)
    elif args.problem == 'schrodinger':
        run_schrodinger_inference(args.method, device, args.num_epochs, args.noise,
                                  args.nsamples, args.checkpoint, args.output_dir)
    else:
        print(f'Inference for {args.problem} / {args.method} not yet implemented', flush=True)


if __name__ == '__main__':
    main()

python -u inference.py --problem burgers --method PID-GAN \
    --checkpoint checkpoints/burgers_PID-GAN_epochs30000_noise0.1.pth \
    --use-npu 1 --nsamples 500 \
    --output-dir results/burgers_pid_gan_noisy