Skip to content

Commit

Permalink
update catalog selection
Browse files Browse the repository at this point in the history
  • Loading branch information
larsbuntemeyer committed Jan 25, 2025
1 parent bfbe160 commit 2c1499e
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 7 deletions.
12 changes: 9 additions & 3 deletions evaltools/source.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,12 @@ def open_catalog(url=None):


def get_source_collection(
variable_id, frequency, driving_source_id="ERA5", add_fx=None, catalog=None
variable_id,
frequency,
driving_source_id="ERA5",
add_fx=None,
catalog=None,
**kwargs,
):
"""
Search the catalog for datasets matching the specified variable_id, frequency, and driving_source_id.
Expand All @@ -50,15 +55,16 @@ def get_source_collection(
frequency=frequency,
driving_source_id=driving_source_id,
require_all_on=["source_id"],
**kwargs,
)
source_ids = list(subset.df.source_id.unique())
print(f"Found: {source_ids} for variables: {variable_id}")
if add_fx:
if add_fx is True:
fx = catalog.search(source_id=source_ids, frequency="fx")
fx = catalog.search(source_id=source_ids, frequency="fx", **kwargs)
else:
fx = catalog.search(
source_id=source_ids, frequency="fx", variable_id=add_fx
source_id=source_ids, frequency="fx", variable_id=add_fx, **kwargs
)
if fx.df.empty:
warn(f"static variables not found: {variable_id}")
Expand Down
10 changes: 6 additions & 4 deletions evaltools/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ def iid_to_dict(iid, attrs=None):
return dict(zip(attrs, values))


def dict_to_iid(attrs, drop=None):
def dict_to_iid(attrs, drop=None, delimiter="."):
"""
Convert a dictionary of dataset attributes to a dataset ID.
Expand All @@ -45,10 +45,10 @@ def dict_to_iid(attrs, drop=None):
"""
if drop is None:
drop = []
return ".".join(v for k, v in attrs.items() if k not in drop)
return delimiter.join(v for k, v in attrs.items() if k not in drop)


def short_iid(iid, attrs=None):
def short_iid(iid, attrs=None, delimiter="."):
"""
Convert a dataset ID to a short ID.
Expand All @@ -61,7 +61,9 @@ def short_iid(iid, attrs=None):
"""
if attrs is None:
attrs = ["institution_id", "source_id", "driving_source_id", "experiment_id"]
return dict_to_iid({k: v for k, v in iid_to_dict(iid).items() if k in attrs})
return dict_to_iid(
{k: v for k, v in iid_to_dict(iid).items() if k in attrs}, delimiter=delimiter
)


def sort_by_grid_mapping(dsets):
Expand Down

0 comments on commit 2c1499e

Please sign in to comment.