NTopo: Mesh-free Topology Optimization using Implicit Neural Representations

模型训练命令模型评估命令

python ntopo.py

python ntopo.py mode=eval PROBLEM=Beam2D EVAL.pretrained_model_path_density=https://paddle-org.bj.bcebos.com/paddlescience/models/ntopo/beam2d_pretrained.pdparams
python ntopo.py mode=eval PROBLEM=Bridge2D EVAL.pretrained_model_path_density=https://paddle-org.bj.bcebos.com/paddlescience/models/ntopo/bridge2d_pretrained.pdparams
python ntopo.py mode=eval PROBLEM=Distributed2D EVAL.pretrained_model_path_density=https://paddle-org.bj.bcebos.com/paddlescience/models/ntopo/distributed2d_pretrained.pdparams
python ntopo.py mode=eval PROBLEM=LongBeam2D EVAL.pretrained_model_path_density=https://paddle-org.bj.bcebos.com/paddlescience/models/ntopo/longbeam2d_pretrained.pdparams
python ntopo.py mode=eval PROBLEM=LShape2D EVAL.pretrained_model_path_density=https://paddle-org.bj.bcebos.com/paddlescience/models/ntopo/lshape2d_pretrained.pdparams
python ntopo.py mode=eval PROBLEM=Triangle2D EVAL.pretrained_model_path_density=https://paddle-org.bj.bcebos.com/paddlescience/models/ntopo/triangle2d_pretrained.pdparams
python ntopo.py mode=eval PROBLEM=TriangleVariants2D EVAL.pretrained_model_path_density=https://paddle-org.bj.bcebos.com/paddlescience/models/ntopo/trianglevariants2d_pretrained.pdparams
python ntopo.py --config-name ntopo.yaml mode=eval PROBLEM=Beam3D EVAL.pretrained_model_path_density=https://paddle-org.bj.bcebos.com/paddlescience/models/ntopo/beam3d_pretrained.pdparams
python ntopo.py --config-name ntopo.yaml mode=eval PROBLEM=Bridge3D EVAL.pretrained_model_path_density=https://paddle-org.bj.bcebos.com/paddlescience/models/ntopo/bridge3d_pretrained.pdparams

注：由于该案例训练方法较为特殊，没有参考指标，训练完成后直接生成可视化结果，以判断训练效果。

1. 背景简介

在拓扑优化问题中有一种常见的方法 SIMP(Solid Isotropic Material with Penalization)，它是一种基于密度法的拓扑优化方法，通过连续设计变量（材料密度）描述设计域内各点的材料分布，最终逼近理想的“0-1”二元分布（材料存在或不存在）。其核心目标是优化材料布局，在满足约束条件（如体积、刚度）下实现特定性能目标（如最小柔顺度）。

在传统的数值计算中，SIMP 将设计域离散为有限元网格，每个单元分配连续密度变量 \(\rho \in [0,1]\)，然后引入幂律插值函数和惩罚因子，通过数学公式将中间密度值向0或1方向惩罚，抑制灰度区域，例如，弹性模量插值公式为：\(E(\rho) = \rho^p E_1\)

该案例提出了一种基于隐式神经表示的新型机器学习方法，用于解决‌拓扑优化‌这一高难度逆问题。传统方法依赖网格化处理，而该案例通过 MLP 无网格化地参数化‌密度场‌和‌位移场‌，利用神经网络的连续可微特性生成高细节解。实验表明，该方法在‌结构顺应性目标优化‌中表现优异，并能通过自监督学习探索拓扑优化问题的连续解空间，克服了传统方法在高维参数空间和非线性目标函数中的局限性。核心创新在于将神经表示与无网格优化结合，为复杂逆问题提供了高效且灵活的解决方案。

2. 问题定义

拓扑优化（TO）的目标是在给定的边界条件、作用力和目标材料体积比下，找到使结构最刚硬的材料分布。这一问题可以形式化为一个带约束的双层最小化问题：

\[ \begin{cases} \min_{\rho}L_{comp}(\rho)=\int_{\Omega}e(\rho,u(\rho),\omega)d\omega \\ s.t. \quad u(\rho)=\arg\min_{u} L_{sim}(u,\rho) \\ \rho(\omega) \in {0, 1}, \quad \frac{1}{|\Omega|} \int_{\Omega} \rho d\omega = \hat{V} \\ \end{cases} \]

其中 \(L_{comp}(\rho)\) 是顺应性损失函数; \(e(\rho, u(\rho), \omega)\) 是点顺应性，与内部能量成正比; \(u(\rho)\) 是位移场，满足力平衡条件，通过最小化模拟损失 \(L_{sim}(u, \rho)\) 得到; \(\rho(\omega)\) 是材料密度场，理论上取值为 \(0\) 或 \(1\)（表示有无材料），但在实际优化中允许连续值，并鼓励收敛到二进制解; \(\Omega\) 是设计域，\(\omega\) 是空间坐标; \(\hat{V}\) 是目标材料体积比。

3. 问题求解

接下来开始讲解如何将问题一步一步地转化为 PaddleScience 代码，用深度学习的方法求解该问题。 为了快速理解 PaddleScience，接下来仅对模型构建、方程构建、计算域构建等关键步骤进行阐述，而其余细节请参考 API文档。

3.1 模型构建

训练整体流程

上图为训练整体流程，通过交替训练位移网络和密度网络这两个神经网络，分别将空间坐标 \(ω\) 映射到平衡位移 \(u\) 和最优密度 \(\rho\) 来计算最优材料分布。在每次迭代中，首先通过最小化系统的总势能来更新位移网络，随后进行灵敏度分析以计算密度空间梯度，并通过灵敏度滤波生成目标密度场 \(\hat{\rho}\)。最后，通过最小化基于当前密度与目标密度均方误差的凸优化目标函数来更新密度网络。

位移网络和密度网络都是使用 SIREN 激活函数的 MLP 网络，具体代码请参考 完整代码 中 model.py 文件。

3.2 参数和超参数设定

我们需要指定问题相关的参数，如要优化的问题的名称（几何类型）、材料参数、优化目标（体积百分比）等：

# set problem parameters
PROBLEM: "Beam2D"

# set working condition
NU: 0.3
E: 1.0
VOLUME_RATIO: 0.5
SIGMOID_ALPHA: 5.0
ENERGY_EXP: 3.0  # TODO: trainable parameter
PENALTY: 10
MAX_MOVE: 0.2
DAMPING: 0.5

USE_MMSE: true
USE_OC: true
FILTER: "Gaussian"  # "null" "Gaussian"
FILTER_RADIUS: 2.0

另外需要在配置文件中指定训练轮数、batch_size 等其他训练所需参数，注意这里两个网络单独的 epochs 参数需要设置为 \(1\)：

# training settings
TRAIN:
  st_epoch: 1
  epochs: 100  # times for for-loop
  disp_net:
    epochs: 1
    iters_per_epoch: 1000
    optimizer:
      learning_rate: 3.0e-4
      beta2: 0.99
      epsilon: 1e-7
  density_net:
    epochs: 1
    iters_per_epoch: 50
    optimizer:
      learning_rate: 3.0e-4
      beta1: 0.8
      beta2: 0.9
      epsilon: 1e-7
  batch_size:
    constraint: 7500
    visualizer: 7500
  save_freq: 10
  eval_during_train: false
  eval_freq: 100
  # Due to the particularity of this case training process,
  # please use `pretrained_model_path` but not `checkpoint_path` to resume training
  # and manually adjust the above `st_epoch`
  pretrained_model_path_disp: null
  pretrained_model_path_density: null
  checkpoint_path_disp: null
  checkpoint_path_density: null
  enable_parallel: false # whether to enable parallel training

特别需要注意的是该案例中使用了一些技巧，如使用移动平方误差(MMSE)、基于多批次的最佳标准方法(OC)、滤波，它们相关的参数以及训练的 iters_per_epoch 需要谨慎设定，并非某个参数越大越好，不同的参数设置可能导向不同的优化结果。

3.3 优化器构建

训练过程会调用优化器来更新模型参数，此处选择 Adam 优化器。

optimizer_density = ppsci.optimizer.Adam(**cfg.TRAIN.density_net.optimizer)(
    problem.density_net
)

# set stratified sampler
bounds = (
    (problem.geo_origin[0], problem.geo_dim[0]),

3.4 方程构建

如 问题定义 中所述，模型训练过程中需要使用弹性模量插值公式等公式，因此需要定义方程，具体代码请参考 完整代码 中 equation.py 文件。

3.5 问题构建（包含 loss）

本问题的计算域为初始几何结构，本案例中提供了部分 2D 及 3D 问题的类，其中包含了对计算域、边界条件、受力情况等多种参数和条件的定义，具体代码请参考 完整代码 中 problems.py 文件。

其中值得注意的是，该案例中使用了一些技巧，如基于多批次的最佳标准方法(OC)，这种方法需要一个批次的输入、输出数据，然后以某种方式计算得到该批次的 loss。

3.6 约束构建

本案例代码中存在 1 种内部点约束和 1 种监督约束（但实际上 loss 计算时没有用到标签，由于调用了 API， 这里按照约束的方式介绍）。

3.6.1 内部点约束

几何内部的点存在约束 InteriorConstraint：

        **train_dataloader_cfg,
        "batch_size": cfg.TRAIN.batch_size.constraint,
        "iters_per_epoch": cfg.TRAIN.disp_net.iters_per_epoch,
    },
    ppsci.loss.FunctionalLoss(problem.disp_loss_func),
    name="INTERIOR_DISP",
)

