Ensemble historgram

src/valenspy/_utilities/_datatree.py

@@ -110,12 +110,19 @@ def datatree_to_dataframe(dt: xr.DataTree, add_attributes=False):
             df["id"] = str(key)
 
             if add_attributes:
-                for attr in add_attributes:
-                    #if attr is a substring of any attribute in ds.attrs:
-                    for ds_attr in ds.attrs:
-                        if attr in ds_attr:
-                            df[attr] = ds.attrs[ds_attr]
-                            break
+                if isinstance(add_attributes, bool): #If add_attributes is True, we add all attributes. If it's a list, we only add the specified attributes.
+                    for attr in ds.attrs:
+                        df[attr] = ds.attrs[attr]
+                    for var in ds.data_vars: #Add all the variable related attributes
+                        for attr in ds[var].attrs:
+                            df[f"{var}_{attr}"] = ds[var].attrs[attr]
+                else:
+                    for attr in add_attributes:
+                        #if attr is a substring of any attribute in ds.attrs:
+                        for ds_attr in ds.attrs:
+                            if attr in ds_attr:
+                                df[attr] = ds.attrs[ds_attr]
+                                break #Break the loop after finding the first match to avoid adding multiple attributes with the same substring.
             data_frames.append(df)
 
     return pd.concat(data_frames, axis=0).reset_index()

src/valenspy/diagnostic/_ensemble2self.py

@@ -2,10 +2,19 @@
 from valenspy.diagnostic.functions import *
 from valenspy.diagnostic.visualizations import *
 
+import seaborn as sns
+
 __all__ = [
     "Ensemble_Quantile_Spatial_Mean",
-    "Ensemble_Quantile_Closest_Member_Spatial_Mean"
-]
+    "Ensemble_Quantile_Closest_Member_Spatial_Mean",
+    "Ensemble_Histogram"
+    ]
+Ensemble_Histogram = Ensemble2Self(
+    ensemble_member_means,
+    sns.histplot,
+    "Ensemble member means",
+    "The histogram of the ensemble member means."
+)
 
 Ensemble_Quantile_Spatial_Mean = Ensemble2Self(
     ensemble_quantile_of_spatial_mean,
@@ -20,4 +29,4 @@
     "Ensemble quantiles of closest member spatial mean",
     "The ensemble members that are closest to the quantiles of the spatial mean, and their spatial mean.",
     plot_type="facetted"
-)
+)

src/valenspy/diagnostic/functions.py

@@ -265,6 +265,23 @@ def calc_metrics_ds(ds_mod: xr.Dataset, ds_obs: xr.Dataset, metrics=None, pss_bi
 # Ensemble2Self diagnostic functions #
 ######################################
 
+def ensemble_member_means(dt: DataTree, add_attributes=False):
+    """
+    Calculate the mean of each ensemble member in the datatree. If lat, lon and/or time dimensions are present, the data is averaged over these dimensions.
+
+    Parameters
+    ----------
+    dt : DataTree
+        The data to calculate the mean of.
+
+    Returns
+    -------
+    xr.Dataset
+        The mean of each ensemble member in the datatree.
+    """
+    dt_m = dt.map_over_datasets(_average_over_dims, ["lat", "lon", "time"])
+    return datatree_to_dataframe(dt_m, add_attributes=add_attributes)
+
 def ensemble_quantile_of_spatial_mean(dt: DataTree, quantile: float | list[float]):
     """
     Calculate the ensemble quantile of the spatial mean of the data. If the time dimension is present, the data is averaged over the time dimension before calculating the percentiles.