Skip to content

Chip Heat Simulation

AI Studio Quick Experience

python chip_heat.py

python chip_heat.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/ChipHeat/chip_heat_pretrained.pdparams
Pretrained Model Metrics
chip_heat_pretrained.pdparams MSE.chip(down_mse): 0.04177
MSE.chip(left_mse): 0.01783
MSE.chip(right_mse): 0.03767
MSE.chip(top_mse): 0.05034

1. Background Introduction

Chip thermal simulation research mainly focuses on predicting and analyzing the temperature distribution of integrated circuits (ICs) during operation, as well as the impact of thermal effects on chip performance, power consumption, reliability, and lifespan. As electronic devices evolve towards higher performance, higher density, and smaller sizes, thermal management has become a critical challenge in chip design and manufacturing.

Chip thermal simulation research provides important tools and methods for understanding and solving chip thermal management problems, playing a crucial role in improving chip performance, reducing power consumption, ensuring reliability, and extending lifespan. As electronic devices develop towards higher performance and compactness, the importance of thermal simulation research will further increase.

Chip thermal simulation has multiple importances in engineering and scientific fields, mainly reflected in the following aspects:

  • Design optimization and validation: Chip thermal simulation can help engineers and scientists evaluate the thermal characteristics of different structures and materials in the early stages of design to optimize design and verify its reliability. By simulating temperature distribution and heat conduction effects under different workloads, potential thermal problems can be discovered in advance and targeted improvements can be made, thereby reducing later development costs and risks.
  • Thermal management and heat dissipation design: Chip thermal simulation can help design effective thermal management systems and heat dissipation schemes to ensure that the chip remains within a safe operating temperature range during long-term high-load operation. By analyzing the heat dissipation structure, fan configuration, heat sink design, etc. around the chip, heat conduction and heat dissipation efficiency can be optimized, improving system stability and reliability.
  • Performance prediction and optimization: Temperature has a significant impact on chip performance and stability. Chip thermal simulation can help predict chip performance under different workloads and environmental conditions, including processor speed, power consumption, and lifespan of electronic devices. By modeling and analyzing thermal effects, chip design and operating conditions can be optimized to achieve better performance and reliability.
  • Energy saving and environmental protection: Effective thermal management and heat dissipation design can reduce system energy consumption and improve energy utilization efficiency, thereby achieving energy saving and environmental protection goals. By reducing heat loss and waste in the system, energy consumption and carbon emissions can be reduced, minimizing negative impacts on the environment.

In summary, chip thermal simulation plays an important role and value in engineering and scientific fields, helping to achieve positive results in design optimization, performance improvement, cost reduction, environmental protection, etc.

2. Problem Definition

2.1 Problem Description

To build a general thermal simulation model, we first briefly describe the thermal simulation problem in general. Thermal simulation aims to predict the temperature field of a given object by globally solving the heat conduction equation, which can usually be represented by the following governing equation:

\[ k \Delta T(x,t) + S(x,t) = \rho c_p \dfrac{\partial T(x,t)}{\partial t},\quad \text { in } \Omega\times (0,t_{*}), \]

where \(\Omega\subset \mathbb{R}^{n},~n=1,2,3\) is the simulation area of the given object material, as shown in the figure for a 2D chip simulation area with random heat source distribution. \(T(x,t),~S(x,t)\) represent the temperature and heat source distribution at any spatio-temporal position \((x,t)\), respectively, and \(t_*\) is the temperature threshold. Here \(k\), \(\rho\), \(c_p\) are all material properties of the given object, representing thermal conductivity, mass density, and specific heat capacity, respectively. For convenience, we focus on the static temperature field of the given object material and simplify the equation by setting \(\frac{dT}{dt}=0\):

\[ \tag{1} k \Delta T(x) + S(x) = 0,\quad \text { in } \Omega. \]

domain_chip.pdf

2D chip simulation area with random heat source distribution inside, and arbitrary boundary conditions on the boundary.

For a general thermal simulation model of a given object material, in addition to satisfying the governing equation (1), its temperature field also depends on some key PDE configurations, including but not limited to material properties and geometric parameters.

The first type of PDE configuration is the boundary conditions of the given object material:

  • Dirichlet boundary condition: The temperature field on the surface is fixed at \(q_d\):
\[ T = q_d. \]
  • Neumann boundary condition: The temperature flux on the surface is fixed at \(q_n\). When \(q_n =0\), it indicates that the surface is completely insulated, called adiabatic boundary condition.
\[ \tag{2} -k \dfrac{\partial T}{\partial n} = q_n. \]
  • Convection boundary condition: Also known as Newton boundary condition, this boundary condition corresponds to the balance between heat conduction and convection in the same direction on the surface, where \(h\) and \(T_{amb}\) represent the surface convection coefficient and ambient temperature.
\[ -k \dfrac{\partial T}{\partial n} = h(T-T_{amb}). \]
  • Radiation boundary condition: This boundary condition corresponds to electromagnetic radiation generated by temperature difference on the surface, where \(\epsilon\) and \(\sigma\) represent thermal radiation coefficient and Stefan-Boltzmann coefficient, respectively.
\[ -k \dfrac{\partial T}{\partial n} = \epsilon \sigma (T^4-T_{amb}^4). \]

The second type of PDE configuration is the position and intensity of boundary or internal heat sources of the given object material. This work considers the following two types of heat sources:

  • Boundary random heat source: Defined by Neumann boundary condition (2), where \(q_n\) is a function of \(x\), i.e., any given temperature flux distribution;
  • Internal random heat source: Defined by governing equation (1), where \(S(x)\) is a function of \(x\), i.e., any given heat source distribution.

Our goal is to obtain the corresponding temperature field distribution on the general thermal simulation model of a given object material by inputting any first or second type of design configuration, where we arbitrarily specify the boundary type and parameters on the boundary. It is worth noting that the PI-DeepONet method for general thermal simulation developed in this work is not limited to the conditions of the first or second type of design configuration and regular geometric shapes. With further code modifications beyond the scope of current work, they can be applied to various loads, material properties, and even various irregular geometric shapes.

2.2 PI-DeepONet Model

The PI-DeepONet model combines DeepONet and PINN methods, which is a deep neural network model combining physical information and operator learning. This model can enhance the DeepONet model through the physical information of governing equations, and can use different PDE configurations as input data for different branch networks, so it can be effectively used for ultra-fast model prediction under various (parametric and non-parametric) PDE configurations.