# re-assign to ITERS_PER_EPOCH_DISP
if cfg.TRAIN.enable_parallel:
    cfg.TRAIN.disp_net.iters_per_epoch = len(interior_disp.data_loader)

# wrap constraints together

InteriorConstraint 的第一个参数是方程（组）表达式，用于描述如何计算约束目标，此处填入在 3.4 方程构建 章节中实例化好的 problem.equation["EEquation"].equations；

第二个参数是约束变量的目标值，在本问题按照中希望与方程相关的 \(E\) 值 E_xyz 或 E_xy 被优化至 0；

第三个参数是约束方程作用的计算域，此处填入在 3.5 问题构建 章节实例化好的相应问题的计算域 problem.geom["geo"] ；

第四个参数是在计算域上的采样配置。

第五个参数是损失函数，此处通过 ppsci.loss.FunctionalLoss 传入自定义损失函数 problem.disp_loss_func；

第六个是约束条件的名字，需要给每一个约束条件命名，方便后续对其索引。此处命名为 "INTERIOR_DISP"。

3.6.2 监督约束

由于该案例中自定义了采样方法，因此此处调用监督约束 SupervisedConstraint，并将采样点以 input 的形式传递给它：

        "dataset": {"name": "NamedArrayDataset", "input": input_},
        "sampler": {
            "name": "BatchSampler",
            "drop_last": True,
            "shuffle": False,
        },
        "num_workers": 0,
        "batch_size": int(np.prod(problem.batch_size) * problem.batch_raito),
    },
    func_module.FunctionalLossBatch(problem.density_loss_func),
    output_expr=problem.equation["EEquation"].equations,
    name="INTERIOR_DENSITY",
)

constraint_density = {interior_density.name: interior_density}

# set visualizer(optional)
pred_input_keys = ("x", "y")
if problem.dim == 3:
    pred_input_keys += ("z",)

SupervisedConstraint 的第一个参数是监督约束的读取配置，其中 dataset 字段表示使用的训练数据集信息，各个字段分别表示：

1. name： 数据集类型，此处 NamedArrayDataset 表示从 Array 中读取的数据集；
2. input： Array 类型的输入数据；

注意，其中没有标签值 label。

sampler 字段表示采样方法，其中各个字段表示：

1. name： 采样器类型，此处 BatchSampler 表示批采样器；
2. drop_last： 是否需要丢弃最后无法凑整一个 mini-batch 的样本，设为 False；
3. shuffle： 是否需要在生成样本下标时打乱顺序，设为 True；

num_workers 字段表示输入加载时的线程数；

batch_size 字段表示 batch 的大小；

第二个参数是损失函数，这里自定义了一个损失函数类，用于接收该案例特殊的批次损失函数 problem.density_loss_func；

第三个参数是约束条件的名字，我们需要给每一个约束条件命名，方便后续对其索引。此处命名为 "INTERIOR_DENSITY"。

3.7 可视化器构建

该案例每隔一定训练间隔，或通过可视化器 ppsci.visualize.VisualizerVtu，将优化结果保存为vtu文件：

    "vis_density": ppsci.visualize.VisualizerVtu(
        pred_input_dict,
        {
            "density": lambda out: out["densities"],
        },
        batch_size=cfg.TRAIN.batch_size.visualizer,
        prefix="vtu_density",
    ),
}

