model_configs.py

recipes/aero_cfd/showcase/model_configs.py

 1#  Copyright © 2026 Emmi AI GmbH. All rights reserved.
 2
 3from __future__ import annotations
 4
 5from dataclasses import dataclass
 6from enum import Enum
 7from typing import Any
 8
 9
10class ModelSize(str, Enum):
11    """Available AB-UPT model sizes."""
12
13    small = "small"
14    scaled = "scaled"
15    scaled_mps = "scaled_mps"
16
17
18ABUPT_MODEL_KIND = "noether.modeling.models.aerodynamics.AeroABUPT"
19TRAINER_KIND = "noether.training.trainers.WeightedLossTrainer"
20
21FIELD_WEIGHTS = {
22    "surface_pressure": 1.0,
23    "surface_friction": 1.0,
24    "volume_pressure": 1.0,
25    "volume_velocity": 1.0,
26    "volume_vorticity": 1.0,
27}
28
29
30@dataclass
31class ABUPTSizeConfig:
32    """AB-UPT model and pipeline configuration for a given size tier."""
33
34    model_params: dict[str, Any]
35    pipeline_overrides: dict[str, Any]
36
37
38# Physics block patterns for the AB-UPT architecture.
39# Default: compact 6-block pattern, more efficient with comparable performance.
40PHYSICS_BLOCKS_DEFAULT: list[str] = ["perceiver", "self", "cross", "self", "cross", "self"]
41# Original: 11-block pattern from the AB-UPT paper (arxiv:2502.09587).
42PHYSICS_BLOCKS_PAPER: list[str] = ["perceiver"] + ["self", "cross"] * 5 + ["self"]
43
44MODEL_SIZES: dict[str, ABUPTSizeConfig] = {
45    "small": ABUPTSizeConfig(
46        model_params=dict(
47            hidden_dim=192,
48            num_heads=3,
49            geometry_depth=1,
50            physics_blocks=PHYSICS_BLOCKS_DEFAULT,
51            num_domain_decoder_blocks={"surface": 2, "volume": 2},
52            radius=0.25,
53        ),
54        pipeline_overrides=dict(
55            num_geometry_points=16384,
56            num_geometry_supernodes=1024,
57            num_surface_anchor_points=512,
58            num_volume_anchor_points=512,
59        ),
60    ),
61    "scaled": ABUPTSizeConfig(
62        model_params=dict(
63            hidden_dim=384,
64            num_heads=6,
65            geometry_depth=1,
66            physics_blocks=PHYSICS_BLOCKS_DEFAULT,
67            num_domain_decoder_blocks={"surface": 6, "volume": 6},
68            radius=0.1,
69        ),
70        pipeline_overrides=dict(
71            num_geometry_points=125000,
72            num_geometry_supernodes=32000,
73            num_surface_anchor_points=16000,
74            num_volume_anchor_points=32000,
75        ),
76    ),
77    # Scaled model with reduced point budget for MPS (Apple Silicon)
78    "scaled_mps": ABUPTSizeConfig(
79        model_params=dict(
80            hidden_dim=384,
81            num_heads=6,
82            geometry_depth=1,
83            physics_blocks=PHYSICS_BLOCKS_DEFAULT,
84            num_domain_decoder_blocks={"surface": 6, "volume": 6},
85            radius=0.1,
86        ),
87        pipeline_overrides=dict(
88            num_geometry_points=16384,
89            num_geometry_supernodes=4096,
90            num_surface_anchor_points=2048,
91            num_volume_anchor_points=4096,
92        ),
93    ),
94}