From 3176ab69fd78ebee7feb4a03335acc8b009901a4 Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Tue, 7 Oct 2025 23:28:01 +0800 Subject: [PATCH 01/10] init --- examples/symbolic_gcn/conf/config.yaml | 84 ++++++++++++++++++++++++++ examples/symbolic_gcn/gcn.py | 0 2 files changed, 84 insertions(+) create mode 100644 examples/symbolic_gcn/conf/config.yaml create mode 100644 examples/symbolic_gcn/gcn.py diff --git a/examples/symbolic_gcn/conf/config.yaml b/examples/symbolic_gcn/conf/config.yaml new file mode 100644 index 0000000000..01d4da56b1 --- /dev/null +++ b/examples/symbolic_gcn/conf/config.yaml @@ -0,0 +1,84 @@ +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: + # dynamic output directory according to running time and override name + dir: outputs_symbolic_gcn/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} + job: + name: ${mode} # name of logfile + chdir: false # keep current working directory unchanged + callbacks: + init_callback: + _target_: ppsci.utils.callbacks.InitCallback + sweep: + # output directory for multirun + dir: ${hydra.run.dir} + subdir: ./ + +# general settings +mode: train # running mode: train/eval +seed: 42 +output_dir: ${hydra:run.dir} +log_freq: 20 +# dataset setting +DATASET: + label_keys: [label] + data_dir: ./dataset/train_data.pkl + +MODEL: + input_keys: [aq_train_data, mete_train_data] + output_keys: [label] + output_attention: true + seq_len: 72 + pred_len: 48 + aq_gat_node_features: 7 + aq_gat_node_num: 35 + mete_gat_node_features: 7 + mete_gat_node_num: 18 + gat_hidden_dim: 32 + gat_edge_dim: 3 + e_layers: 1 + enc_in: 7 + dec_in: 7 + c_out: 7 + d_model: 16 + embed: fixed + freq: t + dropout: 0.05 + factor: 3 + n_heads: 4 + d_ff: 32 + num_kernels: 6 + top_k: 4 + +# training settings +TRAIN: + epochs: 100 + iters_per_epoch: 400 + save_freq: 10 + eval_during_train: true + eval_freq: 10 + batch_size: 32 + learning_rate: 0.0001 + lr_scheduler: + epochs: ${TRAIN.epochs} + iters_per_epoch: ${TRAIN.iters_per_epoch} + learning_rate: 0.0005 + step_size: 20 + gamma: 0.95 + pretrained_model_path: null + checkpoint_path: null + +EVAL: + eval_data_path: ./dataset/val_data.pkl + pretrained_model_path: null + compute_metric_by_batch: false + eval_with_no_grad: true + batch_size: 32 diff --git a/examples/symbolic_gcn/gcn.py b/examples/symbolic_gcn/gcn.py new file mode 100644 index 0000000000..e69de29bb2 From f33c07fbac582e9786278175d031370ff504dfba Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Sun, 19 Oct 2025 22:52:49 +0800 Subject: [PATCH 02/10] Implement symbolic graph networks for physics discovery, including OGN, VarOGN, and HGN architectures. Add simulation dataset generation and configuration files. Remove outdated symbolic GCN files. --- docs/zh/examples/symbolic_gn.md | 833 +++++++++++++++++++++++++ examples/symbolic_gcn/conf/config.yaml | 84 --- examples/symbolic_gcn/gcn.py | 0 examples/symbolic_gn/conf/config.yaml | 138 ++++ examples/symbolic_gn/gn.py | 519 +++++++++++++++ examples/symbolic_gn/simulate.py | 752 ++++++++++++++++++++++ ppsci/arch/__init__.py | 5 + ppsci/arch/symbolic_gn.py | 667 ++++++++++++++++++++ 8 files changed, 2914 insertions(+), 84 deletions(-) create mode 100644 docs/zh/examples/symbolic_gn.md delete mode 100644 examples/symbolic_gcn/conf/config.yaml delete mode 100644 examples/symbolic_gcn/gcn.py create mode 100644 examples/symbolic_gn/conf/config.yaml create mode 100644 examples/symbolic_gn/gn.py create mode 100644 examples/symbolic_gn/simulate.py create mode 100644 ppsci/arch/symbolic_gn.py diff --git a/docs/zh/examples/symbolic_gn.md b/docs/zh/examples/symbolic_gn.md new file mode 100644 index 0000000000..2bb58c7163 --- /dev/null +++ b/docs/zh/examples/symbolic_gn.md @@ -0,0 +1,833 @@ +# Graph Networks for Physics Discovery (GN) + +AI Studio快速体验 + +开始训练、评估前，请先确保已安装依赖 `scipy`、`scikit-learn`、`celluloid`、`pgl`，请运行安装命令： +```bash +pip install scipy scikit-learn celluloid pgl +``` + +=== "模型训练命令" + + ``` sh + python gn.py mode=train MODEL.arch=OGN DATA.type=spring + ``` +=== "模型评估命令" + + ``` sh + python gn.py mode=eval MODEL.arch=OGN EVAL.pretrained_model_path="./outputs_symbolic_gcn/.../checkpoints/latest" + ``` +=== "模型推理命令" + + ``` sh + python gn.py mode=infer MODEL.arch=OGN INFER.pretrained_model_path="./outputs_symbolic_gcn/.../checkpoints/latest" + ``` +=== "模型导出命令" + + ``` sh + python gn.py mode=export MODEL.arch=OGN INFER.pretrained_model_path="./outputs_symbolic_gcn/.../checkpoints/latest" + ``` + +## 1. 背景简介 + +深度学习方法在物理系统建模中展现出强大能力，但传统神经网络往往是黑盒模型，难以提供物理解释。图神经网络（GNNs）通过其固有的归纳偏置，特别适合建模粒子系统中的相互作用。然而，如何从训练好的GNN中提取可解释的物理定律仍然是一个挑战。 + +基于论文 **[Discovering Symbolic Models from Deep Learning with Inductive Biases](https://arxiv.org/abs/2006.11287)** 的思想，我们开发了一个完整的物理发现框架。该框架通过在GNN训练过程中引入强归纳偏置（如L1正则化鼓励稀疏表示），然后对学习到的模型内部组件应用符号回归，从而提取显式的物理关系。这种方法不仅能够重新发现已知的力定律和哈密顿量，还能从复杂数据中发现新的解析公式。 + +本案例使用图网络（Graph Networks, GN）对多种物理系统（如弹簧系统、引力系统、电荷系统等）的动力学进行建模，并通过符号回归提取潜在的物理定律。 + +## 2. 模型原理 + +本章节基于论文 **[Discovering Symbolic Models from Deep Learning with Inductive Biases](https://arxiv.org/abs/2006.11287)** (NeurIPS 2020) 的核心思想进行介绍。 + +### 2.1 核心问题 + +**传统困境**：深度学习模型（如全连接网络、CNN）虽然强大，但缺乏物理解释性。我们能否既利用深度学习的拟合能力，又能提取可解释的物理定律？ + +**本文方案**： +1. 使用具有**归纳偏置**的架构（图网络）学习物理系统 +2. 在训练中引入**稀疏性约束**，迫使模型学习简单表示 +3. 对学到的内部函数进行**符号回归**，提取解析表达式 +4. 用符号表达式替换神经网络组件，获得可解释模型 + +### 2.2 图网络架构 + +图网络（GN）是天然适合物理建模的架构，因为： +- **排列不变性**：粒子顺序不影响结果（符合物理对称性） +- **局部作用**：粒子间通过边交互（符合力的成对性） +- **可扩展性**：可以处理不同数量的粒子 + +GN的核心思想是将物理系统建模为图结构： +- **节点**：代表物理实体（如粒子），包含状态 xᵢ = [位置, 速度, 质量, 电荷] +- **边**：代表实体间的相互作用，通过消息传递机制计算 +- **全局属性**：可选的系统级属性（如总能量） + +GN内部包含三个可分离的函数组件： +- **边函数（消息函数）φₑ(xᵢ, xⱼ)**：计算从节点j到节点i的消息向量 +- **节点函数 φᵥ(xᵢ, Σmᵢⱼ)**：根据聚合的消息更新节点状态 +- **全局函数 φᵤ(Σxᵢ)**：计算系统的全局属性（如总能量） + +**关键**：这些函数组件（φₑ, φᵥ, φᵤ）天然对应物理中的基本概念： +- φₑ 对应**成对作用力**（如引力、库仑力） +- φᵥ 对应**运动方程**（如牛顿第二定律） +- φᵤ 对应**能量函数**（如哈密顿量） + +因此，如果我们能对这些函数进行符号回归，就能直接提取物理定律！ + +### 2.3 物理发现框架 + +我们的物理发现框架包含以下步骤： + +1. **设计具有可分离内部结构的深度学习模型**：使用图网络作为核心归纳偏置，特别适合粒子系统建模 +2. **端到端训练模型**：使用可用数据训练GN模型 +3. **对模型内部函数进行符号回归**：对φₑ、φᵥ、φᵤ等组件拟合符号表达式 +4. **用符号表达式替换神经网络组件**：创建完全解析的物理模型 + +### 2.4 稀疏表示与正则化 + +**核心理论**：如果真实物理系统可以用线性潜在空间完美描述，且我们用足够大的消息维度训练GNN，那么学到的消息向量将是真实力向量的**线性变换**。 + +例如，对于引力系统 `F = -G·m₁m₂/r²·r̂`，消息向量可能学到： +``` +message = [c₁·dx, c₂·dy, c₃·m₁m₂/r², c₄·..., 0, 0, ...] +``` +其中只有少数几个通道是活跃的，其余为零。 + +为了鼓励这种稀疏性，我们在训练中使用正则化策略： + +| 正则化方法 | 适用模型 | 作用机制 | 配置参数 | +|-----------|---------|---------|---------| +| **L1正则化** | OGN | 对节点特征施加L1惩罚，迫使网络学习稀疏激活 | `l1_strength: 1e-2` | +| **KL正则化** | VarOGN | 使后验分布接近稀疏先验，自然产生稀疏性 | `kl_weight: 1e-3` | +| **瓶颈架构** | 所有 | 限制消息维度（msg_dim），迫使信息压缩 | `msg_dim: 100` | + +**为什么稀疏性重要**： +1. **符号回归可行性**：稀疏表示意味着只有少数变量参与，符号回归搜索空间大幅减小 +2. **物理解释性**：活跃的通道对应真实的物理量组合 +3. **泛化能力**：简单模型通常泛化更好（奥卡姆剃刀原则） + +## 3. 实现细节 + +### 3.1 数据集介绍 + +我们的数据生成器支持多种物理系统，包括： + +| 系统类型 | 势能函数 | 物理意义 | +|---------|---------|---------| +| `r2` | `-m₁m₂/r` | 引力/库仑力（吸引）| +| `r1` | `m₁m₂·log(r)` | 2D引力（涡旋）| +| `spring` | `(r-1)²` | 胡克定律（平衡长度=1）| +| `damped` | `(r-1)² + damping·v·x/n` | 带阻尼的弹簧 | +| `string` | `(r-1)² + y·q` | 弦 + 重力 | +| `charge` | `q₁q₂/r` | 库仑力（排斥/吸引）| +| `superposition` | `q₁q₂/r - m₁m₂/r` | 电磁+引力叠加 | +| `discontinuous` | 分段函数 | 模拟非光滑力（如碰撞）| +| `string_ball` | 弹簧 + 球体排斥 | 复杂约束系统 | + +数据生成器使用高精度ODE求解器（`scipy.integrate.odeint`）生成轨迹数据，确保物理准确性。 + +### 3.2 数据生成器详解 + +数据生成器（`simulate.py`）使用 PaddlePaddle 自动微分计算力，核心流程如下： + +1. **势能函数定义**：为每种物理系统定义势能函数 U(r) +2. **自动微分计算力**：`F = -∇U`，使用 `paddle.grad` 计算势能梯度 +3. **ODE求解**：使用 `scipy.integrate.odeint` 求解 `d²x/dt² = F/m` +4. **数据格式**： + - 输入：[x, y, vx, vy, charge, mass]（2D系统） + - 标签：加速度 [ax, ay] 或导数 [dq/dt, dp/dt] + +### 3.3 模型构建 + +我们实现了三种核心模型架构： + +#### OGN (Object-based Graph Network) +- **物理基础**：牛顿力学 (F=ma) +- **输出**：加速度 a = F/m +- **网络结构**： + - 消息函数（Message Function）：`φₑ: (xᵢ, xⱼ) → mᵢⱼ`，将源节点和目标节点特征映射到消息向量（维度100） + - 节点更新函数（Node Function）：`φᵥ: (xᵢ, Σmᵢⱼ) → aᵢ`，聚合所有消息并输出加速度 +- **适用场景**：通用动力学系统（弹簧、引力、电荷等） +- **正则化**：L1正则化鼓励消息向量稀疏性，使符号回归更容易 + +#### VarOGN (Variational ODE Graph Network) +- **物理基础**：变分推断框架 +- **输出**：加速度均值和方差（不确定性量化） +- **网络结构**： + - 消息函数输出 μ 和 logσ² + - 训练时采样：`m = μ + ε·exp(logσ²/2)`，其中 ε ~ N(0,1) + - 推理时使用均值：`m = μ` +- **适用场景**：噪声数据/不确定性建模/鲁棒性评估 +- **正则化**：KL散度正则化使后验分布接近先验分布 + +#### HGN (Hamiltonian Graph Network) +- **物理基础**：哈密顿力学（能量守恒原理） +- **输出**：[dq/dt, dp/dt]（广义坐标和广义动量的时间导数） +- **网络结构**： + - 成对能量函数：`Eᵢⱼ = φₑ(xᵢ, xⱼ)` + - 自身能量函数：`Eᵢ = φᵥ(xᵢ)` + - 哈密顿量：`H = Σᵢ Eᵢ + Σᵢⱼ Eᵢⱼ` + - 哈密顿方程：`dq/dt = ∂H/∂p`, `dp/dt = -∂H/∂q` +- **适用场景**：保守系统（无耗散）、长时间预测 +- **特点**：通过自动微分实现哈密顿方程，天然保证能量守恒（在无耗散系统中） + +模型通过配置文件中的 `MODEL.arch` 参数动态选择： + +```yaml +MODEL: + arch: "OGN" # 可选: "OGN", "VarOGN", "HGN" + input_keys: ["node_features", "edge_index"] + output_keys: ["acceleration"] # HGN: ["derivative"] + n_f: auto # 自动设置为 dimension * 2 + 2 + msg_dim: 100 # 消息维度（OGN/VarOGN） + ndim: auto # 自动设置为 dimension + hidden: 300 # 隐藏层维度 +``` + +**输入输出格式**： + +| 模型 | 输入格式 | 输出格式 | 说明 | +|------|---------|---------|------| +| OGN | `node_features`: [batch, n_nodes, n_f]
`edge_index`: [batch, 2, n_edges] | `acceleration`: [batch, n_nodes, ndim] | 直接预测加速度 | +| VarOGN | 同上 | `acceleration`: [batch, n_nodes, ndim] | 预测加速度（均值） | +| HGN | 同上 | `derivative`: [batch, n_nodes, 2*ndim] | 前半部分为dq/dt，后半部分为dp/dt | + +其中 `n_f = 2*ndim + 2`（位置、速度、电荷、质量）。 + +### 3.4 约束构建 + +我们使用PPSCI的`SupervisedConstraint`构建监督约束。数据流程如下： + +1. **数据集准备**： + ```python + train_dataset = OGNDataset(cfg, mode="train") + train_data, train_labels = train_dataset.get_data_and_labels() + ``` + +2. **约束构建**： + ```python + train_constraint = ppsci.constraint.SupervisedConstraint( + { + "dataset": { + "name": "NamedArrayDataset", + "input": train_data, + "label": train_labels + }, + "batch_size": cfg.TRAIN.batch_size, + "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": True}, + }, + loss=loss_fn, + name="train_constraint", + ) + ``` + +3. **损失函数**：根据模型类型自动调整 + - **OGN/VarOGN**：MAE或MSE损失 + L1正则化（作用于节点特征） + - **HGN**：仅对加速度部分（derivative的后半部分）计算损失 + +**自动化配置**： +- `batch_size = auto`：自动计算为 `int(64 * (4 / num_nodes)²)` +- 训练/验证集划分：默认 0.8/0.2 +- 数据下采样：`downsample_factor = 5`，减少计算量 + +### 3.5 优化器构建 + +训练使用Adam优化器配合OneCycleLR学习率调度： + +```python +# 动态计算每个epoch的批次数 +batch_per_epoch = int(1000*10 / (cfg.TRAIN.batch_size/32.0)) + +# OneCycleLR学习率调度器 +lr_scheduler = paddle.optimizer.lr.OneCycleLR( + max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, + total_steps=cfg.TRAIN.epochs * batch_per_epoch, + divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor +) + +optimizer = paddle.optimizer.Adam( + learning_rate=lr_scheduler, + parameters=model.parameters(), + weight_decay=cfg.TRAIN.optimizer.weight_decay +) +``` + +关键超参数： + +```yaml +TRAIN: + epochs: 1000 + batch_size: auto # 自动计算 + optimizer: + learning_rate: 1e-3 + weight_decay: 1e-8 + lr_scheduler: + name: "OneCycleLR" + max_learning_rate: 1e-3 + final_div_factor: 1e5 # 最终学习率 = max_lr / final_div_factor + loss: + type: "MAE" # 可选 MAE 或 MSE +``` + +**学习率调度策略**： +- OneCycleLR：学习率先从 max_lr/divide_factor 增加到 max_lr，再逐渐降低到 max_lr/final_div_factor +- 适合快速收敛且避免过拟合 + +### 3.6 符号回归准备 + +训练GNN的目的是为符号回归做准备。根据论文思想，我们鼓励稀疏的潜在表示： + +**理论基础**： +- 如果真实物理系统有完美的线性潜在空间描述，训练后的GNN消息向量将是真实力向量的线性变换 +- L1正则化迫使网络学习稀疏表示，使得消息向量中只有少数维度是活跃的 +- 这些活跃的维度通常对应真实的物理量（如距离、质量乘积等） + +**实践流程**： + +1. **训练稀疏GNN**： + ```yaml + MODEL: + msg_dim: 100 # 消息维度 + regularization_type: "l1" + l1_strength: 1e-2 + ``` + +2. **提取消息数据**：训练过程中可以记录消息向量（需要修改代码添加钩子） + ```python + messages = model.msg_fnc(edge_features) # [num_edges, msg_dim] + # 保存消息及对应的物理特征（dx, dy, r, m1, m2等） + ``` + +3. **符号回归**：使用工具如 [PySR](https://github.com/MilesCranmer/PySR) 拟合符号表达式 + ```python + from pysr import PySRRegressor + + # 选择最活跃的消息通道 + active_channels = np.argsort(np.std(messages, axis=0))[-5:] + + # 对每个活跃通道进行符号回归 + for ch in active_channels: + model = PySRRegressor( + niterations=40, + binary_operators=["+", "*", "/", "-"], + unary_operators=["square", "sqrt", "neg"], + ) + model.fit(physical_features, messages[:, ch]) + print(f"Channel {ch}: {model.sympy()}") + ``` + +4. **发现物理定律**： + - 对于弹簧系统：可能发现 `m ~ (r-1)` （胡克定律） + - 对于引力系统：可能发现 `m ~ -m1*m2/r²` （万有引力定律） + - 对于电荷系统：可能发现 `m ~ q1*q2/r²` （库仑定律） + +## 4. 完整代码 + +完整实现包含以下文件： + +```python + + +class OGNDataset: + def __init__(self, cfg, mode="train"): + self.cfg = cfg + self.mode = mode + self.data = None + self.labels = None + self.input_keys = tuple(cfg.MODEL.input_keys) + self.label_keys = tuple(cfg.MODEL.output_keys) + self._prepare_data() + + def _prepare_data(self): + # Get time step size for physics system + dt = 1e-2 + for sim_set in self.cfg.DATATYPE_TYPE: + if sim_set['sim'] == self.cfg.DATA.type: + dt = sim_set['dt'][0] + break + + # Generate simulation data + sim_dataset = SimulationDataset( + sim=self.cfg.DATA.type, + n=self.cfg.DATA.num_nodes, + dim=self.cfg.DATA.dimension, + dt=dt, + nt=self.cfg.DATA.time_steps, + seed=self.cfg.seed + ) + sim_dataset.simulate(self.cfg.DATA.num_samples) + + # Compute target data based on model architecture + if self.cfg.MODEL.arch == "HGN": + raw_target_data = sim_dataset.get_derivative() # [v, a] for HGN + else: + raw_target_data = sim_dataset.get_acceleration() # Acceleration for OGN/VarOGN + + # Downsample data + downsample_factor = self.cfg.DATA.downsample_factor + input_data = np.concatenate([sim_dataset.data[:, i] for i in range(0, sim_dataset.data.shape[1], downsample_factor)]) + target_data = np.concatenate([raw_target_data[:, i] for i in range(0, sim_dataset.data.shape[1], downsample_factor)]) + + # Split train/validation set + from sklearn.model_selection import train_test_split + test_size = 1.0 - getattr(self.cfg.TRAIN, 'train_split', 0.8) + + if self.mode == "train": + input_data, _, target_data, _ = train_test_split( + input_data, target_data, test_size=test_size, shuffle=False + ) + elif self.mode == "eval": + _, input_data, _, target_data = train_test_split( + input_data, target_data, test_size=test_size, shuffle=False + ) + + # Generate edge indices + edge_index = get_edge_index(self.cfg.DATA.num_nodes, self.cfg.DATA.type) + + self.data = { + "node_features": input_data.astype(np.float32), + "edge_index": np.tile(edge_index.T.numpy(), (len(input_data), 1, 1)) + } + self.labels = { + list(self.cfg.MODEL.output_keys)[0]: target_data.astype(np.float32) + } + + def __len__(self): + return len(self.data["node_features"]) + + def __getitem__(self, idx): + return { + "input": { + "node_features": self.data["node_features"][idx], + "edge_index": self.data["edge_index"][idx] + }, + "label": { + list(self.cfg.MODEL.output_keys)[0]: self.labels[list(self.cfg.MODEL.output_keys)[0]][idx] + } + } + + def get_data_and_labels(self): + return self.data, self.labels + + def __call__(self): + return self + + +def create_loss_function(cfg): + def loss_function(input_dict, label_dict, model_output): + # Get loss type from config (MAE or MSE) + loss_type = getattr(cfg.TRAIN.loss, 'type', 'MAE') + + # Get output key + output_key = list(cfg.MODEL.output_keys)[0] if cfg.MODEL.arch != "HGN" else "derivative" + + # Compute base loss + if cfg.MODEL.arch == "HGN": + output_key = "derivative" + pred_accel = model_output[output_key][:, cfg.DATA.dimension:] + target_accel = label_dict[output_key][:, cfg.DATA.dimension:] + if loss_type == "MSE": + base_loss = paddle.mean(paddle.square(pred_accel - target_accel)) + else: # MAE + base_loss = paddle.mean(paddle.abs(pred_accel - target_accel)) + else: + pred = model_output[output_key] + target = label_dict[output_key] + if loss_type == "MSE": + base_loss = paddle.mean(paddle.square(pred - target)) + else: # MAE + base_loss = paddle.mean(paddle.abs(pred - target)) + + # Add regularization if configured + if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"]: + reg_loss = cfg.MODEL.l1_strength * paddle.mean(paddle.abs(input_dict["node_features"])) + return base_loss + reg_loss + + return base_loss + + return loss_function + + +def train(cfg): + # Set random seed for reproducibility + ppsci.utils.misc.set_random_seed(cfg.seed) + + # Parse automatic configuration parameters + if cfg.MODEL.n_f == "auto": + cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 + if cfg.MODEL.ndim == "auto": + cfg.MODEL.ndim = cfg.DATA.dimension + if cfg.TRAIN.batch_size == "auto": + cfg.TRAIN.batch_size = int(64 * (4 / cfg.DATA.num_nodes) ** 2) + + # Update output keys for HGN + if cfg.MODEL.arch == "HGN": + cfg.MODEL.output_keys = ["derivative"] + + # Create datasets + logger.message(f"Creating {cfg.DATA.type} dataset using simulate.py...") + train_dataset = OGNDataset(cfg, mode="train") + eval_dataset = OGNDataset(cfg, mode="eval") + + # Get training and evaluation data + train_data, train_labels = train_dataset.get_data_and_labels() + eval_data, eval_labels = eval_dataset.get_data_and_labels() + + # Create model based on architecture + logger.message(f"Creating model: {cfg.MODEL.arch}") + if cfg.MODEL.arch == "OGN": + model = OGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "VarOGN": + model = VarOGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "HGN": + model = HGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + else: + raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + + # Create loss function and constraints + loss_fn = create_loss_function(cfg) + train_constraint = ppsci.constraint.SupervisedConstraint( + { + "dataset": { + "name": "NamedArrayDataset", + "input": train_data, + "label": train_labels + }, + "batch_size": cfg.TRAIN.batch_size, + "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": True}, + }, + loss=loss_fn, + name="train_constraint", + ) + + eval_constraint = ppsci.constraint.SupervisedConstraint( + { + "dataset": { + "name": "NamedArrayDataset", + "input": eval_data, + "label": eval_labels + }, + "batch_size": cfg.EVAL.batch_size, + "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": False}, + }, + loss=loss_fn, + name="eval_constraint", + ) + + constraint = { + train_constraint.name: train_constraint, + eval_constraint.name: eval_constraint, + } + + # Calculate iterations per epoch + batch_per_epoch = int(1000*10 / (cfg.TRAIN.batch_size/32.0)) + + # Create optimizer and learning rate scheduler + if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": + lr_scheduler = paddle.optimizer.lr.OneCycleLR( + max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, + total_steps=cfg.TRAIN.epochs*batch_per_epoch, + divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor + ) + else: + lr_scheduler = paddle.optimizer.lr.ExponentialDecay( + learning_rate=cfg.TRAIN.optimizer.learning_rate, + gamma=0.9 + ) + + optimizer = paddle.optimizer.Adam( + learning_rate=lr_scheduler, + parameters=model.parameters(), + weight_decay=cfg.TRAIN.optimizer.weight_decay + ) + + # Create PaddleScience Solver and start training + solver = ppsci.solver.Solver( + model, + constraint, + optimizer=optimizer, + cfg=cfg, + ) + solver.train() + + +def evaluate(cfg: DictConfig): + # Parse automatic configuration parameters + if cfg.MODEL.n_f == "auto": + cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 + if cfg.MODEL.ndim == "auto": + cfg.MODEL.ndim = cfg.DATA.dimension + if cfg.MODEL.arch == "HGN": + cfg.MODEL.output_keys = ["derivative"] + + # Create model based on architecture + if cfg.MODEL.arch == "OGN": + model = OGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "VarOGN": + model = VarOGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "HGN": + model = HGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + else: + raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + + # Create evaluation dataset + eval_dataset = OGNDataset(cfg, mode="eval") + eval_data, eval_labels = eval_dataset.get_data_and_labels() + + # Create loss function and constraint + loss_fn = create_loss_function(cfg) + eval_constraint = ppsci.constraint.SupervisedConstraint( + { + "dataset": { + "name": "NamedArrayDataset", + "input": eval_data, + "label": eval_labels + }, + "batch_size": cfg.EVAL.batch_size, + "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": False}, + }, + loss=loss_fn, + name="eval_constraint", + ) + + constraint = { + eval_constraint.name: eval_constraint, + } + + # Create PaddleScience Solver + solver = ppsci.solver.Solver(model, constraint, cfg=cfg) + + # Run evaluation + logger.message("Starting evaluation...") + eval_result = solver.eval() + logger.message(f"Evaluation completed. Result: {eval_result}") + + +def inference(cfg: DictConfig): + # Parse automatic configuration parameters + if cfg.MODEL.n_f == "auto": + cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 + if cfg.MODEL.ndim == "auto": + cfg.MODEL.ndim = cfg.DATA.dimension + if cfg.MODEL.arch == "HGN": + cfg.MODEL.output_keys = ["derivative"] + + # Create model based on architecture + if cfg.MODEL.arch == "OGN": + model = OGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "VarOGN": + model = VarOGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "HGN": + model = HGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + else: + raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + + # Create PaddleScience Solver + solver = ppsci.solver.Solver(model, cfg=cfg) + + # Generate test data + logger.message("Generating inference data using simulate.py...") + dt = 1e-2 + for sim_set in cfg.DATATYPE_TYPE: + if sim_set['sim'] == cfg.DATA.type: + dt = sim_set['dt'][0] + break + + sim_dataset = SimulationDataset( + sim=cfg.DATA.type, + n=cfg.DATA.num_nodes, + dim=cfg.DATA.dimension, + dt=dt, + nt=100, + seed=cfg.seed + ) + sim_dataset.simulate(1) + + # Prepare input data (first time step) + node_features = sim_dataset.data[0, 0] + edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) + + input_dict = { + "node_features": paddle.to_tensor(node_features, dtype='float32').unsqueeze(0), + "edge_index": edge_index.unsqueeze(0) + } + + # Run inference + logger.message("Running inference...") + output_dict = solver.predict(input_dict, return_numpy=True) + + # Display results + if cfg.MODEL.arch == "HGN": + dq_dt = output_dict["derivative"][0, :cfg.DATA.dimension] + dv_dt = output_dict["derivative"][0, cfg.DATA.dimension:] + logger.message(f"HGN inference - dq/dt: {dq_dt}, dv/dt (acceleration): {dv_dt}") + else: + acceleration = output_dict[list(cfg.MODEL.output_keys)[0]][0] + logger.message(f"{cfg.MODEL.arch} inference - acceleration: {acceleration}") + + +def export(cfg: DictConfig): + # Parse automatic configuration parameters + if cfg.MODEL.n_f == "auto": + cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 + if cfg.MODEL.ndim == "auto": + cfg.MODEL.ndim = cfg.DATA.dimension + if cfg.MODEL.arch == "HGN": + cfg.MODEL.output_keys = ["derivative"] + + # Create model based on architecture + if cfg.MODEL.arch == "OGN": + model = OGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "VarOGN": + model = VarOGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "HGN": + model = HGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + else: + raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + + # Create PaddleScience Solver + solver = ppsci.solver.Solver(model, cfg=cfg) + + # Define input specifications for export + from paddle.static import InputSpec + input_spec = [ + { + "node_features": InputSpec([None, cfg.DATA.num_nodes, cfg.MODEL.n_f], "float32", "node_features"), + "edge_index": InputSpec([None, 2, cfg.DATA.num_nodes * (cfg.DATA.num_nodes - 1)], "int64", "edge_index") + } + ] + + # Export model to static graph + logger.message(f"Exporting model to {cfg.INFER.export_path}") + solver.export(input_spec, cfg.INFER.export_path, with_onnx=False) + logger.message("Model export completed!") + + +@hydra.main(version_base=None, config_path="./conf", config_name="config.yaml") +def main(cfg: DictConfig): + if cfg.mode == "train": + train(cfg) + elif cfg.mode == "eval": + evaluate(cfg) + elif cfg.mode == "infer": + inference(cfg) + elif cfg.mode == "export": + export(cfg) + else: + raise ValueError(f"Unsupported mode: {cfg.mode}") + + +if __name__ == "__main__": + main() +``` + + +## 5. 预期结果 + + +## 6. 参考资料 + +- **论文**：[Discovering Symbolic Models from Deep Learning with Inductive Biases (NeurIPS 2020)](https://arxiv.org/abs/2006.11287) +- **作者**：Miles Cranmer, Alvaro Sanchez-Gonzalez, Peter Battaglia, Rui Xu, Kyle Cranmer, David Spergel, Shirley Ho +- **代码仓库**：[github.com/MilesCranmer/symbolic_deep_learning](https://github.com/MilesCranmer/symbolic_deep_learning) +- **符号回归工具**：[PySR - Python Symbolic Regression](https://github.com/MilesCranmer/PySR) diff --git a/examples/symbolic_gcn/conf/config.yaml b/examples/symbolic_gcn/conf/config.yaml deleted file mode 100644 index 01d4da56b1..0000000000 --- a/examples/symbolic_gcn/conf/config.yaml +++ /dev/null @@ -1,84 +0,0 @@ -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: - # dynamic output directory according to running time and override name - dir: outputs_symbolic_gcn/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} - job: - name: ${mode} # name of logfile - chdir: false # keep current working directory unchanged - callbacks: - init_callback: - _target_: ppsci.utils.callbacks.InitCallback - sweep: - # output directory for multirun - dir: ${hydra.run.dir} - subdir: ./ - -# general settings -mode: train # running mode: train/eval -seed: 42 -output_dir: ${hydra:run.dir} -log_freq: 20 -# dataset setting -DATASET: - label_keys: [label] - data_dir: ./dataset/train_data.pkl - -MODEL: - input_keys: [aq_train_data, mete_train_data] - output_keys: [label] - output_attention: true - seq_len: 72 - pred_len: 48 - aq_gat_node_features: 7 - aq_gat_node_num: 35 - mete_gat_node_features: 7 - mete_gat_node_num: 18 - gat_hidden_dim: 32 - gat_edge_dim: 3 - e_layers: 1 - enc_in: 7 - dec_in: 7 - c_out: 7 - d_model: 16 - embed: fixed - freq: t - dropout: 0.05 - factor: 3 - n_heads: 4 - d_ff: 32 - num_kernels: 6 - top_k: 4 - -# training settings -TRAIN: - epochs: 100 - iters_per_epoch: 400 - save_freq: 10 - eval_during_train: true - eval_freq: 10 - batch_size: 32 - learning_rate: 0.0001 - lr_scheduler: - epochs: ${TRAIN.epochs} - iters_per_epoch: ${TRAIN.iters_per_epoch} - learning_rate: 0.0005 - step_size: 20 - gamma: 0.95 - pretrained_model_path: null - checkpoint_path: null - -EVAL: - eval_data_path: ./dataset/val_data.pkl - pretrained_model_path: null - compute_metric_by_batch: false - eval_with_no_grad: true - batch_size: 32 diff --git a/examples/symbolic_gcn/gcn.py b/examples/symbolic_gcn/gcn.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/examples/symbolic_gn/conf/config.yaml b/examples/symbolic_gn/conf/config.yaml new file mode 100644 index 0000000000..2886423a6b --- /dev/null +++ b/examples/symbolic_gn/conf/config.yaml @@ -0,0 +1,138 @@ +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: + # dynamic output directory according to running time and override name + dir: outputs_symbolic_gcn/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} + job: + name: ${mode} # name of logfile + chdir: false # keep current working directory unchanged + callbacks: + init_callback: + _target_: ppsci.utils.callbacks.InitCallback + sweep: + # output directory for multirun + dir: ${hydra.run.dir} + subdir: ./ + +# general settings +mode: train # running mode: train/eval +seed: 42 +output_dir: ${hydra:run.dir} +log_freq: 20 + + +TRAIN: + epochs: 1000 + batch_size: auto # auto = int(64 * (4 / num_nodes) ** 2) + save_freq: 10 + eval_during_train: true + eval_freq: 5 + checkpoint_path: null + pretrained_model_path: null + train_split: 0.8 # Train/validation split ratio + optimizer: + learning_rate: 1e-3 + weight_decay: 1e-8 + lr_scheduler: + name: "OneCycleLR" + max_learning_rate: 1e-3 + final_div_factor: 1e5 + loss: + type: "MAE" # MAE or MSE + physics_weight: 0.0 # Physical regularization weight + message_analysis: + enabled: true + record_interval: 1 + test_sample_size: 1000 + + +EVAL: + pretrained_model_path: null + eval_with_no_grad: true + batch_size: 1024 + +INFER: + pretrained_model_path: null + export_path: ./inference/ogn_model + device: gpu + engine: native + precision: fp32 + + +MODEL: + arch: "OGN" # OGN, varOGN, HGN + input_keys: ["node_features", "edge_index"] + output_keys: ["acceleration"] + n_f: auto # auto = dimension * 2 + 2 + msg_dim: 100 + ndim: auto # auto = dimension + dt: 0.1 + hidden: 300 + aggr: "add" # add, mean, max, min + regularization_type: "l1" # l1, kl, energy, none + l1_strength: 1e-2 + +DATA: + type: "spring" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball + num_samples: 10000 + num_nodes: 4 + dimension: 2 + time_steps: 1000 + time_step_size: auto + downsample_factor: 5 + sample_interval: 5 # Sampling interval for training data + +DATATYPE_TYPE: + - sim: "r1" + dt: [5e-3] + nt: [1000] + n: [4, 8] + dim: [2, 3] + - sim: "r2" + dt: [1e-3] + nt: [1000] + n: [4, 8] + dim: [2, 3] + - sim: "spring" + dt: [1e-2] + nt: [1000] + n: [4, 8] + dim: [2, 3] + - sim: "string" + dt: [1e-2] + nt: [1000] + n: [30] + dim: [2] + - sim: "charge" + dt: [1e-3] + nt: [1000] + n: [4, 8] + dim: [2, 3] + - sim: "superposition" + dt: [1e-3] + nt: [1000] + n: [4, 8] + dim: [2, 3] + - sim: "damped" + dt: [2e-2] + nt: [1000] + n: [4, 8] + dim: [2, 3] + - sim: "discontinuous" + dt: [1e-2] + nt: [1000] + n: [4, 8] + dim: [2, 3] + - sim: "string_ball" + dt: [1e-2] + nt: [1000] + n: [30] + dim: [2] \ No newline at end of file diff --git a/examples/symbolic_gn/gn.py b/examples/symbolic_gn/gn.py new file mode 100644 index 0000000000..7127c21164 --- /dev/null +++ b/examples/symbolic_gn/gn.py @@ -0,0 +1,519 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import numpy as np +import paddle +from omegaconf import DictConfig +import hydra +import ppsci +from ppsci.utils import logger + +from simulate import SimulationDataset +from ppsci.arch.symbolic_gn import OGN, VarOGN, HGN, get_edge_index + + + +class OGNDataset: + def __init__(self, cfg, mode="train"): + self.cfg = cfg + self.mode = mode + self.data = None + self.labels = None + self.input_keys = tuple(cfg.MODEL.input_keys) + self.label_keys = tuple(cfg.MODEL.output_keys) + self._prepare_data() + + def _prepare_data(self): + # Get time step size for physics system + dt = 1e-2 + for sim_set in self.cfg.DATATYPE_TYPE: + if sim_set['sim'] == self.cfg.DATA.type: + dt = sim_set['dt'][0] + break + + # Generate simulation data + sim_dataset = SimulationDataset( + sim=self.cfg.DATA.type, + n=self.cfg.DATA.num_nodes, + dim=self.cfg.DATA.dimension, + dt=dt, + nt=self.cfg.DATA.time_steps, + seed=self.cfg.seed + ) + sim_dataset.simulate(self.cfg.DATA.num_samples) + + # Compute target data based on model architecture + if self.cfg.MODEL.arch == "HGN": + raw_target_data = sim_dataset.get_derivative() # [v, a] for HGN + else: + raw_target_data = sim_dataset.get_acceleration() # Acceleration for OGN/VarOGN + + # Downsample data + downsample_factor = self.cfg.DATA.downsample_factor + input_data = np.concatenate([sim_dataset.data[:, i] for i in range(0, sim_dataset.data.shape[1], downsample_factor)]) + target_data = np.concatenate([raw_target_data[:, i] for i in range(0, sim_dataset.data.shape[1], downsample_factor)]) + + # Split train/validation set + from sklearn.model_selection import train_test_split + test_size = 1.0 - getattr(self.cfg.TRAIN, 'train_split', 0.8) + + if self.mode == "train": + input_data, _, target_data, _ = train_test_split( + input_data, target_data, test_size=test_size, shuffle=False + ) + elif self.mode == "eval": + _, input_data, _, target_data = train_test_split( + input_data, target_data, test_size=test_size, shuffle=False + ) + + # Generate edge indices + edge_index = get_edge_index(self.cfg.DATA.num_nodes, self.cfg.DATA.type) + + self.data = { + "node_features": input_data.astype(np.float32), + "edge_index": np.tile(edge_index.T.numpy(), (len(input_data), 1, 1)) + } + self.labels = { + list(self.cfg.MODEL.output_keys)[0]: target_data.astype(np.float32) + } + + def __len__(self): + return len(self.data["node_features"]) + + def __getitem__(self, idx): + return { + "input": { + "node_features": self.data["node_features"][idx], + "edge_index": self.data["edge_index"][idx] + }, + "label": { + list(self.cfg.MODEL.output_keys)[0]: self.labels[list(self.cfg.MODEL.output_keys)[0]][idx] + } + } + + def get_data_and_labels(self): + return self.data, self.labels + + def __call__(self): + return self + + +def create_loss_function(cfg): + def loss_function(input_dict, label_dict, model_output): + # Get loss type from config (MAE or MSE) + loss_type = getattr(cfg.TRAIN.loss, 'type', 'MAE') + + # Get output key + output_key = list(cfg.MODEL.output_keys)[0] if cfg.MODEL.arch != "HGN" else "derivative" + + # Compute base loss + if cfg.MODEL.arch == "HGN": + output_key = "derivative" + pred_accel = model_output[output_key][:, cfg.DATA.dimension:] + target_accel = label_dict[output_key][:, cfg.DATA.dimension:] + if loss_type == "MSE": + base_loss = paddle.mean(paddle.square(pred_accel - target_accel)) + else: # MAE + base_loss = paddle.mean(paddle.abs(pred_accel - target_accel)) + else: + pred = model_output[output_key] + target = label_dict[output_key] + if loss_type == "MSE": + base_loss = paddle.mean(paddle.square(pred - target)) + else: # MAE + base_loss = paddle.mean(paddle.abs(pred - target)) + + # Add regularization if configured + if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"]: + reg_loss = cfg.MODEL.l1_strength * paddle.mean(paddle.abs(input_dict["node_features"])) + return base_loss + reg_loss + + return base_loss + + return loss_function + + +def train(cfg): + # Set random seed for reproducibility + ppsci.utils.misc.set_random_seed(cfg.seed) + + # Parse automatic configuration parameters + if cfg.MODEL.n_f == "auto": + cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 + if cfg.MODEL.ndim == "auto": + cfg.MODEL.ndim = cfg.DATA.dimension + if cfg.TRAIN.batch_size == "auto": + cfg.TRAIN.batch_size = int(64 * (4 / cfg.DATA.num_nodes) ** 2) + + # Update output keys for HGN + if cfg.MODEL.arch == "HGN": + cfg.MODEL.output_keys = ["derivative"] + + # Create datasets + logger.message(f"Creating {cfg.DATA.type} dataset using simulate.py...") + train_dataset = OGNDataset(cfg, mode="train") + eval_dataset = OGNDataset(cfg, mode="eval") + + # Get training and evaluation data + train_data, train_labels = train_dataset.get_data_and_labels() + eval_data, eval_labels = eval_dataset.get_data_and_labels() + + # Create model based on architecture + logger.message(f"Creating model: {cfg.MODEL.arch}") + if cfg.MODEL.arch == "OGN": + model = OGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "VarOGN": + model = VarOGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "HGN": + model = HGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + else: + raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + + # Create loss function and constraints + loss_fn = create_loss_function(cfg) + train_constraint = ppsci.constraint.SupervisedConstraint( + { + "dataset": { + "name": "NamedArrayDataset", + "input": train_data, + "label": train_labels + }, + "batch_size": cfg.TRAIN.batch_size, + "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": True}, + }, + loss=loss_fn, + name="train_constraint", + ) + + eval_constraint = ppsci.constraint.SupervisedConstraint( + { + "dataset": { + "name": "NamedArrayDataset", + "input": eval_data, + "label": eval_labels + }, + "batch_size": cfg.EVAL.batch_size, + "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": False}, + }, + loss=loss_fn, + name="eval_constraint", + ) + + constraint = { + train_constraint.name: train_constraint, + eval_constraint.name: eval_constraint, + } + + # Calculate iterations per epoch + batch_per_epoch = int(1000*10 / (cfg.TRAIN.batch_size/32.0)) + + # Create optimizer and learning rate scheduler + if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": + lr_scheduler = paddle.optimizer.lr.OneCycleLR( + max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, + total_steps=cfg.TRAIN.epochs*batch_per_epoch, + divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor + ) + else: + lr_scheduler = paddle.optimizer.lr.ExponentialDecay( + learning_rate=cfg.TRAIN.optimizer.learning_rate, + gamma=0.9 + ) + + optimizer = paddle.optimizer.Adam( + learning_rate=lr_scheduler, + parameters=model.parameters(), + weight_decay=cfg.TRAIN.optimizer.weight_decay + ) + + # Create PaddleScience Solver and start training + solver = ppsci.solver.Solver( + model, + constraint, + optimizer=optimizer, + cfg=cfg, + ) + solver.train() + + +def evaluate(cfg: DictConfig): + # Parse automatic configuration parameters + if cfg.MODEL.n_f == "auto": + cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 + if cfg.MODEL.ndim == "auto": + cfg.MODEL.ndim = cfg.DATA.dimension + if cfg.MODEL.arch == "HGN": + cfg.MODEL.output_keys = ["derivative"] + + # Create model based on architecture + if cfg.MODEL.arch == "OGN": + model = OGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "VarOGN": + model = VarOGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "HGN": + model = HGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + else: + raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + + # Create evaluation dataset + eval_dataset = OGNDataset(cfg, mode="eval") + eval_data, eval_labels = eval_dataset.get_data_and_labels() + + # Create loss function and constraint + loss_fn = create_loss_function(cfg) + eval_constraint = ppsci.constraint.SupervisedConstraint( + { + "dataset": { + "name": "NamedArrayDataset", + "input": eval_data, + "label": eval_labels + }, + "batch_size": cfg.EVAL.batch_size, + "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": False}, + }, + loss=loss_fn, + name="eval_constraint", + ) + + constraint = { + eval_constraint.name: eval_constraint, + } + + # Create PaddleScience Solver + solver = ppsci.solver.Solver(model, constraint, cfg=cfg) + + # Run evaluation + logger.message("Starting evaluation...") + eval_result = solver.eval() + logger.message(f"Evaluation completed. Result: {eval_result}") + + +def inference(cfg: DictConfig): + # Parse automatic configuration parameters + if cfg.MODEL.n_f == "auto": + cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 + if cfg.MODEL.ndim == "auto": + cfg.MODEL.ndim = cfg.DATA.dimension + if cfg.MODEL.arch == "HGN": + cfg.MODEL.output_keys = ["derivative"] + + # Create model based on architecture + if cfg.MODEL.arch == "OGN": + model = OGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "VarOGN": + model = VarOGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "HGN": + model = HGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + else: + raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + + # Create PaddleScience Solver + solver = ppsci.solver.Solver(model, cfg=cfg) + + # Generate test data + logger.message("Generating inference data using simulate.py...") + dt = 1e-2 + for sim_set in cfg.DATATYPE_TYPE: + if sim_set['sim'] == cfg.DATA.type: + dt = sim_set['dt'][0] + break + + sim_dataset = SimulationDataset( + sim=cfg.DATA.type, + n=cfg.DATA.num_nodes, + dim=cfg.DATA.dimension, + dt=dt, + nt=100, + seed=cfg.seed + ) + sim_dataset.simulate(1) + + # Prepare input data (first time step) + node_features = sim_dataset.data[0, 0] + edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) + + input_dict = { + "node_features": paddle.to_tensor(node_features, dtype='float32').unsqueeze(0), + "edge_index": edge_index.unsqueeze(0) + } + + # Run inference + logger.message("Running inference...") + output_dict = solver.predict(input_dict, return_numpy=True) + + # Display results + if cfg.MODEL.arch == "HGN": + dq_dt = output_dict["derivative"][0, :cfg.DATA.dimension] + dv_dt = output_dict["derivative"][0, cfg.DATA.dimension:] + logger.message(f"HGN inference - dq/dt: {dq_dt}, dv/dt (acceleration): {dv_dt}") + else: + acceleration = output_dict[list(cfg.MODEL.output_keys)[0]][0] + logger.message(f"{cfg.MODEL.arch} inference - acceleration: {acceleration}") + + +def export(cfg: DictConfig): + # Parse automatic configuration parameters + if cfg.MODEL.n_f == "auto": + cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 + if cfg.MODEL.ndim == "auto": + cfg.MODEL.ndim = cfg.DATA.dimension + if cfg.MODEL.arch == "HGN": + cfg.MODEL.output_keys = ["derivative"] + + # Create model based on architecture + if cfg.MODEL.arch == "OGN": + model = OGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "VarOGN": + model = VarOGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + dt=cfg.MODEL.dt, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + elif cfg.MODEL.arch == "HGN": + model = HGN( + input_keys=tuple(cfg.MODEL.input_keys), + output_keys=tuple(cfg.MODEL.output_keys), + n_f=cfg.MODEL.n_f, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + aggr=cfg.MODEL.aggr + ) + else: + raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + + # Create PaddleScience Solver + solver = ppsci.solver.Solver(model, cfg=cfg) + + # Define input specifications for export + from paddle.static import InputSpec + input_spec = [ + { + "node_features": InputSpec([None, cfg.DATA.num_nodes, cfg.MODEL.n_f], "float32", "node_features"), + "edge_index": InputSpec([None, 2, cfg.DATA.num_nodes * (cfg.DATA.num_nodes - 1)], "int64", "edge_index") + } + ] + + # Export model to static graph + logger.message(f"Exporting model to {cfg.INFER.export_path}") + solver.export(input_spec, cfg.INFER.export_path, with_onnx=False) + logger.message("Model export completed!") + + +@hydra.main(version_base=None, config_path="./conf", config_name="config.yaml") +def main(cfg: DictConfig): + if cfg.mode == "train": + train(cfg) + elif cfg.mode == "eval": + evaluate(cfg) + elif cfg.mode == "infer": + inference(cfg) + elif cfg.mode == "export": + export(cfg) + else: + raise ValueError(f"Unsupported mode: {cfg.mode}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/symbolic_gn/simulate.py b/examples/symbolic_gn/simulate.py new file mode 100644 index 0000000000..a233733279 --- /dev/null +++ b/examples/symbolic_gn/simulate.py @@ -0,0 +1,752 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np # Standard Numpy for numerical computation +import paddle # PaddlePaddle for automatic differentiation +from matplotlib import pyplot as plt +from scipy.integrate import odeint # Scipy ODE solver +import matplotlib as mpl +from functools import partial +from tqdm import tqdm +from celluloid import Camera # For creating animations + +def make_transparent_color(ntimes, fraction): + """ + Create transparent colors for trajectory visualization. + + This function generates a gradient from white to a specified color, + with transparency increasing from 0 to 1. Used for drawing particle + trajectories where old positions are transparent and new positions are opaque. + + Args: + ntimes (int): Number of time steps, i.e., number of colors to generate. + fraction (float): Color fraction in range [0,1] for selecting color from spectrum. + + Returns: + numpy.ndarray: RGBA color array with shape (ntimes, 4). + Each row represents the color [R, G, B, Alpha] at a time step. + """ + # Initialize RGBA array, all set to 1 (white, fully opaque) + rgba = np.ones((ntimes, 4)) + + # Create transparency gradient from 0 (fully transparent) to 1 (fully opaque) + alpha = np.linspace(0, 1, ntimes)[:, np.newaxis] + + # Select base color from spectrum according to fraction + color = np.array(mpl.colors.to_rgba(mpl.cm.gist_ncar(fraction)))[np.newaxis, :] + + # Linear interpolation: from white (1) to target color + rgba[:, :] = 1*(1-alpha) + color*alpha + + # Set alpha channel + rgba[:, 3] = alpha[:, 0] + + return rgba + + +def get_potential(sim, sim_obj): + """ + Get potential energy function based on simulation type. + + This is a factory function that returns the corresponding potential energy + calculation function based on different physical scenarios. + Supported scenarios include: gravity, springs, charges, damped systems, etc. + + Args: + sim (str): Simulation type, options: + 'r2' - Inverse square law (e.g., universal gravitation) + 'r1' - Logarithmic potential + 'spring' - Spring potential + 'damped' - Damped spring + 'string' - String potential + 'string_ball' - String with ball obstacle + 'charge' - Charge interaction + 'superposition' - Superposition of gravity and charge + 'discontinuous' - Piecewise discontinuous potential + sim_obj (SimulationDataset): Simulation dataset object for obtaining parameters like dimension. + + Returns: + function: Potential energy function that takes two particle states and returns their potential energy. + """ + dim = sim_obj._dim # Spatial dimension (2D or 3D) + + def potential(x1, x2): + """ + Calculate potential energy between two particles. + + Particle state format: + 2D: [x, y, vx, vy, charge, mass] + 3D: [x, y, z, vx, vy, vz, charge, mass] + + Args: + x1 (numpy.ndarray): State vector of the first particle. + x2 (numpy.ndarray): State vector of the second particle. + + Returns: + float: Potential energy value between the two particles. + """ + # Calculate Euclidean distance between two particles + dist = np.sqrt(np.sum(np.square(x1[:dim] - x2[:dim]))) + + # Prevent numerical singularity (division by zero), add minimum distance limit + min_dist = 1e-2 + bounded_dist = dist + min_dist + + # ========== Potential energy functions for different physical scenarios ========== + + if sim == 'r2': + # Inverse square law potential: U = -G*m1*m2/r + # Physical scenarios: universal gravitation, Coulomb attraction + return -x1[-1]*x2[-1]/bounded_dist + + elif sim == 'r1': + # Logarithmic potential: U = m1*m2*ln(r) + return x1[-1]*x2[-1]*np.log(bounded_dist) + + elif sim in ['spring', 'damped']: + # Spring potential: U = (r - r0)^2, where r0=1 is natural length + potential_val = (bounded_dist - 1)**2 + + if sim == 'damped': + # Add damping dissipation term (velocity-dependent) + damping = 1 + potential_val += damping*x1[1]*x1[1+sim_obj._dim]/sim_obj._n # y direction + potential_val += damping*x1[0]*x1[0+sim_obj._dim]/sim_obj._n # x direction + if sim_obj._dim == 3: + potential_val += damping*x1[2]*x1[2+sim_obj._dim]/sim_obj._n # z direction + + return potential_val + + elif sim == 'string': + # String potential: spring potential + gravitational potential + return (bounded_dist - 1)**2 + x1[1]*x1[-1] + + elif sim == 'string_ball': + # String potential with ball obstacle + potential_val = (bounded_dist - 1)**2 + x1[1]*x1[-1] + + # Calculate distance from particle to ball center + r = np.sqrt((x1[1] + 15)**2 + (x1[0] - 5)**2) + radius = 4.0 # Ball radius + + # Soft boundary repulsion potential (using more stable form) + k_repel = 100.0 + potential_val += k_repel / np.maximum(r - radius + 0.5, 0.01) + return potential_val + + elif sim in ['charge', 'superposition']: + # Charge interaction: U = q1*q2/r (Coulomb's law) + charge1 = x1[-2] + charge2 = x2[-2] + potential_val = charge1*charge2/bounded_dist + + if sim in ['superposition']: + # Superposition: contains both gravity and charge interaction + m1 = x1[-1] + m2 = x2[-1] + potential_val += -m1*m2/bounded_dist + + return potential_val + + elif sim in ['discontinuous']: + # Piecewise discontinuous potential function + m1 = x1[-1] + m2 = x2[-1] + + # Define potential in three regions + pot_a = 0.0 # r < 1: no interaction + pot_b = 0.0 # 1 <= r < 2: no interaction + pot_c = (bounded_dist - 1)**2 # r >= 2: spring potential + + # Implement piecewise function using conditional expressions + potential_val = ( + pot_a * (bounded_dist < 1) + + (bounded_dist >= 1) * ( + pot_b * (bounded_dist < 2) + + pot_c * (bounded_dist >= 2)) + ) + return potential_val + + else: + raise NotImplementedError('No such simulation ' + str(sim)) + + return potential + + +class SimulationDataset(object): + """ + Physical system simulation dataset class (PaddlePaddle version). + + This class is used to generate and manage simulation data for many-body physical systems. + It simulates the time evolution of particle systems by numerically solving Newton's + equations of motion. + + Main features: + - Generate simulation data for multiple samples + - Compute system acceleration (using PaddlePaddle automatic differentiation) + - Visualize particle trajectories + - Support custom potential energy functions + + Attributes: + _sim (str): Simulation type. + _n (int): Number of particles. + _dim (int): Spatial dimension. + dt (float): Time step size. + nt (int): Number of time steps. + data (ndarray): Simulation data with shape (num_samples, num_timesteps, num_particles, state_dim). + times (ndarray): Time series. + G (float): Gravitational constant. + pairwise (function): Pairwise potential energy function. + """ + + def __init__(self, sim='r2', n=5, dim=2, + dt=0.01, nt=100, extra_potential=None, + **kwargs): + """ + Initialize simulation dataset. + + Args: + sim (str, optional): Simulation type. Defaults to 'r2' (inverse square law). + n (int, optional): Number of particles. Defaults to 5. + dim (int, optional): Spatial dimension (2 or 3). Defaults to 2. + dt (float, optional): Time step size. Defaults to 0.01. + nt (int, optional): Number of time steps to return. Defaults to 100. + extra_potential (function, optional): Additional potential energy function. Defaults to None. + **kwargs: Other parameters reserved for extension. + """ + self._sim = sim + self._n = n + self._dim = dim + self._kwargs = kwargs + self.dt = dt + self.nt = nt + self.data = None + + # Generate uniform time series + self.times = np.linspace(0, self.dt*self.nt, num=self.nt) + + self.G = 1 # Gravitational constant + self.extra_potential = extra_potential + + # Get pairwise potential function for the simulation type + self.pairwise = get_potential(sim=sim, sim_obj=self) + + def simulate(self, ns, seed=0): + """ + Generate simulation data. + + This method is the core function of the simulator. For each sample: + 1. Randomly initialize particle positions, velocities, masses, and charges + 2. Solve Newton's equations of motion using Scipy ODE solver + 3. Record the entire time evolution process + + Physical equations: + d²x/dt² = F/m = -∇U/m + where U is potential energy and F is force + + Args: + ns (int): Number of samples to generate. + seed (int, optional): Seed for numpy random number generator. Defaults to 0. + + Returns: + None: Results are stored in self.data. + """ + # Set random seed + np.random.seed(seed) + + n = self._n + dim = self._dim + sim = self._sim + params = 2 # charge + mass + total_dim = dim*2 + params + times = self.times + G = self.G + + def compute_pairwise_potential(x1, x2_list): + """Vectorized calculation of potential energy between one particle and multiple particles""" + potentials = [] + for x2 in x2_list: + potentials.append(self.pairwise(x1, x2)) + return np.array(potentials) + + def total_potential(xt): + """ + Calculate total potential energy of the system (using Numpy). + + Args: + xt (numpy.ndarray): System state with shape (n, total_dim). + + Returns: + float: Total potential energy of the system. + """ + sum_potential = 0.0 + + for i in range(n - 1): + if sim in ['string', 'string_ball']: + # String type: only adjacent particles interact + sum_potential += G * self.pairwise(xt[i], xt[i+1]) + else: + # Other types: fully connected + for j in range(i+1, n): + sum_potential += G * self.pairwise(xt[i], xt[j]) + + if self.extra_potential is not None: + for i in range(n): + sum_potential += self.extra_potential(xt[i]) + + return sum_potential + + def force_paddle(xt): + """ + Calculate force using PaddlePaddle automatic differentiation. + + This is a key part of the migration: + JAX: force = -grad(total_potential)(xt) + Paddle: Use paddle.grad to calculate gradient + + Complete implementation for all potential scenarios. + + Args: + xt (numpy.ndarray): System state with shape (n, total_dim). + + Returns: + numpy.ndarray: Force on each particle with shape (n, dim). + """ + # Convert numpy array to Paddle Tensor (improved precision: float32 → float64) + xt_tensor = paddle.to_tensor(xt, dtype='float64', stop_gradient=False) + + # Calculate gradient only for position coordinates + positions = xt_tensor[:, :dim] + positions.stop_gradient = False + + # Rebuild complete state for potential energy calculation + xt_for_potential = paddle.concat([ + positions, + xt_tensor[:, dim:] + ], axis=1) + + # Calculate total potential energy (using float64 for higher precision) + sum_potential = paddle.to_tensor(0.0, dtype='float64') + + # Iterate through all particle pairs to calculate potential energy + for i in range(n - 1): + if sim in ['string', 'string_ball']: + # ========== String type: only calculate adjacent particles ========== + x1 = xt_for_potential[i] + x2 = xt_for_potential[i+1] + dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) + bounded_dist = dist + 1e-2 + + if sim == 'string': + # String potential: spring + gravity + pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] + + elif sim == 'string_ball': + # String potential with ball obstacle + pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] + + # Calculate distance to ball center + r = paddle.sqrt((x1[1] + 15)**2 + (x1[0] - 5)**2) + radius = paddle.to_tensor(4.0, dtype='float64') + + # Soft boundary repulsion potential (reduced strength to avoid numerical instability) + # More stable repulsion potential: strong repulsion when r < radius + k_repel = 100.0 # Reduced repulsion strength + pot = pot + k_repel / paddle.maximum(r - radius + 0.5, paddle.to_tensor(0.01, dtype='float64')) + + sum_potential = sum_potential + G * pot + + else: + # ========== Fully connected type: calculate all particle pairs ========== + for j in range(i+1, n): + x1 = xt_for_potential[i] + x2 = xt_for_potential[j] + dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) + bounded_dist = dist + 1e-2 + + # Calculate potential based on simulation type + if sim == 'r2': + # Inverse square law: universal gravitation + pot = -x1[-1]*x2[-1]/bounded_dist + + elif sim == 'r1': + # Logarithmic potential + pot = x1[-1]*x2[-1]*paddle.log(bounded_dist) + + elif sim in ['spring', 'damped']: + # Spring potential + pot = (bounded_dist - 1)**2 + + # Add damping term + if sim == 'damped': + damping = paddle.to_tensor(1.0, dtype='float64') + # Damping dissipation term: proportional to position-velocity product + pot = pot + damping*x1[1]*x1[1+dim]/n # y direction + pot = pot + damping*x1[0]*x1[0+dim]/n # x direction + if dim == 3: + pot = pot + damping*x1[2]*x1[2+dim]/n # z direction + + elif sim in ['charge', 'superposition']: + # Charge interaction (Coulomb's law) + charge1 = x1[-2] + charge2 = x2[-2] + pot = charge1*charge2/bounded_dist + + if sim == 'superposition': + # Superposition: contains both charge and gravity + m1 = x1[-1] + m2 = x2[-1] + pot = pot + (-m1*m2/bounded_dist) + + elif sim == 'discontinuous': + # Piecewise discontinuous potential + # Define potential in three regions + pot_a = paddle.to_tensor(0.0, dtype='float64') # r < 1 + pot_b = paddle.to_tensor(0.0, dtype='float64') # 1 <= r < 2 + pot_c = (bounded_dist - 1)**2 # r >= 2 + + # Implement piecewise function using conditional expressions + # Paddle's conditional operations automatically handle gradients + cond1 = paddle.cast(bounded_dist < 1, 'float64') + cond2 = paddle.cast(bounded_dist >= 1, 'float64') * paddle.cast(bounded_dist < 2, 'float64') + cond3 = paddle.cast(bounded_dist >= 2, 'float64') + + pot = pot_a * cond1 + pot_b * cond2 + pot_c * cond3 + + else: + # Default: spring potential + pot = (bounded_dist - 1)**2 + + sum_potential = sum_potential + G * pot + + # Add external potential support + if self.extra_potential is not None: + # Note: extra_potential needs to be Paddle-compatible + # If user provides a numpy function, there will be issues + # Suggestion: use numerical gradients or require Paddle version from user + for i in range(n): + # Assume extra_potential returns scalar + # May need adjustment in actual use + try: + extra_pot = self.extra_potential(xt_for_potential[i]) + if not isinstance(extra_pot, paddle.Tensor): + extra_pot = paddle.to_tensor(extra_pot, dtype='float64') + sum_potential = sum_potential + extra_pot + except: + # If extra_potential is not Paddle-compatible, skip + # Should give warning in actual use + pass + + # Calculate gradient: F = -dU/dx + # Use paddle.grad for automatic differentiation + grads = paddle.grad( + outputs=sum_potential, + inputs=positions, + create_graph=False, + retain_graph=False + )[0] + + # Convert back to numpy + force = -grads.numpy() + + return force + + def acceleration(xt): + """ + Calculate acceleration: a = F/m. + + Args: + xt (numpy.ndarray): System state. + + Returns: + numpy.ndarray: Acceleration. + """ + force = force_paddle(xt) + masses = xt[:, -1:] # Shape (n, 1) + return force / masses + + def odefunc(y, t): + """ + Ordinary differential equation function (for scipy.odeint). + + Args: + y (numpy.ndarray): Flattened system state. + t (float): Time. + + Returns: + numpy.ndarray: Time derivative of the state. + """ + # Restore shape + y = y.reshape((n, total_dim)) + + # Calculate acceleration + a = acceleration(y) + + # Construct derivative + dydt = np.concatenate([ + y[:, dim:2*dim], # dx/dt = v + a, # dv/dt = a + np.zeros((n, params)) # d(q,m)/dt = 0 + ], axis=1) + + return dydt.flatten() + + def make_sim(sample_idx): + """ + Generate single simulation sample. + + Args: + sample_idx (int): Sample index (for progress display). + + Returns: + numpy.ndarray: Simulation trajectory with shape (nt, n, total_dim). + """ + if sim in ['string', 'string_ball']: + # String type initialization + x0 = np.random.randn(n, total_dim) + x0[:, -1] = 1.0 # Fixed mass + x0[:, 0] = np.arange(n) + x0[:, 0]*0.5 # Evenly spaced x-coordinates + x0[:, 2:3] = 0.0 # vx = 0 + else: + # General initialization + x0 = np.random.randn(n, total_dim) + x0[:, -1] = np.exp(x0[:, -1]) # Positive mass + + if sim in ['charge', 'superposition']: + x0[:, -2] = np.sign(x0[:, -2]) # Charge ±1 + + # Solve ODE using scipy.odeint + # Migration point: JAX odeint → scipy.integrate.odeint + x_times = odeint( + odefunc, + x0.flatten(), + times, + mxstep=2000 + ).reshape(-1, n, total_dim) + + return x_times + + # Generate samples in batch + data = [] + print(f"Start generating {ns} samples...") + for i in tqdm(range(ns)): + data.append(make_sim(i)) + + # Convert to numpy array + self.data = np.array(data) + print(f"Generation complete! Data shape: {self.data.shape}") + + def get_acceleration(self): + """ + Calculate acceleration for all simulation data. + + This method uses PaddlePaddle automatic differentiation to calculate acceleration. + Uses the same complete potential implementation as in simulate(). + + Returns: + numpy.ndarray: Acceleration data with shape (ns, nt, n, dim). + """ + if self.data is None: + raise ValueError("Please call simulate() first to generate data") + + ns, nt, n, total_dim = self.data.shape + dim = self._dim + sim = self._sim + G = self.G + + # Pre-allocate acceleration array + accelerations = np.zeros((ns, nt, n, dim)) + + print("Calculating acceleration...") + for sample_idx in tqdm(range(ns)): + for time_idx in range(nt): + xt = self.data[sample_idx, time_idx] + + # Calculate acceleration using PaddlePaddle (using float64 for higher precision) + xt_tensor = paddle.to_tensor(xt, dtype='float64', stop_gradient=False) + positions = xt_tensor[:, :dim] + positions.stop_gradient = False + + # Calculate total potential + xt_for_potential = paddle.concat([ + positions, + xt_tensor[:, dim:] + ], axis=1) + + sum_potential = paddle.to_tensor(0.0, dtype='float64') + + # Iterate through all particle pairs to calculate potential (identical to force_paddle) + for i in range(n - 1): + if sim in ['string', 'string_ball']: + # ========== String type ========== + x1 = xt_for_potential[i] + x2 = xt_for_potential[i+1] + dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) + bounded_dist = dist + 1e-2 + + if sim == 'string': + pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] + elif sim == 'string_ball': + # Complete implementation + pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] + r = paddle.sqrt((x1[1] + 15)**2 + (x1[0] - 5)**2) + radius = paddle.to_tensor(4.0, dtype='float64') + # More stable repulsion potential + k_repel = 100.0 + pot = pot + k_repel / paddle.maximum(r - radius + 0.5, paddle.to_tensor(0.01, dtype='float64')) + + sum_potential = sum_potential + G * pot + else: + # ========== Fully connected type ========== + for j in range(i+1, n): + x1 = xt_for_potential[i] + x2 = xt_for_potential[j] + dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) + bounded_dist = dist + 1e-2 + + if sim == 'r2': + pot = -x1[-1]*x2[-1]/bounded_dist + elif sim == 'r1': + pot = x1[-1]*x2[-1]*paddle.log(bounded_dist) + elif sim in ['spring', 'damped']: + pot = (bounded_dist - 1)**2 + + # Add damping term + if sim == 'damped': + damping = paddle.to_tensor(1.0, dtype='float64') + pot = pot + damping*x1[1]*x1[1+dim]/n + pot = pot + damping*x1[0]*x1[0+dim]/n + if dim == 3: + pot = pot + damping*x1[2]*x1[2+dim]/n + + elif sim in ['charge', 'superposition']: + charge1 = x1[-2] + charge2 = x2[-2] + pot = charge1*charge2/bounded_dist + if sim == 'superposition': + m1 = x1[-1] + m2 = x2[-1] + pot = pot + (-m1*m2/bounded_dist) + + elif sim == 'discontinuous': + # Piecewise discontinuous potential + pot_a = paddle.to_tensor(0.0, dtype='float64') + pot_b = paddle.to_tensor(0.0, dtype='float64') + pot_c = (bounded_dist - 1)**2 + + cond1 = paddle.cast(bounded_dist < 1, 'float64') + cond2 = paddle.cast(bounded_dist >= 1, 'float64') * paddle.cast(bounded_dist < 2, 'float64') + cond3 = paddle.cast(bounded_dist >= 2, 'float64') + + pot = pot_a * cond1 + pot_b * cond2 + pot_c * cond3 + + else: + pot = (bounded_dist - 1)**2 + + sum_potential = sum_potential + G * pot + + # Add external potential support + if self.extra_potential is not None: + for i in range(n): + try: + extra_pot = self.extra_potential(xt_for_potential[i]) + if not isinstance(extra_pot, paddle.Tensor): + extra_pot = paddle.to_tensor(extra_pot, dtype='float64') + sum_potential = sum_potential + extra_pot + except: + pass + + # Calculate force + grads = paddle.grad( + outputs=sum_potential, + inputs=positions, + create_graph=False, + retain_graph=False + )[0] + + force = -grads.numpy() + masses = xt[:, -1:] + accel = force / masses + + accelerations[sample_idx, time_idx] = accel + + return accelerations + + def plot(self, i, animate=False, plot_size=True, s_size=1): + """ + Visualize simulation results. + + Args: + i (int): Sample index. + animate (bool, optional): Whether to generate animation. Defaults to False. + plot_size (bool, optional): Whether to adjust point size based on mass. Defaults to True. + s_size (float, optional): Point size scaling factor. Defaults to 1. + + Returns: + None (static plot) or HTML (animation). + """ + if self.data is None: + raise ValueError("Please call simulate() first to generate data") + + n = self._n + times = self.times + x_times = self.data[i] + sim = self._sim + masses = x_times[:, :, -1] + + if not animate: + # Static plot + if sim in ['string', 'string_ball']: + rgba = make_transparent_color(len(times), 0) + for idx in range(0, len(times), len(times)//10): + ctimes = x_times[idx] + plt.plot(ctimes[:, 0], ctimes[:, 1], color=rgba[idx]) + plt.xlim(-5, 20) + plt.ylim(-20, 5) + else: + for j in range(n): + rgba = make_transparent_color(len(times), j/n) + if plot_size: + plt.scatter(x_times[:, j, 0], x_times[:, j, 1], + color=rgba, s=3*masses[:, j]*s_size) + else: + plt.scatter(x_times[:, j, 0], x_times[:, j, 1], + color=rgba, s=s_size) + else: + # Animation + if sim in ['string', 'string_ball']: + raise NotImplementedError("Animation mode not yet supported for string type") + + fig = plt.figure() + camera = Camera(fig) + d_idx = 20 + + for t_idx in range(d_idx, len(times), d_idx): + start = max([0, t_idx-300]) + ctimes = times[start:t_idx] + cx_times = x_times[start:t_idx] + + for j in range(n): + rgba = make_transparent_color(len(ctimes), j/n) + if plot_size: + plt.scatter(cx_times[:, j, 0], cx_times[:, j, 1], + color=rgba, s=3*masses[:, j]) + else: + plt.scatter(cx_times[:, j, 0], cx_times[:, j, 1], + color=rgba, s=s_size) + + camera.snap() + + from IPython.display import HTML + return HTML(camera.animate().to_jshtml()) \ No newline at end of file diff --git a/ppsci/arch/__init__.py b/ppsci/arch/__init__.py index 78f381b68e..f473aba392 100644 --- a/ppsci/arch/__init__.py +++ b/ppsci/arch/__init__.py @@ -52,6 +52,7 @@ from ppsci.arch.physx_transformer import PhysformerGPT2 # isort:skip from ppsci.arch.sfnonet import SFNONet # isort:skip from ppsci.arch.spinn import SPINN # isort:skip +from ppsci.arch.symbolic_gn import OGN, VarOGN, HGN, get_edge_index # isort:skip from ppsci.arch.tfnonet import TFNO1dNet, TFNO2dNet, TFNO3dNet # isort:skip from ppsci.arch.transformer import Transformer # isort:skip from ppsci.arch.unetex import UNetEx # isort:skip @@ -95,8 +96,10 @@ "ExtFormerMoECuboid", "FNO1d", "Generator", + "get_edge_index", "GraphCastNet", "HEDeepONets", + "HGN", "LorenzEmbedding", "LatentNO", "LatentNO_time", @@ -105,6 +108,7 @@ "ModelList", "ModifiedMLP", "NowcastNet", + "OGN", "PhyCRNet", "PhysformerGPT2", "PirateNet", @@ -119,6 +123,7 @@ "UNetEx", "UNONet", "USCNN", + "VarOGN", "VelocityDiscriminator", "VelocityGenerator", "RegDGCNN", diff --git a/ppsci/arch/symbolic_gn.py b/ppsci/arch/symbolic_gn.py new file mode 100644 index 0000000000..2ce05c4865 --- /dev/null +++ b/ppsci/arch/symbolic_gn.py @@ -0,0 +1,667 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import paddle +import paddle.nn as nn +import numpy as np +from typing import Dict, List, Optional, Tuple + +from ppsci.arch import base + +def get_edge_index(n: int, sim: str) -> np.ndarray: + """ + Generate edge indices for graph connectivity. + + Args: + n (int): Number of nodes. + sim (str): Simulation type. + + Returns: + numpy.ndarray: Edge indices with shape [2, num_edges]. + """ + if sim in ['string', 'string_ball']: + # Chain connection + top = np.arange(0, n-1) + bottom = np.arange(1, n) + edge_index = np.concatenate([ + np.concatenate([top, bottom])[None, :], + np.concatenate([bottom, top])[None, :] + ], axis=0) + else: + # Full connection + adj = (np.ones((n, n)) - np.eye(n)).astype(int) + edge_index = np.array(np.where(adj)) + + return edge_index + + +class OGN(base.Arch): + def __init__( + self, + input_keys: List[str], + output_keys: List[str], + n_f: int = 6, + msg_dim: int = 100, + ndim: int = 2, + hidden: int = 300, + edge_index: Optional[np.ndarray] = None, + aggr: str = 'sum' + ): + """ + Initialize Object-based Graph Network (OGN). + + Args: + input_keys (List[str]): List of input keys, e.g., ['x', 'edge_index']. + output_keys (List[str]): List of output keys, e.g., ['acceleration']. + n_f (int, optional): Node feature dimension. Defaults to 6. + msg_dim (int, optional): Message dimension. Defaults to 100. + ndim (int, optional): Spatial dimension. Defaults to 2. + hidden (int, optional): Hidden layer size. Defaults to 300. + edge_index (np.ndarray, optional): Edge indices (can also be provided in forward). Defaults to None. + aggr (str, optional): Aggregation method, 'sum' or 'mean'. Defaults to 'sum'. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> model = ppsci.arch.OGN( + ... input_keys=["x"], + ... output_keys=["acceleration"], + ... n_f=6, + ... ndim=2 + ... ) + >>> n_nodes = 5 + >>> edge_index = np.array([[0, 1, 1, 2, 2, 3, 3, 4], [1, 0, 2, 1, 3, 2, 4, 3]]) + >>> input_dict = { + ... "x": paddle.randn([n_nodes, 6]), + ... "edge_index": edge_index + ... } + >>> output = model(input_dict) + >>> print(output["acceleration"].shape) + [5, 2] + """ + super().__init__() + + # PaddleScience standard: store input/output keys + self.input_keys = input_keys + self.output_keys = output_keys + + # Model parameters + self.n_f = n_f + self.msg_dim = msg_dim + self.ndim = ndim + self.hidden = hidden + self.aggr = aggr + + # Edge index (optional) + if edge_index is not None: + self.register_buffer('edge_index_buffer', + paddle.to_tensor(edge_index, dtype='int64')) + else: + self.edge_index_buffer = None + + # Message function network + self.msg_fnc = nn.Sequential( + nn.Linear(2*n_f, hidden), + nn.ReLU(), + nn.Linear(hidden, hidden), + nn.ReLU(), + nn.Linear(hidden, hidden), + nn.ReLU(), + nn.Linear(hidden, msg_dim) + ) + + # Node update network + self.node_fnc = nn.Sequential( + nn.Linear(msg_dim + n_f, hidden), + nn.ReLU(), + nn.Linear(hidden, hidden), + nn.ReLU(), + nn.Linear(hidden, hidden), + nn.ReLU(), + nn.Linear(hidden, ndim) + ) + + def message_passing(self, x: paddle.Tensor, edge_index: np.ndarray) -> paddle.Tensor: + """ + Execute message passing. + + Args: + x: Node features with shape [n, n_f]. + edge_index: Edge indices with shape [2, num_edges]. + + Returns: + paddle.Tensor: Updated node features with shape [n, ndim]. + """ + # Get source and target nodes + row, col = edge_index[0], edge_index[1] + + # Collect neighbor features + x_i = x[col] # Target nodes + x_j = x[row] # Source nodes + + # Compute messages + msg_input = paddle.concat([x_i, x_j], axis=1) + msg = self.msg_fnc(msg_input) + + # Aggregate messages to target nodes + num_nodes = x.shape[0] + aggr_out = paddle.zeros([num_nodes, self.msg_dim], dtype=msg.dtype) + + for i in range(len(col)): + aggr_out[col[i]] += msg[i] + + # Update nodes + node_input = paddle.concat([x, aggr_out], axis=1) + out = self.node_fnc(node_input) + + return out + + def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: + """ + Forward propagation (PaddleScience standard interface). + + Args: + inputs (Dict): Input dictionary containing: + - 'x': Node features [batch*n, n_f] or [n, n_f]. + - 'edge_index': Edge index information (optional if provided at initialization). + - Other possible inputs. + + Returns: + Dict: Output dictionary containing: + - 'acceleration': Predicted acceleration [batch*n, ndim]. + """ + # Extract input + x = inputs['x'] + + # Get edge index + if 'edge_index' in inputs: + # Get from input + edge_index = inputs['edge_index'] + if isinstance(edge_index, paddle.Tensor): + edge_index = edge_index.numpy() + elif self.edge_index_buffer is not None: + # Use stored edge index + edge_index = self.edge_index_buffer.numpy() + else: + raise ValueError("Must provide edge_index") + + # Execute message passing + acceleration = self.message_passing(x, edge_index) + + # Return dictionary format (PaddleScience standard) + return {self.output_keys[0]: acceleration} + + def compute_loss( + self, + inputs: Dict[str, paddle.Tensor], + labels: Dict[str, paddle.Tensor] + ) -> paddle.Tensor: + """ + Compute loss function (PaddleScience standard interface). + + Args: + inputs: Input dictionary. + labels: Label dictionary containing 'acceleration_true'. + + Returns: + paddle.Tensor: Loss scalar. + """ + # Forward propagation + outputs = self.forward(inputs) + + # Compute loss + pred = outputs[self.output_keys[0]] + true = labels.get('acceleration_true', labels.get(self.output_keys[0])) + + # MAE loss + loss = paddle.mean(paddle.abs(pred - true)) + + return loss + + +class HGN(base.Arch): + + def __init__( + self, + input_keys: List[str], + output_keys: List[str], + n_f: int = 6, + ndim: int = 2, + hidden: int = 300, + edge_index: Optional[np.ndarray] = None + ): + """ + Initialize Hamiltonian Graph Network (HGN). + + Args: + input_keys: List of input keys. + output_keys: List of output keys, e.g., ['velocity_derivative', 'acceleration']. + n_f: Node feature dimension. + ndim: Spatial dimension. + hidden: Hidden layer size. + edge_index: Edge indices (optional). + + Examples: + >>> import ppsci + >>> import numpy as np + >>> model = ppsci.arch.HGN( + ... input_keys=["x"], + ... output_keys=["acceleration"], + ... n_f=6, + ... ndim=2 + ... ) + >>> n_nodes = 5 + >>> edge_index = np.array([[0, 1, 1, 2, 2, 3, 3, 4], [1, 0, 2, 1, 3, 2, 4, 3]]) + >>> input_dict = { + ... "x": paddle.randn([n_nodes, 6]), + ... "edge_index": edge_index + ... } + >>> output = model(input_dict) + >>> print(output["acceleration"].shape) + [5, 2] + """ + super().__init__() + + self.input_keys = input_keys + self.output_keys = output_keys + self.n_f = n_f + self.ndim = ndim + self.hidden = hidden + + # Edge index + if edge_index is not None: + self.register_buffer('edge_index_buffer', + paddle.to_tensor(edge_index, dtype='int64')) + else: + self.edge_index_buffer = None + + # Pairwise energy network + self.pair_energy = nn.Sequential( + nn.Linear(2*n_f, hidden), + nn.Softplus(), + nn.Linear(hidden, hidden), + nn.Softplus(), + nn.Linear(hidden, hidden), + nn.Softplus(), + nn.Linear(hidden, 1) + ) + + # Self energy network + self.self_energy = nn.Sequential( + nn.Linear(n_f, hidden), + nn.Softplus(), + nn.Linear(hidden, hidden), + nn.Softplus(), + nn.Linear(hidden, hidden), + nn.Softplus(), + nn.Linear(hidden, 1) + ) + + def compute_energy(self, x: paddle.Tensor, edge_index: np.ndarray) -> paddle.Tensor: + """ + Compute the total energy (Hamiltonian) of the system. + + Args: + x: Node features. + edge_index: Edge indices. + + Returns: + paddle.Tensor: Energy per node. + """ + row, col = edge_index[0], edge_index[1] + + # Compute pairwise energy + x_i = x[col] + x_j = x[row] + edge_input = paddle.concat([x_i, x_j], axis=1) + pair_energies = self.pair_energy(edge_input) + + # Aggregate pairwise energy + num_nodes = x.shape[0] + aggr_pair = paddle.zeros([num_nodes, 1], dtype=pair_energies.dtype) + for i in range(len(col)): + aggr_pair[col[i]] += pair_energies[i] + + # Compute self energy + self_energies = self.self_energy(x) + + # Total energy + total_energy = aggr_pair + self_energies + + return total_energy + + def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: + """ + Forward propagation (PaddleScience standard). + + Compute dynamics derivatives using Hamilton's equations. + + Args: + inputs: Input dictionary. + + Returns: + Dict: Output dictionary. + """ + # Extract input + x_input = inputs['x'].clone() + + # Get edge index + if 'edge_index' in inputs: + edge_index = inputs['edge_index'] + if isinstance(edge_index, paddle.Tensor): + edge_index = edge_index.numpy() + elif self.edge_index_buffer is not None: + edge_index = self.edge_index_buffer.numpy() + else: + raise ValueError("Must provide edge_index") + + # Construct Hamiltonian coordinates + # Input: [q, v, other, m] + # Needed: [q, p=m*v, other] + q = x_input[:, :self.ndim] + v = x_input[:, self.ndim:2*self.ndim] + other = x_input[:, 2*self.ndim:] + + # Extract mass + m_scalar = other[:, -1:] + m_vec = paddle.tile(m_scalar, [1, self.ndim]) + + # Compute momentum + p = v * m_vec + + # Construct Hamiltonian coordinates + x_hamilton = paddle.concat([q, p, other], axis=1) + x_hamilton.stop_gradient = False + + # Compute total energy + total_energy = self.compute_energy(x_hamilton, edge_index) + total_energy_scalar = paddle.sum(total_energy) + + # Compute Hamiltonian gradients + dH = paddle.grad( + outputs=total_energy_scalar, + inputs=x_hamilton, + create_graph=False, + retain_graph=False + )[0] + + # Extract gradients + dH_dq = dH[:, :self.ndim] + dH_dp = dH[:, self.ndim:2*self.ndim] + + # Hamilton's equations + dq_dt = dH_dp + dp_dt = -dH_dq + dv_dt = dp_dt / m_vec + + # Construct output dictionary + outputs = {} + if 'velocity_derivative' in self.output_keys: + outputs['velocity_derivative'] = dq_dt + if 'acceleration' in self.output_keys: + outputs['acceleration'] = dv_dt + + # If only one output key, return directly + if len(self.output_keys) == 1: + outputs[self.output_keys[0]] = paddle.concat([dq_dt, dv_dt], axis=1) + + return outputs + + def compute_loss( + self, + inputs: Dict[str, paddle.Tensor], + labels: Dict[str, paddle.Tensor], + reg_weight: float = 1e-6 + ) -> paddle.Tensor: + """ + Compute loss with physical regularization. + + Args: + inputs: Input dictionary. + labels: Label dictionary. + reg_weight: Regularization weight. + + Returns: + paddle.Tensor: Total loss. + """ + # Forward propagation + outputs = self.forward(inputs) + + # Base loss + pred = outputs.get('acceleration', outputs.get(self.output_keys[0])) + true = labels.get('acceleration_true', labels.get('acceleration')) + + base_loss = paddle.mean(paddle.abs(pred - true)) + + # Physical regularization: energy should not depend on non-physical quantities + x_input = inputs['x'].clone() + edge_index = inputs.get('edge_index', self.edge_index_buffer.numpy()) + if isinstance(edge_index, paddle.Tensor): + edge_index = edge_index.numpy() + + # Construct Hamiltonian coordinates + q = x_input[:, :self.ndim] + v = x_input[:, self.ndim:2*self.ndim] + other = x_input[:, 2*self.ndim:] + + m_scalar = other[:, -1:] + m_vec = paddle.tile(m_scalar, [1, self.ndim]) + p = v * m_vec + + x_hamilton = paddle.concat([q, p, other], axis=1) + x_hamilton.stop_gradient = False + + # Compute energy + total_energy = self.compute_energy(x_hamilton, edge_index) + + # Regularization: penalize dependence on non-physical quantities + regularization = reg_weight * paddle.mean(total_energy**2) + + return base_loss + regularization + + +class VarOGN(base.Arch): + + def __init__( + self, + input_keys: List[str], + output_keys: List[str], + n_f: int = 6, + msg_dim: int = 100, + ndim: int = 2, + hidden: int = 300, + edge_index: Optional[np.ndarray] = None, + enable_sampling: bool = True + ): + """ + Initialize Variational Graph Network (VarGN). + + Args: + input_keys: List of input keys. + output_keys: List of output keys. + n_f: Node feature dimension. + msg_dim: Message dimension. + ndim: Spatial dimension. + hidden: Hidden layer size. + edge_index: Edge indices. + enable_sampling: Whether to enable sampling. + + Examples: + >>> import ppsci + >>> import numpy as np + >>> model = ppsci.arch.VarGN( + ... input_keys=["x"], + ... output_keys=["acceleration_mean", "acceleration_std"], + ... n_f=6, + ... ndim=2 + ... ) + >>> n_nodes = 5 + >>> edge_index = np.array([[0, 1, 1, 2, 2, 3, 3, 4], [1, 0, 2, 1, 3, 2, 4, 3]]) + >>> input_dict = { + ... "x": paddle.randn([n_nodes, 6]), + ... "edge_index": edge_index + ... } + >>> output = model(input_dict) + >>> print(output["acceleration_mean"].shape) + [5, 2] + """ + super().__init__() + + self.input_keys = input_keys + self.output_keys = output_keys + self.n_f = n_f + self.msg_dim = msg_dim + self.ndim = ndim + self.hidden = hidden + self.enable_sampling = enable_sampling + + # Edge index + if edge_index is not None: + self.register_buffer('edge_index_buffer', + paddle.to_tensor(edge_index, dtype='int64')) + else: + self.edge_index_buffer = None + + # Message function: outputs mu and logvar + self.msg_fnc = nn.Sequential( + nn.Linear(2*n_f, hidden), + nn.ReLU(), + nn.Linear(hidden, hidden), + nn.ReLU(), + nn.Linear(hidden, hidden), + nn.ReLU(), + nn.Linear(hidden, msg_dim*2) # mu and logvar + ) + + # Node update function + self.node_fnc = nn.Sequential( + nn.Linear(msg_dim + n_f, hidden), + nn.ReLU(), + nn.Linear(hidden, hidden), + nn.ReLU(), + nn.Linear(hidden, hidden), + nn.ReLU(), + nn.Linear(hidden, ndim) + ) + + def message_passing(self, x: paddle.Tensor, edge_index: np.ndarray) -> Tuple[paddle.Tensor, paddle.Tensor]: + """ + Variational message passing. + + Returns: + Tuple[paddle.Tensor, paddle.Tensor]: (mean output, variance output). + """ + row, col = edge_index[0], edge_index[1] + + x_i = x[col] + x_j = x[row] + + # Compute mu and logvar for messages + msg_input = paddle.concat([x_i, x_j], axis=1) + raw_msg = self.msg_fnc(msg_input) + + mu = raw_msg[:, 0::2] + logvar = raw_msg[:, 1::2] + + # Sample messages + if self.enable_sampling and self.training: + epsilon = paddle.randn(mu.shape) + msg = mu + epsilon * paddle.exp(logvar / 2.0) + else: + msg = mu + + # Aggregate + num_nodes = x.shape[0] + aggr_out = paddle.zeros([num_nodes, self.msg_dim], dtype=msg.dtype) + + for i in range(len(col)): + aggr_out[col[i]] += msg[i] + + # Update nodes + node_input = paddle.concat([x, aggr_out], axis=1) + out_mean = self.node_fnc(node_input) + + # Compute output variance (simplified version) + out_std = paddle.exp(logvar.mean(axis=0, keepdim=True)) + + return out_mean, out_std + + def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: + """ + Forward propagation (PaddleScience standard). + + Args: + inputs: Input dictionary. + + Returns: + Dict: Output dictionary containing mean and standard deviation. + """ + x = inputs['x'] + + # Get edge index + if 'edge_index' in inputs: + edge_index = inputs['edge_index'] + if isinstance(edge_index, paddle.Tensor): + edge_index = edge_index.numpy() + elif self.edge_index_buffer is not None: + edge_index = self.edge_index_buffer.numpy() + else: + raise ValueError("Must provide edge_index") + + # Execute variational message passing + accel_mean, accel_std = self.message_passing(x, edge_index) + + # Construct output + outputs = {} + if 'acceleration_mean' in self.output_keys: + outputs['acceleration_mean'] = accel_mean + if 'acceleration_std' in self.output_keys: + outputs['acceleration_std'] = accel_std + + # Default output + if len(outputs) == 0: + outputs[self.output_keys[0]] = accel_mean + + return outputs + + def compute_loss( + self, + inputs: Dict[str, paddle.Tensor], + labels: Dict[str, paddle.Tensor], + kl_weight: float = 1e-3 + ) -> paddle.Tensor: + """ + Compute variational loss including KL divergence. + + Args: + inputs: Input dictionary. + labels: Label dictionary. + kl_weight: KL divergence weight. + + Returns: + paddle.Tensor: Total loss. + """ + outputs = self.forward(inputs) + + # Reconstruction loss + pred = outputs.get('acceleration_mean', outputs.get(self.output_keys[0])) + true = labels.get('acceleration_true', labels.get('acceleration')) + + recon_loss = paddle.mean(paddle.abs(pred - true)) + + # KL divergence (simplified version) + # Full version needs to be computed in message_passing + kl_loss = paddle.to_tensor(0.0) # Placeholder + + return recon_loss + kl_weight * kl_loss \ No newline at end of file From 2a9c50ee7c7289c14fac0eb45b07037ddaa3c535 Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Wed, 29 Oct 2025 23:56:11 +0800 Subject: [PATCH 03/10] Symbolic GNN --- examples/symbolic_gn/conf/config.yaml | 44 +- examples/symbolic_gn/conf/config_hgn.yaml | 139 ++++ examples/symbolic_gn/conf/config_varogn.yaml | 139 ++++ examples/symbolic_gn/gn.py | 703 +++++++++---------- ppsci/arch/symbolic_gn.py | 403 +++++++---- 5 files changed, 875 insertions(+), 553 deletions(-) create mode 100644 examples/symbolic_gn/conf/config_hgn.yaml create mode 100644 examples/symbolic_gn/conf/config_varogn.yaml diff --git a/examples/symbolic_gn/conf/config.yaml b/examples/symbolic_gn/conf/config.yaml index 2886423a6b..f38cf0a82f 100644 --- a/examples/symbolic_gn/conf/config.yaml +++ b/examples/symbolic_gn/conf/config.yaml @@ -23,18 +23,19 @@ hydra: subdir: ./ # general settings -mode: train # running mode: train/eval +mode: train # running mode: train/eval/export/inference seed: 42 output_dir: ${hydra:run.dir} log_freq: 20 +device: gpu TRAIN: epochs: 1000 - batch_size: auto # auto = int(64 * (4 / num_nodes) ** 2) - save_freq: 10 + batch_size: 32 + save_freq: 100 eval_during_train: true - eval_freq: 5 + eval_freq: 100 checkpoint_path: null pretrained_model_path: null train_split: 0.8 # Train/validation split ratio @@ -48,28 +49,39 @@ TRAIN: loss: type: "MAE" # MAE or MSE physics_weight: 0.0 # Physical regularization weight - message_analysis: - enabled: true - record_interval: 1 - test_sample_size: 1000 - +# evaluation settings EVAL: pretrained_model_path: null eval_with_no_grad: true - batch_size: 1024 + batch_size: 32 + message_analysis: + enabled: true + +# export settings +export_path: "./exported_model" INFER: pretrained_model_path: null export_path: ./inference/ogn_model - device: gpu + pdmodel_path: ./inference/ogn_model.pdmodel + pdiparams_path: ./inference/ogn_model.pdiparams + onnx_path: ${INFER.export_path}.onnx + device: cpu engine: native precision: fp32 + batch_size: 1 + ir_optim: false + min_subgraph_size: 1 + gpu_mem: 2000 + gpu_id: 0 + max_batch_size: 1024 + num_cpu_threads: 10 MODEL: arch: "OGN" # OGN, varOGN, HGN - input_keys: ["node_features", "edge_index"] + input_keys: ["x", "edge_index"] output_keys: ["acceleration"] n_f: auto # auto = dimension * 2 + 2 msg_dim: 100 @@ -82,11 +94,11 @@ MODEL: DATA: type: "spring" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball - num_samples: 10000 - num_nodes: 4 + num_samples: 1000 + num_nodes: 6 dimension: 2 - time_steps: 1000 - time_step_size: auto + time_steps: 10 + time_step_size: 0.01 downsample_factor: 5 sample_interval: 5 # Sampling interval for training data diff --git a/examples/symbolic_gn/conf/config_hgn.yaml b/examples/symbolic_gn/conf/config_hgn.yaml new file mode 100644 index 0000000000..fd2f97c5ff --- /dev/null +++ b/examples/symbolic_gn/conf/config_hgn.yaml @@ -0,0 +1,139 @@ +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: + # dynamic output directory according to running time and override name + dir: outputs_symbolic_gcn/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} + job: + name: ${mode} # name of logfile + chdir: false # keep current working directory unchanged + callbacks: + init_callback: + _target_: ppsci.utils.callbacks.InitCallback + sweep: + # output directory for multirun + dir: ${hydra.run.dir} + subdir: ./ + +# general settings +mode: train # running mode: train/eval/export/inference +seed: 42 +output_dir: ${hydra:run.dir} +log_freq: 20 +device: gpu + +TRAIN: + epochs: 10 + batch_size: 32 + save_freq: 10 + eval_during_train: true + eval_freq: 5 + checkpoint_path: null + pretrained_model_path: null + train_split: 0.8 # Train/validation split ratio + optimizer: + learning_rate: 1e-3 + weight_decay: 1e-8 + lr_scheduler: + name: "OneCycleLR" + max_learning_rate: 1e-3 + final_div_factor: 1e5 + loss: + type: "MAE" # MAE or MSE + physics_weight: 0.0 # Physical regularization weight + message_analysis: + enabled: true + record_interval: 1 + test_sample_size: 10 + +EVAL: + pretrained_model_path: null + eval_with_no_grad: true + batch_size: 1024 + +# export settings +export_path: "./exported_model_hgn" + +INFER: + pretrained_model_path: null + export_path: ./inference/hgn_model + device: gpu + engine: native + precision: fp32 + +MODEL: + arch: "HGN" # Hamiltonian Graph Network + input_keys: ["x", "edge_index"] + output_keys: ["acceleration"] # HGN outputs acceleration (same as VarOGN) + n_f: auto # auto = dimension * 2 + 2 + msg_dim: 100 # Not used in HGN, but kept for compatibility + ndim: auto # auto = dimension + dt: 0.1 + hidden: 300 + aggr: "add" # Not used in HGN, but kept for compatibility + regularization_type: "energy" # HGN uses energy regularization + l1_strength: 0.0 # Not used in HGN + +DATA: + type: "spring" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball + num_samples: 100 + num_nodes: 4 + dimension: 2 + time_steps: 10 + time_step_size: 0.01 + downsample_factor: 5 + sample_interval: 5 # Sampling interval for training data + +DATATYPE_TYPE: + - sim: "r1" + dt: [5e-3] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "r2" + dt: [1e-3] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "spring" + dt: [1e-2] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "string" + dt: [1e-2] + nt: [100] + n: [30] + dim: [2] + - sim: "charge" + dt: [1e-3] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "superposition" + dt: [1e-3] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "damped" + dt: [2e-2] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "discontinuous" + dt: [1e-2] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "string_ball" + dt: [1e-2] + nt: [100] + n: [30] + dim: [2] diff --git a/examples/symbolic_gn/conf/config_varogn.yaml b/examples/symbolic_gn/conf/config_varogn.yaml new file mode 100644 index 0000000000..f1ad28016b --- /dev/null +++ b/examples/symbolic_gn/conf/config_varogn.yaml @@ -0,0 +1,139 @@ +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: + # dynamic output directory according to running time and override name + dir: outputs_symbolic_gcn/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} + job: + name: ${mode} # name of logfile + chdir: false # keep current working directory unchanged + callbacks: + init_callback: + _target_: ppsci.utils.callbacks.InitCallback + sweep: + # output directory for multirun + dir: ${hydra.run.dir} + subdir: ./ + +# general settings +mode: train # running mode: train/eval/export/inference +seed: 42 +output_dir: ${hydra:run.dir} +log_freq: 20 +device: gpu + +TRAIN: + epochs: 10 + batch_size: 32 + save_freq: 10 + eval_during_train: true + eval_freq: 5 + checkpoint_path: null + pretrained_model_path: null + train_split: 0.8 # Train/validation split ratio + optimizer: + learning_rate: 1e-3 + weight_decay: 1e-8 + lr_scheduler: + name: "OneCycleLR" + max_learning_rate: 1e-3 + final_div_factor: 1e5 + loss: + type: "MAE" # MAE or MSE + physics_weight: 0.0 # Physical regularization weight + message_analysis: + enabled: true + record_interval: 1 + test_sample_size: 10 + +EVAL: + pretrained_model_path: null + eval_with_no_grad: true + batch_size: 1024 + +# export settings +export_path: "./exported_model_varogn" + +INFER: + pretrained_model_path: null + export_path: ./inference/varogn_model + device: gpu + engine: native + precision: fp32 + +MODEL: + arch: "VarOGN" # Variational Object-based Graph Network + input_keys: ["x", "edge_index"] + output_keys: ["acceleration"] # VarOGN outputs mean acceleration + n_f: auto # auto = dimension * 2 + 2 + msg_dim: 100 + ndim: auto # auto = dimension + dt: 0.1 + hidden: 300 + aggr: "add" # add, mean, max, min + regularization_type: "kl" # VarOGN uses KL divergence regularization + l1_strength: 1e-3 # L1 regularization strength + +DATA: + type: "spring" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball + num_samples: 100 + num_nodes: 4 + dimension: 2 + time_steps: 10 + time_step_size: 0.01 + downsample_factor: 5 + sample_interval: 5 # Sampling interval for training data + +DATATYPE_TYPE: + - sim: "r1" + dt: [5e-3] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "r2" + dt: [1e-3] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "spring" + dt: [1e-2] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "string" + dt: [1e-2] + nt: [100] + n: [30] + dim: [2] + - sim: "charge" + dt: [1e-3] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "superposition" + dt: [1e-3] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "damped" + dt: [2e-2] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "discontinuous" + dt: [1e-2] + nt: [100] + n: [4, 8] + dim: [2, 3] + - sim: "string_ball" + dt: [1e-2] + nt: [100] + n: [30] + dim: [2] diff --git a/examples/symbolic_gn/gn.py b/examples/symbolic_gn/gn.py index 7127c21164..331a43416b 100644 --- a/examples/symbolic_gn/gn.py +++ b/examples/symbolic_gn/gn.py @@ -14,505 +14,440 @@ import os import sys +from typing import List, Optional import numpy as np import paddle from omegaconf import DictConfig import hydra import ppsci from ppsci.utils import logger +import matplotlib.pyplot as plt from simulate import SimulationDataset from ppsci.arch.symbolic_gn import OGN, VarOGN, HGN, get_edge_index - -class OGNDataset: - def __init__(self, cfg, mode="train"): - self.cfg = cfg - self.mode = mode - self.data = None - self.labels = None - self.input_keys = tuple(cfg.MODEL.input_keys) - self.label_keys = tuple(cfg.MODEL.output_keys) - self._prepare_data() - - def _prepare_data(self): - # Get time step size for physics system - dt = 1e-2 - for sim_set in self.cfg.DATATYPE_TYPE: - if sim_set['sim'] == self.cfg.DATA.type: - dt = sim_set['dt'][0] - break - - # Generate simulation data - sim_dataset = SimulationDataset( - sim=self.cfg.DATA.type, - n=self.cfg.DATA.num_nodes, - dim=self.cfg.DATA.dimension, - dt=dt, - nt=self.cfg.DATA.time_steps, - seed=self.cfg.seed +def create_ppsci_model( + model_type: str = 'OGN', + input_keys: List[str] = ['x', 'edge_index'], + output_keys: List[str] = ['acceleration'], + n_f: int = 6, + msg_dim: int = 100, + ndim: int = 2, + hidden: int = 300, + edge_index: Optional[np.ndarray] = None, + l1_strength: float = 0.0 +): + if model_type == 'OGN': + model = OGN( + input_keys=input_keys, + output_keys=output_keys, + n_f=n_f, + msg_dim=msg_dim, + ndim=ndim, + hidden=hidden, + edge_index=edge_index, + l1_strength=l1_strength ) - sim_dataset.simulate(self.cfg.DATA.num_samples) - - # Compute target data based on model architecture - if self.cfg.MODEL.arch == "HGN": - raw_target_data = sim_dataset.get_derivative() # [v, a] for HGN - else: - raw_target_data = sim_dataset.get_acceleration() # Acceleration for OGN/VarOGN - - # Downsample data - downsample_factor = self.cfg.DATA.downsample_factor - input_data = np.concatenate([sim_dataset.data[:, i] for i in range(0, sim_dataset.data.shape[1], downsample_factor)]) - target_data = np.concatenate([raw_target_data[:, i] for i in range(0, sim_dataset.data.shape[1], downsample_factor)]) - - # Split train/validation set - from sklearn.model_selection import train_test_split - test_size = 1.0 - getattr(self.cfg.TRAIN, 'train_split', 0.8) - - if self.mode == "train": - input_data, _, target_data, _ = train_test_split( - input_data, target_data, test_size=test_size, shuffle=False - ) - elif self.mode == "eval": - _, input_data, _, target_data = train_test_split( - input_data, target_data, test_size=test_size, shuffle=False - ) - - # Generate edge indices - edge_index = get_edge_index(self.cfg.DATA.num_nodes, self.cfg.DATA.type) - - self.data = { - "node_features": input_data.astype(np.float32), - "edge_index": np.tile(edge_index.T.numpy(), (len(input_data), 1, 1)) - } - self.labels = { - list(self.cfg.MODEL.output_keys)[0]: target_data.astype(np.float32) - } + elif model_type == 'HGN': + model = HGN( + input_keys=input_keys, + output_keys=output_keys, + n_f=n_f, + ndim=ndim, + hidden=hidden, + edge_index=edge_index + ) + elif model_type == 'VarOGN': + model = VarOGN( + input_keys=input_keys, + output_keys=output_keys, + n_f=n_f, + msg_dim=msg_dim, + ndim=ndim, + hidden=hidden, + edge_index=edge_index, + l1_strength=l1_strength + ) + else: + raise ValueError(f"未知的模型类型: {model_type}") + + return model - def __len__(self): - return len(self.data["node_features"]) - def __getitem__(self, idx): - return { - "input": { - "node_features": self.data["node_features"][idx], - "edge_index": self.data["edge_index"][idx] - }, - "label": { - list(self.cfg.MODEL.output_keys)[0]: self.labels[list(self.cfg.MODEL.output_keys)[0]][idx] - } - } - - def get_data_and_labels(self): - return self.data, self.labels - - def __call__(self): - return self +def replicate_array(arr, n): + return np.tile(arr, (n,) + (1,) * arr.ndim) def create_loss_function(cfg): - def loss_function(input_dict, label_dict, model_output): - # Get loss type from config (MAE or MSE) + def loss_function(output_dict, label_dict, weight_dict=None): loss_type = getattr(cfg.TRAIN.loss, 'type', 'MAE') - # Get output key - output_key = list(cfg.MODEL.output_keys)[0] if cfg.MODEL.arch != "HGN" else "derivative" + output_key = list(cfg.MODEL.output_keys)[0] - # Compute base loss - if cfg.MODEL.arch == "HGN": - output_key = "derivative" - pred_accel = model_output[output_key][:, cfg.DATA.dimension:] - target_accel = label_dict[output_key][:, cfg.DATA.dimension:] - if loss_type == "MSE": - base_loss = paddle.mean(paddle.square(pred_accel - target_accel)) - else: # MAE - base_loss = paddle.mean(paddle.abs(pred_accel - target_accel)) + target = label_dict[output_key] + pred = output_dict[output_key] + if loss_type == "MSE": + base_loss = paddle.mean(paddle.square(pred - target)) else: - pred = model_output[output_key] - target = label_dict[output_key] - if loss_type == "MSE": - base_loss = paddle.mean(paddle.square(pred - target)) - else: # MAE - base_loss = paddle.mean(paddle.abs(pred - target)) - - # Add regularization if configured - if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"]: - reg_loss = cfg.MODEL.l1_strength * paddle.mean(paddle.abs(input_dict["node_features"])) - return base_loss + reg_loss - - return base_loss + base_loss = paddle.mean(paddle.abs(pred - target)) + + if 'l1_regularization' in label_dict: + return base_loss + label_dict['l1_regularization'] + + return {output_key: base_loss} return loss_function def train(cfg): - # Set random seed for reproducibility ppsci.utils.misc.set_random_seed(cfg.seed) - # Parse automatic configuration parameters + if cfg.MODEL.arch == "HGN": + cfg.MODEL.output_keys = ["acceleration"] + + # Create simulation dataset + sim = SimulationDataset( + sim=cfg.DATA.type, + n=cfg.DATA.num_nodes, + dim=cfg.DATA.dimension, + nt=cfg.DATA.time_steps, + dt=cfg.DATA.time_step_size + ) + sim.simulate(cfg.DATA.num_samples) + accel_data = sim.get_acceleration() + + X_list = [] + y_list = [] + for sample_idx in range(cfg.DATA.num_samples): + for t in range(0, sim.data.shape[1], cfg.DATA.sample_interval): + X_list.append(sim.data[sample_idx, t]) + y_list.append(accel_data[sample_idx, t]) + + X = np.array(X_list, dtype=np.float32) + y = np.array(y_list, dtype=np.float32) + if cfg.MODEL.n_f == "auto": cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 if cfg.MODEL.ndim == "auto": cfg.MODEL.ndim = cfg.DATA.dimension - if cfg.TRAIN.batch_size == "auto": - cfg.TRAIN.batch_size = int(64 * (4 / cfg.DATA.num_nodes) ** 2) - - # Update output keys for HGN - if cfg.MODEL.arch == "HGN": - cfg.MODEL.output_keys = ["derivative"] - - # Create datasets - logger.message(f"Creating {cfg.DATA.type} dataset using simulate.py...") - train_dataset = OGNDataset(cfg, mode="train") - eval_dataset = OGNDataset(cfg, mode="eval") - - # Get training and evaluation data - train_data, train_labels = train_dataset.get_data_and_labels() - eval_data, eval_labels = eval_dataset.get_data_and_labels() - # Create model based on architecture - logger.message(f"Creating model: {cfg.MODEL.arch}") - if cfg.MODEL.arch == "OGN": - model = OGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "VarOGN": - model = VarOGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "HGN": - model = HGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - ndim=cfg.MODEL.ndim, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - else: - raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + train_size = int(len(X) * 0.8) + X_train, X_val = X[:train_size], X[train_size:] + y_train, y_val = y[:train_size], y[train_size:] + edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) + edge_index_train = replicate_array(edge_index, len(X_train)) + edge_index_val = replicate_array(edge_index, len(X_val)) - # Create loss function and constraints - loss_fn = create_loss_function(cfg) train_constraint = ppsci.constraint.SupervisedConstraint( { "dataset": { "name": "NamedArrayDataset", - "input": train_data, - "label": train_labels + "input": {"x": X_train, "edge_index": edge_index_train}, + "label": {cfg.MODEL.output_keys[0]: y_train}, }, "batch_size": cfg.TRAIN.batch_size, - "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": True}, + "sampler": { + "name": "BatchSampler", + "drop_last": False, + "shuffle": True, + }, }, - loss=loss_fn, - name="train_constraint", + create_loss_function(cfg), + name="sup_constraint", ) - eval_constraint = ppsci.constraint.SupervisedConstraint( { "dataset": { "name": "NamedArrayDataset", - "input": eval_data, - "label": eval_labels + "input": {"x": X_val, "edge_index": edge_index_val}, + "label": {cfg.MODEL.output_keys[0]: y_val}, + }, + "batch_size": cfg.TRAIN.batch_size, + "sampler": { + "name": "BatchSampler", + "drop_last": False, + "shuffle": False, }, - "batch_size": cfg.EVAL.batch_size, - "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": False}, }, - loss=loss_fn, + create_loss_function(cfg), name="eval_constraint", ) - constraint = { train_constraint.name: train_constraint, eval_constraint.name: eval_constraint, } + l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 + model = create_ppsci_model( + model_type=cfg.MODEL.arch, + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + edge_index=edge_index, + l1_strength=l1_strength + ) - # Calculate iterations per epoch - batch_per_epoch = int(1000*10 / (cfg.TRAIN.batch_size/32.0)) - - # Create optimizer and learning rate scheduler - if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": - lr_scheduler = paddle.optimizer.lr.OneCycleLR( - max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, - total_steps=cfg.TRAIN.epochs*batch_per_epoch, - divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor - ) - else: - lr_scheduler = paddle.optimizer.lr.ExponentialDecay( - learning_rate=cfg.TRAIN.optimizer.learning_rate, - gamma=0.9 - ) + # batch_per_epoch = int(1000*10 / (cfg.TRAIN.batch_size/32.0)) + # if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": + # lr_scheduler = paddle.optimizer.lr.OneCycleLR( + # max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, + # total_steps=cfg.TRAIN.epochs*batch_per_epoch, + # divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor + # ) + # lr_scheduler.by_epoch = False + # else: + # lr_scheduler = paddle.optimizer.lr.ExponentialDecay( + # learning_rate=cfg.TRAIN.optimizer.learning_rate, + # gamma=0.9 + # ) + # lr_scheduler.by_epoch = True optimizer = paddle.optimizer.Adam( - learning_rate=lr_scheduler, + learning_rate=cfg.TRAIN.optimizer.learning_rate,#lr_scheduler, parameters=model.parameters(), weight_decay=cfg.TRAIN.optimizer.weight_decay ) - # Create PaddleScience Solver and start training solver = ppsci.solver.Solver( model, constraint, - optimizer=optimizer, - cfg=cfg, + cfg.output_dir, + optimizer, + epochs=cfg.TRAIN.epochs, + save_freq=cfg.TRAIN.save_freq, ) + solver.train() + solver.plot_loss_history(by_epoch=True, smooth_step=1) def evaluate(cfg: DictConfig): - # Parse automatic configuration parameters + # Set model if cfg.MODEL.n_f == "auto": cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 if cfg.MODEL.ndim == "auto": cfg.MODEL.ndim = cfg.DATA.dimension - if cfg.MODEL.arch == "HGN": - cfg.MODEL.output_keys = ["derivative"] - # Create model based on architecture - if cfg.MODEL.arch == "OGN": - model = OGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "VarOGN": - model = VarOGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "HGN": - model = HGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - ndim=cfg.MODEL.ndim, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - else: - raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) + l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 - # Create evaluation dataset - eval_dataset = OGNDataset(cfg, mode="eval") - eval_data, eval_labels = eval_dataset.get_data_and_labels() + model = create_ppsci_model( + model_type=cfg.MODEL.arch, + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + edge_index=edge_index, + l1_strength=l1_strength + ) - # Create loss function and constraint - loss_fn = create_loss_function(cfg) - eval_constraint = ppsci.constraint.SupervisedConstraint( - { - "dataset": { - "name": "NamedArrayDataset", - "input": eval_data, - "label": eval_labels - }, - "batch_size": cfg.EVAL.batch_size, - "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": False}, - }, - loss=loss_fn, - name="eval_constraint", + # Load pretrained model + ppsci.utils.save_load.load_pretrain( + model, + cfg.EVAL.pretrained_model_path, ) + + # Generate test data + sim = SimulationDataset( + sim=cfg.DATA.type, + n=cfg.DATA.num_nodes, + dim=cfg.DATA.dimension, + nt=cfg.DATA.time_steps, + dt=cfg.DATA.time_step_size + ) + sim.simulate(cfg.DATA.num_samples) + accel_data = sim.get_acceleration() - constraint = { - eval_constraint.name: eval_constraint, - } + # Prepare test data + X_list = [] + y_list = [] + for sample_idx in range(cfg.DATA.num_samples): + for t in range(0, sim.data.shape[1], cfg.DATA.sample_interval): + X_list.append(sim.data[sample_idx, t]) + y_list.append(accel_data[sample_idx, t]) + + X = np.array(X_list, dtype=np.float32) + y = np.array(y_list, dtype=np.float32) - # Create PaddleScience Solver - solver = ppsci.solver.Solver(model, constraint, cfg=cfg) + # Evaluate on selected samples + sample_indices = [0, 1] if cfg.DATA.num_samples > 1 else [0] - # Run evaluation - logger.message("Starting evaluation...") - eval_result = solver.eval() - logger.message(f"Evaluation completed. Result: {eval_result}") + for sample_idx in sample_indices: + # Get data for this sample + sample_data = X[sample_idx:sample_idx+1] # Shape: [1, num_nodes, n_f] + true_accel = y[sample_idx:sample_idx+1] # Shape: [1, num_nodes, dim] + + # Prepare input - convert to PaddlePaddle tensors + input_dict = { + "x": paddle.to_tensor(sample_data, dtype="float32"), + "edge_index": paddle.to_tensor(edge_index, dtype="int64") + } + # Model prediction + with paddle.no_grad(): + pred_output = model(input_dict) + pred_accel = pred_output[cfg.MODEL.output_keys[0]] # Shape: [1, num_nodes, dim] + + # Calculate error + error = np.mean(np.abs(pred_accel.numpy() - true_accel)) + + # Calculate relative error + rel_error = np.linalg.norm(pred_accel.numpy() - true_accel) / np.linalg.norm(true_accel) + + # Visualization using simulate.py plot function + plt.figure(figsize=(10, 8)) + sim.plot(sample_idx, animate=False, plot_size=True, s_size=2) + plt.title(f'{cfg.DATA.type.capitalize()} System - Sample {sample_idx}\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}') + plt.tight_layout() + + # Save plot + plot_path = os.path.join(cfg.output_dir, f"evaluation_sample_{sample_idx}.png") + plt.savefig(plot_path, dpi=300, bbox_inches='tight') + plt.close() + logger.info(f"Evaluation plot saved to {plot_path}") -def inference(cfg: DictConfig): - # Parse automatic configuration parameters +def export(cfg: DictConfig): + if cfg.MODEL.arch == "HGN": + raise ValueError("HGN is not supported for export") + + ppsci.utils.misc.set_random_seed(cfg.seed) + if cfg.MODEL.n_f == "auto": cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 if cfg.MODEL.ndim == "auto": cfg.MODEL.ndim = cfg.DATA.dimension - if cfg.MODEL.arch == "HGN": - cfg.MODEL.output_keys = ["derivative"] + edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) + l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 + model = create_ppsci_model( + model_type=cfg.MODEL.arch, + n_f=cfg.MODEL.n_f, + msg_dim=cfg.MODEL.msg_dim, + ndim=cfg.MODEL.ndim, + hidden=cfg.MODEL.hidden, + edge_index=edge_index, + l1_strength=l1_strength + ) + + solver = ppsci.solver.Solver( + model, + pretrained_model_path=cfg.INFER.pretrained_model_path, + ) - # Create model based on architecture - if cfg.MODEL.arch == "OGN": - model = OGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "VarOGN": - model = VarOGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "HGN": - model = HGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - ndim=cfg.MODEL.ndim, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - else: - raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + from paddle.static import InputSpec - # Create PaddleScience Solver - solver = ppsci.solver.Solver(model, cfg=cfg) + input_spec = [ + { + key: InputSpec([None, cfg.DATA.num_nodes, cfg.MODEL.n_f], "float32", name=key) + if key == "x" + else InputSpec([2, 30], "int64", name=key) + for key in model.input_keys + }, + ] + solver.export(input_spec, cfg.INFER.export_path) + + +def inference(cfg: DictConfig): + from deploy import python_infer + + try: + predictor = python_infer.GeneralPredictor(cfg) + use_predictor = True + except Exception as e: + logger.error(f"GeneralPredictor failed: {e}") + use_predictor = False # Generate test data - logger.message("Generating inference data using simulate.py...") - dt = 1e-2 - for sim_set in cfg.DATATYPE_TYPE: - if sim_set['sim'] == cfg.DATA.type: - dt = sim_set['dt'][0] - break - - sim_dataset = SimulationDataset( + sim = SimulationDataset( sim=cfg.DATA.type, n=cfg.DATA.num_nodes, dim=cfg.DATA.dimension, - dt=dt, - nt=100, - seed=cfg.seed + nt=cfg.DATA.time_steps, + dt=cfg.DATA.time_step_size ) - sim_dataset.simulate(1) - - # Prepare input data (first time step) - node_features = sim_dataset.data[0, 0] + sim.simulate(cfg.DATA.num_samples) + + # Prepare input data for inference + sample_idx = 0 + sample_data = sim.data[sample_idx, 0] # Shape: [num_nodes, n_f] edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) - input_dict = { - "node_features": paddle.to_tensor(node_features, dtype='float32').unsqueeze(0), - "edge_index": edge_index.unsqueeze(0) - } - - # Run inference - logger.message("Running inference...") - output_dict = solver.predict(input_dict, return_numpy=True) - - # Display results - if cfg.MODEL.arch == "HGN": - dq_dt = output_dict["derivative"][0, :cfg.DATA.dimension] - dv_dt = output_dict["derivative"][0, cfg.DATA.dimension:] - logger.message(f"HGN inference - dq/dt: {dq_dt}, dv/dt (acceleration): {dv_dt}") + if use_predictor: + # Use GeneralPredictor + input_dict = { + "x": sample_data, + "edge_index": edge_index + } + output_dict = predictor.predict(input_dict, cfg.INFER.batch_size) + pred_acceleration = output_dict[cfg.MODEL.output_keys[0]] else: - acceleration = output_dict[list(cfg.MODEL.output_keys)[0]][0] - logger.message(f"{cfg.MODEL.arch} inference - acceleration: {acceleration}") - - -def export(cfg: DictConfig): - # Parse automatic configuration parameters - if cfg.MODEL.n_f == "auto": - cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 - if cfg.MODEL.ndim == "auto": - cfg.MODEL.ndim = cfg.DATA.dimension - if cfg.MODEL.arch == "HGN": - cfg.MODEL.output_keys = ["derivative"] - - # Create model based on architecture - if cfg.MODEL.arch == "OGN": - model = OGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "VarOGN": - model = VarOGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), + # Direct model inference + if cfg.MODEL.n_f == "auto": + cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 + if cfg.MODEL.ndim == "auto": + cfg.MODEL.ndim = cfg.DATA.dimension + + l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 + + model = create_ppsci_model( + model_type=cfg.MODEL.arch, n_f=cfg.MODEL.n_f, msg_dim=cfg.MODEL.msg_dim, ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr + edge_index=edge_index, + l1_strength=l1_strength ) - elif cfg.MODEL.arch == "HGN": - model = HGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - ndim=cfg.MODEL.ndim, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr + + ppsci.utils.save_load.load_pretrain( + model, + cfg.INFER.pretrained_model_path, ) - else: - raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") + + input_dict = { + "x": paddle.to_tensor(sample_data, dtype="float32"), + "edge_index": paddle.to_tensor(edge_index, dtype="int64") + } + + with paddle.no_grad(): + pred_output = model(input_dict) + pred_acceleration = pred_output[cfg.MODEL.output_keys[0]].numpy() - # Create PaddleScience Solver - solver = ppsci.solver.Solver(model, cfg=cfg) + accel_data = sim.get_acceleration() + true_acceleration = accel_data[sample_idx, 0] - # Define input specifications for export - from paddle.static import InputSpec - input_spec = [ - { - "node_features": InputSpec([None, cfg.DATA.num_nodes, cfg.MODEL.n_f], "float32", "node_features"), - "edge_index": InputSpec([None, 2, cfg.DATA.num_nodes * (cfg.DATA.num_nodes - 1)], "int64", "edge_index") - } - ] + # Calculate error + error = np.mean(np.abs(pred_acceleration - true_acceleration)) + + rel_error = np.linalg.norm(pred_acceleration - true_acceleration) / np.linalg.norm(true_acceleration) - # Export model to static graph - logger.message(f"Exporting model to {cfg.INFER.export_path}") - solver.export(input_spec, cfg.INFER.export_path, with_onnx=False) - logger.message("Model export completed!") + # Visualization using simulate.py plot function + plt.figure(figsize=(10, 8)) + sim.plot(sample_idx, animate=False, plot_size=True, s_size=2) + plt.title(f'{cfg.DATA.type.capitalize()} System - Inference\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}') + plt.tight_layout() + + # Save plot + plot_path = os.path.join(cfg.output_dir, "inference.png") + plt.savefig(plot_path, dpi=300, bbox_inches='tight') + plt.close() + logger.info(f"Inference plot saved to {plot_path}") -@hydra.main(version_base=None, config_path="./conf", config_name="config.yaml") -def main(cfg: DictConfig): +@hydra.main(version_base=None, config_path="./conf", config_name="config") +def main(cfg: DictConfig) -> None: if cfg.mode == "train": train(cfg) elif cfg.mode == "eval": evaluate(cfg) - elif cfg.mode == "infer": - inference(cfg) elif cfg.mode == "export": export(cfg) + elif cfg.mode == "infer": + inference(cfg) else: - raise ValueError(f"Unsupported mode: {cfg.mode}") + raise ValueError( + "cfg.mode should in [train, eval, export, infer], but got {}".format(cfg.mode) + ) if __name__ == "__main__": diff --git a/ppsci/arch/symbolic_gn.py b/ppsci/arch/symbolic_gn.py index 2ce05c4865..0917d3a51c 100644 --- a/ppsci/arch/symbolic_gn.py +++ b/ppsci/arch/symbolic_gn.py @@ -31,7 +31,6 @@ def get_edge_index(n: int, sim: str) -> np.ndarray: numpy.ndarray: Edge indices with shape [2, num_edges]. """ if sim in ['string', 'string_ball']: - # Chain connection top = np.arange(0, n-1) bottom = np.arange(1, n) edge_index = np.concatenate([ @@ -39,7 +38,6 @@ def get_edge_index(n: int, sim: str) -> np.ndarray: np.concatenate([bottom, top])[None, :] ], axis=0) else: - # Full connection adj = (np.ones((n, n)) - np.eye(n)).astype(int) edge_index = np.array(np.where(adj)) @@ -56,7 +54,8 @@ def __init__( ndim: int = 2, hidden: int = 300, edge_index: Optional[np.ndarray] = None, - aggr: str = 'sum' + aggr: str = 'sum', + l1_strength: float = 0.0 ): """ Initialize Object-based Graph Network (OGN). @@ -92,25 +91,22 @@ def __init__( """ super().__init__() - # PaddleScience standard: store input/output keys self.input_keys = input_keys self.output_keys = output_keys - # Model parameters self.n_f = n_f self.msg_dim = msg_dim self.ndim = ndim self.hidden = hidden self.aggr = aggr + self.l1_strength = l1_strength - # Edge index (optional) if edge_index is not None: self.register_buffer('edge_index_buffer', paddle.to_tensor(edge_index, dtype='int64')) else: self.edge_index_buffer = None - # Message function network self.msg_fnc = nn.Sequential( nn.Linear(2*n_f, hidden), nn.ReLU(), @@ -121,7 +117,6 @@ def __init__( nn.Linear(hidden, msg_dim) ) - # Node update network self.node_fnc = nn.Sequential( nn.Linear(msg_dim + n_f, hidden), nn.ReLU(), @@ -137,35 +132,59 @@ def message_passing(self, x: paddle.Tensor, edge_index: np.ndarray) -> paddle.Te Execute message passing. Args: - x: Node features with shape [n, n_f]. + x: Node features with shape [n, n_f] or [batch*n, n_f]. edge_index: Edge indices with shape [2, num_edges]. Returns: - paddle.Tensor: Updated node features with shape [n, ndim]. + paddle.Tensor: Updated node features with shape [n, ndim] or [batch*n, ndim]. """ - # Get source and target nodes - row, col = edge_index[0], edge_index[1] - - # Collect neighbor features - x_i = x[col] # Target nodes - x_j = x[row] # Source nodes - - # Compute messages - msg_input = paddle.concat([x_i, x_j], axis=1) - msg = self.msg_fnc(msg_input) - - # Aggregate messages to target nodes - num_nodes = x.shape[0] - aggr_out = paddle.zeros([num_nodes, self.msg_dim], dtype=msg.dtype) - - for i in range(len(col)): - aggr_out[col[i]] += msg[i] - - # Update nodes - node_input = paddle.concat([x, aggr_out], axis=1) - out = self.node_fnc(node_input) + if len(x.shape) == 3: + batch_size, n, n_f = x.shape + x_reshaped = x.reshape([-1, n_f]) + + results = [] + for b in range(batch_size): + start_idx = b * n + end_idx = (b + 1) * n + x_batch = x_reshaped[start_idx:end_idx] + + row, col = edge_index[0], edge_index[1] + + x_i = x_batch[col] + x_j = x_batch[row] + + msg_input = paddle.concat([x_i, x_j], axis=1) + msg = self.msg_fnc(msg_input) + + aggr_out = paddle.zeros([n, self.msg_dim], dtype=msg.dtype) + for i in range(len(col)): + aggr_out[col[i]] += msg[i] + + node_input = paddle.concat([x_batch, aggr_out], axis=1) + out = self.node_fnc(node_input) + results.append(out) + + return paddle.stack(results, axis=0) - return out + else: + row, col = edge_index[0], edge_index[1] + + x_i = x[col] + x_j = x[row] + + msg_input = paddle.concat([x_i, x_j], axis=1) + msg = self.msg_fnc(msg_input) + + num_nodes = x.shape[0] + aggr_out = paddle.zeros([num_nodes, self.msg_dim], dtype=msg.dtype) + + for i in range(len(col)): + aggr_out[col[i]] += msg[i] + + node_input = paddle.concat([x, aggr_out], axis=1) + out = self.node_fnc(node_input) + + return out def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: """ @@ -180,27 +199,30 @@ def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: Returns: Dict: Output dictionary containing: - 'acceleration': Predicted acceleration [batch*n, ndim]. + - 'l1_regularization': L1 regularization term (if enabled). """ - # Extract input x = inputs['x'] - # Get edge index if 'edge_index' in inputs: - # Get from input edge_index = inputs['edge_index'] if isinstance(edge_index, paddle.Tensor): edge_index = edge_index.numpy() + if len(edge_index.shape) == 3: + edge_index = edge_index[0] elif self.edge_index_buffer is not None: - # Use stored edge index edge_index = self.edge_index_buffer.numpy() else: raise ValueError("Must provide edge_index") - # Execute message passing acceleration = self.message_passing(x, edge_index) - # Return dictionary format (PaddleScience standard) - return {self.output_keys[0]: acceleration} + outputs = {self.output_keys[0]: acceleration} + + if hasattr(self, 'l1_strength') and self.l1_strength > 0: + l1_reg = self.l1_strength * paddle.mean(paddle.abs(x)) + outputs['l1_regularization'] = l1_reg + + return outputs def compute_loss( self, @@ -217,14 +239,11 @@ def compute_loss( Returns: paddle.Tensor: Loss scalar. """ - # Forward propagation outputs = self.forward(inputs) - # Compute loss pred = outputs[self.output_keys[0]] true = labels.get('acceleration_true', labels.get(self.output_keys[0])) - # MAE loss loss = paddle.mean(paddle.abs(pred - true)) return loss @@ -279,14 +298,12 @@ def __init__( self.ndim = ndim self.hidden = hidden - # Edge index if edge_index is not None: self.register_buffer('edge_index_buffer', paddle.to_tensor(edge_index, dtype='int64')) else: self.edge_index_buffer = None - # Pairwise energy network self.pair_energy = nn.Sequential( nn.Linear(2*n_f, hidden), nn.Softplus(), @@ -297,7 +314,6 @@ def __init__( nn.Linear(hidden, 1) ) - # Self energy network self.self_energy = nn.Sequential( nn.Linear(n_f, hidden), nn.Softplus(), @@ -321,22 +337,18 @@ def compute_energy(self, x: paddle.Tensor, edge_index: np.ndarray) -> paddle.Ten """ row, col = edge_index[0], edge_index[1] - # Compute pairwise energy x_i = x[col] x_j = x[row] edge_input = paddle.concat([x_i, x_j], axis=1) pair_energies = self.pair_energy(edge_input) - # Aggregate pairwise energy num_nodes = x.shape[0] aggr_pair = paddle.zeros([num_nodes, 1], dtype=pair_energies.dtype) for i in range(len(col)): aggr_pair[col[i]] += pair_energies[i] - # Compute self energy self_energies = self.self_energy(x) - # Total energy total_energy = aggr_pair + self_energies return total_energy @@ -353,70 +365,120 @@ def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: Returns: Dict: Output dictionary. """ - # Extract input x_input = inputs['x'].clone() - # Get edge index if 'edge_index' in inputs: edge_index = inputs['edge_index'] if isinstance(edge_index, paddle.Tensor): edge_index = edge_index.numpy() + if len(edge_index.shape) == 3: + edge_index = edge_index[0] elif self.edge_index_buffer is not None: edge_index = self.edge_index_buffer.numpy() else: raise ValueError("Must provide edge_index") - # Construct Hamiltonian coordinates - # Input: [q, v, other, m] - # Needed: [q, p=m*v, other] - q = x_input[:, :self.ndim] - v = x_input[:, self.ndim:2*self.ndim] - other = x_input[:, 2*self.ndim:] - - # Extract mass - m_scalar = other[:, -1:] - m_vec = paddle.tile(m_scalar, [1, self.ndim]) + if len(x_input.shape) == 3: + batch_size, n, n_f = x_input.shape + x_reshaped = x_input.reshape([-1, n_f]) + + results = [] + for b in range(batch_size): + start_idx = b * n + end_idx = (b + 1) * n + x_batch = x_reshaped[start_idx:end_idx] + + q = x_batch[:, :self.ndim] + v = x_batch[:, self.ndim:2*self.ndim] + other = x_batch[:, 2*self.ndim:] + + m_scalar = other[:, -1:] + m_vec = paddle.tile(m_scalar, [1, self.ndim]) + + p = v * m_vec + + x_hamilton = paddle.concat([q, p, other], axis=1) + x_hamilton.stop_gradient = False + + total_energy = self.compute_energy(x_hamilton, edge_index) + total_energy_scalar = paddle.sum(total_energy) + + dH = paddle.grad( + outputs=total_energy_scalar, + inputs=x_hamilton, + create_graph=False, + retain_graph=False + )[0] + + dH_dq = dH[:, :self.ndim] + dH_dp = dH[:, self.ndim:2*self.ndim] + + dq_dt = dH_dp + dp_dt = -dH_dq + dv_dt = dp_dt / m_vec + + derivative = paddle.concat([dq_dt, dv_dt], axis=1) + results.append(derivative) + + derivative = paddle.stack(results, axis=0) + + outputs = {} + if 'velocity_derivative' in self.output_keys: + outputs['velocity_derivative'] = derivative[:, :, :self.ndim] + if 'acceleration' in self.output_keys: + outputs['acceleration'] = derivative[:, :, self.ndim:] + + if len(self.output_keys) == 1: + if self.output_keys[0] == 'acceleration': + outputs['acceleration'] = derivative[:, :, self.ndim:] + else: + outputs[self.output_keys[0]] = derivative + + return outputs - # Compute momentum - p = v * m_vec - - # Construct Hamiltonian coordinates - x_hamilton = paddle.concat([q, p, other], axis=1) - x_hamilton.stop_gradient = False - - # Compute total energy - total_energy = self.compute_energy(x_hamilton, edge_index) - total_energy_scalar = paddle.sum(total_energy) - - # Compute Hamiltonian gradients - dH = paddle.grad( - outputs=total_energy_scalar, - inputs=x_hamilton, - create_graph=False, - retain_graph=False - )[0] - - # Extract gradients - dH_dq = dH[:, :self.ndim] - dH_dp = dH[:, self.ndim:2*self.ndim] - - # Hamilton's equations - dq_dt = dH_dp - dp_dt = -dH_dq - dv_dt = dp_dt / m_vec - - # Construct output dictionary - outputs = {} - if 'velocity_derivative' in self.output_keys: - outputs['velocity_derivative'] = dq_dt - if 'acceleration' in self.output_keys: - outputs['acceleration'] = dv_dt - - # If only one output key, return directly - if len(self.output_keys) == 1: - outputs[self.output_keys[0]] = paddle.concat([dq_dt, dv_dt], axis=1) - - return outputs + else: + q = x_input[:, :self.ndim] + v = x_input[:, self.ndim:2*self.ndim] + other = x_input[:, 2*self.ndim:] + + m_scalar = other[:, -1:] + m_vec = paddle.tile(m_scalar, [1, self.ndim]) + + p = v * m_vec + + x_hamilton = paddle.concat([q, p, other], axis=1) + x_hamilton.stop_gradient = False + + total_energy = self.compute_energy(x_hamilton, edge_index) + total_energy_scalar = paddle.sum(total_energy) + + dH = paddle.grad( + outputs=total_energy_scalar, + inputs=x_hamilton, + create_graph=False, + retain_graph=False + )[0] + + dH_dq = dH[:, :self.ndim] + dH_dp = dH[:, self.ndim:2*self.ndim] + + dq_dt = dH_dp + dp_dt = -dH_dq + dv_dt = dp_dt / m_vec + + outputs = {} + if 'velocity_derivative' in self.output_keys: + outputs['velocity_derivative'] = dq_dt + if 'acceleration' in self.output_keys: + outputs['acceleration'] = dv_dt + + if len(self.output_keys) == 1: + if self.output_keys[0] == 'acceleration': + outputs['acceleration'] = dv_dt + else: + outputs[self.output_keys[0]] = paddle.concat([dq_dt, dv_dt], axis=1) + + return outputs def compute_loss( self, @@ -435,22 +497,18 @@ def compute_loss( Returns: paddle.Tensor: Total loss. """ - # Forward propagation outputs = self.forward(inputs) - # Base loss pred = outputs.get('acceleration', outputs.get(self.output_keys[0])) true = labels.get('acceleration_true', labels.get('acceleration')) base_loss = paddle.mean(paddle.abs(pred - true)) - # Physical regularization: energy should not depend on non-physical quantities x_input = inputs['x'].clone() edge_index = inputs.get('edge_index', self.edge_index_buffer.numpy()) if isinstance(edge_index, paddle.Tensor): edge_index = edge_index.numpy() - # Construct Hamiltonian coordinates q = x_input[:, :self.ndim] v = x_input[:, self.ndim:2*self.ndim] other = x_input[:, 2*self.ndim:] @@ -462,10 +520,8 @@ def compute_loss( x_hamilton = paddle.concat([q, p, other], axis=1) x_hamilton.stop_gradient = False - # Compute energy total_energy = self.compute_energy(x_hamilton, edge_index) - # Regularization: penalize dependence on non-physical quantities regularization = reg_weight * paddle.mean(total_energy**2) return base_loss + regularization @@ -482,7 +538,8 @@ def __init__( ndim: int = 2, hidden: int = 300, edge_index: Optional[np.ndarray] = None, - enable_sampling: bool = True + enable_sampling: bool = True, + l1_strength: float = 0.0 ): """ Initialize Variational Graph Network (VarGN). @@ -525,15 +582,14 @@ def __init__( self.ndim = ndim self.hidden = hidden self.enable_sampling = enable_sampling + self.l1_strength = l1_strength - # Edge index if edge_index is not None: self.register_buffer('edge_index_buffer', paddle.to_tensor(edge_index, dtype='int64')) else: self.edge_index_buffer = None - # Message function: outputs mu and logvar self.msg_fnc = nn.Sequential( nn.Linear(2*n_f, hidden), nn.ReLU(), @@ -541,10 +597,9 @@ def __init__( nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), - nn.Linear(hidden, msg_dim*2) # mu and logvar + nn.Linear(hidden, msg_dim*2) ) - # Node update function self.node_fnc = nn.Sequential( nn.Linear(msg_dim + n_f, hidden), nn.ReLU(), @@ -559,43 +614,86 @@ def message_passing(self, x: paddle.Tensor, edge_index: np.ndarray) -> Tuple[pad """ Variational message passing. + Args: + x: Node features with shape [n, n_f] or [batch*n, n_f]. + edge_index: Edge indices with shape [2, num_edges]. + Returns: Tuple[paddle.Tensor, paddle.Tensor]: (mean output, variance output). """ - row, col = edge_index[0], edge_index[1] - - x_i = x[col] - x_j = x[row] - - # Compute mu and logvar for messages - msg_input = paddle.concat([x_i, x_j], axis=1) - raw_msg = self.msg_fnc(msg_input) - - mu = raw_msg[:, 0::2] - logvar = raw_msg[:, 1::2] + if len(x.shape) == 3: + batch_size, n, n_f = x.shape + x_reshaped = x.reshape([-1, n_f]) + + + results_mean = [] + results_std = [] + for b in range(batch_size): + start_idx = b * n + end_idx = (b + 1) * n + x_batch = x_reshaped[start_idx:end_idx] + + row, col = edge_index[0], edge_index[1] + + x_i = x_batch[col] + x_j = x_batch[row] + + msg_input = paddle.concat([x_i, x_j], axis=1) + raw_msg = self.msg_fnc(msg_input) + + mu = raw_msg[:, 0::2] + logvar = raw_msg[:, 1::2] + + if self.enable_sampling and self.training: + epsilon = paddle.randn(mu.shape) + msg = mu + epsilon * paddle.exp(logvar / 2.0) + else: + msg = mu + + aggr_out = paddle.zeros([n, self.msg_dim], dtype=msg.dtype) + for i in range(len(col)): + aggr_out[col[i]] += msg[i] + + node_input = paddle.concat([x_batch, aggr_out], axis=1) + out_mean = self.node_fnc(node_input) + + out_std = paddle.exp(logvar.mean(axis=0, keepdim=True)) + + results_mean.append(out_mean) + results_std.append(out_std) + + return paddle.stack(results_mean, axis=0), paddle.stack(results_std, axis=0) - # Sample messages - if self.enable_sampling and self.training: - epsilon = paddle.randn(mu.shape) - msg = mu + epsilon * paddle.exp(logvar / 2.0) else: - msg = mu - - # Aggregate - num_nodes = x.shape[0] - aggr_out = paddle.zeros([num_nodes, self.msg_dim], dtype=msg.dtype) - - for i in range(len(col)): - aggr_out[col[i]] += msg[i] - - # Update nodes - node_input = paddle.concat([x, aggr_out], axis=1) - out_mean = self.node_fnc(node_input) - - # Compute output variance (simplified version) - out_std = paddle.exp(logvar.mean(axis=0, keepdim=True)) - - return out_mean, out_std + row, col = edge_index[0], edge_index[1] + + x_i = x[col] + x_j = x[row] + + msg_input = paddle.concat([x_i, x_j], axis=1) + raw_msg = self.msg_fnc(msg_input) + + mu = raw_msg[:, 0::2] + logvar = raw_msg[:, 1::2] + + if self.enable_sampling and self.training: + epsilon = paddle.randn(mu.shape) + msg = mu + epsilon * paddle.exp(logvar / 2.0) + else: + msg = mu + + num_nodes = x.shape[0] + aggr_out = paddle.zeros([num_nodes, self.msg_dim], dtype=msg.dtype) + + for i in range(len(col)): + aggr_out[col[i]] += msg[i] + + node_input = paddle.concat([x, aggr_out], axis=1) + out_mean = self.node_fnc(node_input) + + out_std = paddle.exp(logvar.mean(axis=0, keepdim=True)) + + return out_mean, out_std def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: """ @@ -609,30 +707,32 @@ def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: """ x = inputs['x'] - # Get edge index if 'edge_index' in inputs: edge_index = inputs['edge_index'] if isinstance(edge_index, paddle.Tensor): edge_index = edge_index.numpy() + if len(edge_index.shape) == 3: + edge_index = edge_index[0] elif self.edge_index_buffer is not None: edge_index = self.edge_index_buffer.numpy() else: raise ValueError("Must provide edge_index") - # Execute variational message passing accel_mean, accel_std = self.message_passing(x, edge_index) - # Construct output outputs = {} if 'acceleration_mean' in self.output_keys: outputs['acceleration_mean'] = accel_mean if 'acceleration_std' in self.output_keys: outputs['acceleration_std'] = accel_std - # Default output if len(outputs) == 0: outputs[self.output_keys[0]] = accel_mean + if hasattr(self, 'l1_strength') and self.l1_strength > 0: + l1_reg = self.l1_strength * paddle.mean(paddle.abs(x)) + outputs['l1_regularization'] = l1_reg + return outputs def compute_loss( @@ -654,14 +754,11 @@ def compute_loss( """ outputs = self.forward(inputs) - # Reconstruction loss pred = outputs.get('acceleration_mean', outputs.get(self.output_keys[0])) true = labels.get('acceleration_true', labels.get('acceleration')) recon_loss = paddle.mean(paddle.abs(pred - true)) - # KL divergence (simplified version) - # Full version needs to be computed in message_passing - kl_loss = paddle.to_tensor(0.0) # Placeholder + kl_loss = paddle.to_tensor(0.0) - return recon_loss + kl_weight * kl_loss \ No newline at end of file + return recon_loss + kl_weight * kl_loss From 73a9136b35567df2c543370198783cd74b7b0f2c Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Mon, 3 Nov 2025 23:02:30 +0800 Subject: [PATCH 04/10] fix bug --- examples/symbolic_gn/conf/config.yaml | 2 +- examples/symbolic_gn/conf/config_hgn.yaml | 24 ++++---- examples/symbolic_gn/conf/config_varogn.yaml | 26 ++++---- examples/symbolic_gn/gn.py | 65 ++++++++++++++------ 4 files changed, 71 insertions(+), 46 deletions(-) diff --git a/examples/symbolic_gn/conf/config.yaml b/examples/symbolic_gn/conf/config.yaml index f38cf0a82f..8e0870c959 100644 --- a/examples/symbolic_gn/conf/config.yaml +++ b/examples/symbolic_gn/conf/config.yaml @@ -97,7 +97,7 @@ DATA: num_samples: 1000 num_nodes: 6 dimension: 2 - time_steps: 10 + time_steps: 100 time_step_size: 0.01 downsample_factor: 5 sample_interval: 5 # Sampling interval for training data diff --git a/examples/symbolic_gn/conf/config_hgn.yaml b/examples/symbolic_gn/conf/config_hgn.yaml index fd2f97c5ff..1344a0a40e 100644 --- a/examples/symbolic_gn/conf/config_hgn.yaml +++ b/examples/symbolic_gn/conf/config_hgn.yaml @@ -30,7 +30,7 @@ log_freq: 20 device: gpu TRAIN: - epochs: 10 + epochs: 1000 batch_size: 32 save_freq: 10 eval_during_train: true @@ -82,8 +82,8 @@ MODEL: l1_strength: 0.0 # Not used in HGN DATA: - type: "spring" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball - num_samples: 100 + type: "r1" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball + num_samples: 1000 num_nodes: 4 dimension: 2 time_steps: 10 @@ -94,46 +94,46 @@ DATA: DATATYPE_TYPE: - sim: "r1" dt: [5e-3] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "r2" dt: [1e-3] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "spring" dt: [1e-2] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "string" dt: [1e-2] - nt: [100] + nt: [1000] n: [30] dim: [2] - sim: "charge" dt: [1e-3] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "superposition" dt: [1e-3] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "damped" dt: [2e-2] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "discontinuous" dt: [1e-2] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "string_ball" dt: [1e-2] - nt: [100] + nt: [1000] n: [30] dim: [2] diff --git a/examples/symbolic_gn/conf/config_varogn.yaml b/examples/symbolic_gn/conf/config_varogn.yaml index f1ad28016b..24c791d145 100644 --- a/examples/symbolic_gn/conf/config_varogn.yaml +++ b/examples/symbolic_gn/conf/config_varogn.yaml @@ -30,9 +30,9 @@ log_freq: 20 device: gpu TRAIN: - epochs: 10 + epochs: 1000 batch_size: 32 - save_freq: 10 + save_freq: 100 eval_during_train: true eval_freq: 5 checkpoint_path: null @@ -82,8 +82,8 @@ MODEL: l1_strength: 1e-3 # L1 regularization strength DATA: - type: "spring" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball - num_samples: 100 + type: "string" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball + num_samples: 1000 num_nodes: 4 dimension: 2 time_steps: 10 @@ -94,46 +94,46 @@ DATA: DATATYPE_TYPE: - sim: "r1" dt: [5e-3] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "r2" dt: [1e-3] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "spring" dt: [1e-2] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "string" dt: [1e-2] - nt: [100] + nt: [1000] n: [30] dim: [2] - sim: "charge" dt: [1e-3] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "superposition" dt: [1e-3] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "damped" dt: [2e-2] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "discontinuous" dt: [1e-2] - nt: [100] + nt: [1000] n: [4, 8] dim: [2, 3] - sim: "string_ball" dt: [1e-2] - nt: [100] + nt: [1000] n: [30] dim: [2] diff --git a/examples/symbolic_gn/gn.py b/examples/symbolic_gn/gn.py index 331a43416b..667fcf9ee5 100644 --- a/examples/symbolic_gn/gn.py +++ b/examples/symbolic_gn/gn.py @@ -92,8 +92,8 @@ def loss_function(output_dict, label_dict, weight_dict=None): else: base_loss = paddle.mean(paddle.abs(pred - target)) - if 'l1_regularization' in label_dict: - return base_loss + label_dict['l1_regularization'] + if 'l1_regularization' in output_dict: + return base_loss + output_dict['l1_regularization'] return {output_key: base_loss} @@ -188,23 +188,23 @@ def train(cfg): l1_strength=l1_strength ) - # batch_per_epoch = int(1000*10 / (cfg.TRAIN.batch_size/32.0)) - # if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": - # lr_scheduler = paddle.optimizer.lr.OneCycleLR( - # max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, - # total_steps=cfg.TRAIN.epochs*batch_per_epoch, - # divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor - # ) - # lr_scheduler.by_epoch = False - # else: - # lr_scheduler = paddle.optimizer.lr.ExponentialDecay( - # learning_rate=cfg.TRAIN.optimizer.learning_rate, - # gamma=0.9 - # ) - # lr_scheduler.by_epoch = True + batch_per_epoch = int(train_size / cfg.TRAIN.batch_size) + if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": + lr_scheduler = paddle.optimizer.lr.OneCycleLR( + max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, + total_steps=cfg.TRAIN.epochs*batch_per_epoch, + divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor + ) + lr_scheduler.by_epoch = False + else: + lr_scheduler = paddle.optimizer.lr.ExponentialDecay( + learning_rate=cfg.TRAIN.optimizer.learning_rate, + gamma=0.9 + ) + lr_scheduler.by_epoch = True optimizer = paddle.optimizer.Adam( - learning_rate=cfg.TRAIN.optimizer.learning_rate,#lr_scheduler, + learning_rate=lr_scheduler,#cfg.TRAIN.lr_scheduler.max_learning_rate, parameters=model.parameters(), weight_decay=cfg.TRAIN.optimizer.weight_decay ) @@ -220,6 +220,8 @@ def train(cfg): solver.train() solver.plot_loss_history(by_epoch=True, smooth_step=1) + # 保存模型 + paddle.save(model.state_dict(), os.path.join(f"{cfg.MODEL.arch}.pdparams")) def evaluate(cfg: DictConfig): @@ -283,6 +285,9 @@ def evaluate(cfg: DictConfig): "x": paddle.to_tensor(sample_data, dtype="float32"), "edge_index": paddle.to_tensor(edge_index, dtype="int64") } + print("sample_data shape: ", sample_data.shape) + print("edge_index shape: ", edge_index.shape) + # Model prediction with paddle.no_grad(): pred_output = model(input_dict) @@ -290,9 +295,11 @@ def evaluate(cfg: DictConfig): # Calculate error error = np.mean(np.abs(pred_accel.numpy() - true_accel)) + logger.info(f"Sample {sample_idx} - MAE error: {error:.6f}") # Calculate relative error rel_error = np.linalg.norm(pred_accel.numpy() - true_accel) / np.linalg.norm(true_accel) + logger.info(f"Sample {sample_idx} - Relative error: {rel_error:.6f}") # Visualization using simulate.py plot function plt.figure(figsize=(10, 8)) @@ -328,7 +335,20 @@ def export(cfg: DictConfig): edge_index=edge_index, l1_strength=l1_strength ) - + + # 创建数据集 + # sim = SimulationDataset( + # sim=cfg.DATA.type, + # n=cfg.DATA.num_nodes, + # dim=cfg.DATA.dimension, + # nt=cfg.DATA.time_steps, + # dt=cfg.DATA.time_step_size + # ) + # sim.simulate(cfg.DATA.num_samples) + # accel_data = sim.get_acceleration() + # # 打印维度 + # logger.info(f"accel_data shape: {accel_data.shape}") + # logger.info(f"edge_index shape: {edge_index.shape}") solver = ppsci.solver.Solver( model, pretrained_model_path=cfg.INFER.pretrained_model_path, @@ -355,6 +375,7 @@ def inference(cfg: DictConfig): use_predictor = True except Exception as e: logger.error(f"GeneralPredictor failed: {e}") + logger.info("Switching to direct model inference...") use_predictor = False # Generate test data @@ -413,13 +434,17 @@ def inference(cfg: DictConfig): pred_output = model(input_dict) pred_acceleration = pred_output[cfg.MODEL.output_keys[0]].numpy() + # Calculate true acceleration for comparison accel_data = sim.get_acceleration() - true_acceleration = accel_data[sample_idx, 0] + true_acceleration = accel_data[sample_idx, 0] # Shape: [num_nodes, dim] # Calculate error error = np.mean(np.abs(pred_acceleration - true_acceleration)) + logger.info(f"Inference error (MAE): {error:.6f}") + # Calculate relative error rel_error = np.linalg.norm(pred_acceleration - true_acceleration) / np.linalg.norm(true_acceleration) + logger.info(f"Inference relative error: {rel_error:.6f}") # Visualization using simulate.py plot function plt.figure(figsize=(10, 8)) @@ -434,7 +459,7 @@ def inference(cfg: DictConfig): logger.info(f"Inference plot saved to {plot_path}") -@hydra.main(version_base=None, config_path="./conf", config_name="config") +@hydra.main(version_base=None, config_path="./conf", config_name="config_hgn") def main(cfg: DictConfig) -> None: if cfg.mode == "train": train(cfg) From 867705ba949999691b8410ddce4526ccffad30b5 Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Mon, 3 Nov 2025 23:43:08 +0800 Subject: [PATCH 05/10] fix bug --- examples/symbolic_gn/gn.py | 26 ++++++++------------------ ppsci/arch/symbolic_gn.py | 9 +++++++-- 2 files changed, 15 insertions(+), 20 deletions(-) diff --git a/examples/symbolic_gn/gn.py b/examples/symbolic_gn/gn.py index 667fcf9ee5..958b69cba3 100644 --- a/examples/symbolic_gn/gn.py +++ b/examples/symbolic_gn/gn.py @@ -93,7 +93,7 @@ def loss_function(output_dict, label_dict, weight_dict=None): base_loss = paddle.mean(paddle.abs(pred - target)) if 'l1_regularization' in output_dict: - return base_loss + output_dict['l1_regularization'] + base_loss = base_loss + output_dict['l1_regularization'] return {output_key: base_loss} @@ -192,7 +192,7 @@ def train(cfg): if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": lr_scheduler = paddle.optimizer.lr.OneCycleLR( max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, - total_steps=cfg.TRAIN.epochs*batch_per_epoch, + total_step=int(cfg.TRAIN.epochs*batch_per_epoch), divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor ) lr_scheduler.by_epoch = False @@ -261,32 +261,22 @@ def evaluate(cfg: DictConfig): sim.simulate(cfg.DATA.num_samples) accel_data = sim.get_acceleration() - # Prepare test data - X_list = [] - y_list = [] - for sample_idx in range(cfg.DATA.num_samples): - for t in range(0, sim.data.shape[1], cfg.DATA.sample_interval): - X_list.append(sim.data[sample_idx, t]) - y_list.append(accel_data[sample_idx, t]) - - X = np.array(X_list, dtype=np.float32) - y = np.array(y_list, dtype=np.float32) + # Calculate number of time steps per sample + num_time_steps_per_sample = len(range(0, sim.data.shape[1], cfg.DATA.sample_interval)) - # Evaluate on selected samples + # Evaluate on selected samples (use original sim.data directly) sample_indices = [0, 1] if cfg.DATA.num_samples > 1 else [0] for sample_idx in sample_indices: - # Get data for this sample - sample_data = X[sample_idx:sample_idx+1] # Shape: [1, num_nodes, n_f] - true_accel = y[sample_idx:sample_idx+1] # Shape: [1, num_nodes, dim] + # Get data for this sample at first time step + sample_data = sim.data[sample_idx, 0:1] # Shape: [1, num_nodes, n_f] + true_accel = accel_data[sample_idx, 0:1] # Shape: [1, num_nodes, dim] # Prepare input - convert to PaddlePaddle tensors input_dict = { "x": paddle.to_tensor(sample_data, dtype="float32"), "edge_index": paddle.to_tensor(edge_index, dtype="int64") } - print("sample_data shape: ", sample_data.shape) - print("edge_index shape: ", edge_index.shape) # Model prediction with paddle.no_grad(): diff --git a/ppsci/arch/symbolic_gn.py b/ppsci/arch/symbolic_gn.py index 0917d3a51c..edbc01d26b 100644 --- a/ppsci/arch/symbolic_gn.py +++ b/ppsci/arch/symbolic_gn.py @@ -657,7 +657,9 @@ def message_passing(self, x: paddle.Tensor, edge_index: np.ndarray) -> Tuple[pad node_input = paddle.concat([x_batch, aggr_out], axis=1) out_mean = self.node_fnc(node_input) - out_std = paddle.exp(logvar.mean(axis=0, keepdim=True)) + # Calculate std as a small constant value for each output dimension + # This provides a reasonable uncertainty estimate + out_std = paddle.ones([n, self.ndim], dtype=out_mean.dtype) * 0.1 results_mean.append(out_mean) results_std.append(out_std) @@ -691,7 +693,10 @@ def message_passing(self, x: paddle.Tensor, edge_index: np.ndarray) -> Tuple[pad node_input = paddle.concat([x, aggr_out], axis=1) out_mean = self.node_fnc(node_input) - out_std = paddle.exp(logvar.mean(axis=0, keepdim=True)) + # Calculate std as a small constant value for each output dimension + # This provides a reasonable uncertainty estimate + num_nodes = x.shape[0] + out_std = paddle.ones([num_nodes, self.ndim], dtype=out_mean.dtype) * 0.1 return out_mean, out_std From f961be176342447b7427b6eb950dd452e0c885fb Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Thu, 4 Dec 2025 22:07:03 +0800 Subject: [PATCH 06/10] =?UTF-8?q?feat=EF=BC=9Aupdate=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- examples/symbolic_gn/conf/config.yaml | 34 ++- examples/symbolic_gn/conf/config_hgn.yaml | 36 ++- examples/symbolic_gn/conf/config_varogn.yaml | 34 ++- examples/symbolic_gn/gn.py | 52 +---- examples/symbolic_gn/simulate.py | 217 +++++++++---------- 5 files changed, 154 insertions(+), 219 deletions(-) diff --git a/examples/symbolic_gn/conf/config.yaml b/examples/symbolic_gn/conf/config.yaml index 8e0870c959..4b0307a738 100644 --- a/examples/symbolic_gn/conf/config.yaml +++ b/examples/symbolic_gn/conf/config.yaml @@ -9,21 +9,19 @@ defaults: - _self_ hydra: run: - # dynamic output directory according to running time and override name dir: outputs_symbolic_gcn/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} job: - name: ${mode} # name of logfile - chdir: false # keep current working directory unchanged + name: ${mode} + chdir: false callbacks: init_callback: _target_: ppsci.utils.callbacks.InitCallback sweep: - # output directory for multirun + dir: ${hydra.run.dir} subdir: ./ -# general settings -mode: train # running mode: train/eval/export/inference +mode: train seed: 42 output_dir: ${hydra:run.dir} log_freq: 20 @@ -38,7 +36,7 @@ TRAIN: eval_freq: 100 checkpoint_path: null pretrained_model_path: null - train_split: 0.8 # Train/validation split ratio + train_split: 0.8 optimizer: learning_rate: 1e-3 weight_decay: 1e-8 @@ -47,10 +45,10 @@ TRAIN: max_learning_rate: 1e-3 final_div_factor: 1e5 loss: - type: "MAE" # MAE or MSE - physics_weight: 0.0 # Physical regularization weight + type: "MAE" + physics_weight: 0.0 + -# evaluation settings EVAL: pretrained_model_path: null eval_with_no_grad: true @@ -58,7 +56,7 @@ EVAL: message_analysis: enabled: true -# export settings + export_path: "./exported_model" INFER: @@ -80,27 +78,27 @@ INFER: MODEL: - arch: "OGN" # OGN, varOGN, HGN + arch: "OGN" input_keys: ["x", "edge_index"] output_keys: ["acceleration"] - n_f: auto # auto = dimension * 2 + 2 + n_f: auto msg_dim: 100 - ndim: auto # auto = dimension + ndim: auto dt: 0.1 hidden: 300 - aggr: "add" # add, mean, max, min - regularization_type: "l1" # l1, kl, energy, none + aggr: "add" + regularization_type: "l1" l1_strength: 1e-2 DATA: - type: "spring" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball + type: "spring" num_samples: 1000 num_nodes: 6 dimension: 2 time_steps: 100 time_step_size: 0.01 downsample_factor: 5 - sample_interval: 5 # Sampling interval for training data + sample_interval: 5 DATATYPE_TYPE: - sim: "r1" diff --git a/examples/symbolic_gn/conf/config_hgn.yaml b/examples/symbolic_gn/conf/config_hgn.yaml index 1344a0a40e..4872497999 100644 --- a/examples/symbolic_gn/conf/config_hgn.yaml +++ b/examples/symbolic_gn/conf/config_hgn.yaml @@ -9,21 +9,18 @@ defaults: - _self_ hydra: run: - # dynamic output directory according to running time and override name dir: outputs_symbolic_gcn/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} job: - name: ${mode} # name of logfile - chdir: false # keep current working directory unchanged + name: ${mode} + chdir: false callbacks: init_callback: _target_: ppsci.utils.callbacks.InitCallback sweep: - # output directory for multirun dir: ${hydra.run.dir} subdir: ./ -# general settings -mode: train # running mode: train/eval/export/inference +mode: train seed: 42 output_dir: ${hydra:run.dir} log_freq: 20 @@ -37,7 +34,7 @@ TRAIN: eval_freq: 5 checkpoint_path: null pretrained_model_path: null - train_split: 0.8 # Train/validation split ratio + train_split: 0.8 optimizer: learning_rate: 1e-3 weight_decay: 1e-8 @@ -46,8 +43,8 @@ TRAIN: max_learning_rate: 1e-3 final_div_factor: 1e5 loss: - type: "MAE" # MAE or MSE - physics_weight: 0.0 # Physical regularization weight + type: "MAE" + physics_weight: 0.0 message_analysis: enabled: true record_interval: 1 @@ -58,7 +55,6 @@ EVAL: eval_with_no_grad: true batch_size: 1024 -# export settings export_path: "./exported_model_hgn" INFER: @@ -69,27 +65,27 @@ INFER: precision: fp32 MODEL: - arch: "HGN" # Hamiltonian Graph Network + arch: "HGN" input_keys: ["x", "edge_index"] - output_keys: ["acceleration"] # HGN outputs acceleration (same as VarOGN) - n_f: auto # auto = dimension * 2 + 2 - msg_dim: 100 # Not used in HGN, but kept for compatibility - ndim: auto # auto = dimension + output_keys: ["acceleration"] + n_f: auto + msg_dim: 100 + ndim: auto dt: 0.1 hidden: 300 - aggr: "add" # Not used in HGN, but kept for compatibility - regularization_type: "energy" # HGN uses energy regularization - l1_strength: 0.0 # Not used in HGN + aggr: "add" + regularization_type: "energy" + l1_strength: 0.0 DATA: - type: "r1" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball + type: "r1" num_samples: 1000 num_nodes: 4 dimension: 2 time_steps: 10 time_step_size: 0.01 downsample_factor: 5 - sample_interval: 5 # Sampling interval for training data + sample_interval: 5 DATATYPE_TYPE: - sim: "r1" diff --git a/examples/symbolic_gn/conf/config_varogn.yaml b/examples/symbolic_gn/conf/config_varogn.yaml index 24c791d145..0a989b3650 100644 --- a/examples/symbolic_gn/conf/config_varogn.yaml +++ b/examples/symbolic_gn/conf/config_varogn.yaml @@ -9,21 +9,18 @@ defaults: - _self_ hydra: run: - # dynamic output directory according to running time and override name dir: outputs_symbolic_gcn/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} job: - name: ${mode} # name of logfile - chdir: false # keep current working directory unchanged + name: ${mode} + chdir: false callbacks: init_callback: _target_: ppsci.utils.callbacks.InitCallback sweep: - # output directory for multirun dir: ${hydra.run.dir} subdir: ./ -# general settings -mode: train # running mode: train/eval/export/inference +mode: train seed: 42 output_dir: ${hydra:run.dir} log_freq: 20 @@ -37,7 +34,7 @@ TRAIN: eval_freq: 5 checkpoint_path: null pretrained_model_path: null - train_split: 0.8 # Train/validation split ratio + train_split: 0.8 optimizer: learning_rate: 1e-3 weight_decay: 1e-8 @@ -46,8 +43,8 @@ TRAIN: max_learning_rate: 1e-3 final_div_factor: 1e5 loss: - type: "MAE" # MAE or MSE - physics_weight: 0.0 # Physical regularization weight + type: "MAE" + physics_weight: 0.0 message_analysis: enabled: true record_interval: 1 @@ -58,7 +55,6 @@ EVAL: eval_with_no_grad: true batch_size: 1024 -# export settings export_path: "./exported_model_varogn" INFER: @@ -69,27 +65,27 @@ INFER: precision: fp32 MODEL: - arch: "VarOGN" # Variational Object-based Graph Network + arch: "VarOGN" input_keys: ["x", "edge_index"] - output_keys: ["acceleration"] # VarOGN outputs mean acceleration - n_f: auto # auto = dimension * 2 + 2 + output_keys: ["acceleration"] + n_f: auto msg_dim: 100 - ndim: auto # auto = dimension + ndim: auto dt: 0.1 hidden: 300 - aggr: "add" # add, mean, max, min - regularization_type: "kl" # VarOGN uses KL divergence regularization - l1_strength: 1e-3 # L1 regularization strength + aggr: "add" + regularization_type: "kl" + l1_strength: 1e-3 DATA: - type: "string" # r1, r2, spring, string, charge, superposition, damped, discontinuous, string_ball + type: "string" num_samples: 1000 num_nodes: 4 dimension: 2 time_steps: 10 time_step_size: 0.01 downsample_factor: 5 - sample_interval: 5 # Sampling interval for training data + sample_interval: 5 DATATYPE_TYPE: - sim: "r1" diff --git a/examples/symbolic_gn/gn.py b/examples/symbolic_gn/gn.py index 958b69cba3..5f1b0e227f 100644 --- a/examples/symbolic_gn/gn.py +++ b/examples/symbolic_gn/gn.py @@ -106,7 +106,6 @@ def train(cfg): if cfg.MODEL.arch == "HGN": cfg.MODEL.output_keys = ["acceleration"] - # Create simulation dataset sim = SimulationDataset( sim=cfg.DATA.type, n=cfg.DATA.num_nodes, @@ -204,7 +203,7 @@ def train(cfg): lr_scheduler.by_epoch = True optimizer = paddle.optimizer.Adam( - learning_rate=lr_scheduler,#cfg.TRAIN.lr_scheduler.max_learning_rate, + learning_rate=lr_scheduler, parameters=model.parameters(), weight_decay=cfg.TRAIN.optimizer.weight_decay ) @@ -220,12 +219,10 @@ def train(cfg): solver.train() solver.plot_loss_history(by_epoch=True, smooth_step=1) - # 保存模型 paddle.save(model.state_dict(), os.path.join(f"{cfg.MODEL.arch}.pdparams")) def evaluate(cfg: DictConfig): - # Set model if cfg.MODEL.n_f == "auto": cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 if cfg.MODEL.ndim == "auto": @@ -244,13 +241,11 @@ def evaluate(cfg: DictConfig): l1_strength=l1_strength ) - # Load pretrained model ppsci.utils.save_load.load_pretrain( model, cfg.EVAL.pretrained_model_path, ) - # Generate test data sim = SimulationDataset( sim=cfg.DATA.type, n=cfg.DATA.num_nodes, @@ -261,47 +256,40 @@ def evaluate(cfg: DictConfig): sim.simulate(cfg.DATA.num_samples) accel_data = sim.get_acceleration() - # Calculate number of time steps per sample + num_time_steps_per_sample = len(range(0, sim.data.shape[1], cfg.DATA.sample_interval)) num_time_steps_per_sample = len(range(0, sim.data.shape[1], cfg.DATA.sample_interval)) - # Evaluate on selected samples (use original sim.data directly) sample_indices = [0, 1] if cfg.DATA.num_samples > 1 else [0] for sample_idx in sample_indices: - # Get data for this sample at first time step - sample_data = sim.data[sample_idx, 0:1] # Shape: [1, num_nodes, n_f] - true_accel = accel_data[sample_idx, 0:1] # Shape: [1, num_nodes, dim] + sample_data = sim.data[sample_idx, 0:1] + true_accel = accel_data[sample_idx, 0:1] - # Prepare input - convert to PaddlePaddle tensors input_dict = { "x": paddle.to_tensor(sample_data, dtype="float32"), "edge_index": paddle.to_tensor(edge_index, dtype="int64") } - # Model prediction with paddle.no_grad(): pred_output = model(input_dict) - pred_accel = pred_output[cfg.MODEL.output_keys[0]] # Shape: [1, num_nodes, dim] + pred_accel = pred_output[cfg.MODEL.output_keys[0]] - # Calculate error error = np.mean(np.abs(pred_accel.numpy() - true_accel)) logger.info(f"Sample {sample_idx} - MAE error: {error:.6f}") - # Calculate relative error rel_error = np.linalg.norm(pred_accel.numpy() - true_accel) / np.linalg.norm(true_accel) logger.info(f"Sample {sample_idx} - Relative error: {rel_error:.6f}") - # Visualization using simulate.py plot function plt.figure(figsize=(10, 8)) sim.plot(sample_idx, animate=False, plot_size=True, s_size=2) plt.title(f'{cfg.DATA.type.capitalize()} System - Sample {sample_idx}\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}') plt.tight_layout() - # Save plot plot_path = os.path.join(cfg.output_dir, f"evaluation_sample_{sample_idx}.png") plt.savefig(plot_path, dpi=300, bbox_inches='tight') plt.close() logger.info(f"Evaluation plot saved to {plot_path}") + def export(cfg: DictConfig): @@ -326,19 +314,6 @@ def export(cfg: DictConfig): l1_strength=l1_strength ) - # 创建数据集 - # sim = SimulationDataset( - # sim=cfg.DATA.type, - # n=cfg.DATA.num_nodes, - # dim=cfg.DATA.dimension, - # nt=cfg.DATA.time_steps, - # dt=cfg.DATA.time_step_size - # ) - # sim.simulate(cfg.DATA.num_samples) - # accel_data = sim.get_acceleration() - # # 打印维度 - # logger.info(f"accel_data shape: {accel_data.shape}") - # logger.info(f"edge_index shape: {edge_index.shape}") solver = ppsci.solver.Solver( model, pretrained_model_path=cfg.INFER.pretrained_model_path, @@ -368,7 +343,6 @@ def inference(cfg: DictConfig): logger.info("Switching to direct model inference...") use_predictor = False - # Generate test data sim = SimulationDataset( sim=cfg.DATA.type, n=cfg.DATA.num_nodes, @@ -378,13 +352,11 @@ def inference(cfg: DictConfig): ) sim.simulate(cfg.DATA.num_samples) - # Prepare input data for inference sample_idx = 0 - sample_data = sim.data[sample_idx, 0] # Shape: [num_nodes, n_f] + sample_data = sim.data[sample_idx, 0] edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) if use_predictor: - # Use GeneralPredictor input_dict = { "x": sample_data, "edge_index": edge_index @@ -392,7 +364,6 @@ def inference(cfg: DictConfig): output_dict = predictor.predict(input_dict, cfg.INFER.batch_size) pred_acceleration = output_dict[cfg.MODEL.output_keys[0]] else: - # Direct model inference if cfg.MODEL.n_f == "auto": cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 if cfg.MODEL.ndim == "auto": @@ -424,30 +395,25 @@ def inference(cfg: DictConfig): pred_output = model(input_dict) pred_acceleration = pred_output[cfg.MODEL.output_keys[0]].numpy() - # Calculate true acceleration for comparison accel_data = sim.get_acceleration() - true_acceleration = accel_data[sample_idx, 0] # Shape: [num_nodes, dim] + true_acceleration = accel_data[sample_idx, 0] - # Calculate error error = np.mean(np.abs(pred_acceleration - true_acceleration)) logger.info(f"Inference error (MAE): {error:.6f}") - # Calculate relative error rel_error = np.linalg.norm(pred_acceleration - true_acceleration) / np.linalg.norm(true_acceleration) logger.info(f"Inference relative error: {rel_error:.6f}") - # Visualization using simulate.py plot function plt.figure(figsize=(10, 8)) sim.plot(sample_idx, animate=False, plot_size=True, s_size=2) plt.title(f'{cfg.DATA.type.capitalize()} System - Inference\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}') plt.tight_layout() - # Save plot plot_path = os.path.join(cfg.output_dir, "inference.png") plt.savefig(plot_path, dpi=300, bbox_inches='tight') plt.close() logger.info(f"Inference plot saved to {plot_path}") - + @hydra.main(version_base=None, config_path="./conf", config_name="config_hgn") def main(cfg: DictConfig) -> None: diff --git a/examples/symbolic_gn/simulate.py b/examples/symbolic_gn/simulate.py index a233733279..c854ef8989 100644 --- a/examples/symbolic_gn/simulate.py +++ b/examples/symbolic_gn/simulate.py @@ -12,14 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -import numpy as np # Standard Numpy for numerical computation -import paddle # PaddlePaddle for automatic differentiation +import numpy as np +import paddle from matplotlib import pyplot as plt -from scipy.integrate import odeint # Scipy ODE solver +from scipy.integrate import odeint import matplotlib as mpl from functools import partial from tqdm import tqdm -from celluloid import Camera # For creating animations +from celluloid import Camera def make_transparent_color(ntimes, fraction): """ @@ -37,19 +37,14 @@ def make_transparent_color(ntimes, fraction): numpy.ndarray: RGBA color array with shape (ntimes, 4). Each row represents the color [R, G, B, Alpha] at a time step. """ - # Initialize RGBA array, all set to 1 (white, fully opaque) rgba = np.ones((ntimes, 4)) - # Create transparency gradient from 0 (fully transparent) to 1 (fully opaque) alpha = np.linspace(0, 1, ntimes)[:, np.newaxis] - # Select base color from spectrum according to fraction color = np.array(mpl.colors.to_rgba(mpl.cm.gist_ncar(fraction)))[np.newaxis, :] - # Linear interpolation: from white (1) to target color rgba[:, :] = 1*(1-alpha) + color*alpha - # Set alpha channel rgba[:, 3] = alpha[:, 0] return rgba @@ -79,7 +74,7 @@ def get_potential(sim, sim_obj): Returns: function: Potential energy function that takes two particle states and returns their potential energy. """ - dim = sim_obj._dim # Spatial dimension (2D or 3D) + dim = sim_obj._dim def potential(x1, x2): """ @@ -96,63 +91,50 @@ def potential(x1, x2): Returns: float: Potential energy value between the two particles. """ - # Calculate Euclidean distance between two particles dist = np.sqrt(np.sum(np.square(x1[:dim] - x2[:dim]))) - # Prevent numerical singularity (division by zero), add minimum distance limit min_dist = 1e-2 bounded_dist = dist + min_dist - # ========== Potential energy functions for different physical scenarios ========== - if sim == 'r2': - # Inverse square law potential: U = -G*m1*m2/r - # Physical scenarios: universal gravitation, Coulomb attraction return -x1[-1]*x2[-1]/bounded_dist elif sim == 'r1': - # Logarithmic potential: U = m1*m2*ln(r) return x1[-1]*x2[-1]*np.log(bounded_dist) elif sim in ['spring', 'damped']: - # Spring potential: U = (r - r0)^2, where r0=1 is natural length potential_val = (bounded_dist - 1)**2 if sim == 'damped': - # Add damping dissipation term (velocity-dependent) damping = 1 - potential_val += damping*x1[1]*x1[1+sim_obj._dim]/sim_obj._n # y direction - potential_val += damping*x1[0]*x1[0+sim_obj._dim]/sim_obj._n # x direction + potential_val += damping*x1[1]*x1[1+sim_obj._dim]/sim_obj._n + potential_val += damping*x1[0]*x1[0+sim_obj._dim]/sim_obj._n if sim_obj._dim == 3: - potential_val += damping*x1[2]*x1[2+sim_obj._dim]/sim_obj._n # z direction + potential_val += damping*x1[2]*x1[2+sim_obj._dim]/sim_obj._n return potential_val elif sim == 'string': - # String potential: spring potential + gravitational potential + return (bounded_dist - 1)**2 + x1[1]*x1[-1] elif sim == 'string_ball': - # String potential with ball obstacle potential_val = (bounded_dist - 1)**2 + x1[1]*x1[-1] - # Calculate distance from particle to ball center r = np.sqrt((x1[1] + 15)**2 + (x1[0] - 5)**2) - radius = 4.0 # Ball radius - - # Soft boundary repulsion potential (using more stable form) + radius = 4.0 k_repel = 100.0 potential_val += k_repel / np.maximum(r - radius + 0.5, 0.01) return potential_val elif sim in ['charge', 'superposition']: - # Charge interaction: U = q1*q2/r (Coulomb's law) + charge1 = x1[-2] charge2 = x2[-2] potential_val = charge1*charge2/bounded_dist if sim in ['superposition']: - # Superposition: contains both gravity and charge interaction + m1 = x1[-1] m2 = x2[-1] potential_val += -m1*m2/bounded_dist @@ -160,16 +142,13 @@ def potential(x1, x2): return potential_val elif sim in ['discontinuous']: - # Piecewise discontinuous potential function + m1 = x1[-1] m2 = x2[-1] - # Define potential in three regions - pot_a = 0.0 # r < 1: no interaction - pot_b = 0.0 # 1 <= r < 2: no interaction - pot_c = (bounded_dist - 1)**2 # r >= 2: spring potential - - # Implement piecewise function using conditional expressions + pot_a = 0.0 + pot_b = 0.0 + pot_c = (bounded_dist - 1)**2 potential_val = ( pot_a * (bounded_dist < 1) + (bounded_dist >= 1) * ( @@ -233,13 +212,13 @@ def __init__(self, sim='r2', n=5, dim=2, self.nt = nt self.data = None - # Generate uniform time series + self.times = np.linspace(0, self.dt*self.nt, num=self.nt) - self.G = 1 # Gravitational constant + self.G = 1 self.extra_potential = extra_potential - # Get pairwise potential function for the simulation type + self.pairwise = get_potential(sim=sim, sim_obj=self) def simulate(self, ns, seed=0): @@ -262,13 +241,13 @@ def simulate(self, ns, seed=0): Returns: None: Results are stored in self.data. """ - # Set random seed + np.random.seed(seed) n = self._n dim = self._dim sim = self._sim - params = 2 # charge + mass + params = 2 total_dim = dim*2 + params times = self.times G = self.G @@ -294,10 +273,10 @@ def total_potential(xt): for i in range(n - 1): if sim in ['string', 'string_ball']: - # String type: only adjacent particles interact + sum_potential += G * self.pairwise(xt[i], xt[i+1]) else: - # Other types: fully connected + for j in range(i+1, n): sum_potential += G * self.pairwise(xt[i], xt[j]) @@ -323,101 +302,101 @@ def force_paddle(xt): Returns: numpy.ndarray: Force on each particle with shape (n, dim). """ - # Convert numpy array to Paddle Tensor (improved precision: float32 → float64) + xt_tensor = paddle.to_tensor(xt, dtype='float64', stop_gradient=False) - # Calculate gradient only for position coordinates + positions = xt_tensor[:, :dim] positions.stop_gradient = False - # Rebuild complete state for potential energy calculation + xt_for_potential = paddle.concat([ positions, xt_tensor[:, dim:] ], axis=1) - # Calculate total potential energy (using float64 for higher precision) + sum_potential = paddle.to_tensor(0.0, dtype='float64') - # Iterate through all particle pairs to calculate potential energy + for i in range(n - 1): if sim in ['string', 'string_ball']: - # ========== String type: only calculate adjacent particles ========== + x1 = xt_for_potential[i] x2 = xt_for_potential[i+1] dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) bounded_dist = dist + 1e-2 if sim == 'string': - # String potential: spring + gravity + pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] elif sim == 'string_ball': - # String potential with ball obstacle + pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] - # Calculate distance to ball center + r = paddle.sqrt((x1[1] + 15)**2 + (x1[0] - 5)**2) radius = paddle.to_tensor(4.0, dtype='float64') - # Soft boundary repulsion potential (reduced strength to avoid numerical instability) - # More stable repulsion potential: strong repulsion when r < radius - k_repel = 100.0 # Reduced repulsion strength + + + k_repel = 100.0 pot = pot + k_repel / paddle.maximum(r - radius + 0.5, paddle.to_tensor(0.01, dtype='float64')) sum_potential = sum_potential + G * pot else: - # ========== Fully connected type: calculate all particle pairs ========== + for j in range(i+1, n): x1 = xt_for_potential[i] x2 = xt_for_potential[j] dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) bounded_dist = dist + 1e-2 - # Calculate potential based on simulation type + if sim == 'r2': - # Inverse square law: universal gravitation + pot = -x1[-1]*x2[-1]/bounded_dist elif sim == 'r1': - # Logarithmic potential + pot = x1[-1]*x2[-1]*paddle.log(bounded_dist) elif sim in ['spring', 'damped']: - # Spring potential + pot = (bounded_dist - 1)**2 - # Add damping term + if sim == 'damped': damping = paddle.to_tensor(1.0, dtype='float64') - # Damping dissipation term: proportional to position-velocity product - pot = pot + damping*x1[1]*x1[1+dim]/n # y direction - pot = pot + damping*x1[0]*x1[0+dim]/n # x direction + + pot = pot + damping*x1[1]*x1[1+dim]/n + pot = pot + damping*x1[0]*x1[0+dim]/n if dim == 3: - pot = pot + damping*x1[2]*x1[2+dim]/n # z direction + pot = pot + damping*x1[2]*x1[2+dim]/n elif sim in ['charge', 'superposition']: - # Charge interaction (Coulomb's law) + charge1 = x1[-2] charge2 = x2[-2] pot = charge1*charge2/bounded_dist if sim == 'superposition': - # Superposition: contains both charge and gravity + m1 = x1[-1] m2 = x2[-1] pot = pot + (-m1*m2/bounded_dist) elif sim == 'discontinuous': - # Piecewise discontinuous potential - # Define potential in three regions - pot_a = paddle.to_tensor(0.0, dtype='float64') # r < 1 - pot_b = paddle.to_tensor(0.0, dtype='float64') # 1 <= r < 2 - pot_c = (bounded_dist - 1)**2 # r >= 2 + + + pot_a = paddle.to_tensor(0.0, dtype='float64') + pot_b = paddle.to_tensor(0.0, dtype='float64') + pot_c = (bounded_dist - 1)**2 - # Implement piecewise function using conditional expressions - # Paddle's conditional operations automatically handle gradients + + cond1 = paddle.cast(bounded_dist < 1, 'float64') cond2 = paddle.cast(bounded_dist >= 1, 'float64') * paddle.cast(bounded_dist < 2, 'float64') cond3 = paddle.cast(bounded_dist >= 2, 'float64') @@ -425,31 +404,31 @@ def force_paddle(xt): pot = pot_a * cond1 + pot_b * cond2 + pot_c * cond3 else: - # Default: spring potential + pot = (bounded_dist - 1)**2 sum_potential = sum_potential + G * pot - # Add external potential support + if self.extra_potential is not None: - # Note: extra_potential needs to be Paddle-compatible - # If user provides a numpy function, there will be issues - # Suggestion: use numerical gradients or require Paddle version from user + + + for i in range(n): - # Assume extra_potential returns scalar - # May need adjustment in actual use + + try: extra_pot = self.extra_potential(xt_for_potential[i]) if not isinstance(extra_pot, paddle.Tensor): extra_pot = paddle.to_tensor(extra_pot, dtype='float64') sum_potential = sum_potential + extra_pot except: - # If extra_potential is not Paddle-compatible, skip - # Should give warning in actual use + + pass - # Calculate gradient: F = -dU/dx - # Use paddle.grad for automatic differentiation + + grads = paddle.grad( outputs=sum_potential, inputs=positions, @@ -457,7 +436,7 @@ def force_paddle(xt): retain_graph=False )[0] - # Convert back to numpy + force = -grads.numpy() return force @@ -473,7 +452,7 @@ def acceleration(xt): numpy.ndarray: Acceleration. """ force = force_paddle(xt) - masses = xt[:, -1:] # Shape (n, 1) + masses = xt[:, -1:] return force / masses def odefunc(y, t): @@ -487,17 +466,17 @@ def odefunc(y, t): Returns: numpy.ndarray: Time derivative of the state. """ - # Restore shape + y = y.reshape((n, total_dim)) - # Calculate acceleration + a = acceleration(y) - # Construct derivative + dydt = np.concatenate([ - y[:, dim:2*dim], # dx/dt = v - a, # dv/dt = a - np.zeros((n, params)) # d(q,m)/dt = 0 + y[:, dim:2*dim], + a, + np.zeros((n, params)) ], axis=1) return dydt.flatten() @@ -513,21 +492,21 @@ def make_sim(sample_idx): numpy.ndarray: Simulation trajectory with shape (nt, n, total_dim). """ if sim in ['string', 'string_ball']: - # String type initialization + x0 = np.random.randn(n, total_dim) - x0[:, -1] = 1.0 # Fixed mass - x0[:, 0] = np.arange(n) + x0[:, 0]*0.5 # Evenly spaced x-coordinates - x0[:, 2:3] = 0.0 # vx = 0 + x0[:, -1] = 1.0 + x0[:, 0] = np.arange(n) + x0[:, 0]*0.5 + x0[:, 2:3] = 0.0 else: - # General initialization + x0 = np.random.randn(n, total_dim) - x0[:, -1] = np.exp(x0[:, -1]) # Positive mass + x0[:, -1] = np.exp(x0[:, -1]) if sim in ['charge', 'superposition']: - x0[:, -2] = np.sign(x0[:, -2]) # Charge ±1 + x0[:, -2] = np.sign(x0[:, -2]) - # Solve ODE using scipy.odeint - # Migration point: JAX odeint → scipy.integrate.odeint + + x_times = odeint( odefunc, x0.flatten(), @@ -537,13 +516,13 @@ def make_sim(sample_idx): return x_times - # Generate samples in batch + data = [] print(f"Start generating {ns} samples...") for i in tqdm(range(ns)): data.append(make_sim(i)) - # Convert to numpy array + self.data = np.array(data) print(f"Generation complete! Data shape: {self.data.shape}") @@ -565,7 +544,7 @@ def get_acceleration(self): sim = self._sim G = self.G - # Pre-allocate acceleration array + accelerations = np.zeros((ns, nt, n, dim)) print("Calculating acceleration...") @@ -573,12 +552,12 @@ def get_acceleration(self): for time_idx in range(nt): xt = self.data[sample_idx, time_idx] - # Calculate acceleration using PaddlePaddle (using float64 for higher precision) + xt_tensor = paddle.to_tensor(xt, dtype='float64', stop_gradient=False) positions = xt_tensor[:, :dim] positions.stop_gradient = False - # Calculate total potential + xt_for_potential = paddle.concat([ positions, xt_tensor[:, dim:] @@ -586,10 +565,10 @@ def get_acceleration(self): sum_potential = paddle.to_tensor(0.0, dtype='float64') - # Iterate through all particle pairs to calculate potential (identical to force_paddle) + for i in range(n - 1): if sim in ['string', 'string_ball']: - # ========== String type ========== + x1 = xt_for_potential[i] x2 = xt_for_potential[i+1] dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) @@ -598,17 +577,17 @@ def get_acceleration(self): if sim == 'string': pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] elif sim == 'string_ball': - # Complete implementation + pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] r = paddle.sqrt((x1[1] + 15)**2 + (x1[0] - 5)**2) radius = paddle.to_tensor(4.0, dtype='float64') - # More stable repulsion potential + k_repel = 100.0 pot = pot + k_repel / paddle.maximum(r - radius + 0.5, paddle.to_tensor(0.01, dtype='float64')) sum_potential = sum_potential + G * pot else: - # ========== Fully connected type ========== + for j in range(i+1, n): x1 = xt_for_potential[i] x2 = xt_for_potential[j] @@ -622,7 +601,7 @@ def get_acceleration(self): elif sim in ['spring', 'damped']: pot = (bounded_dist - 1)**2 - # Add damping term + if sim == 'damped': damping = paddle.to_tensor(1.0, dtype='float64') pot = pot + damping*x1[1]*x1[1+dim]/n @@ -640,7 +619,7 @@ def get_acceleration(self): pot = pot + (-m1*m2/bounded_dist) elif sim == 'discontinuous': - # Piecewise discontinuous potential + pot_a = paddle.to_tensor(0.0, dtype='float64') pot_b = paddle.to_tensor(0.0, dtype='float64') pot_c = (bounded_dist - 1)**2 @@ -656,7 +635,7 @@ def get_acceleration(self): sum_potential = sum_potential + G * pot - # Add external potential support + if self.extra_potential is not None: for i in range(n): try: @@ -667,7 +646,7 @@ def get_acceleration(self): except: pass - # Calculate force + grads = paddle.grad( outputs=sum_potential, inputs=positions, @@ -706,7 +685,7 @@ def plot(self, i, animate=False, plot_size=True, s_size=1): masses = x_times[:, :, -1] if not animate: - # Static plot + if sim in ['string', 'string_ball']: rgba = make_transparent_color(len(times), 0) for idx in range(0, len(times), len(times)//10): @@ -724,7 +703,7 @@ def plot(self, i, animate=False, plot_size=True, s_size=1): plt.scatter(x_times[:, j, 0], x_times[:, j, 1], color=rgba, s=s_size) else: - # Animation + if sim in ['string', 'string_ball']: raise NotImplementedError("Animation mode not yet supported for string type") From ab879ea46a18484997ee3068eeb6b55eac99a3c3 Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Fri, 5 Dec 2025 15:20:11 +0000 Subject: [PATCH 07/10] =?UTF-8?q?feat=EF=BC=9Aupdate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/zh/examples/symbolic_gn.md | 528 +-------------------------- examples/symbolic_gn/gn.py | 92 ++--- examples/symbolic_gn/simulate.py | 588 +++++++++++++++---------------- ppsci/arch/symbolic_gn.py | 474 +++++++++++++------------ 4 files changed, 573 insertions(+), 1109 deletions(-) diff --git a/docs/zh/examples/symbolic_gn.md b/docs/zh/examples/symbolic_gn.md index 2bb58c7163..b6caf40417 100644 --- a/docs/zh/examples/symbolic_gn.md +++ b/docs/zh/examples/symbolic_gn.md @@ -200,25 +200,12 @@ MODEL: 1. **数据集准备**： ```python - train_dataset = OGNDataset(cfg, mode="train") - train_data, train_labels = train_dataset.get_data_and_labels() + examples/symbolic_gn/gn.py:113:119 ``` 2. **约束构建**： ```python - train_constraint = ppsci.constraint.SupervisedConstraint( - { - "dataset": { - "name": "NamedArrayDataset", - "input": train_data, - "label": train_labels - }, - "batch_size": cfg.TRAIN.batch_size, - "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": True}, - }, - loss=loss_fn, - name="train_constraint", - ) + examples/symbolic_gn/gn.py:145:182 ``` 3. **损失函数**：根据模型类型自动调整 @@ -235,21 +222,7 @@ MODEL: 训练使用Adam优化器配合OneCycleLR学习率调度： ```python -# 动态计算每个epoch的批次数 -batch_per_epoch = int(1000*10 / (cfg.TRAIN.batch_size/32.0)) - -# OneCycleLR学习率调度器 -lr_scheduler = paddle.optimizer.lr.OneCycleLR( - max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, - total_steps=cfg.TRAIN.epochs * batch_per_epoch, - divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor -) - -optimizer = paddle.optimizer.Adam( - learning_rate=lr_scheduler, - parameters=model.parameters(), - weight_decay=cfg.TRAIN.optimizer.weight_decay -) +examples/symbolic_gn/gn.py:194:213 ``` 关键超参数： @@ -323,502 +296,9 @@ TRAIN: ## 4. 完整代码 -完整实现包含以下文件： ```python - - -class OGNDataset: - def __init__(self, cfg, mode="train"): - self.cfg = cfg - self.mode = mode - self.data = None - self.labels = None - self.input_keys = tuple(cfg.MODEL.input_keys) - self.label_keys = tuple(cfg.MODEL.output_keys) - self._prepare_data() - - def _prepare_data(self): - # Get time step size for physics system - dt = 1e-2 - for sim_set in self.cfg.DATATYPE_TYPE: - if sim_set['sim'] == self.cfg.DATA.type: - dt = sim_set['dt'][0] - break - - # Generate simulation data - sim_dataset = SimulationDataset( - sim=self.cfg.DATA.type, - n=self.cfg.DATA.num_nodes, - dim=self.cfg.DATA.dimension, - dt=dt, - nt=self.cfg.DATA.time_steps, - seed=self.cfg.seed - ) - sim_dataset.simulate(self.cfg.DATA.num_samples) - - # Compute target data based on model architecture - if self.cfg.MODEL.arch == "HGN": - raw_target_data = sim_dataset.get_derivative() # [v, a] for HGN - else: - raw_target_data = sim_dataset.get_acceleration() # Acceleration for OGN/VarOGN - - # Downsample data - downsample_factor = self.cfg.DATA.downsample_factor - input_data = np.concatenate([sim_dataset.data[:, i] for i in range(0, sim_dataset.data.shape[1], downsample_factor)]) - target_data = np.concatenate([raw_target_data[:, i] for i in range(0, sim_dataset.data.shape[1], downsample_factor)]) - - # Split train/validation set - from sklearn.model_selection import train_test_split - test_size = 1.0 - getattr(self.cfg.TRAIN, 'train_split', 0.8) - - if self.mode == "train": - input_data, _, target_data, _ = train_test_split( - input_data, target_data, test_size=test_size, shuffle=False - ) - elif self.mode == "eval": - _, input_data, _, target_data = train_test_split( - input_data, target_data, test_size=test_size, shuffle=False - ) - - # Generate edge indices - edge_index = get_edge_index(self.cfg.DATA.num_nodes, self.cfg.DATA.type) - - self.data = { - "node_features": input_data.astype(np.float32), - "edge_index": np.tile(edge_index.T.numpy(), (len(input_data), 1, 1)) - } - self.labels = { - list(self.cfg.MODEL.output_keys)[0]: target_data.astype(np.float32) - } - - def __len__(self): - return len(self.data["node_features"]) - - def __getitem__(self, idx): - return { - "input": { - "node_features": self.data["node_features"][idx], - "edge_index": self.data["edge_index"][idx] - }, - "label": { - list(self.cfg.MODEL.output_keys)[0]: self.labels[list(self.cfg.MODEL.output_keys)[0]][idx] - } - } - - def get_data_and_labels(self): - return self.data, self.labels - - def __call__(self): - return self - - -def create_loss_function(cfg): - def loss_function(input_dict, label_dict, model_output): - # Get loss type from config (MAE or MSE) - loss_type = getattr(cfg.TRAIN.loss, 'type', 'MAE') - - # Get output key - output_key = list(cfg.MODEL.output_keys)[0] if cfg.MODEL.arch != "HGN" else "derivative" - - # Compute base loss - if cfg.MODEL.arch == "HGN": - output_key = "derivative" - pred_accel = model_output[output_key][:, cfg.DATA.dimension:] - target_accel = label_dict[output_key][:, cfg.DATA.dimension:] - if loss_type == "MSE": - base_loss = paddle.mean(paddle.square(pred_accel - target_accel)) - else: # MAE - base_loss = paddle.mean(paddle.abs(pred_accel - target_accel)) - else: - pred = model_output[output_key] - target = label_dict[output_key] - if loss_type == "MSE": - base_loss = paddle.mean(paddle.square(pred - target)) - else: # MAE - base_loss = paddle.mean(paddle.abs(pred - target)) - - # Add regularization if configured - if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"]: - reg_loss = cfg.MODEL.l1_strength * paddle.mean(paddle.abs(input_dict["node_features"])) - return base_loss + reg_loss - - return base_loss - - return loss_function - - -def train(cfg): - # Set random seed for reproducibility - ppsci.utils.misc.set_random_seed(cfg.seed) - - # Parse automatic configuration parameters - if cfg.MODEL.n_f == "auto": - cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 - if cfg.MODEL.ndim == "auto": - cfg.MODEL.ndim = cfg.DATA.dimension - if cfg.TRAIN.batch_size == "auto": - cfg.TRAIN.batch_size = int(64 * (4 / cfg.DATA.num_nodes) ** 2) - - # Update output keys for HGN - if cfg.MODEL.arch == "HGN": - cfg.MODEL.output_keys = ["derivative"] - - # Create datasets - logger.message(f"Creating {cfg.DATA.type} dataset using simulate.py...") - train_dataset = OGNDataset(cfg, mode="train") - eval_dataset = OGNDataset(cfg, mode="eval") - - # Get training and evaluation data - train_data, train_labels = train_dataset.get_data_and_labels() - eval_data, eval_labels = eval_dataset.get_data_and_labels() - - # Create model based on architecture - logger.message(f"Creating model: {cfg.MODEL.arch}") - if cfg.MODEL.arch == "OGN": - model = OGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "VarOGN": - model = VarOGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "HGN": - model = HGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - ndim=cfg.MODEL.ndim, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - else: - raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") - - # Create loss function and constraints - loss_fn = create_loss_function(cfg) - train_constraint = ppsci.constraint.SupervisedConstraint( - { - "dataset": { - "name": "NamedArrayDataset", - "input": train_data, - "label": train_labels - }, - "batch_size": cfg.TRAIN.batch_size, - "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": True}, - }, - loss=loss_fn, - name="train_constraint", - ) - - eval_constraint = ppsci.constraint.SupervisedConstraint( - { - "dataset": { - "name": "NamedArrayDataset", - "input": eval_data, - "label": eval_labels - }, - "batch_size": cfg.EVAL.batch_size, - "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": False}, - }, - loss=loss_fn, - name="eval_constraint", - ) - - constraint = { - train_constraint.name: train_constraint, - eval_constraint.name: eval_constraint, - } - - # Calculate iterations per epoch - batch_per_epoch = int(1000*10 / (cfg.TRAIN.batch_size/32.0)) - - # Create optimizer and learning rate scheduler - if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": - lr_scheduler = paddle.optimizer.lr.OneCycleLR( - max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, - total_steps=cfg.TRAIN.epochs*batch_per_epoch, - divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor - ) - else: - lr_scheduler = paddle.optimizer.lr.ExponentialDecay( - learning_rate=cfg.TRAIN.optimizer.learning_rate, - gamma=0.9 - ) - - optimizer = paddle.optimizer.Adam( - learning_rate=lr_scheduler, - parameters=model.parameters(), - weight_decay=cfg.TRAIN.optimizer.weight_decay - ) - - # Create PaddleScience Solver and start training - solver = ppsci.solver.Solver( - model, - constraint, - optimizer=optimizer, - cfg=cfg, - ) - solver.train() - - -def evaluate(cfg: DictConfig): - # Parse automatic configuration parameters - if cfg.MODEL.n_f == "auto": - cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 - if cfg.MODEL.ndim == "auto": - cfg.MODEL.ndim = cfg.DATA.dimension - if cfg.MODEL.arch == "HGN": - cfg.MODEL.output_keys = ["derivative"] - - # Create model based on architecture - if cfg.MODEL.arch == "OGN": - model = OGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "VarOGN": - model = VarOGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "HGN": - model = HGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - ndim=cfg.MODEL.ndim, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - else: - raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") - - # Create evaluation dataset - eval_dataset = OGNDataset(cfg, mode="eval") - eval_data, eval_labels = eval_dataset.get_data_and_labels() - - # Create loss function and constraint - loss_fn = create_loss_function(cfg) - eval_constraint = ppsci.constraint.SupervisedConstraint( - { - "dataset": { - "name": "NamedArrayDataset", - "input": eval_data, - "label": eval_labels - }, - "batch_size": cfg.EVAL.batch_size, - "sampler": {"name": "BatchSampler", "drop_last": False, "shuffle": False}, - }, - loss=loss_fn, - name="eval_constraint", - ) - - constraint = { - eval_constraint.name: eval_constraint, - } - - # Create PaddleScience Solver - solver = ppsci.solver.Solver(model, constraint, cfg=cfg) - - # Run evaluation - logger.message("Starting evaluation...") - eval_result = solver.eval() - logger.message(f"Evaluation completed. Result: {eval_result}") - - -def inference(cfg: DictConfig): - # Parse automatic configuration parameters - if cfg.MODEL.n_f == "auto": - cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 - if cfg.MODEL.ndim == "auto": - cfg.MODEL.ndim = cfg.DATA.dimension - if cfg.MODEL.arch == "HGN": - cfg.MODEL.output_keys = ["derivative"] - - # Create model based on architecture - if cfg.MODEL.arch == "OGN": - model = OGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "VarOGN": - model = VarOGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "HGN": - model = HGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - ndim=cfg.MODEL.ndim, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - else: - raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") - - # Create PaddleScience Solver - solver = ppsci.solver.Solver(model, cfg=cfg) - - # Generate test data - logger.message("Generating inference data using simulate.py...") - dt = 1e-2 - for sim_set in cfg.DATATYPE_TYPE: - if sim_set['sim'] == cfg.DATA.type: - dt = sim_set['dt'][0] - break - - sim_dataset = SimulationDataset( - sim=cfg.DATA.type, - n=cfg.DATA.num_nodes, - dim=cfg.DATA.dimension, - dt=dt, - nt=100, - seed=cfg.seed - ) - sim_dataset.simulate(1) - - # Prepare input data (first time step) - node_features = sim_dataset.data[0, 0] - edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) - - input_dict = { - "node_features": paddle.to_tensor(node_features, dtype='float32').unsqueeze(0), - "edge_index": edge_index.unsqueeze(0) - } - - # Run inference - logger.message("Running inference...") - output_dict = solver.predict(input_dict, return_numpy=True) - - # Display results - if cfg.MODEL.arch == "HGN": - dq_dt = output_dict["derivative"][0, :cfg.DATA.dimension] - dv_dt = output_dict["derivative"][0, cfg.DATA.dimension:] - logger.message(f"HGN inference - dq/dt: {dq_dt}, dv/dt (acceleration): {dv_dt}") - else: - acceleration = output_dict[list(cfg.MODEL.output_keys)[0]][0] - logger.message(f"{cfg.MODEL.arch} inference - acceleration: {acceleration}") - - -def export(cfg: DictConfig): - # Parse automatic configuration parameters - if cfg.MODEL.n_f == "auto": - cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 - if cfg.MODEL.ndim == "auto": - cfg.MODEL.ndim = cfg.DATA.dimension - if cfg.MODEL.arch == "HGN": - cfg.MODEL.output_keys = ["derivative"] - - # Create model based on architecture - if cfg.MODEL.arch == "OGN": - model = OGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "VarOGN": - model = VarOGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - msg_dim=cfg.MODEL.msg_dim, - ndim=cfg.MODEL.ndim, - dt=cfg.MODEL.dt, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - elif cfg.MODEL.arch == "HGN": - model = HGN( - input_keys=tuple(cfg.MODEL.input_keys), - output_keys=tuple(cfg.MODEL.output_keys), - n_f=cfg.MODEL.n_f, - ndim=cfg.MODEL.ndim, - hidden=cfg.MODEL.hidden, - aggr=cfg.MODEL.aggr - ) - else: - raise ValueError(f"Unsupported model architecture: {cfg.MODEL.arch}") - - # Create PaddleScience Solver - solver = ppsci.solver.Solver(model, cfg=cfg) - - # Define input specifications for export - from paddle.static import InputSpec - input_spec = [ - { - "node_features": InputSpec([None, cfg.DATA.num_nodes, cfg.MODEL.n_f], "float32", "node_features"), - "edge_index": InputSpec([None, 2, cfg.DATA.num_nodes * (cfg.DATA.num_nodes - 1)], "int64", "edge_index") - } - ] - - # Export model to static graph - logger.message(f"Exporting model to {cfg.INFER.export_path}") - solver.export(input_spec, cfg.INFER.export_path, with_onnx=False) - logger.message("Model export completed!") - - -@hydra.main(version_base=None, config_path="./conf", config_name="config.yaml") -def main(cfg: DictConfig): - if cfg.mode == "train": - train(cfg) - elif cfg.mode == "eval": - evaluate(cfg) - elif cfg.mode == "infer": - inference(cfg) - elif cfg.mode == "export": - export(cfg) - else: - raise ValueError(f"Unsupported mode: {cfg.mode}") - - -if __name__ == "__main__": - main() +examples/symbolic_gn/gn.py ``` diff --git a/examples/symbolic_gn/gn.py b/examples/symbolic_gn/gn.py index 5f1b0e227f..4d5e99f6f4 100644 --- a/examples/symbolic_gn/gn.py +++ b/examples/symbolic_gn/gn.py @@ -13,18 +13,22 @@ # limitations under the License. import os -import sys -from typing import List, Optional +from typing import List +from typing import Optional + +import hydra +import matplotlib.pyplot as plt import numpy as np import paddle from omegaconf import DictConfig -import hydra +from simulate import SimulationDataset + import ppsci +from ppsci.arch.symbolic_gn import HGN +from ppsci.arch.symbolic_gn import OGN +from ppsci.arch.symbolic_gn import VarOGN +from ppsci.arch.symbolic_gn import get_edge_index from ppsci.utils import logger -import matplotlib.pyplot as plt - -from simulate import SimulationDataset -from ppsci.arch.symbolic_gn import OGN, VarOGN, HGN, get_edge_index def create_ppsci_model( @@ -71,7 +75,7 @@ def create_ppsci_model( ) else: raise ValueError(f"未知的模型类型: {model_type}") - + return model @@ -102,7 +106,7 @@ def loss_function(output_dict, label_dict, weight_dict=None): def train(cfg): ppsci.utils.misc.set_random_seed(cfg.seed) - + if cfg.MODEL.arch == "HGN": cfg.MODEL.output_keys = ["acceleration"] @@ -115,17 +119,17 @@ def train(cfg): ) sim.simulate(cfg.DATA.num_samples) accel_data = sim.get_acceleration() - + X_list = [] y_list = [] for sample_idx in range(cfg.DATA.num_samples): for t in range(0, sim.data.shape[1], cfg.DATA.sample_interval): X_list.append(sim.data[sample_idx, t]) y_list.append(accel_data[sample_idx, t]) - + X = np.array(X_list, dtype=np.float32) y = np.array(y_list, dtype=np.float32) - + if cfg.MODEL.n_f == "auto": cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 if cfg.MODEL.ndim == "auto": @@ -137,7 +141,7 @@ def train(cfg): edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) edge_index_train = replicate_array(edge_index, len(X_train)) edge_index_val = replicate_array(edge_index, len(X_val)) - + train_constraint = ppsci.constraint.SupervisedConstraint( { "dataset": { @@ -186,7 +190,7 @@ def train(cfg): edge_index=edge_index, l1_strength=l1_strength ) - + batch_per_epoch = int(train_size / cfg.TRAIN.batch_size) if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": lr_scheduler = paddle.optimizer.lr.OneCycleLR( @@ -201,13 +205,13 @@ def train(cfg): gamma=0.9 ) lr_scheduler.by_epoch = True - + optimizer = paddle.optimizer.Adam( learning_rate=lr_scheduler, parameters=model.parameters(), weight_decay=cfg.TRAIN.optimizer.weight_decay ) - + solver = ppsci.solver.Solver( model, constraint, @@ -227,10 +231,10 @@ def evaluate(cfg: DictConfig): cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 if cfg.MODEL.ndim == "auto": cfg.MODEL.ndim = cfg.DATA.dimension - + edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 - + model = create_ppsci_model( model_type=cfg.MODEL.arch, n_f=cfg.MODEL.n_f, @@ -240,7 +244,7 @@ def evaluate(cfg: DictConfig): edge_index=edge_index, l1_strength=l1_strength ) - + ppsci.utils.save_load.load_pretrain( model, cfg.EVAL.pretrained_model_path, @@ -255,10 +259,7 @@ def evaluate(cfg: DictConfig): ) sim.simulate(cfg.DATA.num_samples) accel_data = sim.get_acceleration() - - num_time_steps_per_sample = len(range(0, sim.data.shape[1], cfg.DATA.sample_interval)) - num_time_steps_per_sample = len(range(0, sim.data.shape[1], cfg.DATA.sample_interval)) - + sample_indices = [0, 1] if cfg.DATA.num_samples > 1 else [0] for sample_idx in sample_indices: @@ -269,33 +270,32 @@ def evaluate(cfg: DictConfig): "x": paddle.to_tensor(sample_data, dtype="float32"), "edge_index": paddle.to_tensor(edge_index, dtype="int64") } - + with paddle.no_grad(): pred_output = model(input_dict) pred_accel = pred_output[cfg.MODEL.output_keys[0]] - + error = np.mean(np.abs(pred_accel.numpy() - true_accel)) logger.info(f"Sample {sample_idx} - MAE error: {error:.6f}") - + rel_error = np.linalg.norm(pred_accel.numpy() - true_accel) / np.linalg.norm(true_accel) logger.info(f"Sample {sample_idx} - Relative error: {rel_error:.6f}") - + plt.figure(figsize=(10, 8)) sim.plot(sample_idx, animate=False, plot_size=True, s_size=2) plt.title(f'{cfg.DATA.type.capitalize()} System - Sample {sample_idx}\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}') plt.tight_layout() - + plot_path = os.path.join(cfg.output_dir, f"evaluation_sample_{sample_idx}.png") plt.savefig(plot_path, dpi=300, bbox_inches='tight') plt.close() logger.info(f"Evaluation plot saved to {plot_path}") - def export(cfg: DictConfig): if cfg.MODEL.arch == "HGN": raise ValueError("HGN is not supported for export") - + ppsci.utils.misc.set_random_seed(cfg.seed) if cfg.MODEL.n_f == "auto": @@ -313,14 +313,14 @@ def export(cfg: DictConfig): edge_index=edge_index, l1_strength=l1_strength ) - + solver = ppsci.solver.Solver( model, pretrained_model_path=cfg.INFER.pretrained_model_path, ) - + from paddle.static import InputSpec - + input_spec = [ { key: InputSpec([None, cfg.DATA.num_nodes, cfg.MODEL.n_f], "float32", name=key) @@ -351,11 +351,11 @@ def inference(cfg: DictConfig): dt=cfg.DATA.time_step_size ) sim.simulate(cfg.DATA.num_samples) - + sample_idx = 0 sample_data = sim.data[sample_idx, 0] edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) - + if use_predictor: input_dict = { "x": sample_data, @@ -368,9 +368,9 @@ def inference(cfg: DictConfig): cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 if cfg.MODEL.ndim == "auto": cfg.MODEL.ndim = cfg.DATA.dimension - + l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 - + model = create_ppsci_model( model_type=cfg.MODEL.arch, n_f=cfg.MODEL.n_f, @@ -380,40 +380,40 @@ def inference(cfg: DictConfig): edge_index=edge_index, l1_strength=l1_strength ) - + ppsci.utils.save_load.load_pretrain( model, cfg.INFER.pretrained_model_path, ) - + input_dict = { "x": paddle.to_tensor(sample_data, dtype="float32"), "edge_index": paddle.to_tensor(edge_index, dtype="int64") } - + with paddle.no_grad(): pred_output = model(input_dict) pred_acceleration = pred_output[cfg.MODEL.output_keys[0]].numpy() - + accel_data = sim.get_acceleration() true_acceleration = accel_data[sample_idx, 0] - + error = np.mean(np.abs(pred_acceleration - true_acceleration)) logger.info(f"Inference error (MAE): {error:.6f}") - + rel_error = np.linalg.norm(pred_acceleration - true_acceleration) / np.linalg.norm(true_acceleration) logger.info(f"Inference relative error: {rel_error:.6f}") - + plt.figure(figsize=(10, 8)) sim.plot(sample_idx, animate=False, plot_size=True, s_size=2) plt.title(f'{cfg.DATA.type.capitalize()} System - Inference\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}') plt.tight_layout() - + plot_path = os.path.join(cfg.output_dir, "inference.png") plt.savefig(plot_path, dpi=300, bbox_inches='tight') plt.close() logger.info(f"Inference plot saved to {plot_path}") - + @hydra.main(version_base=None, config_path="./conf", config_name="config_hgn") def main(cfg: DictConfig) -> None: diff --git a/examples/symbolic_gn/simulate.py b/examples/symbolic_gn/simulate.py index c854ef8989..0a9c2d99ae 100644 --- a/examples/symbolic_gn/simulate.py +++ b/examples/symbolic_gn/simulate.py @@ -12,52 +12,52 @@ # See the License for the specific language governing permissions and # limitations under the License. +import matplotlib as mpl import numpy as np import paddle +from celluloid import Camera from matplotlib import pyplot as plt from scipy.integrate import odeint -import matplotlib as mpl -from functools import partial from tqdm import tqdm -from celluloid import Camera + def make_transparent_color(ntimes, fraction): """ Create transparent colors for trajectory visualization. - + This function generates a gradient from white to a specified color, with transparency increasing from 0 to 1. Used for drawing particle trajectories where old positions are transparent and new positions are opaque. - + Args: ntimes (int): Number of time steps, i.e., number of colors to generate. fraction (float): Color fraction in range [0,1] for selecting color from spectrum. - + Returns: numpy.ndarray: RGBA color array with shape (ntimes, 4). Each row represents the color [R, G, B, Alpha] at a time step. """ rgba = np.ones((ntimes, 4)) - + alpha = np.linspace(0, 1, ntimes)[:, np.newaxis] - + color = np.array(mpl.colors.to_rgba(mpl.cm.gist_ncar(fraction)))[np.newaxis, :] - - rgba[:, :] = 1*(1-alpha) + color*alpha - + + rgba[:, :] = 1 * (1 - alpha) + color * alpha + rgba[:, 3] = alpha[:, 0] - + return rgba def get_potential(sim, sim_obj): """ Get potential energy function based on simulation type. - + This is a factory function that returns the corresponding potential energy calculation function based on different physical scenarios. Supported scenarios include: gravity, springs, charges, damped systems, etc. - + Args: sim (str): Simulation type, options: 'r2' - Inverse square law (e.g., universal gravitation) @@ -70,7 +70,7 @@ def get_potential(sim, sim_obj): 'superposition' - Superposition of gravity and charge 'discontinuous' - Piecewise discontinuous potential sim_obj (SimulationDataset): Simulation dataset object for obtaining parameters like dimension. - + Returns: function: Potential energy function that takes two particle states and returns their potential energy. """ @@ -79,86 +79,83 @@ def get_potential(sim, sim_obj): def potential(x1, x2): """ Calculate potential energy between two particles. - + Particle state format: 2D: [x, y, vx, vy, charge, mass] 3D: [x, y, z, vx, vy, vz, charge, mass] - + Args: x1 (numpy.ndarray): State vector of the first particle. x2 (numpy.ndarray): State vector of the second particle. - + Returns: float: Potential energy value between the two particles. """ dist = np.sqrt(np.sum(np.square(x1[:dim] - x2[:dim]))) - + min_dist = 1e-2 bounded_dist = dist + min_dist - if sim == 'r2': - return -x1[-1]*x2[-1]/bounded_dist - - elif sim == 'r1': - return x1[-1]*x2[-1]*np.log(bounded_dist) - - elif sim in ['spring', 'damped']: - potential_val = (bounded_dist - 1)**2 - - if sim == 'damped': + if sim == "r2": + return -x1[-1] * x2[-1] / bounded_dist + + elif sim == "r1": + return x1[-1] * x2[-1] * np.log(bounded_dist) + + elif sim in ["spring", "damped"]: + potential_val = (bounded_dist - 1) ** 2 + + if sim == "damped": damping = 1 - potential_val += damping*x1[1]*x1[1+sim_obj._dim]/sim_obj._n - potential_val += damping*x1[0]*x1[0+sim_obj._dim]/sim_obj._n + potential_val += damping * x1[1] * x1[1 + sim_obj._dim] / sim_obj._n + potential_val += damping * x1[0] * x1[0 + sim_obj._dim] / sim_obj._n if sim_obj._dim == 3: - potential_val += damping*x1[2]*x1[2+sim_obj._dim]/sim_obj._n + potential_val += damping * x1[2] * x1[2 + sim_obj._dim] / sim_obj._n return potential_val - - elif sim == 'string': - - return (bounded_dist - 1)**2 + x1[1]*x1[-1] - - elif sim == 'string_ball': - potential_val = (bounded_dist - 1)**2 + x1[1]*x1[-1] - - r = np.sqrt((x1[1] + 15)**2 + (x1[0] - 5)**2) + + elif sim == "string": + + return (bounded_dist - 1) ** 2 + x1[1] * x1[-1] + + elif sim == "string_ball": + potential_val = (bounded_dist - 1) ** 2 + x1[1] * x1[-1] + + r = np.sqrt((x1[1] + 15) ** 2 + (x1[0] - 5) ** 2) radius = 4.0 k_repel = 100.0 potential_val += k_repel / np.maximum(r - radius + 0.5, 0.01) return potential_val - elif sim in ['charge', 'superposition']: - + elif sim in ["charge", "superposition"]: + charge1 = x1[-2] charge2 = x2[-2] - potential_val = charge1*charge2/bounded_dist - - if sim in ['superposition']: - + potential_val = charge1 * charge2 / bounded_dist + + if sim in ["superposition"]: + m1 = x1[-1] m2 = x2[-1] - potential_val += -m1*m2/bounded_dist - + potential_val += -m1 * m2 / bounded_dist + return potential_val - - elif sim in ['discontinuous']: - + + elif sim in ["discontinuous"]: + m1 = x1[-1] m2 = x2[-1] - + pot_a = 0.0 pot_b = 0.0 - pot_c = (bounded_dist - 1)**2 - potential_val = ( - pot_a * (bounded_dist < 1) + - (bounded_dist >= 1) * ( - pot_b * (bounded_dist < 2) + - pot_c * (bounded_dist >= 2)) + pot_c = (bounded_dist - 1) ** 2 + potential_val = pot_a * (bounded_dist < 1) + (bounded_dist >= 1) * ( + pot_b * (bounded_dist < 2) + pot_c * (bounded_dist >= 2) ) return potential_val - + else: - raise NotImplementedError('No such simulation ' + str(sim)) + raise NotImplementedError("No such simulation " + str(sim)) return potential @@ -166,17 +163,17 @@ def potential(x1, x2): class SimulationDataset(object): """ Physical system simulation dataset class (PaddlePaddle version). - + This class is used to generate and manage simulation data for many-body physical systems. It simulates the time evolution of particle systems by numerically solving Newton's equations of motion. - + Main features: - Generate simulation data for multiple samples - Compute system acceleration (using PaddlePaddle automatic differentiation) - Visualize particle trajectories - Support custom potential energy functions - + Attributes: _sim (str): Simulation type. _n (int): Number of particles. @@ -189,12 +186,12 @@ class SimulationDataset(object): pairwise (function): Pairwise potential energy function. """ - def __init__(self, sim='r2', n=5, dim=2, - dt=0.01, nt=100, extra_potential=None, - **kwargs): + def __init__( + self, sim="r2", n=5, dim=2, dt=0.01, nt=100, extra_potential=None, **kwargs + ): """ Initialize simulation dataset. - + Args: sim (str, optional): Simulation type. Defaults to 'r2' (inverse square law). n (int, optional): Number of particles. Defaults to 5. @@ -211,317 +208,283 @@ def __init__(self, sim='r2', n=5, dim=2, self.dt = dt self.nt = nt self.data = None - - - self.times = np.linspace(0, self.dt*self.nt, num=self.nt) - + + self.times = np.linspace(0, self.dt * self.nt, num=self.nt) + self.G = 1 self.extra_potential = extra_potential - - + self.pairwise = get_potential(sim=sim, sim_obj=self) def simulate(self, ns, seed=0): """ Generate simulation data. - + This method is the core function of the simulator. For each sample: 1. Randomly initialize particle positions, velocities, masses, and charges 2. Solve Newton's equations of motion using Scipy ODE solver 3. Record the entire time evolution process - + Physical equations: d²x/dt² = F/m = -∇U/m where U is potential energy and F is force - + Args: ns (int): Number of samples to generate. seed (int, optional): Seed for numpy random number generator. Defaults to 0. - + Returns: None: Results are stored in self.data. """ - + np.random.seed(seed) - + n = self._n dim = self._dim sim = self._sim params = 2 - total_dim = dim*2 + params + total_dim = dim * 2 + params times = self.times G = self.G - + def compute_pairwise_potential(x1, x2_list): """Vectorized calculation of potential energy between one particle and multiple particles""" potentials = [] for x2 in x2_list: potentials.append(self.pairwise(x1, x2)) return np.array(potentials) - + def total_potential(xt): """ Calculate total potential energy of the system (using Numpy). - + Args: xt (numpy.ndarray): System state with shape (n, total_dim). - + Returns: float: Total potential energy of the system. """ sum_potential = 0.0 - + for i in range(n - 1): - if sim in ['string', 'string_ball']: + if sim in ["string", "string_ball"]: - sum_potential += G * self.pairwise(xt[i], xt[i+1]) + sum_potential += G * self.pairwise(xt[i], xt[i + 1]) else: - for j in range(i+1, n): + for j in range(i + 1, n): sum_potential += G * self.pairwise(xt[i], xt[j]) - + if self.extra_potential is not None: for i in range(n): sum_potential += self.extra_potential(xt[i]) - + return sum_potential - + def force_paddle(xt): """ Calculate force using PaddlePaddle automatic differentiation. - + This is a key part of the migration: JAX: force = -grad(total_potential)(xt) Paddle: Use paddle.grad to calculate gradient - + Complete implementation for all potential scenarios. - + Args: xt (numpy.ndarray): System state with shape (n, total_dim). - + Returns: numpy.ndarray: Force on each particle with shape (n, dim). """ - xt_tensor = paddle.to_tensor(xt, dtype='float64', stop_gradient=False) - + xt_tensor = paddle.to_tensor(xt, dtype="float64", stop_gradient=False) positions = xt_tensor[:, :dim] positions.stop_gradient = False - - xt_for_potential = paddle.concat([ - positions, - xt_tensor[:, dim:] - ], axis=1) - + xt_for_potential = paddle.concat([positions, xt_tensor[:, dim:]], axis=1) - sum_potential = paddle.to_tensor(0.0, dtype='float64') - + sum_potential = paddle.to_tensor(0.0, dtype="float64") for i in range(n - 1): - if sim in ['string', 'string_ball']: + if sim in ["string", "string_ball"]: x1 = xt_for_potential[i] - x2 = xt_for_potential[i+1] + x2 = xt_for_potential[i + 1] dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) bounded_dist = dist + 1e-2 - - if sim == 'string': - pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] - - elif sim == 'string_ball': + if sim == "string": + + pot = (bounded_dist - 1) ** 2 + x1[1] * x1[-1] - pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] - + elif sim == "string_ball": - r = paddle.sqrt((x1[1] + 15)**2 + (x1[0] - 5)**2) - radius = paddle.to_tensor(4.0, dtype='float64') - + pot = (bounded_dist - 1) ** 2 + x1[1] * x1[-1] + r = paddle.sqrt((x1[1] + 15) ** 2 + (x1[0] - 5) ** 2) + radius = paddle.to_tensor(4.0, dtype="float64") k_repel = 100.0 - pot = pot + k_repel / paddle.maximum(r - radius + 0.5, paddle.to_tensor(0.01, dtype='float64')) - + pot = pot + k_repel / paddle.maximum( + r - radius + 0.5, paddle.to_tensor(0.01, dtype="float64") + ) + sum_potential = sum_potential + G * pot - + else: - for j in range(i+1, n): + for j in range(i + 1, n): x1 = xt_for_potential[i] x2 = xt_for_potential[j] - dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) + dist = paddle.sqrt( + paddle.sum(paddle.square(x1[:dim] - x2[:dim])) + ) bounded_dist = dist + 1e-2 - - if sim == 'r2': + if sim == "r2": + + pot = -x1[-1] * x2[-1] / bounded_dist - pot = -x1[-1]*x2[-1]/bounded_dist - - elif sim == 'r1': + elif sim == "r1": - pot = x1[-1]*x2[-1]*paddle.log(bounded_dist) - - elif sim in ['spring', 'damped']: + pot = x1[-1] * x2[-1] * paddle.log(bounded_dist) - pot = (bounded_dist - 1)**2 - + elif sim in ["spring", "damped"]: - if sim == 'damped': - damping = paddle.to_tensor(1.0, dtype='float64') + pot = (bounded_dist - 1) ** 2 - pot = pot + damping*x1[1]*x1[1+dim]/n - pot = pot + damping*x1[0]*x1[0+dim]/n + if sim == "damped": + damping = paddle.to_tensor(1.0, dtype="float64") + + pot = pot + damping * x1[1] * x1[1 + dim] / n + pot = pot + damping * x1[0] * x1[0 + dim] / n if dim == 3: - pot = pot + damping*x1[2]*x1[2+dim]/n - - elif sim in ['charge', 'superposition']: + pot = pot + damping * x1[2] * x1[2 + dim] / n + + elif sim in ["charge", "superposition"]: charge1 = x1[-2] charge2 = x2[-2] - pot = charge1*charge2/bounded_dist - - if sim == 'superposition': + pot = charge1 * charge2 / bounded_dist + + if sim == "superposition": m1 = x1[-1] m2 = x2[-1] - pot = pot + (-m1*m2/bounded_dist) - - elif sim == 'discontinuous': + pot = pot + (-m1 * m2 / bounded_dist) + elif sim == "discontinuous": - pot_a = paddle.to_tensor(0.0, dtype='float64') - pot_b = paddle.to_tensor(0.0, dtype='float64') - pot_c = (bounded_dist - 1)**2 - + pot_a = paddle.to_tensor(0.0, dtype="float64") + pot_b = paddle.to_tensor(0.0, dtype="float64") + pot_c = (bounded_dist - 1) ** 2 + cond1 = paddle.cast(bounded_dist < 1, "float64") + cond2 = paddle.cast( + bounded_dist >= 1, "float64" + ) * paddle.cast(bounded_dist < 2, "float64") + cond3 = paddle.cast(bounded_dist >= 2, "float64") - cond1 = paddle.cast(bounded_dist < 1, 'float64') - cond2 = paddle.cast(bounded_dist >= 1, 'float64') * paddle.cast(bounded_dist < 2, 'float64') - cond3 = paddle.cast(bounded_dist >= 2, 'float64') - pot = pot_a * cond1 + pot_b * cond2 + pot_c * cond3 - + else: - pot = (bounded_dist - 1)**2 - + pot = (bounded_dist - 1) ** 2 + sum_potential = sum_potential + G * pot - if self.extra_potential is not None: - - for i in range(n): - - try: - extra_pot = self.extra_potential(xt_for_potential[i]) - if not isinstance(extra_pot, paddle.Tensor): - extra_pot = paddle.to_tensor(extra_pot, dtype='float64') - sum_potential = sum_potential + extra_pot - except: - - - pass - - + extra_pot = self.extra_potential(xt_for_potential[i]) + if not isinstance(extra_pot, paddle.Tensor): + extra_pot = paddle.to_tensor(extra_pot, dtype="float64") + sum_potential = sum_potential + extra_pot grads = paddle.grad( outputs=sum_potential, inputs=positions, create_graph=False, - retain_graph=False + retain_graph=False, )[0] - force = -grads.numpy() - + return force - + def acceleration(xt): """ Calculate acceleration: a = F/m. - + Args: xt (numpy.ndarray): System state. - + Returns: numpy.ndarray: Acceleration. """ force = force_paddle(xt) masses = xt[:, -1:] return force / masses - + def odefunc(y, t): """ Ordinary differential equation function (for scipy.odeint). - + Args: y (numpy.ndarray): Flattened system state. t (float): Time. - + Returns: numpy.ndarray: Time derivative of the state. """ y = y.reshape((n, total_dim)) - a = acceleration(y) - - - dydt = np.concatenate([ - y[:, dim:2*dim], - a, - np.zeros((n, params)) - ], axis=1) - + + dydt = np.concatenate( + [y[:, dim : 2 * dim], a, np.zeros((n, params))], axis=1 + ) + return dydt.flatten() - + def make_sim(sample_idx): """ Generate single simulation sample. - + Args: sample_idx (int): Sample index (for progress display). - + Returns: numpy.ndarray: Simulation trajectory with shape (nt, n, total_dim). """ - if sim in ['string', 'string_ball']: + if sim in ["string", "string_ball"]: x0 = np.random.randn(n, total_dim) x0[:, -1] = 1.0 - x0[:, 0] = np.arange(n) + x0[:, 0]*0.5 + x0[:, 0] = np.arange(n) + x0[:, 0] * 0.5 x0[:, 2:3] = 0.0 else: x0 = np.random.randn(n, total_dim) x0[:, -1] = np.exp(x0[:, -1]) - - if sim in ['charge', 'superposition']: + + if sim in ["charge", "superposition"]: x0[:, -2] = np.sign(x0[:, -2]) - + x_times = odeint(odefunc, x0.flatten(), times, mxstep=2000).reshape( + -1, n, total_dim + ) - x_times = odeint( - odefunc, - x0.flatten(), - times, - mxstep=2000 - ).reshape(-1, n, total_dim) - return x_times - data = [] print(f"Start generating {ns} samples...") for i in tqdm(range(ns)): data.append(make_sim(i)) - self.data = np.array(data) print(f"Generation complete! Data shape: {self.data.shape}") @@ -529,203 +492,214 @@ def make_sim(sample_idx): def get_acceleration(self): """ Calculate acceleration for all simulation data. - + This method uses PaddlePaddle automatic differentiation to calculate acceleration. Uses the same complete potential implementation as in simulate(). - + Returns: numpy.ndarray: Acceleration data with shape (ns, nt, n, dim). """ if self.data is None: raise ValueError("Please call simulate() first to generate data") - + ns, nt, n, total_dim = self.data.shape dim = self._dim sim = self._sim G = self.G - accelerations = np.zeros((ns, nt, n, dim)) - + print("Calculating acceleration...") for sample_idx in tqdm(range(ns)): for time_idx in range(nt): xt = self.data[sample_idx, time_idx] - - xt_tensor = paddle.to_tensor(xt, dtype='float64', stop_gradient=False) + xt_tensor = paddle.to_tensor(xt, dtype="float64", stop_gradient=False) positions = xt_tensor[:, :dim] positions.stop_gradient = False - - xt_for_potential = paddle.concat([ - positions, - xt_tensor[:, dim:] - ], axis=1) - - sum_potential = paddle.to_tensor(0.0, dtype='float64') - + xt_for_potential = paddle.concat( + [positions, xt_tensor[:, dim:]], axis=1 + ) + + sum_potential = paddle.to_tensor(0.0, dtype="float64") for i in range(n - 1): - if sim in ['string', 'string_ball']: + if sim in ["string", "string_ball"]: x1 = xt_for_potential[i] - x2 = xt_for_potential[i+1] - dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) + x2 = xt_for_potential[i + 1] + dist = paddle.sqrt( + paddle.sum(paddle.square(x1[:dim] - x2[:dim])) + ) bounded_dist = dist + 1e-2 - - if sim == 'string': - pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] - elif sim == 'string_ball': - pot = (bounded_dist - 1)**2 + x1[1]*x1[-1] - r = paddle.sqrt((x1[1] + 15)**2 + (x1[0] - 5)**2) - radius = paddle.to_tensor(4.0, dtype='float64') + if sim == "string": + pot = (bounded_dist - 1) ** 2 + x1[1] * x1[-1] + elif sim == "string_ball": + + pot = (bounded_dist - 1) ** 2 + x1[1] * x1[-1] + r = paddle.sqrt((x1[1] + 15) ** 2 + (x1[0] - 5) ** 2) + radius = paddle.to_tensor(4.0, dtype="float64") k_repel = 100.0 - pot = pot + k_repel / paddle.maximum(r - radius + 0.5, paddle.to_tensor(0.01, dtype='float64')) - + pot = pot + k_repel / paddle.maximum( + r - radius + 0.5, + paddle.to_tensor(0.01, dtype="float64"), + ) + sum_potential = sum_potential + G * pot else: - for j in range(i+1, n): + for j in range(i + 1, n): x1 = xt_for_potential[i] x2 = xt_for_potential[j] - dist = paddle.sqrt(paddle.sum(paddle.square(x1[:dim] - x2[:dim]))) + dist = paddle.sqrt( + paddle.sum(paddle.square(x1[:dim] - x2[:dim])) + ) bounded_dist = dist + 1e-2 - - if sim == 'r2': - pot = -x1[-1]*x2[-1]/bounded_dist - elif sim == 'r1': - pot = x1[-1]*x2[-1]*paddle.log(bounded_dist) - elif sim in ['spring', 'damped']: - pot = (bounded_dist - 1)**2 - - - if sim == 'damped': - damping = paddle.to_tensor(1.0, dtype='float64') - pot = pot + damping*x1[1]*x1[1+dim]/n - pot = pot + damping*x1[0]*x1[0+dim]/n + + if sim == "r2": + pot = -x1[-1] * x2[-1] / bounded_dist + elif sim == "r1": + pot = x1[-1] * x2[-1] * paddle.log(bounded_dist) + elif sim in ["spring", "damped"]: + pot = (bounded_dist - 1) ** 2 + + if sim == "damped": + damping = paddle.to_tensor(1.0, dtype="float64") + pot = pot + damping * x1[1] * x1[1 + dim] / n + pot = pot + damping * x1[0] * x1[0 + dim] / n if dim == 3: - pot = pot + damping*x1[2]*x1[2+dim]/n - - elif sim in ['charge', 'superposition']: + pot = pot + damping * x1[2] * x1[2 + dim] / n + + elif sim in ["charge", "superposition"]: charge1 = x1[-2] charge2 = x2[-2] - pot = charge1*charge2/bounded_dist - if sim == 'superposition': + pot = charge1 * charge2 / bounded_dist + if sim == "superposition": m1 = x1[-1] m2 = x2[-1] - pot = pot + (-m1*m2/bounded_dist) - - elif sim == 'discontinuous': - - pot_a = paddle.to_tensor(0.0, dtype='float64') - pot_b = paddle.to_tensor(0.0, dtype='float64') - pot_c = (bounded_dist - 1)**2 - - cond1 = paddle.cast(bounded_dist < 1, 'float64') - cond2 = paddle.cast(bounded_dist >= 1, 'float64') * paddle.cast(bounded_dist < 2, 'float64') - cond3 = paddle.cast(bounded_dist >= 2, 'float64') - + pot = pot + (-m1 * m2 / bounded_dist) + + elif sim == "discontinuous": + + pot_a = paddle.to_tensor(0.0, dtype="float64") + pot_b = paddle.to_tensor(0.0, dtype="float64") + pot_c = (bounded_dist - 1) ** 2 + + cond1 = paddle.cast(bounded_dist < 1, "float64") + cond2 = paddle.cast( + bounded_dist >= 1, "float64" + ) * paddle.cast(bounded_dist < 2, "float64") + cond3 = paddle.cast(bounded_dist >= 2, "float64") + pot = pot_a * cond1 + pot_b * cond2 + pot_c * cond3 - + else: - pot = (bounded_dist - 1)**2 - + pot = (bounded_dist - 1) ** 2 + sum_potential = sum_potential + G * pot - if self.extra_potential is not None: for i in range(n): - try: - extra_pot = self.extra_potential(xt_for_potential[i]) - if not isinstance(extra_pot, paddle.Tensor): - extra_pot = paddle.to_tensor(extra_pot, dtype='float64') - sum_potential = sum_potential + extra_pot - except: - pass - + extra_pot = self.extra_potential(xt_for_potential[i]) + if not isinstance(extra_pot, paddle.Tensor): + extra_pot = paddle.to_tensor(extra_pot, dtype="float64") + sum_potential = sum_potential + extra_pot grads = paddle.grad( outputs=sum_potential, inputs=positions, create_graph=False, - retain_graph=False + retain_graph=False, )[0] - + force = -grads.numpy() masses = xt[:, -1:] accel = force / masses - + accelerations[sample_idx, time_idx] = accel - + return accelerations def plot(self, i, animate=False, plot_size=True, s_size=1): """ Visualize simulation results. - + Args: i (int): Sample index. animate (bool, optional): Whether to generate animation. Defaults to False. plot_size (bool, optional): Whether to adjust point size based on mass. Defaults to True. s_size (float, optional): Point size scaling factor. Defaults to 1. - + Returns: None (static plot) or HTML (animation). """ if self.data is None: raise ValueError("Please call simulate() first to generate data") - + n = self._n times = self.times x_times = self.data[i] sim = self._sim masses = x_times[:, :, -1] - + if not animate: - if sim in ['string', 'string_ball']: + if sim in ["string", "string_ball"]: rgba = make_transparent_color(len(times), 0) - for idx in range(0, len(times), len(times)//10): + for idx in range(0, len(times), len(times) // 10): ctimes = x_times[idx] plt.plot(ctimes[:, 0], ctimes[:, 1], color=rgba[idx]) plt.xlim(-5, 20) plt.ylim(-20, 5) else: for j in range(n): - rgba = make_transparent_color(len(times), j/n) + rgba = make_transparent_color(len(times), j / n) if plot_size: - plt.scatter(x_times[:, j, 0], x_times[:, j, 1], - color=rgba, s=3*masses[:, j]*s_size) + plt.scatter( + x_times[:, j, 0], + x_times[:, j, 1], + color=rgba, + s=3 * masses[:, j] * s_size, + ) else: - plt.scatter(x_times[:, j, 0], x_times[:, j, 1], - color=rgba, s=s_size) + plt.scatter( + x_times[:, j, 0], x_times[:, j, 1], color=rgba, s=s_size + ) else: - if sim in ['string', 'string_ball']: - raise NotImplementedError("Animation mode not yet supported for string type") - + if sim in ["string", "string_ball"]: + raise NotImplementedError( + "Animation mode not yet supported for string type" + ) + fig = plt.figure() camera = Camera(fig) d_idx = 20 - + for t_idx in range(d_idx, len(times), d_idx): - start = max([0, t_idx-300]) + start = max([0, t_idx - 300]) ctimes = times[start:t_idx] cx_times = x_times[start:t_idx] - + for j in range(n): - rgba = make_transparent_color(len(ctimes), j/n) + rgba = make_transparent_color(len(ctimes), j / n) if plot_size: - plt.scatter(cx_times[:, j, 0], cx_times[:, j, 1], - color=rgba, s=3*masses[:, j]) + plt.scatter( + cx_times[:, j, 0], + cx_times[:, j, 1], + color=rgba, + s=3 * masses[:, j], + ) else: - plt.scatter(cx_times[:, j, 0], cx_times[:, j, 1], - color=rgba, s=s_size) - + plt.scatter( + cx_times[:, j, 0], cx_times[:, j, 1], color=rgba, s=s_size + ) + camera.snap() - + from IPython.display import HTML - return HTML(camera.animate().to_jshtml()) \ No newline at end of file + + return HTML(camera.animate().to_jshtml()) diff --git a/ppsci/arch/symbolic_gn.py b/ppsci/arch/symbolic_gn.py index edbc01d26b..c8b230132c 100644 --- a/ppsci/arch/symbolic_gn.py +++ b/ppsci/arch/symbolic_gn.py @@ -12,35 +12,43 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple + +import numpy as np import paddle import paddle.nn as nn -import numpy as np -from typing import Dict, List, Optional, Tuple from ppsci.arch import base + def get_edge_index(n: int, sim: str) -> np.ndarray: """ Generate edge indices for graph connectivity. - + Args: n (int): Number of nodes. sim (str): Simulation type. - + Returns: numpy.ndarray: Edge indices with shape [2, num_edges]. """ - if sim in ['string', 'string_ball']: - top = np.arange(0, n-1) + if sim in ["string", "string_ball"]: + top = np.arange(0, n - 1) bottom = np.arange(1, n) - edge_index = np.concatenate([ - np.concatenate([top, bottom])[None, :], - np.concatenate([bottom, top])[None, :] - ], axis=0) + edge_index = np.concatenate( + [ + np.concatenate([top, bottom])[None, :], + np.concatenate([bottom, top])[None, :], + ], + axis=0, + ) else: adj = (np.ones((n, n)) - np.eye(n)).astype(int) edge_index = np.array(np.where(adj)) - + return edge_index @@ -54,12 +62,12 @@ def __init__( ndim: int = 2, hidden: int = 300, edge_index: Optional[np.ndarray] = None, - aggr: str = 'sum', - l1_strength: float = 0.0 + aggr: str = "sum", + l1_strength: float = 0.0, ): """ Initialize Object-based Graph Network (OGN). - + Args: input_keys (List[str]): List of input keys, e.g., ['x', 'edge_index']. output_keys (List[str]): List of output keys, e.g., ['acceleration']. @@ -69,7 +77,7 @@ def __init__( hidden (int, optional): Hidden layer size. Defaults to 300. edge_index (np.ndarray, optional): Edge indices (can also be provided in forward). Defaults to None. aggr (str, optional): Aggregation method, 'sum' or 'mean'. Defaults to 'sum'. - + Examples: >>> import ppsci >>> import numpy as np @@ -90,33 +98,34 @@ def __init__( [5, 2] """ super().__init__() - + self.input_keys = input_keys self.output_keys = output_keys - + self.n_f = n_f self.msg_dim = msg_dim self.ndim = ndim self.hidden = hidden self.aggr = aggr self.l1_strength = l1_strength - + if edge_index is not None: - self.register_buffer('edge_index_buffer', - paddle.to_tensor(edge_index, dtype='int64')) + self.register_buffer( + "edge_index_buffer", paddle.to_tensor(edge_index, dtype="int64") + ) else: self.edge_index_buffer = None - + self.msg_fnc = nn.Sequential( - nn.Linear(2*n_f, hidden), + nn.Linear(2 * n_f, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), - nn.Linear(hidden, msg_dim) + nn.Linear(hidden, msg_dim), ) - + self.node_fnc = nn.Sequential( nn.Linear(msg_dim + n_f, hidden), nn.ReLU(), @@ -124,87 +133,89 @@ def __init__( nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), - nn.Linear(hidden, ndim) + nn.Linear(hidden, ndim), ) - - def message_passing(self, x: paddle.Tensor, edge_index: np.ndarray) -> paddle.Tensor: + + def message_passing( + self, x: paddle.Tensor, edge_index: np.ndarray + ) -> paddle.Tensor: """ Execute message passing. - + Args: x: Node features with shape [n, n_f] or [batch*n, n_f]. edge_index: Edge indices with shape [2, num_edges]. - + Returns: paddle.Tensor: Updated node features with shape [n, ndim] or [batch*n, ndim]. """ if len(x.shape) == 3: batch_size, n, n_f = x.shape x_reshaped = x.reshape([-1, n_f]) - + results = [] for b in range(batch_size): start_idx = b * n end_idx = (b + 1) * n x_batch = x_reshaped[start_idx:end_idx] - + row, col = edge_index[0], edge_index[1] - + x_i = x_batch[col] x_j = x_batch[row] - + msg_input = paddle.concat([x_i, x_j], axis=1) msg = self.msg_fnc(msg_input) - + aggr_out = paddle.zeros([n, self.msg_dim], dtype=msg.dtype) for i in range(len(col)): aggr_out[col[i]] += msg[i] - + node_input = paddle.concat([x_batch, aggr_out], axis=1) out = self.node_fnc(node_input) results.append(out) - + return paddle.stack(results, axis=0) - + else: row, col = edge_index[0], edge_index[1] - + x_i = x[col] x_j = x[row] - + msg_input = paddle.concat([x_i, x_j], axis=1) msg = self.msg_fnc(msg_input) - + num_nodes = x.shape[0] aggr_out = paddle.zeros([num_nodes, self.msg_dim], dtype=msg.dtype) - + for i in range(len(col)): aggr_out[col[i]] += msg[i] - + node_input = paddle.concat([x, aggr_out], axis=1) out = self.node_fnc(node_input) - + return out - + def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: """ Forward propagation (PaddleScience standard interface). - + Args: inputs (Dict): Input dictionary containing: - 'x': Node features [batch*n, n_f] or [n, n_f]. - 'edge_index': Edge index information (optional if provided at initialization). - Other possible inputs. - + Returns: Dict: Output dictionary containing: - 'acceleration': Predicted acceleration [batch*n, ndim]. - 'l1_regularization': L1 regularization term (if enabled). """ - x = inputs['x'] - - if 'edge_index' in inputs: - edge_index = inputs['edge_index'] + x = inputs["x"] + + if "edge_index" in inputs: + edge_index = inputs["edge_index"] if isinstance(edge_index, paddle.Tensor): edge_index = edge_index.numpy() if len(edge_index.shape) == 3: @@ -213,44 +224,41 @@ def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: edge_index = self.edge_index_buffer.numpy() else: raise ValueError("Must provide edge_index") - + acceleration = self.message_passing(x, edge_index) - + outputs = {self.output_keys[0]: acceleration} - - if hasattr(self, 'l1_strength') and self.l1_strength > 0: + + if hasattr(self, "l1_strength") and self.l1_strength > 0: l1_reg = self.l1_strength * paddle.mean(paddle.abs(x)) - outputs['l1_regularization'] = l1_reg - + outputs["l1_regularization"] = l1_reg + return outputs - + def compute_loss( - self, - inputs: Dict[str, paddle.Tensor], - labels: Dict[str, paddle.Tensor] + self, inputs: Dict[str, paddle.Tensor], labels: Dict[str, paddle.Tensor] ) -> paddle.Tensor: """ Compute loss function (PaddleScience standard interface). - + Args: inputs: Input dictionary. labels: Label dictionary containing 'acceleration_true'. - + Returns: paddle.Tensor: Loss scalar. """ outputs = self.forward(inputs) - + pred = outputs[self.output_keys[0]] - true = labels.get('acceleration_true', labels.get(self.output_keys[0])) - + true = labels.get("acceleration_true", labels.get(self.output_keys[0])) + loss = paddle.mean(paddle.abs(pred - true)) - + return loss class HGN(base.Arch): - def __init__( self, input_keys: List[str], @@ -258,11 +266,11 @@ def __init__( n_f: int = 6, ndim: int = 2, hidden: int = 300, - edge_index: Optional[np.ndarray] = None + edge_index: Optional[np.ndarray] = None, ): """ Initialize Hamiltonian Graph Network (HGN). - + Args: input_keys: List of input keys. output_keys: List of output keys, e.g., ['velocity_derivative', 'acceleration']. @@ -270,7 +278,7 @@ def __init__( ndim: Spatial dimension. hidden: Hidden layer size. edge_index: Edge indices (optional). - + Examples: >>> import ppsci >>> import numpy as np @@ -291,29 +299,30 @@ def __init__( [5, 2] """ super().__init__() - + self.input_keys = input_keys self.output_keys = output_keys self.n_f = n_f self.ndim = ndim self.hidden = hidden - + if edge_index is not None: - self.register_buffer('edge_index_buffer', - paddle.to_tensor(edge_index, dtype='int64')) + self.register_buffer( + "edge_index_buffer", paddle.to_tensor(edge_index, dtype="int64") + ) else: self.edge_index_buffer = None - + self.pair_energy = nn.Sequential( - nn.Linear(2*n_f, hidden), + nn.Linear(2 * n_f, hidden), nn.Softplus(), nn.Linear(hidden, hidden), nn.Softplus(), nn.Linear(hidden, hidden), nn.Softplus(), - nn.Linear(hidden, 1) + nn.Linear(hidden, 1), ) - + self.self_energy = nn.Sequential( nn.Linear(n_f, hidden), nn.Softplus(), @@ -321,54 +330,54 @@ def __init__( nn.Softplus(), nn.Linear(hidden, hidden), nn.Softplus(), - nn.Linear(hidden, 1) + nn.Linear(hidden, 1), ) - + def compute_energy(self, x: paddle.Tensor, edge_index: np.ndarray) -> paddle.Tensor: """ Compute the total energy (Hamiltonian) of the system. - + Args: x: Node features. edge_index: Edge indices. - + Returns: paddle.Tensor: Energy per node. """ row, col = edge_index[0], edge_index[1] - + x_i = x[col] x_j = x[row] edge_input = paddle.concat([x_i, x_j], axis=1) pair_energies = self.pair_energy(edge_input) - + num_nodes = x.shape[0] aggr_pair = paddle.zeros([num_nodes, 1], dtype=pair_energies.dtype) for i in range(len(col)): aggr_pair[col[i]] += pair_energies[i] - + self_energies = self.self_energy(x) - + total_energy = aggr_pair + self_energies - + return total_energy - + def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: """ Forward propagation (PaddleScience standard). - + Compute dynamics derivatives using Hamilton's equations. - + Args: inputs: Input dictionary. - + Returns: Dict: Output dictionary. """ - x_input = inputs['x'].clone() - - if 'edge_index' in inputs: - edge_index = inputs['edge_index'] + x_input = inputs["x"].clone() + + if "edge_index" in inputs: + edge_index = inputs["edge_index"] if isinstance(edge_index, paddle.Tensor): edge_index = edge_index.numpy() if len(edge_index.shape) == 3: @@ -377,158 +386,157 @@ def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: edge_index = self.edge_index_buffer.numpy() else: raise ValueError("Must provide edge_index") - + if len(x_input.shape) == 3: batch_size, n, n_f = x_input.shape x_reshaped = x_input.reshape([-1, n_f]) - + results = [] for b in range(batch_size): start_idx = b * n end_idx = (b + 1) * n x_batch = x_reshaped[start_idx:end_idx] - - q = x_batch[:, :self.ndim] - v = x_batch[:, self.ndim:2*self.ndim] - other = x_batch[:, 2*self.ndim:] - + + q = x_batch[:, : self.ndim] + v = x_batch[:, self.ndim : 2 * self.ndim] + other = x_batch[:, 2 * self.ndim :] + m_scalar = other[:, -1:] m_vec = paddle.tile(m_scalar, [1, self.ndim]) - + p = v * m_vec - + x_hamilton = paddle.concat([q, p, other], axis=1) x_hamilton.stop_gradient = False - + total_energy = self.compute_energy(x_hamilton, edge_index) total_energy_scalar = paddle.sum(total_energy) - + dH = paddle.grad( outputs=total_energy_scalar, inputs=x_hamilton, create_graph=False, - retain_graph=False + retain_graph=False, )[0] - - dH_dq = dH[:, :self.ndim] - dH_dp = dH[:, self.ndim:2*self.ndim] - + + dH_dq = dH[:, : self.ndim] + dH_dp = dH[:, self.ndim : 2 * self.ndim] + dq_dt = dH_dp dp_dt = -dH_dq dv_dt = dp_dt / m_vec - + derivative = paddle.concat([dq_dt, dv_dt], axis=1) results.append(derivative) - + derivative = paddle.stack(results, axis=0) - + outputs = {} - if 'velocity_derivative' in self.output_keys: - outputs['velocity_derivative'] = derivative[:, :, :self.ndim] - if 'acceleration' in self.output_keys: - outputs['acceleration'] = derivative[:, :, self.ndim:] - + if "velocity_derivative" in self.output_keys: + outputs["velocity_derivative"] = derivative[:, :, : self.ndim] + if "acceleration" in self.output_keys: + outputs["acceleration"] = derivative[:, :, self.ndim :] + if len(self.output_keys) == 1: - if self.output_keys[0] == 'acceleration': - outputs['acceleration'] = derivative[:, :, self.ndim:] + if self.output_keys[0] == "acceleration": + outputs["acceleration"] = derivative[:, :, self.ndim :] else: outputs[self.output_keys[0]] = derivative - + return outputs - + else: - q = x_input[:, :self.ndim] - v = x_input[:, self.ndim:2*self.ndim] - other = x_input[:, 2*self.ndim:] - + q = x_input[:, : self.ndim] + v = x_input[:, self.ndim : 2 * self.ndim] + other = x_input[:, 2 * self.ndim :] + m_scalar = other[:, -1:] m_vec = paddle.tile(m_scalar, [1, self.ndim]) - + p = v * m_vec - + x_hamilton = paddle.concat([q, p, other], axis=1) x_hamilton.stop_gradient = False - + total_energy = self.compute_energy(x_hamilton, edge_index) total_energy_scalar = paddle.sum(total_energy) - + dH = paddle.grad( outputs=total_energy_scalar, inputs=x_hamilton, create_graph=False, - retain_graph=False + retain_graph=False, )[0] - - dH_dq = dH[:, :self.ndim] - dH_dp = dH[:, self.ndim:2*self.ndim] - + + dH_dq = dH[:, : self.ndim] + dH_dp = dH[:, self.ndim : 2 * self.ndim] + dq_dt = dH_dp dp_dt = -dH_dq dv_dt = dp_dt / m_vec - + outputs = {} - if 'velocity_derivative' in self.output_keys: - outputs['velocity_derivative'] = dq_dt - if 'acceleration' in self.output_keys: - outputs['acceleration'] = dv_dt - + if "velocity_derivative" in self.output_keys: + outputs["velocity_derivative"] = dq_dt + if "acceleration" in self.output_keys: + outputs["acceleration"] = dv_dt + if len(self.output_keys) == 1: - if self.output_keys[0] == 'acceleration': - outputs['acceleration'] = dv_dt + if self.output_keys[0] == "acceleration": + outputs["acceleration"] = dv_dt else: outputs[self.output_keys[0]] = paddle.concat([dq_dt, dv_dt], axis=1) - + return outputs - + def compute_loss( self, inputs: Dict[str, paddle.Tensor], labels: Dict[str, paddle.Tensor], - reg_weight: float = 1e-6 + reg_weight: float = 1e-6, ) -> paddle.Tensor: """ Compute loss with physical regularization. - + Args: inputs: Input dictionary. labels: Label dictionary. reg_weight: Regularization weight. - + Returns: paddle.Tensor: Total loss. """ outputs = self.forward(inputs) - - pred = outputs.get('acceleration', outputs.get(self.output_keys[0])) - true = labels.get('acceleration_true', labels.get('acceleration')) - + + pred = outputs.get("acceleration", outputs.get(self.output_keys[0])) + true = labels.get("acceleration_true", labels.get("acceleration")) + base_loss = paddle.mean(paddle.abs(pred - true)) - - x_input = inputs['x'].clone() - edge_index = inputs.get('edge_index', self.edge_index_buffer.numpy()) + + x_input = inputs["x"].clone() + edge_index = inputs.get("edge_index", self.edge_index_buffer.numpy()) if isinstance(edge_index, paddle.Tensor): edge_index = edge_index.numpy() - - q = x_input[:, :self.ndim] - v = x_input[:, self.ndim:2*self.ndim] - other = x_input[:, 2*self.ndim:] - + + q = x_input[:, : self.ndim] + v = x_input[:, self.ndim : 2 * self.ndim] + other = x_input[:, 2 * self.ndim :] + m_scalar = other[:, -1:] m_vec = paddle.tile(m_scalar, [1, self.ndim]) p = v * m_vec - + x_hamilton = paddle.concat([q, p, other], axis=1) x_hamilton.stop_gradient = False - + total_energy = self.compute_energy(x_hamilton, edge_index) - + regularization = reg_weight * paddle.mean(total_energy**2) - + return base_loss + regularization class VarOGN(base.Arch): - def __init__( self, input_keys: List[str], @@ -539,11 +547,11 @@ def __init__( hidden: int = 300, edge_index: Optional[np.ndarray] = None, enable_sampling: bool = True, - l1_strength: float = 0.0 + l1_strength: float = 0.0, ): """ Initialize Variational Graph Network (VarGN). - + Args: input_keys: List of input keys. output_keys: List of output keys. @@ -553,7 +561,7 @@ def __init__( hidden: Hidden layer size. edge_index: Edge indices. enable_sampling: Whether to enable sampling. - + Examples: >>> import ppsci >>> import numpy as np @@ -574,7 +582,7 @@ def __init__( [5, 2] """ super().__init__() - + self.input_keys = input_keys self.output_keys = output_keys self.n_f = n_f @@ -583,23 +591,24 @@ def __init__( self.hidden = hidden self.enable_sampling = enable_sampling self.l1_strength = l1_strength - + if edge_index is not None: - self.register_buffer('edge_index_buffer', - paddle.to_tensor(edge_index, dtype='int64')) + self.register_buffer( + "edge_index_buffer", paddle.to_tensor(edge_index, dtype="int64") + ) else: self.edge_index_buffer = None - + self.msg_fnc = nn.Sequential( - nn.Linear(2*n_f, hidden), + nn.Linear(2 * n_f, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), - nn.Linear(hidden, msg_dim*2) + nn.Linear(hidden, msg_dim * 2), ) - + self.node_fnc = nn.Sequential( nn.Linear(msg_dim + n_f, hidden), nn.ReLU(), @@ -607,24 +616,25 @@ def __init__( nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), - nn.Linear(hidden, ndim) + nn.Linear(hidden, ndim), ) - - def message_passing(self, x: paddle.Tensor, edge_index: np.ndarray) -> Tuple[paddle.Tensor, paddle.Tensor]: + + def message_passing( + self, x: paddle.Tensor, edge_index: np.ndarray + ) -> Tuple[paddle.Tensor, paddle.Tensor]: """ Variational message passing. - + Args: x: Node features with shape [n, n_f] or [batch*n, n_f]. edge_index: Edge indices with shape [2, num_edges]. - + Returns: Tuple[paddle.Tensor, paddle.Tensor]: (mean output, variance output). """ if len(x.shape) == 3: batch_size, n, n_f = x.shape - x_reshaped = x.reshape([-1, n_f]) - + x_reshaped = x.reshape([-1, n_f]) results_mean = [] results_std = [] @@ -632,88 +642,88 @@ def message_passing(self, x: paddle.Tensor, edge_index: np.ndarray) -> Tuple[pad start_idx = b * n end_idx = (b + 1) * n x_batch = x_reshaped[start_idx:end_idx] - + row, col = edge_index[0], edge_index[1] - + x_i = x_batch[col] x_j = x_batch[row] - + msg_input = paddle.concat([x_i, x_j], axis=1) raw_msg = self.msg_fnc(msg_input) - + mu = raw_msg[:, 0::2] logvar = raw_msg[:, 1::2] - + if self.enable_sampling and self.training: epsilon = paddle.randn(mu.shape) msg = mu + epsilon * paddle.exp(logvar / 2.0) else: msg = mu - + aggr_out = paddle.zeros([n, self.msg_dim], dtype=msg.dtype) for i in range(len(col)): aggr_out[col[i]] += msg[i] - + node_input = paddle.concat([x_batch, aggr_out], axis=1) out_mean = self.node_fnc(node_input) - + # Calculate std as a small constant value for each output dimension # This provides a reasonable uncertainty estimate out_std = paddle.ones([n, self.ndim], dtype=out_mean.dtype) * 0.1 - + results_mean.append(out_mean) results_std.append(out_std) - + return paddle.stack(results_mean, axis=0), paddle.stack(results_std, axis=0) - + else: row, col = edge_index[0], edge_index[1] - + x_i = x[col] x_j = x[row] - + msg_input = paddle.concat([x_i, x_j], axis=1) raw_msg = self.msg_fnc(msg_input) - + mu = raw_msg[:, 0::2] logvar = raw_msg[:, 1::2] - + if self.enable_sampling and self.training: epsilon = paddle.randn(mu.shape) msg = mu + epsilon * paddle.exp(logvar / 2.0) else: msg = mu - + num_nodes = x.shape[0] aggr_out = paddle.zeros([num_nodes, self.msg_dim], dtype=msg.dtype) - + for i in range(len(col)): aggr_out[col[i]] += msg[i] - + node_input = paddle.concat([x, aggr_out], axis=1) out_mean = self.node_fnc(node_input) - + # Calculate std as a small constant value for each output dimension # This provides a reasonable uncertainty estimate num_nodes = x.shape[0] out_std = paddle.ones([num_nodes, self.ndim], dtype=out_mean.dtype) * 0.1 - + return out_mean, out_std - + def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: """ Forward propagation (PaddleScience standard). - + Args: inputs: Input dictionary. - + Returns: Dict: Output dictionary containing mean and standard deviation. """ - x = inputs['x'] - - if 'edge_index' in inputs: - edge_index = inputs['edge_index'] + x = inputs["x"] + + if "edge_index" in inputs: + edge_index = inputs["edge_index"] if isinstance(edge_index, paddle.Tensor): edge_index = edge_index.numpy() if len(edge_index.shape) == 3: @@ -722,48 +732,48 @@ def forward(self, inputs: Dict[str, paddle.Tensor]) -> Dict[str, paddle.Tensor]: edge_index = self.edge_index_buffer.numpy() else: raise ValueError("Must provide edge_index") - + accel_mean, accel_std = self.message_passing(x, edge_index) - + outputs = {} - if 'acceleration_mean' in self.output_keys: - outputs['acceleration_mean'] = accel_mean - if 'acceleration_std' in self.output_keys: - outputs['acceleration_std'] = accel_std - + if "acceleration_mean" in self.output_keys: + outputs["acceleration_mean"] = accel_mean + if "acceleration_std" in self.output_keys: + outputs["acceleration_std"] = accel_std + if len(outputs) == 0: outputs[self.output_keys[0]] = accel_mean - - if hasattr(self, 'l1_strength') and self.l1_strength > 0: + + if hasattr(self, "l1_strength") and self.l1_strength > 0: l1_reg = self.l1_strength * paddle.mean(paddle.abs(x)) - outputs['l1_regularization'] = l1_reg - + outputs["l1_regularization"] = l1_reg + return outputs - + def compute_loss( self, inputs: Dict[str, paddle.Tensor], labels: Dict[str, paddle.Tensor], - kl_weight: float = 1e-3 + kl_weight: float = 1e-3, ) -> paddle.Tensor: """ Compute variational loss including KL divergence. - + Args: inputs: Input dictionary. labels: Label dictionary. kl_weight: KL divergence weight. - + Returns: paddle.Tensor: Total loss. """ outputs = self.forward(inputs) - - pred = outputs.get('acceleration_mean', outputs.get(self.output_keys[0])) - true = labels.get('acceleration_true', labels.get('acceleration')) - + + pred = outputs.get("acceleration_mean", outputs.get(self.output_keys[0])) + true = labels.get("acceleration_true", labels.get("acceleration")) + recon_loss = paddle.mean(paddle.abs(pred - true)) - + kl_loss = paddle.to_tensor(0.0) - + return recon_loss + kl_weight * kl_loss From d0a201125b7e47925fd6ef8774bb4e78057e6a85 Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Fri, 5 Dec 2025 23:26:57 +0800 Subject: [PATCH 08/10] feat:update --- ppsci/arch/__init__.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ppsci/arch/__init__.py b/ppsci/arch/__init__.py index 06eb28b8d8..67c5b026a4 100644 --- a/ppsci/arch/__init__.py +++ b/ppsci/arch/__init__.py @@ -54,7 +54,7 @@ from ppsci.arch.physx_transformer import PhysformerGPT2 # isort:skip from ppsci.arch.sfnonet import SFNONet # isort:skip from ppsci.arch.spinn import SPINN # isort:skip -from ppsci.arch.symbolic_gn import OGN, VarOGN, HGN, get_edge_index # isort:skip +from ppsci.arch.symbolic_gn import OGN, VarOGN, HGN # isort:skip from ppsci.arch.tfnonet import TFNO1dNet, TFNO2dNet, TFNO3dNet # isort:skip from ppsci.arch.transformer import Transformer # isort:skip from ppsci.arch.unetex import UNetEx # isort:skip @@ -105,7 +105,6 @@ "ExtFormerMoECuboid", "FNO1d", "Generator", - "get_edge_index", "GraphCastNet", "HEDeepONets", "HGN", From 535b278c599a0415799032a44112fab577e75015 Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Sat, 6 Dec 2025 00:13:56 +0800 Subject: [PATCH 09/10] update code style --- examples/symbolic_gn/conf/config.yaml | 7 +- examples/symbolic_gn/gn.py | 139 +++++++++++++++----------- 2 files changed, 82 insertions(+), 64 deletions(-) diff --git a/examples/symbolic_gn/conf/config.yaml b/examples/symbolic_gn/conf/config.yaml index 4b0307a738..5eaececd9c 100644 --- a/examples/symbolic_gn/conf/config.yaml +++ b/examples/symbolic_gn/conf/config.yaml @@ -17,7 +17,6 @@ hydra: init_callback: _target_: ppsci.utils.callbacks.InitCallback sweep: - dir: ${hydra.run.dir} subdir: ./ @@ -27,7 +26,6 @@ output_dir: ${hydra:run.dir} log_freq: 20 device: gpu - TRAIN: epochs: 1000 batch_size: 32 @@ -48,7 +46,6 @@ TRAIN: type: "MAE" physics_weight: 0.0 - EVAL: pretrained_model_path: null eval_with_no_grad: true @@ -56,7 +53,6 @@ EVAL: message_analysis: enabled: true - export_path: "./exported_model" INFER: @@ -76,7 +72,6 @@ INFER: max_batch_size: 1024 num_cpu_threads: 10 - MODEL: arch: "OGN" input_keys: ["x", "edge_index"] @@ -145,4 +140,4 @@ DATATYPE_TYPE: dt: [1e-2] nt: [1000] n: [30] - dim: [2] \ No newline at end of file + dim: [2] diff --git a/examples/symbolic_gn/gn.py b/examples/symbolic_gn/gn.py index 4d5e99f6f4..ec4dd499bc 100644 --- a/examples/symbolic_gn/gn.py +++ b/examples/symbolic_gn/gn.py @@ -13,8 +13,7 @@ # limitations under the License. import os -from typing import List -from typing import Optional +from typing import List, Optional import hydra import matplotlib.pyplot as plt @@ -24,25 +23,22 @@ from simulate import SimulationDataset import ppsci -from ppsci.arch.symbolic_gn import HGN -from ppsci.arch.symbolic_gn import OGN -from ppsci.arch.symbolic_gn import VarOGN -from ppsci.arch.symbolic_gn import get_edge_index +from ppsci.arch.symbolic_gn import HGN, OGN, VarOGN, get_edge_index from ppsci.utils import logger def create_ppsci_model( - model_type: str = 'OGN', - input_keys: List[str] = ['x', 'edge_index'], - output_keys: List[str] = ['acceleration'], + model_type: str = "OGN", + input_keys: List[str] = ["x", "edge_index"], + output_keys: List[str] = ["acceleration"], n_f: int = 6, msg_dim: int = 100, ndim: int = 2, hidden: int = 300, edge_index: Optional[np.ndarray] = None, - l1_strength: float = 0.0 + l1_strength: float = 0.0, ): - if model_type == 'OGN': + if model_type == "OGN": model = OGN( input_keys=input_keys, output_keys=output_keys, @@ -51,18 +47,18 @@ def create_ppsci_model( ndim=ndim, hidden=hidden, edge_index=edge_index, - l1_strength=l1_strength + l1_strength=l1_strength, ) - elif model_type == 'HGN': + elif model_type == "HGN": model = HGN( input_keys=input_keys, output_keys=output_keys, n_f=n_f, ndim=ndim, hidden=hidden, - edge_index=edge_index + edge_index=edge_index, ) - elif model_type == 'VarOGN': + elif model_type == "VarOGN": model = VarOGN( input_keys=input_keys, output_keys=output_keys, @@ -71,7 +67,7 @@ def create_ppsci_model( ndim=ndim, hidden=hidden, edge_index=edge_index, - l1_strength=l1_strength + l1_strength=l1_strength, ) else: raise ValueError(f"未知的模型类型: {model_type}") @@ -85,10 +81,10 @@ def replicate_array(arr, n): def create_loss_function(cfg): def loss_function(output_dict, label_dict, weight_dict=None): - loss_type = getattr(cfg.TRAIN.loss, 'type', 'MAE') - + loss_type = getattr(cfg.TRAIN.loss, "type", "MAE") + output_key = list(cfg.MODEL.output_keys)[0] - + target = label_dict[output_key] pred = output_dict[output_key] if loss_type == "MSE": @@ -96,11 +92,11 @@ def loss_function(output_dict, label_dict, weight_dict=None): else: base_loss = paddle.mean(paddle.abs(pred - target)) - if 'l1_regularization' in output_dict: - base_loss = base_loss + output_dict['l1_regularization'] + if "l1_regularization" in output_dict: + base_loss = base_loss + output_dict["l1_regularization"] return {output_key: base_loss} - + return loss_function @@ -115,7 +111,7 @@ def train(cfg): n=cfg.DATA.num_nodes, dim=cfg.DATA.dimension, nt=cfg.DATA.time_steps, - dt=cfg.DATA.time_step_size + dt=cfg.DATA.time_step_size, ) sim.simulate(cfg.DATA.num_samples) accel_data = sim.get_acceleration() @@ -134,7 +130,7 @@ def train(cfg): cfg.MODEL.n_f = cfg.DATA.dimension * 2 + 2 if cfg.MODEL.ndim == "auto": cfg.MODEL.ndim = cfg.DATA.dimension - + train_size = int(len(X) * 0.8) X_train, X_val = X[:train_size], X[train_size:] y_train, y_val = y[:train_size], y[train_size:] @@ -180,7 +176,11 @@ def train(cfg): train_constraint.name: train_constraint, eval_constraint.name: eval_constraint, } - l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 + l1_strength = ( + cfg.MODEL.l1_strength + if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] + else 0.0 + ) model = create_ppsci_model( model_type=cfg.MODEL.arch, n_f=cfg.MODEL.n_f, @@ -188,28 +188,27 @@ def train(cfg): ndim=cfg.MODEL.ndim, hidden=cfg.MODEL.hidden, edge_index=edge_index, - l1_strength=l1_strength + l1_strength=l1_strength, ) batch_per_epoch = int(train_size / cfg.TRAIN.batch_size) if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": lr_scheduler = paddle.optimizer.lr.OneCycleLR( max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, - total_step=int(cfg.TRAIN.epochs*batch_per_epoch), - divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor + total_step=int(cfg.TRAIN.epochs * batch_per_epoch), + divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor, ) lr_scheduler.by_epoch = False else: lr_scheduler = paddle.optimizer.lr.ExponentialDecay( - learning_rate=cfg.TRAIN.optimizer.learning_rate, - gamma=0.9 + learning_rate=cfg.TRAIN.optimizer.learning_rate, gamma=0.9 ) lr_scheduler.by_epoch = True optimizer = paddle.optimizer.Adam( learning_rate=lr_scheduler, parameters=model.parameters(), - weight_decay=cfg.TRAIN.optimizer.weight_decay + weight_decay=cfg.TRAIN.optimizer.weight_decay, ) solver = ppsci.solver.Solver( @@ -233,7 +232,11 @@ def evaluate(cfg: DictConfig): cfg.MODEL.ndim = cfg.DATA.dimension edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) - l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 + l1_strength = ( + cfg.MODEL.l1_strength + if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] + else 0.0 + ) model = create_ppsci_model( model_type=cfg.MODEL.arch, @@ -242,7 +245,7 @@ def evaluate(cfg: DictConfig): ndim=cfg.MODEL.ndim, hidden=cfg.MODEL.hidden, edge_index=edge_index, - l1_strength=l1_strength + l1_strength=l1_strength, ) ppsci.utils.save_load.load_pretrain( @@ -255,20 +258,20 @@ def evaluate(cfg: DictConfig): n=cfg.DATA.num_nodes, dim=cfg.DATA.dimension, nt=cfg.DATA.time_steps, - dt=cfg.DATA.time_step_size + dt=cfg.DATA.time_step_size, ) sim.simulate(cfg.DATA.num_samples) accel_data = sim.get_acceleration() sample_indices = [0, 1] if cfg.DATA.num_samples > 1 else [0] - + for sample_idx in sample_indices: sample_data = sim.data[sample_idx, 0:1] true_accel = accel_data[sample_idx, 0:1] - + input_dict = { "x": paddle.to_tensor(sample_data, dtype="float32"), - "edge_index": paddle.to_tensor(edge_index, dtype="int64") + "edge_index": paddle.to_tensor(edge_index, dtype="int64"), } with paddle.no_grad(): @@ -278,16 +281,20 @@ def evaluate(cfg: DictConfig): error = np.mean(np.abs(pred_accel.numpy() - true_accel)) logger.info(f"Sample {sample_idx} - MAE error: {error:.6f}") - rel_error = np.linalg.norm(pred_accel.numpy() - true_accel) / np.linalg.norm(true_accel) + rel_error = np.linalg.norm(pred_accel.numpy() - true_accel) / np.linalg.norm( + true_accel + ) logger.info(f"Sample {sample_idx} - Relative error: {rel_error:.6f}") plt.figure(figsize=(10, 8)) sim.plot(sample_idx, animate=False, plot_size=True, s_size=2) - plt.title(f'{cfg.DATA.type.capitalize()} System - Sample {sample_idx}\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}') + plt.title( + f"{cfg.DATA.type.capitalize()} System - Sample {sample_idx}\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}" + ) plt.tight_layout() plot_path = os.path.join(cfg.output_dir, f"evaluation_sample_{sample_idx}.png") - plt.savefig(plot_path, dpi=300, bbox_inches='tight') + plt.savefig(plot_path, dpi=300, bbox_inches="tight") plt.close() logger.info(f"Evaluation plot saved to {plot_path}") @@ -303,7 +310,11 @@ def export(cfg: DictConfig): if cfg.MODEL.ndim == "auto": cfg.MODEL.ndim = cfg.DATA.dimension edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) - l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 + l1_strength = ( + cfg.MODEL.l1_strength + if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] + else 0.0 + ) model = create_ppsci_model( model_type=cfg.MODEL.arch, n_f=cfg.MODEL.n_f, @@ -311,7 +322,7 @@ def export(cfg: DictConfig): ndim=cfg.MODEL.ndim, hidden=cfg.MODEL.hidden, edge_index=edge_index, - l1_strength=l1_strength + l1_strength=l1_strength, ) solver = ppsci.solver.Solver( @@ -323,9 +334,13 @@ def export(cfg: DictConfig): input_spec = [ { - key: InputSpec([None, cfg.DATA.num_nodes, cfg.MODEL.n_f], "float32", name=key) - if key == "x" - else InputSpec([2, 30], "int64", name=key) + key: ( + InputSpec( + [None, cfg.DATA.num_nodes, cfg.MODEL.n_f], "float32", name=key + ) + if key == "x" + else InputSpec([2, 30], "int64", name=key) + ) for key in model.input_keys }, ] @@ -348,7 +363,7 @@ def inference(cfg: DictConfig): n=cfg.DATA.num_nodes, dim=cfg.DATA.dimension, nt=cfg.DATA.time_steps, - dt=cfg.DATA.time_step_size + dt=cfg.DATA.time_step_size, ) sim.simulate(cfg.DATA.num_samples) @@ -357,10 +372,7 @@ def inference(cfg: DictConfig): edge_index = get_edge_index(cfg.DATA.num_nodes, cfg.DATA.type) if use_predictor: - input_dict = { - "x": sample_data, - "edge_index": edge_index - } + input_dict = {"x": sample_data, "edge_index": edge_index} output_dict = predictor.predict(input_dict, cfg.INFER.batch_size) pred_acceleration = output_dict[cfg.MODEL.output_keys[0]] else: @@ -369,7 +381,12 @@ def inference(cfg: DictConfig): if cfg.MODEL.ndim == "auto": cfg.MODEL.ndim = cfg.DATA.dimension - l1_strength = cfg.MODEL.l1_strength if cfg.MODEL.regularization_type == "l1" and cfg.MODEL.arch in ["OGN", "VarOGN"] else 0.0 + l1_strength = ( + cfg.MODEL.l1_strength + if cfg.MODEL.regularization_type == "l1" + and cfg.MODEL.arch in ["OGN", "VarOGN"] + else 0.0 + ) model = create_ppsci_model( model_type=cfg.MODEL.arch, @@ -378,7 +395,7 @@ def inference(cfg: DictConfig): ndim=cfg.MODEL.ndim, hidden=cfg.MODEL.hidden, edge_index=edge_index, - l1_strength=l1_strength + l1_strength=l1_strength, ) ppsci.utils.save_load.load_pretrain( @@ -388,7 +405,7 @@ def inference(cfg: DictConfig): input_dict = { "x": paddle.to_tensor(sample_data, dtype="float32"), - "edge_index": paddle.to_tensor(edge_index, dtype="int64") + "edge_index": paddle.to_tensor(edge_index, dtype="int64"), } with paddle.no_grad(): @@ -401,16 +418,20 @@ def inference(cfg: DictConfig): error = np.mean(np.abs(pred_acceleration - true_acceleration)) logger.info(f"Inference error (MAE): {error:.6f}") - rel_error = np.linalg.norm(pred_acceleration - true_acceleration) / np.linalg.norm(true_acceleration) + rel_error = np.linalg.norm(pred_acceleration - true_acceleration) / np.linalg.norm( + true_acceleration + ) logger.info(f"Inference relative error: {rel_error:.6f}") plt.figure(figsize=(10, 8)) sim.plot(sample_idx, animate=False, plot_size=True, s_size=2) - plt.title(f'{cfg.DATA.type.capitalize()} System - Inference\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}') + plt.title( + f"{cfg.DATA.type.capitalize()} System - Inference\nMAE: {error:.4f}, Rel Error: {rel_error:.4f}" + ) plt.tight_layout() plot_path = os.path.join(cfg.output_dir, "inference.png") - plt.savefig(plot_path, dpi=300, bbox_inches='tight') + plt.savefig(plot_path, dpi=300, bbox_inches="tight") plt.close() logger.info(f"Inference plot saved to {plot_path}") @@ -427,9 +448,11 @@ def main(cfg: DictConfig) -> None: inference(cfg) else: raise ValueError( - "cfg.mode should in [train, eval, export, infer], but got {}".format(cfg.mode) + "cfg.mode should in [train, eval, export, infer], but got {}".format( + cfg.mode + ) ) if __name__ == "__main__": - main() \ No newline at end of file + main() From b499b7ad3b4ff4831af161eb8cc4d17ad585d7b6 Mon Sep 17 00:00:00 2001 From: ADream-ki <2085127827@qq.com> Date: Sun, 14 Dec 2025 22:06:37 +0800 Subject: [PATCH 10/10] update code --- docs/zh/examples/symbolic_gn.md | 6 +++--- examples/symbolic_gn/gn.py | 10 +++++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/docs/zh/examples/symbolic_gn.md b/docs/zh/examples/symbolic_gn.md index b6caf40417..bdd5620881 100644 --- a/docs/zh/examples/symbolic_gn.md +++ b/docs/zh/examples/symbolic_gn.md @@ -265,7 +265,7 @@ TRAIN: l1_strength: 1e-2 ``` -2. **提取消息数据**：训练过程中可以记录消息向量（需要修改代码添加钩子） +2. **提取消息数据**：训练过程中可以记录消息向量 ```python messages = model.msg_fnc(edge_features) # [num_edges, msg_dim] # 保存消息及对应的物理特征（dx, dy, r, m1, m2等） @@ -274,10 +274,10 @@ TRAIN: 3. **符号回归**：使用工具如 [PySR](https://github.com/MilesCranmer/PySR) 拟合符号表达式 ```python from pysr import PySRRegressor - + # 选择最活跃的消息通道 active_channels = np.argsort(np.std(messages, axis=0))[-5:] - + # 对每个活跃通道进行符号回归 for ch in active_channels: model = PySRRegressor( diff --git a/examples/symbolic_gn/gn.py b/examples/symbolic_gn/gn.py index ec4dd499bc..b453dc60da 100644 --- a/examples/symbolic_gn/gn.py +++ b/examples/symbolic_gn/gn.py @@ -13,7 +13,8 @@ # limitations under the License. import os -from typing import List, Optional +from typing import List +from typing import Optional import hydra import matplotlib.pyplot as plt @@ -23,7 +24,10 @@ from simulate import SimulationDataset import ppsci -from ppsci.arch.symbolic_gn import HGN, OGN, VarOGN, get_edge_index +from ppsci.arch.symbolic_gn import HGN +from ppsci.arch.symbolic_gn import OGN +from ppsci.arch.symbolic_gn import VarOGN +from ppsci.arch.symbolic_gn import get_edge_index from ppsci.utils import logger @@ -195,7 +199,7 @@ def train(cfg): if cfg.TRAIN.lr_scheduler.name == "OneCycleLR": lr_scheduler = paddle.optimizer.lr.OneCycleLR( max_learning_rate=cfg.TRAIN.lr_scheduler.max_learning_rate, - total_step=int(cfg.TRAIN.epochs * batch_per_epoch), + total_steps=int(cfg.TRAIN.epochs * batch_per_epoch), divide_factor=cfg.TRAIN.lr_scheduler.final_div_factor, ) lr_scheduler.by_epoch = False