# initialize solver
solver_disp = ppsci.solver.Solver(
    model=model_list,
    constraint=constraint_disp,
    output_dir=cfg.output_dir_disp,
    optimizer=optimizer_disp,
    epochs=cfg.TRAIN.disp_net.epochs,

3.8 其他函数

如上所述，该案例中需要交替训练两个模型，同时增加了一些技巧，如使用移动平方误差(MMSE)、基于多批次的最佳标准方法(OC)、滤波，因此该案例的训练过程与单一模型训练差距较大。

因此该案例中根据 PaddleScience 代码，自定义了：

1. Trainer 类，该类根据接收的 solver 中的信息，定义新的训练过程；
2. FunctionalLossBatch 类，该类基于 ppsci.loss.base.Loss，重新定义 loss 处理方式，并在 Trainer 中被调用；
3. Sampler 类，该类定义了案例所需的采样方法；
4. Plot 类，对于形状对称的几何，该案例选择仅定义对称部分的一半，然后根据问题中的 mirror参数还原完整结果。该类提供了相关处理函数；

具体代码请参考 完整代码 中 functions.py 文件。

3.9 模型训练、评估

完成上述设置之后，将上述实例化的对象按顺序传递给 ppsci.solver.Solver 后，按照自定以训练过程进行训练，具体代码请参考 完整代码 中 ntopo.py 文件。

由于拓扑优化问题没有标签，多种优化结果可能都行之有效，因此需要根据可视化结果，人为评估训练结果。

4. 完整代码

"""
Paper: https://arxiv.org/abs/2102.10782
Reference: https://github.com/JonasZehn/ntopo
"""
from os import makedirs
from os import path as osp

import functions as func_module
import hydra
import model as model_module
import numpy as np
import paddle
import problems as problems_module
from omegaconf import DictConfig

import ppsci
from ppsci.utils import save_load


def train(cfg: DictConfig):
    # make dirs
    makedirs(cfg.output_dir_disp, exist_ok=True)
    makedirs(cfg.output_dir_density, exist_ok=True)

    # set problem
    problem = getattr(problems_module, cfg.PROBLEM)(cfg)

    # set model
    problem.disp_net = model_module.DenseSIRENModel(**cfg.MODEL.disp_net)
    problem.density_net = model_module.DenseSIRENModel(**cfg.MODEL.density_net)

    # set transforms
    problem.disp_net.register_input_transform(problem.transform_in)
    problem.disp_net.register_output_transform(problem.transform_out_disp)
    problem.density_net.register_input_transform(problem.transform_in)
    problem.density_net.register_output_transform(problem.transform_out_density)

    model_list = ppsci.arch.ModelList((problem.disp_net, problem.density_net))

    # set optimizer
    optimizer_disp = ppsci.optimizer.Adam(**cfg.TRAIN.disp_net.optimizer)(
        problem.disp_net
    )
    optimizer_density = ppsci.optimizer.Adam(**cfg.TRAIN.density_net.optimizer)(
        problem.density_net
    )

    # set stratified sampler
    bounds = (
        (problem.geo_origin[0], problem.geo_dim[0]),
        (problem.geo_origin[1], problem.geo_dim[1]),
    )
    if problem.dim == 3:
        bounds += ((problem.geo_origin[2], problem.geo_dim[2]),)
    sampler = func_module.Sampler(problem.geom["geo"], bounds=bounds)

    # set dataloader config
    train_dataloader_cfg = {
        "dataset": "NamedArrayDataset",
        "sampler": {
            "name": "BatchSampler",
            "drop_last": True,
            "shuffle": False,
        },
        "num_workers": 0,
    }

    # set constraint
    interior_disp = ppsci.constraint.InteriorConstraint(
        problem.equation["EEquation"].equations,
        {"E_xyz": 0} if problem.dim == 3 else {"E_xy": 0},
        problem.geom["geo"],
        {
            **train_dataloader_cfg,
            "batch_size": cfg.TRAIN.batch_size.constraint,
            "iters_per_epoch": cfg.TRAIN.disp_net.iters_per_epoch,
        },
        ppsci.loss.FunctionalLoss(problem.disp_loss_func),
        name="INTERIOR_DISP",
    )

    # re-assign to ITERS_PER_EPOCH_DISP
    if cfg.TRAIN.enable_parallel:
        cfg.TRAIN.disp_net.iters_per_epoch = len(interior_disp.data_loader)

    # wrap constraints together
    constraint_disp = {interior_disp.name: interior_disp}

    input_, mask = sampler.sample_interior_stratified(
        n_samples=problem.batch_size,
        n_iter=cfg.TRAIN.density_net.iters_per_epoch,
    )
    interior_density = ppsci.constraint.SupervisedConstraint(
        {
            "dataset": {"name": "NamedArrayDataset", "input": input_},
            "sampler": {
                "name": "BatchSampler",
                "drop_last": True,
                "shuffle": False,
            },
            "num_workers": 0,
            "batch_size": int(np.prod(problem.batch_size) * problem.batch_raito),
        },
        func_module.FunctionalLossBatch(problem.density_loss_func),
        output_expr=problem.equation["EEquation"].equations,
        name="INTERIOR_DENSITY",
    )

    constraint_density = {interior_density.name: interior_density}

    # set visualizer(optional)
    pred_input_keys = ("x", "y")
    if problem.dim == 3:
        pred_input_keys += ("z",)

    # add inferencer data
    samplers = problem.geom["geo"].sample_interior(cfg.TRAIN.batch_size.visualizer)
    pred_input_dict = {}
    for key in pred_input_keys:
        pred_input_dict.update({key: samplers[key]})

    visualizer_disp = {
        "vis_disp": ppsci.visualize.VisualizerVtu(
            pred_input_dict,
            {key: lambda out, k=key: out[k] for key in cfg.MODEL.disp_net.output_keys},
            prefix="vtu_disp",
        ),
    }
    visualizer_density = {
        "vis_density": ppsci.visualize.VisualizerVtu(
            pred_input_dict,
            {
                "density": lambda out: out["densities"],
            },
            batch_size=cfg.TRAIN.batch_size.visualizer,
            prefix="vtu_density",
        ),
    }

    # initialize solver
    solver_disp = ppsci.solver.Solver(
        model=model_list,
        constraint=constraint_disp,
        output_dir=cfg.output_dir_disp,
        optimizer=optimizer_disp,
        epochs=cfg.TRAIN.disp_net.epochs,
        iters_per_epoch=cfg.TRAIN.disp_net.iters_per_epoch,
        seed=cfg.seed,
        equation=problem.equation,
        geom=problem.geom,
        log_freq=cfg.log_freq,
        save_freq=cfg.TRAIN.save_freq,
        eval_during_train=cfg.TRAIN.eval_during_train,
        eval_freq=cfg.TRAIN.eval_freq,
        visualizer=visualizer_disp,
        pretrained_model_path=cfg.TRAIN.pretrained_model_path_disp,
        checkpoint_path=cfg.TRAIN.checkpoint_path_disp,
    )

    solver_density = ppsci.solver.Solver(
        model=model_list,
        constraint=constraint_density,
        output_dir=cfg.output_dir_density,
        optimizer=optimizer_density,
        epochs=cfg.TRAIN.density_net.epochs,
        iters_per_epoch=cfg.TRAIN.density_net.iters_per_epoch,
        equation=problem.equation,
        geom=problem.geom,
        log_freq=cfg.log_freq,
        save_freq=cfg.TRAIN.save_freq,
        eval_during_train=cfg.TRAIN.eval_during_train,
        eval_freq=cfg.TRAIN.eval_freq,
        visualizer=visualizer_density,
        pretrained_model_path=cfg.TRAIN.pretrained_model_path_density,
        checkpoint_path=cfg.TRAIN.checkpoint_path_density,
    )

    # initialize density trainer
    if problem.use_mmse:
        density_trainer = func_module.Trainer(solver_density)

    # training
    solver_disp.train()

    for i in range(cfg.TRAIN.st_epoch, cfg.TRAIN.epochs + 1):
        ppsci.utils.logger.info(f"[Total Train][Epoch {i}/{cfg.TRAIN.epochs}]")

        input_, _ = sampler.sample_interior_stratified(
            n_samples=problem.batch_size,
            n_iter=cfg.TRAIN.density_net.iters_per_epoch,
        )

        interior_density = ppsci.constraint.SupervisedConstraint(
            {
                "dataset": {"name": "NamedArrayDataset", "input": input_},
                "sampler": {
                    "name": "BatchSampler",
                    "drop_last": True,
                    "shuffle": False,
                },
                "num_workers": 0,
                "batch_size": int(np.prod(problem.batch_size) * problem.batch_raito),
            },
            func_module.FunctionalLossBatch(problem.density_loss_func),
            output_expr=problem.equation["EEquation"].equations,
            name="INTERIOR_DENSITY",
        )
        solver_density.constraint["INTERIOR_DENSITY"] = interior_density

        solver_disp.train()
        if problem.use_mmse:
            density_trainer.train_batch()
        else:
            solver_density.train()

        # plotting during training
        if i == 1 or i % cfg.TRAIN.save_freq == 0 or i == cfg.TRAIN.epochs:
            visualizer_density["vis_density"].prefix = f"vtu_density_e{i}"
            solver_density.visualize()

            visualizer_disp["vis_disp"].prefix = f"vtu_disp_e{i}"
            solver_disp.visualize()

            save_load.save_checkpoint(
                solver_density.model,
                solver_density.optimizer,
                {"metric": "dummy", "epoch": i},
                solver_density.scaler,
                solver_density.output_dir,
                f"epoch_{i}",
                solver_density.equation,
            )


def evaluate(cfg: DictConfig):
    # set problem
    problem = getattr(problems_module, cfg.PROBLEM)(cfg)

    # set model
    problem.density_net = model_module.DenseSIRENModel(**cfg.MODEL.density_net)

    # set transforms
    problem.density_net.register_input_transform(problem.transform_in)
    problem.density_net.register_output_transform(problem.transform_out_density)

    if problem.dim == 2:
        # add inferencer data
        samplers = problem.geom["geo"].sample_interior(cfg.EVAL.num_sample)
        pred_input_dict = {}
        if problem.mirror:
            if problem.mirror[0]:
                pred_input_dict["x"] = np.concatenate(
                    [samplers["x"], 2 * problem.geo_dim[0] - samplers["x"]]
                )
                pred_input_dict["y"] = np.concatenate([samplers["y"], samplers["y"]])
            if problem.mirror[1]:
                pred_input_dict["x"] = np.concatenate([samplers["x"], samplers["x"]])
                pred_input_dict["y"] = np.concatenate(
                    [samplers["y"], 2 * problem.geo_dim[1] - samplers["y"]]
                )
        else:
            pred_input_dict["x"] = samplers["x"]
            pred_input_dict["y"] = samplers["y"]

        def compute_mirror_density(problem, out):
            densities = out["densities"][: cfg.EVAL.num_sample]
            if problem.mirror:
                if problem.mirror[0]:
                    densities = paddle.concat([densities, densities])
                if problem.mirror[1]:
                    densities = paddle.concat([densities, densities])
            return densities

        visualizer_density = {
            "vis_density": ppsci.visualize.VisualizerVtu(
                pred_input_dict,
                {"density": lambda out: compute_mirror_density(problem, out)},
                batch_size=pred_input_dict["x"].shape[0],
                prefix="vtu_density",
            ),
        }

        solver_density = ppsci.solver.Solver(
            model=problem.density_net,
            output_dir=cfg.output_dir,
            visualizer=visualizer_density,
            pretrained_model_path=cfg.EVAL.pretrained_model_path_density,
        )
        solver_density.visualize()
    elif problem.dim == 3:
        # load pretrained model
        save_load.load_pretrain(
            problem.density_net, cfg.EVAL.pretrained_model_path_density
        )
        # plotting
        plot = func_module.Plot(
            osp.join(cfg.output_dir, "density.obj"),
            problem,
            cfg.EVAL.n_cells,
            0.5,
        )
        plot.plot_3d()


@hydra.main(version_base=None, config_path="./conf", config_name="ntopo_2d.yaml")
def main(cfg: DictConfig):
    if cfg.mode == "train":
        train(cfg)
    elif cfg.mode == "eval":
        evaluate(cfg)
    else:
        raise ValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'")


if __name__ == "__main__":
    main()

# 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.

from typing import Tuple

import numpy as np
import paddle
import paddle.nn as nn

from ppsci.arch import activation as act_mod
from ppsci.arch import base
from ppsci.utils import initializer


class DenseBlock(nn.Layer):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        w0: float = 1.0,
        coef1: float = 6.0,
        coef2: float = 1.0,
        weight_init=True,
        bias_init=True,
        use_act=False,
        use_sqrt=False,
    ) -> None:
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.act = act_mod.Siren(w0) if use_act else None

        if weight_init:
            self.init_param(self.linear.weight, coef1, coef2, use_sqrt)
        else:
            self.init_zeros(self.linear.weight)
        if bias_init:
            self.init_param(self.linear.bias, coef1, coef2)
        else:
            self.init_zeros(self.linear.bias)

    def init_param(self, param, coef1: float = 6.0, coef2: float = 1.0, use_sqrt=True):
        in_features = param.shape[0]
        with paddle.no_grad():
            initializer.uniform_(
                param,
                -np.sqrt(coef1 / in_features) * coef2,
                np.sqrt(coef1 / in_features) * coef2,
            ) if use_sqrt else initializer.uniform_(
                param,
                -coef1 / in_features * coef2,
                coef1 / in_features * coef2,
            )

    def init_zeros(self, param):
        initializer.zeros_(param)

    def forward(self, x):
        y = x
        y = self.linear(y)
        if self.act:
            y = self.act(y)
        return y


class DenseSIRENModel(base.Arch):
    """DenseSIRENModel network."""

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        num_layers: int,
        hidden_size: int,
        last_layer_init_scale: float,
    ):
        super().__init__()
        self.input_keys = input_keys
        self.output_keys = output_keys
        self.linears = []
        self.skip = (False, False, True, False, True, False)
        in_features = (
            len(input_keys),
            hidden_size,
            hidden_size,
            len(input_keys) + hidden_size,
            hidden_size,
            len(input_keys) + 2 * hidden_size,
        )
        out_features = (
            hidden_size,
            hidden_size,
            hidden_size,
            hidden_size,
            hidden_size,
            len(output_keys),
        )
        w0 = (60.0, 1.0, 1.0, 1.0, 1.0, 1.0)
        coef1 = (1.0, 6.0, 6.0, 6.0, 6.0, 6.0)
        coef2 = (1.0, 1.0, 1.0, 1.0, 1.0, last_layer_init_scale)
        weight_init = (True, True, True, True, True, True)
        bias_init = (True, True, True, True, True, False)
        use_act = (True, True, True, True, True, False)
        use_sqrt = (False, True, True, True, True, True)

        # initialize layers
        for i in range(num_layers):
            self.linears.append(
                DenseBlock(
                    in_features[i],
                    out_features[i],
                    w0[i],
                    coef1[i],
                    coef2[i],
                    weight_init[i],
                    bias_init[i],
                    use_act[i],
                    use_sqrt[i],
                )
            )
        self.linears = nn.LayerList(self.linears)

    def forward_tensor(self, x):
        y = x
        short = y
        for i, layer in enumerate(self.linears):
            y = layer(y)
            if self.skip[i]:
                y = paddle.concat([short, y], axis=-1)
                short = y
        return y

    def forward(self, x):
        if self._input_transform is not None:
            x = self._input_transform(x)

        y = self.concat_to_tensor(x, self.input_keys, axis=-1)
        y = self.forward_tensor(y)
        y = self.split_to_dict(y, self.output_keys, axis=-1)

        if self._output_transform is not None:
            y = self._output_transform(x, y)
        return y