For the chip thermal simulation problem, the PI-DeepONet model can be represented as the model structure shown in the figure:

pi_deeponet.pdf

As shown in the figure, we used a total of 3 branch networks and one trunk network. The branch networks input boundary type index, random heat source distribution \(S(x, y)\) and boundary function \(Q(x, y)\) respectively, and the trunk network inputs 2D coordinate point information. Each branch network and trunk network outputs a \(q\)-dimensional feature vector. All these output features are combined through Hadamard (element-wise) product, and then the resulting vectors are summed as the scalar output of the predicted temperature field.

3. Problem Solving

Next, we will explain how to convert this problem into PaddleScience code step by step and solve the heat exchanger thermal simulation problem using deep learning methods. In order to quickly understand PaddleScience, only key steps such as model construction and constraint construction are described below, while other details please refer to API Documentation.

3.1 Model Construction

In the chip thermal simulation problem, each known coordinate point \((x, y)\) and each set of boundary type \(bt\), random heat source distribution \(S(x, y)\) and boundary function \(Q(x, y)\) correspond to a set of chip temperature distribution \(T\), an unknown quantity to be solved. Here we use 3 branch networks and one trunk network, all 4 networks are MLP (Multilayer Perceptron). The 3 branch networks represent the mapping functions \(f_1, f_2, f_3: \mathbb{R}^3 \to \mathbb{R}^{q}\) from \((bt, S, Q)\) to output functions \((b_1, b_2, b_3)\) respectively, i.e.:

\[ \begin{aligned} b_1 &= f_1(bt),\\ b_2 &= f_2(S),\\ b_3 &= f_3(Q). \end{aligned} \]

In the above formula, \(f_1, f_2, f_3\) are all MLP models, \((b_1, b_2, b_3)\) are the output functions of the three branch networks respectively, and \(q\) is the dimension of the output function. The trunk network represents the mapping function \(f_4: \mathbb{R} \to \mathbb{R}^{q}\) from \((x, y)\) to output function \(t_0\), i.e.:

\[ \begin{aligned} t_0 &= f_4(x, y). \end{aligned} \]

In the above formula, \(f_4\) is an MLP model, \((t_0)\) is the output function of the trunk network, and \(q\) is the dimension of the output function. We can perform Hadamard (element-wise) product on the output functions of the three branch networks and the trunk network \((b_1, b_2, b_3, t_0)\) and then sum them up to obtain the scalar temperature field, i.e.:

\[ T = \sum_{i=1}^q b_1^ib_2^ib_3^it_0^i. \]

We define the ChipHeats model class built in PaddleScience and call it. The PaddleScience code is as follows:

77
78
# set model
model = ppsci.arch.ChipDeepONets(**cfg.MODEL)

In this way, we instantiated a ChipHeats model with 4 MLP models. Each branch network contains 9 hidden layers with 256 neurons per layer. The trunk network contains 6 hidden layers with 128 neurons per layer. "Swish" is used as the activation function. The neural network model model contains an output function \(T\). For more relevant content, please refer to A fast general thermal simulation model based on MultiBranch Physics-Informed deep operator neural network.

3.2 Computational Domain Construction

Construct the training area for the chip thermal simulation problem in this article, which is a 2D area of \([0, 1]\times[0, 1]\). This area can directly use the spatial geometry Rectangle built in PaddleScience to construct the computational domain. The code is as follows:

79
80
81
# set geometry
NPOINT = cfg.NL * cfg.NW
geom = {"rect": ppsci.geometry.Rectangle((0, 0), (cfg.DL, cfg.DW))}
Tip

Rectangle and TimeDomain are two Geometry derived classes that can be used independently.

If the input data only comes from a two-dimensional rectangular geometric domain, you can directly use ppsci.geometry.Rectangle(...) to create a spatial geometric domain object;

If the input data only comes from a one-dimensional time domain, you can directly use ppsci.geometry.TimeDomain(...) to construct a time domain object.

3.3 Input Data Construction

Use 2D correlated and scale-invariant Gaussian random fields to generate random heat source distribution \(S(x)\) and boundary function \(Q(x)\). We refer to the Python implementation described in gaussian-random-fields, where correlation is explained by scale-free spectrum, i.e.:

\[ P(k) \sim \dfrac{1}{|k|^{\alpha/2}}. \]

The smoothness of the sampling function is determined by the length scale coefficient \(\alpha\). The larger the \(\alpha\) value, the smoother the random heat source distribution \(S(x)\) and boundary function \(Q(x)\) obtained. In this article we use \(\alpha = 4\). This parameter can also be adjusted to generate heat source distribution \(S(x)\) and boundary function \(Q(x)\) similar to specific optimization tasks.

Generate training and test input data for random heat source distribution \(S(x)\) and boundary function \(Q(x)\) through Gaussian random fields. The code is as follows:

84
85
86
87
88
89
90
91
92
93
94
95
# generate training data and validation data
data_u = np.ones([1, (cfg.NL - 2) * (cfg.NW - 2)])
data_BC = np.ones([1, NPOINT])
data_u = np.vstack((data_u, np.zeros([1, (cfg.NL - 2) * (cfg.NW - 2)])))
data_BC = np.vstack((data_BC, np.zeros([1, NPOINT])))
for i in range(cfg.NU - 2):
    data_u = np.vstack((data_u, GRF(alpha=cfg.GRF.alpha, size=cfg.NL - 2)))
for i in range(cfg.NBC - 2):
    data_BC = np.vstack((data_BC, GRF(alpha=cfg.GRF.alpha, size=cfg.NL)))
data_u = data_u.astype("float32")
data_BC = data_BC.astype("float32")
test_u = GRF(alpha=4, size=cfg.NL).astype("float32")[0]

Then classify the training data and test data according to spatial coordinates, classifying them into left, right, top, bottom and internal data. The code is as follows:

 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
boundary_indices = np.where(
    (
        (points["x"] == 0)
        | (points["x"] == cfg.DW)
        | (points["y"] == 0)
        | (points["y"] == cfg.DL)
    )
)
interior_indices = np.where(
    (
        (points["x"] != 0)
        & (points["x"] != cfg.DW)
        & (points["y"] != 0)
        & (points["y"] != cfg.DL)
    )
)

