Add mf timeseries for multiple sites

fluxy/io.py

@@ -482,15 +482,6 @@ def read_model_output(
         if add_sites_to_flux and file_type == DataTypes.FLUX:
             ds_all[m] = add_sites_var(ds_all[m], filepath, m, period[i], config_data)
 
-        # Overwrite species attributes
-        current_species = ds_all[m].attrs.get("species", "not set")
-        ds_all[m].attrs["species"] = current_species
-        if species is not None and current_species != species:
-            logger.info(
-                f"'species' attribute in dataset {m} ({current_species}) differs from species {species}. It is overwritten."
-            )
-            ds_all[m].attrs["species"] = species
-
     return ds_all
 
 
@@ -780,9 +771,10 @@ def edit_vars_and_attributes(
             xarray dataset with updated variables and attributes.
     """
 
-    # Add inversion frequency to global attributes
+    # Add inversion frequency and exp name to global attributes
     if "frequency" not in ds.attrs:
         ds.attrs["frequency"] = frequency
+    ds.attrs["exp_name"] = model
 
     # Rename legacy variables
     name_dict = {
@@ -799,14 +791,14 @@ def edit_vars_and_attributes(
     filename_tags = os.path.basename(model)
     m0 = filename_tags.split("_")[0].lower()
 
-    # check the species
+    # check (and overwrite) species attribute
     if species is not None:
-        if "species" not in ds.attrs:
-            ds.attrs["species"] = species
-        elif ds.attrs["species"] != species:
+        current_species = ds.attrs.get("species", "not set")
+        if current_species != species:
             logger.info(
-                f"Species {ds.attrs['species']} in dataset does not match species {species} in model {model}."
+                f"'species' attribute in dataset {model} ({current_species}) differs from species {species}. It is overwritten."
             )
+            ds.attrs["species"] = species
 
     file_type = DataTypes(file_type)
 

fluxy/operators/mf.py

@@ -74,6 +74,18 @@ def compute_mf_difference(
     unique_platforms, platform_indices = np.unique(
         common_platforms, return_inverse=True
     )
+
+    # check species
+    species_set = {ds_left.attrs["species"], ds_right.attrs["species"]}
+    if len(species_set) != 1:
+        logger.warning(
+            f"Different species found {species_set} between the 2 compared datasets."
+        )
+        species_set = "mix"
+    else:
+        species = list(species_set)[0]
+
+    # create dataset with coord and dim
     ds_diff[key_name] = xr.Dataset(
         coords={
             "time": ("index", common_index.get_level_values("time").values),
@@ -82,6 +94,8 @@ def compute_mf_difference(
         },
         attrs={
             "description": f"Difference between {model_left} and {model_right}",
+            "exp_name": f"{ds_left.attrs['exp_name']} - {ds_right.attrs['exp_name']}",
+            "species": species,
         },
     )
 

fluxy/operators/select.py

@@ -55,7 +55,7 @@ def slice_flux(
 
     """
     ds_all_sliced = dict()
-    species_info = config_data.get("species_info",{}).get(species, None)
+    species_info = config_data.get("species_info", {}).get(species, None)
 
     if type(start_date) is str:
         start_date = [start_date] * len(ds_all.keys())
@@ -94,7 +94,7 @@ def slice_mf(
     ds_all: dict[str, xr.Dataset],
     start_date: str = None,
     end_date: str = None,
-    site: str = None,
+    site: str | list[str] | None = None,
     baseline_site: str = None,
     baseline_filename: str = "InTEM_baseline_timestamps",
     data_dir: os.PathLike | None = None,
@@ -114,7 +114,7 @@ def slice_mf(
         end_date (str):
             Date to slice data to, e.g. '2022-01-01' would include all
             data up to 2021-12-31.
-        site (str):
+        site (str | list[str] | None):
             Obs site to select data from, e.g. 'MHD'.
         baseline_site (str):
             Site used to define baseline at, options for 'MHD', 'JFJ', or 'CMN'.
@@ -225,29 +225,51 @@ def slice_mf(
     return ds_all
 
 
-def slice_site(ds: xr.Dataset, site: str) -> xr.Dataset:
+def slice_site(
+    ds: xr.Dataset,
+    site: str | list[str],
+    raise_error: bool = True,
+) -> xr.Dataset | None:
     """
     Slices the dataset to only include data for a given site.
 
     Args:
         ds (xarray dataset):
             Dataset with mf data of a given model.
-        site (str):
-            Site of interest.
+        site (str | list[str]):
+            Site(s) of interest.
+        raise_error:
+            if True, raise an error if the site is not found in the dataset.
     Returns:
         ds (xarray dataset):
-            Dataset with mf data of a given model, sliced to only include data for the given site.
+            Dataset with mf data of a given model, sliced to only include data for the given site(s).
     """
+    if isinstance(site,str):
+        sites = [site]
+    else:
+        sites = site
 
-    site_index = get_site_index(ds, site)
-
-    if site_index is None:
-        raise ValueError(f"Site {site} not found in dataset.")
+    mask = ds["number_of_identifier"] * False
+    site_indices = []
+    for site in sites:
+        site_index = get_site_index(ds, site)
+        site_indices.append(site_index)
 
-    mask = ds["number_of_identifier"] == site_index
-    ds = ds.where(mask, drop=True)
+        if site_index is not None:
+            mask += ds["number_of_identifier"] == site_index
+        else:
+            logger.warning(f"Site {site} not found for model {ds.attrs['exp_name']}.")
 
-    return ds
+    if mask.any():
+        ds = ds.where(mask, drop=True)
+        return ds
+    else:
+        msg = f"No data for any sites {sites} with indices {site_indices} in model {ds.attrs['exp_name']}."
+        if raise_error:
+            raise ValueError(msg)
+        else:
+            logger.warning(msg)
+            return None
 
 
 def slice_height(ds: xr.Dataset, intake_height: float) -> xr.Dataset:
@@ -293,11 +315,11 @@ def get_site_index(ds: xr.Dataset, site: str) -> int | None:
             Index of site of interest in the dataset.
             Returns None if site does not exist.
     """
-
+    
     if site in ds["platform"]:
         index = np.where(ds["platform"] == site)[0][0]
         return index
-
+    
     return None
 
 
@@ -496,3 +518,55 @@ def clean_timeseries_missing_data(
     ds = ds.sortby("time")
 
     return ds
+
+
+def slice_site_dict_of_datasets(
+    ds_all: dict[str, xr.Dataset],
+    site: str,
+) -> dict[str, xr.Dataset]:
+    """
+    Slices all datasets in a dictionary to only include data for a given site.
+    Args:
+        ds_all (dictionary of datasets):
+            xarray datasets read directly from each model's flux netCDF.
+        site (str):
+            Site of interest.
+    Returns:
+        ds_all_site (dictionary of datasets):
+            xarray datasets, sliced to only include data for the given site.
+    """
+
+    ds_all_site = dict()
+
+    for m, ds in ds_all.items():
+        logger.info(f"Slicing site {site} from {m}.")
+
+        if site in ds["platform"].values:
+            ds_all_site[m] = slice_site(ds, site)
+        else:
+            logger.warning(
+                f"Site {site} not found in dataset for {m}. "
+                f"Continuing without {m} - {site}."
+            )
+
+    return ds_all_site
+
+
+def check_site_list(site_list: list | None, ds_all: dict[str, xr.Dataset]) -> list:
+    """
+    Check that every site in the list exists. If None, set it to all the sites available.
+    Args:
+        site_list: list of sites to check
+        ds_all: datasets into which check for the sites
+    Returns:
+        site_list: list of sites
+    """
+    if site_list is None:
+        return get_unique_sites(ds_all)
+    available_sites = get_unique_sites(ds_all)
+    for site in site_list:
+        if site not in available_sites:
+            raise ValueError(
+                f"Site {site} not found in the datasets provided. Available sites are {available_sites}."
+            )
+    return site_list

fluxy/operators/stats.py

@@ -2,7 +2,7 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
-from fluxy.operators.select import get_unique_sites, get_site_index
+from fluxy.operators.select import get_unique_sites, slice_site
 
 
 def stats_observed_vs_simulated(
@@ -77,17 +77,9 @@ def stats_observed_vs_simulated(
     # Compute stats for all sites and all models
     for site in sites_all:
         for model, ds in ds_all.items():
-            site_index = get_site_index(ds, site)
-            if site_index is None:
-                logger.warning(f"Site {site} not found in model {model}.")
+            ds_site = slice_site(ds, site, raise_error=False) 
+            if not ds_site:
                 continue
-            mask_site = ds["number_of_identifier"] == site_index
-            if not mask_site.any():
-                logger.warning(
-                    f"No data for site {site} with index {site_index} in model {model}."
-                )
-                continue
-            ds_site = ds.where(mask_site, drop=True)
 
             # select what to compare
             obs = ds_site[obs_var]