JavaProjectRepo/python-ml/type_config.py
2026-01-05 15:18:21 +08:00

91 lines
2.8 KiB
Python

import numpy as np
import pandas as pd
def cfg_cylindrical():
return {
"cols": ["diameter", "height", "fissile_concentration", "isotopic_abundance"],
"label": "keff",
"aliases": {
"直径": "diameter",
"高度": "height",
"铀浓度": "fissile_concentration",
"U浓度": "fissile_concentration",
"铀富集度": "isotopic_abundance",
"U富集度": "isotopic_abundance",
# 兼容旧英文键
"u_concentration": "fissile_concentration",
"u_enrichment": "isotopic_abundance"
}
}
def cfg_ring():
return {
"cols": ["diameter", "height", "fissile_concentration", "isotopic_abundance"],
"label": "keff",
"aliases": {
"外径": "diameter",
"高度": "height",
"Pu浓度": "fissile_concentration",
"Pu240占比": "isotopic_abundance",
# 兼容旧英文键
"pu_concentration": "fissile_concentration",
"pu_isotope": "isotopic_abundance"
}
}
TYPE_CONFIG = {
"CylindricalTank": cfg_cylindrical(),
"AnnularTank": cfg_ring()
}
def feature_names(device_type):
cfg = TYPE_CONFIG[device_type]
base = cfg["cols"]
if device_type in ("CylindricalTank", "AnnularTank"):
derived = [f"{base[0]}*{base[1]}", f"{base[2]}*{base[3]}"]
return base + derived
return base
def derive_features(device_type, X):
if device_type in ("CylindricalTank", "AnnularTank"):
return np.hstack([X, (X[:, 0] * X[:, 1]).reshape(-1, 1), (X[:, 2] * X[:, 3]).reshape(-1, 1)])
return X
def _resolve_col(df, key, aliases):
if key in df.columns:
return df[key].values
for a, k in aliases.items():
if k == key and a in df.columns:
return df[a].values
raise KeyError(key)
def load_from_excel(device_type, path):
cfg = TYPE_CONFIG[device_type]
df = pd.read_excel(path)
cols = cfg["cols"]
aliases = cfg.get("aliases", {})
arrs = []
for k in cols:
v = _resolve_col(df, k, aliases)
arrs.append(v.astype(float))
X = np.vstack(arrs).T
y = df[cfg["label"]].values.astype(float) if cfg["label"] in df.columns else df[cfg["label"]].values.astype(float)
Xd = derive_features(device_type, X)
return Xd, y
def features_to_vector(device_type, features):
cfg = TYPE_CONFIG[device_type]
aliases = cfg.get("aliases", {})
vals = []
for k in cfg["cols"]:
v = features.get(k)
if v is None:
for a, ck in aliases.items():
if ck == k and a in features:
v = features[a]
break
vals.append(float(v))
X = np.array(vals, dtype=float).reshape(1, -1)
Xd = derive_features(device_type, X)
return Xd