points["u"] = np.tile(test_u[interior_indices[0]], (NPOINT, 1))
points["u_one"] = test_u.T.reshape([-1, 1])
points["bc_data"] = np.tile(test_u[boundary_indices[0]], (NPOINT, 1))
points["bc"] = np.zeros((NPOINT, 1), dtype="float32")

top_indices = np.where(points["x"] == cfg.DW)
down_indices = np.where(points["x"] == 0)
left_indices = np.where(
    (points["y"] == 0) & (points["x"] != 0) & (points["x"] != cfg.DW)
)
right_indices = np.where(
    ((points["y"] == cfg.DL) & (points["x"] != 0) & (points["x"] != cfg.DW))
)

# generate validation data
(
    test_top_data,
    test_down_data,
    test_left_data,
    test_right_data,
    test_interior_data,
) = [
    {
        "x": points["x"][indices_[0]],
        "y": points["y"][indices_[0]],
        "u": points["u"][indices_[0]],
        "u_one": points["u_one"][indices_[0]],
        "bc": points["bc"][indices_[0]],
        "bc_data": points["bc_data"][indices_[0]],
    }
    for indices_ in (
        top_indices,
        down_indices,
        left_indices,
        right_indices,
        interior_indices,
    )
]
# generate train data
top_data = {
    "x": test_top_data["x"],
    "y": test_top_data["y"],
    "u": data_u,
    "u_one": data_BC[:, top_indices[0]].T.reshape([-1, 1]),
    "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
    "bc_data": data_BC[:, boundary_indices[0]],
}
down_data = {
    "x": test_down_data["x"],
    "y": test_down_data["y"],
    "u": data_u,
    "u_one": data_BC[:, down_indices[0]].T.reshape([-1, 1]),
    "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
    "bc_data": data_BC[:, boundary_indices[0]],
}
left_data = {
    "x": test_left_data["x"],
    "y": test_left_data["y"],
    "u": data_u,
    "u_one": data_BC[:, left_indices[0]].T.reshape([-1, 1]),
    "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
    "bc_data": data_BC[:, boundary_indices[0]],
}
right_data = {
    "x": test_right_data["x"],
    "y": test_right_data["y"],
    "u": data_u,
    "u_one": data_BC[:, right_indices[0]].T.reshape([-1, 1]),
    "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
    "bc_data": data_BC[:, boundary_indices[0]],
}
interior_data = {
    "x": test_interior_data["x"],
    "y": test_interior_data["y"],
    "u": data_u,
    "u_one": data_u.T.reshape([-1, 1]),
    "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
    "bc_data": data_BC[:, boundary_indices[0]],
}

3.4 Constraint Construction

Before constructing constraints, we need to introduce ChipHeatDataset, which inherits from Dataset class and can iteratively read array datasets composed of different numpy.ndarray. Due to the large number of model branch networks used, the amount of data used is large. If the data is combined first, the memory occupied by the input data will be large, so ChipHeatDataset is used to iteratively read data.

The chip thermal simulation problem consists of equations described in 2.1 Problem Description. At this time, we set five constraint conditions for left, right, top, bottom and internal data respectively. Next, use SupervisedConstraint built in PaddleScience to construct the above four constraint conditions. The code is as follows:

