diff --git a/docs/zh/examples/images/ML/psc_nn_overview.png b/docs/zh/examples/images/ML/psc_nn_overview.png new file mode 100644 index 0000000000..f8afd96416 Binary files /dev/null and b/docs/zh/examples/images/ML/psc_nn_overview.png differ diff --git a/docs/zh/examples/params/solar_cell_pretrained.pdparams b/docs/zh/examples/params/solar_cell_pretrained.pdparams new file mode 100644 index 0000000000..b22937a19b Binary files /dev/null and b/docs/zh/examples/params/solar_cell_pretrained.pdparams differ diff --git a/docs/zh/examples/perovskite_solar_cells_nn.md b/docs/zh/examples/perovskite_solar_cells_nn.md new file mode 100644 index 0000000000..fb26ad53d1 --- /dev/null +++ b/docs/zh/examples/perovskite_solar_cells_nn.md @@ -0,0 +1,173 @@ +# psc_nn(Machine Learning for Perovskite Solar Cells: An Open-Source Pipeline) + +!!! note "注意事项" + + 1. 开始训练前，请确保数据集已正确放置在 `data/cleaned/` 目录下。 + 2. 训练和评估需要安装额外的依赖包，请使用 `pip install -r requirements.txt` 安装。 + 3. 为获得最佳性能，建议使用 GPU 进行训练。 + +=== "模型训练命令" + + ``` sh + python psc_nn.py mode=train + ``` + +=== "模型评估命令" + + ``` sh + # 使用本地预训练模型 + python psc_nn.py mode=eval eval.pretrained_model_path="Your pdparams path" + ``` + + ``` sh + # 或使用远程预训练模型 + python psc_nn.py mode=eval eval.pretrained_model_path="https://paddle-org.bj.bcebos.com/paddlescience/models/PerovskiteSolarCells/solar_cell_pretrained.pdparams" + ``` + +| 预训练模型 | 指标 | +|:--| :--| +| [solar_cell_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/PerovskiteSolarCells/solar_cell_pretrained.pdparams) | RMSE: 3.91798 | + +## 1. 背景简介 + +太阳能电池是一种通过光电效应将光能直接转换为电能的关键能源器件，其性能预测是优化和设计太阳能电池的重要环节。然而，传统的性能预测方法往往依赖于复杂的物理模拟和大量的实验测试，不仅成本高昂，且耗时较长，制约了研究与开发的效率。 + +近年来，深度学习和机器学习技术的快速发展，为太阳能电池性能预测提供了创新的方法。通过机器学习技术，可以显著加快开发速度，同时实现与实验结果相当的预测精度。特别是在钙钛矿太阳能电池研究中，材料的化学组成和结构多样性为模型训练带来了新的挑战。为了解决这一问题，研究者们通常将材料的特性转换为固定长度的特征向量，以适配机器学习模型。尽管如此，不同性能指标的特征表示设计仍需不断优化，同时对模型预测结果的可解释性要求也更为严格。 + +本研究中，通过利用包含钙钛矿太阳能电池特性信息的全面数据库（PDP），我们构建并评估了包括 XGBoost、psc_nn 在内的多种机器学习模型，专注于预测短路电流密度（Jsc）。研究结果表明，结合深度学习与超参数优化工具（如 Optuna）能够显著提升太阳能电池设计的效率，为新型太阳能电池研发提供了更精确且高效的解决方案。 + +## 2. 模型原理 + +本章节仅对太阳能电池性能预测模型的原理进行简单地介绍，详细的理论推导请阅读 [Machine Learning for Perovskite Solar Cells: An Open-Source Pipeline](https://onlinelibrary.wiley.com/doi/10.1002/apxr.202400060)。 + +该方法的主要思想是通过人工神经网络建立光谱响应数据与短路电流密度（Jsc）之间的非线性映射关系。人工神经网络模型的总体结构如下图所示： + +![psc_nn_overview](../images/ML/psc_nn_overview.png) + +本案例采用多层感知机（MLP）作为基础模型架构，主要包括以下几个部分： + +1. 输入层：接收 2808 维的光谱响应数据 +2. 隐藏层：4-6 层全连接层，每层的神经元数量通过 Optuna 优化 +3. 激活函数：使用 ReLU 激活函数引入非线性特性 +4. 输出层：输出预测的 Jsc 值 + +通过这种方式，我们可以自动找到最适合当前任务的模型配置，提高模型的预测性能。 + +## 3. 模型实现 + +本章节我们讲解如何基于 PaddleScience 代码实现钙钛矿太阳能电池性能预测模型。本案例结合 Optuna 框架进行超参数优化，并使用 PaddleScience 内置的各种功能模块。为了快速理解 PaddleScience，接下来仅对模型构建、约束构建、评估器构建等关键步骤进行阐述，而其余细节请参考 [API文档](../api/arch.md)。 + +### 3.1 数据集介绍 + +本案例使用的数据集包含 Perovskite Database Project(PDP) 数据。数据集分为以下几个部分： + +1. 训练集： + - 特征数据：`data/cleaned/training.csv` + - 标签数据：`data/cleaned/training_labels.csv` +2. 验证集： + - 特征数据：`data/cleaned/validation.csv` + - 标签数据：`data/cleaned/validation_labels.csv` + +为了方便数据处理，我们实现了一个辅助函数 `create_tensor_dict` 来创建输入和标签的 tensor 字典： + +``` py linenums="31" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py:31:37 +--8<-- +``` + +数据集的读取和预处理代码如下： + +``` py linenums="118" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py:118:129 +--8<-- +``` + +为了进行超参数优化，我们将训练集进一步划分为训练集和验证集： + +``` py linenums="130" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py:130:135 +--8<-- +``` + +### 3.2 模型构建 + +本案例使用 PaddleScience 内置的 `ppsci.arch.MLP` 构建多层感知机模型。模型的超参数通过 Optuna 框架进行优化，主要包括： + +1. 网络层数：4-6层 +2. 每层神经元数量：10-input_dim/2 +3. 激活函数：ReLU +4. 输入维度：2808（光谱响应数据维度） +5. 输出维度：1（Jsc 预测值） + +模型定义代码如下： + +``` py linenums="98" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py:98:115 +--8<-- +``` + +### 3.3 损失函数设计 + +考虑到数据集中不同样本的重要性可能不同，我们设计了一个加权均方误差损失函数。该函数对较大的 Jsc 值赋予更高的权重，以提高模型在高性能太阳能电池上的预测准确性： + +``` py linenums="20" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py:20:30 +--8<-- +``` + +### 3.4 约束构建 + +本案例基于数据驱动的方法求解问题，因此使用 PaddleScience 内置的 `SupervisedConstraint` 构建监督约束。为了减少代码重复，我们实现了 `create_constraint` 函数来创建监督约束： + +``` py linenums="38" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py:38:58 +--8<-- +``` + +### 3.5 评估器构建 + +为了实时监测模型的训练情况，我们实现了 `create_validator` 函数来创建评估器： + +``` py linenums="59" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py:59:78 +--8<-- +``` + +### 3.6 优化器构建 + +为了统一管理优化器和学习率调度器的创建，我们实现了 `create_optimizer` 函数： + +``` py linenums="79" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py:79:97 +--8<-- +``` + +### 3.7 模型训练与评估 + +在训练过程中，我们使用上述封装的函数来创建数据字典、约束、评估器和优化器： + +``` py linenums="230" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py:230:280 +--8<-- +``` + +## 4. 完整代码 + +``` py linenums="1" title="examples/perovskite_solar_cells/psc_nn.py" +--8<-- +examples/perovskite_solar_cells/psc_nn.py +--8<-- +``` + +## 5. 参考文献 + +- [Machine Learning for Perovskite Solar Cells: An Open-Source Pipeline](https://onlinelibrary.wiley.com/doi/10.1002/apxr.202400060) diff --git a/examples/perovskite_solar_cells/conf/psc_nn.yaml b/examples/perovskite_solar_cells/conf/psc_nn.yaml new file mode 100644 index 0000000000..ff649cb0c7 --- /dev/null +++ b/examples/perovskite_solar_cells/conf/psc_nn.yaml @@ -0,0 +1,59 @@ +defaults: + - ppsci_default + - TRAIN: train_default + - TRAIN/ema: ema_default + - TRAIN/swa: swa_default + - EVAL: eval_default + - INFER: infer_default + - hydra/job/config/override_dirname/exclude_keys: exclude_keys_default + - _self_ + +hydra: + run: + dir: outputs_allen_cahn_piratenet/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} + job: + name: ${mode} + chdir: false + callbacks: + init_callback: + _target_: ppsci.utils.callbacks.InitCallback + sweep: + dir: ${hydra.run.dir} + subdir: ./ + +mode: "train" +seed: 42 +output_dir: "./output" + +data: + train_features_path: "./data/cleaned/training.csv" + train_labels_path: "./data/cleaned/training_labels.csv" + val_features_path: "./data/cleaned/validation.csv" + val_labels_path: "./data/cleaned/validation_labels.csv" + +model: + num_layers: 4 + hidden_size: [128, 96, 64, 32] + activation: "relu" + input_dim: 2808 + output_dim: 1 + +TRAIN: + epochs: 10 + batch_size: 64 + learning_rate: 0.001 + eval_during_train: true + eval_freq: 5 + save_freq: 10 + log_freq: 50 + lr_scheduler: + gamma: 0.95 + decay_steps: 5 + warmup_epoch: 2 + warmup_start_lr: 1.0e-6 + +eval: + batch_size: 64 + eval_with_no_grad: true + pretrained_model_path: null + log_freq: 50 diff --git a/examples/perovskite_solar_cells/psc_nn.py b/examples/perovskite_solar_cells/psc_nn.py new file mode 100644 index 0000000000..d608e0867f --- /dev/null +++ b/examples/perovskite_solar_cells/psc_nn.py @@ -0,0 +1,364 @@ +import os +from os import path as osp + +import hydra +import numpy as np +import optuna +import paddle +import pandas as pd +from matplotlib import pyplot as plt +from omegaconf import DictConfig +from sklearn.metrics import mean_absolute_percentage_error +from sklearn.metrics import mean_squared_error +from sklearn.metrics import r2_score +from sklearn.model_selection import train_test_split + +import ppsci +from ppsci.constraint import SupervisedConstraint +from ppsci.optimizer import lr_scheduler +from ppsci.optimizer import optimizer +from ppsci.solver import Solver +from ppsci.validate import SupervisedValidator + + +def weighted_loss(output_dict, target_dict, weight_dict=None): + pred = output_dict["target"] + true = target_dict["target"] + epsilon = 1e-06 + n = len(true) + weights = true / (paddle.sum(x=true) + epsilon) + squared = (true - pred) ** 2 + weighted = squared * weights + loss = paddle.sum(x=weighted) / n + return {"weighted_mse": loss} + + +def create_tensor_dict(X, y): + """Create Tensor Dictionary for Input and Labels""" + return { + "input": paddle.to_tensor(X.values, dtype="float32"), + "label": {"target": paddle.to_tensor(y.values, dtype="float32")}, + } + + +def create_constraint(input_dict, batch_size, shuffle=True): + """Create supervision constraints""" + return SupervisedConstraint( + dataloader_cfg={ + "dataset": { + "name": "NamedArrayDataset", + "input": {"input": input_dict["input"]}, + "label": input_dict["label"], + }, + "batch_size": batch_size, + "sampler": { + "name": "BatchSampler", + "drop_last": False, + "shuffle": shuffle, + }, + }, + loss=weighted_loss, + output_expr={"target": lambda out: out["target"]}, + name="train_constraint", + ) + + +def create_validator(input_dict, batch_size, name="validator"): + """Create an evaluator""" + return SupervisedValidator( + dataloader_cfg={ + "dataset": { + "name": "NamedArrayDataset", + "input": {"input": input_dict["input"]}, + "label": input_dict["label"], + }, + "batch_size": batch_size, + }, + loss=weighted_loss, + output_expr={"target": lambda out: out["target"]}, + metric={"RMSE": ppsci.metric.RMSE(), "MAE": ppsci.metric.MAE()}, + name=name, + ) + + +def create_optimizer(model, optimizer_name, lr, train_cfg, data_size): + """Create optimizer and learning rate scheduler""" + schedule = lr_scheduler.ExponentialDecay( + epochs=train_cfg.epochs, + iters_per_epoch=data_size // train_cfg.batch_size, + learning_rate=lr, + gamma=train_cfg.lr_scheduler.gamma, + decay_steps=train_cfg.lr_scheduler.decay_steps, + warmup_epoch=train_cfg.lr_scheduler.warmup_epoch, + warmup_start_lr=train_cfg.lr_scheduler.warmup_start_lr, + )() + + if optimizer_name == "Adam": + return optimizer.Adam(learning_rate=schedule)(model) + elif optimizer_name == "RMSProp": + return optimizer.RMSProp(learning_rate=schedule)(model) + else: + return optimizer.SGD(learning_rate=schedule)(model) + + +def define_model(trial, input_dim, output_dim): + n_layers = trial.suggest_int("n_layers", 4, 6) + hidden_sizes = [] + for i in range(n_layers): + out_features = trial.suggest_int(f"n_units_l{i}", 10, input_dim // 2) + hidden_sizes.append(out_features) + + model = ppsci.arch.MLP( + input_keys=("input",), + output_keys=("target",), + num_layers=None, + hidden_size=hidden_sizes, + activation="relu", + input_dim=input_dim, + output_dim=output_dim, + ) + return model + + +def train(cfg: DictConfig): + # Read and preprocess data + X_train = pd.read_csv(cfg.data.train_features_path) + y_train = pd.read_csv(cfg.data.train_labels_path) + X_val = pd.read_csv(cfg.data.val_features_path) + y_val = pd.read_csv(cfg.data.val_labels_path) + + for col in X_train.columns: + if "[" in col or "]" in col: + old_name = col + new_name = col.replace("[", "(").replace("]", ")") + X_train = X_train.rename(columns={old_name: new_name}) + X_val = X_val.rename(columns={old_name: new_name}) + + X_train, X_verif, y_train, y_verif = train_test_split( + X_train, y_train, test_size=0.1, random_state=42 + ) + + for df in [X_train, y_train, X_verif, y_verif, X_val, y_val]: + df.reset_index(drop=True, inplace=True) + + def objective(trial): + model = define_model(trial, cfg.model.input_dim, cfg.model.output_dim) + + optimizer_name = trial.suggest_categorical( + "optimizer", ["Adam", "RMSProp", "SGD"] + ) + lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True) + + train_dict = create_tensor_dict(X_train, y_train) + verif_dict = create_tensor_dict(X_verif, y_verif) + + opt = create_optimizer(model, optimizer_name, lr, cfg.TRAIN, len(X_train)) + + train_constraint = create_constraint(train_dict, cfg.TRAIN.batch_size) + verif_validator = create_validator( + verif_dict, cfg.eval.batch_size, "verif_validator" + ) + + solver = Solver( + model=model, + constraint={"train": train_constraint}, + optimizer=opt, + validator={"verif": verif_validator}, + output_dir=cfg.output_dir, + epochs=cfg.TRAIN.epochs, + iters_per_epoch=len(X_train) // cfg.TRAIN.batch_size, + eval_during_train=cfg.TRAIN.eval_during_train, + eval_freq=cfg.TRAIN.eval_freq, + save_freq=cfg.TRAIN.save_freq, + eval_with_no_grad=cfg.eval.eval_with_no_grad, + log_freq=cfg.TRAIN.log_freq, + ) + + solver.train() + + verif_preds = solver.predict({"input": verif_dict["input"]}, return_numpy=True)[ + "target" + ] + + verif_rmse = np.sqrt(mean_squared_error(y_verif.values, verif_preds)) + + return verif_rmse + + study = optuna.create_study() + study.optimize(objective, n_trials=50) + + best_params = study.best_trial.params + print("\nBest hyperparameters: " + str(best_params)) + + # Save the optimal model structure + hidden_sizes = [] + for i in range(best_params["n_layers"]): + hidden_sizes.append(best_params[f"n_units_l{i}"]) + + # Create and train the final model + final_model = define_model( + study.best_trial, cfg.model.input_dim, cfg.model.output_dim + ) + opt = create_optimizer( + final_model, + best_params["optimizer"], + best_params["lr"], + cfg.TRAIN, + len(X_train), + ) + + train_dict = create_tensor_dict(X_train, y_train) + val_dict = create_tensor_dict(X_val, y_val) + + train_constraint = create_constraint(train_dict, cfg.TRAIN.batch_size) + val_validator = create_validator(val_dict, cfg.eval.batch_size, "val_validator") + + solver = Solver( + model=final_model, + constraint={"train": train_constraint}, + optimizer=opt, + validator={"valid": val_validator}, + output_dir=cfg.output_dir, + epochs=cfg.TRAIN.epochs, + iters_per_epoch=len(X_train) // cfg.TRAIN.batch_size, + eval_during_train=cfg.TRAIN.eval_during_train, + eval_freq=cfg.TRAIN.eval_freq, + save_freq=cfg.TRAIN.save_freq, + eval_with_no_grad=cfg.eval.eval_with_no_grad, + log_freq=cfg.TRAIN.log_freq, + ) + + solver.train() + + # Save model structure and weights + model_dict = { + "state_dict": final_model.state_dict(), + "hidden_size": hidden_sizes, + "n_layers": best_params["n_layers"], + "optimizer": best_params["optimizer"], + "lr": best_params["lr"], + } + paddle.save( + model_dict, os.path.join(cfg.output_dir, "checkpoints", "best_model.pdparams") + ) + print( + "Saved model structure and weights to " + + os.path.join(cfg.output_dir, "checkpoints", "best_model.pdparams") + ) + + solver.plot_loss_history(by_epoch=True, smooth_step=1) + solver.eval() + + visualize_results(solver, X_val, y_val, cfg.output_dir) + + +def evaluate(cfg: DictConfig): + # Read and preprocess data + X_val = pd.read_csv(cfg.data.val_features_path) + y_val = pd.read_csv(cfg.data.val_labels_path) + + for col in X_val.columns: + if "[" in col or "]" in col: + old_name = col + new_name = col.replace("[", "(").replace("]", ")") + X_val = X_val.rename(columns={old_name: new_name}) + + # Loading model structure and weights + print(f"Loading model from {cfg.eval.pretrained_model_path}") + model_dict = paddle.load(cfg.eval.pretrained_model_path) + hidden_size = model_dict["hidden_size"] + print(f"Loaded model structure with hidden sizes: {hidden_size}") + + model = ppsci.arch.MLP( + input_keys=("input",), + output_keys=("target",), + num_layers=None, + hidden_size=hidden_size, + activation="relu", + input_dim=cfg.model.input_dim, + output_dim=cfg.model.output_dim, + ) + + # Load model weights + model.set_state_dict(model_dict["state_dict"]) + print("Successfully loaded model weights") + + valid_dict = create_tensor_dict(X_val, y_val) + valid_validator = create_validator( + valid_dict, cfg.eval.batch_size, "valid_validator" + ) + + solver = Solver( + model=model, + output_dir=cfg.output_dir, + validator={"valid": valid_validator}, + eval_with_no_grad=cfg.eval.eval_with_no_grad, + ) + + # evaluation model + print("Evaluating model...") + solver.eval() + + # Generate prediction results + predictions = solver.predict({"input": valid_dict["input"]}, return_numpy=True)[ + "target" + ] + + # Calculate multiple evaluation indicators + rmse = np.sqrt(mean_squared_error(y_val.values, predictions)) + r2 = r2_score(y_val.values, predictions) + mape = mean_absolute_percentage_error(y_val.values, predictions) + + print("Evaluation metrics:") + print(f"RMSE: {rmse:.5f}") + print(f"R2 Score: {r2:.5f}") + print(f"MAPE: {mape:.5f}") + + # Visualization results + print("Generating visualization...") + visualize_results(solver, X_val, y_val, cfg.output_dir) + print("Evaluation completed.") + + +def visualize_results(solver, X_val, y_val, output_dir): + pred_dict = solver.predict( + {"input": paddle.to_tensor(X_val.values, dtype="float32")}, return_numpy=True + ) + val_preds = pred_dict["target"] + val_true = y_val.values + + plt.figure(figsize=(10, 6)) + plt.grid(True, linestyle="--", alpha=0.7) + plt.hist(val_true, bins=30, alpha=0.6, label="True Jsc", color="tab:blue") + plt.hist(val_preds, bins=30, alpha=0.6, label="Predicted Jsc", color="orange") + + pred_mean = np.mean(val_preds) + pred_std = np.std(val_preds) + plt.axvline(pred_mean, color="black", linestyle="--") + plt.axvline(pred_mean + pred_std, color="red", linestyle="--") + plt.axvline(pred_mean - pred_std, color="red", linestyle="--") + + val_rmse = np.sqrt(mean_squared_error(val_true, val_preds)) + plt.title(f"Distribution of True Jsc vs Pred Jsc: RMSE {val_rmse:.5f}", pad=20) + plt.xlabel("Jsc (mA/cm²)") + plt.ylabel("Counts") + plt.legend(fontsize=10) + plt.tight_layout() + plt.savefig( + osp.join(output_dir, "jsc_distribution.png"), dpi=300, bbox_inches="tight" + ) + plt.close() + + +@hydra.main(version_base=None, config_path="./conf", config_name="psc_nn.yaml") +def main(cfg: DictConfig): + if cfg.mode == "train": + train(cfg) + elif cfg.mode == "eval": + evaluate(cfg) + else: + raise ValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'") + + +if __name__ == "__main__": + main() diff --git a/examples/perovskite_solar_cells/psc_nn_overview.png b/examples/perovskite_solar_cells/psc_nn_overview.png new file mode 100644 index 0000000000..f8afd96416 Binary files /dev/null and b/examples/perovskite_solar_cells/psc_nn_overview.png differ diff --git a/examples/perovskite_solar_cells/solar_cell_pretrained.pdparams b/examples/perovskite_solar_cells/solar_cell_pretrained.pdparams new file mode 100644 index 0000000000..b22937a19b Binary files /dev/null and b/examples/perovskite_solar_cells/solar_cell_pretrained.pdparams differ diff --git a/jointContribution/DU_CNN/requirements.txt b/jointContribution/DU_CNN/requirements.txt index df2b9a3181..5a07c6511b 100644 --- a/jointContribution/DU_CNN/requirements.txt +++ b/jointContribution/DU_CNN/requirements.txt @@ -1,7 +1,7 @@ +mat73 +matplotlib numpy +paddlepaddle-gpu pyaml -tqdm -matplotlib sklearn -mat73 -paddlepaddle-gpu +tqdm diff --git a/jointContribution/PINO/PINO_paddle/download_data.py b/jointContribution/PINO/PINO_paddle/download_data.py index 5651617381..32f41d9ad4 100644 --- a/jointContribution/PINO/PINO_paddle/download_data.py +++ b/jointContribution/PINO/PINO_paddle/download_data.py @@ -1,40 +1,44 @@ import os from argparse import ArgumentParser + import requests from tqdm import tqdm _url_dict = { - 'NS-T4000': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fft_Re500_T4000.npy', - 'NS-Re500Part0': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part0.npy', - 'NS-Re500Part1': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part1.npy', - 'NS-Re500Part2': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part2.npy', - 'NS-Re100Part0': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re100_T128_part0.npy', - 'burgers': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/burgers_pino.mat', - 'NS-Re500_T300_id0': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/NS-Re500_T300_id0.npy', - 'NS-Re500_T300_id0-shuffle': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS-Re500_T300_id0-shuffle.npy', - 'darcy-train': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/piececonst_r421_N1024_smooth1.mat', - 'Re500-1_8s-800-pino-140k': 'https://hkzdata.s3.us-west-2.amazonaws.com/PINO/checkpoints/Re500-1_8s-800-PINO-140000.pt', + "NS-T4000": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fft_Re500_T4000.npy", + "NS-Re500Part0": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part0.npy", + "NS-Re500Part1": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part1.npy", + "NS-Re500Part2": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re500_T128_part2.npy", + "NS-Re100Part0": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS_fine_Re100_T128_part0.npy", + "burgers": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/burgers_pino.mat", + "NS-Re500_T300_id0": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/NS-Re500_T300_id0.npy", + "NS-Re500_T300_id0-shuffle": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/NS-Re500_T300_id0-shuffle.npy", + "darcy-train": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/data/piececonst_r421_N1024_smooth1.mat", + "Re500-1_8s-800-pino-140k": "https://hkzdata.s3.us-west-2.amazonaws.com/PINO/checkpoints/Re500-1_8s-800-PINO-140000.pt", } + def download_file(url, file_path): - print('Start downloading...') + print("Start downloading...") with requests.get(url, stream=True) as r: r.raise_for_status() - with open(file_path, 'wb') as f: + with open(file_path, "wb") as f: for chunk in tqdm(r.iter_content(chunk_size=256 * 1024 * 1024)): f.write(chunk) - print('Complete') + print("Complete") + def main(args): url = _url_dict[args.name] - file_name = url.split('/')[-1] + file_name = url.split("/")[-1] os.makedirs(args.outdir, exist_ok=True) file_path = os.path.join(args.outdir, file_name) download_file(url, file_path) -if __name__ == '__main__': - parser = ArgumentParser(description='Parser for downloading assets') - parser.add_argument('--name', type=str, default='NS-T4000') - parser.add_argument('--outdir', type=str, default='../data') + +if __name__ == "__main__": + parser = ArgumentParser(description="Parser for downloading assets") + parser.add_argument("--name", type=str, default="NS-T4000") + parser.add_argument("--outdir", type=str, default="../data") args = parser.parse_args() - main(args) \ No newline at end of file + main(args) diff --git a/jointContribution/PINO/PINO_paddle/eval_operator.py b/jointContribution/PINO/PINO_paddle/eval_operator.py index 40991f86c9..6996381302 100644 --- a/jointContribution/PINO/PINO_paddle/eval_operator.py +++ b/jointContribution/PINO/PINO_paddle/eval_operator.py @@ -1,78 +1,94 @@ -import yaml +from argparse import ArgumentParser import paddle +import yaml +from models import FNO2d +from models import FNO3d from paddle.io import DataLoader -from models import FNO3d, FNO2d -from train_utils import NSLoader, get_forcing, DarcyFlow - -from train_utils.eval_3d import eval_ns +from train_utils import DarcyFlow +from train_utils import NSLoader +from train_utils import get_forcing from train_utils.eval_2d import eval_darcy +from train_utils.eval_3d import eval_ns -from argparse import ArgumentParser def test_3d(config): - device = 0 if paddle.cuda.is_available() else 'cpu' - data_config = config['data'] - loader = NSLoader(datapath1=data_config['datapath'], - nx=data_config['nx'], nt=data_config['nt'], - sub=data_config['sub'], sub_t=data_config['sub_t'], - N=data_config['total_num'], - t_interval=data_config['time_interval']) + device = 0 if paddle.cuda.is_available() else "cpu" + data_config = config["data"] + loader = NSLoader( + datapath1=data_config["datapath"], + nx=data_config["nx"], + nt=data_config["nt"], + sub=data_config["sub"], + sub_t=data_config["sub_t"], + N=data_config["total_num"], + t_interval=data_config["time_interval"], + ) - eval_loader = loader.make_loader(n_sample=data_config['n_sample'], - batch_size=config['test']['batchsize'], - start=data_config['offset'], - train=data_config['shuffle']) - model = FNO3d(modes1=config['model']['modes1'], - modes2=config['model']['modes2'], - modes3=config['model']['modes3'], - fc_dim=config['model']['fc_dim'], - layers=config['model']['layers']).to(device) + eval_loader = loader.make_loader( + n_sample=data_config["n_sample"], + batch_size=config["test"]["batchsize"], + start=data_config["offset"], + train=data_config["shuffle"], + ) + model = FNO3d( + modes1=config["model"]["modes1"], + modes2=config["model"]["modes2"], + modes3=config["model"]["modes3"], + fc_dim=config["model"]["fc_dim"], + layers=config["model"]["layers"], + ).to(device) - if 'ckpt' in config['test']: - ckpt_path = config['test']['ckpt'] + if "ckpt" in config["test"]: + ckpt_path = config["test"]["ckpt"] ckpt = paddle.load(ckpt_path) - model.load_state_dict(ckpt['model']) - print('Weights loaded from %s' % ckpt_path) - print(f'Resolution : {loader.S}x{loader.S}x{loader.T}') + model.load_state_dict(ckpt["model"]) + print("Weights loaded from %s" % ckpt_path) + print(f"Resolution : {loader.S}x{loader.S}x{loader.T}") forcing = get_forcing(loader.S).to(device) - eval_ns(model, - loader, - eval_loader, - forcing, - config, - device=device) + eval_ns(model, loader, eval_loader, forcing, config, device=device) + def test_2d(config): - data_config = config['data'] - dataset = DarcyFlow(data_config['datapath'], - nx=data_config['nx'], sub=data_config['sub'], - offset=data_config['offset'], num=data_config['n_sample']) - dataloader = DataLoader(dataset, batch_size=config['test']['batchsize'], shuffle=False) - model = FNO2d(modes1=config['model']['modes1'], - modes2=config['model']['modes2'], - fc_dim=config['model']['fc_dim'], - layers=config['model']['layers'], - act=config['model']['act']) + data_config = config["data"] + dataset = DarcyFlow( + data_config["datapath"], + nx=data_config["nx"], + sub=data_config["sub"], + offset=data_config["offset"], + num=data_config["n_sample"], + ) + dataloader = DataLoader( + dataset, batch_size=config["test"]["batchsize"], shuffle=False + ) + model = FNO2d( + modes1=config["model"]["modes1"], + modes2=config["model"]["modes2"], + fc_dim=config["model"]["fc_dim"], + layers=config["model"]["layers"], + act=config["model"]["act"], + ) # Load from checkpoint - if 'ckpt' in config['test']: - ckpt_path = config['test']['ckpt'] + if "ckpt" in config["test"]: + ckpt_path = config["test"]["ckpt"] ckpt = paddle.load(ckpt_path) - model.set_dict(ckpt['model']) - print('Weights loaded from %s' % ckpt_path) + model.set_dict(ckpt["model"]) + print("Weights loaded from %s" % ckpt_path) eval_darcy(model, dataloader, config) -if __name__ == '__main__': - parser = ArgumentParser(description='Basic paser') - parser.add_argument('--config_path', type=str, help='Path to the configuration file') - parser.add_argument('--log', action='store_true', help='Turn on the wandb') + +if __name__ == "__main__": + parser = ArgumentParser(description="Basic paser") + parser.add_argument( + "--config_path", type=str, help="Path to the configuration file" + ) + parser.add_argument("--log", action="store_true", help="Turn on the wandb") options = parser.parse_args() config_file = options.config_path - with open(config_file, 'r') as stream: + with open(config_file, "r") as stream: config = yaml.load(stream, yaml.FullLoader) - if 'name' in config['data'] and config['data']['name'] == 'Darcy': + if "name" in config["data"] and config["data"]["name"] == "Darcy": test_2d(config) else: test_3d(config) - diff --git a/jointContribution/PINO/PINO_paddle/models/FCN.py b/jointContribution/PINO/PINO_paddle/models/FCN.py index d429d0f43f..7435c76eca 100644 --- a/jointContribution/PINO/PINO_paddle/models/FCN.py +++ b/jointContribution/PINO/PINO_paddle/models/FCN.py @@ -1,29 +1,31 @@ import paddle.nn as nn + def linear_block(in_channel, out_channel): - block = nn.Sequential( - nn.Linear(in_channel, out_channel), - nn.Tanh() - ) + block = nn.Sequential(nn.Linear(in_channel, out_channel), nn.Tanh()) return block + class FCNet(nn.Layer): - ''' + """ Fully connected layers with Tanh as nonlinearity Reproduced from PINNs Burger equation - ''' + """ def __init__(self, layers=[2, 10, 1]): super(FCNet, self).__init__() - fc_list = [linear_block(in_size, out_size) - for in_size, out_size in zip(layers, layers[1:-1])] + fc_list = [ + linear_block(in_size, out_size) + for in_size, out_size in zip(layers, layers[1:-1]) + ] fc_list.append(nn.Linear(layers[-2], layers[-1])) self.fc = nn.Sequential(*fc_list) def forward(self, x): return self.fc(x) + class DenseNet(nn.Layer): def __init__(self, layers, nonlinearity, out_nonlinearity=None, normalize=False): super(DenseNet, self).__init__() @@ -31,20 +33,20 @@ def __init__(self, layers, nonlinearity, out_nonlinearity=None, normalize=False) self.n_layers = len(layers) - 1 assert self.n_layers >= 1 if isinstance(nonlinearity, str): - if nonlinearity == 'tanh': + if nonlinearity == "tanh": nonlinearity = nn.Tanh - elif nonlinearity == 'relu': + elif nonlinearity == "relu": nonlinearity == nn.ReLU else: - raise ValueError(f'{nonlinearity} is not supported') + raise ValueError(f"{nonlinearity} is not supported") self.layers = nn.ModuleList() for j in range(self.n_layers): - self.layers.append(nn.Linear(layers[j], layers[j+1])) + self.layers.append(nn.Linear(layers[j], layers[j + 1])) if j != self.n_layers - 1: if normalize: - self.layers.append(nn.BatchNorm1d(layers[j+1])) + self.layers.append(nn.BatchNorm1d(layers[j + 1])) self.layers.append(nonlinearity()) @@ -56,4 +58,3 @@ def forward(self, x): x = l(x) return x - diff --git a/jointContribution/PINO/PINO_paddle/models/FNO_blocks.py b/jointContribution/PINO/PINO_paddle/models/FNO_blocks.py index c158f4f217..edd4432a0b 100644 --- a/jointContribution/PINO/PINO_paddle/models/FNO_blocks.py +++ b/jointContribution/PINO/PINO_paddle/models/FNO_blocks.py @@ -1,9 +1,11 @@ +import itertools + import paddle import paddle.nn as nn -import itertools einsum_symbols = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" + def _contract_dense(x, weight, separable=False): # order = tl.ndim(x) order = len(x.shape) @@ -28,6 +30,7 @@ def _contract_dense(x, weight, separable=False): return paddle.einsum(eq, x, weight) + def _contract_dense_trick(x, weight, separable=False): # the same as above function, but do the complex multiplication manually to avoid the einsum bug in paddle weight_real = weight.data.real() @@ -59,11 +62,13 @@ def _contract_dense_trick(x, weight, separable=False): x = paddle.complex(o1_real, o1_imag) return x + def _contract_dense_separable(x, weight, separable=True): if separable == False: raise ValueError("This function is only for separable=True") return x * weight + def _contract_cp(x, cp_weight, separable=False): # order = tl.ndim(x) order = len(x.shape) @@ -84,6 +89,7 @@ def _contract_cp(x, cp_weight, separable=False): return paddle.einsum(eq, x, cp_weight.weights, *cp_weight.factors) + def _contract_tucker(x, tucker_weight, separable=False): # order = tl.ndim(x) order = len(x.shape) @@ -118,6 +124,7 @@ def _contract_tucker(x, tucker_weight, separable=False): print(eq) # 'abcd,fghi,bf,eg,ch,di->aecd' return paddle.einsum(eq, x, tucker_weight.core, *tucker_weight.factors) + def _contract_tt(x, tt_weight, separable=False): # order = tl.ndim(x) order = len(x.shape) @@ -144,16 +151,17 @@ def _contract_tt(x, tt_weight, separable=False): return paddle.einsum(eq, x, *tt_weight.factors) + def get_contract_fun(weight, implementation="reconstructed", separable=False): """Generic ND implementation of Fourier Spectral Conv contraction - + Parameters ---------- weight : tensorl-paddle's FactorizedTensor implementation : {'reconstructed', 'factorized'}, default is 'reconstructed' whether to reconstruct the weight and do a forward pass (reconstructed) or contract directly the factors of the factorized weight with the input (factorized) - + Returns ------- function : (x, weight) -> x * weight in Fourier space @@ -176,6 +184,7 @@ def get_contract_fun(weight, implementation="reconstructed", separable=False): f'Got implementation = {implementation}, expected "reconstructed" or "factorized"' ) + class FactorizedTensor(nn.Layer): def __init__(self, shape, init_scale): super().__init__() @@ -195,6 +204,7 @@ def __repr__(self): def data(self): return paddle.complex(self.real, self.imag) + class FactorizedSpectralConv(nn.Layer): """Generic N-Dimensional Fourier Neural Operator @@ -324,14 +334,14 @@ def __init__( def forward(self, x, indices=0): """Generic forward pass for the Factorized Spectral Conv - - Parameters + + Parameters ---------- x : paddle.Tensor input activation of size (batch_size, channels, d1, ..., dN) indices : int, default is 0 if joint_factorization, index of the layers for n_layers > 1 - + Returns ------- tensorized_spectral_conv(x) @@ -348,7 +358,8 @@ def forward(self, x, indices=0): x = paddle.fft.rfftn(x_float, norm=self.fft_norm, axes=fft_dims) out_fft = paddle.zeros( - [batchsize, self.out_channels, *fft_size], dtype=paddle.complex64, + [batchsize, self.out_channels, *fft_size], + dtype=paddle.complex64, ) # [1,32,16,9], all zeros, complex # We contract all corners of the Fourier coefs @@ -402,13 +413,14 @@ def get_conv(self, indices): def __getitem__(self, indices): return self.get_conv(indices) + class SubConv2d(nn.Layer): """Class representing one of the convolutions from the mother joint factorized convolution Notes ----- This relies on the fact that nn.Parameters are not duplicated: - if the same nn.Parameter is assigned to multiple modules, they all point to the same data, + if the same nn.Parameter is assigned to multiple modules, they all point to the same data, which is shared. """ @@ -420,6 +432,7 @@ def __init__(self, main_conv, indices): def forward(self, x): return self.main_conv.forward(x, self.indices) + class FactorizedSpectralConv1d(FactorizedSpectralConv): def __init__( self, @@ -462,7 +475,8 @@ def forward(self, x, indices=0): x = paddle.fft.rfft(x, norm=self.fft_norm) out_fft = paddle.zeros( - [batchsize, self.out_channels, width // 2 + 1], dtype=paddle.complex64, + [batchsize, self.out_channels, width // 2 + 1], + dtype=paddle.complex64, ) out_fft[:, :, : self.half_modes_height] = self._contract( x[:, :, : self.half_modes_height], @@ -477,6 +491,7 @@ def forward(self, x, indices=0): return x + class FactorizedSpectralConv2d(FactorizedSpectralConv): def __init__( self, @@ -522,7 +537,8 @@ def forward(self, x, indices=0): # The output will be of size (batch_size, self.out_channels, x.size(-2), x.size(-1)//2 + 1) out_fft = paddle.zeros( - [batchsize, self.out_channels, height, width // 2 + 1], dtype=x.dtype, + [batchsize, self.out_channels, height, width // 2 + 1], + dtype=x.dtype, ) # upper block (truncate high freq) @@ -551,6 +567,7 @@ def forward(self, x, indices=0): return x + class FactorizedSpectralConv3d(FactorizedSpectralConv): def __init__( self, @@ -679,6 +696,7 @@ def forward(self, x, indices=0): return x + if __name__ == "__main__": # let x be a complex tensor of size (32, 32, 8, 8) x = paddle.randn([32, 32, 8, 8]).astype("complex64") diff --git a/jointContribution/PINO/PINO_paddle/models/basics.py b/jointContribution/PINO/PINO_paddle/models/basics.py index 8a8fea46f3..dfb216b5ef 100644 --- a/jointContribution/PINO/PINO_paddle/models/basics.py +++ b/jointContribution/PINO/PINO_paddle/models/basics.py @@ -1,33 +1,37 @@ import numpy as np - import paddle import paddle.nn as nn import paddle.nn.initializer as Initializer + def compl_mul1d(a: paddle.Tensor, b: paddle.Tensor) -> paddle.Tensor: # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) res = paddle.einsum("bix,iox->box", a, b) return res + def compl_mul2d(a: paddle.Tensor, b: paddle.Tensor) -> paddle.Tensor: # (batch, in_channel, x,y,t ), (in_channel, out_channel, x,y,t) -> (batch, out_channel, x,y,t) - res = paddle.einsum("bixy,ioxy->boxy", a, b) + res = paddle.einsum("bixy,ioxy->boxy", a, b) return res + def compl_mul3d(a: paddle.Tensor, b: paddle.Tensor) -> paddle.Tensor: res = paddle.einsum("bixyz,ioxyz->boxyz", a, b) return res + ################################################################ # 1d fourier layer ################################################################ + class SpectralConv1d(nn.Layer): def __init__(self, in_channels, out_channels, modes1): super(SpectralConv1d, self).__init__() """ - 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. + 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. """ self.in_channels = in_channels @@ -35,8 +39,14 @@ def __init__(self, in_channels, out_channels, modes1): # Number of Fourier modes to multiply, at most floor(N/2) + 1 self.modes1 = modes1 - self.scale = (1 / (in_channels*out_channels)) - self.weights1 = paddle.create_parameter(shape=(in_channels, out_channels), dtype=paddle.complex64, attr=Initializer.Assign(self.scale * paddle.rand(in_channels, out_channels, self.modes1))) + self.scale = 1 / (in_channels * out_channels) + self.weights1 = paddle.create_parameter( + shape=(in_channels, out_channels), + dtype=paddle.complex64, + attr=Initializer.Assign( + self.scale * paddle.rand(in_channels, out_channels, self.modes1) + ), + ) def forward(self, x): batchsize = x.shape[0] @@ -44,17 +54,23 @@ def forward(self, x): x_ft = paddle.fft.rfftn(x, dim=[2]) # Multiply relevant Fourier modes - out_ft = paddle.zeros(batchsize, self.in_channels, x.shape[-1]//2 + 1, dtype=paddle.complex64) - out_ft[:, :, :self.modes1] = compl_mul1d(x_ft[:, :, :self.modes1], self.weights1) + out_ft = paddle.zeros( + batchsize, self.in_channels, x.shape[-1] // 2 + 1, dtype=paddle.complex64 + ) + out_ft[:, :, : self.modes1] = compl_mul1d( + x_ft[:, :, : self.modes1], self.weights1 + ) # Return to physical space x = paddle.fft.irfftn(out_ft, s=[x.shape[-1]], axes=[2]) return x + ################################################################ # 2d fourier layer ################################################################ + class SpectralConv2d(nn.Layer): def __init__(self, in_channels, out_channels, modes1, modes2): super(SpectralConv2d, self).__init__() @@ -64,12 +80,32 @@ def __init__(self, in_channels, out_channels, modes1, modes2): self.modes1 = modes1 self.modes2 = modes2 - self.scale = (1 / (in_channels * out_channels)) - weights1_attr = Initializer.Assign(self.scale * paddle.rand([in_channels, out_channels, self.modes1, self.modes2], dtype=paddle.float32)) - self.weights1_real = paddle.create_parameter(shape=(in_channels, out_channels), dtype=paddle.float32, attr=weights1_attr) - self.weights1_imag = paddle.create_parameter(shape=(in_channels, out_channels), dtype=paddle.float32, attr=weights1_attr) + self.scale = 1 / (in_channels * out_channels) + weights1_attr = Initializer.Assign( + self.scale + * paddle.rand( + [in_channels, out_channels, self.modes1, self.modes2], + dtype=paddle.float32, + ) + ) + self.weights1_real = paddle.create_parameter( + shape=(in_channels, out_channels), dtype=paddle.float32, attr=weights1_attr + ) + self.weights1_imag = paddle.create_parameter( + shape=(in_channels, out_channels), dtype=paddle.float32, attr=weights1_attr + ) self.weights1 = paddle.concat([self.weights1_real, self.weights1_imag], axis=0) - self.weights2 = paddle.create_parameter(shape=(in_channels, out_channels), dtype=paddle.float32, attr=Initializer.Assign(self.scale * paddle.rand([in_channels, out_channels, self.modes1, self.modes2], dtype=paddle.float32))) + self.weights2 = paddle.create_parameter( + shape=(in_channels, out_channels), + dtype=paddle.float32, + attr=Initializer.Assign( + self.scale + * paddle.rand( + [in_channels, out_channels, self.modes1, self.modes2], + dtype=paddle.float32, + ) + ), + ) def forward(self, x): batchsize = x.shape[0] @@ -79,87 +115,177 @@ def forward(self, x): x_ft = paddle.fft.rfftn(x, axes=[2, 3]) # Multiply relevant Fourier modes - out_ft = paddle.zeros([batchsize, self.out_channels, x.shape[-2], x.shape[-1] // 2 + 1], - dtype=paddle.float32) - out_ft[:, :, :self.modes1, :self.modes2] = \ - compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1) - out_ft[:, :, -self.modes1:, :self.modes2] = \ - compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2) + out_ft = paddle.zeros( + [batchsize, self.out_channels, x.shape[-2], x.shape[-1] // 2 + 1], + dtype=paddle.float32, + ) + out_ft[:, :, : self.modes1, : self.modes2] = compl_mul2d( + x_ft[:, :, : self.modes1, : self.modes2], self.weights1 + ) + out_ft[:, :, -self.modes1 :, : self.modes2] = compl_mul2d( + x_ft[:, :, -self.modes1 :, : self.modes2], self.weights2 + ) # Return to physical space x = paddle.fft.irfftn(out_ft, s=(x.shape[-2], x.shape[-1]), axes=[2, 3]) return x + class SpectralConv3d(nn.Layer): def __init__(self, in_channels, out_channels, modes1, modes2, modes3): super(SpectralConv3d, self).__init__() self.in_channels = in_channels self.out_channels = out_channels - self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 + self.modes1 = ( + modes1 # Number of Fourier modes to multiply, at most floor(N/2) + 1 + ) self.modes2 = modes2 self.modes3 = modes3 - self.scale = (1 / (in_channels * out_channels)) - self.weights1 = paddle.create_parameter(shape=(in_channels, out_channels), dtype=paddle.complex64, attr=Initializer.Assign(self.scale * paddle.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3))) - self.weights2 = paddle.create_parameter(shape=(in_channels, out_channels), dtype=paddle.complex64, attr=Initializer.Assign(self.scale * paddle.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3))) - self.weights3 = paddle.create_parameter(shape=(in_channels, out_channels), dtype=paddle.complex64, attr=Initializer.Assign(self.scale * paddle.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3))) - self.weights4 = paddle.create_parameter(shape=(in_channels, out_channels), dtype=paddle.complex64, attr=Initializer.Assign(self.scale * paddle.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3))) + self.scale = 1 / (in_channels * out_channels) + self.weights1 = paddle.create_parameter( + shape=(in_channels, out_channels), + dtype=paddle.complex64, + attr=Initializer.Assign( + self.scale + * paddle.rand( + in_channels, out_channels, self.modes1, self.modes2, self.modes3 + ) + ), + ) + self.weights2 = paddle.create_parameter( + shape=(in_channels, out_channels), + dtype=paddle.complex64, + attr=Initializer.Assign( + self.scale + * paddle.rand( + in_channels, out_channels, self.modes1, self.modes2, self.modes3 + ) + ), + ) + self.weights3 = paddle.create_parameter( + shape=(in_channels, out_channels), + dtype=paddle.complex64, + attr=Initializer.Assign( + self.scale + * paddle.rand( + in_channels, out_channels, self.modes1, self.modes2, self.modes3 + ) + ), + ) + self.weights4 = paddle.create_parameter( + shape=(in_channels, out_channels), + dtype=paddle.complex64, + attr=Initializer.Assign( + self.scale + * paddle.rand( + in_channels, out_channels, self.modes1, self.modes2, self.modes3 + ) + ), + ) def forward(self, x): batchsize = x.shape[0] # Compute Fourier coeffcients up to factor of e^(- something constant) - x_ft = paddle.fft.rfftn(x, dim=[2,3,4]) - + x_ft = paddle.fft.rfftn(x, dim=[2, 3, 4]) + z_dim = min(x_ft.shape[4], self.modes3) - + # Multiply relevant Fourier modes - out_ft = paddle.zeros(batchsize, self.out_channels, x_ft.shape[2], x_ft.shape[3], self.modes3, dtype=paddle.complex64) - - # if x_ft.shape[4] > self.modes3, truncate; if x_ft.shape[4] < self.modes3, add zero padding - coeff = paddle.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, dtype=paddle.complex64) - coeff[..., :z_dim] = x_ft[:, :, :self.modes1, :self.modes2, :z_dim] - out_ft[:, :, :self.modes1, :self.modes2, :] = compl_mul3d(coeff, self.weights1) - - coeff = paddle.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, dtype=paddle.complex64) - coeff[..., :z_dim] = x_ft[:, :, -self.modes1:, :self.modes2, :z_dim] - out_ft[:, :, -self.modes1:, :self.modes2, :] = compl_mul3d(coeff, self.weights2) - - coeff = paddle.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, dtype=paddle.complex64) - coeff[..., :z_dim] = x_ft[:, :, :self.modes1, -self.modes2:, :z_dim] - out_ft[:, :, :self.modes1, -self.modes2:, :] = compl_mul3d(coeff, self.weights3) - - coeff = paddle.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, dtype=paddle.complex64) - coeff[..., :z_dim] = x_ft[:, :, -self.modes1:, -self.modes2:, :z_dim] - out_ft[:, :, -self.modes1:, -self.modes2:, :] = compl_mul3d(coeff, self.weights4) - - #Return to physical space - x = paddle.fft.irfftn(out_ft, s=(x.shape[2], x.shape[3], x.shape[4]), axes=[2,3,4]) + out_ft = paddle.zeros( + batchsize, + self.out_channels, + x_ft.shape[2], + x_ft.shape[3], + self.modes3, + dtype=paddle.complex64, + ) + + # if x_ft.shape[4] > self.modes3, truncate; if x_ft.shape[4] < self.modes3, add zero padding + coeff = paddle.zeros( + batchsize, + self.in_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=paddle.complex64, + ) + coeff[..., :z_dim] = x_ft[:, :, : self.modes1, : self.modes2, :z_dim] + out_ft[:, :, : self.modes1, : self.modes2, :] = compl_mul3d( + coeff, self.weights1 + ) + + coeff = paddle.zeros( + batchsize, + self.in_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=paddle.complex64, + ) + coeff[..., :z_dim] = x_ft[:, :, -self.modes1 :, : self.modes2, :z_dim] + out_ft[:, :, -self.modes1 :, : self.modes2, :] = compl_mul3d( + coeff, self.weights2 + ) + + coeff = paddle.zeros( + batchsize, + self.in_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=paddle.complex64, + ) + coeff[..., :z_dim] = x_ft[:, :, : self.modes1, -self.modes2 :, :z_dim] + out_ft[:, :, : self.modes1, -self.modes2 :, :] = compl_mul3d( + coeff, self.weights3 + ) + + coeff = paddle.zeros( + batchsize, + self.in_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=paddle.complex64, + ) + coeff[..., :z_dim] = x_ft[:, :, -self.modes1 :, -self.modes2 :, :z_dim] + out_ft[:, :, -self.modes1 :, -self.modes2 :, :] = compl_mul3d( + coeff, self.weights4 + ) + + # Return to physical space + x = paddle.fft.irfftn( + out_ft, s=(x.shape[2], x.shape[3], x.shape[4]), axes=[2, 3, 4] + ) return x + class FourierBlock(nn.Layer): - def __init__(self, in_channels, out_channels, modes1, modes2, modes3, act='tanh'): + def __init__(self, in_channels, out_channels, modes1, modes2, modes3, act="tanh"): super(FourierBlock, self).__init__() self.in_channel = in_channels self.out_channel = out_channels self.speconv = SpectralConv3d(in_channels, out_channels, modes1, modes2, modes3) self.linear = nn.Conv1D(in_channels, out_channels, 1) - if act == 'tanh': + if act == "tanh": self.act = paddle.tanh_ - elif act == 'gelu': + elif act == "gelu": self.act = nn.GELU - elif act == 'none': + elif act == "none": self.act = None else: - raise ValueError(f'{act} is not supported') + raise ValueError(f"{act} is not supported") def forward(self, x): - ''' + """ input x: (batchsize, channel width, x_grid, y_grid, t_grid) - ''' + """ x1 = self.speconv(x) x2 = self.linear(x.reshape(x.shape[0], self.in_channel, -1)) - out = x1 + x2.reshape(x.shape[0], self.out_channel, x.shape[2], x.shape[3], x.shape[4]) + out = x1 + x2.reshape( + x.shape[0], self.out_channel, x.shape[2], x.shape[3], x.shape[4] + ) if self.act is not None: out = self.act(out) return out - diff --git a/jointContribution/PINO/PINO_paddle/models/fourier2d.py b/jointContribution/PINO/PINO_paddle/models/fourier2d.py index 327369aed1..dffcebbbb4 100644 --- a/jointContribution/PINO/PINO_paddle/models/fourier2d.py +++ b/jointContribution/PINO/PINO_paddle/models/fourier2d.py @@ -1,34 +1,44 @@ import paddle.nn as nn + from .basics import SpectralConv2d from .FNO_blocks import FactorizedSpectralConv2d -from .utils import _get_act, add_padding2, remove_padding2 +from .utils import _get_act +from .utils import add_padding2 +from .utils import remove_padding2 + class FNO2d(nn.Layer): - def __init__(self, modes1, modes2, - width=64, fc_dim=128, - layers=None, - in_dim=3, out_dim=1, - act='gelu', - pad_ratio=[0., 0.]): + def __init__( + self, + modes1, + modes2, + width=64, + fc_dim=128, + layers=None, + in_dim=3, + out_dim=1, + act="gelu", + pad_ratio=[0.0, 0.0], + ): super(FNO2d, self).__init__() """ Args: - modes1: list of int, number of modes in first dimension in each layer - modes2: list of int, number of modes in second dimension in each layer - - width: int, optional, if layers is None, it will be initialized as [width] * [len(modes1) + 1] + - width: int, optional, if layers is None, it will be initialized as [width] * [len(modes1) + 1] - in_dim: number of input channels - out_dim: number of output channels - act: activation function, {tanh, gelu, relu, leaky_relu}, default: gelu - - pad_ratio: list of float, or float; portion of domain to be extended. If float, paddings are added to the right. - If list, paddings are added to both sides. pad_ratio[0] pads left, pad_ratio[1] pads right. + - pad_ratio: list of float, or float; portion of domain to be extended. If float, paddings are added to the right. + If list, paddings are added to both sides. pad_ratio[0] pads left, pad_ratio[1] pads right. """ if isinstance(pad_ratio, float): pad_ratio = [pad_ratio, pad_ratio] else: - assert len(pad_ratio) == 2, 'Cannot add padding in more than 2 directions' + assert len(pad_ratio) == 2, "Cannot add padding in more than 2 directions" self.modes1 = modes1 self.modes2 = modes2 - + self.pad_ratio = pad_ratio # input channel is 3: (a(x, y), x, y) if layers is None: @@ -37,13 +47,21 @@ def __init__(self, modes1, modes2, self.layers = layers self.fc0 = nn.Linear(in_dim, layers[0]) - self.sp_convs = nn.LayerList([FactorizedSpectralConv2d( - in_size, out_size, mode1_num, mode2_num) - for in_size, out_size, mode1_num, mode2_num - in zip(self.layers, self.layers[1:], self.modes1, self.modes2)]) + self.sp_convs = nn.LayerList( + [ + FactorizedSpectralConv2d(in_size, out_size, mode1_num, mode2_num) + for in_size, out_size, mode1_num, mode2_num in zip( + self.layers, self.layers[1:], self.modes1, self.modes2 + ) + ] + ) - self.ws = nn.LayerList([nn.Conv1D(in_size, out_size, 1) - for in_size, out_size in zip(self.layers, self.layers[1:])]) + self.ws = nn.LayerList( + [ + nn.Conv1D(in_size, out_size, 1) + for in_size, out_size in zip(self.layers, self.layers[1:]) + ] + ) self.fc1 = nn.Linear(layers[-1], fc_dim) self.fc2 = nn.Linear(fc_dim, layers[-1]) @@ -51,28 +69,30 @@ def __init__(self, modes1, modes2, self.act = _get_act(act) def forward(self, x): - ''' + """ Args: - x : (batch size, x_grid, y_grid, 2) Returns: - x: (batch size, x_grid, y_grid, 1) - ''' + """ size_1, size_2 = x.shape[1], x.shape[2] if max(self.pad_ratio) > 0: num_pad1 = [round(i * size_1) for i in self.pad_ratio] num_pad2 = [round(i * size_2) for i in self.pad_ratio] else: - num_pad1 = num_pad2 = [0.] + num_pad1 = num_pad2 = [0.0] length = len(self.ws) batchsize = x.shape[0] x = self.fc0(x) - x = x.transpose([0, 3, 1, 2]) # B, C, X, Y + x = x.transpose([0, 3, 1, 2]) # B, C, X, Y x = add_padding2(x, num_pad1, num_pad2) size_x, size_y = x.shape[-2], x.shape[-1] for i, (speconv, w) in enumerate(zip(self.sp_convs, self.ws)): x1 = speconv(x) - x2 = w(x.reshape([batchsize, self.layers[i], -1])).reshape([batchsize, self.layers[i+1], size_x, size_y]) + x2 = w(x.reshape([batchsize, self.layers[i], -1])).reshape( + [batchsize, self.layers[i + 1], size_x, size_y] + ) x = x1 + x2 if i != length - 1: x = self.act(x) diff --git a/jointContribution/PINO/PINO_paddle/models/fourier3d.py b/jointContribution/PINO/PINO_paddle/models/fourier3d.py index 077794f645..ed5c4b6093 100644 --- a/jointContribution/PINO/PINO_paddle/models/fourier3d.py +++ b/jointContribution/PINO/PINO_paddle/models/fourier3d.py @@ -1,18 +1,27 @@ import paddle.nn as nn + from .basics import SpectralConv3d from .FNO_blocks import FactorizedSpectralConv3d -from .utils import add_padding, remove_padding, _get_act +from .utils import _get_act +from .utils import add_padding +from .utils import remove_padding + class FNO3d(nn.Layer): - def __init__(self, - modes1, modes2, modes3, - width=16, - fc_dim=128, - layers=None, - in_dim=4, out_dim=1, - act='gelu', - pad_ratio=[0., 0.]): - ''' + def __init__( + self, + modes1, + modes2, + modes3, + width=16, + fc_dim=128, + layers=None, + in_dim=4, + out_dim=1, + act="gelu", + pad_ratio=[0.0, 0.0], + ): + """ Args: modes1: list of int, first dimension maximal modes for each layer modes2: list of int, second dimension maximal modes for each layer @@ -23,13 +32,13 @@ def __init__(self, out_dim: int, output dimension act: {tanh, gelu, relu, leaky_relu}, activation function pad_ratio: the ratio of the extended domain - ''' + """ super(FNO3d, self).__init__() if isinstance(pad_ratio, float): pad_ratio = [pad_ratio, pad_ratio] else: - assert len(pad_ratio) == 2, 'Cannot add padding in more than 2 directions.' + assert len(pad_ratio) == 2, "Cannot add padding in more than 2 directions." self.pad_ratio = pad_ratio self.modes1 = modes1 @@ -43,35 +52,45 @@ def __init__(self, self.layers = layers self.fc0 = nn.Linear(in_dim, layers[0]) - self.sp_convs = nn.LayerList([FactorizedSpectralConv3d( - in_size, out_size, mode1_num, mode2_num, mode3_num) - for in_size, out_size, mode1_num, mode2_num, mode3_num - in zip(self.layers, self.layers[1:], self.modes1, self.modes2, self.modes3)]) + self.sp_convs = nn.LayerList( + [ + FactorizedSpectralConv3d( + in_size, out_size, mode1_num, mode2_num, mode3_num + ) + for in_size, out_size, mode1_num, mode2_num, mode3_num in zip( + self.layers, self.layers[1:], self.modes1, self.modes2, self.modes3 + ) + ] + ) - self.ws = nn.LayerList([nn.Conv1D(in_size, out_size, 1) - for in_size, out_size in zip(self.layers, self.layers[1:])]) + self.ws = nn.LayerList( + [ + nn.Conv1D(in_size, out_size, 1) + for in_size, out_size in zip(self.layers, self.layers[1:]) + ] + ) self.fc1 = nn.Linear(layers[-1], fc_dim) self.fc2 = nn.Linear(fc_dim, out_dim) self.act = _get_act(act) def forward(self, x): - ''' + """ Args: x: (batchsize, x_grid, y_grid, t_grid, 3) Returns: u: (batchsize, x_grid, y_grid, t_grid, 1) - ''' + """ size_z = x.shape[-2] if max(self.pad_ratio) > 0: num_pad = [round(size_z * i) for i in self.pad_ratio] else: - num_pad = [0., 0.] + num_pad = [0.0, 0.0] length = len(self.ws) batchsize = x.shape[0] - + x = self.fc0(x) x = x.transpose([0, 4, 1, 2, 3]) x = add_padding(x, num_pad=num_pad) @@ -79,7 +98,9 @@ def forward(self, x): for i, (speconv, w) in enumerate(zip(self.sp_convs, self.ws)): x1 = speconv(x) - x2 = w(x.reshape([batchsize, self.layers[i], -1])).reshape([batchsize, self.layers[i+1], size_x, size_y, size_z]) + x2 = w(x.reshape([batchsize, self.layers[i], -1])).reshape( + [batchsize, self.layers[i + 1], size_x, size_y, size_z] + ) x = x1 + x2 if i != length - 1: x = self.act(x) @@ -88,4 +109,4 @@ def forward(self, x): x = self.fc1(x) x = self.act(x) x = self.fc2(x) - return x \ No newline at end of file + return x diff --git a/jointContribution/PINO/PINO_paddle/models/utils.py b/jointContribution/PINO/PINO_paddle/models/utils.py index 0302c5ab3c..2ae65ef56c 100644 --- a/jointContribution/PINO/PINO_paddle/models/utils.py +++ b/jointContribution/PINO/PINO_paddle/models/utils.py @@ -1,45 +1,51 @@ import paddle.nn.functional as F + def add_padding(x, num_pad): if max(num_pad) > 0: - res = F.pad(x, (num_pad[0], num_pad[1]), 'constant', 0) + res = F.pad(x, (num_pad[0], num_pad[1]), "constant", 0) else: res = x return res + def add_padding2(x, num_pad1, num_pad2): if max(num_pad1) > 0 or max(num_pad2) > 0: - res = F.pad(x, (num_pad2[0], num_pad2[1], num_pad1[0], num_pad1[1]), 'constant', 0.) + res = F.pad( + x, (num_pad2[0], num_pad2[1], num_pad1[0], num_pad1[1]), "constant", 0.0 + ) else: res = x return res + def remove_padding(x, num_pad): if max(num_pad) > 0: - res = x[..., num_pad[0]:-num_pad[1]] + res = x[..., num_pad[0] : -num_pad[1]] else: res = x return res + def remove_padding2(x, num_pad1, num_pad2): if max(num_pad1) > 0 or max(num_pad2) > 0: - res = x[..., num_pad1[0]:-num_pad1[1], num_pad2[0]:-num_pad2[1]] + res = x[..., num_pad1[0] : -num_pad1[1], num_pad2[0] : -num_pad2[1]] else: res = x return res + def _get_act(act): - if act == 'tanh': + if act == "tanh": func = F.tanh - elif act == 'gelu': + elif act == "gelu": func = F.gelu - elif act == 'relu': + elif act == "relu": func = F.relu_ - elif act == 'elu': + elif act == "elu": func = F.elu_ - elif act == 'leaky_relu': + elif act == "leaky_relu": func = F.leaky_relu else: - raise ValueError(f'{act} is not supported') + raise ValueError(f"{act} is not supported") return func - diff --git a/jointContribution/PINO/PINO_paddle/prepare_data.py b/jointContribution/PINO/PINO_paddle/prepare_data.py index 74a14694e1..c7a7c55bba 100644 --- a/jointContribution/PINO/PINO_paddle/prepare_data.py +++ b/jointContribution/PINO/PINO_paddle/prepare_data.py @@ -1,33 +1,37 @@ -import numpy as np import matplotlib.pyplot as plt +import numpy as np + def shuffle_data(datapath): data = np.load(datapath) rng = np.random.default_rng(123) rng.shuffle(data, axis=0) - savepath = datapath.replace('.npy', '-shuffle.npy') + savepath = datapath.replace(".npy", "-shuffle.npy") np.save(savepath, data) + def test_data(datapath): - raw = np.load(datapath, mmap_mode='r') + raw = np.load(datapath, mmap_mode="r") print(raw[0, 0, 0, 0:10]) - newpath = datapath.replace('.npy', '-shuffle.npy') - new = np.load(newpath, mmap_mode='r') + newpath = datapath.replace(".npy", "-shuffle.npy") + new = np.load(newpath, mmap_mode="r") print(new[0, 0, 0, 0:10]) + def get_slice(datapath): - raw = np.load(datapath, mmap_mode='r') + raw = np.load(datapath, mmap_mode="r") data = raw[-10:] print(data.shape) - savepath = 'data/Re500-5x513x256x256.npy' + savepath = "data/Re500-5x513x256x256.npy" np.save(savepath, data) + def plot_test(datapath): duration = 0.125 - raw = np.load(datapath, mmap_mode='r') - + raw = np.load(datapath, mmap_mode="r") + -if __name__ == '__main__': - datapath = '/raid/hongkai/NS-Re500_T300_id0-shuffle.npy' - get_slice(datapath) \ No newline at end of file +if __name__ == "__main__": + datapath = "/raid/hongkai/NS-Re500_T300_id0-shuffle.npy" + get_slice(datapath) diff --git a/jointContribution/PINO/PINO_paddle/solver/random_fields.py b/jointContribution/PINO/PINO_paddle/solver/random_fields.py index 448567db3e..458c1a0e72 100644 --- a/jointContribution/PINO/PINO_paddle/solver/random_fields.py +++ b/jointContribution/PINO/PINO_paddle/solver/random_fields.py @@ -1,58 +1,102 @@ +import math + import paddle -import math class GaussianRF(object): - def __init__(self, dim, size, length=1.0, alpha=2.0, tau=3.0, sigma=None, boundary="periodic", constant_eig=False): + def __init__( + self, + dim, + size, + length=1.0, + alpha=2.0, + tau=3.0, + sigma=None, + boundary="periodic", + constant_eig=False, + ): self.dim = dim if sigma is None: - sigma = tau**(0.5*(2*alpha - self.dim)) + sigma = tau ** (0.5 * (2 * alpha - self.dim)) - k_max = size//2 + k_max = size // 2 - const = (4*(math.pi**2))/(length**2) + const = (4 * (math.pi**2)) / (length**2) if dim == 1: - k = paddle.concat((paddle.arange(start=0, end=k_max, step=1), \ - paddle.arange(start=-k_max, end=0, step=1)), 0) - - self.sqrt_eig = size*math.sqrt(2.0)*sigma*((const*(k**2) + tau**2)**(-alpha/2.0)) + k = paddle.concat( + ( + paddle.arange(start=0, end=k_max, step=1), + paddle.arange(start=-k_max, end=0, step=1), + ), + 0, + ) + + self.sqrt_eig = ( + size + * math.sqrt(2.0) + * sigma + * ((const * (k**2) + tau**2) ** (-alpha / 2.0)) + ) if constant_eig: - self.sqrt_eig[0] = size*sigma*(tau**(-alpha)) + self.sqrt_eig[0] = size * sigma * (tau ** (-alpha)) else: self.sqrt_eig[0] = 0.0 elif dim == 2: - wavenumers = paddle.concat((paddle.arange(start=0, end=k_max, step=1), \ - paddle.arange(start=-k_max, end=0, step=1)), 0).repeat(size,1) - - k_x = wavenumers.transpose(0,1) + wavenumers = paddle.concat( + ( + paddle.arange(start=0, end=k_max, step=1), + paddle.arange(start=-k_max, end=0, step=1), + ), + 0, + ).repeat(size, 1) + + k_x = wavenumers.transpose(0, 1) k_y = wavenumers - self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((const*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0)) + self.sqrt_eig = ( + (size**2) + * math.sqrt(2.0) + * sigma + * ((const * (k_x**2 + k_y**2) + tau**2) ** (-alpha / 2.0)) + ) if constant_eig: - self.sqrt_eig[0,0] = (size**2)*sigma*(tau**(-alpha)) + self.sqrt_eig[0, 0] = (size**2) * sigma * (tau ** (-alpha)) else: - self.sqrt_eig[0,0] = 0.0 + self.sqrt_eig[0, 0] = 0.0 elif dim == 3: - wavenumers = paddle.concat((paddle.arange(start=0, end=k_max, step=1), \ - paddle.arange(start=-k_max, end=0, step=1)), 0).repeat(size,size,1) - - k_x = wavenumers.transpose(1,2) + wavenumers = paddle.concat( + ( + paddle.arange(start=0, end=k_max, step=1), + paddle.arange(start=-k_max, end=0, step=1), + ), + 0, + ).repeat(size, size, 1) + + k_x = wavenumers.transpose(1, 2) k_y = wavenumers - k_z = wavenumers.transpose(0,2) - - self.sqrt_eig = (size**3)*math.sqrt(2.0)*sigma*((const*(k_x**2 + k_y**2 + k_z**2) + tau**2)**(-alpha/2.0)) + k_z = wavenumers.transpose(0, 2) + + self.sqrt_eig = ( + (size**3) + * math.sqrt(2.0) + * sigma + * ( + (const * (k_x**2 + k_y**2 + k_z**2) + tau**2) + ** (-alpha / 2.0) + ) + ) if constant_eig: - self.sqrt_eig[0,0,0] = (size**3)*sigma*(tau**(-alpha)) + self.sqrt_eig[0, 0, 0] = (size**3) * sigma * (tau ** (-alpha)) else: - self.sqrt_eig[0,0,0] = 0.0 + self.sqrt_eig[0, 0, 0] = 0.0 self.size = [] for j in range(self.dim): @@ -63,14 +107,26 @@ def __init__(self, dim, size, length=1.0, alpha=2.0, tau=3.0, sigma=None, bounda def sample(self, N): coeff = paddle.randn(N, *self.size, dtype=paddle.float32) - coeff = self.sqrt_eig*coeff + coeff = self.sqrt_eig * coeff u = paddle.fft.irfftn(coeff, self.size, norm="backward") return u -class GaussianRF2d(object): - def __init__(self, s1, s2, L1=2*math.pi, L2=2*math.pi, alpha=2.0, tau=3.0, sigma=None, mean=None, boundary="periodic", dtype=paddle.float64): +class GaussianRF2d(object): + def __init__( + self, + s1, + s2, + L1=2 * math.pi, + L2=2 * math.pi, + alpha=2.0, + tau=3.0, + sigma=None, + mean=None, + boundary="periodic", + dtype=paddle.float64, + ): self.s1 = s1 self.s2 = s2 @@ -80,34 +136,44 @@ def __init__(self, s1, s2, L1=2*math.pi, L2=2*math.pi, alpha=2.0, tau=3.0, sigma self.dtype = dtype if sigma is None: - self.sigma = tau**(0.5*(2*alpha - 2.0)) + self.sigma = tau ** (0.5 * (2 * alpha - 2.0)) else: self.sigma = sigma - const1 = (4*(math.pi**2))/(L1**2) - const2 = (4*(math.pi**2))/(L2**2) + const1 = (4 * (math.pi**2)) / (L1**2) + const2 = (4 * (math.pi**2)) / (L2**2) - freq_list1 = paddle.concat((paddle.arange(start=0, end=s1//2, step=1),\ - paddle.arange(start=-s1//2, end=0, step=1)), 0) - k1 = freq_list1.reshape([-1,1]).repeat([1, s2//2 + 1]).type(dtype) + freq_list1 = paddle.concat( + ( + paddle.arange(start=0, end=s1 // 2, step=1), + paddle.arange(start=-s1 // 2, end=0, step=1), + ), + 0, + ) + k1 = freq_list1.reshape([-1, 1]).repeat([1, s2 // 2 + 1]).type(dtype) - freq_list2 = paddle.arange(start=0, end=s2//2 + 1, step=1) + freq_list2 = paddle.arange(start=0, end=s2 // 2 + 1, step=1) - k2 = freq_list2.view(1,-1).repeat(s1, 1).type(dtype) + k2 = freq_list2.view(1, -1).repeat(s1, 1).type(dtype) - self.sqrt_eig = s1*s2*self.sigma*((const1*k1**2 + const2*k2**2 + tau**2)**(-alpha/2.0)) - self.sqrt_eig[0,0] = 0.0 + self.sqrt_eig = ( + s1 + * s2 + * self.sigma + * ((const1 * k1**2 + const2 * k2**2 + tau**2) ** (-alpha / 2.0)) + ) + self.sqrt_eig[0, 0] = 0.0 def sample(self, N, xi=None): if xi is None: - xi = paddle.randn(N, self.s1, self.s2//2 + 1, 2, dtype=self.dtype) - - xi[...,0] = self.sqrt_eig*xi [...,0] - xi[...,1] = self.sqrt_eig*xi [...,1] - + xi = paddle.randn(N, self.s1, self.s2 // 2 + 1, 2, dtype=self.dtype) + + xi[..., 0] = self.sqrt_eig * xi[..., 0] + xi[..., 1] = self.sqrt_eig * xi[..., 1] + u = paddle.fft.irfft2(paddle.reshape(xi), s=(self.s1, self.s2)) if self.mean is not None: u += self.mean - - return u \ No newline at end of file + + return u diff --git a/jointContribution/PINO/PINO_paddle/train_burgers.py b/jointContribution/PINO/PINO_paddle/train_burgers.py index 7b5d2d614b..a309349f94 100644 --- a/jointContribution/PINO/PINO_paddle/train_burgers.py +++ b/jointContribution/PINO/PINO_paddle/train_burgers.py @@ -1,85 +1,113 @@ -import yaml from argparse import ArgumentParser + import paddle -from paddle.optimizer.lr import MultiStepDecay -from paddle.optimizer import Adam +import yaml from models import FNO2d +from paddle.optimizer import Adam +from paddle.optimizer.lr import MultiStepDecay from train_utils.datasets import BurgersLoader -from train_utils.train_2d import train_2d_burger from train_utils.eval_2d import eval_burgers +from train_utils.train_2d import train_2d_burger + def run(args, config): - data_config = config['data'] - dataset = BurgersLoader(data_config['datapath'], - nx=data_config['nx'], nt=data_config['nt'], - sub=data_config['sub'], sub_t=data_config['sub_t'], new=True) - train_loader = dataset.make_loader(n_sample=data_config['n_sample'], - batch_size=config['train']['batchsize'], - start=data_config['offset']) + data_config = config["data"] + dataset = BurgersLoader( + data_config["datapath"], + nx=data_config["nx"], + nt=data_config["nt"], + sub=data_config["sub"], + sub_t=data_config["sub_t"], + new=True, + ) + train_loader = dataset.make_loader( + n_sample=data_config["n_sample"], + batch_size=config["train"]["batchsize"], + start=data_config["offset"], + ) - model = FNO2d(modes1=config['model']['modes1'], - modes2=config['model']['modes2'], - fc_dim=config['model']['fc_dim'], - layers=config['model']['layers'], - act=config['model']['act']) - param_state_dict = paddle.load('init_param/init_burgers.pdparams') + model = FNO2d( + modes1=config["model"]["modes1"], + modes2=config["model"]["modes2"], + fc_dim=config["model"]["fc_dim"], + layers=config["model"]["layers"], + act=config["model"]["act"], + ) + param_state_dict = paddle.load("init_param/init_burgers.pdparams") model.set_dict(param_state_dict) # Load from checkpoint - if 'ckpt' in config['train']: - ckpt_path = config['train']['ckpt'] + if "ckpt" in config["train"]: + ckpt_path = config["train"]["ckpt"] ckpt = paddle.load(ckpt_path) - model.load_state_dict(ckpt['model']) - print('Weights loaded from %s' % ckpt_path) + model.load_state_dict(ckpt["model"]) + print("Weights loaded from %s" % ckpt_path) - scheduler = MultiStepDecay(learning_rate=config['train']['base_lr'], - milestones=config['train']['milestones'], - gamma=config['train']['scheduler_gamma']) + scheduler = MultiStepDecay( + learning_rate=config["train"]["base_lr"], + milestones=config["train"]["milestones"], + gamma=config["train"]["scheduler_gamma"], + ) optimizer = Adam(learning_rate=scheduler, parameters=model.parameters()) - - train_2d_burger(model, - train_loader, - dataset.v, - optimizer, - scheduler, - config, - rank=0, - log=args.log, - project=config['log']['project'], - group=config['log']['group']) + + train_2d_burger( + model, + train_loader, + dataset.v, + optimizer, + scheduler, + config, + rank=0, + log=args.log, + project=config["log"]["project"], + group=config["log"]["group"], + ) + def test(config): - data_config = config['data'] - dataset = BurgersLoader(data_config['datapath'], - nx=data_config['nx'], nt=data_config['nt'], - sub=data_config['sub'], sub_t=data_config['sub_t'], new=True) - dataloader = dataset.make_loader(n_sample=data_config['n_sample'], - batch_size=config['test']['batchsize'], - start=data_config['offset']) + data_config = config["data"] + dataset = BurgersLoader( + data_config["datapath"], + nx=data_config["nx"], + nt=data_config["nt"], + sub=data_config["sub"], + sub_t=data_config["sub_t"], + new=True, + ) + dataloader = dataset.make_loader( + n_sample=data_config["n_sample"], + batch_size=config["test"]["batchsize"], + start=data_config["offset"], + ) - model = FNO2d(modes1=config['model']['modes1'], - modes2=config['model']['modes2'], - fc_dim=config['model']['fc_dim'], - layers=config['model']['layers'], - act=config['model']['act']) + model = FNO2d( + modes1=config["model"]["modes1"], + modes2=config["model"]["modes2"], + fc_dim=config["model"]["fc_dim"], + layers=config["model"]["layers"], + act=config["model"]["act"], + ) # Load from checkpoint - if 'ckpt' in config['test']: - ckpt_path = config['test']['ckpt'] + if "ckpt" in config["test"]: + ckpt_path = config["test"]["ckpt"] ckpt = paddle.load(ckpt_path) - model.load_state_dict(ckpt['model']) - print('Weights loaded from %s' % ckpt_path) + model.load_state_dict(ckpt["model"]) + print("Weights loaded from %s" % ckpt_path) eval_burgers(model, dataloader, dataset.v, config) -if __name__ == '__main__': - parser = ArgumentParser(description='Basic paser') - parser.add_argument('--config_path', type=str, help='Path to the configuration file') - parser.add_argument('--log', action='store_true', help='Turn on the wandb') - parser.add_argument('--mode', type=str, help='train or test') + +if __name__ == "__main__": + parser = ArgumentParser(description="Basic paser") + parser.add_argument( + "--config_path", type=str, help="Path to the configuration file" + ) + parser.add_argument("--log", action="store_true", help="Turn on the wandb") + parser.add_argument("--mode", type=str, help="train or test") args = parser.parse_args() config_file = args.config_path - with open(config_file, 'r') as stream: + with open(config_file, "r") as stream: config = yaml.load(stream, yaml.FullLoader) - if args.mode == 'train': + if args.mode == "train": run(args, config) else: test(config) diff --git a/jointContribution/PINO/PINO_paddle/train_operator.py b/jointContribution/PINO/PINO_paddle/train_operator.py index 5f9fd6ae1c..93e93d1083 100644 --- a/jointContribution/PINO/PINO_paddle/train_operator.py +++ b/jointContribution/PINO/PINO_paddle/train_operator.py @@ -1,123 +1,165 @@ -import yaml -from argparse import ArgumentParser import math +from argparse import ArgumentParser import paddle +import yaml +from models import FNO2d +from models import FNO3d from paddle.io import DataLoader from paddle.optimizer.lr import MultiStepDecay - from solver.random_fields import GaussianRF from train_utils import Adam -from train_utils.datasets import NSLoader, online_loader, DarcyFlow, DarcyCombo -from train_utils.train_3d import mixed_train +from train_utils.datasets import DarcyCombo +from train_utils.datasets import DarcyFlow +from train_utils.datasets import NSLoader +from train_utils.datasets import online_loader from train_utils.train_2d import train_2d_operator -from models import FNO3d, FNO2d +from train_utils.train_3d import mixed_train + def train_3d(args, config): - data_config = config['data'] + data_config = config["data"] # prepare dataloader for training with data - if 'datapath2' in data_config: - loader = NSLoader(datapath1=data_config['datapath'], datapath2=data_config['datapath2'], - nx=data_config['nx'], nt=data_config['nt'], - sub=data_config['sub'], sub_t=data_config['sub_t'], - N=data_config['total_num'], - t_interval=data_config['time_interval']) + if "datapath2" in data_config: + loader = NSLoader( + datapath1=data_config["datapath"], + datapath2=data_config["datapath2"], + nx=data_config["nx"], + nt=data_config["nt"], + sub=data_config["sub"], + sub_t=data_config["sub_t"], + N=data_config["total_num"], + t_interval=data_config["time_interval"], + ) else: - loader = NSLoader(datapath1=data_config['datapath'], - nx=data_config['nx'], nt=data_config['nt'], - sub=data_config['sub'], sub_t=data_config['sub_t'], - N=data_config['total_num'], - t_interval=data_config['time_interval']) + loader = NSLoader( + datapath1=data_config["datapath"], + nx=data_config["nx"], + nt=data_config["nt"], + sub=data_config["sub"], + sub_t=data_config["sub_t"], + N=data_config["total_num"], + t_interval=data_config["time_interval"], + ) - train_loader = loader.make_loader(data_config['n_sample'], - batch_size=config['train']['batchsize'], - start=data_config['offset'], - train=data_config['shuffle']) + train_loader = loader.make_loader( + data_config["n_sample"], + batch_size=config["train"]["batchsize"], + start=data_config["offset"], + train=data_config["shuffle"], + ) # prepare dataloader for training with only equations - gr_sampler = GaussianRF(2, data_config['S2'], 2 * math.pi, alpha=2.5, tau=7) - a_loader = online_loader(gr_sampler, - S=data_config['S2'], - T=data_config['T2'], - time_scale=data_config['time_interval'], - batchsize=config['train']['batchsize']) + gr_sampler = GaussianRF(2, data_config["S2"], 2 * math.pi, alpha=2.5, tau=7) + a_loader = online_loader( + gr_sampler, + S=data_config["S2"], + T=data_config["T2"], + time_scale=data_config["time_interval"], + batchsize=config["train"]["batchsize"], + ) # create model - model = FNO3d(modes1=config['model']['modes1'], - modes2=config['model']['modes2'], - modes3=config['model']['modes3'], - fc_dim=config['model']['fc_dim'], - layers=config['model']['layers'], - act=config['model']['act']) + model = FNO3d( + modes1=config["model"]["modes1"], + modes2=config["model"]["modes2"], + modes3=config["model"]["modes3"], + fc_dim=config["model"]["fc_dim"], + layers=config["model"]["layers"], + act=config["model"]["act"], + ) # Load from checkpoint - if 'ckpt' in config['train']: - ckpt_path = config['train']['ckpt'] + if "ckpt" in config["train"]: + ckpt_path = config["train"]["ckpt"] ckpt = paddle.load(ckpt_path) - model.load_state_dict(ckpt['model']) - print('Weights loaded from %s' % ckpt_path) + model.load_state_dict(ckpt["model"]) + print("Weights loaded from %s" % ckpt_path) # create optimizer and learning rate scheduler - scheduler = MultiStepDecay(learning_rate=config['train']['base_lr'], - milestones=config['train']['milestones'], - gamma=config['train']['scheduler_gamma']) + scheduler = MultiStepDecay( + learning_rate=config["train"]["base_lr"], + milestones=config["train"]["milestones"], + gamma=config["train"]["scheduler_gamma"], + ) optimizer = Adam(learning_rate=scheduler, parameters=model.parameters()) - mixed_train(model, - train_loader, - loader.S, loader.T, - a_loader, - data_config['S2'], data_config['T2'], - optimizer, - scheduler, - config, - log=args.log, - project=config['log']['project'], - group=config['log']['group']) + mixed_train( + model, + train_loader, + loader.S, + loader.T, + a_loader, + data_config["S2"], + data_config["T2"], + optimizer, + scheduler, + config, + log=args.log, + project=config["log"]["project"], + group=config["log"]["group"], + ) + def train_2d(args, config): - data_config = config['data'] + data_config = config["data"] - dataset = DarcyCombo(datapath=data_config['datapath'], - nx=data_config['nx'], - sub=data_config['sub'], - pde_sub=data_config['pde_sub'], - num=data_config['n_sample'], - offset=data_config['offset']) - train_loader = DataLoader(dataset, batch_size=config['train']['batchsize'], shuffle=True) - model = FNO2d(modes1=config['model']['modes1'], - modes2=config['model']['modes2'], - fc_dim=config['model']['fc_dim'], - layers=config['model']['layers'], - act=config['model']['act'], - pad_ratio=config['model']['pad_ratio']) + dataset = DarcyCombo( + datapath=data_config["datapath"], + nx=data_config["nx"], + sub=data_config["sub"], + pde_sub=data_config["pde_sub"], + num=data_config["n_sample"], + offset=data_config["offset"], + ) + train_loader = DataLoader( + dataset, batch_size=config["train"]["batchsize"], shuffle=True + ) + model = FNO2d( + modes1=config["model"]["modes1"], + modes2=config["model"]["modes2"], + fc_dim=config["model"]["fc_dim"], + layers=config["model"]["layers"], + act=config["model"]["act"], + pad_ratio=config["model"]["pad_ratio"], + ) # Load from checkpoint - if 'ckpt' in config['train']: - ckpt_path = config['train']['ckpt'] + if "ckpt" in config["train"]: + ckpt_path = config["train"]["ckpt"] ckpt = paddle.load(ckpt_path) - model.load_state_dict(ckpt['model']) - print('Weights loaded from %s' % ckpt_path) + model.load_state_dict(ckpt["model"]) + print("Weights loaded from %s" % ckpt_path) - scheduler = MultiStepDecay(learning_rate=config['train']['base_lr'], - milestones=config['train']['milestones'], - gamma=config['train']['scheduler_gamma']) + scheduler = MultiStepDecay( + learning_rate=config["train"]["base_lr"], + milestones=config["train"]["milestones"], + gamma=config["train"]["scheduler_gamma"], + ) optimizer = Adam(learning_rate=scheduler, parameters=model.parameters()) - train_2d_operator(model, - train_loader, - optimizer, scheduler, - config, rank=0, log=args.log, - project=config['log']['project'], - group=config['log']['group']) + train_2d_operator( + model, + train_loader, + optimizer, + scheduler, + config, + rank=0, + log=args.log, + project=config["log"]["project"], + group=config["log"]["group"], + ) + -if __name__ == '__main__': +if __name__ == "__main__": # parse options - parser = ArgumentParser(description='Basic paser') - parser.add_argument('--config_path', type=str, help='Path to the configuration file') - parser.add_argument('--log', action='store_true', help='Turn on the wandb') + parser = ArgumentParser(description="Basic paser") + parser.add_argument( + "--config_path", type=str, help="Path to the configuration file" + ) + parser.add_argument("--log", action="store_true", help="Turn on the wandb") args = parser.parse_args() config_file = args.config_path - with open(config_file, 'r') as stream: + with open(config_file, "r") as stream: config = yaml.load(stream, yaml.FullLoader) - if 'name' in config['data'] and config['data']['name'] == 'Darcy': + if "name" in config["data"] and config["data"]["name"] == "Darcy": train_2d(args, config) else: train_3d(args, config) diff --git a/jointContribution/PINO/PINO_paddle/train_pino.py b/jointContribution/PINO/PINO_paddle/train_pino.py index 828ae24702..7fbf25af53 100644 --- a/jointContribution/PINO/PINO_paddle/train_pino.py +++ b/jointContribution/PINO/PINO_paddle/train_pino.py @@ -1,17 +1,25 @@ import os -import yaml import random from argparse import ArgumentParser -from tqdm import tqdm + import numpy as np import paddle -from paddle.optimizer import Adam +import yaml +from models import FNO3d from paddle.io import DataLoader +from paddle.optimizer import Adam from paddle.optimizer.lr import MultiStepDecay -from models import FNO3d -from train_utils.losses import LpLoss, PINO_loss3d, get_forcing -from train_utils.datasets import KFDataset, KFaDataset, sample_data -from train_utils.utils import save_ckpt, count_params, dict2str +from tqdm import tqdm +from train_utils.datasets import KFaDataset +from train_utils.datasets import KFDataset +from train_utils.datasets import sample_data +from train_utils.losses import LpLoss +from train_utils.losses import PINO_loss3d +from train_utils.losses import get_forcing +from train_utils.utils import count_params +from train_utils.utils import dict2str +from train_utils.utils import save_ckpt + @paddle.no_grad() def eval_ns(model, val_loader, criterion): @@ -29,34 +37,38 @@ def eval_ns(model, val_loader, criterion): std_err = np.std(val_err, ddof=1) / np.sqrt(N) return avg_err, std_err -def train_ns(model, - train_u_loader, # training data - train_a_loader, # initial conditions - val_loader, # validation data - optimizer, - scheduler, - config, args): - start_iter = config['train']['start_iter'] - v = 1/ config['data']['Re'] - t_duration = config['data']['t_duration'] - save_step = config['train']['save_step'] - eval_step = config['train']['eval_step'] - - ic_weight = config['train']['ic_loss'] - f_weight = config['train']['f_loss'] - xy_weight = config['train']['xy_loss'] + +def train_ns( + model, + train_u_loader, # training data + train_a_loader, # initial conditions + val_loader, # validation data + optimizer, + scheduler, + config, + args, +): + start_iter = config["train"]["start_iter"] + v = 1 / config["data"]["Re"] + t_duration = config["data"]["t_duration"] + save_step = config["train"]["save_step"] + eval_step = config["train"]["eval_step"] + + ic_weight = config["train"]["ic_loss"] + f_weight = config["train"]["f_loss"] + xy_weight = config["train"]["xy_loss"] # set up directory - base_dir = os.path.join('exp', config['log']['logdir']) - ckpt_dir = os.path.join(base_dir, 'ckpts') + base_dir = os.path.join("exp", config["log"]["logdir"]) + ckpt_dir = os.path.join(base_dir, "ckpts") os.makedirs(ckpt_dir, exist_ok=True) # loss fn lploss = LpLoss(size_average=True) - - S = config['data']['pde_res'][0] + + S = config["data"]["pde_res"][0] forcing = get_forcing(S) - - pbar = range(start_iter, config['train']['num_iter']) + + pbar = range(start_iter, config["train"]["num_iter"]) if args.tqdm: pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.2) @@ -75,18 +87,18 @@ def train_ns(model, out = model(a_in) data_loss = lploss(out, u) else: - data_loss = paddle.zeros(1, dtype='float32') + data_loss = paddle.zeros(1, dtype="float32") if f_weight != 0.0: # pde loss a = next(a_loader) a = a out = model(a) - - u0 = a[:, :, :, 0, -1] + + u0 = a[:, :, :, 0, -1] loss_ic, loss_f = PINO_loss3d(out, u0, forcing, v, t_duration) - log_dict['IC'] = loss_ic.item() - log_dict['PDE'] = loss_f.item() + log_dict["IC"] = loss_ic.item() + log_dict["PDE"] = loss_f.item() else: loss_ic = loss_f = 0.0 loss = data_loss * xy_weight + loss_f * f_weight + loss_ic * ic_weight @@ -95,125 +107,131 @@ def train_ns(model, optimizer.step() scheduler.step() - log_dict['train loss'] = loss.item() - log_dict['data'] = data_loss.item() + log_dict["train loss"] = loss.item() + log_dict["data"] = data_loss.item() if e % eval_step == 0: eval_err, std_err = eval_ns(model, val_loader, lploss) - log_dict['val error'] = eval_err - + log_dict["val error"] = eval_err + if args.tqdm: logstr = dict2str(log_dict) - pbar.set_description( - ( - logstr - ) - ) + pbar.set_description((logstr)) if e % save_step == 0 and e > 0: - ckpt_path = os.path.join(ckpt_dir, f'model-{e}.pt') + ckpt_path = os.path.join(ckpt_dir, f"model-{e}.pt") save_ckpt(ckpt_path, model, optimizer, scheduler) + def subprocess(args): - with open(args.config, 'r') as f: + with open(args.config, "r") as f: config = yaml.load(f, yaml.FullLoader) # set random seed - config['seed'] = args.seed + config["seed"] = args.seed seed = args.seed paddle.seed(seed) random.seed(seed) if paddle.cuda.is_available(): paddle.cuda.manual_seed_all(seed) - # create model - model = FNO3d(modes1=config['model']['modes1'], - modes2=config['model']['modes2'], - modes3=config['model']['modes3'], - fc_dim=config['model']['fc_dim'], - layers=config['model']['layers'], - act=config['model']['act'], - pad_ratio=config['model']['pad_ratio']) + # create model + model = FNO3d( + modes1=config["model"]["modes1"], + modes2=config["model"]["modes2"], + modes3=config["model"]["modes3"], + fc_dim=config["model"]["fc_dim"], + layers=config["model"]["layers"], + act=config["model"]["act"], + pad_ratio=config["model"]["pad_ratio"], + ) num_params = count_params(model) - config['num_params'] = num_params - print(f'Number of parameters: {num_params}') + config["num_params"] = num_params + print(f"Number of parameters: {num_params}") # Load from checkpoint if args.ckpt: ckpt_path = args.ckpt ckpt = paddle.load(ckpt_path) - model.load_state_dict(ckpt['model']) - print('Weights loaded from %s' % ckpt_path) - + model.load_state_dict(ckpt["model"]) + print("Weights loaded from %s" % ckpt_path) + if args.test: - batchsize = config['test']['batchsize'] - testset = KFDataset(paths=config['data']['paths'], - raw_res=config['data']['raw_res'], - data_res=config['test']['data_res'], - pde_res=config['test']['data_res'], - n_samples=config['data']['n_test_samples'], - offset=config['data']['testoffset'], - t_duration=config['data']['t_duration']) + batchsize = config["test"]["batchsize"] + testset = KFDataset( + paths=config["data"]["paths"], + raw_res=config["data"]["raw_res"], + data_res=config["test"]["data_res"], + pde_res=config["test"]["data_res"], + n_samples=config["data"]["n_test_samples"], + offset=config["data"]["testoffset"], + t_duration=config["data"]["t_duration"], + ) testloader = DataLoader(testset, batch_size=batchsize, num_workers=4) criterion = LpLoss() test_err, std_err = eval_ns(model, testloader, criterion) - print(f'Averaged test relative L2 error: {test_err}; Standard error: {std_err}') + print(f"Averaged test relative L2 error: {test_err}; Standard error: {std_err}") else: # training set - batchsize = config['train']['batchsize'] - u_set = KFDataset(paths=config['data']['paths'], - raw_res=config['data']['raw_res'], - data_res=config['data']['data_res'], - pde_res=config['data']['data_res'], - n_samples=config['data']['n_data_samples'], - offset=config['data']['offset'], - t_duration=config['data']['t_duration']) + batchsize = config["train"]["batchsize"] + u_set = KFDataset( + paths=config["data"]["paths"], + raw_res=config["data"]["raw_res"], + data_res=config["data"]["data_res"], + pde_res=config["data"]["data_res"], + n_samples=config["data"]["n_data_samples"], + offset=config["data"]["offset"], + t_duration=config["data"]["t_duration"], + ) u_loader = DataLoader(u_set, batch_size=batchsize, num_workers=4, shuffle=True) - a_set = KFaDataset(paths=config['data']['paths'], - raw_res=config['data']['raw_res'], - pde_res=config['data']['pde_res'], - n_samples=config['data']['n_a_samples'], - offset=config['data']['a_offset'], - t_duration=config['data']['t_duration']) + a_set = KFaDataset( + paths=config["data"]["paths"], + raw_res=config["data"]["raw_res"], + pde_res=config["data"]["pde_res"], + n_samples=config["data"]["n_a_samples"], + offset=config["data"]["a_offset"], + t_duration=config["data"]["t_duration"], + ) a_loader = DataLoader(a_set, batch_size=batchsize, num_workers=4, shuffle=True) # val set - valset = KFDataset(paths=config['data']['paths'], - raw_res=config['data']['raw_res'], - data_res=config['test']['data_res'], - pde_res=config['test']['data_res'], - n_samples=config['data']['n_test_samples'], - offset=config['data']['testoffset'], - t_duration=config['data']['t_duration']) + valset = KFDataset( + paths=config["data"]["paths"], + raw_res=config["data"]["raw_res"], + data_res=config["test"]["data_res"], + pde_res=config["test"]["data_res"], + n_samples=config["data"]["n_test_samples"], + offset=config["data"]["testoffset"], + t_duration=config["data"]["t_duration"], + ) val_loader = DataLoader(valset, batch_size=batchsize, num_workers=4) - print(f'Train set: {len(u_set)}; Test set: {len(valset)}; IC set: {len(a_set)}') + print(f"Train set: {len(u_set)}; Test set: {len(valset)}; IC set: {len(a_set)}") - scheduler = MultiStepDecay(learning_rate=config['train']['base_lr'], - milestones=config['train']['milestones'], - gamma=config['train']['scheduler_gamma']) + scheduler = MultiStepDecay( + learning_rate=config["train"]["base_lr"], + milestones=config["train"]["milestones"], + gamma=config["train"]["scheduler_gamma"], + ) optimizer = Adam(learning_rate=scheduler, parameters=model.parameters()) if args.ckpt: ckpt = paddle.load(ckpt_path) - optimizer.load_state_dict(ckpt['optim']) - scheduler.load_state_dict(ckpt['scheduler']) - config['train']['start_iter'] = scheduler.last_epoch - train_ns(model, - u_loader, a_loader, - val_loader, - optimizer, scheduler, - config, args) - print('Done!') - - - -if __name__ == '__main__': + optimizer.load_state_dict(ckpt["optim"]) + scheduler.load_state_dict(ckpt["scheduler"]) + config["train"]["start_iter"] = scheduler.last_epoch + train_ns( + model, u_loader, a_loader, val_loader, optimizer, scheduler, config, args + ) + print("Done!") + + +if __name__ == "__main__": paddle.backends.cudnn.benchmark = True # parse options - parser = ArgumentParser(description='Basic paser') - parser.add_argument('--config', type=str, help='Path to the configuration file') - parser.add_argument('--log', action='store_true', help='Turn on the wandb') - parser.add_argument('--seed', type=int, default=None) - parser.add_argument('--ckpt', type=str, default=None) - parser.add_argument('--test', action='store_true', help='Test') - parser.add_argument('--tqdm', action='store_true', help='Turn on the tqdm') + parser = ArgumentParser(description="Basic paser") + parser.add_argument("--config", type=str, help="Path to the configuration file") + parser.add_argument("--log", action="store_true", help="Turn on the wandb") + parser.add_argument("--seed", type=int, default=None) + parser.add_argument("--ckpt", type=str, default=None) + parser.add_argument("--test", action="store_true", help="Test") + parser.add_argument("--tqdm", action="store_true", help="Turn on the tqdm") args = parser.parse_args() if args.seed is None: args.seed = random.randint(0, 100000) - subprocess(args) \ No newline at end of file + subprocess(args) diff --git a/jointContribution/PINO/PINO_paddle/train_utils/__init__.py b/jointContribution/PINO/PINO_paddle/train_utils/__init__.py index 6bcd02cdfa..86b2918b2b 100644 --- a/jointContribution/PINO/PINO_paddle/train_utils/__init__.py +++ b/jointContribution/PINO/PINO_paddle/train_utils/__init__.py @@ -1,2 +1,4 @@ -from .datasets import NSLoader, DarcyFlow -from .losses import get_forcing, LpLoss \ No newline at end of file +from .datasets import DarcyFlow +from .datasets import NSLoader +from .losses import LpLoss +from .losses import get_forcing diff --git a/jointContribution/PINO/PINO_paddle/train_utils/datasets.py b/jointContribution/PINO/PINO_paddle/train_utils/datasets.py index e95988b684..dd5faad422 100644 --- a/jointContribution/PINO/PINO_paddle/train_utils/datasets.py +++ b/jointContribution/PINO/PINO_paddle/train_utils/datasets.py @@ -1,29 +1,34 @@ -import scipy.io import numpy as np +import scipy.io try: from pyDOE import lhs + # Only needed for PINN's dataset except ImportError: lhs = None import paddle from paddle.io import Dataset -from .utils import get_grid3d, convert_ic, paddle2dgrid + +from .utils import convert_ic +from .utils import get_grid3d +from .utils import paddle2dgrid + def online_loader(sampler, S, T, time_scale, batchsize=1): while True: u0 = sampler.sample(batchsize) - a = convert_ic(u0, batchsize, - S, T, - time_scale=time_scale) + a = convert_ic(u0, batchsize, S, T, time_scale=time_scale) yield a + def sample_data(loader): while True: for batch in loader: yield batch + class MatReader(object): def __init__(self, file_path, to_paddle=True, to_cuda=False, to_float=True): super(MatReader, self).__init__() @@ -73,8 +78,9 @@ def set_paddle(self, to_paddle): def set_float(self, to_float): self.to_float = to_float + class BurgersLoader(object): - def __init__(self, datapath, nx=2 ** 10, nt=100, sub=8, sub_t=1, new=False): + def __init__(self, datapath, nx=2**10, nt=100, sub=8, sub_t=1, new=False): dataloader = MatReader(datapath) self.sub = sub self.sub_t = sub_t @@ -83,24 +89,29 @@ def __init__(self, datapath, nx=2 ** 10, nt=100, sub=8, sub_t=1, new=False): self.new = new if new: self.T += 1 - self.x_data = dataloader.read_field('input')[:, ::sub] - self.y_data = dataloader.read_field('output')[:, ::sub_t, ::sub] - self.v = dataloader.read_field('visc').item() + self.x_data = dataloader.read_field("input")[:, ::sub] + self.y_data = dataloader.read_field("output")[:, ::sub_t, ::sub] + self.v = dataloader.read_field("visc").item() def make_loader(self, n_sample, batch_size, start=0, train=True): - Xs = self.x_data[start:start + n_sample] - ys = self.y_data[start:start + n_sample] + Xs = self.x_data[start : start + n_sample] + ys = self.y_data[start : start + n_sample] if self.new: - gridx = paddle.to_tensor(np.linspace(0, 1, self.s + 1)[:-1], dtype='float32') - gridt = paddle.to_tensor(np.linspace(0, 1, self.T), dtype='float32') + gridx = paddle.to_tensor( + np.linspace(0, 1, self.s + 1)[:-1], dtype="float32" + ) + gridt = paddle.to_tensor(np.linspace(0, 1, self.T), dtype="float32") else: - gridx = paddle.to_tensor(np.linspace(0, 1, self.s), dtype='float32') - gridt = paddle.to_tensor(np.linspace(0, 1, self.T + 1)[1:], dtype='float32') + gridx = paddle.to_tensor(np.linspace(0, 1, self.s), dtype="float32") + gridt = paddle.to_tensor(np.linspace(0, 1, self.T + 1)[1:], dtype="float32") gridx = gridx.reshape([1, 1, self.s]) gridt = gridt.reshape([1, self.T, 1]) Xs = Xs.reshape([n_sample, 1, self.s]).tile([1, self.T, 1]) - Xs = paddle.stack([Xs, gridx.tile([n_sample, self.T, 1]), gridt.tile([n_sample, 1, self.s])], axis=3) + Xs = paddle.stack( + [Xs, gridx.tile([n_sample, self.T, 1]), gridt.tile([n_sample, 1, self.s])], + axis=3, + ) dataset = paddle.io.TensorDataset([Xs, ys]) if train: loader = paddle.io.DataLoader(dataset, batch_size=batch_size, shuffle=False) @@ -108,12 +119,12 @@ def make_loader(self, n_sample, batch_size, start=0, train=True): loader = paddle.io.DataLoader(dataset, batch_size=batch_size, shuffle=False) return loader + class NSLoader(object): - def __init__(self, datapath1, - nx, nt, - datapath2=None, sub=1, sub_t=1, - N=100, t_interval=1.0): - ''' + def __init__( + self, datapath1, nx, nt, datapath2=None, sub=1, sub_t=1, N=100, t_interval=1.0 + ): + """ Load data from npy and reshape to (N, X, Y, T) Args: datapath1: path to data @@ -124,16 +135,16 @@ def __init__(self, datapath1, sub_t: N: t_interval: - ''' + """ self.S = nx // sub self.T = int(nt * t_interval) // sub_t + 1 self.time_scale = t_interval data1 = np.load(datapath1) - data1 = paddle.to_tensor(data1, dtype='float32')[..., ::sub_t, ::sub, ::sub] + data1 = paddle.to_tensor(data1, dtype="float32")[..., ::sub_t, ::sub, ::sub] if datapath2 is not None: data2 = np.load(datapath2) - data2 = paddle.to_tensor(data2, dtype='float32')[..., ::sub_t, ::sub, ::sub] + data2 = paddle.to_tensor(data2, dtype="float32")[..., ::sub_t, ::sub, ::sub] if t_interval == 0.5: data1 = self.extract(data1) if datapath2 is not None: @@ -147,46 +158,69 @@ def __init__(self, datapath1, def make_loader(self, n_sample, batch_size, start=0, train=True): if train: - a_data = self.data[start:start + n_sample, :, :, 0].reshape([n_sample, self.S, self.S]) - u_data = self.data[start:start + n_sample].reshape([n_sample, self.S, self.S, self.T]) + a_data = self.data[start : start + n_sample, :, :, 0].reshape( + [n_sample, self.S, self.S] + ) + u_data = self.data[start : start + n_sample].reshape( + [n_sample, self.S, self.S, self.T] + ) else: a_data = self.data[-n_sample:, :, :, 0].reshape([n_sample, self.S, self.S]) u_data = self.data[-n_sample:].reshape([n_sample, self.S, self.S, self.T]) - a_data = a_data.reshape([n_sample, self.S, self.S, 1, 1]).tile([1, 1, 1, self.T, 1]) + a_data = a_data.reshape([n_sample, self.S, self.S, 1, 1]).tile( + [1, 1, 1, self.T, 1] + ) gridx, gridy, gridt = get_grid3d(self.S, self.T, time_scale=self.time_scale) - a_data = paddle.concat((gridx.tile([n_sample, 1, 1, 1, 1]), gridy.tile([n_sample, 1, 1, 1, 1]), - gridt.tile([n_sample, 1, 1, 1, 1]), a_data), axis=-1) + a_data = paddle.concat( + ( + gridx.tile([n_sample, 1, 1, 1, 1]), + gridy.tile([n_sample, 1, 1, 1, 1]), + gridt.tile([n_sample, 1, 1, 1, 1]), + a_data, + ), + axis=-1, + ) dataset = paddle.io.TensorDataset(a_data, u_data) loader = paddle.io.DataLoader(dataset, batch_size=batch_size, shuffle=train) return loader def make_dataset(self, n_sample, start=0, train=True): if train: - a_data = self.data[start:start + n_sample, :, :, 0].reshape([n_sample, self.S, self.S]) - u_data = self.data[start:start + n_sample].reshape([n_sample, self.S, self.S, self.T]) + a_data = self.data[start : start + n_sample, :, :, 0].reshape( + [n_sample, self.S, self.S] + ) + u_data = self.data[start : start + n_sample].reshape( + [n_sample, self.S, self.S, self.T] + ) else: a_data = self.data[-n_sample:, :, :, 0].reshape([n_sample, self.S, self.S]) u_data = self.data[-n_sample:].reshape([n_sample, self.S, self.S, self.T]) - a_data = a_data.reshape([n_sample, self.S, self.S, 1, 1]).tile([1, 1, 1, self.T, 1]) + a_data = a_data.reshape([n_sample, self.S, self.S, 1, 1]).tile( + [1, 1, 1, self.T, 1] + ) gridx, gridy, gridt = get_grid3d(self.S, self.T) - a_data = paddle.concat(( - gridx.tile([n_sample, 1, 1, 1, 1]), - gridy.tile([n_sample, 1, 1, 1, 1]), - gridt.tile([n_sample, 1, 1, 1, 1]), - a_data), axis=-1) + a_data = paddle.concat( + ( + gridx.tile([n_sample, 1, 1, 1, 1]), + gridy.tile([n_sample, 1, 1, 1, 1]), + gridt.tile([n_sample, 1, 1, 1, 1]), + a_data, + ), + axis=-1, + ) dataset = paddle.io.TensorDataset(a_data, u_data) return dataset @staticmethod def extract(data): - ''' + """ Extract data with time range 0-0.5, 0.25-0.75, 0.5-1.0, 0.75-1.25,... Args: data: tensor with size N x 129 x 128 x 128 Returns: output: (4*N-1) x 65 x 128 x 128 - ''' + """ T = data.shape[1] // 2 interval = data.shape[1] // 4 N = data.shape[0] @@ -197,25 +231,34 @@ def extract(data): # reach boundary break if j != 3: - new_data[i * 4 + j] = data[i, interval * j:interval * j + T + 1] + new_data[i * 4 + j] = data[i, interval * j : interval * j + T + 1] else: - new_data[i * 4 + j, 0: interval] = data[i, interval * j:interval * j + interval] - new_data[i * 4 + j, interval: T + 1] = data[i + 1, 0:interval + 1] + new_data[i * 4 + j, 0:interval] = data[ + i, interval * j : interval * j + interval + ] + new_data[i * 4 + j, interval : T + 1] = data[ + i + 1, 0 : interval + 1 + ] return new_data + class KFDataset(Dataset): - def __init__(self, paths, - data_res, pde_res, - raw_res, - n_samples=None, - total_samples=None, - idx=0, - offset=0, - t_duration=1.0): + def __init__( + self, + paths, + data_res, + pde_res, + raw_res, + n_samples=None, + total_samples=None, + idx=0, + offset=0, + t_duration=1.0, + ): super().__init__() - self.data_res = data_res # data resolution - self.pde_res = pde_res # pde loss resolution - self.raw_res = raw_res # raw data resolution + self.data_res = data_res # data resolution + self.pde_res = pde_res # pde loss resolution + self.raw_res = raw_res # raw data resolution self.t_duration = t_duration self.paths = paths self.offset = offset @@ -223,94 +266,111 @@ def __init__(self, paths, if t_duration == 1.0: self.T = self.pde_res[2] else: - self.T = int(self.pde_res[2] * t_duration) + 1 # number of points in time dimension + self.T = ( + int(self.pde_res[2] * t_duration) + 1 + ) # number of points in time dimension self.load() if total_samples is not None: - print(f'Load {total_samples} samples starting from {idx}th sample') - self.data = self.data[idx:idx + total_samples] - self.a_data = self.a_data[idx:idx + total_samples] - + print(f"Load {total_samples} samples starting from {idx}th sample") + self.data = self.data[idx : idx + total_samples] + self.a_data = self.a_data[idx : idx + total_samples] + self.data_s_step = pde_res[0] // data_res[0] self.data_t_step = (pde_res[2] - 1) // (data_res[2] - 1) def load(self): datapath = self.paths[0] - raw_data = np.load(datapath, mmap_mode='r') + raw_data = np.load(datapath, mmap_mode="r") # subsample ratio sub_x = self.raw_res[0] // self.data_res[0] sub_t = (self.raw_res[2] - 1) // (self.data_res[2] - 1) - + a_sub_x = self.raw_res[0] // self.pde_res[0] # load data - data = raw_data[self.offset: self.offset + self.n_samples, ::sub_t, ::sub_x, ::sub_x] + data = raw_data[ + self.offset : self.offset + self.n_samples, ::sub_t, ::sub_x, ::sub_x + ] # divide data - if self.t_duration != 0.: + if self.t_duration != 0.0: end_t = self.raw_res[2] - 1 - K = int(1/self.t_duration) + K = int(1 / self.t_duration) step = end_t // K data = self.partition(data) - a_data = raw_data[self.offset: self.offset + self.n_samples, 0:end_t:step, ::a_sub_x, ::a_sub_x] - a_data = a_data.reshape([self.n_samples * K, 1, self.pde_res[0], self.pde_res[1]]) # 2N x 1 x S x S + a_data = raw_data[ + self.offset : self.offset + self.n_samples, + 0:end_t:step, + ::a_sub_x, + ::a_sub_x, + ] + a_data = a_data.reshape( + [self.n_samples * K, 1, self.pde_res[0], self.pde_res[1]] + ) # 2N x 1 x S x S else: - a_data = raw_data[self.offset: self.offset + self.n_samples, 0:1, ::a_sub_x, ::a_sub_x] + a_data = raw_data[ + self.offset : self.offset + self.n_samples, 0:1, ::a_sub_x, ::a_sub_x + ] # convert into paddle tensor - data = paddle.to_tensor(data, dtype='float32') - a_data = paddle.to_tensor(a_data, dtype='float32').transpose([0, 2, 3, 1]) + data = paddle.to_tensor(data, dtype="float32") + a_data = paddle.to_tensor(a_data, dtype="float32").transpose([0, 2, 3, 1]) self.data = data.transpose([0, 2, 3, 1]) S = self.pde_res[1] - - a_data = a_data[:, :, :, :, None] # N x S x S x 1 x 1 + + a_data = a_data[:, :, :, :, None] # N x S x S x 1 x 1 gridx, gridy, gridt = get_grid3d(S, self.T) - self.grid = paddle.concat((gridx[0], gridy[0], gridt[0]), axis=-1) # S x S x T x 3 + self.grid = paddle.concat( + (gridx[0], gridy[0], gridt[0]), axis=-1 + ) # S x S x T x 3 self.a_data = a_data def partition(self, data): - ''' + """ Args: data: tensor with size N x T x S x S Returns: output: int(1/t_duration) *N x (T//2 + 1) x 128 x 128 - ''' + """ N, T, S = data.shape[:3] K = int(1 / self.t_duration) new_data = np.zeros((K * N, T // K + 1, S, S)) step = T // K for i in range(N): for j in range(K): - new_data[i * K + j] = data[i, j * step: (j+1) * step + 1] + new_data[i * K + j] = data[i, j * step : (j + 1) * step + 1] return new_data def __getitem__(self, idx): - a_data = paddle.concat(( - self.grid, - self.a_data[idx].tile(1, 1, self.T, 1) - ), axis=-1) + a_data = paddle.concat( + (self.grid, self.a_data[idx].tile(1, 1, self.T, 1)), axis=-1 + ) return self.data[idx], a_data - def __len__(self, ): + def __len__( + self, + ): return self.data.shape[0] + class BurgerData(Dataset): - ''' - members: + """ + members: - t, x, Exact: raw data - - X, T: meshgrid + - X, T: meshgrid - X_star, u_star: flattened (x, t), u array - lb, ub: lower bound and upper bound vector - X_u, u: boundary condition data (x, t), u - ''' + """ def __init__(self, datapath): data = scipy.io.loadmat(datapath) # raw 2D data - self.t = data['t'].flatten()[:, None] # (100,1) - self.x = data['x'].flatten()[:, None] # (256, 1) - self.Exact = np.real(data['usol']).T # (100, 256) + self.t = data["t"].flatten()[:, None] # (100,1) + self.x = data["x"].flatten()[:, None] # (256, 1) + self.Exact = np.real(data["usol"]).T # (100, 256) # Flattened sequence self.get_flatten_data() @@ -344,39 +404,40 @@ def get_boundary_data(self): self.u = np.vstack([uu1, uu2, uu3]) def sample_xt(self, N=10000): - ''' + """ Sample (x, t) pairs within the boundary Return: - X_f: (N, 2) array - ''' + """ X_f = self.lb + (self.ub - self.lb) * lhs(2, N) X_f = np.vstack((X_f, self.X_u)) return X_f def sample_xu(self, N=100): - ''' + """ Sample N points from boundary data - Return: - - X_u: (N, 2) array + Return: + - X_u: (N, 2) array - u: (N, 1) array - ''' + """ idx = np.random.choice(self.X_u.shape[0], N, replace=False) X_u = self.X_u[idx, :] u = self.u[idx, :] return X_u, u + class DarcyFlow(Dataset): - def __init__(self, - datapath, - nx, sub, - offset=0, - num=1): + def __init__(self, datapath, nx, sub, offset=0, num=1): self.S = int(nx // sub) + 1 if sub > 1 else nx data = scipy.io.loadmat(datapath) - a = data['coeff'] - u = data['sol'] - self.a = paddle.to_tensor(a[offset: offset + num, ::sub, ::sub], dtype='float32') - self.u = paddle.to_tensor(u[offset: offset + num, ::sub, ::sub], dtype='float32') + a = data["coeff"] + u = data["sol"] + self.a = paddle.to_tensor( + a[offset : offset + num, ::sub, ::sub], dtype="float32" + ) + self.u = paddle.to_tensor( + u[offset : offset + num, ::sub, ::sub], dtype="float32" + ) self.mesh = paddle2dgrid(self.S, self.S) def __len__(self): @@ -386,22 +447,25 @@ def __getitem__(self, item): fa = self.a[item] return paddle.concat([fa.unsqueeze(2), self.mesh], axis=2), self.u[item] + class DarcyIC(Dataset): - def __init__(self, - datapath, - nx, sub, - offset=0, - num=1): + def __init__(self, datapath, nx, sub, offset=0, num=1): self.S = int(nx // sub) + 1 if sub > 1 else nx data = scipy.io.loadmat(datapath) - a = data['coeff'] - self.a = paddle.to_tensor(a[offset: offset + num, ::sub, ::sub], dtype='float32') + a = data["coeff"] + self.a = paddle.to_tensor( + a[offset : offset + num, ::sub, ::sub], dtype="float32" + ) self.mesh = paddle2dgrid(self.S, self.S) data = scipy.io.loadmat(datapath) - a = data['coeff'] - u = data['sol'] - self.a = paddle.to_tensor(a[offset: offset + num, ::sub, ::sub], dtype='float32') - self.u = paddle.to_tensor(u[offset: offset + num, ::sub, ::sub], dtype='float32') + a = data["coeff"] + u = data["sol"] + self.a = paddle.to_tensor( + a[offset : offset + num, ::sub, ::sub], dtype="float32" + ) + self.u = paddle.to_tensor( + u[offset : offset + num, ::sub, ::sub], dtype="float32" + ) self.mesh = paddle2dgrid(self.S, self.S) def __len__(self): @@ -409,24 +473,27 @@ def __len__(self): def __getitem__(self, item): fa = self.a[item] - return paddle.concat([fa.unsqueeze(2), self.mesh], axis=2) + return paddle.concat([fa.unsqueeze(2), self.mesh], axis=2) + class DarcyCombo(Dataset): - def __init__(self, - datapath, - nx, - sub, pde_sub, - num=1000, offset=0) -> None: + def __init__(self, datapath, nx, sub, pde_sub, num=1000, offset=0) -> None: super().__init__() self.S = int(nx // sub) + 1 if sub > 1 else nx self.pde_S = int(nx // pde_sub) + 1 if sub > 1 else nx data = scipy.io.loadmat(datapath) - a = data['coeff'] - u = data['sol'] - self.a = paddle.to_tensor(a[offset: offset + num, ::sub, ::sub], dtype='float32') - self.u = paddle.to_tensor(u[offset: offset + num, ::sub, ::sub], dtype='float32') + a = data["coeff"] + u = data["sol"] + self.a = paddle.to_tensor( + a[offset : offset + num, ::sub, ::sub], dtype="float32" + ) + self.u = paddle.to_tensor( + u[offset : offset + num, ::sub, ::sub], dtype="float32" + ) self.mesh = paddle2dgrid(self.S, self.S) - self.pde_a = paddle.to_tensor(a[offset: offset + num, ::pde_sub, ::pde_sub], dtype='float32') + self.pde_a = paddle.to_tensor( + a[offset : offset + num, ::pde_sub, ::pde_sub], dtype="float32" + ) self.pde_mesh = paddle2dgrid(self.pde_S, self.pde_S) def __len__(self): @@ -439,19 +506,19 @@ def __getitem__(self, item): pde_ic = paddle.concat([pde_a.unsqueeze[2], self.pde_mesh], axis=2) return data_ic, self.u[item], pde_ic -''' + +""" dataset class for loading initial conditions for Komogrov flow -''' +""" + + class KFaDataset(Dataset): - def __init__(self, paths, - pde_res, - raw_res, - n_samples=None, - offset=0, - t_duration=1.0): + def __init__( + self, paths, pde_res, raw_res, n_samples=None, offset=0, t_duration=1.0 + ): super().__init__() - self.pde_res = pde_res # pde loss resolution - self.raw_res = raw_res # raw data resolution + self.pde_res = pde_res # pde loss resolution + self.raw_res = raw_res # raw data resolution self.t_duration = t_duration self.paths = paths self.offset = offset @@ -459,39 +526,53 @@ def __init__(self, paths, if t_duration == 1.0: self.T = self.pde_res[2] else: - self.T = int(self.pde_res[2] * t_duration) + 1 # number of points in time dimension + self.T = ( + int(self.pde_res[2] * t_duration) + 1 + ) # number of points in time dimension self.load() def load(self): datapath = self.paths[0] - raw_data = np.load(datapath, mmap_mode='r') + raw_data = np.load(datapath, mmap_mode="r") # subsample ratio a_sub_x = self.raw_res[0] // self.pde_res[0] # load data - if self.t_duration != 0.: + if self.t_duration != 0.0: end_t = self.raw_res[2] - 1 - K = int(1/self.t_duration) + K = int(1 / self.t_duration) step = end_t // K - a_data = raw_data[self.offset: self.offset + self.n_samples, 0:end_t:step, ::a_sub_x, ::a_sub_x] - a_data = a_data.reshape([self.n_samples * K, 1, self.pde_res[0], self.pde_res[1]]) # 2N x 1 x S x S + a_data = raw_data[ + self.offset : self.offset + self.n_samples, + 0:end_t:step, + ::a_sub_x, + ::a_sub_x, + ] + a_data = a_data.reshape( + [self.n_samples * K, 1, self.pde_res[0], self.pde_res[1]] + ) # 2N x 1 x S x S else: - a_data = raw_data[self.offset: self.offset + self.n_samples, 0:1, ::a_sub_x, ::a_sub_x] + a_data = raw_data[ + self.offset : self.offset + self.n_samples, 0:1, ::a_sub_x, ::a_sub_x + ] # convert into tensor - a_data = paddle.to_tensor(a_data, dtype='float32').permute(0, 2, 3, 1) + a_data = paddle.to_tensor(a_data, dtype="float32").permute(0, 2, 3, 1) S = self.pde_res[1] - a_data = a_data[:, :, :, :, None] # N x S x S x 1 x 1 + a_data = a_data[:, :, :, :, None] # N x S x S x 1 x 1 gridx, gridy, gridt = get_grid3d(S, self.T) - self.grid = paddle.concat((gridx[0], gridy[0], gridt[0]), axis=-1) # S x S x T x 3 + self.grid = paddle.concat( + (gridx[0], gridy[0], gridt[0]), axis=-1 + ) # S x S x T x 3 self.a_data = a_data def __getitem__(self, idx): - a_data = paddle.concat(( - self.grid, - self.a_data[idx].tile(1, 1, self.T, 1) - ), axis=-1) + a_data = paddle.concat( + (self.grid, self.a_data[idx].tile(1, 1, self.T, 1)), axis=-1 + ) return a_data - def __len__(self, ): - return self.a_data.shape[0] \ No newline at end of file + def __len__( + self, + ): + return self.a_data.shape[0] diff --git a/jointContribution/PINO/PINO_paddle/train_utils/eval_2d.py b/jointContribution/PINO/PINO_paddle/train_utils/eval_2d.py index 459c97d3a9..b5e0e08a4a 100644 --- a/jointContribution/PINO/PINO_paddle/train_utils/eval_2d.py +++ b/jointContribution/PINO/PINO_paddle/train_utils/eval_2d.py @@ -1,16 +1,14 @@ -from tqdm import tqdm +import matplotlib.pyplot as plt import numpy as np - import paddle +from tqdm import tqdm -from .losses import LpLoss, darcy_loss, PINO_loss +from .losses import LpLoss +from .losses import PINO_loss +from .losses import darcy_loss -import matplotlib.pyplot as plt -def eval_darcy(model, - dataloader, - config, - use_tqdm=True): +def eval_darcy(model, dataloader, config, use_tqdm=True): model.eval() myloss = LpLoss(size_average=True) if use_tqdm: @@ -19,11 +17,13 @@ def eval_darcy(model, pbar = dataloader mesh = dataloader.dataset.mesh - mollifier = paddle.sin(np.pi * mesh[..., 0]) * paddle.sin(np.pi * mesh[..., 1]) * 0.001 + mollifier = ( + paddle.sin(np.pi * mesh[..., 0]) * paddle.sin(np.pi * mesh[..., 1]) * 0.001 + ) f_val = [] test_err = [] i = 0 - fig, ax = plt.subplots(3,3) + fig, ax = plt.subplots(3, 3) with paddle.no_grad(): for x, y in pbar: @@ -31,19 +31,19 @@ def eval_darcy(model, pred = pred * mollifier if i < 3: ax[2][i].imshow(pred[0, :, :]) - ax[2][i].set_title('prediction') + ax[2][i].set_title("prediction") ax[1][i].imshow(y[0, :, :]) - ax[1][i].set_title('ground truth') + ax[1][i].set_title("ground truth") ax[0][i].imshow(x[0, :, :, 0]) - ax[0][i].set_title('input') + ax[0][i].set_title("input") for k in range(3): - ax[k][i].set_xlabel('x') - ax[k][i].set_ylabel('y') - if i==3: + ax[k][i].set_xlabel("x") + ax[k][i].set_ylabel("y") + if i == 3: plt.tight_layout() - plt.savefig('result.png') - i+=1 + plt.savefig("result.png") + i += 1 data_loss = myloss(pred, y) a = x[..., 0] f_loss = darcy_loss(pred, a) @@ -53,7 +53,7 @@ def eval_darcy(model, if use_tqdm: pbar.set_description( ( - f'Equation error: {f_loss.item():.5f}, test l2 error: {data_loss.item()}' + f"Equation error: {f_loss.item():.5f}, test l2 error: {data_loss.item()}" ) ) mean_f_err = np.mean(f_val) @@ -62,14 +62,13 @@ def eval_darcy(model, mean_err = np.mean(test_err) std_err = np.std(test_err, ddof=1) / np.sqrt(len(test_err)) - print(f'==Averaged relative L2 error mean: {mean_err}, std error: {std_err}==\n' - f'==Averaged equation error mean: {mean_f_err}, std error: {std_f_err}==') + print( + f"==Averaged relative L2 error mean: {mean_err}, std error: {std_err}==\n" + f"==Averaged equation error mean: {mean_f_err}, std error: {std_f_err}==" + ) + -def eval_burgers(model, - dataloader, - v, - config, - use_tqdm=True): +def eval_burgers(model, dataloader, v, config, use_tqdm=True): model.eval() myloss = LpLoss(size_average=True) if use_tqdm: @@ -80,24 +79,24 @@ def eval_burgers(model, test_err = [] f_err = [] i = 0 - fig, ax = plt.subplots(2,3) + fig, ax = plt.subplots(2, 3) for x, y in pbar: x, y = x, y out = model(x).reshape(y.shape) data_loss = myloss(out, y) - if i<3: + if i < 3: ax[0][i].imshow(out[0, :, :]) - ax[0][i].set_xlabel('x') - ax[0][i].set_ylabel('t') - ax[0][i].set_title('prediction') + ax[0][i].set_xlabel("x") + ax[0][i].set_ylabel("t") + ax[0][i].set_title("prediction") ax[1][i].imshow(y[0, :, :]) - ax[1][i].set_xlabel('x') - ax[1][i].set_ylabel('t') - ax[1][i].set_title('ground truth') - if i==3: + ax[1][i].set_xlabel("x") + ax[1][i].set_ylabel("t") + ax[1][i].set_title("ground truth") + if i == 3: plt.tight_layout() - plt.savefig('result.png') - i+=1 + plt.savefig("result.png") + i += 1 loss_u, f_loss = PINO_loss(out, x[:, 0, :, 0], v) test_err.append(data_loss.item()) f_err.append(f_loss.item()) @@ -108,6 +107,7 @@ def eval_burgers(model, mean_err = np.mean(test_err) std_err = np.std(test_err, ddof=1) / np.sqrt(len(test_err)) - print(f'==Averaged relative L2 error mean: {mean_err}, std error: {std_err}==\n' - f'==Averaged equation error mean: {mean_f_err}, std error: {std_f_err}==') - + print( + f"==Averaged relative L2 error mean: {mean_err}, std error: {std_err}==\n" + f"==Averaged equation error mean: {mean_f_err}, std error: {std_f_err}==" + ) diff --git a/jointContribution/PINO/PINO_paddle/train_utils/eval_3d.py b/jointContribution/PINO/PINO_paddle/train_utils/eval_3d.py index f387669d1e..5f7ff56e46 100644 --- a/jointContribution/PINO/PINO_paddle/train_utils/eval_3d.py +++ b/jointContribution/PINO/PINO_paddle/train_utils/eval_3d.py @@ -1,31 +1,35 @@ +from timeit import default_timer + import paddle import paddle.nn.functional as F - from tqdm import tqdm -from timeit import default_timer -from .losses import LpLoss, PINO_loss3d +from .losses import LpLoss +from .losses import PINO_loss3d + -def eval_ns(model, - loader, - dataloader, - forcing, - config, - device, - log=False, - project='PINO-default', - group='FDM', - tags=['Nan'], - use_tqdm=True): - ''' +def eval_ns( + model, + loader, + dataloader, + forcing, + config, + device, + log=False, + project="PINO-default", + group="FDM", + tags=["Nan"], + use_tqdm=True, +): + """ Evaluate the model for Navier Stokes equation - ''' + """ # data parameters - v = 1 / config['data']['Re'] + v = 1 / config["data"]["Re"] S, T = loader.S, loader.T - t_interval = config['data']['time_interval'] + t_interval = config["data"]["time_interval"] # eval settings - batch_size = config['test']['batchsize'] + batch_size = config["test"]["batchsize"] model.eval() myloss = LpLoss(size_average=True) @@ -33,8 +37,7 @@ def eval_ns(model, pbar = tqdm(dataloader, dynamic_ncols=True, smoothing=0.05) else: pbar = dataloader - loss_dict = {'f_error': 0.0, - 'test_l2': 0.0} + loss_dict = {"f_error": 0.0, "test_l2": 0.0} start_time = default_timer() with paddle.no_grad(): for x, y in pbar: @@ -44,19 +47,23 @@ def eval_ns(model, out = out[..., :-5] x = x[:, :, :, 0, -1] loss_l2 = myloss(out.view(batch_size, S, S, T), y.view(batch_size, S, S, T)) - loss_ic, loss_f = PINO_loss3d(out.view(batch_size, S, S, T), x, forcing, v, t_interval) + loss_ic, loss_f = PINO_loss3d( + out.view(batch_size, S, S, T), x, forcing, v, t_interval + ) - loss_dict['f_error'] += loss_f - loss_dict['test_l2'] += loss_l2 + loss_dict["f_error"] += loss_f + loss_dict["test_l2"] += loss_l2 if device == 0 and use_tqdm: pbar.set_description( ( - f'Train f error: {loss_f.item():.5f}; Test l2 error: {loss_l2.item():.5f}' + f"Train f error: {loss_f.item():.5f}; Test l2 error: {loss_l2.item():.5f}" ) ) end_time = default_timer() - test_l2 = loss_dict['test_l2'].item() / len(dataloader) - loss_f = loss_dict['f_error'].item() / len(dataloader) - print(f'==Averaged relative L2 error is: {test_l2}==\n' - f'==Averaged equation error is: {loss_f}==') - print(f'Time cost: {end_time - start_time} s') + test_l2 = loss_dict["test_l2"].item() / len(dataloader) + loss_f = loss_dict["f_error"].item() / len(dataloader) + print( + f"==Averaged relative L2 error is: {test_l2}==\n" + f"==Averaged equation error is: {loss_f}==" + ) + print(f"Time cost: {end_time - start_time} s") diff --git a/jointContribution/PINO/PINO_paddle/train_utils/losses.py b/jointContribution/PINO/PINO_paddle/train_utils/losses.py index 86bc2844f2..b06180583d 100644 --- a/jointContribution/PINO/PINO_paddle/train_utils/losses.py +++ b/jointContribution/PINO/PINO_paddle/train_utils/losses.py @@ -2,6 +2,7 @@ import paddle import paddle.nn.functional as F + def FDM_Darcy(u, a, D=1): batchsize = u.size(0) size = u.size(1) @@ -18,9 +19,10 @@ def FDM_Darcy(u, a, D=1): auy = a * uy auxx = (aux[:, 2:, 1:-1] - aux[:, :-2, 1:-1]) / (2 * dx) auyy = (auy[:, 1:-1, 2:] - auy[:, 1:-1, :-2]) / (2 * dy) - Du = - (auxx + auyy) + Du = -(auxx + auyy) return Du + def darcy_loss(u, a): batchsize = u.shape[0] size = u.shape[1] @@ -33,7 +35,8 @@ def darcy_loss(u, a): loss_f = lploss.rel(Du, f) return loss_f -def FDM_NS_vorticity(w, v=1/40, t_interval=1.0): + +def FDM_NS_vorticity(w, v=1 / 40, t_interval=1.0): batchsize = w.shape[0] nx = w.shape[1] ny = w.shape[2] @@ -42,14 +45,34 @@ def FDM_NS_vorticity(w, v=1/40, t_interval=1.0): w_h = paddle.fft.fft2(w, axes=[1, 2]) # Wavenumbers in y-direction - k_max = nx//2 + k_max = nx // 2 N = nx - k_x = paddle.concat((paddle.arange(start=0, end=k_max, step=1), - paddle.arange(start=-k_max, end=0, step=1)), 0).reshape([N, 1]).tile([1, N]).reshape([1,N,N,1]) - k_y = paddle.concat((paddle.arange(start=0, end=k_max, step=1), - paddle.arange(start=-k_max, end=0, step=1)), 0).reshape([1, N]).tile([N, 1]).reshape([1,N,N,1]) + k_x = ( + paddle.concat( + ( + paddle.arange(start=0, end=k_max, step=1), + paddle.arange(start=-k_max, end=0, step=1), + ), + 0, + ) + .reshape([N, 1]) + .tile([1, N]) + .reshape([1, N, N, 1]) + ) + k_y = ( + paddle.concat( + ( + paddle.arange(start=0, end=k_max, step=1), + paddle.arange(start=-k_max, end=0, step=1), + ), + 0, + ) + .reshape([1, N]) + .tile([N, 1]) + .reshape([1, N, N, 1]) + ) # Negative Laplacian in Fourier space - lap = (k_x ** 2 + k_y ** 2) + lap = k_x**2 + k_y**2 lap[0, 0, 0, 0] = 1.0 f_h = w_h / lap @@ -59,28 +82,31 @@ def FDM_NS_vorticity(w, v=1/40, t_interval=1.0): wy_h = 1j * k_y * w_h wlap_h = -lap * w_h - ux = paddle.fft.irfft2(ux_h[:, :, :k_max + 1], axes=[1, 2]) - uy = paddle.fft.irfft2(uy_h[:, :, :k_max + 1], axes=[1, 2]) - wx = paddle.fft.irfft2(wx_h[:, :, :k_max+1], axes=[1,2]) - wy = paddle.fft.irfft2(wy_h[:, :, :k_max+1], axes=[1,2]) - wlap = paddle.fft.irfft2(wlap_h[:, :, :k_max+1], axes=[1,2]) + ux = paddle.fft.irfft2(ux_h[:, :, : k_max + 1], axes=[1, 2]) + uy = paddle.fft.irfft2(uy_h[:, :, : k_max + 1], axes=[1, 2]) + wx = paddle.fft.irfft2(wx_h[:, :, : k_max + 1], axes=[1, 2]) + wy = paddle.fft.irfft2(wy_h[:, :, : k_max + 1], axes=[1, 2]) + wlap = paddle.fft.irfft2(wlap_h[:, :, : k_max + 1], axes=[1, 2]) - dt = t_interval / (nt-1) + dt = t_interval / (nt - 1) wt = (w[:, :, :, 2:] - w[:, :, :, :-2]) / (2 * dt) - Du1 = wt + (ux*wx + uy*wy - v*wlap)[...,1:-1] #- forcing + Du1 = wt + (ux * wx + uy * wy - v * wlap)[..., 1:-1] # - forcing return Du1 -def Autograd_Burgers(u, grid, v=1/100): + +def Autograd_Burgers(u, grid, v=1 / 100): from paddle.autograd import grad + gridt, gridx = grid ut = grad(u.sum(), gridt, create_graph=True)[0] ux = grad(u.sum(), gridx, create_graph=True)[0] uxx = grad(ux.sum(), gridx, create_graph=True)[0] - Du = ut + ux*u - v*uxx + Du = ut + ux * u - v * uxx return Du, ux, uxx, ut + def AD_loss(u, u0, grid, index_ic=None, p=None, q=None): batchsize = u.size(0) @@ -92,7 +118,9 @@ def AD_loss(u, u0, grid, index_ic=None, p=None, q=None): nx = u.size(2) u = u.reshape(batchsize, nt, nx) - index_t = paddle.zeros(nx,).long() + index_t = paddle.zeros( + nx, + ).long() index_x = paddle.tensor(range(nx)).long() boundary_u = u[:, index_t, index_x] else: @@ -106,14 +134,16 @@ def AD_loss(u, u0, grid, index_ic=None, p=None, q=None): loss_f = F.mse_loss(Du, f) return loss_ic, loss_f + class LpLoss(object): - ''' + """ loss function with rel/abs Lp loss - ''' + """ + def __init__(self, d=2, p=2, size_average=True, reduction=True): super(LpLoss, self).__init__() - #Dimension and Lp-norm type are postive + # Dimension and Lp-norm type are postive assert d > 0 and p > 0 self.d = d @@ -124,10 +154,12 @@ def __init__(self, d=2, p=2, size_average=True, reduction=True): def abs(self, x, y): num_examples = x.size()[0] - #Assume uniform mesh + # Assume uniform mesh h = 1.0 / (x.shape[1] - 1.0) - all_norms = (h**(self.d/self.p))*paddle.norm(x.reshape(num_examples,-1) - y.reshape(num_examples,-1), self.p, 1) + all_norms = (h ** (self.d / self.p)) * paddle.norm( + x.reshape(num_examples, -1) - y.reshape(num_examples, -1), self.p, 1 + ) if self.reduction: if self.size_average: @@ -140,42 +172,51 @@ def abs(self, x, y): def rel(self, x, y): num_examples = x.shape[0] - diff_norms = paddle.norm(x.reshape([num_examples,-1]) - y.reshape([num_examples,-1]), self.p, 1) - y_norms = paddle.norm(y.reshape([num_examples,-1]), self.p, 1) + diff_norms = paddle.norm( + x.reshape([num_examples, -1]) - y.reshape([num_examples, -1]), self.p, 1 + ) + y_norms = paddle.norm(y.reshape([num_examples, -1]), self.p, 1) if self.reduction: if self.size_average: - return paddle.mean(diff_norms/y_norms) + return paddle.mean(diff_norms / y_norms) else: - return paddle.sum(diff_norms/y_norms) + return paddle.sum(diff_norms / y_norms) - return diff_norms/y_norms + return diff_norms / y_norms def __call__(self, x, y): return self.rel(x, y) + def FDM_Burgers(u, v, D=1): batchsize = u.shape[0] nt = u.shape[1] nx = u.shape[2] u = u.reshape([batchsize, nt, nx]) - dt = D / (nt-1) + dt = D / (nt - 1) dx = D / (nx) u_h = paddle.fft.fft(u, axis=2) # Wavenumbers in y-direction - k_max = nx//2 - k_x = paddle.concat((paddle.arange(start=0, end=k_max, step=1, dtype='float32'), - paddle.arange(start=-k_max, end=0, step=1, dtype='float32')), 0).reshape([1,1,nx]) - ux_h = 2j *np.pi*k_x*u_h - uxx_h = 2j *np.pi*k_x*ux_h - ux = paddle.fft.irfft(ux_h[:, :, :k_max+1], axis=2, n=nx) - uxx = paddle.fft.irfft(uxx_h[:, :, :k_max+1], axis=2, n=nx) + k_max = nx // 2 + k_x = paddle.concat( + ( + paddle.arange(start=0, end=k_max, step=1, dtype="float32"), + paddle.arange(start=-k_max, end=0, step=1, dtype="float32"), + ), + 0, + ).reshape([1, 1, nx]) + ux_h = 2j * np.pi * k_x * u_h + uxx_h = 2j * np.pi * k_x * ux_h + ux = paddle.fft.irfft(ux_h[:, :, : k_max + 1], axis=2, n=nx) + uxx = paddle.fft.irfft(uxx_h[:, :, : k_max + 1], axis=2, n=nx) ut = (u[:, 2:, :] - u[:, :-2, :]) / (2 * dt) - Du = ut + (ux*u - v*uxx)[:,1:-1,:] + Du = ut + (ux * u - v * uxx)[:, 1:-1, :] return Du + def PINO_loss(u, u0, v): batchsize = u.shape[0] nt = u.shape[1] @@ -183,8 +224,8 @@ def PINO_loss(u, u0, v): u = u.reshape([batchsize, nt, nx]) - index_t = paddle.zeros(1,'int32') - index_x = paddle.to_tensor(list(range(nx)),'int32') + index_t = paddle.zeros(1, "int32") + index_x = paddle.to_tensor(list(range(nx)), "int32") boundary_u = paddle.index_select(u, index_t, axis=1).squeeze(1) loss_u = F.mse_loss(boundary_u, u0) @@ -194,7 +235,8 @@ def PINO_loss(u, u0, v): return loss_u, loss_f -def PINO_loss3d(u, u0, forcing, v=1/40, t_interval=1.0): + +def PINO_loss3d(u, u0, forcing, v=1 / 40, t_interval=1.0): batchsize = u.shape[0] nx = u.shape[1] ny = u.shape[2] @@ -207,13 +249,14 @@ def PINO_loss3d(u, u0, forcing, v=1/40, t_interval=1.0): loss_ic = lploss(u_in, u0) Du = FDM_NS_vorticity(u, v, t_interval) - f = forcing.tile([batchsize, 1, 1, nt-2]) + f = forcing.tile([batchsize, 1, 1, nt - 2]) loss_f = lploss(Du, f) return loss_ic, loss_f + def PDELoss(model, x, t, nu): - ''' + """ Compute the residual of PDE: residual = u_t + u * u_x - nu * u_{xx} : (N,1) @@ -223,18 +266,35 @@ def PDELoss(model, x, t, nu): - nu: constant of PDE Return: - mean of residual : scalar - ''' + """ u = model(paddle.cat([x, t], dim=1)) # First backward to compute u_x (shape: N x 1), u_t (shape: N x 1) - grad_x, grad_t = paddle.autograd.grad(outputs=[u.sum()], inputs=[x, t], create_graph=True) + grad_x, grad_t = paddle.autograd.grad( + outputs=[u.sum()], inputs=[x, t], create_graph=True + ) # Second backward to compute u_{xx} (shape N x 1) - gradgrad_x, = paddle.autograd.grad(outputs=[grad_x.sum()], inputs=[x], create_graph=True) + (gradgrad_x,) = paddle.autograd.grad( + outputs=[grad_x.sum()], inputs=[x], create_graph=True + ) residual = grad_t + u * grad_x - nu * gradgrad_x return residual + def get_forcing(S): - x1 = paddle.to_tensor(np.linspace(0, 2*np.pi, S, endpoint=False), dtype=paddle.float32).reshape([S, 1]).tile([1, S]) - x2 = paddle.to_tensor(np.linspace(0, 2*np.pi, S, endpoint=False), dtype=paddle.float32).reshape([1, S]).tile([S, 1]) - return -4 * (paddle.cos(4*(x2))).reshape([1,S,S,1]) \ No newline at end of file + x1 = ( + paddle.to_tensor( + np.linspace(0, 2 * np.pi, S, endpoint=False), dtype=paddle.float32 + ) + .reshape([S, 1]) + .tile([1, S]) + ) + x2 = ( + paddle.to_tensor( + np.linspace(0, 2 * np.pi, S, endpoint=False), dtype=paddle.float32 + ) + .reshape([1, S]) + .tile([S, 1]) + ) + return -4 * (paddle.cos(4 * (x2))).reshape([1, S, S, 1]) diff --git a/jointContribution/PINO/PINO_paddle/train_utils/train_2d.py b/jointContribution/PINO/PINO_paddle/train_utils/train_2d.py index 3e41fc590a..3e2f06ca09 100644 --- a/jointContribution/PINO/PINO_paddle/train_utils/train_2d.py +++ b/jointContribution/PINO/PINO_paddle/train_utils/train_2d.py @@ -1,36 +1,52 @@ import numpy as np import paddle from tqdm import tqdm + +from .losses import LpLoss +from .losses import PINO_loss +from .losses import darcy_loss from .utils import save_checkpoint -from .losses import LpLoss, darcy_loss, PINO_loss - -def train_2d_operator(model, - train_loader, - optimizer, scheduler, - config, - rank=0, log=False, - project='PINO-2d-default', - group='default', - tags=['default'], - use_tqdm=True, - profile=False): - - data_weight = config['train']['xy_loss'] - f_weight = config['train']['f_loss'] + + +def train_2d_operator( + model, + train_loader, + optimizer, + scheduler, + config, + rank=0, + log=False, + project="PINO-2d-default", + group="default", + tags=["default"], + use_tqdm=True, + profile=False, +): + + data_weight = config["train"]["xy_loss"] + f_weight = config["train"]["f_loss"] model.train() myloss = LpLoss(size_average=True) - pbar = range(config['train']['epochs']) + pbar = range(config["train"]["epochs"]) if use_tqdm: pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.1) mesh = train_loader.dataset.mesh - mollifier = paddle.sin(np.pi * mesh[..., 0]) * paddle.sin(np.pi * mesh[..., 1]) * 0.001 + mollifier = ( + paddle.sin(np.pi * mesh[..., 0]) * paddle.sin(np.pi * mesh[..., 1]) * 0.001 + ) pde_mesh = train_loader.dataset.pde_mesh - pde_mol = paddle.sin(np.pi * pde_mesh[..., 0]) * paddle.sin(np.pi * pde_mesh[..., 1]) * 0.001 + pde_mol = ( + paddle.sin(np.pi * pde_mesh[..., 0]) + * paddle.sin(np.pi * pde_mesh[..., 1]) + * 0.001 + ) for e in pbar: - loss_dict = {'train_loss': 0.0, - 'data_loss': 0.0, - 'f_loss': 0.0, - 'test_error': 0.0} + loss_dict = { + "train_loss": 0.0, + "data_loss": 0.0, + "f_loss": 0.0, + "test_error": 0.0, + } for data_ic, u, pde_ic in train_loader: data_ic, u, pde_ic = data_ic.to(rank), u.to(rank), pde_ic.to(rank) @@ -49,44 +65,50 @@ def train_2d_operator(model, loss.backward() optimizer.step() - loss_dict['train_loss'] += loss.item() * y.shape[0] - loss_dict['f_loss'] += f_loss.item() * y.shape[0] - loss_dict['data_loss'] += data_loss.item() * y.shape[0] + loss_dict["train_loss"] += loss.item() * y.shape[0] + loss_dict["f_loss"] += f_loss.item() * y.shape[0] + loss_dict["data_loss"] += data_loss.item() * y.shape[0] scheduler.step() - train_loss_val = loss_dict['train_loss'] / len(train_loader.dataset) - f_loss_val = loss_dict['f_loss'] / len(train_loader.dataset) - data_loss_val = loss_dict['data_loss'] / len(train_loader.dataset) + train_loss_val = loss_dict["train_loss"] / len(train_loader.dataset) + f_loss_val = loss_dict["f_loss"] / len(train_loader.dataset) + data_loss_val = loss_dict["data_loss"] / len(train_loader.dataset) if use_tqdm: pbar.set_description( ( - f'Epoch: {e}, train loss: {train_loss_val:.5f}, ' - f'f_loss: {f_loss_val:.5f}, ' - f'data loss: {data_loss_val:.5f}' + f"Epoch: {e}, train loss: {train_loss_val:.5f}, " + f"f_loss: {f_loss_val:.5f}, " + f"data loss: {data_loss_val:.5f}" ) ) - save_checkpoint(config['train']['save_dir'], - config['train']['save_name'], - model, optimizer) - print('Done!') - -def train_2d_burger(model, - train_loader, v, - optimizer, scheduler, - config, - rank=0, log=False, - project='PINO-2d-default', - group='default', - tags=['default'], - use_tqdm=True): - - data_weight = config['train']['xy_loss'] - f_weight = config['train']['f_loss'] - ic_weight = config['train']['ic_loss'] + save_checkpoint( + config["train"]["save_dir"], config["train"]["save_name"], model, optimizer + ) + print("Done!") + + +def train_2d_burger( + model, + train_loader, + v, + optimizer, + scheduler, + config, + rank=0, + log=False, + project="PINO-2d-default", + group="default", + tags=["default"], + use_tqdm=True, +): + + data_weight = config["train"]["xy_loss"] + f_weight = config["train"]["f_loss"] + ic_weight = config["train"]["ic_loss"] model.train() myloss = LpLoss(size_average=True) - pbar = range(config['train']['epochs']) + pbar = range(config["train"]["epochs"]) if use_tqdm: pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.1) @@ -102,8 +124,10 @@ def train_2d_burger(model, data_loss = myloss(out, y) loss_u, loss_f = PINO_loss(out, x[:, 0, :, 0], v) - total_loss = loss_u * ic_weight + loss_f * f_weight + data_loss * data_weight - + total_loss = ( + loss_u * ic_weight + loss_f * f_weight + data_loss * data_weight + ) + optimizer.clear_grad() total_loss.backward() optimizer.step() @@ -111,7 +135,7 @@ def train_2d_burger(model, data_l2 += data_loss.item() train_pino += loss_f.item() train_loss += total_loss.item() - + scheduler.step() data_l2 /= len(train_loader) train_pino /= len(train_loader) @@ -119,17 +143,20 @@ def train_2d_burger(model, if use_tqdm: pbar.set_description( ( - f'Epoch {e}, train loss: {train_loss:.5f} ' - f'train f error: {train_pino:.5f}; ' - f'data l2 error: {data_l2:.5f}' + f"Epoch {e}, train loss: {train_loss:.5f} " + f"train f error: {train_pino:.5f}; " + f"data l2 error: {data_l2:.5f}" ) ) if e % 100 == 0: - save_checkpoint(config['train']['save_dir'], - config['train']['save_name'].replace('.pt', f'_{e}.pt'), - model, optimizer) - save_checkpoint(config['train']['save_dir'], - config['train']['save_name'], - model, optimizer) - print('Done!') \ No newline at end of file + save_checkpoint( + config["train"]["save_dir"], + config["train"]["save_name"].replace(".pt", f"_{e}.pt"), + model, + optimizer, + ) + save_checkpoint( + config["train"]["save_dir"], config["train"]["save_name"], model, optimizer + ) + print("Done!") diff --git a/jointContribution/PINO/PINO_paddle/train_utils/train_3d.py b/jointContribution/PINO/PINO_paddle/train_utils/train_3d.py index 87b2c6c485..e67db1e8bf 100644 --- a/jointContribution/PINO/PINO_paddle/train_utils/train_3d.py +++ b/jointContribution/PINO/PINO_paddle/train_utils/train_3d.py @@ -1,53 +1,58 @@ -import paddle -from tqdm import tqdm from timeit import default_timer + +import paddle import paddle.nn.functional as F -from .utils import save_checkpoint -from .losses import LpLoss, PINO_loss3d, get_forcing -from .distributed import reduce_loss_dict +from tqdm import tqdm + from .data_utils import sample_data +from .distributed import reduce_loss_dict +from .losses import LpLoss +from .losses import PINO_loss3d +from .losses import get_forcing +from .utils import save_checkpoint - -def train(model, - loader, train_loader, - optimizer, scheduler, - forcing, config, - rank=0, - log=False, - project='PINO-default', - group='FDM', - tags=['Nan'], - use_tqdm=True, - profile=False): +def train( + model, + loader, + train_loader, + optimizer, + scheduler, + forcing, + config, + rank=0, + log=False, + project="PINO-default", + group="FDM", + tags=["Nan"], + use_tqdm=True, + profile=False, +): # data parameters - v = 1 / config['data']['Re'] + v = 1 / config["data"]["Re"] S, T = loader.S, loader.T - t_interval = config['data']['time_interval'] + t_interval = config["data"]["time_interval"] # training settings - batch_size = config['train']['batchsize'] - ic_weight = config['train']['ic_loss'] - f_weight = config['train']['f_loss'] - xy_weight = config['train']['xy_loss'] + batch_size = config["train"]["batchsize"] + ic_weight = config["train"]["ic_loss"] + f_weight = config["train"]["f_loss"] + xy_weight = config["train"]["xy_loss"] model.train() myloss = LpLoss(size_average=True) - pbar = range(config['train']['epochs']) + pbar = range(config["train"]["epochs"]) if use_tqdm: pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.05) - zero = paddle.zeros(1, dtype='float32').to(rank) + zero = paddle.zeros(1, dtype="float32").to(rank) for ep in pbar: - loss_dict = {'train_loss': 0.0, - 'train_ic': 0.0, - 'train_f': 0.0, - 'test_l2': 0.0} + loss_dict = {"train_loss": 0.0, "train_ic": 0.0, "train_f": 0.0, "test_l2": 0.0} log_dict = {} if rank == 0 and profile: - paddle.cuda.synchronize() - t1 = default_timer() + paddle.cuda.synchronize() + t1 = default_timer() # start solving for x, y in train_loader: x, y = x.to(rank), y.to(rank) @@ -61,7 +66,9 @@ def train(model, loss_l2 = myloss(out.view(batch_size, S, S, T), y.view(batch_size, S, S, T)) if ic_weight != 0 or f_weight != 0: - loss_ic, loss_f = PINO_loss3d(out.view(batch_size, S, S, T), x, forcing, v, t_interval) + loss_ic, loss_f = PINO_loss3d( + out.view(batch_size, S, S, T), x, forcing, v, t_interval + ) else: loss_ic, loss_f = zero, zero @@ -70,75 +77,80 @@ def train(model, total_loss.backward() optimizer.step() - loss_dict['train_ic'] += loss_ic - loss_dict['test_l2'] += loss_l2 - loss_dict['train_loss'] += total_loss - loss_dict['train_f'] += loss_f + loss_dict["train_ic"] += loss_ic + loss_dict["test_l2"] += loss_l2 + loss_dict["train_loss"] += total_loss + loss_dict["train_f"] += loss_f if rank == 0 and profile: paddle.cuda.synchronize() t2 = default_timer() - log_dict['Time cost'] = t2 - t1 + log_dict["Time cost"] = t2 - t1 scheduler.step() loss_reduced = reduce_loss_dict(loss_dict) - train_ic = loss_reduced['train_ic'].item() / len(train_loader) - train_f = loss_reduced['train_f'].item() / len(train_loader) - train_loss = loss_reduced['train_loss'].item() / len(train_loader) - test_l2 = loss_reduced['test_l2'].item() / len(train_loader) + train_ic = loss_reduced["train_ic"].item() / len(train_loader) + train_f = loss_reduced["train_f"].item() / len(train_loader) + train_loss = loss_reduced["train_loss"].item() / len(train_loader) + test_l2 = loss_reduced["test_l2"].item() / len(train_loader) log_dict = { - 'Train f error': train_f, - 'Train L2 error': train_ic, - 'Train loss': train_loss, - 'Test L2 error': test_l2 - } + "Train f error": train_f, + "Train L2 error": train_ic, + "Train loss": train_loss, + "Test L2 error": test_l2, + } if rank == 0: if use_tqdm: pbar.set_description( ( - f'Train f error: {train_f:.5f}; Train ic l2 error: {train_ic:.5f}. ' - f'Train loss: {train_loss:.5f}; Test l2 error: {test_l2:.5f}' + f"Train f error: {train_f:.5f}; Train ic l2 error: {train_ic:.5f}. " + f"Train loss: {train_loss:.5f}; Test l2 error: {test_l2:.5f}" ) ) if rank == 0: - save_checkpoint(config['train']['save_dir'], - config['train']['save_name'], - model, optimizer) - -def mixed_train(model, # model of neural operator - train_loader, # dataloader for training with data - S1, T1, # spacial and time dimension for training with data - a_loader, # generator for ICs - S2, T2, # spacial and time dimension for training with equation only - optimizer, # optimizer - scheduler, # learning rate scheduler - config, # configuration dict - log=False, # turn on the wandb - project='PINO-default', # project name - group='FDM', # group name - tags=['Nan'], # tags - use_tqdm=True): # turn on tqdm + save_checkpoint( + config["train"]["save_dir"], config["train"]["save_name"], model, optimizer + ) + + +def mixed_train( + model, # model of neural operator + train_loader, # dataloader for training with data + S1, + T1, # spacial and time dimension for training with data + a_loader, # generator for ICs + S2, + T2, # spacial and time dimension for training with equation only + optimizer, # optimizer + scheduler, # learning rate scheduler + config, # configuration dict + log=False, # turn on the wandb + project="PINO-default", # project name + group="FDM", # group name + tags=["Nan"], # tags + use_tqdm=True, +): # turn on tqdm # data parameters - v = 1 / config['data']['Re'] - t_interval = config['data']['time_interval'] + v = 1 / config["data"]["Re"] + t_interval = config["data"]["time_interval"] forcing_1 = get_forcing(S1) forcing_2 = get_forcing(S2) # training settings - batch_size = config['train']['batchsize'] - ic_weight = config['train']['ic_loss'] - f_weight = config['train']['f_loss'] - xy_weight = config['train']['xy_loss'] - num_data_iter = config['train']['data_iter'] - num_eqn_iter = config['train']['eqn_iter'] + batch_size = config["train"]["batchsize"] + ic_weight = config["train"]["ic_loss"] + f_weight = config["train"]["f_loss"] + xy_weight = config["train"]["xy_loss"] + num_data_iter = config["train"]["data_iter"] + num_eqn_iter = config["train"]["eqn_iter"] model.train() myloss = LpLoss(size_average=True) - pbar = range(config['train']['epochs']) + pbar = range(config["train"]["epochs"]) if use_tqdm: pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.05) - zero = paddle.zeros(1, dtype='float32') + zero = paddle.zeros(1, dtype="float32") train_loader = sample_data(train_loader) for ep in pbar: model.train() @@ -158,13 +170,15 @@ def mixed_train(model, # model of neural operator out = out[..., :-5] x = x[:, :, :, 0, -1] - loss_l2 = myloss(out.reshape([batch_size, S1, S1, T1]), - y.reshape([batch_size, S1, S1, T1])) + loss_l2 = myloss( + out.reshape([batch_size, S1, S1, T1]), + y.reshape([batch_size, S1, S1, T1]), + ) if ic_weight != 0 or f_weight != 0: - loss_ic, loss_f = PINO_loss3d(out.reshape([batch_size, S1, S1, T1]), - x, forcing_1, - v, t_interval) + loss_ic, loss_f = PINO_loss3d( + out.reshape([batch_size, S1, S1, T1]), x, forcing_1, v, t_interval + ) else: loss_ic, loss_f = zero, zero @@ -191,9 +205,9 @@ def mixed_train(model, # model of neural operator out = model(x_in).reshape([batch_size, S2, S2, T2 + 5]) out = out[..., :-5] new_a = new_a[:, :, :, 0, -1] - loss_ic, loss_f = PINO_loss3d(out.reshape([batch_size, S2, S2, T2]), - new_a, forcing_2, - v, t_interval) + loss_ic, loss_f = PINO_loss3d( + out.reshape([batch_size, S2, S2, T2]), new_a, forcing_2, v, t_interval + ) eqn_loss = loss_f * f_weight + loss_ic * ic_weight eqn_loss.backward() optimizer.step() @@ -207,46 +221,52 @@ def mixed_train(model, # model of neural operator if use_tqdm: pbar.set_description( ( - f'Data f error: {train_f:.5f}; Data ic l2 error: {train_ic:.5f}. ' - f'Data train loss: {train_loss:.5f}; Data l2 error: {test_l2:.5f}' - f'Eqn loss: {err_eqn:.5f}' + f"Data f error: {train_f:.5f}; Data ic l2 error: {train_ic:.5f}. " + f"Data train loss: {train_loss:.5f}; Data l2 error: {test_l2:.5f}" + f"Eqn loss: {err_eqn:.5f}" ) ) - save_checkpoint(config['train']['save_dir'], - config['train']['save_name'], - model, optimizer) - -def progressive_train(model, - loader, train_loader, - optimizer, scheduler, - milestones, config, - log=False, - project='PINO-default', - group='FDM', - tags=['Nan'], - use_tqdm=True): + save_checkpoint( + config["train"]["save_dir"], config["train"]["save_name"], model, optimizer + ) + + +def progressive_train( + model, + loader, + train_loader, + optimizer, + scheduler, + milestones, + config, + log=False, + project="PINO-default", + group="FDM", + tags=["Nan"], + use_tqdm=True, +): # data parameters - v = 1 / config['data']['Re'] + v = 1 / config["data"]["Re"] T = loader.T - t_interval = config['data']['time_interval'] + t_interval = config["data"]["time_interval"] # training settings - batch_size = config['train']['batchsize'] - ic_weight = config['train']['ic_loss'] - f_weight = config['train']['f_loss'] - xy_weight = config['train']['xy_loss'] + batch_size = config["train"]["batchsize"] + ic_weight = config["train"]["ic_loss"] + f_weight = config["train"]["f_loss"] + xy_weight = config["train"]["xy_loss"] model.train() myloss = LpLoss(size_average=True) - zero = paddle.zeros(1, dtype='float32') - for milestone, epochs in zip(milestones, config['train']['epochs']): + zero = paddle.zeros(1, dtype="float32") + for milestone, epochs in zip(milestones, config["train"]["epochs"]): pbar = range(epochs) if use_tqdm: pbar = tqdm(pbar, dynamic_ncols=True, smoothing=0.05) S = loader.S // milestone - print(f'Resolution :{S}') + print(f"Resolution :{S}") forcing = get_forcing(S) for ep in pbar: model.train() @@ -265,15 +285,20 @@ def progressive_train(model, out = out[..., :-5] x = x[:, :, :, 0, -1] - loss_l2 = myloss(out.view(batch_size, S, S, T), y.view(batch_size, S, S, T)) + loss_l2 = myloss( + out.view(batch_size, S, S, T), y.view(batch_size, S, S, T) + ) if ic_weight != 0 or f_weight != 0: - loss_ic, loss_f = PINO_loss3d(out.view(batch_size, S, S, T), - x, forcing, v, t_interval) + loss_ic, loss_f = PINO_loss3d( + out.view(batch_size, S, S, T), x, forcing, v, t_interval + ) else: loss_ic, loss_f = zero, zero - total_loss = loss_l2 * xy_weight + loss_f * f_weight + loss_ic * ic_weight + total_loss = ( + loss_l2 * xy_weight + loss_f * f_weight + loss_ic * ic_weight + ) total_loss.backward() @@ -292,12 +317,11 @@ def progressive_train(model, if use_tqdm: pbar.set_description( ( - f'Train f error: {train_f:.5f}; Train ic l2 error: {train_ic:.5f}. ' - f'Train loss: {train_loss:.5f}; Test l2 error: {test_l2:.5f}' + f"Train f error: {train_f:.5f}; Train ic l2 error: {train_ic:.5f}. " + f"Train loss: {train_loss:.5f}; Test l2 error: {test_l2:.5f}" ) ) - save_checkpoint(config['train']['save_dir'], - config['train']['save_name'], - model, optimizer) - + save_checkpoint( + config["train"]["save_dir"], config["train"]["save_name"], model, optimizer + ) diff --git a/jointContribution/PINO/PINO_paddle/train_utils/utils.py b/jointContribution/PINO/PINO_paddle/train_utils/utils.py index 7b116aeb69..9dd352a30f 100644 --- a/jointContribution/PINO/PINO_paddle/train_utils/utils.py +++ b/jointContribution/PINO/PINO_paddle/train_utils/utils.py @@ -1,16 +1,18 @@ import os + import numpy as np import paddle + def vor2vel(w, L=2 * np.pi): - ''' + """ Convert vorticity into velocity Args: w: vorticity with shape (batchsize, num_x, num_y, num_t) Returns: ux, uy with the same shape - ''' + """ batchsize = w.size(0) nx = w.size(1) ny = w.size(2) @@ -22,107 +24,157 @@ def vor2vel(w, L=2 * np.pi): # Wavenumbers in y-direction k_max = nx // 2 N = nx - k_x = paddle.cat((paddle.arange(start=0, end=k_max, step=1, device=device), - paddle.arange(start=-k_max, end=0, step=1, device=device)), 0) \ - .reshape(N, 1).repeat(1, N).reshape(1, N, N, 1) - k_y = paddle.cat((paddle.arange(start=0, end=k_max, step=1, device=device), - paddle.arange(start=-k_max, end=0, step=1, device=device)), 0) \ - .reshape(1, N).repeat(N, 1).reshape(1, N, N, 1) + k_x = ( + paddle.cat( + ( + paddle.arange(start=0, end=k_max, step=1, device=device), + paddle.arange(start=-k_max, end=0, step=1, device=device), + ), + 0, + ) + .reshape(N, 1) + .repeat(1, N) + .reshape(1, N, N, 1) + ) + k_y = ( + paddle.cat( + ( + paddle.arange(start=0, end=k_max, step=1, device=device), + paddle.arange(start=-k_max, end=0, step=1, device=device), + ), + 0, + ) + .reshape(1, N) + .repeat(N, 1) + .reshape(1, N, N, 1) + ) # Negative Laplacian in Fourier space - lap = (k_x ** 2 + k_y ** 2) + lap = k_x**2 + k_y**2 lap[0, 0, 0, 0] = 1.0 f_h = w_h / lap ux_h = 2 * np.pi / L * 1j * k_y * f_h uy_h = -2 * np.pi / L * 1j * k_x * f_h - ux = paddle.fft.irfft2(ux_h[:, :, :k_max + 1], dim=[1, 2]) - uy = paddle.fft.irfft2(uy_h[:, :, :k_max + 1], dim=[1, 2]) + ux = paddle.fft.irfft2(ux_h[:, :, : k_max + 1], dim=[1, 2]) + uy = paddle.fft.irfft2(uy_h[:, :, : k_max + 1], dim=[1, 2]) return ux, uy + def get_sample(N, T, s, p, q): # sample p nodes from Initial Condition, p nodes from Boundary Condition, q nodes from Interior # sample IC index_ic = paddle.randint(s, size=(N, p)) sample_ic_t = paddle.zeros(N, p) - sample_ic_x = index_ic/s + sample_ic_x = index_ic / s # sample BC - sample_bc = paddle.rand(size=(N, p//2)) - sample_bc_t = paddle.cat([sample_bc, sample_bc],dim=1) - sample_bc_x = paddle.cat([paddle.zeros(N, p//2), paddle.ones(N, p//2)],dim=1) + sample_bc = paddle.rand(size=(N, p // 2)) + sample_bc_t = paddle.cat([sample_bc, sample_bc], dim=1) + sample_bc_x = paddle.cat([paddle.zeros(N, p // 2), paddle.ones(N, p // 2)], dim=1) - sample_i_t = -paddle.cos(paddle.rand(size=(N, q))*np.pi/2) + 1 - sample_i_x = paddle.rand(size=(N,q)) + sample_i_t = -paddle.cos(paddle.rand(size=(N, q)) * np.pi / 2) + 1 + sample_i_x = paddle.rand(size=(N, q)) sample_t = paddle.cat([sample_ic_t, sample_bc_t, sample_i_t], dim=1).cuda() sample_t.requires_grad = True sample_x = paddle.cat([sample_ic_x, sample_bc_x, sample_i_x], dim=1).cuda() sample_x.requires_grad = True - sample = paddle.stack([sample_t, sample_x], dim=-1).reshape(N, (p+p+q), 2) + sample = paddle.stack([sample_t, sample_x], dim=-1).reshape(N, (p + p + q), 2) return sample, sample_t, sample_x, index_ic.long() + def get_grid(N, T, s): - gridt = paddle.tensor(np.linspace(0, 1, T), dtype=paddle.float).reshape(1, T, 1).repeat(N, 1, s).cuda() + gridt = ( + paddle.tensor(np.linspace(0, 1, T), dtype=paddle.float) + .reshape(1, T, 1) + .repeat(N, 1, s) + .cuda() + ) gridt.requires_grad = True - gridx = paddle.tensor(np.linspace(0, 1, s+1)[:-1], dtype=paddle.float).reshape(1, 1, s).repeat(N, T, 1).cuda() + gridx = ( + paddle.tensor(np.linspace(0, 1, s + 1)[:-1], dtype=paddle.float) + .reshape(1, 1, s) + .repeat(N, T, 1) + .cuda() + ) gridx.requires_grad = True - grid = paddle.stack([gridt, gridx], dim=-1).reshape(N, T*s, 2) + grid = paddle.stack([gridt, gridx], dim=-1).reshape(N, T * s, 2) return grid, gridt, gridx + def get_2dgrid(S): - ''' + """ get array of points on 2d grid in (0,1)^2 Args: S: resolution Returns: points: flattened grid, ndarray (N, 2) - ''' + """ xarr = np.linspace(0, 1, S) yarr = np.linspace(0, 1, S) - xx, yy = np.meshgrid(xarr, yarr, indexing='ij') + xx, yy = np.meshgrid(xarr, yarr, indexing="ij") points = np.stack([xx.ravel(), yy.ravel()], axis=0).T return points -def paddle2dgrid(num_x, num_y, bot=(0,0), top=(1,1)): + +def paddle2dgrid(num_x, num_y, bot=(0, 0), top=(1, 1)): x_bot, y_bot = bot x_top, y_top = top x_arr = paddle.linspace(x_bot, x_top, num=num_x) y_arr = paddle.linspace(y_bot, y_top, num=num_y) - xx, yy = paddle.meshgrid(x_arr, y_arr, indexing='ij') + xx, yy = paddle.meshgrid(x_arr, y_arr, indexing="ij") mesh = paddle.stack([xx, yy], dim=2) return mesh -def get_grid3d(S, T, time_scale=1.0, device='cpu'): - gridx = paddle.tensor(np.linspace(0, 1, S + 1)[:-1], dtype=paddle.float, device=device) + +def get_grid3d(S, T, time_scale=1.0, device="cpu"): + gridx = paddle.tensor( + np.linspace(0, 1, S + 1)[:-1], dtype=paddle.float, device=device + ) gridx = gridx.reshape(1, S, 1, 1, 1).repeat([1, 1, S, T, 1]) - gridy = paddle.tensor(np.linspace(0, 1, S + 1)[:-1], dtype=paddle.float, device=device) + gridy = paddle.tensor( + np.linspace(0, 1, S + 1)[:-1], dtype=paddle.float, device=device + ) gridy = gridy.reshape(1, 1, S, 1, 1).repeat([1, S, 1, T, 1]) - gridt = paddle.tensor(np.linspace(0, 1 * time_scale, T), dtype=paddle.float, device=device) + gridt = paddle.tensor( + np.linspace(0, 1 * time_scale, T), dtype=paddle.float, device=device + ) gridt = gridt.reshape(1, 1, 1, T, 1).repeat([1, S, S, 1, 1]) return gridx, gridy, gridt + def convert_ic(u0, N, S, T, time_scale=1.0): u0 = u0.reshape(N, S, S, 1, 1).repeat([1, 1, 1, T, 1]) gridx, gridy, gridt = get_grid3d(S, T, time_scale=time_scale, device=u0.device) - a_data = paddle.cat((gridx.repeat([N, 1, 1, 1, 1]), gridy.repeat([N, 1, 1, 1, 1]), - gridt.repeat([N, 1, 1, 1, 1]), u0), dim=-1) + a_data = paddle.cat( + ( + gridx.repeat([N, 1, 1, 1, 1]), + gridy.repeat([N, 1, 1, 1, 1]), + gridt.repeat([N, 1, 1, 1, 1]), + u0, + ), + dim=-1, + ) return a_data + def requires_grad(model, flag=True): for p in model.parameters(): p.requires_grad = flag + def set_grad(tensors, flag=True): for p in tensors: p.requires_grad = flag + def zero_grad(params): - ''' + """ set grad field to 0 - ''' + """ if isinstance(params, paddle.Tensor): if params.grad is not None: params.grad.zero_() @@ -131,14 +183,16 @@ def zero_grad(params): if p.grad is not None: p.grad.zero_() + def count_params(net): count = 0 for p in net.parameters(): count += p.numel() return count + def save_checkpoint(path, name, model, optimizer=None): - ckpt_dir = 'checkpoints/%s/' % path + ckpt_dir = "checkpoints/%s/" % path if not os.path.exists(ckpt_dir): os.makedirs(ckpt_dir) try: @@ -151,11 +205,9 @@ def save_checkpoint(path, name, model, optimizer=None): else: optim_dict = 0.0 - paddle.save({ - 'model': model_state_dict, - 'optim': optim_dict - }, ckpt_dir + name) - print('Checkpoint is saved at %s' % ckpt_dir + name) + paddle.save({"model": model_state_dict, "optim": optim_dict}, ckpt_dir + name) + print("Checkpoint is saved at %s" % ckpt_dir + name) + def save_ckpt(path, model, optimizer=None, scheduler=None): model_state = model.state_dict() @@ -163,20 +215,19 @@ def save_ckpt(path, model, optimizer=None, scheduler=None): optim_state = optimizer.state_dict() else: optim_state = None - + if scheduler: scheduler_state = scheduler.state_dict() else: scheduler_state = None - paddle.save({ - 'model': model_state, - 'optim': optim_state, - 'scheduler': scheduler_state - }, path) - print(f'Checkpoint is saved to {path}') + paddle.save( + {"model": model_state, "optim": optim_state, "scheduler": scheduler_state}, path + ) + print(f"Checkpoint is saved to {path}") + def dict2str(log_dict): - res = '' + res = "" for key, value in log_dict.items(): - res += f'{key}: {value}|' - return res \ No newline at end of file + res += f"{key}: {value}|" + return res diff --git a/jointContribution/PINO/rfcs/PINO.md b/jointContribution/PINO/rfcs/PINO.md index 345896a5ef..2b13409440 100644 --- a/jointContribution/PINO/rfcs/PINO.md +++ b/jointContribution/PINO/rfcs/PINO.md @@ -56,7 +56,7 @@ class FNO2d(nn.Layer): width=64, fc_dim=128, layers=None, in_dim=3, out_dim=1, - act='gelu', + act='gelu', pad_ratio=[0., 0.]): super(FNO2d, self).__init__() @@ -127,4 +127,4 @@ def FDM_Burgers(u, v, D=1): - 202308 : 调研 - 202309 ：基于Paddle API的复现 -- 202310 ：整理项目产出，撰写案例文档 \ No newline at end of file +- 202310 ：整理项目产出，撰写案例文档 diff --git a/jointContribution/XPINNs/requirements.txt b/jointContribution/XPINNs/requirements.txt index 26a2b438ff..74fa65e6b5 100644 --- a/jointContribution/XPINNs/requirements.txt +++ b/jointContribution/XPINNs/requirements.txt @@ -1,3 +1,3 @@ matplotlib -scipy numpy +scipy diff --git a/jointContribution/graphGalerkin/rfcs/graphGalerkin.md b/jointContribution/graphGalerkin/rfcs/graphGalerkin.md index c969876d56..a5025ce992 100644 --- a/jointContribution/graphGalerkin/rfcs/graphGalerkin.md +++ b/jointContribution/graphGalerkin/rfcs/graphGalerkin.md @@ -22,7 +22,7 @@ PaddlePaddle目前无相关模型实现。 ## 3.1 解决的问题 本项目基于嵌入物理信息的图神经网络求解偏微分方程，方程主要以以下形式出现： $$ -\nabla \cdot F(u, \nabla u; \boldsymbol\mu) = S(u, \nabla u; \boldsymbol\mu) \: in \: \Omega +\nabla \cdot F(u, \nabla u; \boldsymbol\mu) = S(u, \nabla u; \boldsymbol\mu) \: in \: \Omega $$ 边界条件为： $$ @@ -78,4 +78,4 @@ $$ - 202307 : 调研 - 202308 ：基于Paddle API的复现 -- 202309 ：整理项目产出，撰写案例文档 \ No newline at end of file +- 202309 ：整理项目产出，撰写案例文档 diff --git a/jointContribution/graphGalerkin/source/FEM_ForwardModel.py b/jointContribution/graphGalerkin/source/FEM_ForwardModel.py index c3cbbb7ce7..661e43243a 100644 --- a/jointContribution/graphGalerkin/source/FEM_ForwardModel.py +++ b/jointContribution/graphGalerkin/source/FEM_ForwardModel.py @@ -1,18 +1,20 @@ -import numpy as np import pdb -import matplotlib.pyplot as plt -from pyCaMOtk.create_mesh_hsphere import mesh_hsphere -from pyCaMOtk.setup_linelptc_sclr_base_handcode import setup_linelptc_sclr_base_handcode +import matplotlib.pyplot as plt +import numpy as np from pyCaMOtk.create_dbc_strct import create_dbc_strct from pyCaMOtk.create_femsp_cg import create_femsp_cg +from pyCaMOtk.create_mesh_hsphere import mesh_hsphere +from pyCaMOtk.setup_linelptc_sclr_base_handcode import setup_linelptc_sclr_base_handcode from pyCaMOtk.solve_fem import solve_fem from pyCaMOtk.visualize_fem import visualize_fem -def analyticalPossion(xcg,Tc,Tb=0): - Ue=Tc*(1-xcg[0,:]**2-xcg[1,:]**2)/4+Tb + +def analyticalPossion(xcg, Tc, Tb=0): + Ue = Tc * (1 - xcg[0, :] ** 2 - xcg[1, :] ** 2) / 4 + Tb return Ue.flatten() -def analyticalConeInterpolation(xcg,Tc,Tb=0): - Ue=Tc*(1-np.sqrt(xcg[0,:]**2+xcg[1,:]**2))/4+Tb - return Ue.flatten() \ No newline at end of file + +def analyticalConeInterpolation(xcg, Tc, Tb=0): + Ue = Tc * (1 - np.sqrt(xcg[0, :] ** 2 + xcg[1, :] ** 2)) / 4 + Tb + return Ue.flatten() diff --git a/jointContribution/graphGalerkin/source/GCNNModel.py b/jointContribution/graphGalerkin/source/GCNNModel.py index f7b8c54fd5..fd07eb85e8 100644 --- a/jointContribution/graphGalerkin/source/GCNNModel.py +++ b/jointContribution/graphGalerkin/source/GCNNModel.py @@ -1,13 +1,15 @@ # from pgl.nn import GCNConv +import sys + import numpy as np import paddle import paddle.nn.initializer as Initializer -import sys sys.path.insert(0, "utils") from ChebConv import ChebConv +from paddle.nn import Layer +from paddle.nn import Linear from paddle.nn.functional import relu -from paddle.nn import Layer, Linear place = paddle.CUDAPlace(0) diff --git a/jointContribution/graphGalerkin/source/setup_prob_eqn_handcode.py b/jointContribution/graphGalerkin/source/setup_prob_eqn_handcode.py index c5c3e83a56..30003a55d5 100644 --- a/jointContribution/graphGalerkin/source/setup_prob_eqn_handcode.py +++ b/jointContribution/graphGalerkin/source/setup_prob_eqn_handcode.py @@ -1,228 +1,264 @@ -import paddle -import numpy as np import pdb -from TensorFEMCore import Double, ReshapeFix +import numpy as np +import paddle +from TensorFEMCore import Double +from TensorFEMCore import ReshapeFix """ ####Possion Equation """ + + class setup_linelptc_sclr_base_handcode(object): - """docstring for setup_linelptc_sclr_base_handcode""" - def __init__(self,ndim,K,f,Qb,bnd2nbc): - self.ndim=ndim - self.K=K - self.f=f - self.Qb=Qb - self.bnd2nbc=bnd2nbc - - self.I=np.eye(self.ndim) - if self.K==None: - self.K=lambda x,el: self.I.reshape(self.ndim**2,1,order='F') #Fortan like - if self.f==None: - self.f=lambda x,el: 0 - if self.Qb==None: - self.Qb=lambda x,n,bnd,el,fc: 0 - - self.eqn=LinearEllipticScalarBaseHandcode() - self.vol_pars_fcn=lambda x,el:np.vstack((self.K(x, el),self.f(x, el),np.nan)) - self.bnd_pars_fcn=lambda x,n,bnd,el,fc:np.vstack((self.K(x,el), - self.f(x,el), - self.Qb(x,n,bnd,el,fc))) - - + """docstring for setup_linelptc_sclr_base_handcode""" + + def __init__(self, ndim, K, f, Qb, bnd2nbc): + self.ndim = ndim + self.K = K + self.f = f + self.Qb = Qb + self.bnd2nbc = bnd2nbc + + self.I = np.eye(self.ndim) + if self.K == None: + self.K = lambda x, el: self.I.reshape( + self.ndim**2, 1, order="F" + ) # Fortan like + if self.f == None: + self.f = lambda x, el: 0 + if self.Qb == None: + self.Qb = lambda x, n, bnd, el, fc: 0 + + self.eqn = LinearEllipticScalarBaseHandcode() + self.vol_pars_fcn = lambda x, el: np.vstack( + (self.K(x, el), self.f(x, el), np.nan) + ) + self.bnd_pars_fcn = lambda x, n, bnd, el, fc: np.vstack( + (self.K(x, el), self.f(x, el), self.Qb(x, n, bnd, el, fc)) + ) + class LinearEllipticScalarBaseHandcode(object): - """docstring for LinearEllipticScalarBaseHandcode""" - def __init__(self): - self.neqn=1 - self.nvar=1 - self.ncomp=1 - - def srcflux(self,UQ,pars,x,model=None): - """ - eval_linelptc_base_handcode_srcflux - """ - # Extract information from input - q=UQ[0,1:] - q=ReshapeFix(q,[len(q),1]) - self.ndim=len(q) - try: - k=np.reshape(pars[0:self.ndim**2], - (self.ndim,self.ndim),order='F') - except: - k=paddle.reshape(pars[0:self.ndim**2], - (self.ndim,self.ndim)) - f=pars[self.ndim**2] - try: - temp_flag=(f.requires_grad) - f=f.reshape([1,1]) - except: - f=paddle.to_tensor(f, dtype='float32').reshape([1,1]) - - k_ml=paddle.to_tensor(k, dtype='float32') - # Define flux and source - SF=paddle.concat((f,-1*paddle.mm(k_ml,q)),axis=0) - - # Define partial derivative - dSFdU=np.zeros([self.neqn, self.ndim+1, self.ncomp,self.ndim+1]) - try: - dSFdU[:,1:,:,1:]=np.reshape(-1*k,[self.neqn, self.ndim,self.ncomp,self.ndim]) - except: - k=k.detach().cpu().numpy() - dSFdU[:,1:,:,1:]=np.reshape(-1*k,[self.neqn, self.ndim,self.ncomp,self.ndim]) - dSFdU=paddle.to_tensor(dSFdU, dtype='float32') - return SF, dSFdU - - def bndstvcflux(self,nbcnbr,UQ,pars,x,n): - nvar=UQ.shape[0] - ndim=UQ.shape[1]-1 - - Ub=UQ[:,0] - dUb=np.zeros([nvar,nvar,self.ndim+1]) - dUb[:,:,0]=np.eye(nvar) - - Fn=pars[ndim**2+1] - dFn=np.zeros([nvar,nvar,self.ndim+1]) - dUb=Double(dUb) - Fn=Double(Fn) - dFn=Double(dFn) - return Ub,dUb,Fn,dFn + """docstring for LinearEllipticScalarBaseHandcode""" + + def __init__(self): + self.neqn = 1 + self.nvar = 1 + self.ncomp = 1 + + def srcflux(self, UQ, pars, x, model=None): + """ + eval_linelptc_base_handcode_srcflux + """ + # Extract information from input + q = UQ[0, 1:] + q = ReshapeFix(q, [len(q), 1]) + self.ndim = len(q) + try: + k = np.reshape(pars[0 : self.ndim**2], (self.ndim, self.ndim), order="F") + except: + k = paddle.reshape(pars[0 : self.ndim**2], (self.ndim, self.ndim)) + f = pars[self.ndim**2] + try: + temp_flag = f.requires_grad + f = f.reshape([1, 1]) + except: + f = paddle.to_tensor(f, dtype="float32").reshape([1, 1]) + + k_ml = paddle.to_tensor(k, dtype="float32") + # Define flux and source + SF = paddle.concat((f, -1 * paddle.mm(k_ml, q)), axis=0) + + # Define partial derivative + dSFdU = np.zeros([self.neqn, self.ndim + 1, self.ncomp, self.ndim + 1]) + try: + dSFdU[:, 1:, :, 1:] = np.reshape( + -1 * k, [self.neqn, self.ndim, self.ncomp, self.ndim] + ) + except: + k = k.detach().cpu().numpy() + dSFdU[:, 1:, :, 1:] = np.reshape( + -1 * k, [self.neqn, self.ndim, self.ncomp, self.ndim] + ) + dSFdU = paddle.to_tensor(dSFdU, dtype="float32") + return SF, dSFdU + + def bndstvcflux(self, nbcnbr, UQ, pars, x, n): + nvar = UQ.shape[0] + ndim = UQ.shape[1] - 1 + + Ub = UQ[:, 0] + dUb = np.zeros([nvar, nvar, self.ndim + 1]) + dUb[:, :, 0] = np.eye(nvar) + + Fn = pars[ndim**2 + 1] + dFn = np.zeros([nvar, nvar, self.ndim + 1]) + dUb = Double(dUb) + Fn = Double(Fn) + dFn = Double(dFn) + return Ub, dUb, Fn, dFn + """ ####Linear Elasticity Equation """ + class setup_linelast_base_handcode(object): - """docstring for setup_linelast_base_handcode""" - def __init__(self,ndim,lam,mu,f,tb,bnd2nbc): - self.bnd2nbc=bnd2nbc - self.eqn=LinearElasticityBaseHandcode(ndim) - self.vol_pars_fcn=lambda x, el: np.vstack((lam(x,el), - mu(x,el), - f(x,el), - np.zeros([ndim,1])+np.nan)) - self.bnd_pars_fcn=lambda x,n,bnd,el,fc:np.vstack((lam(x, el), - mu(x, el), - f(x, el), - tb(x, n, bnd, el, fc))) - - + """docstring for setup_linelast_base_handcode""" + + def __init__(self, ndim, lam, mu, f, tb, bnd2nbc): + self.bnd2nbc = bnd2nbc + self.eqn = LinearElasticityBaseHandcode(ndim) + self.vol_pars_fcn = lambda x, el: np.vstack( + (lam(x, el), mu(x, el), f(x, el), np.zeros([ndim, 1]) + np.nan) + ) + self.bnd_pars_fcn = lambda x, n, bnd, el, fc: np.vstack( + (lam(x, el), mu(x, el), f(x, el), tb(x, n, bnd, el, fc)) + ) + class LinearElasticityBaseHandcode(object): - """docstring for LinearElasticityBaseHandcode""" - def __init__(self,ndim): - self.neqn=ndim - self.nvar=ndim - self.bndstvcflux=\ - lambda nbcnbr, UQ, pars, x, n:\ - eval_linelast_base_handcode_bndstvc_intr_bndflux_pars(UQ, pars, x, n) - self.srcflux=lambda UQ,pars,x:\ - eval_linelast_base_handcode_srcflux(UQ, pars, x) + """docstring for LinearElasticityBaseHandcode""" + + def __init__(self, ndim): + self.neqn = ndim + self.nvar = ndim + self.bndstvcflux = lambda nbcnbr, UQ, pars, x, n: eval_linelast_base_handcode_bndstvc_intr_bndflux_pars( + UQ, pars, x, n + ) + self.srcflux = lambda UQ, pars, x: eval_linelast_base_handcode_srcflux( + UQ, pars, x + ) + def eval_linelast_base_handcode_srcflux(UQ, pars, x): - q=UQ[:,1:] - ndim=q.shape[0] - # Define information regarding size of the system - neqn=ndim - ncomp=ndim - - # Extract parameters - lam=pars[0] - mu=pars[1] - f=pars[2:2+ndim] - F=-lam*paddle.trace(q)*(Double(np.eye(ndim)))-mu*(q+q.T) - try: - S=Double(f.reshape([ndim,1],order='F')) - except: - S=f.reshape([ndim,1]) - SF=paddle.concat((S,F),axis=1) - dSFdU=Double(np.zeros([neqn,ndim+1,ncomp,ndim+1])) - for i in range(ndim): - for j in range(ndim): - dSFdU[i,1+i,j,1+j]=dSFdU[i,1+i,j,1+j]-lam - dSFdU[i,1+j,i,1+j]=dSFdU[i,1+j,i,1+j]-mu - dSFdU[i,1+j,j,1+i]=dSFdU[i,1+j,j,1+i]-mu - return SF, dSFdU - -def eval_linelast_base_handcode_bndstvc_intr_bndflux_pars(UQ,pars,x,n): - nvar=UQ.shape[0] - ndim=UQ.shape[1]-1 - - Ub=UQ[:,0] - dUb=np.zeros([nvar,nvar,ndim+1]) - dUb[:,:,0]=np.eye(nvar) - Fn=-pars[-ndim:] - dFn=np.zeros([nvar,nvar,ndim+1]) - dUb=Double(dUb) - Fn=ReshapeFix(Double(Fn),[len(Fn),1],order='F') - dFn=Double(dFn) - return Ub,dUb,Fn,dFn + q = UQ[:, 1:] + ndim = q.shape[0] + # Define information regarding size of the system + neqn = ndim + ncomp = ndim + + # Extract parameters + lam = pars[0] + mu = pars[1] + f = pars[2 : 2 + ndim] + F = -lam * paddle.trace(q) * (Double(np.eye(ndim))) - mu * (q + q.T) + try: + S = Double(f.reshape([ndim, 1], order="F")) + except: + S = f.reshape([ndim, 1]) + SF = paddle.concat((S, F), axis=1) + dSFdU = Double(np.zeros([neqn, ndim + 1, ncomp, ndim + 1])) + for i in range(ndim): + for j in range(ndim): + dSFdU[i, 1 + i, j, 1 + j] = dSFdU[i, 1 + i, j, 1 + j] - lam + dSFdU[i, 1 + j, i, 1 + j] = dSFdU[i, 1 + j, i, 1 + j] - mu + dSFdU[i, 1 + j, j, 1 + i] = dSFdU[i, 1 + j, j, 1 + i] - mu + return SF, dSFdU + + +def eval_linelast_base_handcode_bndstvc_intr_bndflux_pars(UQ, pars, x, n): + nvar = UQ.shape[0] + ndim = UQ.shape[1] - 1 + + Ub = UQ[:, 0] + dUb = np.zeros([nvar, nvar, ndim + 1]) + dUb[:, :, 0] = np.eye(nvar) + Fn = -pars[-ndim:] + dFn = np.zeros([nvar, nvar, ndim + 1]) + dUb = Double(dUb) + Fn = ReshapeFix(Double(Fn), [len(Fn), 1], order="F") + dFn = Double(dFn) + return Ub, dUb, Fn, dFn + """ #### Inconpressible Navier Stokes Equation """ + + class setup_ins_base_handcode(object): - """docstring for setup_ins_base_handcode""" - def __init__(self,ndim,rho,nu,tb,bnd2nbc): - self.eqn=IncompressibleNavierStokes(ndim) - self.bnd2nbc=bnd2nbc - self.vol_pars_fcn=lambda x,el:np.vstack([rho(x, el), - nu(x, el), - np.zeros([ndim+1,1])+np.nan]) - self.bnd_pars_fcn=lambda x,n,bnd,el,fc:np.vstack([rho(x,el), - nu(x,el), - tb(x,n,bnd,el,fc)]) + """docstring for setup_ins_base_handcode""" + + def __init__(self, ndim, rho, nu, tb, bnd2nbc): + self.eqn = IncompressibleNavierStokes(ndim) + self.bnd2nbc = bnd2nbc + self.vol_pars_fcn = lambda x, el: np.vstack( + [rho(x, el), nu(x, el), np.zeros([ndim + 1, 1]) + np.nan] + ) + self.bnd_pars_fcn = lambda x, n, bnd, el, fc: np.vstack( + [rho(x, el), nu(x, el), tb(x, n, bnd, el, fc)] + ) + class IncompressibleNavierStokes(object): - """docstring for IncompressibleNavierStokes""" - def __init__(self,ndim): - self.ndim=ndim - self.nvar=ndim+1 - self.srcflux=lambda UQ,pars,x:\ - eval_ins_base_handcode_srcflux(UQ,pars,x) - self.bndstvcflux=lambda nbcnbr,UQ,pars,x,n:\ - eval_ins_base_handcode_bndstvc_intr_bndflux_pars(UQ,pars,x,n) - -def eval_ins_base_handcode_srcflux(UQ,pars,x): - u=UQ[:,0]; q=UQ[:,1:] - ndim=u.shape[0]-1 - neqn=ndim+1 - ncomp=ndim+1 - rho=pars[0] - nu=pars[1] - v=u[0:ndim] - - v=ReshapeFix(v,[len(v),1],'F') - - p=u[-1] - dv=q[0:ndim,:] - S=paddle.concat([-rho*paddle.mm(dv,v),-paddle.trace(dv).reshape([1,1])],axis=0) - - F=paddle.concat([-rho*nu*dv+p*paddle.eye(ndim, dtype='float32'), - paddle.zeros([1,ndim], dtype='float32')],axis=0) - - - SF=paddle.concat([S,F],axis=1) - - dSFdUQ=np.zeros([neqn,ndim+1,ncomp,ndim+1]) - dSFdUQ[:,0,:,0]=np.vstack([np.hstack([-rho*dv.detach().cpu().numpy(),np.zeros([ndim,1])]), np.zeros([1,ndim+1])]) - for i in range(ndim): - dSFdUQ[i,0,i,1:]=-rho*v.detach().cpu().numpy().reshape(dSFdUQ[i,0,i,1:].shape,order='F') - dSFdUQ[-1,0,0:-1,1:]=np.reshape(-np.eye(ndim),[1,ndim,ndim],order='F') - dSFdUQ[0:-1,1:,-1,0]=np.eye(ndim) - for i in range(ndim): - for j in range(ndim): - dSFdUQ[i,1+j,i,1+j]=dSFdUQ[i,1+j,i,1+j]-rho*nu - dSFdUQ=Double(dSFdUQ) - return SF,dSFdUQ - -def eval_ins_base_handcode_bndstvc_intr_bndflux_pars(UQ,pars,x,n): - nvar=UQ.shape[0] - ndim=UQ.shape[1]-1 - Ub=UQ[:,0] - dUb=np.zeros([nvar,nvar,ndim+1]) - dUb[:,:,0]=np.eye(nvar) - Fn=-pars[-ndim-1:].reshape([-1,1]) - dFn=np.zeros([nvar,nvar,ndim+1]) - return Ub,Double(dUb),Double(Fn),Double(dFn) + """docstring for IncompressibleNavierStokes""" + + def __init__(self, ndim): + self.ndim = ndim + self.nvar = ndim + 1 + self.srcflux = lambda UQ, pars, x: eval_ins_base_handcode_srcflux(UQ, pars, x) + self.bndstvcflux = lambda nbcnbr, UQ, pars, x, n: eval_ins_base_handcode_bndstvc_intr_bndflux_pars( + UQ, pars, x, n + ) + + +def eval_ins_base_handcode_srcflux(UQ, pars, x): + u = UQ[:, 0] + q = UQ[:, 1:] + ndim = u.shape[0] - 1 + neqn = ndim + 1 + ncomp = ndim + 1 + rho = pars[0] + nu = pars[1] + v = u[0:ndim] + + v = ReshapeFix(v, [len(v), 1], "F") + + p = u[-1] + dv = q[0:ndim, :] + S = paddle.concat( + [-rho * paddle.mm(dv, v), -paddle.trace(dv).reshape([1, 1])], axis=0 + ) + + F = paddle.concat( + [ + -rho * nu * dv + p * paddle.eye(ndim, dtype="float32"), + paddle.zeros([1, ndim], dtype="float32"), + ], + axis=0, + ) + + SF = paddle.concat([S, F], axis=1) + + dSFdUQ = np.zeros([neqn, ndim + 1, ncomp, ndim + 1]) + dSFdUQ[:, 0, :, 0] = np.vstack( + [ + np.hstack([-rho * dv.detach().cpu().numpy(), np.zeros([ndim, 1])]), + np.zeros([1, ndim + 1]), + ] + ) + for i in range(ndim): + dSFdUQ[i, 0, i, 1:] = -rho * v.detach().cpu().numpy().reshape( + dSFdUQ[i, 0, i, 1:].shape, order="F" + ) + dSFdUQ[-1, 0, 0:-1, 1:] = np.reshape(-np.eye(ndim), [1, ndim, ndim], order="F") + dSFdUQ[0:-1, 1:, -1, 0] = np.eye(ndim) + for i in range(ndim): + for j in range(ndim): + dSFdUQ[i, 1 + j, i, 1 + j] = dSFdUQ[i, 1 + j, i, 1 + j] - rho * nu + dSFdUQ = Double(dSFdUQ) + return SF, dSFdUQ + + +def eval_ins_base_handcode_bndstvc_intr_bndflux_pars(UQ, pars, x, n): + nvar = UQ.shape[0] + ndim = UQ.shape[1] - 1 + Ub = UQ[:, 0] + dUb = np.zeros([nvar, nvar, ndim + 1]) + dUb[:, :, 0] = np.eye(nvar) + Fn = -pars[-ndim - 1 :].reshape([-1, 1]) + dFn = np.zeros([nvar, nvar, ndim + 1]) + return Ub, Double(dUb), Double(Fn), Double(dFn) diff --git a/jointContribution/graphGalerkin/utils/init.py b/jointContribution/graphGalerkin/utils/init.py index e7f7ba175d..3cbdebe676 100644 --- a/jointContribution/graphGalerkin/utils/init.py +++ b/jointContribution/graphGalerkin/utils/init.py @@ -1,74 +1,82 @@ -from typing import Any - import math +from typing import Any import paddle from paddle import Tensor from paddle.nn.initializer import Orthogonal + def uniform(size: int, value: Any): if isinstance(value, Tensor): bound = 1.0 / math.sqrt(size) value.data.uniform_(-bound, bound) else: - for v in value.parameters() if hasattr(value, 'parameters') else []: + for v in value.parameters() if hasattr(value, "parameters") else []: uniform(size, v) - for v in value.buffers() if hasattr(value, 'buffers') else []: + for v in value.buffers() if hasattr(value, "buffers") else []: uniform(size, v) + def kaiming_uniform(value: Any, fan: int, a: float): if isinstance(value, Tensor): bound = math.sqrt(6 / ((1 + a**2) * fan)) value.data.uniform_(-bound, bound) else: - for v in value.parameters() if hasattr(value, 'parameters') else []: + for v in value.parameters() if hasattr(value, "parameters") else []: kaiming_uniform(v, fan, a) - for v in value.buffers() if hasattr(value, 'buffers') else []: + for v in value.buffers() if hasattr(value, "buffers") else []: kaiming_uniform(v, fan, a) + def glorot(value: Any): if isinstance(value, Tensor): stdv = math.sqrt(6.0 / (value.shape[-2] + value.shape[-1])) value = paddle.uniform(value.shape, value.dtype, -stdv, stdv) else: - for v in value.parameters() if hasattr(value, 'parameters') else []: + for v in value.parameters() if hasattr(value, "parameters") else []: glorot(v) - for v in value.buffers() if hasattr(value, 'buffers') else []: + for v in value.buffers() if hasattr(value, "buffers") else []: glorot(v) + def glorot_orthogonal(tensor, scale): if tensor is not None: tensor = paddle.create_parameter(tensor.shape, attr=Orthogonal()) - scale /= ((tensor.size(-2) + tensor.size(-1)) * tensor.var()) + scale /= (tensor.size(-2) + tensor.size(-1)) * tensor.var() tensor.data *= scale.sqrt() + def constant(value: Any, fill_value: float): if isinstance(value, Tensor): value = paddle.full(value.shape, fill_value, value.dtype) else: - for v in value.parameters() if hasattr(value, 'parameters') else []: + for v in value.parameters() if hasattr(value, "parameters") else []: constant(v, fill_value) - for v in value.buffers() if hasattr(value, 'buffers') else []: + for v in value.buffers() if hasattr(value, "buffers") else []: constant(v, fill_value) + def zeros(value: Any): - constant(value, 0.) + constant(value, 0.0) + def ones(tensor: Any): - constant(tensor, 1.) + constant(tensor, 1.0) + def normal(value: Any, mean: float, std: float): if isinstance(value, Tensor): value.data.normal_(mean, std) else: - for v in value.parameters() if hasattr(value, 'parameters') else []: + for v in value.parameters() if hasattr(value, "parameters") else []: normal(v, mean, std) - for v in value.buffers() if hasattr(value, 'buffers') else []: + for v in value.buffers() if hasattr(value, "buffers") else []: normal(v, mean, std) + def reset(value: Any): - if hasattr(value, 'reset_parameters'): + if hasattr(value, "reset_parameters"): value.reset_parameters() else: - for child in value.children() if hasattr(value, 'children') else []: + for child in value.children() if hasattr(value, "children") else []: reset(child) diff --git a/jointContribution/graphGalerkin/utils/inspector.py b/jointContribution/graphGalerkin/utils/inspector.py index 249187a549..9d5b50f422 100644 --- a/jointContribution/graphGalerkin/utils/inspector.py +++ b/jointContribution/graphGalerkin/utils/inspector.py @@ -1,16 +1,23 @@ -import re import inspect +import re from collections import OrderedDict -from typing import Dict, List, Any, Optional, Callable, Set, Tuple +from typing import Any +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Set +from typing import Tuple + import pyparsing as pp + class Inspector(object): def __init__(self, base_class: Any): self.base_class: Any = base_class self.params: Dict[str, Dict[str, Any]] = {} - def inspect(self, func: Callable, - pop_first: bool = False) -> Dict[str, Any]: + def inspect(self, func: Callable, pop_first: bool = False) -> Dict[str, Any]: params = inspect.signature(func).parameters params = OrderedDict(params) if pop_first: @@ -24,7 +31,7 @@ def keys(self, func_names: Optional[List[str]] = None) -> Set[str]: return set(keys) def __implements__(self, cls, func_name: str) -> bool: - if cls.__name__ == 'MessagePassing': + if cls.__name__ == "MessagePassing": return False if func_name in cls.__dict__.keys(): return True @@ -39,30 +46,31 @@ def distribute(self, func_name, kwargs: Dict[str, Any]): data = kwargs.get(key, inspect.Parameter.empty) if data is inspect.Parameter.empty: if param.default is inspect.Parameter.empty: - raise TypeError(f'Required parameter {key} is empty.') + raise TypeError(f"Required parameter {key} is empty.") data = param.default out[key] = data return out + def func_header_repr(func: Callable, keep_annotation: bool = True) -> str: source = inspect.getsource(func) signature = inspect.signature(func) if keep_annotation: - return ''.join(re.split(r'(\).*?:.*?\n)', source, - maxsplit=1)[:2]).strip() + return "".join(re.split(r"(\).*?:.*?\n)", source, maxsplit=1)[:2]).strip() - params_repr = ['self'] + params_repr = ["self"] for param in signature.parameters.values(): params_repr.append(param.name) if param.default is not inspect.Parameter.empty: - params_repr[-1] += f'={param.default}' + params_repr[-1] += f"={param.default}" return f'def {func.__name__}({", ".join(params_repr)}):' + def func_body_repr(func: Callable, keep_annotation: bool = True) -> str: source = inspect.getsource(func) - body_repr = re.split(r'\).*?:.*?\n', source, maxsplit=1)[1] + body_repr = re.split(r"\).*?:.*?\n", source, maxsplit=1)[1] if not keep_annotation: - body_repr = re.sub(r'\s*# type:.*\n', '', body_repr) - return body_repr \ No newline at end of file + body_repr = re.sub(r"\s*# type:.*\n", "", body_repr) + return body_repr diff --git a/jointContribution/graphGalerkin/utils/linear.py b/jointContribution/graphGalerkin/utils/linear.py index b74f953c3e..861008ebf2 100644 --- a/jointContribution/graphGalerkin/utils/linear.py +++ b/jointContribution/graphGalerkin/utils/linear.py @@ -1,16 +1,15 @@ -from typing import Optional -from collections import OrderedDict - import copy import math +from collections import OrderedDict +from typing import Optional +import init as inits import paddle -from paddle import nn -from paddle import Tensor import paddle.nn.functional as F +from paddle import Tensor +from paddle import nn from paddle.nn.initializer import Uniform -import init as inits class Linear(nn.Layer): r"""Applies a linear tranformation to the incoming data @@ -38,9 +37,15 @@ class Linear(nn.Layer): If set to :obj:`None`, will match default bias initialization of :class:`torch.nn.Linear`. (default: :obj:`None`) """ - def __init__(self, in_channels: int, out_channels: int, bias: bool = True, - weight_initializer: Optional[str] = None, - bias_initializer: Optional[str] = None): + + def __init__( + self, + in_channels: int, + out_channels: int, + bias: bool = True, + weight_initializer: Optional[str] = None, + bias_initializer: Optional[str] = None, + ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels @@ -49,26 +54,32 @@ def __init__(self, in_channels: int, out_channels: int, bias: bool = True, self._parameters = OrderedDict() if in_channels > 0: - self.weight = paddle.create_parameter(shape=[in_channels, out_channels], dtype=paddle.float32) + self.weight = paddle.create_parameter( + shape=[in_channels, out_channels], dtype=paddle.float32 + ) else: self.weight = nn.parameter.UninitializedParameter() - self._hook = self.register_forward_pre_hook( - self.initialize_parameters) + self._hook = self.register_forward_pre_hook(self.initialize_parameters) if bias: - self.bias = paddle.create_parameter(shape=[out_channels], dtype=paddle.float32) + self.bias = paddle.create_parameter( + shape=[out_channels], dtype=paddle.float32 + ) else: - self.register_parameter('bias', None) + self.register_parameter("bias", None) - self._load_hook = self.register_state_dict_hook( - self._lazy_load_hook) + self._load_hook = self.register_state_dict_hook(self._lazy_load_hook) self.reset_parameters() def __deepcopy__(self, memo): - out = Linear(self.in_channels, self.out_channels, self.bias - is not None, self.weight_initializer, - self.bias_initializer) + out = Linear( + self.in_channels, + self.out_channels, + self.bias is not None, + self.weight_initializer, + self.bias_initializer, + ) if self.in_channels > 0: out.weight = copy.deepcopy(self.weight, memo) if self.bias is not None: @@ -77,35 +88,41 @@ def __deepcopy__(self, memo): def reset_parameters(self): # if isinstance(self.weight, nn.parameter.UninitializedParameter): - if self.in_channels<=0: + if self.in_channels <= 0: pass - elif self.weight_initializer == 'glorot': + elif self.weight_initializer == "glorot": inits.glorot(self.weight) - elif self.weight_initializer == 'uniform': + elif self.weight_initializer == "uniform": bound = 1.0 / math.sqrt(self.weight.size(-1)) - self.weight = paddle.create_parameter(shape=self.weight.shape, dtype=paddle.float32, attr=paddle.ParamAttr(initializer=Uniform(-bound, bound))) - elif self.weight_initializer == 'kaiming_uniform': - inits.kaiming_uniform(self.weight, fan=self.in_channels, - a=math.sqrt(5)) + self.weight = paddle.create_parameter( + shape=self.weight.shape, + dtype=paddle.float32, + attr=paddle.ParamAttr(initializer=Uniform(-bound, bound)), + ) + elif self.weight_initializer == "kaiming_uniform": + inits.kaiming_uniform(self.weight, fan=self.in_channels, a=math.sqrt(5)) elif self.weight_initializer is None: - inits.kaiming_uniform(self.weight, fan=self.in_channels, - a=math.sqrt(5)) + inits.kaiming_uniform(self.weight, fan=self.in_channels, a=math.sqrt(5)) else: - raise RuntimeError(f"Linear layer weight initializer " - f"'{self.weight_initializer}' is not supported") + raise RuntimeError( + f"Linear layer weight initializer " + f"'{self.weight_initializer}' is not supported" + ) # if isinstance(self.weight, nn.parameter.UninitializedParameter): - if self.in_channels<=0: + if self.in_channels <= 0: pass elif self.bias is None: pass - elif self.bias_initializer == 'zeros': + elif self.bias_initializer == "zeros": inits.zeros(self.bias) elif self.bias_initializer is None: inits.uniform(self.in_channels, self.bias) else: - raise RuntimeError(f"Linear layer bias initializer " - f"'{self.bias_initializer}' is not supported") + raise RuntimeError( + f"Linear layer bias initializer " + f"'{self.bias_initializer}' is not supported" + ) def forward(self, x: Tensor) -> Tensor: """""" @@ -118,41 +135,49 @@ def initialize_parameters(self, module, input): self.weight.materialize((self.out_channels, self.in_channels)) self.reset_parameters() self._hook.remove() - delattr(self, '_hook') + delattr(self, "_hook") def _save_to_state_dict(self, destination, prefix, keep_vars): if isinstance(self.weight, nn.parameter.UninitializedParameter): - destination[prefix + 'weight'] = self.weight + destination[prefix + "weight"] = self.weight else: - destination[prefix + 'weight'] = self.weight.detach() + destination[prefix + "weight"] = self.weight.detach() if self.bias is not None: - destination[prefix + 'bias'] = self.bias.detach() - - def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict, - missing_keys, unexpected_keys, error_msgs): - - weight = state_dict[prefix + 'weight'] + destination[prefix + "bias"] = self.bias.detach() + + def _lazy_load_hook( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + + weight = state_dict[prefix + "weight"] if isinstance(weight, nn.parameter.UninitializedParameter): self.in_channels = -1 self.weight = nn.parameter.UninitializedParameter() - if not hasattr(self, '_hook'): - self._hook = self.register_forward_pre_hook( - self.initialize_parameters) + if not hasattr(self, "_hook"): + self._hook = self.register_forward_pre_hook(self.initialize_parameters) elif isinstance(self.weight, nn.parameter.UninitializedParameter): self.in_channels = weight.size(-1) self.weight.materialize((self.out_channels, self.in_channels)) - if hasattr(self, '_hook'): + if hasattr(self, "_hook"): self._hook.remove() - delattr(self, '_hook') + delattr(self, "_hook") def __repr__(self) -> str: - return (f'{self.__class__.__name__}({self.in_channels}, ' - f'{self.out_channels}, bias={self.bias is not None})') - + return ( + f"{self.__class__.__name__}({self.in_channels}, " + f"{self.out_channels}, bias={self.bias is not None})" + ) + def register_parameter(self, name: str, param) -> None: if param is None: self._parameters[name] = None else: self._parameters[name] = param - diff --git a/jointContribution/graphGalerkin/utils/message_passing.py b/jointContribution/graphGalerkin/utils/message_passing.py index 98ef276f2b..86becbc1a8 100644 --- a/jointContribution/graphGalerkin/utils/message_passing.py +++ b/jointContribution/graphGalerkin/utils/message_passing.py @@ -1,17 +1,20 @@ -from typing import List, Optional, Set, Tuple - -from uuid import uuid1 -from inspect import Parameter from collections import OrderedDict +from inspect import Parameter +from typing import List +from typing import Optional +from typing import Set +from typing import Tuple +from uuid import uuid1 import paddle -from paddle import Tensor - from inspector import Inspector -from scatter import scatter +from paddle import Tensor from paddle.nn import Layer +from scatter import scatter Size = Optional[Tuple[int, int]] + + class MessagePassing(Layer): r"""Base class for creating message passing layers of the form @@ -55,21 +58,33 @@ class MessagePassing(Layer): """ special_args: Set[str] = { - 'edge_index', 'adj_t', 'edge_index_i', 'edge_index_j', 'size', - 'size_i', 'size_j', 'ptr', 'index', 'dim_size' + "edge_index", + "adj_t", + "edge_index_i", + "edge_index_j", + "size", + "size_i", + "size_j", + "ptr", + "index", + "dim_size", } - def __init__(self, aggr: Optional[str] = "add", - flow: str = "source_to_target", node_dim: int = -2, - decomposed_layers: int = 1): + def __init__( + self, + aggr: Optional[str] = "add", + flow: str = "source_to_target", + node_dim: int = -2, + decomposed_layers: int = 1, + ): super().__init__() self.aggr = aggr - assert self.aggr in ['add', 'mean', 'max', None] + assert self.aggr in ["add", "mean", "max", None] self.flow = flow - assert self.flow in ['source_to_target', 'target_to_source'] + assert self.flow in ["source_to_target", "target_to_source"] self.node_dim = node_dim self.decomposed_layers = decomposed_layers @@ -82,14 +97,17 @@ def __init__(self, aggr: Optional[str] = "add", self.inspector.inspect(self.message_and_aggregate, pop_first=True) self.__user_args__ = self.inspector.keys( - ['message', 'aggregate', 'update']).difference(self.special_args) + ["message", "aggregate", "update"] + ).difference(self.special_args) self.__fused_user_args__ = self.inspector.keys( - ['message_and_aggregate', 'update']).difference(self.special_args) - self.__edge_user_args__ = self.inspector.keys( - ['edge_update']).difference(self.special_args) + ["message_and_aggregate", "update"] + ).difference(self.special_args) + self.__edge_user_args__ = self.inspector.keys(["edge_update"]).difference( + self.special_args + ) # Support for "fused" message passing. - self.fuse = self.inspector.implements('message_and_aggregate') + self.fuse = self.inspector.implements("message_and_aggregate") # Support for GNNExplainer. self.__explain__ = False @@ -114,24 +132,27 @@ def __set_size__(self, size: List[Optional[int]], dim: int, src: Tensor): size[dim] = src.shape[self.node_dim] elif the_size != src.shape[self.node_dim]: raise ValueError( - (f'Encountered tensor with size {src.size(self.node_dim)} in ' - f'dimension {self.node_dim}, but expected size {the_size}.')) - + ( + f"Encountered tensor with size {src.size(self.node_dim)} in " + f"dimension {self.node_dim}, but expected size {the_size}." + ) + ) + def __lift__(self, src, edge_index, dim): if isinstance(edge_index, Tensor): index = edge_index[dim] return src.index_select(index, self.node_dim) raise ValueError - + def __collect__(self, args, edge_index, size, kwargs): - i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1) + i, j = (1, 0) if self.flow == "source_to_target" else (0, 1) out = {} for arg in args: - if arg[-2:] not in ['_i', '_j']: + if arg[-2:] not in ["_i", "_j"]: out[arg] = kwargs.get(arg, Parameter.empty) else: - dim = 0 if arg[-2:] == '_j' else 1 + dim = 0 if arg[-2:] == "_j" else 1 data = kwargs.get(arg[:-2], Parameter.empty) if isinstance(data, (tuple, list)): @@ -142,26 +163,25 @@ def __collect__(self, args, edge_index, size, kwargs): if isinstance(data, Tensor): self.__set_size__(size, dim, data) - data = self.__lift__(data, edge_index, - j if arg[-2:] == '_j' else i) + data = self.__lift__(data, edge_index, j if arg[-2:] == "_j" else i) out[arg] = data if isinstance(edge_index, Tensor): - out['adj_t'] = None - out['edge_index'] = edge_index - out['edge_index_i'] = edge_index[i] - out['edge_index_j'] = edge_index[j] - out['ptr'] = None - - out['index'] = out['edge_index_i'] - out['size'] = size - out['size_i'] = size[1] or size[0] - out['size_j'] = size[0] or size[1] - out['dim_size'] = out['size_i'] + out["adj_t"] = None + out["edge_index"] = edge_index + out["edge_index_i"] = edge_index[i] + out["edge_index_j"] = edge_index[j] + out["ptr"] = None + + out["index"] = out["edge_index_i"] + out["size"] = size + out["size_i"] = size[1] or size[0] + out["size_j"] = size[0] or size[1] + out["dim_size"] = out["size_i"] return out - + def __check_input__(self, edge_index, size): the_size: List[Optional[int]] = [None, None] @@ -175,9 +195,12 @@ def __check_input__(self, edge_index, size): return the_size raise ValueError( - ('`MessagePassing.propagate` only supports `int` of ' - 'shape `[2, num_messages]` for ' - 'argument `edge_index`.')) + ( + "`MessagePassing.propagate` only supports `int` of " + "shape `[2, num_messages]` for " + "argument `edge_index`." + ) + ) def propagate(self, edge_index: Tensor, size: Size = None, **kwargs): r"""The initial call to start propagating messages. @@ -216,10 +239,9 @@ def propagate(self, edge_index: Tensor, size: Size = None, **kwargs): if isinstance(edge_index, Tensor) or not self.fuse: if decomposed_layers > 1: user_args = self.__user_args__ - decomp_args = {a[:-2] for a in user_args if a[-2:] == '_j'} + decomp_args = {a[:-2] for a in user_args if a[-2:] == "_j"} decomp_kwargs = { - a: kwargs[a].chunk(decomposed_layers, -1) - for a in decomp_args + a: kwargs[a].chunk(decomposed_layers, -1) for a in decomp_args } decomp_out = [] @@ -228,17 +250,18 @@ def propagate(self, edge_index: Tensor, size: Size = None, **kwargs): for arg in decomp_args: kwargs[arg] = decomp_kwargs[arg][i] - coll_dict = self.__collect__(self.__user_args__, edge_index, - size, kwargs) + coll_dict = self.__collect__( + self.__user_args__, edge_index, size, kwargs + ) - msg_kwargs = self.inspector.distribute('message', coll_dict) + msg_kwargs = self.inspector.distribute("message", coll_dict) for hook in self._message_forward_pre_hooks.values(): - res = hook(self, (msg_kwargs, )) + res = hook(self, (msg_kwargs,)) if res is not None: msg_kwargs = res[0] if isinstance(res, tuple) else res out = self.message(**msg_kwargs) for hook in self._message_forward_hooks.values(): - res = hook(self, (msg_kwargs, ), out) + res = hook(self, (msg_kwargs,), out) if res is not None: out = res @@ -256,19 +279,19 @@ def propagate(self, edge_index: Tensor, size: Size = None, **kwargs): assert out.size(self.node_dim) == edge_mask.size(0) out = out * edge_mask.view([-1] + [1] * (out.dim() - 1)) - aggr_kwargs = self.inspector.distribute('aggregate', coll_dict) + aggr_kwargs = self.inspector.distribute("aggregate", coll_dict) for hook in self._aggregate_forward_pre_hooks.values(): - res = hook(self, (aggr_kwargs, )) + res = hook(self, (aggr_kwargs,)) if res is not None: aggr_kwargs = res[0] if isinstance(res, tuple) else res out = self.aggregate(out, **aggr_kwargs) for hook in self._aggregate_forward_hooks.values(): - res = hook(self, (aggr_kwargs, ), out) + res = hook(self, (aggr_kwargs,), out) if res is not None: out = res - update_kwargs = self.inspector.distribute('update', coll_dict) + update_kwargs = self.inspector.distribute("update", coll_dict) out = self.update(out, **update_kwargs) if decomposed_layers > 1: @@ -283,7 +306,7 @@ def propagate(self, edge_index: Tensor, size: Size = None, **kwargs): out = res return out - + def message(self, x_j: Tensor) -> Tensor: r"""Constructs messages from node :math:`j` to node :math:`i` in analogy to :math:`\phi_{\mathbf{\Theta}}` for each edge in @@ -295,10 +318,14 @@ def message(self, x_j: Tensor) -> Tensor: :obj:`_j` to the variable name, *.e.g.* :obj:`x_i` and :obj:`x_j`. """ return x_j - - def aggregate(self, inputs: Tensor, index: Tensor, - ptr: Optional[Tensor] = None, - dim_size: Optional[int] = None) -> Tensor: + + def aggregate( + self, + inputs: Tensor, + index: Tensor, + ptr: Optional[Tensor] = None, + dim_size: Optional[int] = None, + ) -> Tensor: r"""Aggregates messages from neighbors as :math:`\square_{j \in \mathcal{N}(i)}`. @@ -309,9 +336,10 @@ def aggregate(self, inputs: Tensor, index: Tensor, that support "add", "mean" and "max" operations as specified in :meth:`__init__` by the :obj:`aggr` argument. """ - return scatter(inputs, index, dim=self.node_dim, dim_size=dim_size, - reduce=self.aggr) - + return scatter( + inputs, index, dim=self.node_dim, dim_size=dim_size, reduce=self.aggr + ) + def update(self, inputs: Tensor) -> Tensor: r"""Updates node embeddings in analogy to :math:`\gamma_{\mathbf{\Theta}}` for each node @@ -320,7 +348,7 @@ def update(self, inputs: Tensor) -> Tensor: which was initially passed to :meth:`propagate`. """ return inputs - + def edge_update(self) -> Tensor: r"""Computes or updates features for each edge in the graph. This function can take any argument as input which was initially passed @@ -330,6 +358,7 @@ def edge_update(self) -> Tensor: :obj:`_j` to the variable name, *.e.g.* :obj:`x_i` and :obj:`x_j`. """ raise NotImplementedError + def message_and_aggregate(self, adj_t) -> Tensor: r"""Fuses computations of :func:`message` and :func:`aggregate` into a single function. @@ -338,4 +367,4 @@ def message_and_aggregate(self, adj_t) -> Tensor: This function will only gets called in case it is implemented and propagation takes place based on a :obj:`torch_sparse.SparseTensor`. """ - raise NotImplementedError \ No newline at end of file + raise NotImplementedError diff --git a/jointContribution/graphGalerkin/utils/scatter.py b/jointContribution/graphGalerkin/utils/scatter.py index eebea65be9..4561bbb2bc 100644 --- a/jointContribution/graphGalerkin/utils/scatter.py +++ b/jointContribution/graphGalerkin/utils/scatter.py @@ -1,7 +1,9 @@ -from typing import Optional, Tuple +from typing import Optional +from typing import Tuple import paddle + def broadcast(src: paddle.Tensor, other: paddle.Tensor, dim: int): if dim < 0: dim = other.dim() + dim @@ -13,21 +15,29 @@ def broadcast(src: paddle.Tensor, other: paddle.Tensor, dim: int): src = src.expand_as(other) return src + def scatter_add_(dim, index, src, x): - if x.dim()==1: - output = paddle.scatter_nd_add(x.unsqueeze(-1), index.unsqueeze(-1), src.unsqueeze(-1)).squeeze(-1) + if x.dim() == 1: + output = paddle.scatter_nd_add( + x.unsqueeze(-1), index.unsqueeze(-1), src.unsqueeze(-1) + ).squeeze(-1) else: i, j = index.shape - grid_x , grid_y = paddle.meshgrid(paddle.arange(i), paddle.arange(j)) + grid_x, grid_y = paddle.meshgrid(paddle.arange(i), paddle.arange(j)) index = paddle.stack([index.flatten(), grid_y.flatten()], axis=1) updates_index = paddle.stack([grid_x.flatten(), grid_y.flatten()], axis=1) updates = paddle.gather_nd(src, index=updates_index) output = paddle.scatter_nd_add(x, index, updates) return output -def scatter_sum(src: paddle.Tensor, index: paddle.Tensor, dim: int = -1, - out: Optional[paddle.Tensor] = None, - dim_size: Optional[int] = None) -> paddle.Tensor: + +def scatter_sum( + src: paddle.Tensor, + index: paddle.Tensor, + dim: int = -1, + out: Optional[paddle.Tensor] = None, + dim_size: Optional[int] = None, +) -> paddle.Tensor: index = broadcast(index, src, dim) if out is None: size = list(src.shape) @@ -42,14 +52,24 @@ def scatter_sum(src: paddle.Tensor, index: paddle.Tensor, dim: int = -1, else: return scatter_add_(0, index, src, out) -def scatter_add(src: paddle.Tensor, index: paddle.Tensor, dim: int = -1, - out: Optional[paddle.Tensor] = None, - dim_size: Optional[int] = None) -> paddle.Tensor: + +def scatter_add( + src: paddle.Tensor, + index: paddle.Tensor, + dim: int = -1, + out: Optional[paddle.Tensor] = None, + dim_size: Optional[int] = None, +) -> paddle.Tensor: return scatter_sum(src, index, dim, out, dim_size) -def scatter_mean(src: paddle.Tensor, index: paddle.Tensor, dim: int = -1, - out: Optional[paddle.Tensor] = None, - dim_size: Optional[int] = None) -> paddle.Tensor: + +def scatter_mean( + src: paddle.Tensor, + index: paddle.Tensor, + dim: int = -1, + out: Optional[paddle.Tensor] = None, + dim_size: Optional[int] = None, +) -> paddle.Tensor: out = scatter_sum(src, index, dim, out, dim_size) dim_size = out.size(dim) @@ -70,9 +90,15 @@ def scatter_mean(src: paddle.Tensor, index: paddle.Tensor, dim: int = -1, out.floor_divide_(count) return out -def scatter(src: paddle.Tensor, index: paddle.Tensor, dim: int = -1, - out: Optional[paddle.Tensor] = None, dim_size: Optional[int] = None, - reduce: str = "sum") -> paddle.Tensor: + +def scatter( + src: paddle.Tensor, + index: paddle.Tensor, + dim: int = -1, + out: Optional[paddle.Tensor] = None, + dim_size: Optional[int] = None, + reduce: str = "sum", +) -> paddle.Tensor: r""" Reduces all values from the :attr:`src` tensor into :attr:`out` at the indices specified in the :attr:`index` tensor along a given axis @@ -121,9 +147,9 @@ def scatter(src: paddle.Tensor, index: paddle.Tensor, dim: int = -1, :rtype: :class:`Tensor` """ - if reduce == 'sum' or reduce == 'add': + if reduce == "sum" or reduce == "add": return scatter_sum(src, index, dim, out, dim_size) - elif reduce == 'mean': + elif reduce == "mean": return scatter_mean(src, index, dim, out, dim_size) else: raise ValueError diff --git a/jointContribution/graphGalerkin/utils/utils.py b/jointContribution/graphGalerkin/utils/utils.py index c495408bd6..3d55494035 100644 --- a/jointContribution/graphGalerkin/utils/utils.py +++ b/jointContribution/graphGalerkin/utils/utils.py @@ -1,22 +1,29 @@ -from paddle import Tensor -from typing import Optional, Tuple, Union +from typing import Optional +from typing import Tuple +from typing import Union + import paddle +from paddle import Tensor +from scatter import scatter +from scatter import scatter_add -from scatter import scatter, scatter_add OptTensor = Optional[Tensor] import pgl -class Data(): + +class Data: def __init__(self, x, y, edge_index): self.y = y self.x = x self.edge_index = edge_index + def __call__(self): - return pgl.Graph(edges=self.edge_index, - num_nodes=self.x.shape[0], - node_feat=self.x) - + return pgl.Graph( + edges=self.edge_index, num_nodes=self.x.shape[0], node_feat=self.x + ) + + def maybe_num_nodes(edge_index, num_nodes=None): if num_nodes is not None: return num_nodes @@ -25,8 +32,10 @@ def maybe_num_nodes(edge_index, num_nodes=None): else: return max(edge_index.size(0), edge_index.size(1)) -def remove_self_loops(edge_index: Tensor, - edge_attr: OptTensor = None) -> Tuple[Tensor, OptTensor]: + +def remove_self_loops( + edge_index: Tensor, edge_attr: OptTensor = None +) -> Tuple[Tensor, OptTensor]: r"""Removes every self-loop in the graph given by :attr:`edge_index`, so that :math:`(i,i) \not\in \mathcal{E}` for every :math:`i \in \mathcal{V}`. @@ -39,16 +48,19 @@ def remove_self_loops(edge_index: Tensor, """ mask = edge_index[0] != edge_index[1] for _ in range(edge_index.dim()): - edge_index[_] = paddle.masked_select(edge_index[_], mask) + edge_index[_] = paddle.masked_select(edge_index[_], mask) if edge_attr is None: return edge_index, None else: return edge_index, edge_attr[mask] + def add_self_loops( - edge_index: Tensor, edge_attr: OptTensor = None, - fill_value: Union[float, Tensor, str] = None, - num_nodes: Optional[int] = None) -> Tuple[Tensor, OptTensor]: + edge_index: Tensor, + edge_attr: OptTensor = None, + fill_value: Union[float, Tensor, str] = None, + num_nodes: Optional[int] = None, +) -> Tuple[Tensor, OptTensor]: r"""Adds a self-loop :math:`(i,i) \in \mathcal{E}` to every node :math:`i \in \mathcal{V}` in the graph given by :attr:`edge_index`. In case the graph is weighted or has multi-dimensional edge features @@ -78,10 +90,10 @@ def add_self_loops( loop_index = paddle.tile(loop_index.unsqueeze(0), repeat_times=[2, 1]) if edge_attr is not None: if fill_value is None: - loop_attr = edge_attr.new_full((N, ) + edge_attr.size()[1:], 1.) + loop_attr = edge_attr.new_full((N,) + edge_attr.size()[1:], 1.0) elif isinstance(fill_value, (int, float)): - loop_attr = paddle.full((N, ), fill_value, dtype=edge_attr.dtype) + loop_attr = paddle.full((N,), fill_value, dtype=edge_attr.dtype) elif isinstance(fill_value, Tensor): loop_attr = fill_value.to(edge_attr.device, edge_attr.dtype) if edge_attr.dim() != loop_attr.dim(): @@ -90,21 +102,26 @@ def add_self_loops( loop_attr = loop_attr.repeat(*sizes) elif isinstance(fill_value, str): - loop_attr = scatter(edge_attr, edge_index[1], dim=0, dim_size=N, - reduce=fill_value) + loop_attr = scatter( + edge_attr, edge_index[1], dim=0, dim_size=N, reduce=fill_value + ) else: raise AttributeError("No valid 'fill_value' provided") - + edge_attr = paddle.concat([edge_attr, loop_attr], axis=0) edge_index = paddle.concat([edge_index, loop_index], axis=1) return edge_index, edge_attr - -def get_laplacian(edge_index, edge_weight: Optional[paddle.Tensor] = None, - normalization: Optional[str] = None, - dtype: Optional[int] = None, - num_nodes: Optional[int] = None): - r""" Computes the graph Laplacian of the graph given by :obj:`edge_index` + + +def get_laplacian( + edge_index, + edge_weight: Optional[paddle.Tensor] = None, + normalization: Optional[str] = None, + dtype: Optional[int] = None, + num_nodes: Optional[int] = None, +): + r"""Computes the graph Laplacian of the graph given by :obj:`edge_index` and optional :obj:`edge_weight`. Args: @@ -130,7 +147,7 @@ def get_laplacian(edge_index, edge_weight: Optional[paddle.Tensor] = None, """ if normalization is not None: - assert normalization in ['sym', 'rw'] # 'Invalid normalization' + assert normalization in ["sym", "rw"] # 'Invalid normalization' edge_index, edge_weight = remove_self_loops(edge_index, edge_weight) @@ -144,38 +161,44 @@ def get_laplacian(edge_index, edge_weight: Optional[paddle.Tensor] = None, # L = D - A. edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes) edge_weight = paddle.concat([-edge_weight, deg], dim=0) - elif normalization == 'sym': + elif normalization == "sym": # Compute A_norm = -D^{-1/2} A D^{-1/2}. deg_inv_sqrt = deg.pow(-0.5) - deg_inv_sqrt = masked_fill(deg_inv_sqrt,deg_inv_sqrt == float('inf'), 0) + deg_inv_sqrt = masked_fill(deg_inv_sqrt, deg_inv_sqrt == float("inf"), 0) edge_weight = deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col] # L = I - A_norm. - edge_index, tmp = add_self_loops(edge_index, -edge_weight, - fill_value=1., num_nodes=num_nodes) + edge_index, tmp = add_self_loops( + edge_index, -edge_weight, fill_value=1.0, num_nodes=num_nodes + ) assert tmp is not None edge_weight = tmp else: # Compute A_norm = -D^{-1} A. deg_inv = 1.0 / deg - deg_inv.masked_fill_(deg_inv == float('inf'), 0) + deg_inv.masked_fill_(deg_inv == float("inf"), 0) edge_weight = deg_inv[row] * edge_weight # L = I - A_norm. - edge_index, tmp = add_self_loops(edge_index, -edge_weight, - fill_value=1., num_nodes=num_nodes) + edge_index, tmp = add_self_loops( + edge_index, -edge_weight, fill_value=1.0, num_nodes=num_nodes + ) assert tmp is not None edge_weight = tmp return edge_index, edge_weight + def masked_fill(x, mask, value): y = paddle.full(x.shape, value, x.dtype) return paddle.where(mask, y, x) + def add_remaining_self_loops( - edge_index: Tensor, edge_attr: OptTensor = None, - fill_value: Union[float, Tensor, str] = None, - num_nodes: Optional[int] = None) -> Tuple[Tensor, OptTensor]: + edge_index: Tensor, + edge_attr: OptTensor = None, + fill_value: Union[float, Tensor, str] = None, + num_nodes: Optional[int] = None, +) -> Tuple[Tensor, OptTensor]: r"""Adds remaining self-loop :math:`(i,i) \in \mathcal{E}` to every node :math:`i \in \mathcal{V}` in the graph given by :attr:`edge_index`. In case the graph is weighted or has multi-dimensional edge features @@ -207,10 +230,10 @@ def add_remaining_self_loops( if edge_attr is not None: if fill_value is None: - loop_attr = edge_attr.new_full((N, ) + edge_attr.size()[1:], 1.) + loop_attr = edge_attr.new_full((N,) + edge_attr.size()[1:], 1.0) elif isinstance(fill_value, (int, float)): - loop_attr = paddle.full((N, ), fill_value, dtype=edge_attr.dtype) + loop_attr = paddle.full((N,), fill_value, dtype=edge_attr.dtype) elif isinstance(fill_value, Tensor): loop_attr = fill_value.to(edge_attr.device, edge_attr.dtype) if edge_attr.dim() != loop_attr.dim(): @@ -219,8 +242,9 @@ def add_remaining_self_loops( loop_attr = loop_attr.repeat(*sizes) elif isinstance(fill_value, str): - loop_attr = scatter(edge_attr, edge_index[1], dim=0, dim_size=N, - reduce=fill_value) + loop_attr = scatter( + edge_attr, edge_index[1], dim=0, dim_size=N, reduce=fill_value + ) else: raise AttributeError("No valid 'fill_value' provided") @@ -230,7 +254,8 @@ def add_remaining_self_loops( edge_index = paddle.concat([edge_index, loop_index], axis=1) return edge_index, edge_attr + def expand_left(src: paddle.Tensor, dim: int, dims: int) -> paddle.Tensor: for _ in range(dims + dim if dim < 0 else dim): src = src.unsqueeze(0) - return src \ No newline at end of file + return src diff --git a/mkdocs.yml b/mkdocs.yml index e5dc9bef74..9ff33d59e5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -91,6 +91,7 @@ nav: - 材料科学(AI for Material): - hPINNs: zh/examples/hpinns.md - CGCNN: zh/examples/cgcnn.md + - NN: zh/examples/perovskite_solar_cells_nn.md - 地球科学(AI for Earth Science): - Extformer-MoE: zh/examples/extformer_moe.md - FourCastNet: zh/examples/fourcastnet.md