class AeroTransformerSurrogate(nn.Module):
"""
FlowDesign surrogate wrapper for AeroTransformer ATsurf_S.
Public forward input:
DFlow design variables, usually SuperWing geometry parameters [B, 38].
Public forward output:
Tensor [B, 3] in order [CL, CD, CM]. CL/CD are required by DFlow.
Internal workflow:
- Reconstruct SuperWing reference meshes from the first 38 geometry
columns of `configs.dat`.
- Run the real ATsurf_S checkpoint on `geom0` with condition
`[AoA, Mach]`.
- Integrate predicted `[Cp, Cf_tau, Cf_z]` over `origingeom` to return
`[CL, CD, CMz]`.
"""
def __init__(
self,
config_path: str | None = None,
checkpoint_path: str | None = None,
model_class: str | None = None,
design_dim: int = 38,
field_size: tuple[int, int] = (128, 256),
mach: float = 0.85,
alpha: float = 2.0,
ref_area: float | None = None,
cd_bias: float = 0.02,
strict_load: bool = False,
use_superwing_builder: bool = True,
allow_proxy_coefficients: bool = False,
require_upstream_model: bool = True,
geometry_clamp: bool = True,
device: str | torch.device | None = None,
):
super().__init__()
self.design_dim = int(design_dim)
self.field_size = tuple(field_size)
self.mach = float(mach)
self.alpha = float(alpha)
self.ref_area_override = None if ref_area is None or float(ref_area) <= 0.0 else float(ref_area)
self.cd_bias = float(cd_bias)
self.use_superwing_builder = bool(use_superwing_builder)
self.allow_proxy_coefficients = bool(allow_proxy_coefficients)
self.device_hint = torch.device(device) if device is not None else None
self.geometry_builder = (
TorchSuperWingGeometryBuilder(clamp_inputs=geometry_clamp)
if self.use_superwing_builder
else None
)
self.config = self._load_config(config_path)
init_kwargs = dict(self.config.get("init_kwargs", {}))
init_kwargs.setdefault("fun_dim", 3)
init_kwargs.setdefault("out_dim", 3)
init_kwargs.setdefault("n_hidden", 16)
init_kwargs.setdefault("n_layers", 5)
if model_class is None:
try:
import flowvae.app.wing.models # noqa: F401
model_cls = UpstreamAeroTransformerAdapter
self.uses_upstream_model = True
except Exception as exc:
if require_upstream_model:
raise ImportError(
"The real ATsurf_S path requires the author's flowvae package. "
"Install it with: python -m pip install git+https://github.com/YangYunjia/floGen.git"
) from exc
model_cls = WingPDETransformer
self.uses_upstream_model = False
self.uses_fallback_model = not self.uses_upstream_model
else:
model_cls = _load_class(model_class)
self.uses_upstream_model = False
self.uses_fallback_model = False
self.model = model_cls(*self.config.get("init_args", []), **init_kwargs)
if checkpoint_path:
self.load_checkpoint(checkpoint_path, strict=strict_load)
@staticmethod
def _load_config(config_path: str | None) -> dict:
if config_path is None:
return {
"model_name": "WingPDETransformer",
"schema_version": 1,
"init_args": [],
"init_kwargs": {
"fun_dim": 3,
"mlp_ratio": 4,
"n_hidden": 16,
"n_layers": 5,
"out_dim": 3,
"patch_size": [4, 4],
"type_cond": "inj",
"window_size": [8, 8],
},
}
with open(config_path, "r", encoding="utf-8") as handle:
return json.load(handle)
def load_checkpoint(self, checkpoint_path: str, strict: bool = False) -> None:
state = torch.load(checkpoint_path, map_location="cpu")
if isinstance(state, dict):
state = state.get("state_dict", state.get("model_state_dict", state))
try:
result = self.model.load_state_dict(state, strict=strict)
except RuntimeError:
if strict:
raise
print(
"AeroTransformerSurrogate: checkpoint could not be loaded into "
"the local fallback model. Provide --model-class with the real "
"WingPDETransformer implementation for faithful ATsurf_S weights."
)
return
if not strict and (result.missing_keys or result.unexpected_keys):
print(
"AeroTransformerSurrogate: checkpoint did not exactly match the "
f"local model class. missing={len(result.missing_keys)} "
f"unexpected={len(result.unexpected_keys)}"
)
def _normalize_design(self, x: torch.Tensor) -> torch.Tensor:
if x.dim() == 3 and x.shape[1] == 1:
x = x.squeeze(1)
elif x.dim() == 1:
x = x.reshape(1, -1)
if x.dim() != 2:
raise ValueError(f"Expected design variables [B,D], got {tuple(x.shape)}")
if x.shape[1] < self.design_dim:
raise ValueError(
f"Expected at least {self.design_dim} design variables, got {x.shape[1]}"
)
return x[:, : self.design_dim]
def make_condition(self, batch_size: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
"""AeroTransformer/floGen uses condition order `[AoA, Mach]`."""
return torch.tensor([self.alpha, self.mach], dtype=dtype, device=device).reshape(1, 2).expand(batch_size, 2)
def _proxy_preprocess(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
x = self._normalize_design(x).float()
height, width = self.field_size
bsz = x.shape[0]
# Differentiable proxy geometry channels. They preserve gradients from
# the field/coefficient losses back to the generated design variables.
span = torch.linspace(0.0, 1.0, height, dtype=x.dtype, device=x.device)
chord = torch.linspace(0.0, 1.0, width, dtype=x.dtype, device=x.device)
yy, xx = torch.meshgrid(span, chord, indexing="ij")
xx = xx.expand(bsz, 1, height, width)
yy = yy.expand(bsz, 1, height, width)
planform = x[:, 0:1].reshape(bsz, 1, 1, 1)
twist = x[:, min(13, self.design_dim - 1) : min(14, self.design_dim)].reshape(
bsz, 1, 1, 1
)
thickness = x[:, min(17, self.design_dim - 1) : min(18, self.design_dim)].reshape(
bsz, 1, 1, 1
)
zz = 0.01 * planform * yy + 0.02 * twist * (xx - 0.5) + 0.01 * thickness
field_input = torch.cat([xx, yy, zz.expand_as(xx)], dim=1)
condition = self.make_condition(bsz, x.dtype, x.device)
return field_input, condition
def preprocess(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Convert DFlow design variables to ATsurf_S inputs."""
x = self._normalize_design(x).float()
if self.geometry_builder is None:
return self._proxy_preprocess(x)
geometry = self.geometry_builder(x)
condition = self.make_condition(x.shape[0], geometry.geom0.dtype, geometry.geom0.device)
return geometry.geom0, condition
def build_geometry(self, x: torch.Tensor) -> SuperWingGeometry:
if self.geometry_builder is None:
raise RuntimeError("AeroTransformerSurrogate was created without the SuperWing geometry builder.")
return self.geometry_builder(self._normalize_design(x).float())
def predict_fields(self, x: torch.Tensor) -> torch.Tensor:
field_input, condition = self.preprocess(x)
return self.model(field_input, condition)
def fields_to_coefficients(
self,
fields: torch.Tensor,
geometry: SuperWingGeometry | None = None,
) -> torch.Tensor:
"""
Convert predicted [Cp, Cf_tau, Cf_z] fields to [CL, CD, CM].
"""
if fields.dim() != 4 or fields.shape[1] < 3:
raise ValueError(f"Expected fields [B,3,H,W], got {tuple(fields.shape)}")
if geometry is not None:
batch = fields.shape[0]
ref_area = geometry.ref_area.to(dtype=fields.dtype, device=fields.device)
if self.ref_area_override is not None:
ref_area = torch.full((batch,), self.ref_area_override, dtype=fields.dtype, device=fields.device)
aoa = torch.full((batch,), self.alpha, dtype=fields.dtype, device=fields.device)
return _integrate_surface_fields(
origingeom=geometry.origingeom.to(dtype=fields.dtype, device=fields.device),
fields=fields,
aoa=aoa,
ref_area=ref_area,
ref_chord=geometry.ref_chord.to(dtype=fields.dtype, device=fields.device),
xref=geometry.xref.to(dtype=fields.dtype, device=fields.device),
yref=geometry.yref.to(dtype=fields.dtype, device=fields.device),
)
if not self.allow_proxy_coefficients:
raise ValueError(
"fields_to_coefficients() requires SuperWingGeometry for real surface integration. "
"Pass allow_proxy_coefficients=True only for debugging."
)
cp = fields[:, 0]
cf_tau = fields[:, 1]
cf_z = fields[:, 2]
alpha = torch.as_tensor(self.alpha, dtype=fields.dtype, device=fields.device)
alpha_rad = torch.deg2rad(alpha)
pressure_lift = -cp.mean(dim=(1, 2)) * torch.cos(alpha_rad)
friction_drag = cf_tau.abs().mean(dim=(1, 2)) + 0.1 * cf_z.abs().mean(dim=(1, 2))
pressure_drag = cp.abs().mean(dim=(1, 2)) * torch.sin(alpha_rad).abs()
ref_area = self.ref_area_override if self.ref_area_override is not None else 1.0
cl = pressure_lift / max(ref_area, 1e-8)
cd = self.cd_bias + (friction_drag + pressure_drag) / max(ref_area, 1e-8)
cm = (cp * torch.linspace(0.0, 1.0, cp.shape[-1], device=cp.device)).mean(
dim=(1, 2)
)
return torch.stack([cl, cd, cm], dim=1)
def predict_fields_and_coefficients(self, x: torch.Tensor) -> AeroTransformerOutput:
if self.geometry_builder is None:
fields = self.predict_fields(x)
coefficients = self.fields_to_coefficients(fields)
return AeroTransformerOutput(fields=fields, coefficients=coefficients)
geometry = self.build_geometry(x)
condition = self.make_condition(geometry.geom0.shape[0], geometry.geom0.dtype, geometry.geom0.device)
fields = self.model(geometry.geom0, condition)
coefficients = self.fields_to_coefficients(fields, geometry=geometry)
return AeroTransformerOutput(fields=fields, coefficients=coefficients)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.predict_fields_and_coefficients(x).coefficients