Misunderstood lists and dicts in argsparge -> bugfix

mllam · May 30, 2024 · 2c07814 · 2c07814
1 parent 7ce40d8
commit 2c07814
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -473,7 +473,7 @@ def create_metric_log_dict(self, metric_tensor, prefix, metric_name):
         # Check if metrics are watched, log exact values for specific vars
         if full_log_name in self.args.metrics_watch:
             for var_i, timesteps in self.args.var_leads_metrics_watch.items():
-                var = self.config_loader.dataset.var_nums[var_i]
+                var = self.config_loader.dataset.var_names[var_i]
                 log_dict.update(
                     {
                         f"{full_log_name}_{var}_step_{step}": metric_tensor[

diff --git a/train_model.py b/train_model.py
@@ -197,19 +197,21 @@ def main():
     )
     parser.add_argument(
         "--metrics_watch",
-        type=list,
+        nargs="+",
         default=[],
         help="List of metrics to watch, including any prefix (e.g. val_rmse)",
     )
     parser.add_argument(
         "--var_leads_metrics_watch",
         type=str,
         default="{}",
-        help="JSON string with variables and lead times to log watched metrics"
-        # e.g. '{"var1": [1, 2], "var2": [3, 4]}'
+        help="""JSON string with variable-IDs and lead times to log watched
+             metrics (e.g. '{"1": [1, 2], "3": [3, 4]}')""",
     )
     args = parser.parse_args()
-    args.var_leads_metrics_watch = json.loads(args.var_leads_metrics_watch)
+    args.var_leads_metrics_watch = {
+        int(k): v for k, v in json.loads(args.var_leads_metrics_watch).items()
+    }
 
     config_loader = config.Config.from_file(args.data_config)