diff --git a/.github/workflows/pynorms.yaml b/.github/workflows/pynorms.yaml index d4b9b0e..1b771a8 100644 --- a/.github/workflows/pynorms.yaml +++ b/.github/workflows/pynorms.yaml @@ -5,7 +5,6 @@ jobs: pre-commit: runs-on: ubuntu-latest steps: - - name: Check python coding norms with pre-commit - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: diff --git a/.github/workflows/pytests.yaml b/.github/workflows/pytests.yaml new file mode 100644 index 0000000..f9cd6b9 --- /dev/null +++ b/.github/workflows/pytests.yaml @@ -0,0 +1,72 @@ +name: pytests +on: [push, pull_request] + +env: + g2c_VERSION: 2.2.0 + CC: gcc-14 + +jobs: + cache-g2c: + runs-on: ubuntu-latest + strategy: + fail-fast: true + name: Build and cache g2c library + steps: + - name: apt-get and install dependencies + run: | + sudo apt-get update + sudo apt-get install -y libaec-dev libpng-dev zlib1g-dev libjpeg-dev libopenjp2-7-dev + + - name: Cache g2c installation + id: cache-g2c + uses: actions/cache@v4 + with: + path: | + g2c-${{ env.g2c_VERSION }} + key: g2c-${{ env.g2c_VERSION }} + + - name: Checkout, build and install g2c + if: steps.cache-g2c.outputs.cache-hit != 'true' + run: | + wget https://github.com/NOAA-EMC/NCEPLIBS-g2c/archive/refs/tags/v${{ env.g2c_VERSION }}.tar.gz + tar -xzvf v${{ env.g2c_VERSION }}.tar.gz + cd NCEPLIBS-g2c-${{ env.g2c_VERSION }} + mkdir build && cd build + cmake -DUSE_Jasper=OFF -DUSE_OpenJPEG=ON -DBUILD_PNG=ON -DBUILD_AEC=ON -DBUILD_SHARED_LIBS=OFF \ + -DCMAKE_INSTALL_PREFIX=${{ github.workspace }}/g2c-${{ env.g2c_VERSION }} .. + make -j2 + make install + + run_pytests: + runs-on: ubuntu-latest + name: Install and run tests with pytests + needs: cache-g2c + strategy: + matrix: + #python: ["3.11", "3.12", "3.13"] + python: ["3.12"] # Temporarily only test with 3.12 for faster CI + + steps: + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python }} + + - name: Install (upgrade) python dependencies + run: | + pip install --upgrade pip + + - name: Checkout + uses: actions/checkout@v4 + + - name: Install + run: | + export G2C_STATIC=True + export G2C_DIR=${{ github.workspace }}/g2c-${{ env.g2c_VERSION }}/lib + cd $GITHUB_WORKSPACE + pip install .[dev] + + - name: Run pytests + run: | + cd $GITHUB_WORKSPACE + pytest -v diff --git a/.gitignore b/.gitignore index 5e46d37..eadc0e2 100644 --- a/.gitignore +++ b/.gitignore @@ -99,3 +99,9 @@ dmypy.json # VSCode files *vscode* + +# netCDF and grib files +*.nc +*.grib2 +*.grib2.idx +*pgrb2* diff --git a/.gitmodules b/.gitmodules index fd7d84c..90582fd 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,4 @@ [submodule "graphcast"] path = graphcast url = https://github.com/google-deepmind/graphcast.git + ignore = dirty diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f7ec5b1..fba6bfe 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -15,7 +15,7 @@ repos: hooks: - id: black language_version: python3.12 # Specify your Python version - files: ^(src/|tests/) + files: ^(src/|tests/)$ exclude: ^(oper/|graphcast/|training/) # Flake8 linter for style guide enforcement and error checking diff --git a/config/gefs_varinfo.yaml b/config/gefs_varinfo.yaml new file mode 100644 index 0000000..954d9c4 --- /dev/null +++ b/config/gefs_varinfo.yaml @@ -0,0 +1,20 @@ +time_variant: + pgrb2s.0p25.f000: + - variables: ["TMP"] + levels: ["2 m above ground"] + - variables: ["PRMSL"] + levels: ["mean sea level"] + - variables: ["UGRD", "VGRD"] + levels: ["10 m above ground"] + pgrb2.0p25.f000: + - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"] + levels: ["50 mb", "100 mb", "150 mb", "200 mb", "250 mb", "300 mb", "400 mb", "500 mb", "600 mb", "700 mb", "850 mb", "925 mb", "1000 mb"] +time_invariant: + pgrb2.0p25.f000: + - variables: ["LAND"] + levels: ["surface"] + first_time_step_only: True + pgrb2s.0p25.f000: + - variables: ["HGT"] + levels: ["surface"] + first_time_step_only: True diff --git a/config/gfs_varinfo.yaml b/config/gfs_varinfo.yaml new file mode 100644 index 0000000..4b5c23a --- /dev/null +++ b/config/gfs_varinfo.yaml @@ -0,0 +1,24 @@ +time_variant: + pgrb2.0p25.f000: + - variables: ["TMP"] + levels: ["2 m above ground"] + - variables: ["PRMSL"] + levels: ["mean sea level"] + - variables: ["UGRD", "VGRD"] + levels: ["10 m above ground"] + - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"] + levels: ["50 mb", "100 mb", "150 mb", "200 mb", "250 mb", "300 mb", "400 mb", "500 mb", "600 mb", "700 mb", "850 mb", "925 mb", "1000 mb"] + pgrb2.0p25.f006: + - variables: ["APCP"] + levels: ["surface"] + pgrb2b.0p25.f000: + - variables: ["SPFH", "VVEL", "UGRD", "VGRD", "HGT", "TMP"] + levels: ["125 mb", "175 mb", "225 mb", "775 mb", "825 mb", "875 mb"] +time_invariant: + pgrb2.0p25.f000: + - variables: ["HGT"] + levels: ["surface"] + first_time_step_only: True + - variables: ["LAND"] + levels: ["surface"] + first_time_step_only: True diff --git a/environment.yml b/environment.yml index 07d06cf..f1805e9 100644 --- a/environment.yml +++ b/environment.yml @@ -1,4 +1,4 @@ -name: graphcast +name: condaenv channels: - conda-forge dependencies: @@ -9,7 +9,8 @@ dependencies: - cartopy - jupyterlab - boto3 + - nceplibs-g2c + - grib2io - pip - pip: - - https://github.com/deepmind/graphcast/archive/master.zip - - flax + - git+https://github.com/noaa-emc/graphcast@aea0678cfbd7e866e4a1d364b456861d0a03954b diff --git a/pyproject.toml b/pyproject.toml index 316db50..c0b395c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,8 @@ dependencies = [ "scipy>=1.9.0", "trimesh>=3.15.0", "typing-extensions>=4.4.0", - "graphcast@git+https://github.com/google-deepmind/graphcast@6819a0f19796ca9c34a079855339b37134b8d930" + "grib2io>=2.5.4", + "graphcast@git+https://github.com/noaa-emc/graphcast@aea0678cfbd7e866e4a1d364b456861d0a03954b" ] [project.optional-dependencies] @@ -48,6 +49,7 @@ gpu = ["jax[cuda12]>=0.4.1"] dev = [ "pytest>=7.0.0", "pytest-cov>=4.0.0", + "pytest-mock>=3.10.0", "pre-commit>=4.3.0", "black>=24.4.2", "isort>=6.0.1", @@ -70,13 +72,9 @@ notebooks = [ requires = ["setuptools>=64.0.0", "setuptools-scm"] build-backend = "setuptools.build_meta" - [tool.setuptools] package-dir = {"" = "src"} packages = ["mlglobal"] [project.scripts] -ic-download = "mlglobal.cli.ic_download_main:main" -#ic-preprocess = "mlglobal.cli.gefs_preprocess_main:main" -#graphcast-run = "mlglobal.cli.graphcast_main:main" -#nc2grib2 = "mlglobal.cli.nc2grib2_main:main" +gen_ics = "mlglobal.cli.gen_ics_main:main" diff --git a/src/mlglobal/__init__.py b/src/mlglobal/__init__.py index 6e5f509..5127c24 100644 --- a/src/mlglobal/__init__.py +++ b/src/mlglobal/__init__.py @@ -14,3 +14,5 @@ __license__ = "CC-0" __all__ = ["__version__", "__author__", "__email__", "__license__"] + +from .logger import setup_logging diff --git a/src/mlglobal/cli/gen_ics_main.py b/src/mlglobal/cli/gen_ics_main.py new file mode 100644 index 0000000..ba73c3a --- /dev/null +++ b/src/mlglobal/cli/gen_ics_main.py @@ -0,0 +1,194 @@ +import argparse +import logging +import os +from datetime import datetime + +from mlglobal.ic_downloader import ICDownloader +from mlglobal.ic_processor import ICProcessor +from mlglobal.logger import setup_logging + +logger = logging.getLogger(__name__) + +_here = os.path.abspath(os.path.dirname(__file__)) +_top = os.path.abspath(os.path.join(os.path.abspath(_here), "../../../")) + +# Default bucket and root directory for each model # TODO: store this in a yaml or some config file +DEFAULTS = { + "gfs": { + "bucket_name": "noaa-gfs-bdp-pds", + "bucket_root_directory": "", + "comroot": "/lfs/h1/ops/prod/com/gfs/v16.3", + }, + "gefs": { + "bucket_name": "noaa-ncepdev-none-ca-ufs-cpldcld", + "bucket_root_directory": "Linlin.Cui/gefs_wcoss2", + "comroot": "/lfs/h1/ops/prod/com/gefs/v12.3", + }, +} + +# --- Helper strings for dynamic help text --- +BUCKET_HELP = ( + f"S3 bucket name. [default: {DEFAULTS['gfs']['bucket_name']} (for GFS), " + f"{DEFAULTS['gefs']['bucket_name']} (for GEFS)]" +) + +BUCKET_ROOT_DIR_HELP = ( + f"S3 bucket root directory. [default: {DEFAULTS['gfs']['bucket_root_directory']} (for GFS), " + f"{DEFAULTS['gefs']['bucket_root_directory']} (for GEFS)]" +) + +COMROOT_DIR_HELP = ( + f"Root directory. [default: {DEFAULTS['gfs']['comroot']} (for GFS), " + f"{DEFAULTS['gefs']['comroot']} (for GEFS)]" +) + + +def main(): + + parser = argparse.ArgumentParser( + description="Download IC data for GFS or GEFS", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + subparsers = parser.add_subparsers( + dest="model", + help="Model to download and process initial conditions for [GFS | GEFS]", + required=True, + ) + + def _common_args(inparser, model: str): + dict_in = DEFAULTS[model] + inparser.add_argument( + "--current-cycle", + help="Datetime to download and process initial conditions for in YYYYMMDDHH format", + type=str, + metavar="YYYYMMDDHH", + required=True, + ) + inparser.add_argument( + "--source", + help="Data source for getting model grib2 data", + type=str, + choices=["local", "s3"], + default="local", + required=False, + ) + inparser.add_argument( + "--target", + help="Target directory to store grib2 model data into", + type=str, + default="./input_data", + required=False, + ) + inparser.add_argument( + "--bucket-name", + help=BUCKET_HELP, + type=str, + default=dict_in["bucket_name"], + required=False, + ) + inparser.add_argument( + "--bucket-root-directory", + help=BUCKET_ROOT_DIR_HELP, + type=str, + default=dict_in["bucket_root_directory"], + required=False, + ) + inparser.add_argument( + "--comroot", + help=COMROOT_DIR_HELP, + type=str, + default=dict_in["comroot"], + required=False, + ) + inparser.add_argument( + "--num-levels", + help="Number of vertical levels to download from the model data", + type=int, + default=13, + required=False, + ) + inparser.add_argument( + "--varinfo-yaml", + help="Path to the varinfo YAML file", + type=str, + default=os.path.join(_top, "config", f"{model}_varinfo.yaml"), + required=False, + ) + inparser.add_argument( + "--output", + help="Name of the output NetCDF file", + type=str, + default="graphcast_ic.nc", + required=False, + ) + inparser.add_argument( + "--debug", + help="Set logging level to DEBUG", + action="store_true", + required=False, + ) + inparser.add_argument( + "--download-only", + help="Only download the data, do not process", + action="store_true", + required=False, + ) + + return inparser + + # GFS subparser + gfs_parser = subparsers.add_parser("gfs", help="Download GFS data") + gfs_parser = _common_args(gfs_parser, "gfs") + + # GEFS subparser + gefs_parser = subparsers.add_parser("gefs", help="Download GEFS ensemble data") + gefs_parser = _common_args(gefs_parser, "gefs") + gefs_members = ["c00"] + [f"p{str(i).zfill(2)}" for i in range(1, 31)] + gefs_parser.add_argument( + "--member", + help="Ensemble member", + type=str, + choices=gefs_members, + default="c00", + ) + + args = parser.parse_args() + + setup_logging(debug=args.debug) + + current_cycle = datetime.strptime(args.current_cycle, "%Y%m%d%H") + + logger.info(f"Downloading {args.model.upper()} data for cycle {args.current_cycle}") + downloader = ICDownloader( + current_cycle=current_cycle, + num_levels=args.num_levels, + member=None if args.model == "gfs" else args.member, + download_source=args.source, + local_directory=args.target, + bucket_name=args.bucket_name, + bucket_root_directory=args.bucket_root_directory, + root_directory=args.comroot, + ) + downloader.get_data() + + # TODO: validate the downloaded files before processing + + if args.download_only: + logger.info("Download-only flag set, skipping processing step.") + return + + logger.info(f"Processing {args.model.upper()} data for cycle {args.current_cycle}") + processor = ICProcessor( + current_cycle=current_cycle, + num_levels=args.num_levels, + member=None if args.model == "gfs" else args.member, + varinfo_file=args.varinfo_yaml, + data_dir=args.target, + output_netcdf=args.output, + ) + processor.process_files() + + +if __name__ == "__main__": + main() diff --git a/src/mlglobal/cli/ic_download_main.py b/src/mlglobal/cli/ic_download_main.py deleted file mode 100644 index 4df6886..0000000 --- a/src/mlglobal/cli/ic_download_main.py +++ /dev/null @@ -1,112 +0,0 @@ -import argparse -import os -from datetime import datetime - -from mlglobal.ic_downloader import ICDownloader - -# Default bucket and root directory for each mode # TODO: store this in a yaml or some config file -DEFAULTS = { - "gfs": { - "bucket_name": "noaa-gfs-bdp-pds", - "root_directory": "gdas", - }, - "gefs": { - "bucket_name": "noaa-ncepdev-none-ca-ufs-cpldcld", - "root_directory": "gefs", - }, -} - - -def main(): - parser = argparse.ArgumentParser(description="Download IC data for GFS or GEFS") - - subparsers = parser.add_subparsers( - dest="mode", help="System to download IC data for GFS or GEFS", required=True - ) - - def _common_args(inparser, dict_in): - inparser.add_argument( - "--start_date", - help="Start datetime", - type=str, - metavar="YYYYMMDDHH", - required=True, - ) - inparser.add_argument( - "--end_date", - help="End datetime", - type=str, - metavar="YYYYMMDDHH", - required=True, - ) - inparser.add_argument( - "--levels", - help="number of pressure levels", - type=int, - choices=[13, 37], - default=13, - required=False, - ) - inparser.add_argument( - "--source", - help="Data source", - type=str, - choices=["s3", "local"], - default="s3", - required=False, - ) - inparser.add_argument( - "--target", - help="Target directory to store raw data into", - type=str, - default=os.getcwd(), - required=False, - ) - inparser.add_argument( - "--bucket-name", - help="S3 bucket name", - type=str, - default=dict_in["bucket_name"], - required=False, - ) - inparser.add_argument( - "--root-directory", - help="Root directory", - type=str, - default=dict_in["root_directory"], - required=False, - ) - return inparser - - # GFS subparser - gfs_parser = subparsers.add_parser("gfs", help="Download GFS ensemble data") - gfs_parser = _common_args(gfs_parser, DEFAULTS["gfs"]) - - # GEFS subparser - gefs_parser = subparsers.add_parser("gefs", help="Download GEFS ensemble data") - gefs_parser = _common_args(gefs_parser, DEFAULTS["gefs"]) - gefs_parser.add_argument( - "--member", - help="Ensemble member", - type=int, - choices=list(range(0, 31)), - default=0, - ) - - args = parser.parse_args() - - downloader = ICDownloader( - mode=args.mode, - start_datetime=datetime.strptime(args.start_date, "%Y%m%d%H"), - end_datetime=datetime.strptime(args.end_date, "%Y%m%d%H"), - member=None if args.mode == "gfs" else args.member, - download_source=args.source, - download_directory=args.target, - bucket_name=args.bucket_name, - root_directory=args.root_directory, - ) - downloader.download() - - -if __name__ == "__main__": - main() diff --git a/src/mlglobal/ic_downloader.py b/src/mlglobal/ic_downloader.py index cfdcef3..982725c 100644 --- a/src/mlglobal/ic_downloader.py +++ b/src/mlglobal/ic_downloader.py @@ -1,154 +1,258 @@ -import glob import os import shutil -from datetime import datetime, timedelta +from datetime import timedelta +from logging import getLogger +logger = getLogger(__name__) -class FileFormats: - def __init__(self, mode, num_levels=13): - FILE_FORMATS = {"gfs": self.gfs_file_formats, "gefs": self.gefs_file_formats} +class FileLookup: + def __init__(self, current_cycle, num_levels=13, member=None): + + self.current_cycle = current_cycle self.num_levels = num_levels - self.file_formats = FILE_FORMATS[mode]() + self.member = member # GEFS member values are c00, p01, p02, ..., p30 + + # Look back 6 hours for precip and 2 time-level data + self.current_cycle_m6h = self.current_cycle - timedelta(hours=6) - def gfs_file_formats(self): - # List of file formats to download - if self.num_levels == 13: - file_formats = ["pgrb2.0p25.f000", "pgrb2.0p25.f006"] # , '0p25.f001' + if self.member is not None: + self.get_file_info = self._gefs_file_info else: - file_formats = [ - "pgrb2.0p25.f000", - "pgrb2b.0p25.f000", - "pgrb2.0p25.f006", - ] # , '0p25.f001' + self.get_file_info = self._gfs_file_info - return file_formats + def _gfs_file_info(self): - def gefs_file_formats(self): + # Template for GFS files + template = f"gfs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/gfs.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}" # noqa: F541 - # List of file formats to download - if self.num_levels == 13: - file_formats = ["pgrb2.0p25.f000", "pgrb2s.0p25.f000"] # , '0p25.f001' - else: - file_formats = [ - "pgrb2.0p25.f000", - "pgrb2b.0p25.f000", - "pgrb2.0p25.f006", - ] # , '0p25.f001' + # From current cycle + pgrb2_0p25_f000 = template.format( + cycle=self.current_cycle, fspec="pgrb2.0p25", fhour=0 + ) + if self.num_levels == 37: # Need additional pgrb2b file for 37 level data + pgrb2b_0p25_f000 = template.format( + cycle=self.current_cycle, fspec="pgrb2b.0p25", fhour=0 + ) - return file_formats + # From current cycle - 6 hours + pgrb2_0p25_f000_m6 = template.format( + cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=0 + ) + pgrb2_0p25_f006_m6 = template.format( + cycle=self.current_cycle_m6h, fspec="pgrb2.0p25", fhour=6 + ) + + # Create a flat list of files valid at current and current-6h cycles + file_dict = {} + file_dict[self.current_cycle] = [pgrb2_0p25_f000, pgrb2_0p25_f006_m6] + if self.num_levels == 37: + file_dict[self.current_cycle].append(pgrb2b_0p25_f000) + file_dict[self.current_cycle_m6h] = [pgrb2_0p25_f000_m6] + + return file_dict + + def _gefs_file_info(self): + + # Template for GEFS files + template = f"gefs.{{cycle:%Y%m%d}}/{{cycle:%H}}/atmos/{{fspec_dir}}/ge{{member}}.t{{cycle:%H}}z.{{fspec}}.f{{fhour:03d}}" # noqa: F541 + + # From current cycle + pgrb2_0p25_f000 = template.format( + cycle=self.current_cycle, + fspec_dir="pgrb2p25", + fspec="pgrb2.0p25", + fhour=0, + member=self.member, + ) + pgrb2s_0p25_f000 = template.format( + cycle=self.current_cycle, + fspec_dir="pgrb2sp25", + fspec="pgrb2s.0p25", + fhour=0, + member=self.member, + ) + + # From current cycle - 6 hours + pgrb2_0p25_f000_m6 = template.format( + cycle=self.current_cycle_m6h, + fspec_dir="pgrb2p25", + fspec="pgrb2.0p25", + fhour=0, + member=self.member, + ) + pgrb2s_0p25_f000_m6 = template.format( + cycle=self.current_cycle_m6h, + fspec_dir="pgrb2sp25", + fspec="pgrb2s.0p25", + fhour=0, + member=self.member, + ) + + file_dict = {} + file_dict[self.current_cycle] = [pgrb2_0p25_f000, pgrb2s_0p25_f000] + file_dict[self.current_cycle_m6h] = [pgrb2_0p25_f000_m6, pgrb2s_0p25_f000_m6] + + return file_dict class ICDownloader: def __init__( self, - mode, - start_datetime, - end_datetime, + current_cycle, + num_levels=13, member=None, - num_pressure_levels=13, - download_source="s3", - download_directory=None, + download_source="local", + local_directory="./data", bucket_name=None, + bucket_root_directory=None, root_directory=None, ): - self.mode = mode - self.start_datetime = start_datetime - self.end_datetime = end_datetime + self.current_cycle = current_cycle + self.num_levels = num_levels self.member = member - self.num_pressure_levels = num_pressure_levels self.download_source = download_source - self.download_directory = download_directory + self.local_directory = local_directory self.bucket_name = bucket_name + self.bucket_root_directory = bucket_root_directory self.root_directory = root_directory - ff = FileFormats(mode, num_levels=self.num_pressure_levels) - self.file_formats = ff.file_formats + # Generate the lookup dictionary + lookup = FileLookup( + self.current_cycle, member=self.member, num_levels=self.num_levels + ) + self.file_dict = lookup.get_file_info() + + # Flatten the file_dict to a simple list of files + file_list = [] + for files in self.file_dict.values(): + file_list.extend(files) + self.file_list = file_list + + if self.download_source in ["s3"]: + aws_profile = os.environ.get("AWS_PROFILE", "default") + self.s3 = self.get_s3_client_by_bucket_type( + self.bucket_name, profile_name=aws_profile + ) - self.s3 = self.init_s3_client() if self.download_source == "s3" else None + os.makedirs(self.local_directory, exist_ok=True) @staticmethod - def init_s3_client(): + def get_s3_client_by_bucket_type(bucket_name, profile_name="default"): + """ + Initializes and returns a boto3 S3 client for a given bucket. + + The function first attempts to get a public client. If that fails, it + assumes the bucket is private and creates a client using the specified + AWS profile. + + Args: + bucket_name (str): The name of the S3 bucket. + profile_name (str): The AWS profile to use for private buckets. + Defaults to 'default'. + + Returns: + boto3.client: A configured S3 client. + """ try: import boto3 from botocore import UNSIGNED from botocore.config import Config + from botocore.exceptions import ClientError except ImportError as ee: raise ImportError( "boto3 and botocore are required for S3 operations." ) from ee + # 1. Try to get a client configured for a public (unsigned) bucket + s3 = boto3.client("s3", config=Config(signature_version=UNSIGNED)) try: - # Try to create S3 client using profile method - profile_name = os.environ.get("AWS_PROFILE", "default") - session = boto3.Session(profile_name=profile_name) - current_credentials = session.get_credentials().get_frozen_credentials() - s3 = session.client( - "s3", - aws_access_key_id=current_credentials.access_key, - aws_secret_access_key=current_credentials.secret_key, + # Check if the bucket can be accessed publicly without authentication + s3.head_bucket(Bucket=bucket_name) + logger.info( + f"Bucket '{bucket_name}' is public. Returning an unsigned client." ) - except Exception as e1: - print(f"Failed to create S3 client with profile method: {e1}") - try: - # Try to create S3 client using unsigned method - s3 = boto3.client("s3", config=Config(signature_version=UNSIGNED)) - except Exception as e2: - print(f"Failed to create S3 client with unsigned method: {e2}") - raise RuntimeError( - "Failed to create S3 client with unsigned method." - ) from e2 - - return s3 - - def get_s3_specs(self, ymd, hh, file_format): + return s3 + except ClientError as ee: + error_code = ee.response["Error"]["Code"] + # 2. If it's a 403 Forbidden, the bucket is likely private. + if error_code in ("403", "404"): + logger.warning( + f"Bucket '{bucket_name}' is not public. Returning client with profile '{profile_name}'." + ) - if self.mode == "gefs": + # Create a session with the specified profile + session = boto3.Session(profile_name=profile_name) + current_credentials = session.get_credentials().get_frozen_credentials() + s3 = session.client( + "s3", + aws_access_key_id=current_credentials.access_key, + aws_secret_access_key=current_credentials.secret_key, + ) + return s3 + else: + # Handle other errors, such as a non-existent bucket + logger.error(f"Error accessing bucket '{bucket_name}': {ee}") + return None - s3_prefix = ( - f"Linlin.Cui/gefs_wcoss2/{self.root_directory}.{ymd}/{hh}/atmos/" - ) - s3_file_format = f"{self.member:02d}.t{hh}z.{file_format}" + @staticmethod + def get_s3_objects(s3, bucket_name: str, prefix: str) -> list: + """Retrieve a list of objects from an S3 bucket with a specific prefix. + This method handles pagination to ensure all objects are retrieved. + By default, S3 returns up to 1000 objects per request. - elif self.mode == "gfs": + Parameters + ---------- + s3 : boto3.client + The S3 client to use for the operation. + bucket_name : str + The name of the S3 bucket. + prefix : str + The prefix to filter the objects. - if file_format == "pgrb2.0p25.f006": - # get prefix for precip from the previous cycle - # Convert ymd and hh to datetime object - datetime_obj = datetime.strptime(ymd + hh, "%Y%m%d%H") + Returns + ------- + list + A list of S3 objects that match the prefix. + """ - # Get the datetime 6 hours before - datetime_before = datetime_obj - timedelta(hours=6) + objects = [] + continuation_token = None - # Get the date string and time string from datetime objects - ymd_precip = datetime_before.strftime("%Y%m%d") - hh_precip = datetime_before.strftime("%H") + while True: + list_kwargs = { + "Bucket": bucket_name, + "Prefix": prefix, + "MaxKeys": 1000, # Explicitly set MaxKeys, though it's default + } + if continuation_token: + list_kwargs["ContinuationToken"] = continuation_token - # Construct the S3 prefix for the directory - s3_prefix = f"{self.root_directory}.{ymd_precip}/{hh_precip}/" + response = s3.list_objects_v2(**list_kwargs) - else: + if "Contents" in response: + objects.extend(response["Contents"]) - s3_prefix = f"{self.root_directory}.{ymd}/{hh}/" + if not response.get("IsTruncated"): + # No more objects to retrieve, or less than 1000 objects in total + break - s3_file_format = file_format + continuation_token = response.get("NextContinuationToken") + if not continuation_token: + # Should not happen if 'IsTruncated' is True, but as a safeguard + break - return s3_prefix, s3_file_format + return objects - def get_data_from_s3( - self, path_prefix: str, file_format: str, local_directory: str - ) -> None: + def get_data_from_s3(self, file_list: list, local_directory: str) -> None: """ - Downloads files with a specific format from an S3 bucket to a local directory. - + Download files from S3 bucket to a local directory. Parameters ---------- - prefix : str - The prefix (folder path) in the S3 bucket to filter objects. - file_format : str - The file extension or format to filter files (e.g., '.csv', '.json'). + file_list : list + A list of files that need to be downloaded local_directory : str The local directory path where the downloaded files will be saved. @@ -157,126 +261,81 @@ def get_data_from_s3( None This function does not return anything. Files are downloaded as a side effect. - Notes - ----- - Only files ending with the specified `file_format` will be downloaded. + Raises + ------ + Exception + If the file download operation fails. """ - objects = self.s3.list_objects_v2(Bucket=self.bucket_name, Prefix=path_prefix) - for obj in objects.get("Contents", []): - obj_key = obj["Key"] - if obj_key.endswith(f"{file_format}"): - local_file_path = os.path.join( - local_directory, os.path.basename(obj_key) + logger.info(f"Downloading files from S3 bucket: {self.bucket_name}") + logger.info(f"Downloading files to {local_directory}") + + for file_name in file_list: + local_file_path = os.path.join(local_directory, os.path.basename(file_name)) + if os.path.exists(local_file_path): + logger.warning(f"File already exists, skipping: {file_name}") + continue + file_name_in_bucket = ( + self.bucket_root_directory + "/" + file_name + if self.bucket_root_directory + else file_name + ) + print(f"Downloading {file_name_in_bucket} to {local_file_path}") + try: + self.s3.download_file( + self.bucket_name, file_name_in_bucket, local_file_path ) - if not os.path.exists(local_file_path): - self.s3.download_file(self.bucket_name, obj_key, local_file_path) - print(f"Downloaded {obj_key} to {local_file_path}") - else: - print(f"File {local_file_path} already exists, skipping download.") + logger.info(f"Downloaded: {file_name} -> {local_directory}") + except Exception as ee: + logger.error(f"Error downloading {file_name}: {ee}") - def get_local_specs(self, ymd: str, hh: str, file_format: str) -> tuple[str, str]: - """ - Get the local specifications for a given date, time, and file format. + def get_data_from_local(self, file_list: list, local_directory: str) -> None: + """Copy files from a local directory to another local directory. Parameters ---------- - ymd : str - The date string in the format 'YYYYMMDD'. - hh : str - The time string in the format 'HH'. - file_format : str - The file format string (e.g., 'pgrb2.0p25.f006'). + file_list : list + A list of files that need to be copied. + local_directory : str + The local directory path where the copied files will be saved. Returns ------- - tuple[str, str] - A tuple containing the local path prefix and the local file format. - """ - # TODO: elevate the hard-coded paths to class constructor (or above) - if self.mode == "gefs": - gefs_com_dir = "/lfs/h2/emc/da/noscrub/rahul.mahajan/mldata" # For testing on Acorn - gefs_com_dir = "/lfs/h2/emc/ptmp/jun.wang" # NCO does not mirror all of GEFS data to dev in RT - gefs_com_dir = "/lfs/h1/ops/prod/com/gefs/v12.3" - fprefix = file_format.split('.')[0] - local_prefix = f"{gefs_com_dir}/gefs.{ymd}/{hh}/atmos/{fprefix}p25" - local_file_format = f"{self.member:02d}.t{hh}z.{file_format}" - - elif self.mode == "gfs": - - # TODO: elevate the hard-coded paths to constructor - gfs_com_dir = "/lfs/h2/emc/da/noscrub/rahul.mahajan/mldata" # For testing on Acorn - gfs_com_dir = "/lfs/h1/ops/prod/com/gfs/v16.3" - if file_format == "pgrb2.0p25.f006": - # get prefix for precip from the previous cycle - # Convert ymd and hh to datetime object - datetime_obj = datetime.strptime(ymd + hh, "%Y%m%d%H") - - # Get the datetime 6 hours before - datetime_before = datetime_obj - timedelta(hours=6) - - # Get the date string and time string from datetime objects - ymd_precip = datetime_before.strftime("%Y%m%d") - hh_precip = datetime_before.strftime("%H") - - # Construct the S3 prefix for the directory - local_prefix = f"{gfs_com_dir}/gfs.{ymd_precip}/{hh_precip}/atmos" - local_file_format = f"gfs.t{hh_precip}z.{file_format}" - - else: - - local_prefix = f"{gfs_com_dir}/gfs.{ymd}/{hh}/atmos" - local_file_format = f"gfs.t{hh}z.{file_format}" - - return local_prefix, local_file_format - - def get_data_from_local( - self, path_prefix: str, file_format: str, local_directory: str - ) -> None: - - file_objects = glob.glob(f"{path_prefix}/*") - for obj_key in file_objects: - if obj_key.endswith(f"{file_format}"): + None + This function does not return anything. Files are copied as a side effect. - # Define the local file path - local_file_path = os.path.join( - local_directory, os.path.basename(obj_key) - ) + Raises + ------ + OSError + If the file copy operation fails. + """ - # Copy data to the local path - try: - shutil.copy2(obj_key, local_file_path) - print(f"Copied: {obj_key} -> {local_directory}") - except OSError: - raise OSError(f"Unable to copy {obj_key} to {local_directory}") + logger.info(f"Copying files from directory: {self.root_directory}") + logger.info(f"Copying files to {local_directory}") + + for file_name in file_list: + local_file_path = os.path.join(local_directory, os.path.basename(file_name)) + if os.path.exists(local_file_path): + logger.warning(f"File already exists, skipping: {file_name}") + continue + file_name_in_remote = ( + self.root_directory + "/" + file_name + if self.root_directory + else file_name + ) + try: + shutil.copy2(file_name_in_remote, local_file_path) + logger.info(f"Copied: {file_name} -> {local_directory}") + except OSError: + logger.error(f"Unable to copy {file_name} to {local_directory}") + raise OSError(f"Unable to copy {file_name} to {local_directory}") return - def download(self, loop_interval=6): - - _SPECS_MAP = {"s3": self.get_s3_specs, "local": self.get_local_specs} - _GET_DATA_MAP = {"s3": self.get_data_from_s3, "local": self.get_data_from_local} - - interval_dt = timedelta(hours=loop_interval) - - # Loop through the intervals - current_datetime = self.start_datetime - while current_datetime <= self.end_datetime: - ymd = current_datetime.strftime("%Y%m%d") - hh = current_datetime.strftime("%H") - - # Define the local directory path where the file will be saved - local_directory = os.path.join(self.download_directory, ymd, hh) - - # Create the local directory if it doesn't exist - os.makedirs(local_directory, exist_ok=True) - - # Loop over file formats and download data - for file_format in self.file_formats: - prefix, fformat = _SPECS_MAP[self.download_source](ymd, hh, file_format) - _GET_DATA_MAP[self.download_source](prefix, fformat, local_directory) + def get_data(self) -> None: - # Move to the next interval - current_datetime += interval_dt + GET_DATA_MAP = {"s3": self.get_data_from_s3, "local": self.get_data_from_local} - print("Download completed.") + logger.info("Starting download...") + GET_DATA_MAP[self.download_source](self.file_list, self.local_directory) + logger.info("Download completed.") diff --git a/src/mlglobal/ic_processor.py b/src/mlglobal/ic_processor.py new file mode 100644 index 0000000..25f3e78 --- /dev/null +++ b/src/mlglobal/ic_processor.py @@ -0,0 +1,351 @@ +import fnmatch +import logging +import os +from datetime import datetime + +import grib2io +import numpy as np +import xarray as xr +import yaml + +from mlglobal.ic_downloader import FileLookup + +logger = logging.getLogger(__name__) + + +class ICProcessor: + """Class to prepare initial conditions for graphcast from downloaded data files.""" + + def __init__( + self, + current_cycle: datetime, + varinfo_file: str, + member: str | None = None, + num_levels: int = 13, + data_dir: str = "./data", + output_netcdf: str = "./ic.nc", + ) -> None: + + self.current_cycle = current_cycle + self.varinfo = self.load_yaml(varinfo_file) + self.member = member + self.num_levels = num_levels + self.data_dir = data_dir + self.output_netcdf = output_netcdf + + # Generate the lookup dictionary + lookup = FileLookup( + self.current_cycle, member=self.member, num_levels=self.num_levels + ) + self.file_dict = lookup.get_file_info() + + @staticmethod + def load_yaml(yaml_file: str) -> dict: + """Load a YAML file and return the contents as a dictionary. + Parameters + ---------- + yaml_file : str + Path to the YAML file. + Returns + ------- + dict + Dictionary containing the contents of the YAML file. + Raises + ------ + FileNotFoundError + If the YAML file does not exist. + yaml.YAMLError + If there is an error parsing the YAML file. + """ + + try: + with open(yaml_file, "r") as fh: + yaml_dict = yaml.safe_load(fh) + except FileNotFoundError: + logger.error(f"YAML file {yaml_file} not found.") + raise FileNotFoundError(f"YAML file {yaml_file} not found.") + except yaml.YAMLError as e: + logger.error(f"Error parsing YAML file {yaml_file}: {e}") + raise yaml.YAMLError(f"Error parsing YAML file {yaml_file}: {e}") + + return yaml_dict + + def get_matching_pattern(self, file, patterns, fatal=True): + + filename = os.path.basename(file) + + matched = False + pattern = None + # Check each pattern to see if it matches the filename + for pp in patterns: + if fnmatch.fnmatch(file, os.path.join(self.data_dir, "*" + pp)): + pattern = pp + logger.debug(f"Matched pattern: {pattern} in file: {filename}") + matched = True + break + + if not matched and fatal: + logger.error(f"No pattern matched for file: {filename}") + raise FileNotFoundError(f"No pattern matched for file: {filename}") + + return pattern + + def process_files(self): + """ + Process input files and extract relevant variables. + This function reads GRIB2 files, extracts specified variables, + and saves them into a NetCDF file. + Returns + ------- + None + + Raises + ------ + FileNotFoundError + If any of the input files do not exist. + ValueError + If there is an error processing the files. + IOError + If there is an error saving the NetCDF file. + """ + + logger.info("Starting to process files...") + + file_patterns = list(self.varinfo["time_variant"].keys()) + mergeDSs = [] + for cycle in self.file_dict.keys(): + + mergeDAs = [] + + # Files valid at the same time, so we can merge them along the time dimension + for file in self.file_dict[cycle]: + + logger.info(f"Processing {file=}") + + file = os.path.join(self.data_dir, os.path.basename(file)) + + pattern = self.get_matching_pattern(file, file_patterns) + + gribfh = grib2io.open(file, "r") + for var_dict in self.varinfo["time_variant"][pattern]: + for var_name in var_dict["variables"]: + logger.info( + f"Extracting variable: {var_name} at levels: {var_dict['levels']}" + ) + da = self.get_dataarray(gribfh, var_name, var_dict["levels"]) + mergeDAs.append(da) + gribfh.close() + + ds = xr.merge(mergeDAs, compat="no_conflicts") + mergeDSs.append(ds) + ds.close() + + # Concatenate along the time dimension + ds = xr.concat(mergeDSs, dim="time") + + logger.info("Processing static variables...") + + # Now handle the static variables + file_patterns = list(self.varinfo["time_invariant"].keys()) + mergeDAs = [] + for file in self.file_dict[self.current_cycle]: + + logger.info(f"Processing {file=}") + + file = os.path.join(self.data_dir, os.path.basename(file)) + + pattern = self.get_matching_pattern(file, file_patterns, fatal=False) + if pattern is None: # No pattern matched, so skip this file + continue + + gribfh = grib2io.open(file, "r") + for var_dict in self.varinfo["time_invariant"][pattern]: + for var_name in var_dict["variables"]: + logger.info( + f"Extracting variable: {var_name} at levels: {var_dict['levels']}" + ) + da = self.get_dataarray(gribfh, var_name, var_dict["levels"]) + mergeDAs.append(da) + gribfh.close() + + ds_static = xr.merge(mergeDAs, compat="no_conflicts") + ds = xr.merge([ds, ds_static], compat="no_conflicts") + ds_static.close() + + # Rename variables to match graphcast naming conventions + ds = self.rename_dsvars(ds) + + # Add datetime coordinate + ds = ds.assign_coords(datetime=ds.time) + + # Adjust time values relative to the first time step + ds["time"] = ds["time"] - ds["time"][0] + + # Expand dimensions + ds = ds.expand_dims(dim="batch") + ds["datetime"] = ds["datetime"].expand_dims(dim="batch") + + # Squeeze dimensions + ds["geopotential_at_surface"] = ds["geopotential_at_surface"].squeeze("batch") + ds["land_sea_mask"] = ds["land_sea_mask"].squeeze("batch") + + # Update geopotential unit to m2/s2 by multiplying 9.80665 + ds["geopotential_at_surface"] = ds["geopotential_at_surface"] * 9.80665 + ds["geopotential"] = ds["geopotential"] * 9.80665 + + # For GEFS, if total_precipitation_6hr is missing, create it with zeros + if "total_precipitation_6hr" not in ds: + logger.warning( + "total_precipitation_6hr variable not found. Creating it with zeros." + ) + ds["total_precipitation_6hr"] = xr.zeros_like(ds["2m_temperature"]) + else: + # Update total_precipitation_6hr unit to (m) from (kg/m^2) by dividing it by 1000kg/m³ + ds["total_precipitation_6hr"] = ds["total_precipitation_6hr"] / 1000.0 + + # Save to NetCDF + logger.info(f"Saving processed data to {self.output_netcdf}") + os.makedirs(os.path.dirname(self.output_netcdf), exist_ok=True) + ds.to_netcdf(self.output_netcdf) + ds.close() + + return + + @staticmethod + def rename_dsvars(ds: xr.Dataset, rename_dict: dict = None) -> xr.Dataset: + """Rename dataset variables to match graphcast naming conventions.""" + logger.info("Renaming dataset variables to match graphcast naming conventions.") + + # Default rename dictionary # TODO: move to yaml? + if rename_dict is None: + rename_dict = { + "LAND_surface": "land_sea_mask", + "HGT_surface": "geopotential_at_surface", + "PRMSL_meansealevel": "mean_sea_level_pressure", + "TMP_2maboveground": "2m_temperature", + "UGRD_10maboveground": "10m_u_component_of_wind", + "VGRD_10maboveground": "10m_v_component_of_wind", + "HGT": "geopotential", + "TMP": "temperature", + "SPFH": "specific_humidity", + "VVEL": "vertical_velocity", + "UGRD": "u_component_of_wind", + "VGRD": "v_component_of_wind", + } + if "APCP_surface" in ds: + rename_dict["APCP_surface"] = "total_precipitation_6hr" + logger.debug(f"Using rename_dict: {rename_dict}") + + ds = ds.rename(rename_dict) + + return ds + + @staticmethod + def get_dataarray(gh, var_name: str, levels: list | int | str) -> xr.Dataset: + """Get a data array from the GRIB file. + + Parameters + ---------- + gh : Handle of opened grib2io file + var_name : str + Name of the variable to extract + levels : list | int | str + Levels to extract (can be a single level or a list of levels) + + Returns + ------- + xr.Dataset + A dataset containing the extracted data + """ + + if isinstance(levels, list) and len(levels) > 1: + return ICProcessor._get_dataarray_3d(gh, var_name, levels) + else: + return ICProcessor._get_dataarray_2d( + gh, var_name, levels[0] if isinstance(levels, list) else levels + ) + + @staticmethod + def _get_dataarray_2d(gh, var_name, level): + + logger.info(f"Getting 2D data for variable: {var_name} at level: {level}") + + msg = gh.select(shortName=var_name, level=level)[0] + + # create a netcdf dataset using the matching grib message + lats, lons = msg.latlons() + lats = lats[:, 0] + lons = lons[0, :] + + # check latitude range, graphcast needs [-90, 90] + reverse_lat = False + if lats[0] > 0: + reverse_lat = True + + steps = msg.validDate + data = msg.data + if reverse_lat: + data = data[::-1, :] + lats = lats[::-1] + + if len(data.shape) == 2: + var_name2 = f'{var_name}_{"".join(level.split())}' + da = xr.Dataset( + data_vars={var_name2: (["lat", "lon"], data.astype("float32"))}, + coords={ + "lon": lons.astype("float32"), + "lat": lats.astype("float32"), + "time": steps, + }, + ) + elif len(data.shape) == 3: + logger.warning("3D data found in 2D function") + da = xr.Dataset( + data_vars={var_name: (["level", "lat", "lon"], data.astype("float32"))}, + coords={ + "lon": lons.astype("float32"), + "lat": lats.astype("float32"), + "level": np.array(level).astype("int32"), + "time": steps, + }, + ) + + return da + + @staticmethod + def _get_dataarray_3d(gh, var_name, levels): + + logger.info(f"Getting 3D data for variable: {var_name} at levels: {levels}") + data, vlevs = [], [] + for ii, level in enumerate(levels): + msg = gh.select(shortName=var_name, level=level)[0] + + if ii == 0: + lats, lons = msg.latlons() + lats = lats[:, 0] + lons = lons[0, :] + + # check latitude range, graphcast needs [-90, 90] + reverse_lat = False + if lats[0] > 0: + reverse_lat = True + steps = msg.validDate + data.append(msg.data) + vlevs.append(int(level.split(" ")[0])) + + data = np.array(data) + if reverse_lat: + data = data[:, ::-1, :] + lats = lats[::-1] + + da = xr.Dataset( + data_vars={var_name: (["level", "lat", "lon"], data.astype("float32"))}, + coords={ + "lon": lons.astype("float32"), + "lat": lats.astype("float32"), + "level": np.array(vlevs).astype("int32"), + "time": steps, + }, + ) + + return da diff --git a/src/mlglobal/logger.py b/src/mlglobal/logger.py new file mode 100644 index 0000000..9554c95 --- /dev/null +++ b/src/mlglobal/logger.py @@ -0,0 +1,37 @@ +import logging + +""" +Logger setup module for the mlglobal project. + +This module configures the logging settings for the entire project using Python's built-in +logging module. It sets up a default logging format, date format, and log level, and provides +a logger instance named "mlglobal" for use throughout the project. + +Functions +--------- +setup_logger(level=logging.INFO, fmt="[%(asctime)s] %(levelname)8s - %(message)s", datefmt="%Y-%m-%dT%H:%M:%S") + Configures the logging system with the specified log level, format, and date format. + +Attributes +---------- +logger : logging.Logger +""" + + +log = logging.getLogger() + + +def setup_logging( + fmt: str = "[%(asctime)s] %(levelname)8s - %(message)s", + datefmt: str = "%Y-%m-%dT%H:%M:%S", + debug: bool = False, +): + + logger = logging.getLogger() + for handler in logger.handlers: + logger.removeHandler(handler) + + level = logging.DEBUG if debug else logging.INFO + kwargs: dict = {"datefmt": datefmt, "format": fmt, "level": level} + + logging.basicConfig(**kwargs) diff --git a/tests/test_file_lookup.py b/tests/test_file_lookup.py new file mode 100644 index 0000000..1ad71b0 --- /dev/null +++ b/tests/test_file_lookup.py @@ -0,0 +1,144 @@ +from datetime import datetime, timedelta + +import pytest + +from mlglobal.ic_downloader import FileLookup + + +class TestFileLookup: + """Test cases for the FileLookup class.""" + + def test_init_gfs(self): + """Test FileLookup initialization for GFS mode.""" + current_cycle = datetime(2023, 1, 1, 12, 0, 0) + lookup = FileLookup(current_cycle, num_levels=13, member=None) + + assert lookup.current_cycle == current_cycle + assert lookup.num_levels == 13 + assert lookup.member is None + assert lookup.current_cycle_m6h == current_cycle - timedelta(hours=6) + assert lookup.get_file_info == lookup._gfs_file_info + + def test_init_gefs(self): + """Test FileLookup initialization for GEFS mode.""" + current_cycle = datetime(2023, 1, 1, 12, 0, 0) + lookup = FileLookup(current_cycle, num_levels=13, member="c00") + + assert lookup.current_cycle == current_cycle + assert lookup.num_levels == 13 + assert lookup.member == "c00" + assert lookup.current_cycle_m6h == current_cycle - timedelta(hours=6) + assert lookup.get_file_info == lookup._gefs_file_info + + def test_gfs_file_info_13_levels(self): + """Test GFS file info generation for 13 levels.""" + current_cycle = datetime(2023, 1, 1, 12, 0, 0) + lookup = FileLookup(current_cycle, num_levels=13, member=None) + + file_dict = lookup._gfs_file_info() + + # Check structure + assert len(file_dict) == 2 + assert current_cycle in file_dict + assert lookup.current_cycle_m6h in file_dict + + # Check current cycle files + current_files = file_dict[current_cycle] + assert len(current_files) == 2 + assert "gfs.20230101/12/atmos/gfs.t12z.pgrb2.0p25.f000" in current_files[0] + assert "gfs.20230101/06/atmos/gfs.t06z.pgrb2.0p25.f006" in current_files[1] + + # Check previous cycle files + prev_files = file_dict[lookup.current_cycle_m6h] + assert len(prev_files) == 1 + assert "gfs.20230101/06/atmos/gfs.t06z.pgrb2.0p25.f000" in prev_files[0] + + def test_gfs_file_info_37_levels(self): + """Test GFS file info generation for 37 levels.""" + current_cycle = datetime(2023, 1, 1, 12, 0, 0) + lookup = FileLookup(current_cycle, num_levels=37, member=None) + + file_dict = lookup._gfs_file_info() + + # Check current cycle has additional pgrb2b file + current_files = file_dict[current_cycle] + assert len(current_files) == 3 + assert any("pgrb2b.0p25" in f for f in current_files) + + def test_gefs_file_info(self): + """Test GEFS file info generation.""" + current_cycle = datetime(2023, 1, 1, 12, 0, 0) + lookup = FileLookup(current_cycle, num_levels=13, member="c00") + + file_dict = lookup._gefs_file_info() + + # Check structure + assert len(file_dict) == 2 + assert current_cycle in file_dict + assert lookup.current_cycle_m6h in file_dict + + # Check current cycle files + current_files = file_dict[current_cycle] + assert len(current_files) == 2 + assert "gec00.t12z.pgrb2.0p25.f000" in current_files[0] + assert "gec00.t12z.pgrb2s.0p25.f000" in current_files[1] + + # Check previous cycle files + prev_files = file_dict[lookup.current_cycle_m6h] + assert len(prev_files) == 2 + assert "gec00.t06z.pgrb2.0p25.f000" in prev_files[0] + assert "gec00.t06z.pgrb2s.0p25.f000" in prev_files[1] + + @pytest.mark.parametrize("num_levels,expected_current_files", [ + (13, 2), # For 13 levels: pgrb2.0p25 files only + (37, 3), # For 37 levels: includes pgrb2b.0p25 file + ]) + def test_gfs_file_count_by_levels(self, num_levels, expected_current_files): + """Test that GFS mode generates correct number of files based on levels.""" + current_cycle = datetime(2023, 1, 1, 12, 0, 0) + lookup = FileLookup(current_cycle, num_levels=num_levels, member=None) + file_dict = lookup._gfs_file_info() + + # Check current cycle files + current_files = file_dict[current_cycle] + assert len(current_files) == expected_current_files + + @pytest.mark.parametrize("member", ["c00", "p01", "p15", "p30"]) + def test_gefs_members(self, member): + """Test GEFS mode with different member values.""" + current_cycle = datetime(2023, 1, 1, 12, 0, 0) + lookup = FileLookup(current_cycle, num_levels=13, member=member) + file_dict = lookup._gefs_file_info() + + # Verify member appears in file paths + for files in file_dict.values(): + for file_path in files: + assert f"ge{member}" in file_path + + def test_current_cycle_m6h_calculation(self): + """Test that current_cycle_m6h is calculated correctly.""" + current_cycle = datetime(2023, 1, 1, 18, 0, 0) + lookup = FileLookup(current_cycle, num_levels=13, member=None) + + expected_m6h = datetime(2023, 1, 1, 12, 0, 0) + assert lookup.current_cycle_m6h == expected_m6h + + def test_file_info_delegation_gfs(self): + """Test that get_file_info correctly delegates to GFS method.""" + current_cycle = datetime(2023, 1, 1, 12, 0, 0) + lookup = FileLookup(current_cycle, num_levels=13, member=None) + + gfs_result = lookup._gfs_file_info() + delegated_result = lookup.get_file_info() + + assert gfs_result == delegated_result + + def test_file_info_delegation_gefs(self): + """Test that get_file_info correctly delegates to GEFS method.""" + current_cycle = datetime(2023, 1, 1, 12, 0, 0) + lookup = FileLookup(current_cycle, num_levels=13, member="c00") + + gefs_result = lookup._gefs_file_info() + delegated_result = lookup.get_file_info() + + assert gefs_result == delegated_result diff --git a/tests/test_ic_downloader.py b/tests/test_ic_downloader.py new file mode 100644 index 0000000..3290ec5 --- /dev/null +++ b/tests/test_ic_downloader.py @@ -0,0 +1,429 @@ +import os +import shutil +import tempfile +from datetime import datetime + +import pytest +from botocore.exceptions import ClientError + +from mlglobal.ic_downloader import ICDownloader + + +class TestICDownloader: + """Test cases for the ICDownloader class.""" + + def setup_method(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + self.current_cycle = datetime(2023, 1, 1, 12, 0, 0) + + def teardown_method(self): + """Clean up test fixtures.""" + if os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir) + + def test_init_gfs(self): + """Test ICDownloader initialization for GFS mode.""" + downloader = ICDownloader( + current_cycle=self.current_cycle, + num_levels=13, + member=None, + download_source="local", + local_directory=self.test_dir, + root_directory="/test/root" + ) + + assert downloader.current_cycle == self.current_cycle + assert downloader.num_levels == 13 + assert downloader.member is None + assert downloader.download_source == "local" + assert downloader.local_directory == self.test_dir + assert downloader.root_directory == "/test/root" + assert len(downloader.file_list) > 0 + assert os.path.exists(self.test_dir) + + def test_init_gefs(self): + """Test ICDownloader initialization for GEFS mode.""" + downloader = ICDownloader( + current_cycle=self.current_cycle, + num_levels=13, + member="c00", + download_source="local", + local_directory=self.test_dir + ) + + assert downloader.member == "c00" + assert len(downloader.file_list) > 0 + + def test_init_s3(self, mocker): + """Test ICDownloader initialization for S3 source.""" + mock_s3_client = mocker.MagicMock() + + # Mock the get_s3_client_by_bucket_type method since boto3 is imported inside it + mock_get_s3_client = mocker.patch.object( + ICDownloader, 'get_s3_client_by_bucket_type', + return_value=mock_s3_client + ) + + # Mock the environment variable + mocker.patch.dict(os.environ, {'AWS_PROFILE': 'test-profile'}) + + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="s3", + bucket_name="test-bucket" + ) + + assert downloader.s3 == mock_s3_client + mock_get_s3_client.assert_called_once_with("test-bucket", profile_name="test-profile") + + def test_get_s3_client_by_bucket_type_public(self, mocker): + """Test S3 client creation for public bucket.""" + mock_client = mocker.MagicMock() + mock_boto3_client = mocker.patch('boto3.client') + mock_boto3_client.return_value = mock_client + mock_client.head_bucket.return_value = None # Success + + result = ICDownloader.get_s3_client_by_bucket_type("public-bucket") + + assert result == mock_client + mock_client.head_bucket.assert_called_once_with(Bucket="public-bucket") + + def test_get_s3_client_by_bucket_type_private(self, mocker): + """Test S3 client creation for private bucket.""" + # Mock public client failing + mock_public_client = mocker.MagicMock() + mock_boto3_client = mocker.patch('boto3.client') + mock_boto3_client.return_value = mock_public_client + mock_public_client.head_bucket.side_effect = ClientError( + {"Error": {"Code": "403"}}, "head_bucket" + ) + + # Mock private client + mock_private_client = mocker.MagicMock() + mock_session_instance = mocker.MagicMock() + mock_session = mocker.patch('boto3.Session') + mock_session.return_value = mock_session_instance + mock_session_instance.client.return_value = mock_private_client + mock_session_instance.get_credentials.return_value.get_frozen_credentials.return_value = mocker.MagicMock( + access_key="test-key", secret_key="test-secret" + ) + + result = ICDownloader.get_s3_client_by_bucket_type("private-bucket", "test-profile") + + assert result == mock_private_client + mock_session.assert_called_once_with(profile_name="test-profile") + + def test_get_s3_client_by_bucket_type_error(self, mocker): + """Test S3 client creation with error.""" + mock_client = mocker.MagicMock() + mock_boto3_client = mocker.patch('boto3.client') + mock_boto3_client.return_value = mock_client + mock_client.head_bucket.side_effect = ClientError( + {"Error": {"Code": "NoSuchBucket"}}, "head_bucket" + ) + + result = ICDownloader.get_s3_client_by_bucket_type("nonexistent-bucket") + + assert result is None + + def test_get_s3_objects(self, mocker): + """Test S3 objects retrieval with pagination.""" + mock_s3 = mocker.MagicMock() + + # Mock paginated response + mock_s3.list_objects_v2.side_effect = [ + { + "Contents": [{"Key": "file1"}, {"Key": "file2"}], + "IsTruncated": True, + "NextContinuationToken": "token123" + }, + { + "Contents": [{"Key": "file3"}], + "IsTruncated": False + } + ] + + result = ICDownloader.get_s3_objects(mock_s3, "test-bucket", "prefix/") + + assert len(result) == 3 + assert result[0]["Key"] == "file1" + assert result[1]["Key"] == "file2" + assert result[2]["Key"] == "file3" + + # Verify pagination calls + assert mock_s3.list_objects_v2.call_count == 2 + + def test_get_s3_objects_no_contents(self, mocker): + """Test S3 objects retrieval with no contents.""" + mock_s3 = mocker.MagicMock() + mock_s3.list_objects_v2.return_value = {"IsTruncated": False} + + result = ICDownloader.get_s3_objects(mock_s3, "test-bucket", "prefix/") + + assert len(result) == 0 + + def test_get_data_from_s3(self, mocker): + """Test downloading data from S3.""" + mock_s3 = mocker.MagicMock() + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="s3", + local_directory=self.test_dir, + bucket_name="test-bucket" + ) + downloader.s3 = mock_s3 + + file_list = ["path/to/file1.grib2", "path/to/file2.grib2"] + + downloader.get_data_from_s3(file_list, self.test_dir) + + # Verify download calls + assert mock_s3.download_file.call_count == 2 + mock_s3.download_file.assert_any_call( + "test-bucket", "path/to/file1.grib2", + os.path.join(self.test_dir, "file1.grib2") + ) + mock_s3.download_file.assert_any_call( + "test-bucket", "path/to/file2.grib2", + os.path.join(self.test_dir, "file2.grib2") + ) + + def test_get_data_from_s3_with_bucket_root(self, mocker): + """Test downloading data from S3 with bucket root directory.""" + mock_s3 = mocker.MagicMock() + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="s3", + local_directory=self.test_dir, + bucket_name="test-bucket", + bucket_root_directory="root/dir" + ) + downloader.s3 = mock_s3 + + file_list = ["file1.grib2"] + + downloader.get_data_from_s3(file_list, self.test_dir) + + mock_s3.download_file.assert_called_once_with( + "test-bucket", "root/dir/file1.grib2", + os.path.join(self.test_dir, "file1.grib2") + ) + + def test_get_data_from_s3_file_exists(self, mocker): + """Test S3 download when local file already exists.""" + mock_s3 = mocker.MagicMock() + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="s3", + local_directory=self.test_dir, + bucket_name="test-bucket" + ) + downloader.s3 = mock_s3 + + # Create existing file + existing_file = os.path.join(self.test_dir, "file1.grib2") + with open(existing_file, 'w') as f: + f.write("test") + + file_list = ["path/to/file1.grib2"] + + downloader.get_data_from_s3(file_list, self.test_dir) + + # Should not attempt download + mock_s3.download_file.assert_not_called() + + def test_get_data_from_s3_download_error(self, mocker): + """Test S3 download with error.""" + mock_s3 = mocker.MagicMock() + mock_s3.download_file.side_effect = Exception("Download failed") + + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="s3", + local_directory=self.test_dir, + bucket_name="test-bucket" + ) + downloader.s3 = mock_s3 + + file_list = ["path/to/file1.grib2"] + + # Should not raise exception, just log error + downloader.get_data_from_s3(file_list, self.test_dir) + + mock_s3.download_file.assert_called_once() + + def test_get_data_from_local(self): + """Test copying data from local directory.""" + # Create source directory and files + source_dir = tempfile.mkdtemp() + try: + source_file1 = os.path.join(source_dir, "file1.grib2") + source_file2 = os.path.join(source_dir, "file2.grib2") + + with open(source_file1, 'w') as f: + f.write("test data 1") + with open(source_file2, 'w') as f: + f.write("test data 2") + + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="local", + local_directory=self.test_dir, + root_directory=source_dir + ) + + file_list = ["file1.grib2", "file2.grib2"] + + downloader.get_data_from_local(file_list, self.test_dir) + + # Verify files were copied + assert os.path.exists(os.path.join(self.test_dir, "file1.grib2")) + assert os.path.exists(os.path.join(self.test_dir, "file2.grib2")) + + # Verify content + with open(os.path.join(self.test_dir, "file1.grib2"), 'r') as f: + assert f.read() == "test data 1" + + finally: + shutil.rmtree(source_dir) + + def test_get_data_from_local_file_exists(self): + """Test local copy when destination file already exists.""" + source_dir = tempfile.mkdtemp() + try: + source_file = os.path.join(source_dir, "file1.grib2") + with open(source_file, 'w') as f: + f.write("source data") + + # Create existing destination file + dest_file = os.path.join(self.test_dir, "file1.grib2") + with open(dest_file, 'w') as f: + f.write("existing data") + + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="local", + local_directory=self.test_dir, + root_directory=source_dir + ) + + file_list = ["file1.grib2"] + + downloader.get_data_from_local(file_list, self.test_dir) + + # File should not be overwritten + with open(dest_file, 'r') as f: + assert f.read() == "existing data" + + finally: + shutil.rmtree(source_dir) + + def test_get_data_from_local_copy_error(self): + """Test local copy with error.""" + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="local", + local_directory=self.test_dir, + root_directory="/nonexistent/path" + ) + + file_list = ["file1.grib2"] + + with pytest.raises(OSError, match="Unable to copy"): + downloader.get_data_from_local(file_list, self.test_dir) + + def test_get_data_s3(self, mocker): + """Test get_data method with S3 source.""" + mock_s3 = mocker.MagicMock() + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="s3", + local_directory=self.test_dir, + bucket_name="test-bucket" + ) + downloader.s3 = mock_s3 + + mock_get_s3 = mocker.patch.object(downloader, 'get_data_from_s3') + downloader.get_data() + mock_get_s3.assert_called_once_with(downloader.file_list, self.test_dir) + + def test_get_data_local(self, mocker): + """Test get_data method with local source.""" + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="local", + local_directory=self.test_dir, + root_directory="/test/root" + ) + + mock_get_local = mocker.patch.object(downloader, 'get_data_from_local') + downloader.get_data() + mock_get_local.assert_called_once_with(downloader.file_list, self.test_dir) + + def test_get_data_invalid_source(self): + """Test get_data method with invalid source.""" + downloader = ICDownloader( + current_cycle=self.current_cycle, + download_source="invalid", + local_directory=self.test_dir + ) + + with pytest.raises(KeyError): + downloader.get_data() + + +@pytest.fixture +def sample_datetime(): + """Fixture providing a sample datetime for testing.""" + return datetime(2023, 1, 1, 12, 0, 0) + + +@pytest.fixture +def temp_directory(): + """Fixture providing a temporary directory for testing.""" + temp_dir = tempfile.mkdtemp() + yield temp_dir + shutil.rmtree(temp_dir) + + +class TestIntegration: + """Integration tests for ICDownloader.""" + + def test_full_workflow_local(self, sample_datetime, temp_directory, mocker): + """Test full workflow with local data source.""" + # Create mock source data + source_dir = os.path.join(temp_directory, "source") + os.makedirs(source_dir, exist_ok=True) + + # Create some mock files that match the expected patterns + test_files = [ + "gfs.20230101/12/atmos/gfs.t12z.pgrb2.0p25.f000", + "gfs.20230101/12/atmos/gfs.t12z.pgrb2.0p25.f006", + ] + + for file_path in test_files: + full_path = os.path.join(source_dir, file_path) + os.makedirs(os.path.dirname(full_path), exist_ok=True) + with open(full_path, 'w') as f: + f.write(f"mock data for {file_path}") + + dest_dir = os.path.join(temp_directory, "dest") + + downloader = ICDownloader( + current_cycle=sample_datetime, + num_levels=13, + member=None, + download_source="local", + local_directory=dest_dir, + root_directory=source_dir + ) + + # Mock the file lookup to return our test files + downloader.file_list = test_files + + # This should work without errors + downloader.get_data() + + # Verify destination directory was created + assert os.path.exists(dest_dir)