# 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.

from typing import Dict
from typing import Union

import numpy as np

from ppsci.autodiff import jacobian
from ppsci.equation.pde import base


class EEquation(base.PDE):
    r"""Linear elasticity equations.
    Use either (E, nu) or (lambda_, mu) to define the material properties.

    2D:
    $$
    \begin{cases}
        \sigma_{xx} = \mu \dfrac{\partial u}{\partial x} + \lambda(\dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y})
        \sigma_{xy} = \dfrac{\mu}{2}(\dfrac{\partial u}{\partial y} + \dfrac{\partial v}{\partial x})
        \sigma_{yy} = \mu \dfrac{\partial v}{\partial y} + \lambda(\dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y})
        E_{xy} = \dfrac{1}{2}(\dfrac{\partial u}{\partial x}\sigma_{xx} + (\dfrac{\partial u}{\partial y} + \dfrac{\partial v}{\partial x})\sigma_{xy} + \dfrac{\partial v}{\partial y}\sigma_{yy}) \\
    \end{cases}
    $$

    3D:
    $$
    \begin{cases}
        \epsilon_{xy} = \dfrac{1}{2}(\dfrac{\partial u}{\partial y} + \dfrac{\partial v}{\partial x})
        \epsilon_{xz} = \dfrac{1}{2}(\dfrac{\partial u}{\partial z} + \dfrac{\partial w}{\partial x})
        \epsilon_{zy} = \dfrac{1}{2}(\dfrac{\partial v}{\partial z} + \dfrac{\partial w}{\partial y})
        E_{xyz} = \dfrac{\lambda}{2}(\dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} + \dfrac{\partial w}{\partial z})^2 + \mu (\dfrac{\partial^2 u}{\partial x^2} + \dfrac{\partial^2 v}{\partial y^2} + \dfrac{\partial^2 w}{\partial z^2}) + 2\mu (\epsilon_{xy}^2 + \epsilon_{xz}^2 + \epsilon_{yz}^2) \\
    \end{cases}
    $$

    Args:
        E (Optional[float]): The Young's modulus. Defaults to None.
        nu (Optional[float]): The Poisson's ratio. Defaults to None.
        lambda_ (Optional[float]): Lamé's first parameter. Defaults to None.
        mu (Optional[float]): Lamé's second parameter (shear modulus). Defaults to None.
        rho (float, optional): Mass density. Defaults to 1.
        dim (int, optional): Dimension of the linear elasticity (2 or 3). Defaults to 3.
        time (bool, optional): Whether contains time data. Defaults to False.

    Examples:
        >>> import ppsci
        >>> pde = ppsci.equation.LinearElasticity(
        ...     E=None, nu=None, lambda_=1e4, mu=100, dim=3
        ... )
    """

    def __init__(
        self,
        param_dict: Dict[str, Union[float, np.ndarray]],
        dim: int = 3,
        time: bool = False,
    ):
        super().__init__()
        self.param_keys = (
            "E",
            "nu",
            "lambda_",
            "mu",
            "u__x",
            "u__y",
            "u__z",
            "v__x",
            "v__y",
            "v__z",
            "w__x",
            "w__y",
            "w__z",
        )
        self.dim = dim
        self.time = time
        self.param_dict = param_dict
        for key in self.param_keys:
            if key in self.param_dict:
                setattr(self, key, self.param_dict[key])

        def init_params(out):
            for key in self.param_keys:
                if key not in self.param_dict:
                    setattr(self, key, out[key] if key in out else None)

            if self.u__x is None:
                self.u__x = jacobian(out["u"], out["x"])
            if self.u__y is None:
                self.u__y = jacobian(out["u"], out["y"])
            if self.v__x is None:
                self.v__x = jacobian(out["v"], out["x"])
            if self.v__y is None:
                self.v__y = jacobian(out["v"], out["y"])

            if self.dim == 3:
                if self.u__z is None:
                    self.u__z = jacobian(out["u"], out["z"])
                if self.v__z is None:
                    self.v__z = jacobian(out["v"], out["z"])
                if self.w__x is None:
                    self.w__x = jacobian(out["w"], out["x"])
                if self.w__y is None:
                    self.w__y = jacobian(out["w"], out["y"])
                if self.w__z is None:
                    self.w__z = jacobian(out["w"], out["z"])

        # Energy equations
        def E_xy_compute_func(out):
            init_params(out)
            sigma_xx = self.mu * self.u__x + self.lambda_ * (self.u__x + self.v__y)
            sigma_xy = 0.5 * self.mu * (self.u__y + self.v__x)
            sigma_yy = self.mu * self.v__y + self.lambda_ * (self.u__x + self.v__y)
            E_xy = 0.5 * (
                self.u__x * sigma_xx
                + (self.u__y + self.v__x) * sigma_xy
                + self.v__y * sigma_yy
            )
            return E_xy

        self.add_equation("E_xy", E_xy_compute_func)

        if self.dim == 3:

            def E_xyz_compute_func(out):
                init_params(out)
                eps12 = 0.5 * self.u__y + 0.5 * self.v__x
                eps13 = 0.5 * self.u__z + 0.5 * self.w__x
                eps23 = 0.5 * self.v__z + 0.5 * self.w__y

                trace_strain = self.u__x + self.v__y + self.w__z
                squared_diagonal = (
                    self.u__x * self.u__x
                    + self.v__y * self.v__y
                    + self.w__z * self.w__z
                )
                E_xyz = 0.5 * self.lambda_ * trace_strain * trace_strain + self.mu * (
                    squared_diagonal
                    + 2.0 * eps12 * eps12
                    + 2.0 * eps13 * eps13
                    + 2.0 * eps23 * eps23
                )
                return E_xyz

            self.add_equation("E_xyz", E_xyz_compute_func)

import equation as pdb_module
import numpy as np
import paddle
import paddle.nn.functional as F
from skimage.filters import gaussian

import ppsci
from ppsci.utils import logger


class GradientReversalLayer(paddle.autograd.PyLayer):
    @staticmethod
    def forward(ctx, x):
        return x.clone()

    @staticmethod
    def backward(ctx, grad_output):
        return -grad_output


