some arg cleaning in era_downloader

pyproject.toml

@@ -41,7 +41,7 @@ dependencies = [
  "pytest>=5.2",
  "scipy>=1.0.0",
  "sphinx>=7.0",
- "tensorflow>2.4,<2.16",
+ "tensorflow>2.4,<2.10",
  "xarray>=2023.0",
 ]
 

sup3r/utilities/era_downloader.py

@@ -94,7 +94,8 @@ def __init__(self,
  run_interp=True,
  overwrite=False,
  variables=None,
- check_files=False):
+ check_files=False,
+ product_type='reanalysis'):
  """Initialize the class.
 
  Parameters
@@ -123,6 +124,9 @@ def __init__(self,
  and wind components.
  check_files : bool
  Check existing files. Remove and redownload if checks fail.
+ product_type : str
+ Can be 'reanalysis', 'ensemble_mean', 'ensemble_spread',
+ 'ensemble_members'
  """
  self.year = year
  self.month = month
@@ -141,12 +145,23 @@ def __init__(self,
  self.sfc_file_variables = ['geopotential']
  self.level_file_variables = ['geopotential']
  self.prep_var_lists(self.variables)
+ self.product_type = product_type
+ self.hours = self.get_hours()
 
  msg = ('Initialized EraDownloader with: '
  f'year={self.year}, month={self.month}, area={self.area}, '
  f'levels={self.levels}, variables={self.variables}')
  logger.info(msg)
 
+ def get_hours(self):
+ """ERA5 is hourly and EDA is 3-hourly. Check and warn for incompatible
+ requests."""
+ if self.product_type == 'reanalysis':
+ hours = [str(n).zfill(2) + ":00" for n in range(0, 24)]
+ else:
+ hours = [str(n).zfill(2) + ":00" for n in range(0, 24, 3)]
+ return hours
+
  @property
  def variables(self):
  """Get list of requested variables"""
@@ -275,18 +290,20 @@ def download_process_combine(self):
  if sfc_check:
  self.download_file(self.sfc_file_variables, time_dict=time_dict,
  area=self.area, out_file=self.surface_file,
- level_type='single', overwrite=self.overwrite)
+ level_type='single', overwrite=self.overwrite,
+ product_type=self.product_type)
  if level_check:
  self.download_file(self.level_file_variables, time_dict=time_dict,
  area=self.area, out_file=self.level_file,
  level_type='pressure', levels=self.levels,
- overwrite=self.overwrite)
+ overwrite=self.overwrite,
+ product_type=self.product_type)
  if sfc_check or level_check:
  self.process_and_combine()
 
  @classmethod
  def download_file(cls, variables, time_dict, area, out_file, level_type,
- levels=None, overwrite=False):
+ levels=None, product_type='reanalysis', overwrite=False):
  """Download either single-level or pressure-level file
 
  Parameters
@@ -304,6 +321,9 @@ def download_file(cls, variables, time_dict, area, out_file, level_type,
  Either 'single' or 'pressure'
  levels : list
  List of pressure levels to download, if level_type == 'pressure'
+ product_type : str
+ Can be 'reanalysis', 'ensemble_mean', 'ensemble_spread',
+ 'ensemble_members'
  overwrite : bool
  Whether to overwrite existing file
  """
@@ -555,7 +575,7 @@ def run_interpolation(self, max_workers=None, **kwargs):
  overwrite=self.overwrite,
  **kwargs)
 
