""" Shared utilities for the Curl 2026 AITF POD-NN scripts. Provides: - Constants (paths, field keys, hyperparameters) - Data loading from sim_*.npz + model_setup.json - POD-NN model build / train / predict - Save / load trained artifacts (Keras h5 + pickled sklearn objects) - Thermal-resistance computation - Common matplotlib rcParams """ import json import pickle import numpy as np import matplotlib.pyplot as plt from pathlib import Path from sklearn.decomposition import PCA from sklearn.preprocessing import StandardScaler import tensorflow as tf from tensorflow import keras plt.rcParams.update({"font.family": "serif", "font.serif": ["Times New Roman"]}) SCRIPT_DIR = Path(__file__).resolve().parent PACKAGE_DIR = SCRIPT_DIR.parent DATASET_DIR = PACKAGE_DIR / "1_Dataset" ANALYZED_DIR = PACKAGE_DIR / "2_AnalyzedData" MODEL_DIR = ANALYZED_DIR / "trained_model" PLOTS_DIR = ANALYZED_DIR / "plots" FIELD_KEY = "yz-mid|temperature" COORD_KEY = "yz-mid|coordinates" TRAIN_SIZE = 100 TEST_SIZE = 400 N_MODES = 4 EPOCHS = 500 RANDOM_SEED = 42 CHIP_AREA = 0.0016 # m^2 # Raw parameter index map (sorted DOE order) PARAM_IDX = { "Q1": 0, # chip1.thermal.heat_flux "Q2": 1, # chip2.thermal.heat_flux "m1": 2, # inlet1.momentum.mass_flow_rate "T_in1": 3, # inlet1.thermal.total_temperature "m2": 4, # inlet2.momentum.mass_flow_rate "T_in2": 5, # inlet2.thermal.total_temperature } def load_data(): """Load DOE parameters, fields, and yz-mid spatial coordinates. Returns ------- parameters : (n_samples, 6) ndarray fields : (n_samples, n_points) ndarray, yz-mid temperature coordinates: (n_points, 3) ndarray param_names: list[str], len 6 """ with open(DATASET_DIR / "model_setup.json", "r") as f: setup = json.load(f) doe = setup["doe_configuration"] param_names, columns = [], [] for bc in sorted(doe.keys()): for pname in sorted(doe[bc].keys()): vals = doe[bc][pname] if isinstance(vals, list): param_names.append(f"{bc}.{pname}") columns.append(vals) parameters = np.array(columns).T sim_files = sorted(DATASET_DIR.glob("sim_*.npz")) assert len(sim_files) == parameters.shape[0] fields = [] coords = None for f in sim_files: d = np.load(f, allow_pickle=True) fields.append(d[FIELD_KEY]) if coords is None: coords = d[COORD_KEY] fields = np.array(fields) return parameters, fields, coords, param_names def build_pod_nn(input_dim, n_modes): """3-layer NN: 64 -> 64 -> 32 -> n_modes.""" model = keras.Sequential([ keras.layers.Input(shape=(input_dim,)), keras.layers.Dense(64, activation="relu", kernel_regularizer=keras.regularizers.l2(0.001)), keras.layers.Dropout(0.1), keras.layers.Dense(64, activation="relu", kernel_regularizer=keras.regularizers.l2(0.001)), keras.layers.Dropout(0.1), keras.layers.Dense(32, activation="relu", kernel_regularizer=keras.regularizers.l2(0.001)), keras.layers.Dense(n_modes), ]) model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.001), loss="mse", metrics=["mae"]) return model def train_pod_nn(X_train, Y_train, n_modes=N_MODES, validation_split=0.2, epochs=EPOCHS, verbose=0): """Fit POD then NN. Returns dict of artifacts + history.""" n_train = int(len(X_train) * (1 - validation_split)) n_modes = min(n_modes, n_train, Y_train.shape[1]) pca = PCA(n_components=n_modes) modes = pca.fit_transform(Y_train) param_scaler = StandardScaler().fit(X_train) mode_scaler = StandardScaler().fit(modes) X_s = param_scaler.transform(X_train) M_s = mode_scaler.transform(modes) nn = build_pod_nn(X_train.shape[1], n_modes) monitor = "loss" if validation_split == 0.0 else "val_loss" callbacks = [ keras.callbacks.EarlyStopping(monitor=monitor, patience=80, restore_best_weights=True, min_delta=1e-7, start_from_epoch=30), keras.callbacks.ReduceLROnPlateau(monitor=monitor, factor=0.5, patience=15, min_lr=1e-7), ] history = nn.fit(X_s, M_s, validation_split=validation_split, epochs=epochs, batch_size=8, callbacks=callbacks, verbose=verbose) return { "pca": pca, "param_scaler": param_scaler, "mode_scaler": mode_scaler, "nn_model": nn, "n_modes": n_modes, "history": history.history, } def predict_field(trained, X): """X (n, 6) -> Y (n, n_points).""" X_s = trained["param_scaler"].transform(X) M_s = trained["nn_model"].predict(X_s, verbose=0) M = trained["mode_scaler"].inverse_transform(M_s) return trained["pca"].inverse_transform(M) def save_artifacts(trained, model_dir=MODEL_DIR): """Persist NN (h5), PCA + scalers (pkl), history (json).""" model_dir = Path(model_dir) model_dir.mkdir(parents=True, exist_ok=True) trained["nn_model"].save(model_dir / "pod_nn.h5") with open(model_dir / "pca.pkl", "wb") as f: pickle.dump(trained["pca"], f) with open(model_dir / "param_scaler.pkl", "wb") as f: pickle.dump(trained["param_scaler"], f) with open(model_dir / "mode_scaler.pkl", "wb") as f: pickle.dump(trained["mode_scaler"], f) with open(model_dir / "training_history.json", "w") as f: json.dump(trained["history"], f) print(f"Saved trained artifacts to {model_dir}") def load_artifacts(model_dir=MODEL_DIR): """Inverse of save_artifacts.""" model_dir = Path(model_dir) nn = keras.models.load_model(model_dir / "pod_nn.h5", compile=False) with open(model_dir / "pca.pkl", "rb") as f: pca = pickle.load(f) with open(model_dir / "param_scaler.pkl", "rb") as f: param_scaler = pickle.load(f) with open(model_dir / "mode_scaler.pkl", "rb") as f: mode_scaler = pickle.load(f) with open(model_dir / "training_history.json", "r") as f: history = json.load(f) return { "pca": pca, "param_scaler": param_scaler, "mode_scaler": mode_scaler, "nn_model": nn, "n_modes": pca.n_components_, "history": history, } def compute_rth(fields, coords, parameters): """Per-chip thermal resistance Rth = (Tmax_chip - T_inlet) / Q. Geometry: chip1 = right (z >= z_mid), chip2 = left (z < z_mid). """ z = coords[:, 2] z_mid = (z.min() + z.max()) / 2 right = z >= z_mid left = z < z_mid Q1 = parameters[:, PARAM_IDX["Q1"]] * CHIP_AREA Q2 = parameters[:, PARAM_IDX["Q2"]] * CHIP_AREA T_in1 = parameters[:, PARAM_IDX["T_in1"]] T_in2 = parameters[:, PARAM_IDX["T_in2"]] n = len(fields) tmax1 = np.array([fields[i, right].max() for i in range(n)]) tmax2 = np.array([fields[i, left].max() for i in range(n)]) rth1 = (tmax1 - T_in1) / Q1 rth2 = (tmax2 - T_in2) / Q2 return rth1, rth2, tmax1, tmax2 def yz_mm(coords): """Centered Y/Z coordinates in mm for plotting.""" y = coords[:, 1] z = coords[:, 2] y = np.round((y - (y.min() + y.max()) / 2) * 1000, 1) z = np.round((z - (z.min() + z.max()) / 2) * 1000, 1) return y, z def seed_all(seed=RANDOM_SEED): np.random.seed(seed) tf.random.set_seed(seed)