class Problems:
    def __init__(self, cfg, geo_origin, geo_dim):
        self.cfg = cfg
        self.geo_origin = geo_origin
        self.geo_dim = geo_dim
        self.dim = len(geo_dim)
        self.volume = self.comput_volume()
        self.batch_raito = 1.0

        self.lambda_ = (
            (cfg.NU * cfg.E / ((1 - cfg.NU * cfg.NU)))
            if self.dim == 2
            else cfg.NU * cfg.E / ((1 + cfg.NU) * (1 - 2 * cfg.NU))
        )
        self.mu = cfg.E / (1 + cfg.NU) if self.dim == 2 else cfg.E / (2 * (1 + cfg.NU))
        self.equation = {
            "EEquation": pdb_module.EEquation(
                param_dict={"lambda_": self.lambda_, "mu": self.mu}, dim=self.dim
            ),
        }

        self.disp_net = None
        self.density_net = None

        self.volume_ratio = getattr(cfg, "VOLUME_RATIO", 0.5)
        self.alpha = getattr(cfg, "SIGMOID_ALPHA", 5.0)
        self.exponent = getattr(cfg, "ENERGY_EXP", 3.0)
        self.vol_penalty_strength = getattr(cfg, "PENALTY", 10.0)
        self.use_mmse = getattr(cfg, "USE_MMSE", False)
        self.use_oc = getattr(cfg, "USE_OC", False)
        self.max_move = getattr(cfg, "MAX_MOVE", 0.2)
        self.damping_parameter = getattr(cfg, "DAMPING", 0.5)
        self.filter = getattr(cfg, "FILTER", None)
        self.filter_radius = getattr(cfg, "FILTER_RADIUS", 2.0)

    def comput_volume(self):
        return np.prod(self.geo_dim)

    # transforms
    def transform_in(self, _in):
        x, y = _in["x"], _in["y"]
        x_scaled = 2.0 / self.geo_dim[0] * x + (
            -1.0 - 2.0 * self.geo_origin[0] / self.geo_dim[0]
        )
        y_scaled = 2.0 / self.geo_dim[1] * y + (
            -1.0 - 2.0 * self.geo_origin[1] / self.geo_dim[1]
        )

        sin_x_scaled, sin_y_scaled = paddle.sin(x_scaled), paddle.sin(y_scaled)

        in_trans = {
            "x_scaled": x_scaled,
            "y_scaled": y_scaled,
            "sin_x_scaled": sin_x_scaled,
            "sin_y_scaled": sin_y_scaled,
        }

        if self.dim == 3:
            z = _in["z"]
            z_scaled = 2.0 / self.geo_dim[2] * z + (
                -1.0 - 2.0 * self.geo_origin[2] / self.geo_dim[2]
            )
            sin_z_scaled = paddle.sin(z_scaled)
            in_trans["z_scaled"] = z_scaled
            in_trans["sin_z_scaled"] = sin_z_scaled

        return in_trans

    def transform_out_disp(self, _in, _out):
        "Different for each problems because of different boundary constraints."
        # logger.info("In default out transform of disp net")
        return _out

    def transform_out_density(self, _in, _out):
        density = _out["density"]
        offset = np.log(self.volume_ratio / (1.0 - self.volume_ratio))
        densities = F.sigmoid(self.alpha * density + offset)
        return {"densities": densities}

    # functions
    def compute_E(self, densities, E, reverse_grad=False):
        if reverse_grad:
            densities = GradientReversalLayer.apply(densities)
        E_densities = paddle.pow(densities, self.exponent) * E
        return self.volume * paddle.mean(E_densities, keepdim=True)

    def compute_force(self):
        "Different for each problems because of different force."
        logger.info("In default force compute function")
        return 0.0

    def compute_penalty(self, densities):
        target_volume = self.volume_ratio * self.volume
        volume_estimate = self.volume * paddle.mean(densities, keepdim=True)
        return (
            self.vol_penalty_strength
            * (volume_estimate - target_volume)
            * (volume_estimate - target_volume)
            / target_volume
        )

    # oc
    def compute_target_densities(self, densities_list, sensitivities_list):
        if self.use_oc:
            return self.compute_oc_multi_batch(densities_list, sensitivities_list)
        else:
            return self.compute_target_densities_gradient_descent(
                densities_list, sensitivities_list
            )

    def compute_oc_multi_batch(self, old_densities, sensitivities):
        target_volume = self.volume_ratio * self.volume
        logger.info(f"target_volume: {target_volume}")

        lagrange_lower_estimate = 0
        lagrange_upper_estimate = 1e9
        conv_threshold = 1e-3
        total_samples = len(old_densities) * old_densities[0].shape[0]
        dv = self.volume / total_samples

        density_lower_bound = [
            paddle.maximum(paddle.to_tensor(0.0), od - self.max_move)
            for od in old_densities
        ]
        density_upper_bound = [
            paddle.minimum(paddle.to_tensor(1.0), od + self.max_move)
            for od in old_densities
        ]

        while (lagrange_upper_estimate - lagrange_lower_estimate) / (
            lagrange_lower_estimate + lagrange_upper_estimate
        ) > conv_threshold:
            lagrange_current = 0.5 * (lagrange_upper_estimate + lagrange_lower_estimate)

            target_densities_div = [
                paddle.divide(
                    di,
                    paddle.to_tensor(
                        -dv * lagrange_current, dtype=paddle.get_default_dtype()
                    ),
                )
                for di in sensitivities
            ]
            di_div_max = float(paddle.max(paddle.concat(target_densities_div)))
            if np.isinf(di_div_max):
                logger.info("Warning! target_densities_div is nan")
                exit()

            target_densities_mul = [
                paddle.multiply(
                    old_densities[i],
                    paddle.pow(target_densities_div[i], self.damping_parameter),
                )
                for i in range(len(old_densities))
            ]

            target_densities_clip = [
                paddle.maximum(
                    density_lower_bound[i],
                    paddle.minimum(density_upper_bound[i], target_densities_mul[i]),
                )
                for i in range(len(old_densities))
            ]

            new_volume = self.volume * np.mean(
                [paddle.mean(di) for di in target_densities_clip]
            )

            if new_volume > target_volume:
                lagrange_lower_estimate = lagrange_current
            else:
                lagrange_upper_estimate = lagrange_current
        logger.info(f"new_volume: {new_volume}")
        return target_densities_clip

    def compute_target_densities_gradient_descent(
        self, densities_list, sensitivities_list
    ):
        projected_sensitivities = [
            (
                paddle.maximum(
                    paddle.to_tensor(0.0),
                    paddle.minimum(
                        paddle.to_tensor(1.0), densities_list[i] - sensitivities_list[i]
                    ),
                )
                - densities_list[i]
            )
            for i in range(len(densities_list))
        ]

        step_size = 0.05 / paddle.mean(
            paddle.to_tensor([paddle.abs(si) for si in projected_sensitivities]),
            keepdim=True,
        )
        return [
            densities_list[i] - step_size * sensitivities_list[i]
            for i in range(len(densities_list))
        ]

    # filter
    def gaussian_filter(self, sensitivities):
        sensitivities = paddle.reshape(sensitivities, self.batch_size)
        sensitivities_blur = gaussian(sensitivities.numpy(), 3).reshape([-1, 1])
        return paddle.to_tensor(sensitivities_blur, dtype=paddle.get_default_dtype())

    # loss functions
    def disp_loss_func(
        self,
        output_dict,
        label_dict=None,
        weight_dict={},
    ):
        densities = output_dict["densities"].detach().clone()
        E = output_dict["E_xy"] if self.dim == 2 else output_dict["E_xyz"]
        loss_E = self.compute_E(densities, E)
        loss_force = self.compute_force()
        return {"loss_E": loss_E, "loss_force": loss_force}

    def density_loss_func(
        self,
        output_dicts_list,
        label_dicts_list=None,
        weight_dicts_list=None,
        input_dicts_list=None,
    ):
        loss_list = []
        densities_list = []
        sensitivities_list = []

        if not isinstance(output_dicts_list, list):
            output_dicts_list = [output_dicts_list]

        for i, output_dict in enumerate(output_dicts_list):
            if isinstance(output_dict, list):
                output_dict = output_dict[0]

            densities = output_dict["densities"]
            E = output_dict["E_xy"] if self.dim == 2 else output_dict["E_xyz"]
            ppsci.autodiff.clear()

            loss_E = self.compute_E(densities, E, reverse_grad=True)
            loss = loss_E
            if not self.use_oc:
                loss_penalty = self.compute_penalty(densities)
                loss += loss_penalty
            loss_list.append(loss)

            sensitivities = paddle.grad(loss, densities)[0].detach().clone()
            densities = densities.detach().clone()

            # filter
            if self.filter == "Gaussian":
                sensitivities = self.gaussian_filter(sensitivities)

            densities_list.append(densities)
            sensitivities_list.append(sensitivities)
            ppsci.autodiff.clear()

        if not self.use_mmse:
            return {"loss_density": loss_list[0]}
        else:
            if input_dicts_list is None:
                raise ValueError("When use mmse, 'input_dicts_list' should not be None")

            target_densities_list = self.compute_target_densities(
                densities_list, sensitivities_list
            )
            logger.info(
                f"target density: {np.mean([td.numpy().mean() for td in target_densities_list])}"
            )

            def oc_loss_func(model, i):
                # cannot use `output_dict["densities"]` to get densities_i
                # because model is updated every time before entering this function
                input_dict = input_dicts_list[i]
                if isinstance(input_dict, list):
                    input_dict = input_dict[0]
                densities_i = model(input_dict)["densities"]
                loss = F.mse_loss(densities_i, target_densities_list[i], "mean")
                return loss

            return oc_loss_func

    # eval metric functions
    def density_metric_func(self, output_dict, *args):
        density = output_dict["densities"]
        logger.info(f"mean: {float(paddle.mean(density))}")
        logger.info(f"max: {float(paddle.max(density))}")
        logger.info(f"min: {float(paddle.min(density))}")
        metric_dict = {"densities": density.mean() - self.volume_ratio}
        return metric_dict


class Beam2D(Problems):
    def __init__(self, cfg):
        geo_origin = (0.0, 0.0)
        geo_dim = (1.5, 0.5)
        super().__init__(cfg, geo_origin, geo_dim)

        beam = ppsci.geometry.Rectangle((0.0, 0.0), (1.5, 0.5))
        self.geom = {"geo": beam}
        self.force = -0.0025
        self.mirror = None
        self.batch_size = (150, 50)

    # bc
    def transform_out_disp(self, _in, _out):
        x_scaled = _in["x_scaled"]
        x = self.geo_dim[0] / 2 * (1 + x_scaled) + self.geo_origin[0]
        u, v = x * _out["u"], x * _out["v"]
        return {"u": u, "v": v}

    # force
    def compute_force(self):
        input_pos = {
            "x": paddle.to_tensor([[1.5]], dtype=paddle.get_default_dtype()),
            "y": paddle.to_tensor([[0.0]], dtype=paddle.get_default_dtype()),
        }
        output_force = self.disp_net(input_pos)
        v = output_force["v"]
        return -paddle.mean(v * self.force, keepdim=True)


