Fix nodeweights

MeteoSwiss · Dec 11, 2024 · 44f070e · 44f070e
1 parent 2df18a7
commit 44f070e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/anemoi/training/diagnostics/plots.py b/src/anemoi/training/diagnostics/plots.py
@@ -817,8 +817,8 @@ def get_scatter_frame(
         vmin=vmin,
         vmax=vmax,
     )
-    ax.set_xlim((-np.pi, np.pi))
-    ax.set_ylim((-np.pi / 2, np.pi / 2))
+    # ax.set_xlim((-np.pi, np.pi))
+    # ax.set_ylim((-np.pi / 2, np.pi / 2))
     continents.plot_continents(ax)
     ax.set_aspect("auto", adjustable=None)
     _hide_axes_ticks(ax)

diff --git a/src/anemoi/training/losses/nodeweights.py b/src/anemoi/training/losses/nodeweights.py
@@ -60,7 +60,7 @@ def area_weights(self, graph_data: HeteroData) -> torch.Tensor:
         torch.Tensor
             area weights of the target nodes
         """
-        return AreaWeights(norm="unit-max", fill_value=0).compute(graph_data, self.target)
+        return AreaWeights(flat=True, norm="unit-max").compute(graph_data, self.target)
 
     def weights(self, graph_data: HeteroData) -> torch.Tensor:
         """Returns weight of type self.node_attribute for nodes self.target.