clean up code

mllam · Feb 11, 2025 · 0948f9a · 0948f9a
1 parent 8698f3c
commit 0948f9a
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py
@@ -322,21 +322,23 @@ def get_standardization_dataarray(self, category: str) -> xr.Dataset:
         stats_variables = {
             f"{category}__{split}__{op}": f"{category}_{op}" for op in ops
         }
+
+        # Add state diff stats
         if category == "state":
             stats_variables.update(
                 {f"state__{split}__diff_{op}": f"state_diff_{op}" for op in ops}
             )
 
         ds_stats = self._ds[stats_variables.keys()].rename(stats_variables)
 
+        # Add standardized state diff stats
         if category == "state":
-            standardized_state_diff = {
-                op: ds_stats[f"state_diff_{op}"] / ds_stats["state_std"]
-                for op in ops
-            }
             ds_stats = ds_stats.assign(
                 **{
-                    f"state_diff_{op}_standardized": standardized_state_diff[op]
+                    f"state_diff_{op}_standardized": ds_stats[
+                        f"state_diff_{op}"
+                    ]
+                    / ds_stats["state_std"]
                     for op in ops
                 }
             )