class Distributed2D(Problems):
    def __init__(self, cfg):
        geo_origin = (0.0, 0.0)
        geo_dim = (1.5, 0.5)
        super().__init__(cfg, geo_origin, geo_dim)

        beam = ppsci.geometry.Rectangle((0.0, 0.0), (1.5, 0.5))
        self.geom = {"geo": beam}
        self.force = -0.0025
        self.mirror = None
        self.batch_size = (150, 50)

    # bc
    def transform_out_disp(self, _in, _out):
        x_scaled = _in["x_scaled"]
        x = self.geo_dim[0] / 2 * (1 + x_scaled) + self.geo_origin[0]
        u, v = x * _out["u"], x * _out["v"]
        return {"u": u, "v": v}

    # force
    def get_force_pos(self):
        sample_num = 400
        input_pos_np = self.geom["geo"].sample_boundary(
            n=sample_num,
            criteria=lambda x, y: y >= self.geo_dim[1] - 1e-3,
        )
        return {
            "x": paddle.to_tensor(input_pos_np["x"], dtype=paddle.get_default_dtype()),
            "y": paddle.full((sample_num, 1), 0.5, dtype=paddle.get_default_dtype()),
        }

    def compute_force(self):
        input_pos = self.get_force_pos()
        output_force = self.disp_net(input_pos)
        v = output_force["v"]
        return -paddle.mean(v * self.force, keepdim=True)


class LongBeam2D(Problems):
    def __init__(self, cfg):
        geo_origin = (0.0, 0.0)
        geo_dim = (1.0, 0.5)
        super().__init__(cfg, geo_origin, geo_dim)

        long_beam = ppsci.geometry.Rectangle(geo_origin, geo_dim)
        self.geom = {"geo": long_beam}
        self.force = -0.0025
        self.mirror = [True, False]
        self.batch_size = (122, 61)

    # bc
    def transform_out_disp(self, _in, _out):
        x_scaled = _in["x_scaled"]
        x = self.geo_dim[0] / 2 * (1 + x_scaled) + self.geo_origin[0]
        u, v = (
            x * (x - 1) * _out["u"],
            x * _out["v"],
        )
        return {"u": u, "v": v}

    # force
    def compute_force(self):
        input_pos = {
            "x": paddle.to_tensor([[1.0]], dtype=paddle.get_default_dtype()),
            "y": paddle.to_tensor([[0.0]], dtype=paddle.get_default_dtype()),
        }
        output_force = self.disp_net(input_pos)
        v = output_force["v"]
        return -paddle.mean(v * self.force, keepdim=True)


class Bridge2D(Problems):
    def __init__(self, cfg):
        geo_origin = (0.0, 0.0)
        geo_dim = (1.0, 0.5)
        super().__init__(cfg, geo_origin, geo_dim)

        bridge = ppsci.geometry.Rectangle(geo_origin, geo_dim)
        self.geom = {"geo": bridge}
        self.force = -0.0025
        self.mirror = [True, False]
        self.batch_size = (122, 61)

    # bc
    def transform_out_disp(self, _in, _out):
        x_scaled = _in["x_scaled"]
        x = self.geo_dim[0] / 2 * (1 + x_scaled) + self.geo_origin[0]
        u, v = x * (x - 1) * _out["u"], x * _out["v"]
        return {"u": u, "v": v}

    # force
    def get_force_pos(self):
        sample_num = 400
        input_pos_np = self.geom["geo"].sample_boundary(
            n=sample_num,
            criteria=lambda x, y: y <= self.geo_origin[1] + 1e-3,
        )
        return {
            "x": paddle.to_tensor(input_pos_np["x"], dtype=paddle.get_default_dtype()),
            "y": paddle.full((sample_num, 1), 0.0, dtype=paddle.get_default_dtype()),
        }

    def compute_force(self):
        input_pos = self.get_force_pos()
        output_force = self.disp_net(input_pos)
        v = output_force["v"]
        return -paddle.mean(v * self.force, keepdim=True)


class Triangle2D(Problems):
    def __init__(self, cfg):
        geo_origin = (0.0, 0.0)
        geo_dim = (2.0, 3**0.5)
        super().__init__(cfg, geo_origin, geo_dim)

        triangle = ppsci.geometry.Triangle((0.0, 0.0), (2.0, 0.0), (1.0, 3**0.5))
        self.geom = {"geo": triangle}
        force = 0.0025
        self.mirror = None
        self.force = [
            [-(3**0.5) * 0.5 * force, -0.5 * force],
            [(3**0.5) * 0.5 * force, -0.5 * force],
            [0.0, 1 * force],
        ]
        self.batch_size = (200, 170)
        self.volume = 3**0.5
        self.batch_raito = 0.5
        self.filter = (
            None  # Gaussian filtering cannot be used for non-rectangular shapes
        )

    # force
    def compute_force(self):
        input_pos = {
            "x": paddle.to_tensor(
                [[0.0], [2.0], [1.0]], dtype=paddle.get_default_dtype()
            ),
            "y": paddle.to_tensor(
                [[0.0], [0.0], [3**0.5]], dtype=paddle.get_default_dtype()
            ),
        }
        output_force = self.disp_net(input_pos)
        u, v = output_force["u"], output_force["v"]
        force = paddle.to_tensor(self.force)
        return -paddle.mean(
            paddle.multiply(force[:, 0], u[:, 0])
            + paddle.multiply(force[:, 1], v[:, 0]),
            keepdim=True,
        )


class TriangleVariants2D(Problems):
    def __init__(self, cfg):
        geo_origin = (0.0, 0.0)
        geo_dim = (2.0, 3**0.5)
        super().__init__(cfg, geo_origin, geo_dim)

        triangle = ppsci.geometry.Triangle((0.0, 0.0), (2.0, 0.0), (1.0, 3**0.5))
        disk = ppsci.geometry.Disk((1.0, 3**0.5 / 3), 0.1)
        self.geom = {"geo": triangle - disk}
        force = 0.0025
        self.mirror = None
        self.force = [
            [-(3**0.5) * 0.5 * force, -0.5 * force],
            [(3**0.5) * 0.5 * force, -0.5 * force],
            [0.0, 1.0 * force],
        ]
        self.batch_size = (200, 170)
        self.volume = 3**0.5 - np.pi * 0.01
        self.batch_raito = 0.4
        self.filter = (
            None  # Gaussian filtering cannot be used for non-rectangular shapes
        )

    # bc
    def transform_out_disp(self, _in, _out):
        x_scaled, y_scaled = _in["x_scaled"], _in["y_scaled"]
        x = self.geo_dim[0] / 2 * (1 + x_scaled) + self.geo_origin[0]
        y = self.geo_dim[1] / 2 * (1 + y_scaled) + self.geo_origin[1]
        constraint = (x - 1) ** 2 + (y - 3**0.5 / 3) ** 2 - 0.15**2
        u, v = constraint * _out["u"], constraint * _out["v"]
        return {"u": u, "v": v}

    # force
    def compute_force(self):
        input_pos = {
            "x": paddle.to_tensor(
                [[0.0], [2.0], [1.0]], dtype=paddle.get_default_dtype()
            ),
            "y": paddle.to_tensor(
                [[0.0], [0.0], [3**0.5]], dtype=paddle.get_default_dtype()
            ),
        }
        output_force = self.disp_net(input_pos)
        u, v = output_force["u"], output_force["v"]
        force = paddle.to_tensor(self.force)
        return -paddle.mean(
            paddle.multiply(force[:, 0], u[:, 0])
            + paddle.multiply(force[:, 1], v[:, 0]),
            keepdim=True,
        )


class LShape2D(Problems):
    def __init__(self, cfg):
        geo_origin = (0.0, 0.0)
        geo_dim = (1.5, 1.5)
        super().__init__(cfg, geo_origin, geo_dim)

        rec_1 = ppsci.geometry.Rectangle((0.0, 0.0), (1.5, 1.5))
        rec_2 = ppsci.geometry.Rectangle((0.5, 0.5), (1.5, 1.5))
        custom_geo = rec_1 - rec_2
        self.geom = {"geo": custom_geo}

        self.force = -0.0025
        self.mirror = None
        self.batch_size = (150, 150)  # 12500 = 150 * 150 - 100 * 100
        self.volume = 1.25
        self.batch_raito = 5.0 / 9.0
        self.filter = (
            None  # Gaussian filtering cannot be used for non-rectangular shapes
        )

    # bc
    def transform_out_disp(self, _in, _out):
        y_scaled = _in["y_scaled"]
        y = self.geo_dim[0] / 2 * (1 + y_scaled) + self.geo_origin[0]
        u, v = (y - 1.5) * _out["u"], (y - 1.5) * _out["v"]
        return {"u": u, "v": v}

    # force
    def compute_force(self):
        input_pos = {
            "x": paddle.to_tensor([[1.5]], dtype=paddle.get_default_dtype()),
            "y": paddle.to_tensor([[0.5]], dtype=paddle.get_default_dtype()),
        }
        output_force = self.disp_net(input_pos)
        v = output_force["v"]
        return -paddle.mean(v * self.force, keepdim=True)


