From c0c2102237567de9988c8f4dce22b49abc8ee9c6 Mon Sep 17 00:00:00 2001
From: Corey Adams <6619961+coreyjadams@users.noreply.github.com>
Date: Tue, 11 Nov 2025 02:46:45 +0000
Subject: [PATCH] Ensure training sets factors to none if not used, and doesn't
 look for them if not needed.

---
 .../external_aerodynamics/domino/src/utils.py | 76 ++++++++++++-------
 1 file changed, 47 insertions(+), 29 deletions(-)

diff --git a/examples/cfd/external_aerodynamics/domino/src/utils.py b/examples/cfd/external_aerodynamics/domino/src/utils.py
index 5ec93d2381..adae03c1fd 100644
--- a/examples/cfd/external_aerodynamics/domino/src/utils.py
+++ b/examples/cfd/external_aerodynamics/domino/src/utils.py
@@ -336,40 +336,58 @@ def load_scaling_factors(
         )
 
     if cfg.model.normalization == "min_max_scaling":
-        vol_factors = np.asarray(
-            [
-                scaling_factors.max_val["volume_fields"],
-                scaling_factors.min_val["volume_fields"],
-            ]
-        )
-        surf_factors = np.asarray(
-            [
-                scaling_factors.max_val["surface_fields"],
-                scaling_factors.min_val["surface_fields"],
-            ]
-        )
+        if cfg.model.model_type == "volume" or cfg.model.model_type == "combined":
+            vol_factors = np.asarray(
+                [
+                    scaling_factors.max_val["volume_fields"],
+                    scaling_factors.min_val["volume_fields"],
+                ]
+            )
+        else:
+            vol_factors = None
+        if cfg.model.model_type == "surface" or cfg.model.model_type == "combined":
+            surf_factors = np.asarray(
+                [
+                    scaling_factors.max_val["surface_fields"],
+                    scaling_factors.min_val["surface_fields"],
+                ]
+            )
+        else:
+            surf_factors = None
     elif cfg.model.normalization == "mean_std_scaling":
-        vol_factors = np.asarray(
-            [
-                scaling_factors.mean["volume_fields"],
-                scaling_factors.std["volume_fields"],
-            ]
-        )
-        surf_factors = np.asarray(
-            [
-                scaling_factors.mean["surface_fields"],
-                scaling_factors.std["surface_fields"],
-            ]
-        )
+        if cfg.model.model_type == "volume" or cfg.model.model_type == "combined":
+            vol_factors = np.asarray(
+                [
+                    scaling_factors.mean["volume_fields"],
+                    scaling_factors.std["volume_fields"],
+                ]
+            )
+        else:
+            vol_factors = None
+        if cfg.model.model_type == "surface" or cfg.model.model_type == "combined":
+            surf_factors = np.asarray(
+                [
+                    scaling_factors.mean["surface_fields"],
+                    scaling_factors.std["surface_fields"],
+                ]
+            )
+        else:
+            surf_factors = None
     else:
         raise ValueError(f"Invalid normalization mode: {cfg.model.normalization}")
 
-    vol_factors_tensor = torch.from_numpy(vol_factors)
-    surf_factors_tensor = torch.from_numpy(surf_factors)
-
     dm = DistributedManager()
-    vol_factors_tensor = vol_factors_tensor.to(dm.device, dtype=torch.float32)
-    surf_factors_tensor = surf_factors_tensor.to(dm.device, dtype=torch.float32)
+
+    if cfg.model.model_type == "volume" or cfg.model.model_type == "combined":
+        vol_factors_tensor = torch.from_numpy(vol_factors)
+        vol_factors_tensor = vol_factors_tensor.to(dm.device, dtype=torch.float32)
+    else:
+        vol_factors_tensor = None
+    if cfg.model.model_type == "surface" or cfg.model.model_type == "combined":
+        surf_factors_tensor = torch.from_numpy(surf_factors)
+        surf_factors_tensor = surf_factors_tensor.to(dm.device, dtype=torch.float32)
+    else:
+        surf_factors_tensor = None
 
     return vol_factors_tensor, surf_factors_tensor