83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
from joblib import load as joblib_load
|
||
|
|
|
||
|
|
# ---- load NN model once (at import time, not every iteration) ----
|
||
|
|
class MyNet(nn.Module):
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.input = nn.Linear(2, 10)
|
||
|
|
self.hidden1 = nn.Linear(10, 10)
|
||
|
|
self.hidden2 = nn.Linear(10, 1)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
x = nn.functional.relu(self.input(x))
|
||
|
|
x = nn.functional.relu(self.hidden1(x))
|
||
|
|
return self.hidden2(x)
|
||
|
|
|
||
|
|
NN_bool = True # set False to use standard AKN analytic f2
|
||
|
|
|
||
|
|
if NN_bool:
|
||
|
|
_NN_f2 = torch.load('model-neural-k-omega-f_2.pth', weights_only=False)
|
||
|
|
_scaler_yp = joblib_load('scaler-yplus-k-omega-f_2.bin')
|
||
|
|
_scaler_ys = joblib_load('scaler-ystar-k-omega-f_2.bin')
|
||
|
|
_yplus_min, _yplus_max, _ystar_min, _ystar_max, _f2_min, _f2_max = \
|
||
|
|
np.loadtxt('min-max-model-f_2.txt')
|
||
|
|
|
||
|
|
|
||
|
|
def calceps(su2d, sp2d, eps2d, gen):
|
||
|
|
if iter == 0:
|
||
|
|
print('calceps (NN version) called')
|
||
|
|
|
||
|
|
ueps = (eps2d * viscos) ** 0.25
|
||
|
|
ystar = ueps * dist / viscos
|
||
|
|
rt = k2d ** 2 / eps2d / viscos
|
||
|
|
|
||
|
|
# ---- compute f2 ------------------------------------------------
|
||
|
|
if NN_bool:
|
||
|
|
# friction velocity from south wall (same as rans-k-eps-NN.py)
|
||
|
|
ustar = (viscos * u2d[:, 0] / yp2d[:, 0]) ** 0.5 # shape (ni,)
|
||
|
|
|
||
|
|
# yplus and ystar for every cell, clipped to training range
|
||
|
|
yplus_2d = np.minimum(yp2d, yp2d[:, -1:] - yp2d) * ustar[:, None] / viscos
|
||
|
|
ystar_2d = ystar.copy()
|
||
|
|
|
||
|
|
yplus_2d = np.clip(yplus_2d, _yplus_min, _yplus_max)
|
||
|
|
ystar_2d = np.clip(ystar_2d, _ystar_min, _ystar_max)
|
||
|
|
|
||
|
|
# scale and build input matrix
|
||
|
|
X = np.zeros((ni * nj, 2))
|
||
|
|
X[:, 0] = _scaler_yp.transform(yplus_2d.reshape(-1, 1))[:, 0]
|
||
|
|
X[:, 1] = _scaler_ys.transform(ystar_2d.reshape(-1, 1))[:, 0]
|
||
|
|
|
||
|
|
with torch.no_grad():
|
||
|
|
preds = _NN_f2(torch.tensor(X, dtype=torch.float32))
|
||
|
|
|
||
|
|
f2 = preds.numpy().reshape(ni, nj)
|
||
|
|
f2 = np.clip(f2, _f2_min, _f2_max)
|
||
|
|
else:
|
||
|
|
f2 = ((1 - np.exp(-ystar / 3.1)) ** 2) * \
|
||
|
|
(1. - 0.3 * np.exp(-(rt / 6.5) ** 2))
|
||
|
|
|
||
|
|
# ---- fmu (always analytic) -------------------------------------
|
||
|
|
fmu2d = ((1 - np.exp(-ystar / 14)) ** 2) * \
|
||
|
|
(1 + 5 / rt ** 0.75 * np.exp(-(rt / 200) ** 2))
|
||
|
|
fmu2d = np.minimum(fmu2d, 1)
|
||
|
|
|
||
|
|
# ---- production term -------------------------------------------
|
||
|
|
vist = vis2d - viscos
|
||
|
|
su2d = su2d + c_eps_1 * cmu * fmu2d * gen * k2d * vol
|
||
|
|
|
||
|
|
# ---- dissipation term ------------------------------------------
|
||
|
|
sp2d = sp2d - c_eps_2 * f2 * eps2d * vol / k2d
|
||
|
|
|
||
|
|
# ---- modify su & sp (case-specific hooks) ----------------------
|
||
|
|
su2d, sp2d = modify_eps(su2d, sp2d)
|
||
|
|
|
||
|
|
ap2d = aw2d + ae2d + as2d + an2d - sp2d
|
||
|
|
|
||
|
|
# ---- under-relaxation ------------------------------------------
|
||
|
|
ap2d = ap2d / urf_eps
|
||
|
|
su2d = su2d + (1 - urf_eps) * ap2d * eps2d
|
||
|
|
|
||
|
|
return su2d, sp2d, ap2d, fmu2d
|