class Beam3D(Problems):
    def __init__(self, cfg):
        geo_origin = (0.0, 0.0, 0.0)
        geo_dim = (1.0, 0.5, 0.25)
        super().__init__(cfg, geo_origin, geo_dim)

        beam = ppsci.geometry.Cuboid(geo_origin, geo_dim)
        self.geom = {"geo": beam}
        self.force = (0.0, -0.0005, 0.0)
        self.mirror = [False, False, True]
        self.batch_size = (40, 20, 10)  # (80, 40, 20)
        self.input_force = self.get_input_force(sample_num=400)
        self.force_volume = self.geo_dim[2] - self.geo_origin[2]

    # bc
    def transform_out_disp(self, _in, _out):
        x_scaled, z_scaled = _in["x_scaled"], _in["z_scaled"]
        x = self.geo_dim[0] / 2 * (1 + x_scaled) + self.geo_origin[0]
        z = self.geo_dim[2] / 2 * (1 + z_scaled) + self.geo_origin[2]
        u, v, w = x * _out["u"], x * _out["v"], x * (z - self.geo_dim[2]) * _out["w"]
        return {"u": u, "v": v, "w": w}

    # force
    def get_input_force(self, sample_num):
        input_pos_np = self.geom["geo"].sample_boundary(
            n=sample_num,
            criteria=lambda x, y, z: np.logical_and(
                x >= self.geo_dim[0] - 1e-5,
                y <= self.geo_origin[1] + 1e-5,
            ),
        )
        return {
            "x": paddle.to_tensor(input_pos_np["x"], dtype=paddle.get_default_dtype()),
            "y": paddle.to_tensor(input_pos_np["y"], dtype=paddle.get_default_dtype()),
            "z": paddle.to_tensor(input_pos_np["z"], dtype=paddle.get_default_dtype()),
        }

    def compute_force(self):
        output_force = self.disp_net(self.input_force)
        u, v, w = output_force["u"], output_force["v"], output_force["w"]
        return -self.force_volume * paddle.mean(
            (u * self.force[0] + v * self.force[1] + w * self.force[2]), keepdim=True
        )


class Bridge3D(Problems):
    def __init__(self, cfg):
        geo_origin = (0.0, 0.0, 0.0)
        geo_dim = (1.0, 0.5, 0.25)
        super().__init__(cfg, geo_origin, geo_dim)

        bridge = ppsci.geometry.Cuboid(geo_origin, geo_dim)
        self.geom = {"geo": bridge}
        self.force = (0.0, -0.0025, 0.0)
        self.mirror = [True, False, True]
        self.batch_size = (40, 20, 10)  # (80, 40, 20)
        self.input_force = self.get_input_force(sample_num=500 * 1 * 125)
        self.force_volume = (geo_dim[0] - geo_origin[0]) * (geo_dim[2] - geo_origin[2])

    # bc
    def transform_out_disp(self, _in, _out):
        x_scaled, z_scaled = _in["x_scaled"], _in["z_scaled"]
        x = self.geo_dim[0] / 2 * (1 + x_scaled) + self.geo_origin[0]
        z = self.geo_dim[2] / 2 * (1 + z_scaled) + self.geo_origin[2]
        u, v, w = (
            x * (x - self.geo_dim[0]) * _out["u"],
            x * _out["v"],
            x * (z - self.geo_dim[2]) * _out["w"],
        )
        return {"u": u, "v": v, "w": w}

    # force
    def get_input_force(self, sample_num):
        input_pos_np = self.geom["geo"].sample_boundary(
            n=sample_num,
            criteria=lambda x, y, z: y == self.geo_origin[0],
        )
        return {
            "x": paddle.to_tensor(input_pos_np["x"], dtype=paddle.get_default_dtype()),
            "y": paddle.to_tensor(input_pos_np["y"], dtype=paddle.get_default_dtype()),
            "z": paddle.to_tensor(input_pos_np["z"], dtype=paddle.get_default_dtype()),
        }

    def compute_force(self):
        output_force = self.disp_net(self.input_force)
        u, v, w = output_force["u"], output_force["v"], output_force["w"]
        return -self.force_volume * paddle.mean(
            (u * self.force[0] + v * self.force[1] + w * self.force[2]), keepdim=True
        )

# 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.

from typing import Callable
from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
from typing import Tuple
from typing import Union

import numpy as np
import paddle
from paddle import nn
from paddle.distributed.fleet.utils import hybrid_parallel_util as hpu
from skimage import measure

import ppsci
from ppsci import geometry
from ppsci.autodiff import clear
from ppsci.solver import Solver
from ppsci.utils import logger
from ppsci.utils import misc
from ppsci.utils import save_load


class Trainer:
    def __init__(self, solver: "Solver") -> None:
        self.solver = solver

    def run_forward(
        self,
        expr_dicts: Tuple[Dict[str, Callable], ...],
        input_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
        model: nn.Layer,
    ) -> Tuple["paddle.Tensor", ...]:
        """Forward computation for training, including model forward and equation
        forward.

        Args:
            expr_dicts (Tuple[Dict[str, Callable], ...]): Tuple of expression dicts.
            input_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of input dicts.
            model (nn.Layer): NN model.

        Returns:
            Tuple[paddle.Tensor, ...]: Tuple of output for each expression.
        """
        output_dicts = []
        for i, expr_dict in enumerate(expr_dicts):
            # model forward
            output_dict = model(input_dicts[i])

            # equation forward
            data_dict = {k: v for k, v in input_dicts[i].items()}
            data_dict.update(output_dict)
            for name, expr in expr_dict.items():
                output_dict[name] = expr(data_dict)

            # put field 'area' into output_dict
            if "area" in input_dicts[i]:
                output_dict["area"] = input_dicts[i]["area"]

            output_dicts.append(output_dict)

            # clear differentiation cache
            clear()
        return output_dicts

    def train_forward(self):
        input_dicts_list = []
        label_dicts_list = []
        weight_dicts_list = []
        output_dicts_list = []
        for _ in range(1, self.solver.iters_per_epoch + 1):
            loss_dict = misc.Prettydefaultdict(float)
            loss_dict["loss"] = 0.0

            input_dicts = []
            label_dicts = []
            weight_dicts = []
            for _, _constraint in self.solver.constraint.items():
                try:
                    input_dict, label_dict, weight_dict = next(_constraint.data_iter)
                except StopIteration:
                    _constraint.data_iter = iter(_constraint.data_loader)
                    input_dict, label_dict, weight_dict = next(_constraint.data_iter)

                for v in input_dict.values():
                    v.stop_gradient = False

                # gather each constraint's input, label, weight to a list
                input_dicts.append(input_dict)
                label_dicts.append(label_dict)
                weight_dicts.append(weight_dict)

            output_dicts = self.run_forward(
                (
                    _constraint.output_expr
                    for _constraint in self.solver.constraint.values()
                ),
                input_dicts,
                self.solver.model,
            )
            input_dicts_list.append(input_dicts)
            label_dicts_list.append(label_dicts)
            weight_dicts_list.append(weight_dicts)
            output_dicts_list.append(output_dicts)
        return input_dicts_list, label_dicts_list, weight_dicts_list, output_dicts_list

    def train_backward(self, constraint_losses_list):
        for iter_id in range(1, self.solver.iters_per_epoch + 1):
            # compute loss for each constraint according to its' own output, label and weight
            total_loss = 0
            for i, _ in enumerate(self.solver.constraint.values()):
                loss = constraint_losses_list[i]
                if isinstance(loss, Callable):
                    total_loss += loss(self.solver.model, iter_id - 1)
                elif isinstance(loss, List):
                    total_loss += loss[iter_id - 1]

            if self.solver.use_amp:
                total_loss_scaled = self.solver.scaler.scale(total_loss)
                total_loss_scaled.backward()
            else:
                total_loss.backward()

            # update parameters
            if (
                iter_id % self.solver.update_freq == 0
                or iter_id == self.solver.iters_per_epoch
            ):
                if self.solver.world_size > 1:
                    # fuse + allreduce manually before optimization if use DDP + no_sync
                    # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622
                    hpu.fused_allreduce_gradients(
                        list(self.solver.model.parameters()), None
                    )
                if self.solver.use_amp:
                    self.solver.scaler.minimize(
                        self.solver.optimizer, total_loss_scaled
                    )
                else:
                    self.solver.optimizer.step()
                self.solver.optimizer.clear_grad()

            # update learning rate by step
            if (
                self.solver.lr_scheduler is not None
                and not self.solver.lr_scheduler.by_epoch
            ):
                self.solver.lr_scheduler.step()

            # update and log training information
            self.solver.global_step += 1

    def train_batch(self):
        """Training."""
        self.solver.global_step = (
            self.solver.best_metric["epoch"] * self.solver.iters_per_epoch
        )

        for epoch_id in range(
            self.solver.best_metric["epoch"] + 1, self.solver.epochs + 1
        ):
            # forward
            (
                input_dicts_list,
                label_dicts_list,
                weight_dicts_list,
                output_dicts_list,
            ) = self.train_forward()

            # batch loss
            constraint_losses_list = []
            for _, _constraint in enumerate(self.solver.constraint.values()):
                if not isinstance(_constraint.loss, FunctionalLossBatch):
                    raise TypeError(
                        "Loss function of constraint should be FunctionalLossBatch when using train_batch"
                    )
                constraint_loss_list = _constraint.loss(
                    output_dicts_list,
                    label_dicts_list,
                    weight_dicts_list,
                    input_dicts_list,
                )
                constraint_losses_list.append(constraint_loss_list)

            # backward
            self.train_backward(constraint_losses_list)

            # log training summation at end of a epoch
            metric_msg = ", ".join(
                [
                    self.solver.train_output_info[key].avg_info
                    for key in self.solver.train_output_info
                ]
            )
            logger.info(
                f"[Train][Epoch {epoch_id}/{self.solver.epochs}][Avg] {metric_msg}"
            )
            self.solver.train_output_info.clear()

            cur_metric = float("inf")
            # evaluate during training
            if (
                self.solver.eval_during_train
                and epoch_id % self.solver.eval_freq == 0
                and epoch_id >= self.solver.start_eval_epoch
            ):
                cur_metric, metric_dict_group = self.eval(epoch_id)
                if cur_metric < self.solver.best_metric["metric"]:
                    self.solver.best_metric["metric"] = cur_metric
                    self.solver.best_metric["epoch"] = epoch_id
                    save_load.save_checkpoint(
                        self.solver.model,
                        self.solver.optimizer,
                        self.solver.best_metric,
                        self.solver.scaler,
                        self.solver.output_dir,
                        "best_model",
                        self.solver.equation,
                    )
                logger.info(
                    f"[Eval][Epoch {epoch_id}]"
                    f"[best metric: {self.solver.best_metric['metric']}]"
                    f"[best epoch: {self.solver.best_metric['epoch']}]"
                    f"[current metric: {cur_metric}]"
                )
                for metric_dict in metric_dict_group.values():
                    logger.scaler(
                        {f"eval/{k}": v for k, v in metric_dict.items()},
                        epoch_id,
                        self.vdl_writer,
                        self.wandb_writer,
                    )

                # visualize after evaluation
                if self.solver.visualizer is not None:
                    self.solver.visualize(epoch_id)

            # update learning rate by epoch
            if (
                self.solver.lr_scheduler is not None
                and self.solver.lr_scheduler.by_epoch
            ):
                self.solver.lr_scheduler.step()

            # save epoch model every save_freq epochs
            if self.solver.save_freq > 0 and epoch_id % self.solver.save_freq == 0:
                save_load.save_checkpoint(
                    self.solver.model,
                    self.solver.optimizer,
                    {"metric": cur_metric, "epoch": epoch_id},
                    self.solver.scaler,
                    self.solver.output_dir,
                    f"epoch_{epoch_id}",
                    self.solver.equation,
                )

            # save the latest model for convenient resume training
            save_load.save_checkpoint(
                self.solver.model,
                self.solver.optimizer,
                {"metric": cur_metric, "epoch": epoch_id},
                self.solver.scaler,
                self.solver.output_dir,
                "latest",
                self.solver.equation,
            )

        # close VisualDL
        if self.solver.vdl_writer is not None:
            self.solver.vdl_writer.close()


