| """Shared feature + target construction for EZFlow v2 (OpenFOAM-RANS surrogate). |
| |
| THIS IS THE SINGLE SOURCE OF TRUTH for node features and target transforms. |
| Both the training ETL (etl.py) and the serving path (infer_v5.py) import from |
| here, so a feature added in one place can never silently diverge from the other |
| (that duplication bug existed between the v1 data.py and inference.py). |
| |
| Convention (non-dimensional, matches the CFD): U_inf=1, rho=1, L_ref=1, nu=1/Re. |
| |
| Node features (NODE_FEATURE_DIM = 12), all per-node: |
| [0:3] inflow_dir = [cos(yaw), sin(yaw), 0] |
| [3] d_wall = nearest-surface distance / L_ref(=1) (TRUE scale, |
| NOT per-case-max normalized -- absolute BL scale) |
| [4] log1p(d_wall) = near-wall resolution where BL gradients are steep |
| [5] d_plus_proxy = log1p(clip(d_wall*sqrt(Re), 0, 3000)) Re-aware y+ surrogate |
| [6:9] nearest_wall_normal = outward unit normal of closest wall face (0 far away) |
| [9] normal_dot_inflow = stagnation(+1)/leeward(-1) indicator |
| [10] is_wall = 1 if node lies on the body surface |
| [11] local_mesh_scale = mean incident edge length (resolution awareness) |
| |
| Global features (GLOBAL_DIM = 2): [log10(Re), yaw_rad] |
| |
| Targets (TARGET_DIM = 7): [u, v, w, p, log1p(k), log(omega), log(nut*Re)] |
| velocity/pressure are z-scored downstream; the 3 turbulence channels are put in |
| log space here (heavy-tailed, span decades) and nut is made ~Re-invariant. |
| """ |
| from __future__ import annotations |
| import numpy as np |
| from scipy.spatial import cKDTree |
|
|
| NODE_FEATURE_DIM = 12 |
| GLOBAL_DIM = 2 |
| TARGET_DIM = 7 |
| TARGET_NAMES = ["u", "v", "w", "p", "log1p_k", "log_omega", "log_nutRe"] |
| L_REF = 1.0 |
| _EPS = 1e-12 |
|
|
|
|
| def knn_edges(pos: np.ndarray, k: int = 8) -> np.ndarray: |
| """Bidirectional kNN edge_index (2, E) via cKDTree (no torch_cluster dep).""" |
| tree = cKDTree(pos) |
| kq = min(k + 1, len(pos)) |
| _, nbr = tree.query(pos, k=kq) |
| nbr = np.atleast_2d(nbr) |
| src = np.repeat(np.arange(len(pos)), nbr.shape[1] - 1) |
| dst = nbr[:, 1:].reshape(-1) |
| return np.stack([np.concatenate([src, dst]), np.concatenate([dst, src])]).astype(np.int64) |
|
|
|
|
| def local_mesh_scale(pos: np.ndarray, edge_index: np.ndarray) -> np.ndarray: |
| """Mean incident edge length per node (resolution-awareness feature).""" |
| src, dst = edge_index |
| d = np.linalg.norm(pos[dst] - pos[src], axis=1) |
| n = pos.shape[0] |
| s = np.zeros(n); c = np.zeros(n) |
| np.add.at(s, dst, d); np.add.at(c, dst, 1.0) |
| return s / np.maximum(c, 1.0) |
|
|
|
|
| def build_node_features(pos: np.ndarray, wall_pts: np.ndarray, wall_normals: np.ndarray, |
| yaw_rad: float, Re: float, edge_index: np.ndarray, |
| is_wall: np.ndarray | None = None) -> np.ndarray: |
| """Construct the (N, NODE_FEATURE_DIM) node feature matrix. Geometry is encoded |
| via the nearest wall point/normal (scale-consistent, L_ref=1).""" |
| n = pos.shape[0] |
| dx, dy = float(np.cos(yaw_rad)), float(np.sin(yaw_rad)) |
| inflow = np.tile([dx, dy, 0.0], (n, 1)).astype(np.float32) |
| if wall_pts is not None and len(wall_pts) > 0: |
| tree = cKDTree(wall_pts) |
| d, idx = tree.query(pos, k=1) |
| d_wall = (d / L_REF).astype(np.float32) |
| nrm = wall_normals[idx].astype(np.float32) if wall_normals is not None \ |
| else np.zeros((n, 3), np.float32) |
| else: |
| d_wall = np.zeros(n, np.float32) |
| nrm = np.zeros((n, 3), np.float32) |
| log_d = np.log1p(d_wall) |
| d_plus = np.log1p(np.clip(d_wall * np.sqrt(max(Re, 1.0)), 0.0, 3000.0)).astype(np.float32) |
| ndot = (nrm @ np.array([dx, dy, 0.0], np.float32)).astype(np.float32) |
| if is_wall is None: |
| is_wall = (d_wall < 1e-4).astype(np.float32) |
| scale = local_mesh_scale(pos, edge_index).astype(np.float32) |
| x = np.concatenate([ |
| inflow, d_wall[:, None], log_d[:, None], d_plus[:, None], |
| nrm, ndot[:, None], is_wall.astype(np.float32)[:, None], scale[:, None], |
| ], axis=1).astype(np.float32) |
| return x |
|
|
|
|
| def global_features(Re: float, yaw_rad: float) -> np.ndarray: |
| return np.array([[np.log10(max(Re, 1.0)), yaw_rad]], dtype=np.float32) |
|
|
|
|
| def transform_targets(U: np.ndarray, p: np.ndarray, k: np.ndarray, |
| omega: np.ndarray, nut: np.ndarray, Re: float) -> np.ndarray: |
| """Raw CFD fields -> (N, 7) target array (turbulence channels in log space).""" |
| k = np.clip(k, 0.0, None) |
| omega = np.clip(omega, _EPS, None) |
| nut = np.clip(nut, 0.0, None) |
| return np.concatenate([ |
| U.reshape(-1, 3), |
| p.reshape(-1, 1), |
| np.log1p(k).reshape(-1, 1), |
| np.log(omega).reshape(-1, 1), |
| np.log(nut * max(Re, 1.0) + _EPS).reshape(-1, 1), |
| ], axis=1).astype(np.float32) |
|
|
|
|
| def inverse_targets(y: np.ndarray, Re: float) -> dict: |
| """(N, 7) prediction -> physical fields dict {U, p, k, omega, nut}.""" |
| y = np.asarray(y) |
| return { |
| "U": y[:, 0:3], |
| "p": y[:, 3], |
| "k": np.expm1(y[:, 4]).clip(0.0), |
| "omega": np.exp(y[:, 5]), |
| "nut": (np.exp(y[:, 6]) / max(Re, 1.0)).clip(0.0), |
| } |
|
|
