Composable neural emulators accelerate thermoelectric generator design.

1. #!/usr/bin/env python
2. # coding: utf-8
4. import numpy as np
5. import pandas as pd
6. import os
7. import torch
8. import math
9. import torch.nn as nn
10. import torch.optim as optim
11. import matplotlib.pyplot as plt
12. from scipy.interpolate import interp1d
13. from scipy.interpolate import CubicSpline
14. from scipy.integrate import quad
15. from torch.optim.lr_scheduler import StepLR, CosineAnnealingLR, ReduceLROnPlateau, MultiStepLR, ExponentialLR
16. import seaborn as sns
17. import pandas as pd
18. from sklearn.model_selection import train_test_split
19. from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
21. # ==================== Data load and process ====================
22. def load_dataset_from_csv(csv_file):
23. if not os.path.exists(csv_file):
24. raise FileNotFoundError(f"CSV file {csv_file} not existed！")
26. X_columns = ['a', 'b', 'c', 'Tc', 'Th', 'Ih']
27. Y_columns = ['V', 'Q']
29. df = pd.read_csv(csv_file)
30. X_cached = df[X_columns].copy().values
31. Y_cached = df[Y_columns].copy().values
33. X_cached[:, 5] = np.log1p(X_cached[:, 5])
34. Y_offset = np.zeros((1, Y_cached.shape[1]))
35. for i in range(Y_cached.shape[1]):
36. min_val = Y_cached[:, i].min()
37. if min_val <= 0:
38. Y_offset[0, i] = -min_val + 1e-3
39. Y_cached[:, i] += Y_offset[0, i]
40. Y_cached[:, i] = np.log1p(Y_cached[:, i])
42. X_mean = X_cached.mean(axis=0, keepdims=True)
43. X_std = X_cached.std(axis=0, keepdims=True)
44. Y_mean = Y_cached.mean(axis=0, keepdims=True)
45. Y_std = Y_cached.std(axis=0, keepdims=True)
47. X_norm = (X_cached - X_mean) / X_std
48. Y_norm = (Y_cached - Y_mean) / Y_std
50. X_tensor = torch.tensor(X_norm, dtype=torch.float32)
51. Y_tensor = torch.tensor(Y_norm, dtype=torch.float32)
52. X_cached_tensor = torch.tensor(X_cached, dtype=torch.float32)
54. stats = {
55. 'X_mean': torch.tensor(X_mean, dtype=torch.float32),
56. 'X_std': torch.tensor(X_std, dtype=torch.float32),
57. 'Y_mean': torch.tensor(Y_mean, dtype=torch.float32),
58. 'Y_std': torch.tensor(Y_std, dtype=torch.float32),
59. 'Y_offset': torch.tensor(Y_offset, dtype=torch.float32)
60. }
62. print(f"Load success：{X_tensor.shape[0]} sample size in total")
63. return X_tensor, Y_tensor, stats, X_cached_tensor
65. # ==================== Loss function ====================
66. def loss_fn(model, X_norm, Y_norm, stats, X_raw):
67. Y_pred = model(X_norm)
69. v_true, q_true = Y_norm[:, 0], Y_norm[:, 1]
70. v_pred, q_pred = Y_pred[:, 0], Y_pred[:, 1]
72. v_loss = ((v_pred - v_true) ** 2).mean()
73. q_loss = ((q_pred - q_true) ** 2).mean()
75. loss = v_loss + 5*q_loss
77. return loss, v_loss.item(), q_loss.item()
79. # ==================== TEGNet ====================
80. class TEGNet(nn.Module):
81. def __init__(self):
82. super().__init__()
83. self.net = nn.Sequential(
84. nn.Linear(6, 128),
85. nn.SiLU(),
86. nn.Linear(128, 128),
87. nn.SiLU(),
88. nn.Linear(128, 128),
89. nn.SiLU(),
90. nn.Linear(128, 2)
91. )
92. def forward(self, x):
93. return self.net(x)
95. # ==================== Model training ====================
96. def train_model(model, X_train, Y_train, stats, X_raw, epochs=1000, lr=1e-3):
97. optimizer = optim.Adam(model.parameters(), lr=lr)
99. scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-4)
101. train_loss_history = []
103. for epoch in range(epochs):
104. model.train()
105. optimizer.zero_grad()
106. loss, V_loss, Q_loss = loss_fn(model, X_train, Y_train, stats, X_raw)
107. loss.backward()
108. optimizer.step()
109. train_loss_history.append(loss.item())
111. scheduler.step()
113. if epoch % 100 == 0:
114. print(f"Epoch {epoch:4d} | Train: {loss.item():.4e} ")
116. return train_loss_history, V_loss, Q_loss
118. # ==================== Model evaluation ====================
119. def evaluate_model(model, X_test, Y_test, stats, X_raw):
120. model.eval()
121. with torch.no_grad():
123. test_loss, V_loss, Q_loss = loss_fn(model, X_test, Y_test, stats, X_raw)
126. Y_pred_norm = model(X_test)
127. Y_pred = Y_pred_norm * stats['Y_std'] + stats['Y_mean']
128. Y_true = Y_test * stats['Y_std'] + stats['Y_mean']
131. Y_pred = torch.expm1(Y_pred) - stats['Y_offset']
132. Y_true = torch.expm1(Y_true) - stats['Y_offset']
134. Y_pred_np = Y_pred.numpy()
135. Y_true_np = Y_true.numpy()
137. metrics = {}
138. for i, name in enumerate(['Voltage', 'HeatFlux']):
139. y_pred_i = Y_pred_np[:, i]
140. y_true_i = Y_true_np[:, i]
142. metrics[name] = {
143. 'MSE': mean_squared_error(y_true_i, y_pred_i),
144. 'RMSE': mean_squared_error(y_true_i, y_pred_i) ** 0.5,
145. 'MAE': mean_absolute_error(y_true_i, y_pred_i),
146. 'R2': r2_score(y_true_i, y_pred_i)
147. }
149. print("\n" + "=" * 50)
150. print(f" Loss: {test_loss.item():.3e})")
151. for k in metrics:
152. print(f"- {k:10s}: MSE = {metrics[k]['MSE']:.3e} | RMSE = {metrics[k]['RMSE']:.3e} | MAE = {metrics[k]['MAE']:.3e} | R² = {metrics[k]['R2']:.4f}")
153. print("=" * 50)
155. return test_loss.item(), V_loss, Q_loss, metrics
157. # ==================== Save model ====================
158. def save_model(model, path):
159. torch.save(model.state_dict(), path)
160. print(f"Model is saved to {path}")
162. # ==================== Main ====================
163. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
165. parent_dir = os.path.abspath(os.path.join(os.getcwd(), ".."))
167. csv_file = os.path.join(parent_dir, "data", "Ag2TeS.csv")
168. if not os.path.exists(csv_file):
169. raise FileNotFoundError(f"CSV file {csv_file} does not exist!")
171. print("\n" + "=" * 50)
172. X_tensor, Y_tensor, stats, X_raw = load_dataset_from_csv(csv_file)
174. model = TEGNet().to(device)
175. print("\n" + "=" * 50)
176. print("Starting traning...")
178. # 训练模型
179. train_loss, V_train_loss, Q_train_loss = train_model(model, X_tensor, Y_tensor, stats, X_raw, epochs=6000, lr=1e-3)
181. print("\n" + "=" * 50)
182. model_dir = os.path.join(parent_dir, "model")
183. os.makedirs(model_dir, exist_ok=True)
185. model_path = os.path.join(model_dir, "Ag2TeS.pth")
186. save_model(model, path=model_path)

Ag2TeS_model training.py at commit a5c868d, under Apache-2.0 · at the source