class FunctionalLossBatch(ppsci.loss.base.Loss):
    def __init__(
        self,
        loss_expr: Callable,
        reduction: Literal["mean", "sum"] = "mean",
        weight: Optional[Union[float, Dict[str, float]]] = None,
    ):
        if reduction not in ["mean", "sum"]:
            raise ValueError(
                f"reduction should be 'mean' or 'sum', but got {reduction}"
            )
        super().__init__(reduction, weight)
        self.loss_expr = loss_expr

    def forward(
        self,
        output_dicts_list,
        label_dicts_list=None,
        weight_dicts_list=None,
        input_dicts_list=None,
    ):
        return self.loss_expr(
            output_dicts_list, label_dicts_list, weight_dicts_list, input_dicts_list
        )


class Sampler:
    def __init__(
        self,
        geom: geometry.Geometry,
        bounds: Tuple[Tuple[float, float], ...],
        criteria: Optional[Callable] = None,
    ) -> None:
        self.geom = geom
        self.dim = geom.ndim
        self.dim_keys = geom.dim_keys
        self.bounds = bounds
        self.criteria = criteria

    def stratified_random(self, n_samples: Tuple[int, ...]) -> np.ndarray:
        """Stratified random."""
        # Divide the geometry into uniformly sized chunks
        # and randomly sample within each chunk.

        zeros = np.zeros(n_samples)
        grid_points = np.transpose(np.where(zeros == 0)).astype(
            paddle.get_default_dtype()
        )
        all_samples = np.prod(n_samples)
        for i in range(self.dim):
            random = np.random.uniform(0.0, 1.0, (all_samples))
            grid_size = (self.bounds[i][1] - self.bounds[i][0]) / n_samples[i]
            grid_points[:, i] = (grid_points[:, i] + random) * grid_size
        return grid_points

    def sample_interior_stratified(self, n_samples: Tuple[int, ...], n_iter: int):
        """Sample random points in the geometry and return those meet criteria."""
        points = self.stratified_random(n_samples)
        for _ in range(1, n_iter):
            points = np.concatenate((points, self.stratified_random(n_samples)), axis=0)

        if self.criteria is not None:
            criteria_mask = self.criteria(*np.split(points, self.dim, axis=1)).flatten()
            points = points[criteria_mask]
        else:
            criteria_mask = self.geom.is_inside(points)
            points = points[criteria_mask]
        points_dict = {}
        for i in range(self.dim):
            points_dict[self.dim_keys[i]] = points[:, i].reshape([-1, 1])

        return points_dict, criteria_mask


class Plot:
    def __init__(self, filename, problem, n_cells, threshold=0.25) -> None:
        self.filename = filename
        self.problem = problem
        self.n_cells = n_cells
        self.threshold = threshold

    def prepare_data(self):
        cx = 0.5 * self.problem.geo_dim[0] / self.n_cells[0]
        cy = 0.5 * self.problem.geo_dim[1] / self.n_cells[1]
        cz = 0.5 * self.problem.geo_dim[2] / self.n_cells[2]

        x = np.linspace(
            self.problem.geo_origin[0] + cx,
            self.problem.geo_dim[0] - cx,
            num=self.n_cells[0],
            dtype=paddle.get_default_dtype(),
        )
        y = np.linspace(
            self.problem.geo_origin[1] + cy,
            self.problem.geo_dim[1] - cy,
            num=self.n_cells[1],
            dtype=paddle.get_default_dtype(),
        )
        z = np.linspace(
            self.problem.geo_origin[2] + cz,
            self.problem.geo_dim[2] - cz,
            num=self.n_cells[2],
            dtype=paddle.get_default_dtype(),
        )
        xs, ys, zs = np.meshgrid(x, y, z, indexing="ij")

        input_dict = {}
        input_dict["x"] = paddle.to_tensor(
            xs.reshape(-1, 1), dtype=paddle.get_default_dtype()
        )
        input_dict["y"] = paddle.to_tensor(
            ys.reshape(-1, 1), dtype=paddle.get_default_dtype()
        )
        input_dict["z"] = paddle.to_tensor(
            zs.reshape(-1, 1), dtype=paddle.get_default_dtype()
        )

        self.input_dict = input_dict

    def compute_densities(self):
        self.densities = (
            self.problem.density_net(self.input_dict)["densities"]
            .numpy()
            .reshape(self.n_cells)
        )

    def compute_mirror(self):
        if self.problem.mirror[0]:
            self.densities = np.concatenate(
                (
                    self.densities,
                    self.densities[::-1, :, :],
                ),
                axis=0,
            )
        if self.problem.mirror[1]:
            self.densities = np.concatenate(
                (
                    self.densities,
                    self.densities[:, ::-1, :],
                ),
                axis=1,
            )
        if self.problem.mirror[2]:
            self.densities = np.concatenate(
                (
                    self.densities,
                    self.densities[:, :, ::-1],
                ),
                axis=2,
            )

    def pad_with_zeros(self):
        self.density_grid = np.pad(
            self.density_grid, ((1, 1), (1, 1), (1, 1)), "constant", constant_values=0.0
        )

    def save_solid(self):
        if self.problem.mirror:
            for i in range(len(self.problem.mirror)):
                self.n_cells[i] *= 2 if self.problem.mirror[i] else 1
        self.density_grid = np.reshape(
            self.densities, (self.n_cells[0], self.n_cells[1], self.n_cells[2])
        )
        self.pad_with_zeros()

        if (
            np.amax(self.density_grid) < self.threshold
            or np.amin(self.density_grid) > self.threshold
        ):
            print("Warning! Cannot save density grid cause the levelset is empty")
            return

        verts, faces, normals, _ = measure.marching_cubes(
            self.density_grid,
            level=self.threshold,
            spacing=[0.005, 0.005, 0.005],
            gradient_direction="ascent",
            method="lewiner",
        )

        with open(self.filename, "w") as file:
            for item in verts:
                file.write(f"v {item[0]} {item[1]} {item[2]}\n")

            for item in normals:
                file.write(f"vn {item[0]} {item[1]} {item[2]}\n")

            for item in faces:
                idx_0 = item[0] + 1
                idx_1 = item[1] + 1
                idx_2 = item[2] + 1
                file.write(f"f {idx_0}//{idx_0} {idx_1}//{idx_1} {idx_2}//{idx_2}\n")

    def plot_3d(self):
        self.prepare_data()
        self.compute_densities()
        if self.problem.mirror:
            self.compute_mirror()
        self.save_solid()

5. 结果展示

下面展示了不同问题上的优化结果。

序号	问题名称	预训练模型	结果
1	Beam2D	beam2d_pretrained.pdparams
2	Bridge2D	bridge2d_pretrained.pdparams
3	Distributed2D	distributed2d_pretrained.pdparams
4	LongBeam2D	longbeam2d_pretrained.pdparams
5	LShape2D	lshape2d_pretrained.pdparams
6	Triangle2D	triangle2d_pretrained.pdparams
7	TriangleVariants2D	trianglevariants2d_pretrained.pdparams
8	Beam3D	beam3d_pretrained.pdparams
9	Bridge3D	bridge3d_pretrained.pdparams

6. 参考文献

• NTopo: Mesh-free Topology Optimization using Implicit Neural Representations

• 参考代码