File size: 2,383 Bytes
bcf89e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import os, json, numpy as np, torch, torch.nn as nn
OBJECTIVES = ["min_cd", "max_cl", "max_ld"]
class MLPSelector(nn.Module):
def __init__(self, in_dim:int, n_airfoils:int, obj_dim:int=3, af_embed_dim:int=8, hidden:int=128):
super().__init__()
self.af_emb = nn.Embedding(n_airfoils, af_embed_dim)
self.net = nn.Sequential(
nn.Linear(in_dim + obj_dim + af_embed_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, 1),
)
def forward(self, x, obj_id, af_id):
B = x.size(0)
obj_oh = torch.zeros(B, 3, device=x.device)
obj_oh[torch.arange(B), obj_id] = 1.0
af_e = self.af_emb(af_id)
z = torch.cat([x, obj_oh, af_e], dim=1)
return self.net(z).squeeze(1)
def load_selector(local_dir=".", device="cpu"):
ckpt_path = os.path.join(local_dir, "best.pt")
if not os.path.exists(ckpt_path):
ckpt_path = os.path.join(local_dir, "last.pt")
if not os.path.exists(ckpt_path):
raise FileNotFoundError("best.pt/last.pt not found in "+local_dir)
ckpt = torch.load(ckpt_path, map_location=device)
cfg = {
"in_dim": int(ckpt["in_dim"]),
"n_airfoils": int(ckpt["n_airfoils"]),
"feat_stats": {
"means": np.array(ckpt["feat_stats"]["means"], dtype=np.float32),
"stds": np.array(ckpt["feat_stats"]["stds"], dtype=np.float32),
}
}
model = MLPSelector(cfg["in_dim"], cfg["n_airfoils"])
model.load_state_dict(ckpt["model"])
model.to(device).eval()
return model, cfg
def standardize(X_raw: np.ndarray, means: np.ndarray, stds: np.ndarray) -> np.ndarray:
X_imp = np.where(np.isfinite(X_raw), X_raw, means)
return (X_imp - means) / np.where(stds==0, 1.0, stds)
def score_wings(model, X_std: np.ndarray, airfoil_id: int, objective: str, device="cpu"):
obj_id = OBJECTIVES.index(objective)
X = torch.tensor(X_std, dtype=torch.float32, device=device)
obj_ids = torch.full((X.size(0),), obj_id, dtype=torch.long, device=device)
af_ids = torch.full((X.size(0),), airfoil_id, dtype=torch.long, device=device)
with torch.no_grad():
probs = torch.sigmoid(model(X, obj_ids, af_ids)).cpu().numpy()
return probs # higher = better
|