194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
# set constraint
index = ("x", "u", "bc", "bc_data")
label = {"chip": np.array([0], dtype="float32")}
weight = {"chip": np.array([cfg.TRAIN.weight], dtype="float32")}
top_sup_constraint = ppsci.constraint.SupervisedConstraint(
    {
        "dataset": {
            "name": "ChipHeatDataset",
            "input": top_data,
            "label": label,
            "index": index,
            "data_type": "bc_data",
            "weight": weight,
        },
        "batch_size": cfg.TRAIN.batch_size,
        "sampler": {
            "name": "BatchSampler",
            "drop_last": False,
            "shuffle": True,
        },
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={
        "chip": lambda out: paddle.where(
            out["bc"] == 1,
            jacobian(out["T"], out["x"]) - out["u_one"],
            paddle.where(
                out["bc"] == 0,
                out["T"] - out["u_one"],
                paddle.where(
                    out["bc"] == 2,
                    jacobian(out["T"], out["x"]) + out["u_one"] * (out["T"] - 1),
                    jacobian(out["T"], out["x"])
                    + out["u_one"]
                    * (out["T"] ** 2 - 1)
                    * (out["T"] ** 2 + 1)
                    * 5.6
                    / 50000,
                ),
            ),
        )
    },
    name="top_sup",
)
down_sup_constraint = ppsci.constraint.SupervisedConstraint(
    {
        "dataset": {
            "name": "ChipHeatDataset",
            "input": down_data,
            "label": label,
            "index": index,
            "data_type": "bc_data",
            "weight": weight,
        },
        "batch_size": cfg.TRAIN.batch_size,
        "sampler": {
            "name": "BatchSampler",
            "drop_last": False,
            "shuffle": True,
        },
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={
        "chip": lambda out: paddle.where(
            out["bc"] == 1,
            jacobian(out["T"], out["x"]) - out["u_one"],
            paddle.where(
                out["bc"] == 0,
                out["T"] - out["u_one"],
                paddle.where(
                    out["bc"] == 2,
                    jacobian(out["T"], out["x"]) + out["u_one"] * (out["T"] - 1),
                    jacobian(out["T"], out["x"])
                    + out["u_one"]
                    * (out["T"] ** 2 - 1)
                    * (out["T"] ** 2 + 1)
                    * 5.6
                    / 50000,
                ),
            ),
        )
    },
    name="down_sup",
)
left_sup_constraint = ppsci.constraint.SupervisedConstraint(
    {
        "dataset": {
            "name": "ChipHeatDataset",
            "input": left_data,
            "label": label,
            "index": index,
            "data_type": "bc_data",
            "weight": weight,
        },
        "batch_size": cfg.TRAIN.batch_size,
        "sampler": {
            "name": "BatchSampler",
            "drop_last": False,
            "shuffle": True,
        },
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={
        "chip": lambda out: paddle.where(
            out["bc"] == 1,
            jacobian(out["T"], out["y"]) - out["u_one"],
            paddle.where(
                out["bc"] == 0,
                out["T"] - out["u_one"],
                paddle.where(
                    out["bc"] == 2,
                    jacobian(out["T"], out["y"]) + out["u_one"] * (out["T"] - 1),
                    jacobian(out["T"], out["y"])
                    + out["u_one"]
                    * (out["T"] ** 2 - 1)
                    * (out["T"] ** 2 + 1)
                    * 5.6
                    / 50000,
                ),
            ),
        )
    },
    name="left_sup",
)
right_sup_constraint = ppsci.constraint.SupervisedConstraint(
    {
        "dataset": {
            "name": "ChipHeatDataset",
            "input": right_data,
            "label": label,
            "index": index,
            "data_type": "bc_data",
            "weight": weight,
        },
        "batch_size": cfg.TRAIN.batch_size,
        "sampler": {
            "name": "BatchSampler",
            "drop_last": False,
            "shuffle": True,
        },
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={
        "chip": lambda out: paddle.where(
            out["bc"] == 1,
            jacobian(out["T"], out["y"]) - out["u_one"],
            paddle.where(
                out["bc"] == 0,
                out["T"] - out["u_one"],
                paddle.where(
                    out["bc"] == 2,
                    jacobian(out["T"], out["y"]) + out["u_one"] * (out["T"] - 1),
                    jacobian(out["T"], out["y"])
                    + out["u_one"]
                    * (out["T"] ** 2 - 1)
                    * (out["T"] ** 2 + 1)
                    * 5.6
                    / 50000,
                ),
            ),
        )
    },
    name="right_sup",
)
interior_sup_constraint = ppsci.constraint.SupervisedConstraint(
    {
        "dataset": {
            "name": "ChipHeatDataset",
            "input": interior_data,
            "label": label,
            "index": index,
            "data_type": "u",
        },
        "batch_size": cfg.TRAIN.batch_size,
        "sampler": {
            "name": "BatchSampler",
            "drop_last": False,
            "shuffle": True,
        },
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={
        "chip": lambda out: hessian(out["T"], out["x"])
        + hessian(out["T"], out["y"])
        + 100 * out["u_one"]
    },
    name="interior_sup",
)

The first parameter of SupervisedConstraint is the reading configuration of supervised constraint, where the "dataset" field represents the training dataset information used, and each field represents:

  1. name: Dataset type, here ChipHeatDataset means iteratively reading data in batches;
  2. input: Input variable name;
  3. label: Label variable name;
  4. index: Index of input dataset;
  5. data_type: Type of input data;
  6. weight: Weight size.

The "sampler" field defines the Sampler class name used as BatchSampler, and also specifies that the parameters drop_last is False and shuffle is True during initialization of this class.

The second parameter is the loss function. Here we choose the commonly used MSE function, and reduction is "mean", that is, we will sum and average the loss terms generated by all data points involved in the calculation;

The third parameter is the label expression list. Here we use equation expressions corresponding to left, right, top, bottom and internal regions. At the same time, we use \(0, 1, 2, 3\) to represent Dirichlet boundary, Neumann boundary, convection boundary and radiation boundary respectively. Different boundary conditions are set for different boundary types;

The fourth parameter is the name of the constraint condition. We need to name each constraint condition for subsequent indexing.

After the differential equation constraint and supervised constraint are constructed, encapsulate them into a dictionary with the names we just named as keys for subsequent access.

382
383
384
385
386
387
388
389
# wrap constraints together
constraint = {
    down_sup_constraint.name: down_sup_constraint,
    left_sup_constraint.name: left_sup_constraint,
    right_sup_constraint.name: right_sup_constraint,
    interior_sup_constraint.name: interior_sup_constraint,
    top_sup_constraint.name: top_sup_constraint,
}

3.5 Optimizer Construction

Next we need to specify the learning rate, which is set to 0.001. The training process will call the optimizer to update model parameters. Here, the more commonly used Adam optimizer is selected.

391
392
clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=1.0)
# set optimizer

3.6 Validator Construction

Usually during the training process, the training status of the current model is evaluated using the validation set (test set) at a certain epoch interval. We use ppsci.validate.SupervisedValidator to construct the validator.

394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
# set validator
top_down_label = {"chip": np.zeros([cfg.NL, 1], dtype="float32")}
left_right_label = {"chip": np.zeros([(cfg.NL - 2), 1], dtype="float32")}
interior_label = {
    "thermal_condution": np.zeros(
        [test_interior_data["x"].shape[0], 1], dtype="float32"
    )
}
top_validator = ppsci.validate.SupervisedValidator(
    {
        "dataset": {
            "name": "NamedArrayDataset",
            "input": test_top_data,
            "label": top_down_label,
            "weight": {
                "chip": np.full([cfg.NL, 1], cfg.TRAIN.weight, dtype="float32")
            },
        },
        "batch_size": cfg.NL,
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={"chip": lambda out: out["T"] - out["u_one"]},
    metric={"MSE": ppsci.metric.MSE()},
    name="top_mse",
)
down_validator = ppsci.validate.SupervisedValidator(
    {
        "dataset": {
            "name": "NamedArrayDataset",
            "input": test_down_data,
            "label": top_down_label,
            "weight": {
                "chip": np.full([cfg.NL, 1], cfg.TRAIN.weight, dtype="float32")
            },
        },
        "batch_size": cfg.NL,
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={"chip": lambda out: out["T"] - out["u_one"]},
    metric={"MSE": ppsci.metric.MSE()},
    name="down_mse",
)
left_validator = ppsci.validate.SupervisedValidator(
    {
        "dataset": {
            "name": "NamedArrayDataset",
            "input": test_left_data,
            "label": left_right_label,
            "weight": {
                "chip": np.full([cfg.NL - 2, 1], cfg.TRAIN.weight, dtype="float32")
            },
        },
        "batch_size": (cfg.NL - 2),
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={"chip": lambda out: out["T"] - out["u_one"]},
    metric={"MSE": ppsci.metric.MSE()},
    name="left_mse",
)
right_validator = ppsci.validate.SupervisedValidator(
    {
        "dataset": {
            "name": "NamedArrayDataset",
            "input": test_right_data,
            "label": left_right_label,
            "weight": {
                "chip": np.full([cfg.NL - 2, 1], cfg.TRAIN.weight, dtype="float32")
            },
        },
        "batch_size": (cfg.NL - 2),
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={"chip": lambda out: out["T"] - out["u_one"]},
    metric={"MSE": ppsci.metric.MSE()},
    name="right_mse",
)
interior_validator = ppsci.validate.SupervisedValidator(
    {
        "dataset": {
            "name": "NamedArrayDataset",
            "input": test_interior_data,
            "label": interior_label,
        },
        "batch_size": cfg.TRAIN.batch_size,
    },
    ppsci.loss.MSELoss("mean"),
    output_expr={
        "thermal_condution": lambda out: (
            hessian(out["T"], out["x"]) + hessian(out["T"], out["y"])
        )
        + 100 * out["u_one"]
    },
    metric={"MSE": ppsci.metric.MSE()},
    name="interior_mse",
)
validator = {
    down_validator.name: down_validator,
    left_validator.name: left_validator,
    right_validator.name: right_validator,
    top_validator.name: top_validator,
    interior_validator.name: interior_validator,

The configuration is similar to the setting of 3.4 Constraint Construction. It should be noted that since the amount of data used for evaluation is not very large, we do not need to use ChipHeatDataset to iteratively read data, but use NamedArrayDataset to read data here.

3.7 Model Training

After completing the above settings, you only need to pass the instantiated objects to ppsci.solver.Solver in order, and then start training and evaluation.

497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
# initialize solver
solver = ppsci.solver.Solver(
    model,
    constraint,
    cfg.output_dir,
    optimizer,
    None,
    cfg.TRAIN.epochs,
    cfg.TRAIN.iters_per_epoch,
    eval_during_train=cfg.TRAIN.eval_during_train,
    eval_freq=cfg.TRAIN.eval_freq,
    validator=validator,
)
# train model
solver.train()
# evaluate after finished training

3.8 Result Visualization

Finally, prediction and visualization are performed on the given visualization area. The visualization data is a two-dimensional point set in the area. At each coordinate \((x, y)\), the corresponding temperature value \(T\) is plotted. Here we plot the image of \(T\) change on the area. At the same time, different boundary types, random heat source distribution \(S(x)\) and boundary function \(Q(x)\) can be set as needed. The code is as follows:

514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
solver.eval()
# visualize prediction after finished training
pred_points = geom["rect"].sample_interior(NPOINT, evenly=True)
pred_points["u"] = points["u"]
pred_points["bc_data"] = np.zeros_like(points["bc_data"])
pred_points["bc"] = np.repeat(
    np.array([[cfg.EVAL.bc_type]], dtype="float32"), NPOINT, axis=0
)
pred = solver.predict(pred_points)
logger.message("Now saving visual result to: visual/result.vtu, please wait...")
ppsci.visualize.save_vtu_from_dict(
    osp.join(cfg.output_dir, "visual/result.vtu"),
    {
        "x": pred_points["x"],
        "y": pred_points["y"],
        "T": pred["T"],
    },
    (
        "x",
        "y",
    ),
    ("T"),

4. Complete Code

chip_heat.py
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
# Copyright (c) 2023 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 os import path as osp

import hydra
import numpy as np
import paddle
import scipy.fftpack
import scipy.io
from omegaconf import DictConfig

import ppsci
from ppsci.autodiff import hessian
from ppsci.autodiff import jacobian
from ppsci.utils import logger


def fftind(size):
    """
    Returns the momentum indices for the 2D Fast Fourier Transform (FFT).

    Args:
        size (int): Size of the 2D array.

    Returns:
        numpy.ndarray: Array of momentum indices for the 2D FFT.
    """
    k_ind = np.mgrid[:size, :size] - int((size + 1) / 2)
    k_ind = scipy.fftpack.fftshift(k_ind)
    return k_ind


def GRF(alpha=3.0, size=128, flag_normalize=True):
    """
    Generates a Gaussian random field(GRF) with a power law amplitude spectrum.

    Args:
        alpha (float, optional): Power law exponent. Defaults to 3.0.
        size (int, optional): Size of the output field. Defaults to 128.
        flag_normalize (bool, optional): Flag indicating whether to normalize the field. Defaults to True.

    Returns:
        numpy.ndarray: Generated Gaussian random field.
    """
    # Defines momentum indices
    k_idx = fftind(size)
    # Defines the amplitude as a power law 1/|k|^(alpha/2)
    amplitude = np.power(k_idx[0] ** 2 + k_idx[1] ** 2 + 1e-10, -alpha / 4.0)
    amplitude[0, 0] = 0
    # Draws a complex gaussian random noise with normal
    # (circular) distribution
    noise = np.random.normal(size=(size, size)) + 1j * np.random.normal(
        size=(size, size)
    )
    # To real space
    gfield = np.fft.ifft2(noise * amplitude).real
    # Sets the standard deviation to one
    if flag_normalize:
        gfield = gfield - np.mean(gfield)
        gfield = gfield / np.std(gfield)
    return gfield.reshape([1, -1])


def train(cfg: DictConfig):
    # set model
    model = ppsci.arch.ChipDeepONets(**cfg.MODEL)
    # set geometry
    NPOINT = cfg.NL * cfg.NW
    geom = {"rect": ppsci.geometry.Rectangle((0, 0), (cfg.DL, cfg.DW))}
    points = geom["rect"].sample_interior(NPOINT, evenly=True)

    # generate training data and validation data
    data_u = np.ones([1, (cfg.NL - 2) * (cfg.NW - 2)])
    data_BC = np.ones([1, NPOINT])
    data_u = np.vstack((data_u, np.zeros([1, (cfg.NL - 2) * (cfg.NW - 2)])))
    data_BC = np.vstack((data_BC, np.zeros([1, NPOINT])))
    for i in range(cfg.NU - 2):
        data_u = np.vstack((data_u, GRF(alpha=cfg.GRF.alpha, size=cfg.NL - 2)))
    for i in range(cfg.NBC - 2):
        data_BC = np.vstack((data_BC, GRF(alpha=cfg.GRF.alpha, size=cfg.NL)))
    data_u = data_u.astype("float32")
    data_BC = data_BC.astype("float32")
    test_u = GRF(alpha=4, size=cfg.NL).astype("float32")[0]

    boundary_indices = np.where(
        (
            (points["x"] == 0)
            | (points["x"] == cfg.DW)
            | (points["y"] == 0)
            | (points["y"] == cfg.DL)
        )
    )
    interior_indices = np.where(
        (
            (points["x"] != 0)
            & (points["x"] != cfg.DW)
            & (points["y"] != 0)
            & (points["y"] != cfg.DL)
        )
    )

    points["u"] = np.tile(test_u[interior_indices[0]], (NPOINT, 1))
    points["u_one"] = test_u.T.reshape([-1, 1])
    points["bc_data"] = np.tile(test_u[boundary_indices[0]], (NPOINT, 1))
    points["bc"] = np.zeros((NPOINT, 1), dtype="float32")

    top_indices = np.where(points["x"] == cfg.DW)
    down_indices = np.where(points["x"] == 0)
    left_indices = np.where(
        (points["y"] == 0) & (points["x"] != 0) & (points["x"] != cfg.DW)
    )
    right_indices = np.where(
        ((points["y"] == cfg.DL) & (points["x"] != 0) & (points["x"] != cfg.DW))
    )

    # generate validation data
    (
        test_top_data,
        test_down_data,
        test_left_data,
        test_right_data,
        test_interior_data,
    ) = [
        {
            "x": points["x"][indices_[0]],
            "y": points["y"][indices_[0]],
            "u": points["u"][indices_[0]],
            "u_one": points["u_one"][indices_[0]],
            "bc": points["bc"][indices_[0]],
            "bc_data": points["bc_data"][indices_[0]],
        }
        for indices_ in (
            top_indices,
            down_indices,
            left_indices,
            right_indices,
            interior_indices,
        )
    ]
    # generate train data
    top_data = {
        "x": test_top_data["x"],
        "y": test_top_data["y"],
        "u": data_u,
        "u_one": data_BC[:, top_indices[0]].T.reshape([-1, 1]),
        "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
        "bc_data": data_BC[:, boundary_indices[0]],
    }
    down_data = {
        "x": test_down_data["x"],
        "y": test_down_data["y"],
        "u": data_u,
        "u_one": data_BC[:, down_indices[0]].T.reshape([-1, 1]),
        "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
        "bc_data": data_BC[:, boundary_indices[0]],
    }
    left_data = {
        "x": test_left_data["x"],
        "y": test_left_data["y"],
        "u": data_u,
        "u_one": data_BC[:, left_indices[0]].T.reshape([-1, 1]),
        "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
        "bc_data": data_BC[:, boundary_indices[0]],
    }
    right_data = {
        "x": test_right_data["x"],
        "y": test_right_data["y"],
        "u": data_u,
        "u_one": data_BC[:, right_indices[0]].T.reshape([-1, 1]),
        "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
        "bc_data": data_BC[:, boundary_indices[0]],
    }
    interior_data = {
        "x": test_interior_data["x"],
        "y": test_interior_data["y"],
        "u": data_u,
        "u_one": data_u.T.reshape([-1, 1]),
        "bc": np.array([[0], [1], [2], [3]], dtype="float32"),
        "bc_data": data_BC[:, boundary_indices[0]],
    }

    # set constraint
    index = ("x", "u", "bc", "bc_data")
    label = {"chip": np.array([0], dtype="float32")}
    weight = {"chip": np.array([cfg.TRAIN.weight], dtype="float32")}
    top_sup_constraint = ppsci.constraint.SupervisedConstraint(
        {
            "dataset": {
                "name": "ChipHeatDataset",
                "input": top_data,
                "label": label,
                "index": index,
                "data_type": "bc_data",
                "weight": weight,
            },
            "batch_size": cfg.TRAIN.batch_size,
            "sampler": {
                "name": "BatchSampler",
                "drop_last": False,
                "shuffle": True,
            },
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={
            "chip": lambda out: paddle.where(
                out["bc"] == 1,
                jacobian(out["T"], out["x"]) - out["u_one"],
                paddle.where(
                    out["bc"] == 0,
                    out["T"] - out["u_one"],
                    paddle.where(
                        out["bc"] == 2,
                        jacobian(out["T"], out["x"]) + out["u_one"] * (out["T"] - 1),
                        jacobian(out["T"], out["x"])
                        + out["u_one"]
                        * (out["T"] ** 2 - 1)
                        * (out["T"] ** 2 + 1)
                        * 5.6
                        / 50000,
                    ),
                ),
            )
        },
        name="top_sup",
    )
    down_sup_constraint = ppsci.constraint.SupervisedConstraint(
        {
            "dataset": {
                "name": "ChipHeatDataset",
                "input": down_data,
                "label": label,
                "index": index,
                "data_type": "bc_data",
                "weight": weight,
            },
            "batch_size": cfg.TRAIN.batch_size,
            "sampler": {
                "name": "BatchSampler",
                "drop_last": False,
                "shuffle": True,
            },
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={
            "chip": lambda out: paddle.where(
                out["bc"] == 1,
                jacobian(out["T"], out["x"]) - out["u_one"],
                paddle.where(
                    out["bc"] == 0,
                    out["T"] - out["u_one"],
                    paddle.where(
                        out["bc"] == 2,
                        jacobian(out["T"], out["x"]) + out["u_one"] * (out["T"] - 1),
                        jacobian(out["T"], out["x"])
                        + out["u_one"]
                        * (out["T"] ** 2 - 1)
                        * (out["T"] ** 2 + 1)
                        * 5.6
                        / 50000,
                    ),
                ),
            )
        },
        name="down_sup",
    )
    left_sup_constraint = ppsci.constraint.SupervisedConstraint(
        {
            "dataset": {
                "name": "ChipHeatDataset",
                "input": left_data,
                "label": label,
                "index": index,
                "data_type": "bc_data",
                "weight": weight,
            },
            "batch_size": cfg.TRAIN.batch_size,
            "sampler": {
                "name": "BatchSampler",
                "drop_last": False,
                "shuffle": True,
            },
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={
            "chip": lambda out: paddle.where(
                out["bc"] == 1,
                jacobian(out["T"], out["y"]) - out["u_one"],
                paddle.where(
                    out["bc"] == 0,
                    out["T"] - out["u_one"],
                    paddle.where(
                        out["bc"] == 2,
                        jacobian(out["T"], out["y"]) + out["u_one"] * (out["T"] - 1),
                        jacobian(out["T"], out["y"])
                        + out["u_one"]
                        * (out["T"] ** 2 - 1)
                        * (out["T"] ** 2 + 1)
                        * 5.6
                        / 50000,
                    ),
                ),
            )
        },
        name="left_sup",
    )
    right_sup_constraint = ppsci.constraint.SupervisedConstraint(
        {
            "dataset": {
                "name": "ChipHeatDataset",
                "input": right_data,
                "label": label,
                "index": index,
                "data_type": "bc_data",
                "weight": weight,
            },
            "batch_size": cfg.TRAIN.batch_size,
            "sampler": {
                "name": "BatchSampler",
                "drop_last": False,
                "shuffle": True,
            },
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={
            "chip": lambda out: paddle.where(
                out["bc"] == 1,
                jacobian(out["T"], out["y"]) - out["u_one"],
                paddle.where(
                    out["bc"] == 0,
                    out["T"] - out["u_one"],
                    paddle.where(
                        out["bc"] == 2,
                        jacobian(out["T"], out["y"]) + out["u_one"] * (out["T"] - 1),
                        jacobian(out["T"], out["y"])
                        + out["u_one"]
                        * (out["T"] ** 2 - 1)
                        * (out["T"] ** 2 + 1)
                        * 5.6
                        / 50000,
                    ),
                ),
            )
        },
        name="right_sup",
    )
    interior_sup_constraint = ppsci.constraint.SupervisedConstraint(
        {
            "dataset": {
                "name": "ChipHeatDataset",
                "input": interior_data,
                "label": label,
                "index": index,
                "data_type": "u",
            },
            "batch_size": cfg.TRAIN.batch_size,
            "sampler": {
                "name": "BatchSampler",
                "drop_last": False,
                "shuffle": True,
            },
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={
            "chip": lambda out: hessian(out["T"], out["x"])
            + hessian(out["T"], out["y"])
            + 100 * out["u_one"]
        },
        name="interior_sup",
    )
    # wrap constraints together
    constraint = {
        down_sup_constraint.name: down_sup_constraint,
        left_sup_constraint.name: left_sup_constraint,
        right_sup_constraint.name: right_sup_constraint,
        interior_sup_constraint.name: interior_sup_constraint,
        top_sup_constraint.name: top_sup_constraint,
    }

    clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=1.0)
    # set optimizer
    optimizer = ppsci.optimizer.Adam(cfg.TRAIN.learning_rate, grad_clip=clip)(model)

    # set validator
    top_down_label = {"chip": np.zeros([cfg.NL, 1], dtype="float32")}
    left_right_label = {"chip": np.zeros([(cfg.NL - 2), 1], dtype="float32")}
    interior_label = {
        "thermal_condution": np.zeros(
            [test_interior_data["x"].shape[0], 1], dtype="float32"
        )
    }
    top_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_top_data,
                "label": top_down_label,
                "weight": {
                    "chip": np.full([cfg.NL, 1], cfg.TRAIN.weight, dtype="float32")
                },
            },
            "batch_size": cfg.NL,
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={"chip": lambda out: out["T"] - out["u_one"]},
        metric={"MSE": ppsci.metric.MSE()},
        name="top_mse",
    )
    down_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_down_data,
                "label": top_down_label,
                "weight": {
                    "chip": np.full([cfg.NL, 1], cfg.TRAIN.weight, dtype="float32")
                },
            },
            "batch_size": cfg.NL,
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={"chip": lambda out: out["T"] - out["u_one"]},
        metric={"MSE": ppsci.metric.MSE()},
        name="down_mse",
    )
    left_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_left_data,
                "label": left_right_label,
                "weight": {
                    "chip": np.full([cfg.NL - 2, 1], cfg.TRAIN.weight, dtype="float32")
                },
            },
            "batch_size": (cfg.NL - 2),
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={"chip": lambda out: out["T"] - out["u_one"]},
        metric={"MSE": ppsci.metric.MSE()},
        name="left_mse",
    )
    right_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_right_data,
                "label": left_right_label,
                "weight": {
                    "chip": np.full([cfg.NL - 2, 1], cfg.TRAIN.weight, dtype="float32")
                },
            },
            "batch_size": (cfg.NL - 2),
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={"chip": lambda out: out["T"] - out["u_one"]},
        metric={"MSE": ppsci.metric.MSE()},
        name="right_mse",
    )
    interior_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_interior_data,
                "label": interior_label,
            },
            "batch_size": cfg.TRAIN.batch_size,
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={
            "thermal_condution": lambda out: (
                hessian(out["T"], out["x"]) + hessian(out["T"], out["y"])
            )
            + 100 * out["u_one"]
        },
        metric={"MSE": ppsci.metric.MSE()},
        name="interior_mse",
    )
    validator = {
        down_validator.name: down_validator,
        left_validator.name: left_validator,
        right_validator.name: right_validator,
        top_validator.name: top_validator,
        interior_validator.name: interior_validator,
    }

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        constraint,
        cfg.output_dir,
        optimizer,
        None,
        cfg.TRAIN.epochs,
        cfg.TRAIN.iters_per_epoch,
        eval_during_train=cfg.TRAIN.eval_during_train,
        eval_freq=cfg.TRAIN.eval_freq,
        validator=validator,
    )
    # train model
    solver.train()
    # evaluate after finished training
    solver.eval()
    # visualize prediction after finished training
    pred_points = geom["rect"].sample_interior(NPOINT, evenly=True)
    pred_points["u"] = points["u"]
    pred_points["bc_data"] = np.zeros_like(points["bc_data"])
    pred_points["bc"] = np.repeat(
        np.array([[cfg.EVAL.bc_type]], dtype="float32"), NPOINT, axis=0
    )
    pred = solver.predict(pred_points)
    logger.message("Now saving visual result to: visual/result.vtu, please wait...")
    ppsci.visualize.save_vtu_from_dict(
        osp.join(cfg.output_dir, "visual/result.vtu"),
        {
            "x": pred_points["x"],
            "y": pred_points["y"],
            "T": pred["T"],
        },
        (
            "x",
            "y",
        ),
        ("T"),
    )


def evaluate(cfg: DictConfig):
    # set model
    model = ppsci.arch.ChipDeepONets(**cfg.MODEL)
    # set geometry
    NPOINT = cfg.NL * cfg.NW
    geom = {"rect": ppsci.geometry.Rectangle((0, 0), (cfg.DL, cfg.DW))}
    points = geom["rect"].sample_interior(NPOINT, evenly=True)

    # generate validation data
    test_u = GRF(alpha=4, size=cfg.NL).astype("float32")[0]

    boundary_indices = np.where(
        (
            (points["x"] == 0)
            | (points["x"] == cfg.DW)
            | (points["y"] == 0)
            | (points["y"] == cfg.DL)
        )
    )
    interior_indices = np.where(
        (
            (points["x"] != 0)
            & (points["x"] != cfg.DW)
            & (points["y"] != 0)
            & (points["y"] != cfg.DL)
        )
    )

    points["u"] = np.tile(test_u[interior_indices[0]], (NPOINT, 1))
    points["u_one"] = test_u.T.reshape([-1, 1])
    points["bc_data"] = np.tile(test_u[boundary_indices[0]], (NPOINT, 1))
    points["bc"] = np.zeros((NPOINT, 1), dtype="float32")

    top_indices = np.where(points["x"] == cfg.DW)
    down_indices = np.where(points["x"] == 0)
    left_indices = np.where(
        (points["y"] == 0) & (points["x"] != 0) & (points["x"] != cfg.DW)
    )
    right_indices = np.where(
        ((points["y"] == cfg.DL) & (points["x"] != 0) & (points["x"] != cfg.DW))
    )

    # generate validation data
    (
        test_top_data,
        test_down_data,
        test_left_data,
        test_right_data,
        test_interior_data,
    ) = [
        {
            "x": points["x"][indices_[0]],
            "y": points["y"][indices_[0]],
            "u": points["u"][indices_[0]],
            "u_one": points["u_one"][indices_[0]],
            "bc": points["bc"][indices_[0]],
            "bc_data": points["bc_data"][indices_[0]],
        }
        for indices_ in (
            top_indices,
            down_indices,
            left_indices,
            right_indices,
            interior_indices,
        )
    ]

    # set validator
    top_down_label = {"chip": np.zeros([cfg.NL, 1], dtype="float32")}
    left_right_label = {"chip": np.zeros([(cfg.NL - 2), 1], dtype="float32")}
    interior_label = {
        "thermal_condution": np.zeros(
            [test_interior_data["x"].shape[0], 1], dtype="float32"
        )
    }
    top_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_top_data,
                "label": top_down_label,
                "weight": {
                    "chip": np.full([cfg.NL, 1], cfg.TRAIN.weight, dtype="float32")
                },
            },
            "batch_size": cfg.NL,
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={"chip": lambda out: out["T"] - out["u_one"]},
        metric={"MSE": ppsci.metric.MSE()},
        name="top_mse",
    )
    down_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_down_data,
                "label": top_down_label,
                "weight": {
                    "chip": np.full([cfg.NL, 1], cfg.TRAIN.weight, dtype="float32")
                },
            },
            "batch_size": cfg.NL,
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={"chip": lambda out: out["T"] - out["u_one"]},
        metric={"MSE": ppsci.metric.MSE()},
        name="down_mse",
    )
    left_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_left_data,
                "label": left_right_label,
                "weight": {
                    "chip": np.full([cfg.NL - 2, 1], cfg.TRAIN.weight, dtype="float32")
                },
            },
            "batch_size": (cfg.NL - 2),
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={"chip": lambda out: out["T"] - out["u_one"]},
        metric={"MSE": ppsci.metric.MSE()},
        name="left_mse",
    )
    right_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_right_data,
                "label": left_right_label,
                "weight": {
                    "chip": np.full([cfg.NL - 2, 1], cfg.TRAIN.weight, dtype="float32")
                },
            },
            "batch_size": (cfg.NL - 2),
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={"chip": lambda out: out["T"] - out["u_one"]},
        metric={"MSE": ppsci.metric.MSE()},
        name="right_mse",
    )
    interior_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": test_interior_data,
                "label": interior_label,
            },
            "batch_size": cfg.TRAIN.batch_size,
        },
        ppsci.loss.MSELoss("mean"),
        output_expr={
            "thermal_condution": lambda out: (
                hessian(out["T"], out["x"]) + hessian(out["T"], out["y"])
            )
            + 100 * out["u_one"]
        },
        metric={"MSE": ppsci.metric.MSE()},
        name="interior_mse",
    )
    validator = {
        down_validator.name: down_validator,
        left_validator.name: left_validator,
        right_validator.name: right_validator,
        top_validator.name: top_validator,
        interior_validator.name: interior_validator,
    }

    # directly evaluate pretrained model(optional)
    solver = ppsci.solver.Solver(
        model,
        output_dir=cfg.output_dir,
        validator=validator,
        pretrained_model_path=cfg.EVAL.pretrained_model_path,
    )
    solver.eval()
    # visualize prediction result
    pred_points = geom["rect"].sample_interior(NPOINT, evenly=True)
    pred_points["u"] = points["u"]
    pred_points["bc_data"] = np.zeros_like(points["bc_data"])
    pred_points["bc"] = np.full((NPOINT, 1), cfg.EVAL.bc_type, dtype="float32")
    pred = solver.predict(pred_points)
    logger.message("Now saving visual result to: visual/result.vtu, please wait...")
    ppsci.visualize.save_vtu_from_dict(
        osp.join(cfg.output_dir, "visual/result.vtu"),
        {
            "x": pred_points["x"],
            "y": pred_points["y"],
            "T": pred["T"],
        },
        (
            "x",
            "y",
        ),
        ("T"),
    )


@hydra.main(version_base=None, config_path="./conf", config_name="chip_heat.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()

5. Result Display

Three sets of random heat source distributions \(S(x)\) are generated by Gaussian random fields, as shown in the first row of the figure. Next, we can set any boundary condition in the first type of PDE. Here we give five types of boundary conditions, as shown in the boundary equation in the first column of governing equations in the figure. During the test, we set \(k = 100,~h = 100,~T_{amb} = 1,~\epsilon\sigma= 5.6 \times 10^{-7}\). Under different random heat source \(S(x)\) distributions and different boundary conditions, the temperature field distribution tested by the PI-DeepONet model is shown in the figure. From the figure, it can be seen that although there are significant differences in random heat source distribution \(S(x)\) and boundary conditions between test samples, the PI-DeepONet model can correctly predict the two-dimensional diffusion property solutions inside and on the boundary controlled by the heat conduction equation.

chip.png

6. References

Reference: A fast general thermal simulation model based on MultiBranch Physics-Informed deep operator neural network

Reference Code: gaussian-random-fields