- def get_monthly_file(self, interp_workers=None, prune_variables=None,
+ def get_monthly_file(self, interp_workers=None, prune_variables=False,
  **interp_kwargs):
  """Download level and surface files, process variables, and combine
  processed files. Includes checks for shape and variables and option to
@@ -604,8 +624,8 @@ def all_months_exist(cls, year, file_pattern):
  @classmethod
  def already_pruned(cls, infile, prune_variables):
  """Check if file has been pruned already."""
- if prune_variables is None:
- logger.info('Received prune_variables=None. Skipping pruning.')
+ if not prune_variables:
+ logger.info('Received prune_variables=False. Skipping pruning.')
  return
  with xr.open_dataset(infile) as ds:
  check_variables = [var for var in ds.data_vars
@@ -614,15 +634,15 @@ def already_pruned(cls, infile, prune_variables):
  return pruned
 
  @classmethod
- def prune_output(cls, infile, prune_variables=None):
+ def prune_output(cls, infile, prune_variables=False):
  """Prune output file to keep just single level variables"""
- if prune_variables is None:
- logger.info('Received prune_variables=None. Skipping pruning.')
+ if not prune_variables:
+ logger.info('Received prune_variables=False. Skipping pruning.')
  return
  else:
  logger.info(f'Pruning {infile}.')
  tmp_file = cls.get_tmp_file(infile)
- with xr.Dataset(infile) as ds:
+ with xr.open_dataset(infile) as ds:
  keep_vars = {k: v for k, v in dict(ds.data_vars)
  if 'level' not in ds[k].dims}
  new_coords = {k: v for k, v in dict(ds.coords).items()
@@ -645,8 +665,9 @@ def run_month(cls,
  overwrite=False,
  interp_workers=None,
  variables=None,
- prune_variables=None,
+ prune_variables=False,
  check_files=False,
+ product_type='reanalysis',
  **interp_kwargs):
  """Run routine for all months in the requested year.
 
@@ -676,13 +697,16 @@ def run_month(cls,
  variables : list | None
  Variables to download. If None this defaults to just gepotential
  and wind components.
- prune_variables : bool | None
- Variables to remove from final files. This is usually the multi
- pressure level array of a variable which has since been
- interpolated to specific heights.
- pruned.
+ prune_variables : bool
+ Whether to remove 4D variables from data after interpolation. e.g.
+ height interpolation could give u_10m, u_100m, u_120m from a 4D u
+ array. If we only need these heights we could remove the 4D u array
+ from the final data file.
  check_files : bool
  Check existing files. Remove and redownload if checks fail.
+ product_type : str
+ Can be 'reanalysis', 'ensemble_mean', 'ensemble_spread',
+ 'ensemble_members'
  **interp_kwargs : dict
  Keyword args for LogLinInterpolator.run()
  """
@@ -695,7 +719,8 @@ def run_month(cls,
  run_interp=run_interp,
  overwrite=overwrite,
  variables=variables,
- check_files=check_files)
+ check_files=check_files,
+ product_type=product_type)
  downloader.get_monthly_file(interp_workers=interp_workers,
  prune_variables=prune_variables,
  **interp_kwargs)
@@ -714,8 +739,9 @@ def run_year(cls,
  max_workers=None,
  interp_workers=None,
  variables=None,
- prune_variables=None,
+ prune_variables=False,
  check_files=False,
+ product_type='reanalysis',
  **interp_kwargs):
  """Run routine for all months in the requested year.
 
@@ -750,11 +776,16 @@ def run_year(cls,
  variables : list | None
  Variables to download. If None this defaults to just gepotential
  and wind components.
- prune_variables : list | None
- Variables to keep in final files. All other variables will be
- pruned.
+ prune_variables : bool
+ Whether to remove 4D variables from data after interpolation. e.g.
+ height interpolation could give u_10m, u_100m, u_120m from a 4D u
+ array. If we only need these heights we could remove the 4D u array
+ from the final data file.
  check_files : bool
  Check existing files. Remove and redownload if checks fail.
+ product_type : str
+ Can be 'reanalysis', 'ensemble_mean', 'ensemble_spread',
+ 'ensemble_members'
  **interp_kwargs : dict
  Keyword args for LogLinInterpolator.run()
  """
@@ -772,6 +803,7 @@ def run_year(cls,
  variables=variables,
  prune_variables=prune_variables,
  check_files=check_files,
+ product_type=product_type,
  **interp_kwargs)
  else:
  futures = {}
@@ -791,6 +823,7 @@ def run_year(cls,
  prune_variables=prune_variables,
  variables=variables,
  check_files=check_files,
+ product_type=product_type,
  **interp_kwargs)
  futures[future] = {'year': year, 'month': month}
  logger.info(f'Submitted future for year {year} and month '