fix(tools): type hints and API updates for python scripts

Author: openhands-agent

tools/ashrae_140_test_harness.py

@@ -128,7 +128,9 @@ def run_rust_tests(
             duration = (datetime.now() - start_time).total_seconds()
 
             # Parse test output
-            total, passed, failed, skipped = self._parse_cTest_output(result.stdout)
+            total, passed, failed, skipped = self._parse_cargo_test_output(
+                result.stdout
+            )
 
             suite_result = TestSuiteResult(
                 suite_name=test_pattern,

tools/ep_oracle.py

@@ -17,7 +17,7 @@
 import subprocess
 import sys
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 
 class EnergyPlusOracle:
@@ -382,16 +382,16 @@ def run_energyplus(
         Returns:
             Dictionary with simulation results
         """
-        output_dir = Path(output_dir)
-        output_dir.mkdir(parents=True, exist_ok=True)
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
 
         # Build EnergyPlus command
         ep_cmd = [
             str(self.ep_path / "energyplus"),
             "-w",
             str(epw_path),
             "-d",
-            str(output_dir),
+            str(output_path),
             "-r",  # Read variables
             idf_path,
         ]
@@ -412,17 +412,17 @@ def run_energyplus(
             raise RuntimeError(f"EnergyPlus simulation failed: {result.returncode}")
 
         # Parse SQL output
-        sql_file = output_dir / "eplusout.sql"
+        sql_file = output_path / "eplusout.sql"
         if sql_file.exists():
             return self._parse_sql_results(sql_file)
         else:
             raise RuntimeError("EnergyPlus SQL output file not found")
 
-    def _parse_sql_results(self, sql_file: Path) -> Dict:
+    def _parse_sql_results(self, sql_file: Path) -> Dict[str, Any]:
         """Parse EnergyPlus SQL output."""
         import sqlite3
 
-        results = {}
+        results: Dict[str, Any] = {}
 
         try:
             conn = sqlite3.connect(sql_file)
@@ -529,14 +529,14 @@ def compare_results(
         ep_results: Dict,
         tolerance_abs: float = 1.0,
         tolerance_rel: float = 0.05,
-    ) -> Dict:
+    ) -> Dict[str, Any]:
         """
         Compare Fluxion and EnergyPlus results.
 
         Returns:
             Comparison report with metrics
         """
-        report = {
+        report: Dict[str, Any] = {
             "passed": True,
             "metrics": {},
             "details": {},

tools/geometry_extraction.py

@@ -959,7 +959,7 @@ def to_cta_tensors(self, geometry: BuildingGeometry) -> Dict[str, np.ndarray]:
         Returns:
             Dictionary of numpy arrays
         """
-        tensors = {}
+        tensors: Dict[str, np.ndarray] = {}
 
         # Zone coordinates tensor
         # Format: [x1, y1, x2, y2, ..., x8, y8, floor_height, ceiling_height, area, volume, perimeter, zone_id]

tools/integrate_ml_surrogate_pipeline.py

@@ -358,14 +358,13 @@ def forward(self, x):
         # Export to ONNX
         torch.onnx.export(
             model,
-            dummy_input,
-            output_path,
+            (dummy_input,),
+            str(output_path),
             input_names=["input"],
             output_names=["output"],
             dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
             opset_version=17,
         )
-
         logger.info(f"✓ Model exported successfully to {output_path}")
         return True
 

tools/physics_informed_loss.py

@@ -169,6 +169,9 @@ class PhysicsInformedLoss(nn.Module):
     Physics-informed loss function that combines MSE with physics constraints.
     """
 
+    min_temp: torch.Tensor
+    max_temp: torch.Tensor
+
     def __init__(
         self,
         mse_weight: float = 1.0,

tools/tests/test_train_integration.py

@@ -9,8 +9,8 @@
 # Add project root to path
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
 
-from tools import train_surrogate
 from tools.data_gen import geometry, simulation
+from tools.train_ml_surrogate import train_surrogate
 
 
 class TestTrainIntegration(unittest.TestCase):

tools/train_ensemble_surrogates.py

@@ -39,7 +39,7 @@ def __init__(
         torch.manual_seed(seed)
         np.random.seed(seed)
 
-        layers = []
+        layers: List[nn.Module] = []
         prev_dim = input_dim
 
         for hidden_dim in hidden_dims:
@@ -153,8 +153,8 @@ def export_onnx(model: nn.Module, sample_input: np.ndarray, output_path: Path):
 
     torch.onnx.export(
         model,
-        dummy_input,
-        output_path,
+        (dummy_input,),
+        str(output_path),
         input_names=["input"],
         output_names=["output"],
         dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
@@ -253,20 +253,19 @@ def evaluate_ensemble(
     X_t = torch.from_numpy(X)
 
     # Get predictions from all models
-    all_predictions = []
+    all_predictions_list = []
     with torch.no_grad():
         for model in models:
             pred = model(X_t).numpy()
-            all_predictions.append(pred)
+            all_predictions_list.append(pred)
 
-    all_predictions = np.array(
-        all_predictions
+    all_predictions_arr = np.array(
+        all_predictions_list
     )  # Shape: (num_models, num_samples, output_dim)
 
     # Calculate ensemble statistics
-    mean_pred = np.mean(all_predictions, axis=0)
-    std_pred = np.std(all_predictions, axis=0)
-
+    mean_pred = np.mean(all_predictions_arr, axis=0)
+    std_pred = np.std(all_predictions_arr, axis=0)
     # Calculate metrics
     mse = np.mean((y - mean_pred) ** 2)
     mae = np.mean(np.abs(y - mean_pred))

tools/train_ml_surrogate.py

@@ -370,7 +370,7 @@ def train_surrogate(
     # Optimizer and loss
     optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-5)
     scheduler = optim.lr_scheduler.ReduceLROnPlateau(
-        optimizer, "min", patience=10, factor=0.5, verbose=True
+        optimizer, "min", patience=10, factor=0.5
     )
     criterion = PhysicsLoss()
 

tools/train_pinn.py

@@ -33,7 +33,7 @@
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -126,7 +126,7 @@ def __init__(
     ):
         super().__init__()
 
-        layers = []
+        layers: List[nn.Module] = []
         prev_dim = input_dim
 
         # First layer: input -> first hidden
@@ -477,13 +477,13 @@ def generate_training_data(
         Returns:
             Tuple of (inputs_dict, targets_array)
         """
-        X: Dict[str, List[float]] = {
+        X: Dict[str, Any] = {
             "time": [],
             "t_outdoor": [],
             "q_solar": [],
             "q_internal": [],
         }
-        y = []
+        y: List[float] = []
 
         for _ in range(n_samples):
             # Generate outdoor temperature profile (daily cycle with noise)
@@ -521,11 +521,12 @@ def generate_training_data(
             y.extend(t_indoor)
 
         # Convert to arrays
+        X_array: Dict[str, np.ndarray] = {}
         for key in X:
-            X[key] = np.array(X[key], dtype=np.float32)
-        y = np.array(y, dtype=np.float32).reshape(-1, 1)
+            X_array[key] = np.array(X[key], dtype=np.float32)
+        y_array = np.array(y, dtype=np.float32).reshape(-1, 1)
 
-        return X, y
+        return X_array, y_array
 
     def generate_collocation_points(
         self,
@@ -800,8 +801,8 @@ def export_onnx(
     model.eval()
     torch.onnx.export(
         model,
-        sample_input,
-        output_path,
+        (sample_input,),
+        str(output_path),
         input_names=["input"],
         output_names=["t_indoor"],
         dynamic_axes={
@@ -810,6 +811,7 @@ def export_onnx(
         },
         opset_version=17,
     )
+
     logger.info("ONNX export successful")