shifted mean_ess, sd_edd computing to before plot_kwargs check+artist…

… plotting logic
arviz-devs · Aug 26, 2024 · 620ae63 · 620ae63
1 parent 765e878
commit 620ae63
Showing 1 changed file with 11 additions and 14 deletions.
diff --git a/src/arviz_plots/plots/evolutionplot.py b/src/arviz_plots/plots/evolutionplot.py
@@ -396,18 +396,20 @@ def compute_ess_dataset(
 
     # plot mean and sd and annotate them
     if extra_methods is not False:
-        mean_ess = None
-        sd_ess = None
+        # computing mean_ess
+        mean_dims, mean_aes, mean_ignore = filter_aes(plot_collection, aes_map, "mean", sample_dims)
+        mean_ess = distribution.azstats.ess(
+            dims=mean_dims, method="mean", relative=relative, **stats_kwargs.get("mean", {})
+        )
+
+        # computing sd_ess
+        sd_dims, sd_aes, sd_ignore = filter_aes(plot_collection, aes_map, "sd", sample_dims)
+        sd_ess = distribution.azstats.ess(
+            dims=sd_dims, method="sd", relative=relative, **stats_kwargs.get("sd", {})
+        )
 
         mean_kwargs = copy(plot_kwargs.get("mean", {}))
         if mean_kwargs is not False:
-            mean_dims, mean_aes, mean_ignore = filter_aes(
-                plot_collection, aes_map, "mean", sample_dims
-            )
-            mean_ess = distribution.azstats.ess(
-                dims=mean_dims, method="mean", relative=relative, **stats_kwargs.get("mean", {})
-            )
-
             # getting 2nd default linestyle for chosen backend and assigning it by default
             mean_kwargs.setdefault("linestyle", linestyles[1])
 
@@ -425,11 +427,6 @@ def compute_ess_dataset(
 
         sd_kwargs = copy(plot_kwargs.get("sd", {}))
         if sd_kwargs is not False:
-            sd_dims, sd_aes, sd_ignore = filter_aes(plot_collection, aes_map, "sd", sample_dims)
-            sd_ess = distribution.azstats.ess(
-                dims=sd_dims, method="sd", relative=relative, **stats_kwargs.get("sd", {})
-            )
-
             sd_kwargs.setdefault("linestyle", linestyles[2])
 
             if "color" not in sd_aes: